cs.LGMar 18, 2026

Dropout Robustness and Cognitive Profiling of Transformer Models via Stochastic Inference

Authors: Antônio Junior Alves Caiado, Michael Hahsler

Organizations: Lyle School of Engineering, Southern Methodist University, Dallas, TX, USA

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

Transformer-based language models are widely deployed for reasoning, yet their behavior under inference-time stochasticity remains underexplored. While dropout is common during training, its inference-time effects via Monte Carlo sampling lack systematic evaluation across architectures, limiting understanding of model reliability in uncertainty-aware applications. This work analyzes dropout-induced variability across 19 transformer models using MC Dropout with 100 stochastic forward passes per sample. Dropout robustness is defined as maintaining high accuracy and stable predictions under stochastic inference, measured by standard deviation of per-run accuracies. A cognitive decomposition framework disentangles performance into memory and reasoning components. Experiments span five dropout configurations yielding 95 unique evaluations on 1,000 samples. Results reveal substantial architectural variation. Smaller models demonstrate perfect prediction stability while medium-sized models exhibit notable volatility. Mid-sized models achieve the best overall performance; larger models excel at memory tasks. Critically, 53% of models suffer severe accuracy degradation under baseline MC Dropout, with task-specialized models losing up to 24 percentage points, indicating unsuitability for uncertainty quantification in these architectures. Asymmetric effects emerge: high dropout reduces memory accuracy by 27 percentage points while reasoning degrades only 1 point, suggesting memory tasks rely on stable representations that dropout disrupts. 84% of models demonstrate memory-biased performance. This provides the first comprehensive MC Dropout benchmark for transformers, revealing dropout robustness is architecture-dependent and uncorrelated with scale. The cognitive profiling framework offers actionable guidance for model selection in uncertainty-aware applications.

Figures & tables

DatasetCognitive DomainSourceSample Size (Total/Train/Test)Max Input Length
SQuAD v1.1Memory (Factual Recall)Wikipedia Passages500 (400/100)200 chars
HellaSwagReasoning (Commonsense)Event Scenarios500 (400/100)Full context
Table 1: Dataset Sources and Sampling Strategy
Architecture FamilyTypeHuggingFace Checkpoint
BERTEncoderbert-base-uncased
Encoderbert-large-uncased
Encoderprajjwal1/bert-tiny
RoBERTaEncoderroberta-base
Encoderdistilroberta-base
Encoderdeepset/roberta-base-squad2
Table 2: Evaluated Transformer Models
ConfigurationAttention RateFFN RateDescription
Deterministic0.00.0Standard inference mode (dropout disabled)
Baseline0.10.1Typical training-time dropout rates
High Attention0.60.1High stochasticity in attention mechanisms
High FFN0.10.6High stochasticity in feed-forward networks
High Both0.60.6Elevated stochasticity in all components
Table 3: Dropout Configurations
ModelOverallMemoryReasoningOverallMemoryReasoning
 AccuracyAccuracyAccuracyStd DevStd DevStd Dev
deberta-v3-small0.7960.9220.6690.01610.01390.0527
gpt2-medium0.7220.9220.5230.01460.01390.0464
scibert0.7030.8700.5350.01380.01720.0446
spanbert0.7010.9150.4860.01430.01390.0464
deberta-v3-base0.7000.9380.4620.01560.01390.0526
Table 4: Top 5 Performing Models (Baseline Configuration)
ModelDeterministicBaseline (MC)Degradation
 MeanStdMeanStd 
