Bayes-by-Backprop Tutorial¶

# ----------------------------------------------------------------------------
# Core imports + compatibility utilities
# - We import Bayes-by-Backprop primitives from the deepuq library:
#     * `deepuq.methods.vi_elbo_step`
#     * `deepuq.methods.vi.BayesianLinear`
# - `vi_elbo_step_compat` lets this notebook run against both old and new
#   deepuq APIs (`n_batches` vs `num_batches`, optional `mc_samples`).
# ----------------------------------------------------------------------------
import random
import inspect
import copy
import math

import numpy as np
import torch
import torch.nn as nn
from torch import optim
from torch.utils.data import DataLoader, TensorDataset
import matplotlib.pyplot as plt

# DeepUQ VI API (primary objective function for Bayes-by-Backprop training)
from deepuq.methods import vi_elbo_step

# DeepUQ Bayesian layer used to build the custom model in this notebook
from deepuq.methods.vi import BayesianLinear

VI_ELBO_SIG = inspect.signature(vi_elbo_step)
print(f"vi_elbo_step signature: {VI_ELBO_SIG}")

DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')


def set_seed(seed: int) -> None:
    """Set RNG seeds for reproducibility across numpy/torch/cuda."""
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def vi_elbo_step_compat(
    model,
    x,
    y,
    *,
    num_batches,
    criterion,
    kl_weight,
    mc_samples=1,
):
    """
    Compatibility wrapper around deepuq `vi_elbo_step`.

    Why this exists:
    - Some versions expect `num_batches`.
    - Older versions expect `n_batches` and may not expose `mc_samples`.
    """
    params = VI_ELBO_SIG.parameters

    # New API supports `num_batches` directly.
    if 'num_batches' in params:
        kwargs = {
            'num_batches': num_batches,
            'criterion': criterion,
            'kl_weight': kl_weight,
        }
        if 'mc_samples' in params:
            kwargs['mc_samples'] = mc_samples
        return vi_elbo_step(model, x, y, **kwargs)

    # Legacy API path: if mc_samples == 1, pass through once.
    if mc_samples == 1:
        return vi_elbo_step(
            model,
            x,
            y,
            n_batches=num_batches,
            criterion=criterion,
            kl_weight=kl_weight,
        )

    # Legacy API path with manual Monte-Carlo averaging.
    loss_acc = 0.0
    nll_acc = 0.0
    kl_acc = 0.0
    for _ in range(mc_samples):
        loss, nll, kl = vi_elbo_step(
            model,
            x,
            y,
            n_batches=num_batches,
            criterion=criterion,
            kl_weight=kl_weight,
        )
        loss_acc = loss_acc + loss
        nll_acc = nll_acc + nll
        kl_acc = kl_acc + kl

    loss = loss_acc / float(mc_samples)
    nll = nll_acc / float(mc_samples)
    kl = kl_acc / float(mc_samples)
    return loss, nll, kl

NOTEBOOK_QUICK_MODE = os.environ.get('DEEPUQ_NOTEBOOK_QUICK', '0') == '1'
print(f"Notebook quick mode: {NOTEBOOK_QUICK_MODE}")

# ----------------------------------------------------------------------------
# Quick data sanity plot
# - Shows noisy train/val samples plus the true underlying function.
# - Use this to verify the dataset shape before training.
# ----------------------------------------------------------------------------
plt.figure(figsize=(7, 4))
plt.scatter(x_train_raw.squeeze().numpy(), y_train_raw.squeeze().numpy(), s=12, alpha=0.45, label='Train noisy')
plt.scatter(x_val_raw.squeeze().numpy(), y_val_raw.squeeze().numpy(), s=12, alpha=0.35, label='Val noisy')

x_vis = torch.linspace(-3.0, 3.0, 300).unsqueeze(-1)
y_vis = target_fn(x_vis)
plt.plot(x_vis.squeeze().numpy(), y_vis.squeeze().numpy(), color='black', linewidth=2, label='True function')

plt.xlabel('x')
plt.ylabel('y')
plt.title('Synthetic regression data')
plt.legend(frameon=False)
plt.tight_layout()
plt.show()

# ----------------------------------------------------------------------------
# Part 1 training loop
# Stage 1: deterministic warmup (`sample=False`) for stable fit initialization.
# Stage 2: stochastic VI optimization with ELBO from deepuq `vi_elbo_step`.
# We keep 1000 VI epochs as requested and restore the best val-RMSE checkpoint.
# ----------------------------------------------------------------------------
set_seed(13)

