Title: Diffusion-Inspired Reconfiguration of Transformers for Uncertainty Quantification
Primary area: other topics in machine learning (i.e., none of the above)
Keywords: Uncertainty quantification, Transformer

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.

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