Variational Neurons in Transformers for Language Modeling

Yves Ruffenach

cs.LGMar 30, 2026

Variational Neurons in Transformers for Language Modeling

Authors: Yves Ruffenach

Organizations: Conservatoire National des Arts et Métiers Strasbourg, France

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.

Figures & tables

Figure 1: Standard and variational Transformer blocks. The overall Transformer backbone is unchanged. Only the feed-forward computation is replaced by a variational block that infers a local latent distribution, samples a latent state and projects it to model space before the residual update.

Run	Loss $\downarrow$	CE $\downarrow$	PPL $\downarrow$	Acc $\uparrow$	MC NLL $\downarrow$	ECE $\downarrow$	MI	Flip
Run A (19.9k)	4.8118	4.7370	114.10	0.2293	4.8164	0.0445	0.1188	0.2757
Run B (19.9k)	4.7950	4.7217	112.36	0.2307	4.8286	0.0457	0.1168	0.2549

Table 1: 19.9k-example setting. Final-validation objective and task metrics, together with extended-validation probabilistic metrics. Run B is slightly stronger on task-level metrics, while Run A preserves the stronger deep latent regime.

Run	$w_{1}$	$w_{2}$	$w_{3}$	$\mathrm{KL}_{1}$	$\mathrm{KL}_{2}$	$\mathrm{KL}_{3}$	$\mu_{1}^{2}$	$\mu_{2}^{2}$	$\mu_{3}^{2}$
A (19.9k)	0.9120	0.0438	0.0442	0.1608	0.1920	$\sim 1.6{\times}10^{-5}$	0.1551	0.0911	$\sim 1.5{\times}10^{-5}$
B (19.9k)	0.9171	0.0462	0.0367	0.1636	$2.0{\times}10^{-5}$	$1.47{\times}10^{-5}$	0.1522	$1.94{\times}10^{-5}$	$1.33{\times}10^{-5}$

Table 2: 19.9k-example setting internal diagnostics. Run A keeps a genuinely active second latent layer, whereas Run B almost collapses the deeper layers.

Run	CE $\downarrow$	PPL $\downarrow$	Acc $\uparrow$	$\mu^{2}_{\mathrm{eval}}\downarrow$	$\mu^{2}_{\mathrm{std}}\downarrow$	$\mu^{2}_{3}\downarrow$
v23 (10.0k)	4.8157	123.44	0.2256	1.5313	53.45	550.82
v22 (10.0k)	4.8208	124.06	0.2264	0.2119	2.94	13.78

Table 3: 10.0k-example setting. Final-validation task metrics and internal diagnostics. v23 is slightly better on raw LM metrics, while v22 is substantially cleaner internally.

Model	CE $\downarrow$	PPL $\downarrow$	Acc $\uparrow$	MC NLL $\downarrow$	ECE $\downarrow$	MI	Flip	CVaR $\downarrow$
	Final validation			Extended validation
DET	5.5417	255.12	0.1525	5.6673	0.0603	$\approx 0$	0	13.51
EVE	5.2896	198.26	0.1752	5.4709	0.0444	0.1830	0.3288	12.51

Table 4: Lower-budget matched deterministic-versus-variational comparison on the dedicated 2.0k-example setting. The same overall pattern already appears at this scale: EVE is stronger on final-validation task metrics and on extended-validation probabilistic metrics.

Figure 2: Validation CE across epochs on the matched approximately 19,925-raw-example protocol. EVE improves steadily across the full 5-epoch run, whereas DET reaches its best point at epoch 3 and then rises.

Figure 3: Sampling-based epistemic metrics for EVE and the matched deterministic baseline. Mutual information, conditional Monte Carlo variance, top-1 Monte Carlo flip rate, and epistemic ratio are clearly non-zero for EVE, whereas they remain at zero or near-zero values for DET under repeated deterministic forward evaluation. The vertical axis is logarithmic; exact zeros for DET are displayed with a tiny visualization floor only for plotting.

Explore similar work

Uncertainty-Aware Transformers: Conformal Prediction for Language Models
Apr 10, 2026Abhiram Vellore, Niraj K. JhaLLM Uncertainty Estimation Conformal Prediction
Department of Electrical and Computer Engineering, Princeton University
Variational Distributional Neuron
Feb 20, 2026Yves RuffenachVariational Autoencoders Variational Inference
Conservatoire National des Arts et Métiers
Agentic Control in Variational Language Models
Apr 14, 2026Yves RuffenachLanguage Model-Based Control Language Modeling
Conservatoire National des Arts et Métiers, Strasbourg, France

Apr 10, 2026J/K move · Enter open · S save

Uncertainty-Aware Transformers: Conformal Prediction for Language Models

Abhiram Vellore, Niraj K. Jha

Department of Electrical and Computer Engineering, Princeton University

LLM Uncertainty Estimation Conformal Prediction Language Model Calibration

Open paper AskPDF arXiv

Transformers have had a profound impact on the field of artificial intelligence, especially on large language models and their variants. However, as was the case with neural networks, their black-box nature limits trust and deployment in high-stakes settings. For models to be genuinely useful and trustworthy in critical applications, they must provide more than just predictions: they must supply users with a clear understanding of the reasoning that underpins their decisions. This article presents an uncertainty quantification framework for transformer-based language models. This framework, called CONFIDE (CONformal prediction for FIne-tuned DEep language models), applies conformal prediction to the internal embeddings of encoder-only architectures, like BERT and RoBERTa, while enabling hyperparameter tuning. CONFIDE uses either [CLS] token embeddings or flattened hidden states to construct class-conditional nonconformity scores, enabling statistically valid prediction sets with instance-level explanations. Empirically, CONFIDE improves test accuracy by up to 4.09% on BERT-tiny and achieves greater correct efficiency (i.e., the expected size of the prediction set conditioned on it containing the true label) compared to prior methods, including NM2 and VanillaNN. We show that early and intermediate transformer layers often yield better-calibrated and more semantically meaningful representations for conformal prediction. In resource-constrained models and high-stakes tasks with ambiguous labels, CONFIDE offers robustness and interpretability where softmax-based uncertainty fails. We position CONFIDE as a framework for practical diagnostic and efficiency/robustness improvement over prior conformal baselines.

Abstract

Explore similar work

Uncertainty-Aware Transformers: Conformal Prediction for Language Models

Variational Distributional Neuron

Agentic Control in Variational Language Models