Diffusion-Inspired Reconfiguration of Transformers for Uncertainty Calibration
cs.LG
Submitted: 2026-02-09
Updated: 2026-09-06
License: http://creativecommons.org/licenses/by/4.0/
The gist: Uncertainty calibration in pre-trained transformers is critical for their reliable deployment in risk-sensitive applications.
Terminology
Abstract
Uncertainty calibration in pre-trained transformers is critical for their reliable deployment in risk-sensitive applications. Yet, most existing pre-trained transformers do not have a principled mechanism for uncertainty propagation through their feature transformation stack. In this work, we propose a diffusion-inspired reconfiguration of transformers in which each feature transformation block is modeled as a probabilistic mapping. Composing these probabilistic mappings reveals a probability path that mimics the structure of a diffusion process, transporting data mass from the input distribution to the pre-trained feature distribution. This probability path can then be recompiled on a diffusion process with a unified transition model to enable principled propagation of representation uncertainty throughout the pre-trained model's architecture while maintaining its original predictive performance. Empirical results across a variety of vision and language benchmarks demonstrate that our method achieves superior calibration and predictive accuracy compared to existing uncertainty-aware transformers.
Sources
- GPT-4 Technical Report
- Deep Batch Active Learning by Diverse, Uncertain Gradient Lower Bounds
- Benchmarking Bayesian Deep Learning on Diabetic Retinopathy Detection Tasks
- Adversarial Examples, Uncertainty, and Transfer Testing Robustness in Gaussian Process Hybrid Deep Networks
- Revisiting Kernel Attention with Correlated Gaussian Process Representation
- A Geometric View of Posterior Approximation
- Calibrating Transformers via Sparse Gaussian Processes
- Self-Attention through Kernel-Eigen Pair Sparse Variational Gaussian Processes
- An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale
- PaLM-E: An Embodied Multimodal Language Model
- Informative Priors Improve the Reliability of Multimodal Clinical Data Classification
- Evaluating Bayesian Deep Learning Methods for Semantic Segmentation
- Scalable Diffusion Models with Transformers
- Evidential Deep Learning to Quantify Classification Uncertainty
- Gemini: A Family of Highly Capable Multimodal Models
- LLaMA: Open and Efficient Foundation Language Models
- Plex: Towards Reliability using Pretrained Large Model Extensions
- Transformer Dissection: A Unified Understanding of Transformer's Attention via the Lens of Kernel
- A Survey on Bayesian Deep Learning
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: Local-Consistency Optimisation for Bayesian Neural Networks
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks