ベイズニューラルネットワーク — 重みの不確実性を考慮した深層学習

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import matplotlib.pyplot as plt

class BayesianLinear(nn.Module):
    """ベイズ線形層 (Bayes by Backprop)"""
    def __init__(self, in_features, out_features, prior_sigma=1.0):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features

        # 重みの変分パラメータ
        self.weight_mu = nn.Parameter(
            torch.Tensor(out_features, in_features).uniform_(-0.2, 0.2))
        self.weight_rho = nn.Parameter(
            torch.Tensor(out_features, in_features).uniform_(-5, -4))

        # バイアスの変分パラメータ
        self.bias_mu = nn.Parameter(
            torch.Tensor(out_features).uniform_(-0.2, 0.2))
        self.bias_rho = nn.Parameter(
            torch.Tensor(out_features).uniform_(-5, -4))

        # 事前分布のパラメータ
        self.prior_sigma = prior_sigma
        self.log_prior = 0
        self.log_variational_posterior = 0

    def forward(self, x):
        # ソフトプラス変換で正の標準偏差を保証
        weight_sigma = torch.log1p(torch.exp(self.weight_rho))
        bias_sigma = torch.log1p(torch.exp(self.bias_rho))

        # 再パラメータ化トリック
        weight_epsilon = torch.randn_like(weight_sigma)
        bias_epsilon = torch.randn_like(bias_sigma)

        weight = self.weight_mu + weight_sigma * weight_epsilon
        bias = self.bias_mu + bias_sigma * bias_epsilon

        # 対数事前確率
        self.log_prior = (
            -0.5 * torch.sum(weight**2) / self.prior_sigma**2
            - 0.5 * torch.sum(bias**2) / self.prior_sigma**2
        )

        # 対数変分事後確率
        self.log_variational_posterior = (
            -0.5 * torch.sum(((weight - self.weight_mu) / weight_sigma)**2)
            - torch.sum(torch.log(weight_sigma))
            - 0.5 * torch.sum(((bias - self.bias_mu) / bias_sigma)**2)
            - torch.sum(torch.log(bias_sigma))
        )

        return F.linear(x, weight, bias)

class BayesianNN(nn.Module):
    """ベイズニューラルネットワーク"""
    def __init__(self, input_dim, hidden_dim, output_dim):
        super().__init__()
        self.blinear1 = BayesianLinear(input_dim, hidden_dim)
        self.blinear2 = BayesianLinear(hidden_dim, hidden_dim)
        self.blinear3 = BayesianLinear(hidden_dim, output_dim)

    def forward(self, x):
        x = F.relu(self.blinear1(x))
        x = F.relu(self.blinear2(x))
        x = self.blinear3(x)
        return x

    def log_prior(self):
        return (self.blinear1.log_prior
                + self.blinear2.log_prior
                + self.blinear3.log_prior)

    def log_variational_posterior(self):
        return (self.blinear1.log_variational_posterior
                + self.blinear2.log_variational_posterior
                + self.blinear3.log_variational_posterior)

    def elbo_loss(self, x, y, n_samples, n_data):
        """ELBO損失（最小化する）"""
        total_log_lik = 0
        total_log_prior = 0
        total_log_var_post = 0

        for _ in range(n_samples):
            output = self.forward(x)
            log_lik = -0.5 * torch.sum((y - output)**2)
            total_log_lik += log_lik
            total_log_prior += self.log_prior()
            total_log_var_post += self.log_variational_posterior()

        # モンテカルロ推定の平均
        total_log_lik /= n_samples
        total_log_prior /= n_samples
        total_log_var_post /= n_samples

        # ミニバッチスケーリング
        kl = (total_log_var_post - total_log_prior) / n_data
        nll = -total_log_lik / n_data

        return nll + kl

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import matplotlib.pyplot as plt

