Variational Neurons in Transformers for Language Modeling

TL;DR

This paper introduces variational neurons into Transformer models for language modeling, enabling internal uncertainty estimation while maintaining strong predictive performance.

Contribution

It presents a novel integration of variational units within Transformers, enhancing internal uncertainty modeling without sacrificing accuracy.

Findings

Variational neurons integrate stably into Transformers.

They preserve strong predictive performance.

They produce informative uncertainty signals.

Abstract

Transformers for language modeling usually rely on deterministic internal computation, with uncertainty expressed mainly at the output layer. We introduce variational neurons into Transformer feed-forward computation so that uncertainty becomes part of the internal computation itself. Concretely, we replace deterministic feed-forward units with local variational units based on EVE while preserving the overall Transformer backbone. We evaluate this design in compact next-token language-modeling settings. We compare deterministic and variational variants with both predictive and probabilistic criteria. Alongside negative log-likelihood, perplexity and accuracy, we analyze calibration, conditional variance, mutual information and latent-usage statistics. The resulting picture is clear. Variational neurons integrate stably into Transformers, preserve strong predictive performance and produce informative uncertainty signals. The experiments also show that task quality, useful depth and internal stability are distinct properties. These results establish variational Transformers as a practical form of uncertainty-aware language modeling. They show that Transformers can predict with an explicit internal structure of uncertainty, which supports stronger probabilistic evaluation and a more informative analysis of model behavior.