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
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 ( ↑ ), NLL ( ↓ ) and Brier score ( ↓ ); the x -axis is the backbone size (Small, Base, Large, Huge). Points show the mean over ten seeds and error bars ±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 ( ↑ ) and FPR at 95% TPR ( ↓ ). Points show the mean over ten seeds with ±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 ↑
0.8880 ± 0.0011
0.8866 ± 0.0004
0.8904±0.0007
NLL ↓
0.3891 ± 0.0015
0.4034 ± 0.0010
0.3744±0.0010
Brier ↓
0.1669 ± 0.0008
0.1693 ± 0.0006
0.1628±0.0003
Under shift (mean of four sets)
Accuracy ↑
0.8323 ± 0.0017
0.8373 ± 0.0025
0.8395±0.0015
Table 1: Fine-tuning endpoint at 2.56 M samples, the largest evaluation present in every run. Mean ± 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×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×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.
AdamW+CE
VBLL
IBProbit
Selected at 781 steps
completed trials
61
59
52
learning rate / B
4.27×10−6
6.34×10−6
6.13×10−6
weight decay / B
1.60×10−7
2.92×10−4
1.81×10−5
β1 -horizon / β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.
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.
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
While Deep Neural Networks (DNNs) achieve remarkable performance, their tendency to produce overconfident predictions. Evidential Deep Learning (EDL) mitigates this by formulating predictions as a Dirichlet distribution over class probabilities to explicitly quantify epistemic uncertainty. However, we found that the conventional EDL suffers from two fundamental limitations: a Kullback-Leibler (KL) penalty that only suppresses the evidence of negative classes, producing excessively high evidence therefore decreasing the model's ability to quantify uncertainty, and an absence in theoretical guarantee of setting Dirichlet parameter α=e+1. In this paper, we propose a mathematically principled framework, Variational Inference Evidential Deep Learning (VI-EDL). By reformulating evidential learning through the lens of variational inference, we derive an Evidence Lower Bound (ELBO), which prevents the evidence from growing excessively. Theoretically, we rigorously establish a generalization bound and reveal how the predicted uncertainty, feature and network complexity affect this bound, and why setting α=e+1 can minimize it. Extensive experiments on standard visual and medical datasets demonstrate that VI-EDL achieves state-of-the-art performance, showing excellent performance in out-of-distribution detection, noise detection and autonomous driving scenario. The code is available in https://github.com/seutjw/VI-EDL.
Jiawei Tang, Xinyan Du, Hui Liu +2
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
Modern deep learning models remain notoriously prone to overconfidence, limiting their reliability in high-stakes applications. Bayesian methods aim to counter this by learning a distribution over model parameters, and recent advances now make this feasible for large-scale architectures at costs comparable to AdamW. However, a challenge remains at test time: predictions must be averaged across many forward passes with weights sampled from the posterior, which is prohibitively expensive. Variance propagation offers an efficient alternative, computing layer-wise analytical approximations of uncertainty in a single forward pass. While such techniques are effective for MLPs, their extension to modern architectures remains challenging, due to increased depth and diversity of layer types. To fill this gap, we propose Calibrated Variance Propagation (CVP), which introduces a new propagation method for normalization layers, combines it with recent techniques for handling activation functions, and absorbs residual error through a light calibration step. CVP yields comparably accurate uncertainty estimates to MC sampling across transformers and CNNs, at a fraction of the cost. Against prior variance propagation work, CVP improves coverage at 0.5% risk from 8.2% to 14.6% with BEiT-3 on Visual Reasoning (NLVR2) and from 2.6% to 10.8% with ViLT on VQAv2, with gains extending to convolutional architectures.
Tobias Jan Wieczorek, Leon de Andrade, Thomas Möllenhoff +1
TU Darmstadt & hessian.AI, Darmstadt, Germany · RIKEN Center for Advanced Intelligence Project, Tokyo, Japan