model = ComplexBayesianRegressor(
    input_dim=x_train.shape[1],
    width=MODEL_WIDTH,
    depth=MODEL_DEPTH,
    prior_sigma=MODEL_PRIOR_SIGMA,
).to(DEVICE)
initialize_bayesian_posteriors(model, rho_init=-4.5)

# MSE handles regression data-fit (NLL proxy in normalized target space).
criterion = nn.MSELoss(reduction='mean')

# AdamW with mild weight decay and gradient clipping improves stability.
optimizer = optim.AdamW(model.parameters(), lr=4e-4, weight_decay=1e-5)

# Training policy
if NOTEBOOK_QUICK_MODE:
    warmup_epochs = 10
    vi_epochs = 80
    eval_mc_elbo = 2
    eval_mc_rmse = 8
else:
    warmup_epochs = 60
    vi_epochs = 1000
    eval_mc_elbo = 6
    eval_mc_rmse = 32
kl_weight = 1e-4
num_batches = len(train_loader)
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=vi_epochs, eta_min=1e-4)

# History dictionary used by downstream plots/diagnostics.
history = {
    'epoch': [],
    'stage': [],
    'train_elbo': [],
    'val_elbo': [],
    'train_nll': [],
    'val_nll': [],
    'train_kl': [],
    'val_kl': [],
    'train_rmse': [],
    'val_rmse': [],
}

# Track best validation checkpoint; final model is restored to this state.
best_val_rmse = float('inf')
best_epoch = 0
best_state = copy.deepcopy(model.state_dict())

# ------------------------------------------------------------------------
# Stage 1: deterministic warmup (no weight sampling in forward pass)
# ------------------------------------------------------------------------
for epoch_idx in range(warmup_epochs):
    model.train()
    for x_batch, y_batch in train_loader:
        x_batch = x_batch.to(DEVICE)
        y_batch = y_batch.to(DEVICE)

        optimizer.zero_grad(set_to_none=True)
        pred = model(x_batch, sample=False)
        loss = criterion(pred, y_batch)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        optimizer.step()

    # Evaluate after each epoch using MC-averaged metrics.
    train_elbo, train_nll, train_kl = evaluate_elbo(model, train_loader, criterion, kl_weight, num_batches, mc_samples=eval_mc_elbo)
    val_elbo, val_nll, val_kl = evaluate_elbo(model, val_loader, criterion, kl_weight, num_batches, mc_samples=eval_mc_elbo)
    train_rmse = evaluate_rmse(model, train_loader, mc_samples=eval_mc_rmse)
    val_rmse = evaluate_rmse(model, val_loader, mc_samples=eval_mc_rmse)

    if val_rmse < best_val_rmse:
        best_val_rmse = val_rmse
        best_epoch = epoch_idx + 1
        best_state = copy.deepcopy(model.state_dict())

    history['epoch'].append(epoch_idx + 1)
    history['stage'].append('warmup')
    history['train_elbo'].append(train_elbo)
    history['val_elbo'].append(val_elbo)
    history['train_nll'].append(train_nll)
    history['val_nll'].append(val_nll)
    history['train_kl'].append(train_kl)
    history['val_kl'].append(val_kl)
    history['train_rmse'].append(train_rmse)
    history['val_rmse'].append(val_rmse)

    if (epoch_idx + 1) % 20 == 0:
        print(
            f"Warmup {epoch_idx + 1:03d}/{warmup_epochs} | "
            f"val_rmse={val_rmse:.4f} | best_val_rmse={best_val_rmse:.4f} | "
            f"train_elbo={train_elbo:.4f}"
        )

