Adaptive-Prior Bayesian Uncertainty (VIDS-style)¶

← Heteroscedastic Laplace | Methods Overview →

This algorithm implements a variational regressor that trains an input-conditional, context-aware adaptive prior to make model predictions robust and well-calibrated under covariate shifts.

Mathematical Background¶

In standard Bayesian neural networks, the prior \(p(\theta)\) is static and independent of the test distribution. Under covariate shift, this can lead to overconfident or poorly calibrated predictions.

VIDS (Variational Inference under Distribution Shifts) addresses this by learning an adaptive prior \(p(\theta | X_{\text{context}}, x)\) conditioned on the global context of the feature distribution \(X_{\text{context}}\) and the specific query point \(x\).

1. Synthetic Environment Sampling¶

To train the adaptive prior, we generate \(E\) synthetic environments \(e = (X_e, Y_e)\) by bootstrapping subsets of our training data. Each bootstrap sample represents a simulated shifted training/validation environment.

2. Amortized Variational Inference (Guide)¶

For each environment \(e\), we define a variational posterior guide \(q(\theta | X_e, Y_e)\) parameterized by a neural network that maps context summaries of the environment's features and targets to posterior parameters:

3. Adaptive Prior Network¶

The prior network \(p(\theta | X_e, x)\) maps the environment context features \(X_e\) and a query feature \(x\) to prior parameters:

The parameters are optimized by maximizing the Evidence Lower Bound (ELBO) averaged across all synthetic environments:

where:

• \(\log p(Y_e | X_e, \theta)\) is the likelihood (negative log-likelihood loss w.r.t observation noise variance \(\sigma_{\text{noise}}^2\)).
• \(D_{KL}(q || p)\) is the analytical KL divergence between the diagonal Gaussian guide and the diagonal Gaussian adaptive prior.
• \(\beta\) is a KL regularization scaling factor.

4. Test-Time Inference¶

At test time, given the global training context \(X_{\text{train}}\) and a new test point \(x_{\text{query}}\), we evaluate the adaptive prior:

We draw \(S\) samples from this prior to predict \(\mu^{(s)} = \theta^{(s)} \cdot [x_{\text{query}}, 1]\). The predictive variance is decomposed into:

• Epistemic Uncertainty (context-shift variance): $\(\sigma^2_{\text{epistemic}} = \text{Var}\left(\{\mu^{(s)}\}_{s=1}^S\right)\)$
• Aleatoric Uncertainty (observation noise): $\(\sigma^2_{\text{aleatoric}} = \sigma^2_{\text{noise}}\)$
• Total Uncertainty: \(\sigma^2_{\text{total}} = \sigma^2_{\text{epistemic}} + \sigma^2_{\text{aleatoric}}\).

High-Level API: VIDSRegressor¶

from torchregress.algorithms import VIDSRegressor

# 1. Instantiate the regressor
model = VIDSRegressor(
    in_features=5,
    target_dim=1,
    hidden_dim=64,
    noise_variance_init=0.1,
)

# 2. Fit the adaptive prior and guide using bootstrap environments
model.fit(
    x_train_features=x_train,
    y_train=y_train,
    n_environments=32,       # Number of synthetic environments
    bootstrap_fraction=0.3,   # Size of each bootstrap environment
    epochs=50,
    beta=1.0,                 # KL penalty multiplier
)

# 3. Predict distribution under potential shift
pred = model.predict_distribution(x_test, n_samples=30)

Parameters¶

Complete Example¶

import torch
from torchregress.algorithms import VIDSRegressor

# Generate synthetic train data (source distribution)
torch.manual_seed(42)
x_train = torch.randn(100, 3)
y_train = x_train[:, [0]] * 1.5 - x_train[:, [1]] * 0.8 + 0.3 * torch.randn(100, 1)

# Generate shifted test data (covariate shift)
x_test = torch.randn(20, 3) + 2.0  # shifted mean

model = VIDSRegressor(in_features=3, target_dim=1, hidden_dim=32)

# Train using VIDS ELBO
model.fit(
    x_train_features=x_train,
    y_train=y_train,
    n_environments=12,
    bootstrap_fraction=0.4,
    epochs=10,
    lr=1e-2,
)

# Evaluate predictions
pred = model.predict_distribution(x_test, n_samples=20)
print("Shifted Test Means:\n", pred.mean.squeeze(-1))
print("Shifted Test Total Stds:\n", pred.std.squeeze(-1))

When to use this method¶

Limitations¶

1. Linear head only: VIDS fits a linear Bayesian head on top of features. It is not an end-to-end deep BNN — the feature extractor must be pre-trained or frozen. For deep architectures with learnable feature uncertainty, use an ensemble or MC Dropout.
2. Bootstrap environment cost: Training on \(E\) bootstrap environments requires \(E \times\) forward passes per ELBO step. With n_environments=32 and bootstrap_fraction=0.3, training time scales roughly as \(E \times N_{\text{epochs}}\) — significantly slower than a standard MAP regressor.
3. ELBO optimization sensitivity: The KL term \(\beta \cdot D_{KL}(q || p)\) can dominate the likelihood term early in training, causing posterior collapse (\(\sigma^2_{\text{post}} \to 0\)) before the prior adapts. Requires careful \(\beta\) scheduling.
4. \(\beta\) sensitivity: The KL weight \(\beta\) controls the exploration–exploitation balance. Too low → prior ignores context (reverts to static prior), too high → posterior collapses to prior (no data signal). Standard grid search over \(\{0.01, 0.1, 1.0, 5.0\}\) is recommended.
5. Synthetic environments may not cover real shifts: Bootstrap resampling only produces environments within the support of the training distribution. For shifts outside this support (e.g., new feature combinations, domain gaps), the adaptive prior may be overconfident. Complement with conformal prediction for coverage guarantees under arbitrary shifts.

Recommendations¶

1. Use pre-trained or frozen features: Apply VIDS as a last-layer uncertainty head. Train a backbone with a supervised or self-supervised objective first, freeze it, then fit VIDS on the frozen features.
2. Set n_environments \geq 32: Fewer than 16 environments yields insufficient prior diversity; the adaptive prior network cannot learn meaningful context conditioning.
3. Schedule \(\beta\): Start with \(\beta=0.01\) for the first 20% of epochs (warm-up), then linearly anneal to \(\beta=1.0\). Monitor the KL term — it should stabilise around 0.1–1.0 nats.
4. Tune noise_variance_init: Start at 0.1. If aleatoric variance grows with training (loss plateaus early), lower to 0.01. If epistemic variance collapses, raise to 0.5.
5. Validate on shifted hold-out data: Split a shifted validation set (if available) or use a worst-slice evaluation. VIDS can overfit to in-distribution environments if only in-distribution validation is monitored.
6. Complement with Heteroscedastic Laplace when you need both adaptive priors and full heteroscedastic noise modelling — Laplace handles the aleatoric head while VIDS handles the epistemic prior.

Next steps¶

• Effective Bayesian Laplace — alternative post-hoc uncertainty decomposition with lower computational cost
• IVON Optimizer — variational training that maintains a weight posterior during stochastic optimisation
• Conformal prediction — coverage-guaranteed intervals without distributional assumptions
• Test-time adaptation — OT-inspired reweighting for score shift under covariate drift

References¶