roberta-base-squad20.7350.00000.4970.0314−0.238-0.238
albert-base-v20.6400.01290.4890.0339−0.151-0.151
deberta-v3-base0.8000.00000.7000.0156−0.100-0.100
electra-base-discriminator0.7350.00000.6730.0191−0.062-0.062
spanbert-base-cased0.7450.00000.7010.0143−0.044-0.044
Table 5: Models Where Deterministic Inference Outperforms MC Dropout
Figure 1: Overall Accuracy Comparison (Top 5 Degradation). Performance comparison between Deterministic inference (blue) and Baseline MC Dropout (orange, 0.1 rate) across the five models exhibiting the largest overall degradation. Error bars represent the standard deviation across 100 stochastic forward passes.
Figure 2: Reasoning Task Performance. Reasoning accuracy remains relatively stable between deterministic and stochastic modes, with overlapping error bars indicating minimal impact of dropout on inferential capabilities.
Figure 3: Memory Task Performance. Memory accuracy exhibits severe degradation under MC Dropout, particularly for task-specialized models like roberta-base-squad2.
ConfigurationMemory Acc.Reasoning Acc.Gap
Deterministic (0.0/0.0)0.8040.508+0.296+0.296
Baseline (0.1/0.1)0.7920.492+0.301+0.301
High Attention (0.6/0.1)0.5880.479+0.110+0.110
High FFN (0.1/0.6)0.5380.494+0.044+0.044
High Both (0.6/0.6)0.5330.497+0.036+0.036
Table 6: Dropout Configuration Effects on Cognitive Capabilities
More in this topic

Monte Carlo Dropout

See all 11 papers

Explore similar work

CardsList
Apr 10, 2026

Uncertainty-Aware Transformers: Conformal Prediction for Language Models

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.

Abhiram Vellore, Niraj K. Jha

Department of Electrical and Computer Engineering, Princeton University

Conformal Prediction for Language ModelsLLM Uncertainty EstimationConformal PredictionLanguage Model Calibration
AskPDF
Sep 4, 2026

Don't Drop Dropout: Optimizing Layer Sparsity for Efficient LLM Training and Inference

Layer dropout (a.k.a. stochastic depth) has been shown to enable faster training, higher accuracy, and robustness to zero-shot layer pruning in both language and vision transformers. However, as models and datasets have scaled, dropout - particularly layer dropout - has largely disappeared from large language models (LLMs) pre-training recipes. While some prior work has reported that dropout can degrade accuracy, no comprehensive study has quantified, let alone mitigated, this effect. In this study, we show that layer dropout should be used in state-of-the-art LLM training, establishing best practices and scaling analysis for both training and post-training benefits. Concretely, with optimal layer distribution, time schedule, and optimizer hyperparameters, we observe that at the same training FLOPs layer dropout leads to lower loss. For a given number of training steps, LLMs can achieve lower or similar validation loss while saving upto 25% of training FLOPs. Moreover, layer dropout enables significant post-training optimizations, such as early exit, intermediate-layer skipping, and self-speculative decoding, yielding up to 1.5x inference speedup with negligible accuracy loss. Across more than 2400 training experiments, spanning models from 271M to 8.2B parameters and datasets up to 160B tokens, we demonstrate that these findings extend reliably to large-scale training regimes. All pre-training experiments were run on Cerebras CS-3 systems.

Mostafa Elhoushi, Alex Pretko, Nolan Dey +6

Cerebras Systems · MBZUAI

LLM TrainingLLM Inference Acceleration
AskPDF
Oct 3, 2026

Length Generalization Needs Proper Regularization

Length generalization is the ability of sequential models to perform well on context lengths unseen during training. In this work, we show that the challenge of achieving length generalization is related not only to architectural choices such as positional encoding and the attention mechanism but also to the training procedure itself. We study how regularization affects length generalization and find that weight decay can hinder extrapolation. In contrast, dropout improves extrapolation when its placement within the architecture is reconsidered. In that regard, we show that the standard placement of dropout before layer normalization introduces a systematic distributional mismatch, and that applying dropout just before the linear projection resolves this issue. For example, a modified SmolLM3 with sliding window attention, continually pre-trained with dropout, can extrapolate perfectly to 64×\times on Needle-in-a-Haystack and far beyond the pre-training context size on RULER and HELMET. Mamba2 also benefits from dropout, suggesting an architecture-agnostic nature of the problem. We further propose Variance-Preserving Affine Dropout (VPAD), a new dropout strategy that substantially reduces the resulting pre-activation variance mismatch, leading to further extrapolation improvements in transformers.

Pavlo Vasylenko, Matthias Lindemann, André F. T. Martins +1

Instituto Superior Técnico, Universidade de Lisboa · Instituto de Telecomunicações · ELLIS Unit Lisbon +1

Length GeneralizationLanguage Model GeneralizationNeural Network GeneralizationLLM TrainingLong-Context Language Modeling
AskPDF