vbll¶

probly.method.vbll(base: Predictor[In, Out], variant: str = 'discriminative', parameterization: str = 'dense', prior_scale: float = 1.0, noise_init: float = 0.36787944117144233, cov_rank: int = 3, wishart_scale: float = 1.0, dof: float | None = None, noise_prior_scale: float = 0.01) → VBLLPredictor[In, Out][source]¶

Wrap a model with a Variational Bayesian Last Layer (VBLL).

Replaces the model’s last nn.Linear with a variational Bayesian last layer and emits, in closed form, a Gaussian over the network outputs based on [HWS24]. A trailing softmax (if any) is removed, since the layer outputs logits.

Three discriminative variant s are available, all producing a Gaussian over logits and sharing the same predict/representer pipeline:

"discriminative": the standard VBLL classifier (VBLLLayer), also usable for regression; its logit Gaussian is fully closed-form.
"student_t": additionally infers the noise variance via a Gamma posterior, giving a Student-t marginal (TVBLLLayer); the logit Gaussian uses one sampled noise variance per input.
"heteroscedastic": input-dependent noise via a second weight posterior (HetVBLLLayer); the logit Gaussian uses one sampled noise variance per input.

The returned predictor’s predict() yields the closed-form GaussianDistribution (the regression predictive, or the logit Gaussian for classification). For classification, use the registered VBLLRepresenter to obtain a categorical sample and an aleatoric/epistemic decomposition.

Parameters:

base – The model to wrap.
variant – The last-layer variant, one of "discriminative", "student_t" or "heteroscedastic". Defaults to "discriminative".
parameterization – Posterior covariance parametrization. The "discriminative" variant supports "diagonal", "dense" and "lowrank"; the other variants support "diagonal" and "dense". Defaults to "dense".
prior_scale – Scale of the isotropic prior covariance. Defaults to 1.0.
noise_init – Median of the random initial per-output noise standard deviation ("discriminative" variant only). Defaults to exp(-1), matching the reference initialization.
cov_rank – Rank of the low-rank covariance factor (only used when parameterization="lowrank"). Defaults to 3.
wishart_scale – Scale of the Wishart/Gamma prior on the noise precision ("discriminative" and "student_t" variants). Defaults to 1.0.
dof – Degrees of freedom of the Wishart/Gamma prior on the noise precision ("discriminative" and "student_t" variants; must be > 1 for "student_t"). Defaults to the reference values: 1.0 for "discriminative" and 2.0 for "student_t".
noise_prior_scale – Scale of the prior on the input-dependent noise weights ("heteroscedastic" variant only). Defaults to 0.01.

Returns: