Diffusion-Inspired Reconfiguration of Transformers for Uncertainty Quantification

ICLR 2026 Conference Submission16783 Authors

Keywords: Uncertainty quantification, Transformer

TL;DR: We reconfigure a transformer as a probability path, modeling each block probabilistically, with a unified model propagating representation uncertainty throughout the network.

Abstract: Uncertainty quantification 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.

Primary Area: other topics in machine learning (i.e., none of the above)

Submission Number: 16783