# ------------------------------------------------------------------------
# Stage 2: variational training with stochastic forward passes + ELBO
# ------------------------------------------------------------------------
for vi_idx in range(vi_epochs):
    model.train()
    for x_batch, y_batch in train_loader:
        x_batch = x_batch.to(DEVICE)
        y_batch = y_batch.to(DEVICE)

        optimizer.zero_grad(set_to_none=True)
        # Main Bayes-by-Backprop objective call via DeepUQ helper.
        loss, _, _ = vi_elbo_step_compat(
            model,
            x_batch,
            y_batch,
            num_batches=num_batches,
            criterion=criterion,
            kl_weight=kl_weight,
            mc_samples=1,
        )
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        optimizer.step()

    scheduler.step()

    train_elbo, train_nll, train_kl = evaluate_elbo(model, train_loader, criterion, kl_weight, num_batches, mc_samples=eval_mc_elbo)
    val_elbo, val_nll, val_kl = evaluate_elbo(model, val_loader, criterion, kl_weight, num_batches, mc_samples=eval_mc_elbo)
    train_rmse = evaluate_rmse(model, train_loader, mc_samples=eval_mc_rmse)
    val_rmse = evaluate_rmse(model, val_loader, mc_samples=eval_mc_rmse)

    epoch_number = warmup_epochs + vi_idx + 1
    if val_rmse < best_val_rmse:
        best_val_rmse = val_rmse
        best_epoch = epoch_number
        best_state = copy.deepcopy(model.state_dict())

    history['epoch'].append(epoch_number)
    history['stage'].append('vi')
    history['train_elbo'].append(train_elbo)
    history['val_elbo'].append(val_elbo)
    history['train_nll'].append(train_nll)
    history['val_nll'].append(val_nll)
    history['train_kl'].append(train_kl)
    history['val_kl'].append(val_kl)
    history['train_rmse'].append(train_rmse)
    history['val_rmse'].append(val_rmse)

    if (vi_idx + 1) % 50 == 0:
        print(
            f"VI {vi_idx + 1:04d}/{vi_epochs} | beta={kl_weight:.5f} | "
            f"lr={optimizer.param_groups[0]['lr']:.6f} | "
            f"val_rmse={val_rmse:.4f} | best_val_rmse={best_val_rmse:.4f} | "
            f"train_elbo={train_elbo:.4f} val_elbo={val_elbo:.4f}"
        )

# Restore the best checkpoint so downstream plots/results use best generalization.
model.load_state_dict(best_state)
print(f"Restored best checkpoint from epoch {best_epoch} with val RMSE={best_val_rmse:.4f}")

test_elbo, test_nll, test_kl = evaluate_elbo(model, test_loader, criterion, kl_weight, num_batches, mc_samples=8)
test_rmse = evaluate_rmse(model, test_loader, mc_samples=64)
print(f"Test ELBO={test_elbo:.4f} | Test NLL={test_nll:.4f} | Test KL={test_kl:.4f} | Test RMSE={test_rmse:.4f}")

# ----------------------------------------------------------------------------
# Part 1 training curves
# Subplot 1: ELBO (raw + EMA)
# Subplot 2: RMSE (raw + EMA)
# Subplot 3: ELBO components (NLL and KL)
# Vertical dashed line = switch from warmup to VI stage.
# ----------------------------------------------------------------------------
epochs = history['epoch']
train_ema = ema(history['train_elbo'], alpha=0.2)
val_ema = ema(history['val_elbo'], alpha=0.2)
train_rmse_ema = ema(history['train_rmse'], alpha=0.2)
val_rmse_ema = ema(history['val_rmse'], alpha=0.2)

fig, axes = plt.subplots(1, 3, figsize=(17, 4))

axes[0].plot(epochs, history['train_elbo'], color='tab:blue', alpha=0.25, label='Train ELBO (raw)')
axes[0].plot(epochs, history['val_elbo'], color='tab:orange', alpha=0.25, label='Val ELBO (raw)')
axes[0].plot(epochs, train_ema, color='tab:blue', label='Train ELBO (EMA)')
axes[0].plot(epochs, val_ema, color='tab:orange', label='Val ELBO (EMA)')
axes[0].axvline(warmup_epochs, color='tab:gray', linestyle='--', linewidth=1.0)
axes[0].set_xlabel('Epoch')
axes[0].set_ylabel('ELBO')
axes[0].set_title('ELBO trend')
axes[0].grid(True, linestyle='--', linewidth=0.5, alpha=0.4)
axes[0].legend(frameon=False)

axes[1].plot(epochs, history['train_rmse'], color='tab:green', alpha=0.25, label='Train RMSE (raw)')
axes[1].plot(epochs, history['val_rmse'], color='tab:red', alpha=0.25, label='Val RMSE (raw)')
axes[1].plot(epochs, train_rmse_ema, color='tab:green', label='Train RMSE (EMA)')
axes[1].plot(epochs, val_rmse_ema, color='tab:red', label='Val RMSE (EMA)')
axes[1].axvline(warmup_epochs, color='tab:gray', linestyle='--', linewidth=1.0)
axes[1].set_xlabel('Epoch')
axes[1].set_ylabel('RMSE')
axes[1].set_title('Fit quality')
axes[1].grid(True, linestyle='--', linewidth=0.5, alpha=0.4)
axes[1].legend(frameon=False)

