IVON Optimizer (Improved Variational Online Newton)¶

← Error-Aware Encoding | TIC-TAC →

IVON is a natural-gradient-based optimization algorithm derived from the Bayesian Learning Rule framework (Khan & Rue, 2023). It fits a Gaussian variational posterior \(q(\theta) = \mathcal{N}(\theta \mid \mu, \Sigma)\) over model parameters by tracking first and second moments, providing both optimization and lightweight uncertainty quantification in a single algorithm.

Mathematical Background¶

Variational Gaussian Posterior¶

IVON maintains a diagonal Gaussian variational posterior:

At each step, it samples \(\theta \sim q_t\), computes the gradient, and updates both \(\mu\) and the precision \(\Sigma^{-1}\) using natural-gradient-style updates.

Hessian Approximation¶

IVON uses one of two Hessian approximations to update the precision:

where \(\epsilon \sim \mathcal{N}(0, \Sigma)\) is the parameter noise and \(g = \nabla_\theta \mathcal{L}\) is the gradient at the perturbed parameters.

Update Rules¶

For parameter group with learning rate \(\eta\), effective sample size \(N\), and decay rates \(\beta_1, \beta_2\):

where \(d_t = 1 - \beta_1^t\) is a debiasing factor and \(r\) is the clipping radius.

Usage¶

Basic Training Loop¶

The key API pattern: wrap the forward/backward pass in optimizer.sampled_params(train=True) to activate weight sampling:

import torch
import torch.nn as nn
from torchregress.algorithms import IVON

model = nn.Sequential(nn.Linear(10, 64), nn.ReLU(), nn.Linear(64, 1))
optimizer = IVON(
    model.parameters(),
    lr=0.1,
    ess=100.0,          # Effective sample size (typically dataset size)
    mc_samples=1,        # MC samples per step
    hess_approx="price", # Hessian approximation method
    weight_decay=1e-4,
)

for x, y in train_loader:
    with optimizer.sampled_params(train=True):
        optimizer.zero_grad()
        pred = model(x)
        loss = nn.functional.mse_loss(pred, y)
        loss.backward()
    optimizer.step()

Multi-Sample Training¶

For better gradient estimates, increase mc_samples — the optimizer averages gradients over multiple weight samples:

optimizer = IVON(model.parameters(), lr=0.1, ess=100.0, mc_samples=4)

def closure():
    optimizer.zero_grad()
    with optimizer.sampled_params(train=True):
        pred = model(x)
        loss = nn.functional.mse_loss(pred, y)
        loss.backward()
    return loss

loss = optimizer.step(closure)

Uncertainty Estimation¶

The learned posterior enables MC prediction at test time:

model.eval()
predictions = []
for _ in range(30):
    with optimizer.sampled_params(train=False):
        pred = model(x_test)
        predictions.append(pred.detach())

preds = torch.stack(predictions)
mean = preds.mean(0)
epistemic_std = preds.std(0)

Checkpointing¶

optimizer.state_dict() includes current_step, which drives the momentum debias \(1-\beta_1^t\). Load it into a freshly constructed optimizer (same hyper-parameters) together with the model weights to resume exactly where training stopped.

When to Use IVON¶

Comparison with Other Optimizers¶

Limitations¶

Next steps¶

• Ensemble methods — combine multiple IVON-trained models for richer uncertainty
• Effective Bayesian Laplace — post-hoc uncertainty decomposition via natural-parameter Laplace
• Adaptive-Prior VI — context-aware variational inference for covariate-shifted test data
• Gaussian losses — pair IVON's weight uncertainty with a heteroscedastic output head

References¶