Effective Bayesian Heteroscedastic Laplace¶

← TIC-TAC | Adaptive-Prior VI →

This algorithm performs a last-layer Laplace posterior approximation over a natural-parameterized heteroscedastic regression head. It allows for fast, post-hoc uncertainty decomposition into epistemic (model) and aleatoric (data) components.

Mathematical Background¶

Standard deep heteroscedastic regression predicts mean \(\mu\) and log-variance \(\log\sigma^2\) directly. However, the resulting negative log-likelihood (NLL) loss is non-convex w.r.t these parameters, making the Gaussian approximation of the parameter posterior (Laplace approximation) inaccurate.

1. Natural Parameterization¶

Following Immer et al., we parameterize the Gaussian distribution using its exponential family natural parameters:

To ensure \(\sigma^2 > 0\), we predict \(f_1, f_2 \in \mathbb{R}\) and define the natural parameters using a link function. For the exponential link function (link_fn="exp"):

This maps back to standard parameters as:

This parameterization makes the NLL loss convex w.r.t the network outputs \(f\), ensuring optimization stability and making the posterior well-suited for a Laplace approximation.

2. Last-Layer Laplace Approximation¶

We place a Gaussian prior over the natural head weights \(\theta \sim \mathcal{N}(0, \sigma_0^2 I)\). After finding the MAP estimate \(\theta_{MAP}\) on the training set, we approximate the posterior distribution of \(\theta\):

where the diagonal precision matrix \(\Lambda\) is computed using the empirical Fisher information matrix:

where \(\mathcal{L}_n\) is the NLL loss for sample \(n\). This gradient w.r.t head weights is computed efficiently using torch.func.grad and vectorized via torch.vmap.

3. Uncertainty Decomposition¶

At test time, we draw \(S\) weight samples \(\theta^{(s)} \sim q(\theta)\) to compute \(S\) forward passes. For each test point \(x\), this yields \(S\) mean predictions \(\mu^{(s)}(x)\) and variance predictions \(\sigma^{2(s)}(x)\). We decompose the predictive uncertainty into:

• Epistemic Uncertainty (model parameter uncertainty): $\(\sigma^2_{\text{epistemic}}(x) = \text{Var}\left(\{\mu^{(s)}(x)\}_{s=1}^S\right)\)$
• Aleatoric Uncertainty (inherent data noise): $\(\sigma^2_{\text{aleatoric}}(x) = \frac{1}{S} \sum_{s=1}^S \sigma^{2(s)}(x)\)$
• Total Uncertainty: \(\sigma^2_{\text{total}}(x) = \sigma^2_{\text{epistemic}}(x) + \sigma^2_{\text{aleatoric}}(x)\).

High-Level API: HeteroscedasticLaplaceRegressor¶

from torchregress.algorithms import (
    NaturalHeteroscedasticHead,
    HeteroscedasticLaplaceRegressor,
)

# 1. Instantiate the natural parameter head
head = NaturalHeteroscedasticHead(in_features=64, out_features=1, link_fn="exp")

# 2. Wrap the backbone and head
model = HeteroscedasticLaplaceRegressor(
    base_model=my_backbone,
    head=head,
    prior_precision=1.0,  # Prior precision 1 / sigma_0^2
    n_samples=30,         # MC samples for inference
)

# 3. Fit the model and compute the Laplace posterior
model.fit(train_loader, lr=1e-3, epochs=10)

# 4. Predict a PredictiveBatch containing mean, std, samples, and epistemic/aleatoric metrics
pred = model.predict_distribution(x_test)
epistemic = pred.extra["epistemic_variance"]
aleatoric = pred.extra["aleatoric_variance"]

Parameters¶

Complete Example¶

import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset
from torchregress.algorithms import (
    NaturalHeteroscedasticHead,
    HeteroscedasticLaplaceRegressor,
)

# Generate synthetic data
torch.manual_seed(42)
x = torch.randn(200, 5)
# Target has input-dependent noise
y = 2.0 * x[:, [0]] + torch.randn(200, 1) * torch.exp(0.5 * x[:, [1]])

dataset = TensorDataset(x, y)
loader = DataLoader(dataset, batch_size=32, shuffle=True)

# Define backbone and head
backbone = nn.Sequential(nn.Linear(5, 16), nn.ReLU(), nn.Linear(16, 8))
head = NaturalHeteroscedasticHead(in_features=8, out_features=1)

# Fit regressor
regressor = HeteroscedasticLaplaceRegressor(
    base_model=backbone,
    head=head,
    prior_precision=0.1,
)
regressor.fit(loader, lr=0.01, epochs=5)

# Predict distribution
x_test = torch.randn(10, 5)
pred_batch = regressor.predict_distribution(x_test)

print("Mean Predictions:\n", pred_batch.mean.squeeze(-1))
print("Epistemic Variance:\n", pred_batch.extra["epistemic_variance"].squeeze(-1))
print("Aleatoric Variance:\n", pred_batch.extra["aleatoric_variance"].squeeze(-1))

When to use this method¶

Limitations¶

1. Last-layer only — no backbone uncertainty: The Laplace approximation is restricted to the final head weights. Backbone parameters (feature extractor) are treated as fixed point estimates, so epistemic uncertainty from feature-level ambiguity is missed. For full-model uncertainty, use an ensemble or MC Dropout.
2. Diagonal Fisher approximation: The precision matrix \(\Lambda\) uses only diagonal entries of the empirical Fisher. This ignores off-diagonal correlations between head parameters, which can under-report epistemic variance for tightly coupled weights.
3. Natural head is mandatory: Standard \(\mu, \log\sigma^2\) heads produce non-convex NLL objectives where the Laplace approximation is invalid. You must use NaturalHeteroscedasticHead — and it only supports Gaussian likelihoods (no Poisson, Bernoulli, etc.).
4. No prior on backbone weights: The prior \(p(\theta)\) is placed only on the head layer. If the backbone was pre-trained on a different task or domain, the Laplace posterior inherits backbone miscalibration without correction.
5. Post-hoc only — does not influence training: The Laplace approximation is computed after MAP training converges. It adds uncertainty at test time but cannot regularize training or prevent overfitting the way a variational BNN or IVON optimizer can.

Recommendations¶

1. Always use NaturalHeteroscedasticHead with link_fn="exp". Standard \(\mu, \log\sigma^2\) heads produce non-Gaussian posteriors that break Laplace validity.
2. Tune prior_precision via a small validation split. Start at 1.0, sweep {0.01, 0.1, 1.0, 10.0}. Too low → under-regularized MAP, too high → over-smoothed posterior.
3. Set n_samples \geq 30 for stable epistemic variance estimates. Below 20 samples, \(\text{Var}(\{\mu^{(s)}\})\) fluctuates significantly between runs. With n_samples=1 the epistemic variance is reported as 0 (a single draw carries no spread information).
4. Freeze backbone after pre-training — re-fitting the Laplace on a backbone trained from scratch with the head can be numerically unstable. Pre-train backbone, freeze, then fit head + Laplace.
5. Complement with conformal prediction for coverage guarantees. Laplace provides a parametric posterior; conformal adds distribution-free coverage without assumptions on the posterior shape.

Next steps¶

• IVON Optimizer — natural-gradient variational training that maintains uncertainty during optimisation
• Adaptive-Prior VI — context-aware priors for covariate-shifted test data
• Uncertainty decomposition — semantics of epistemic vs aleatoric contracts
• Gaussian losses — pair with heteroscedastic output heads for aleatoric modelling

References¶