axes[2].plot(epochs, history['train_nll'], color='tab:purple', label='Train NLL')
axes[2].plot(epochs, history['val_nll'], color='tab:brown', label='Val NLL')
axes[2].plot(epochs, history['train_kl'], color='tab:olive', label='Train KL / num_batches')
axes[2].axvline(warmup_epochs, color='tab:gray', linestyle='--', linewidth=1.0)
axes[2].set_xlabel('Epoch')
axes[2].set_ylabel('Value')
axes[2].set_title('ELBO components')
axes[2].grid(True, linestyle='--', linewidth=0.5, alpha=0.4)
axes[2].legend(frameon=False)

plt.tight_layout()
plt.show()

# ----------------------------------------------------------------------------
# Part 1 diagnostics summary
# - Checks whether VI-stage EMA ELBO trends downward.
# - Reports RMSE EMA delta and best validation checkpoint stats.
# ----------------------------------------------------------------------------
vi_start = warmup_epochs
ema_train_elbo = ema(history['train_elbo'][vi_start:], alpha=0.2)

delta = ema_train_elbo[-1] - ema_train_elbo[0]
if len(ema_train_elbo) > 1:
    decreasing_ratio = float((np.diff(np.array(ema_train_elbo)) <= 0.0).mean())
else:
    decreasing_ratio = 1.0

delta_ok = delta < 0.0
ratio_ok = decreasing_ratio >= 0.65

ema_val_rmse = ema(history['val_rmse'], alpha=0.2)
rmse_delta = ema_val_rmse[-1] - ema_val_rmse[0]

print(f"VI-stage EMA train ELBO delta (last-first): {delta:.6f}")
print(f"VI-stage EMA train ELBO non-increase ratio: {decreasing_ratio:.3f}")
print(f"Criterion delta < 0.0: {'PASS' if delta_ok else 'FAIL'}")
print(f"Criterion ratio >= 0.65: {'PASS' if ratio_ok else 'FAIL'}")
print(f"Val RMSE EMA delta (last-first): {rmse_delta:.6f}")
print(f"Best val RMSE during training: {best_val_rmse:.4f} at epoch {best_epoch}")
print(f"Last logged val RMSE: {history['val_rmse'][-1]:.4f}")

# ----------------------------------------------------------------------------
# Part 1 posterior predictive plot
# - Draw MC predictions at a dense x-grid.
# - Plot mean and 95% interval from epistemic uncertainty.
# - Compare against noisy samples and true function.
# ----------------------------------------------------------------------------
x_grid_raw = torch.linspace(-3.0, 3.0, 300).unsqueeze(-1)
y_grid_true = target_fn(x_grid_raw)
x_grid_scaled = scale_x(x_grid_raw).to(DEVICE)

pred_mean, pred_std = posterior_predictive(model, x_grid_scaled, n_samples=300)
lower = pred_mean - 1.96 * pred_std
upper = pred_mean + 1.96 * pred_std

plt.figure(figsize=(7, 4))
plt.scatter(x_train_raw.squeeze().numpy(), y_train_raw.squeeze().numpy(), s=10, alpha=0.3, label='Train noisy')
plt.scatter(x_test_raw.squeeze().numpy(), y_test_raw.squeeze().numpy(), s=10, alpha=0.25, label='Test noisy')
plt.plot(x_grid_raw.squeeze().numpy(), y_grid_true.squeeze().numpy(), color='black', linewidth=2, label='True function')
plt.plot(x_grid_raw.squeeze().numpy(), pred_mean.squeeze().numpy(), color='tab:blue', label='Predictive mean')
plt.fill_between(
    x_grid_raw.squeeze().numpy(),
    lower.squeeze().numpy(),
    upper.squeeze().numpy(),
    color='tab:blue',
    alpha=0.2,
    label='95% epistemic interval',
)
plt.xlabel('x')
plt.ylabel('y')
plt.title('Bayes-by-Backprop posterior predictive')
plt.legend(frameon=False)
plt.tight_layout()
plt.show()