Variational Filtering for Bayesian Last Layers with an Independent Binary Probit Surrogate

Dimitrije Markovic; Conor Heins; Alexander Tschantz; Christopher L. Buckley

stat.MLOct 12, 2026

Variational Filtering for Bayesian Last Layers with an Independent Binary Probit Surrogate

Authors: Dimitrije Markovic, Conor Heins, Alexander Tschantz, Christopher L. Buckley

Organizations: Independent Researcher Vienna, Austria · Independent Researcher Berlin, Germany · University of Sussex Brighton, UK · Department of Informatics University of Sussex Brighton, UK

Abstract

Modern deep classifiers trained by gradient descent need many passes over the data to become accurate, and even then are not always well calibrated; post-hoc temperature scaling corrects calibration after training but cannot influence training-time accuracy or convergence. We propose a variational filtering algorithm for the last layer of deep classifiers, whose parameters are treated as the latent state of a linear-Gaussian state-space model. Each mini-batch is assimilated by an exact prediction step followed by a closed-form variational update, made conjugate by an independent binary probit (IBProbit) surrogate for the categorical likelihood. The filter's process noise stabilises the updates under feature drift and keeps the predictive distribution from collapsing to a point estimate, allowing the same procedure to be used for linear probing on a frozen backbone and for full-network fine-tuning. We evaluate on DINOv3 backbones against two baselines tuned under an identical protocol: AdamW with cross-entropy and post-hoc temperature scaling, and Variational Bayesian Last Layers (VBLL). As a linear probe on frozen features, IBProbit attains the best accuracy and negative log-likelihood on all five downstream tasks at the largest backbone, and on two of the five at the smallest. In out-of-distribution detection with CIFAR probes on the largest backbone, IBProbit reduces the false-positive rate at 95% true-positive rate by a factor of 1.4 to 3.2 relative to the better of the two baselines, on all three dataset pairs. When the whole network is fine-tuned on ImageNet-1k, IBProbit reaches the final performance of the cross-entropy baseline in half the wall-clock time.

Figures & tables

Figure 1: Linear probing on frozen DINOv3 features across backbone scale. Columns are the five downstream tasks; rows are accuracy (

\uparrow

), NLL (

\downarrow

) and Brier score (

\downarrow

); the

x

-axis is the backbone size (Small, Base, Large, Huge). Points show the mean over ten seeds and error bars

\pm 1

standard deviation. AdamW+CE is shown without and with temperature scaling (TS); the temperature is fit on the same validation split used for model selection. A black ring marks a setting in which the validation optimum of at least one seed fell in the last 10% of the step budget, so that the plotted value is a lower bound (Appendix C ). Every method is tuned independently per dataset under one shared 30,000-step budget, at DINOv3-Huge only; the other three backbone sizes use the configuration tuned at Huge (Appendix C ), so points left of Huge are not separately optimized for any method.

Figure 2: Out-of-distribution detection from frozen DINOv3-Huge features, scored by maximum softmax probability. Columns are the three ID/OOD pairs, far-OOD for the two pairs with SVHN and near-OOD for CIFAR-10 vs CIFAR-100; rows are AUROC (

\uparrow

) and FPR at 95% TPR (

\downarrow

). Points show the mean over ten seeds with

\pm 1

standard deviation. AdamW+CE is shown with temperature scaling fit on the in-distribution validation split — the same operating point reported in Figure 1 — and the single fitted temperature is applied to both the ID and the OOD logits, as it would be at deployment.

Figure 3: Full-network fine-tuning of DINOv3-Huge on ImageNet-1k: metrics on the held-out

80\%

of the ImageNet validation split against samples seen (log scale). Lines show the mean over five seeds, and bands the range between the smallest and the largest value across seeds. The dashed vertical line marks the

399{,}872

-sample budget the hyperparameters were selected at; results to the right of it lie beyond the budget used for tuning, for every method equally. Accuracy, NLL and Brier are shown; NLL and Brier are proper scoring rules and support the calibration claim, whereas ECE depends on the choice of binning. Brier is the multiclass definition over the full predictive vector. Raw AdamW+CE is omitted: CE+TS is the baseline being compared against, and in the accuracy panel the two coincide exactly.

Figure 4: The same fine-tuning trajectories scored under distribution shift. Columns are the four shift sets; rows are accuracy, NLL, Brier score and the AUROC of maximum-softmax-probability detection against the in-distribution split. Lines, bands and the dashed vertical line are as in Figure 3 , and markers show the individual evaluations. ImageNet-R and -A cover 200 of the 1000 classes and are evaluated with a subset mask; ImageNet-V2 and -Sketch span all 1000 and are not masked. Every method is fit in distribution and then evaluated under shift; in particular the single temperature of the CE baseline is fitted on the in-distribution validation split and applied unchanged to the shifted logits, never refitted on the shifted set. The shift sets are scored at every second evaluation of the trajectory, so these panels contain fewer points than Figure 3 .

Figure 5: Training needed to reach the final performance of cross-entropy with temperature scaling (CE+TS), in wall-clock hours (left) and in epochs (right). The target is the mean over the five CE+TS seeds of the accuracy, NLL and Brier score at the end of the two-epoch run (Table 1 ). The target is joint : a run reaches it at the first evaluation at which all three metrics are at or better than the target. The CE+TS bar is the length of its own runs. For IBProbit the bar shows the mean over the five seeds and circles the individual seeds; all five reach the target, and the number above the bar is the speed-up relative to CE+TS. No VBLL run reaches the target within two epochs, so VBLL is right-censored: the hatched bar with an arrow marks the length of its runs, which is a lower bound. Time is measured from the first evaluation of each run. The times of CE+TS and IBProbit include the evaluation of the shift sets during training, which the VBLL runs delegate to a separate process.

Metric	AdamW+CE+TS	VBLL	IBProbit
In distribution
Accuracy $\uparrow$	0.8880 $\pm$ 0.0011	0.8866 $\pm$ 0.0004	$\bm{0.8904\pm 0.0007}$
NLL $\downarrow$	0.3891 $\pm$ 0.0015	0.4034 $\pm$ 0.0010	$\bm{0.3744\pm 0.0010}$
Brier $\downarrow$	0.1669 $\pm$ 0.0008	0.1693 $\pm$ 0.0006	$\bm{0.1628\pm 0.0003}$
Under shift (mean of four sets)
Accuracy $\uparrow$	0.8323 $\pm$ 0.0017	0.8373 $\pm$ 0.0025	$\bm{0.8395\pm 0.0015}$

Table 1: Fine-tuning endpoint at

2.56

M samples, the largest evaluation present in every run. Mean

\pm

standard deviation over five seeds. The rows under shift average ImageNet-R, -A, -V2 and -Sketch; AUROC is maximum-softmax-probability detection against the in-distribution split at the same step, with the in-distribution reference restricted to the covered classes on the two masked sets. ECE is omitted for the reason given in Figure 3 ’s caption; Brier is the multiclass definition over the full predictive vector.

Appendix figures & tables2 assets

Supplementary material from the paper’s appendix.

Appendix

AdamW + cross-entropy (+ post-hoc temperature scaling)
Dataset	batch	learning rate	weight decay	epochs	trials
Oxford Pets	64	$0.00395$	$2.64\times 10^{-5}$	527	102
Food-101	512	$0.00113$	$0.000223$	205	100
Flowers-102	256	$0.021$	$0.00168$	10000	103
Stanford Cars	64	$0.00101$	$1.6\times 10^{-5}$	237	102
DTD	64	$0.00217$	$0.00103$	518	100

Appendix

Table 2: Selected configurations, one per (method, dataset), tuned at DINOv3-Huge and applied at every backbone size. epochs restates the shared 30,000-step budget at that entry’s batch size and is not an independent hyperparameter. The final column is the number of completed trials the study reached. IBProbit has no learning rate or weight decay: its head update is closed-form.

Selected at $781$ steps
	AdamW+CE	VBLL	IBProbit