class BayesianLinear(nn.Module):
    def __init__(self, in_features, out_features, prior_sigma=1.0):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.weight_mu = nn.Parameter(
            torch.Tensor(out_features, in_features).uniform_(-0.2, 0.2))
        self.weight_rho = nn.Parameter(
            torch.Tensor(out_features, in_features).uniform_(-5, -4))
        self.bias_mu = nn.Parameter(
            torch.Tensor(out_features).uniform_(-0.2, 0.2))
        self.bias_rho = nn.Parameter(
            torch.Tensor(out_features).uniform_(-5, -4))
        self.prior_sigma = prior_sigma
        self.log_prior = 0
        self.log_variational_posterior = 0

    def forward(self, x):
        weight_sigma = torch.log1p(torch.exp(self.weight_rho))
        bias_sigma = torch.log1p(torch.exp(self.bias_rho))
        weight = self.weight_mu + weight_sigma * torch.randn_like(weight_sigma)
        bias = self.bias_mu + bias_sigma * torch.randn_like(bias_sigma)
        self.log_prior = (-0.5 * torch.sum(weight**2) / self.prior_sigma**2
                          - 0.5 * torch.sum(bias**2) / self.prior_sigma**2)
        self.log_variational_posterior = (
            -0.5 * torch.sum(((weight - self.weight_mu) / weight_sigma)**2)
            - torch.sum(torch.log(weight_sigma))
            - 0.5 * torch.sum(((bias - self.bias_mu) / bias_sigma)**2)
            - torch.sum(torch.log(bias_sigma)))
        return F.linear(x, weight, bias)

class BayesianNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.l1 = BayesianLinear(1, 50)
        self.l2 = BayesianLinear(50, 50)
        self.l3 = BayesianLinear(50, 1)

    def forward(self, x):
        x = F.relu(self.l1(x))
        x = F.relu(self.l2(x))
        return self.l3(x)

    def log_prior(self):
        return self.l1.log_prior + self.l2.log_prior + self.l3.log_prior

    def log_variational_posterior(self):
        return (self.l1.log_variational_posterior
                + self.l2.log_variational_posterior
                + self.l3.log_variational_posterior)

# データ生成
torch.manual_seed(42)
np.random.seed(42)

def true_function(x):
    return np.sin(2 * x) + 0.3 * x

N = 40
x_train = np.sort(np.random.uniform(-2, 2, N))
y_train = true_function(x_train) + np.random.normal(0, 0.2, N)

x_tensor = torch.FloatTensor(x_train).reshape(-1, 1)
y_tensor = torch.FloatTensor(y_train).reshape(-1, 1)

# 学習
model = BayesianNN()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

n_epochs = 2000
losses = []

for epoch in range(n_epochs):
    model.train()
    optimizer.zero_grad()

    # ELBO損失
    n_mc = 3  # モンテカルロサンプル数
    total_loss = 0
    for _ in range(n_mc):
        output = model(x_tensor)
        nll = 0.5 * torch.sum((y_tensor - output)**2) / N
        kl = (model.log_variational_posterior() - model.log_prior()) / N
        total_loss += nll + 0.1 * kl  # KLの重みを調整

    loss = total_loss / n_mc
    loss.backward()
    optimizer.step()
    losses.append(loss.item())

# 予測（複数回フォワードパスで不確実性を推定）
model.eval()
x_test = np.linspace(-4, 4, 200)
x_test_tensor = torch.FloatTensor(x_test).reshape(-1, 1)

T = 100  # サンプル数
predictions = np.zeros((T, len(x_test)))

with torch.no_grad():
    for t in range(T):
        predictions[t] = model(x_test_tensor).numpy().flatten()

pred_mean = predictions.mean(axis=0)
pred_std = predictions.std(axis=0)

# 可視化
fig, axes = plt.subplots(1, 2, figsize=(14, 5))

# 左: BNNの予測と不確実性
ax = axes[0]
ax.plot(x_test, pred_mean, 'b-', linewidth=2, label='BNN予測平均')
ax.fill_between(x_test, pred_mean - 2*pred_std, pred_mean + 2*pred_std,
                alpha=0.2, color='blue', label='95%信頼区間')
ax.plot(x_test, true_function(x_test), 'k--', linewidth=1.5,
        label='真の関数')