completed trials	$61$	$59$	$52$
learning rate / $\sqrt{B}$	$4.27\times 10^{-6}$	$6.34\times 10^{-6}$	$6.13\times 10^{-6}$
weight decay / $\sqrt{B}$	$1.60\times 10^{-7}$	$2.92\times 10^{-4}$	$1.81\times 10^{-5}$
$\beta_{1}$ -horizon / $\beta_{2}$ -horizon	$95.9$ / $0.87$	$90.7$ / $2.15$	$84.5$ / $10.7$
head / backbone learning rate	$29.7$	$479$	—

Appendix

Table 3: Selected fine-tuning configurations at DINOv3-Huge. The upper rows are the values the search selected at

781

steps; the lower rows are what each run of Figures 3 – 5 and Table 1 actually used at

5{,}004

steps, after the transfer of Appendix D.3 . Every run uses batch size

B=512

, a constant learning rate and the same layer-wise decay.

Explore similar work

Online Bayesian Node Classification on Inductive Graphs under Distribution Shift
Sep 12, 2026Jinwen Xu, Gonzalo Mateos Buckstein, Qin LuNode Classification Graph Neural Networks
School of Electrical and Computer Engineering, University of Georgia, Athens, GA, USA · Department of Electrical and Computer Engineering, University of Rochester, Rochester, NY, USA
Variational Inference for Evidential Deep Learning
May 26, 2026Jiawei Tang, Xinyan Du, Hui Liu +2Evidential Deep Learning Uncertainty Quantification
School of Computer Science and Engineering Southeast University Nanjing 210096, China · School of Computing Information Sciences Saint Francis University2026 Hong Kong SAR, China · CityDepartmentUniversityof Computerof Hong KongScience Hong Kong SAR, China
Calibrated Sampling-Free Uncertainty Estimation in Bayesian Deep Learning
Jun 15, 2026Tobias Jan Wieczorek, Leon de Andrade, Thomas Möllenhoff +1Bayesian Neural Networks Uncertainty Quantification
TU Darmstadt & hessian.AI, Darmstadt, Germany · RIKEN Center for Advanced Intelligence Project, Tokyo, Japan

Sep 12, 2026J/K move · Enter open · S save

Online Bayesian Node Classification on Inductive Graphs under Distribution Shift

Jinwen Xu, Gonzalo Mateos Buckstein, Qin Lu

School of Electrical and Computer Engineering, University of Georgia, Athens, GA, USA · Department of Electrical and Computer Engineering, University of Rochester, Rochester, NY, USA

Node Classification Graph Neural Networks Bayesian Neural Networks

Open paper AskPDF arXiv

On evolving graphs, node classifiers must satisfy two key requirements: inductive generalization to newly arriving nodes under distribution shift and calibrated uncertainty for safety-sensitive applications. Standard graph neural networks (GNNs) are typically trained once and address neither requirement. We adapt the Bayesian last-layer (BLL) model by placing random last-layer parameters on top of a deterministic GNN encoder for uncertainty quantification. The categorical softmax likelihood required for classification breaks Gaussian conjugacy, so neither the training posterior nor the test-time streaming update has a closed-form solution. To address both challenges, we introduce a variational Bayesian last-layer (VBLL) objective that jointly trains the encoder and an approximate last-layer posterior by maximizing an evidence lower bound with a Monte Carlo expected log-likelihood. At test time, we freeze the encoder and apply an online Laplace update to the last-layer posterior. This update corresponds to a power-prior Bayesian model with exponential forgetting and a Kullback-Leibler anchor to the training posterior. Across five node-classification benchmarks under distribution shift, online GVBLL is the only method to achieve the best accuracy and negative log-likelihood on every dataset. It improves accuracy by up to 17 percentage points on Cora and 14 percentage points on ogbn-arxiv over the strongest non-GVBLL baseline, while remaining competitive in calibration with MC Dropout, Deep Ensembles, Temperature Scaling, and Gaussian-process classifiers.

Abstract

Explore similar work

Online Bayesian Node Classification on Inductive Graphs under Distribution Shift

Variational Inference for Evidential Deep Learning

Calibrated Sampling-Free Uncertainty Estimation in Bayesian Deep Learning