ax.scatter(x_train, y_train, c='red', s=30, zorder=5, label='訓練データ')

# 訓練データの範囲を示す
ax.axvspan(-2, 2, alpha=0.05, color='green')
ax.set_xlabel('x')
ax.set_ylabel('y')
ax.set_title('ベイズニューラルネットワークの予測')
ax.legend(fontsize=9, loc='upper left')
ax.set_ylim(-4, 4)

# 右: 複数のサンプルによる予測
ax = axes[1]
for t in range(20):
    ax.plot(x_test, predictions[t], 'b-', alpha=0.1, linewidth=0.8)
ax.plot(x_test, true_function(x_test), 'k--', linewidth=1.5,
        label='真の関数')
ax.scatter(x_train, y_train, c='red', s=30, zorder=5, label='訓練データ')
ax.axvspan(-2, 2, alpha=0.05, color='green')
ax.set_xlabel('x')
ax.set_ylabel('y')
ax.set_title('重みサンプルごとの予測（20回）')
ax.legend(fontsize=10)
ax.set_ylim(-4, 4)

plt.tight_layout()
plt.savefig("bnn_regression.png", dpi=150, bbox_inches="tight")
plt.show()

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import matplotlib.pyplot as plt

class BayesianLinear(nn.Module):
    def __init__(self, in_features, out_features, prior_sigma=1.0):
        super().__init__()
        self.weight_mu = nn.Parameter(
            torch.Tensor(out_features, in_features).uniform_(-0.2, 0.2))
        self.weight_rho = nn.Parameter(
            torch.Tensor(out_features, in_features).uniform_(-5, -4))
        self.bias_mu = nn.Parameter(
            torch.Tensor(out_features).uniform_(-0.2, 0.2))
        self.bias_rho = nn.Parameter(
            torch.Tensor(out_features).uniform_(-5, -4))
        self.prior_sigma = prior_sigma

    def forward(self, x):
        weight_sigma = torch.log1p(torch.exp(self.weight_rho))
        bias_sigma = torch.log1p(torch.exp(self.bias_rho))
        weight = self.weight_mu + weight_sigma * torch.randn_like(weight_sigma)
        bias = self.bias_mu + bias_sigma * torch.randn_like(bias_sigma)
        return F.linear(x, weight, bias)

class HeteroscedasticBNN(nn.Module):
    """偶然的不確実性も学習するBNN"""
    def __init__(self):
        super().__init__()
        self.shared1 = BayesianLinear(1, 50)
        self.shared2 = BayesianLinear(50, 50)
        self.mean_head = BayesianLinear(50, 1)
        self.var_head = BayesianLinear(50, 1)  # 対数分散を出力

    def forward(self, x):
        h = F.relu(self.shared1(x))
        h = F.relu(self.shared2(h))
        mean = self.mean_head(h)
        log_var = self.var_head(h)
        return mean, log_var

# 不均一ノイズを持つデータ
torch.manual_seed(42)
np.random.seed(42)

N = 80
x_train = np.sort(np.random.uniform(-3, 3, N))

# ノイズの大きさがxに依存（不均一分散）
noise_level = 0.1 + 0.3 * np.abs(x_train)
y_train = np.sin(x_train) + noise_level * np.random.randn(N)

x_tensor = torch.FloatTensor(x_train).reshape(-1, 1)
y_tensor = torch.FloatTensor(y_train).reshape(-1, 1)

# 学習
model = HeteroscedasticBNN()
optimizer = torch.optim.Adam(model.parameters(), lr=0.005)

for epoch in range(3000):
    model.train()
    optimizer.zero_grad()

    mean_pred, log_var_pred = model(x_tensor)
    var_pred = torch.exp(log_var_pred)

    # 異分散正規分布の負の対数尤度
    nll = 0.5 * torch.mean(log_var_pred + (y_tensor - mean_pred)**2 / var_pred)

    nll.backward()
    optimizer.step()

# 予測
model.eval()
x_test = np.linspace(-5, 5, 200)
x_test_tensor = torch.FloatTensor(x_test).reshape(-1, 1)

T = 200
means = np.zeros((T, len(x_test)))
log_vars = np.zeros((T, len(x_test)))

with torch.no_grad():
    for t in range(T):
        m, lv = model(x_test_tensor)
        means[t] = m.numpy().flatten()
        log_vars[t] = lv.numpy().flatten()

# 不確実性の分解
pred_mean = means.mean(axis=0)
aleatoric = np.exp(log_vars).mean(axis=0)  # E[sigma^2(x)]
epistemic = means.var(axis=0)               # Var[mu(x)]
total = aleatoric + epistemic

# 可視化
fig, axes = plt.subplots(1, 2, figsize=(14, 5))

# 左: 予測と2種類の不確実性
ax = axes[0]
ax.plot(x_test, pred_mean, 'b-', linewidth=2, label='予測平均')
ax.fill_between(x_test,
                pred_mean - 2*np.sqrt(epistemic),
                pred_mean + 2*np.sqrt(epistemic),
                alpha=0.4, color='orange', label='認識的不確実性')
ax.fill_between(x_test,
                pred_mean - 2*np.sqrt(total),
                pred_mean + 2*np.sqrt(total),
                alpha=0.2, color='blue', label='全不確実性')
ax.plot(x_test, np.sin(x_test), 'k--', linewidth=1.5, label='真の関数')
ax.scatter(x_train, y_train, c='gray', s=15, alpha=0.5)
ax.axvspan(-3, 3, alpha=0.05, color='green')
ax.set_xlabel('x')
ax.set_ylabel('y')
ax.set_title('予測と不確実性の分解')
ax.legend(fontsize=9)
ax.set_ylim(-4, 4)

# 右: 各不確実性のプロファイル
ax = axes[1]
ax.plot(x_test, np.sqrt(aleatoric), 'r-', linewidth=2, label='偶然的')
ax.plot(x_test, np.sqrt(epistemic), 'orange', linewidth=2, label='認識的')
ax.plot(x_test, np.sqrt(total), 'b-', linewidth=2, label='合計')
ax.axvspan(-3, 3, alpha=0.05, color='green', label='訓練データ範囲')
ax.set_xlabel('x')
ax.set_ylabel('標準偏差')
ax.set_title('不確実性のプロファイル')
ax.legend(fontsize=10)

plt.tight_layout()
plt.savefig("bnn_uncertainty_decomposition.png", dpi=150, bbox_inches="tight")
plt.show()

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import matplotlib.pyplot as plt

class BayesianLinear(nn.Module):
    def __init__(self, in_features, out_features):
        super().__init__()
        self.weight_mu = nn.Parameter(
            torch.Tensor(out_features, in_features).uniform_(-0.2, 0.2))
        self.weight_rho = nn.Parameter(
            torch.Tensor(out_features, in_features).uniform_(-5, -4))
        self.bias_mu = nn.Parameter(
            torch.Tensor(out_features).uniform_(-0.2, 0.2))
        self.bias_rho = nn.Parameter(
            torch.Tensor(out_features).uniform_(-5, -4))

    def forward(self, x):
        ws = torch.log1p(torch.exp(self.weight_rho))
        bs = torch.log1p(torch.exp(self.bias_rho))
        w = self.weight_mu + ws * torch.randn_like(ws)
        b = self.bias_mu + bs * torch.randn_like(bs)
        return F.linear(x, w, b)

class BNNClassifier(nn.Module):
    def __init__(self, input_dim, hidden_dim, n_classes):
        super().__init__()
        self.l1 = BayesianLinear(input_dim, hidden_dim)
        self.l2 = BayesianLinear(hidden_dim, hidden_dim)
        self.l3 = BayesianLinear(hidden_dim, n_classes)

    def forward(self, x):
        x = F.relu(self.l1(x))
        x = F.relu(self.l2(x))
        return self.l3(x)

# 2次元の2クラス分類データを生成
np.random.seed(42)
torch.manual_seed(42)

n_per_class = 100
# クラス0: 中心 (-1, -1)
x0 = np.random.randn(n_per_class, 2) * 0.5 + np.array([-1, -1])
# クラス1: 中心 (1, 1)
x1 = np.random.randn(n_per_class, 2) * 0.5 + np.array([1, 1])

X = np.vstack([x0, x1]).astype(np.float32)
y = np.concatenate([np.zeros(n_per_class), np.ones(n_per_class)]).astype(np.int64)

X_tensor = torch.FloatTensor(X)
y_tensor = torch.LongTensor(y)

# 学習
model = BNNClassifier(2, 50, 2)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

for epoch in range(1000):
    model.train()
    optimizer.zero_grad()

    n_mc = 3
    loss = 0
    for _ in range(n_mc):
        logits = model(X_tensor)
        loss += F.cross_entropy(logits, y_tensor)
    loss = loss / n_mc
    loss.backward()
    optimizer.step()

# OOD検出のための予測不確実性マップ
model.eval()
grid_range = np.linspace(-4, 4, 100)
xx, yy = np.meshgrid(grid_range, grid_range)
grid_points = np.column_stack([xx.ravel(), yy.ravel()]).astype(np.float32)
grid_tensor = torch.FloatTensor(grid_points)

T = 50
all_probs = np.zeros((T, len(grid_points), 2))

with torch.no_grad():
    for t in range(T):
        logits = model(grid_tensor)
        probs = F.softmax(logits, dim=1).numpy()
        all_probs[t] = probs

# 不確実性指標: 予測エントロピー
mean_probs = all_probs.mean(axis=0)
predictive_entropy = -np.sum(mean_probs * np.log(mean_probs + 1e-10), axis=1)

# 認識的不確実性（相互情報量）
entropy_per_sample = -np.sum(all_probs * np.log(all_probs + 1e-10), axis=2)
expected_entropy = entropy_per_sample.mean(axis=0)
mutual_information = predictive_entropy - expected_entropy

# 可視化
fig, axes = plt.subplots(1, 3, figsize=(18, 5))

# 左: 予測クラス
ax = axes[0]
pred_class = mean_probs[:, 1].reshape(xx.shape)
c = ax.contourf(xx, yy, pred_class, levels=20, cmap='RdBu_r', alpha=0.8)
plt.colorbar(c, ax=ax, label='P(class=1)')
ax.scatter(x0[:, 0], x0[:, 1], c='blue', s=15, alpha=0.6, label='Class 0')
ax.scatter(x1[:, 0], x1[:, 1], c='red', s=15, alpha=0.6, label='Class 1')
ax.set_xlabel('x1')
ax.set_ylabel('x2')
ax.set_title('予測クラス確率')
ax.legend()

# 中央: 予測エントロピー（全不確実性）
ax = axes[1]
entropy_map = predictive_entropy.reshape(xx.shape)
c = ax.contourf(xx, yy, entropy_map, levels=20, cmap='hot_r', alpha=0.8)
plt.colorbar(c, ax=ax, label='エントロピー')
ax.scatter(x0[:, 0], x0[:, 1], c='blue', s=15, alpha=0.6)
ax.scatter(x1[:, 0], x1[:, 1], c='red', s=15, alpha=0.6)
ax.set_xlabel('x1')
ax.set_ylabel('x2')
ax.set_title('予測エントロピー（全不確実性）')

# 右: 認識的不確実性（相互情報量）
ax = axes[2]
mi_map = mutual_information.reshape(xx.shape)
c = ax.contourf(xx, yy, mi_map, levels=20, cmap='hot_r', alpha=0.8)
plt.colorbar(c, ax=ax, label='相互情報量')
ax.scatter(x0[:, 0], x0[:, 1], c='blue', s=15, alpha=0.6)
ax.scatter(x1[:, 0], x1[:, 1], c='red', s=15, alpha=0.6)
ax.set_xlabel('x1')
ax.set_ylabel('x2')
ax.set_title('認識的不確実性（OOD検出）')

plt.tight_layout()
plt.savefig("bnn_ood_detection.png", dpi=150, bbox_inches="tight")
plt.show()