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

Mostafa Elhoushi; Alex Pretko; Nolan Dey; Bin Claire Zhang; Gavia Gray; Gurpreet Gosal; Abdulrahman Mahmoud; Shane Bergsma

cs.AISep 4, 2026

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

Authors: Mostafa Elhoushi, Alex Pretko, Nolan Dey, Bin Claire Zhang, Gavia Gray, Gurpreet Gosal, Abdulrahman Mahmoud, Shane Bergsma,

Organizations: Cerebras Systems · MBZUAI

Abstract

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.

Figures & tables

Figure 1: Layer dropout as a unified mechanism for efficient LLM training and inference. ( Left ) Layer dropout skips layers stochastically during pre-training, leading to faster training, and with our proposed configuration, does so without sacrificing validation loss. ( Center ) Trained models gain zero-shot “elastic depth,” degrading gracefully under early exit and layer skipping. ( Right ) This robustness carries over to post-training adapters and self-speculative decoding for lossless inference speedup.

Panel of Figure 2 — Figure 2: Coordinate Check. Scaling with $r^{l,t}_{\text{train}}=1/\rho^{l,t}$ during training with layer dropout yields stable activation scale across depth density. More details in Sec. B .

Panel of Figure 3 — Figure 3: Analysis of hyperparameter transferability on 271M model. We observe that optimal value for each hyperparameter remains similar across most layer dropout rates with the scaling factor $r_{\text{train}}=1/\rho$ .

Panel of Table 1 — Table 1: Ablating model granularities. Models trained at 20 TPP.

Figure 4: Ablating tensor granularity. Models trained at 20 TPP.

Panel of Table 2 — Table 2: Analysis of Dropout Distributions across layers. Models trained at 20 TPP.

Panel of Table 3 — Table 3: Ablating dropout time schedule. We group different dropout configurations that have the same active non-embedding FLOPs reduction induced by layer dropout across all training steps. Models trained at 20 TPP.

Figure 5: Early-exit validation loss for different model sizes trained at 20 TPP: no dropout vs. dropout configurations with 10% FLOPs savings. Additional comparisons in Fig. A.2 .

Figure 6: Intermediate layer skipping loss for models at 20 TPP: baseline vs. dropout configurations with 20% FLOPs savings. Extended results in Fig. A.3 .

Panel of Figure 7 — Figure 7: Early exit losses for models pre-trained with different dropout configurations, followed by freezing their weights and training early exit adapters as proposed by Balcony ( Jamialahmadi et al., 2025 ) . Models pre-trained with dropout always lead to better early exit losses even after adding exit adapters.

Model Size	TPP	Dropout Config.	Self-Speculative Decoding
			Speedup	$\alpha$	$\frac{T_{\text{Draft}}}{T_{\text{Target}}}$
270M	20	0	1.01 $\times$	92%	0.65
		0.2 (ILD)	1.20 $\times$	86%	0.55
		0.4 (ILD)	1.35 $\times$	89%	0.48
		0.4 (ILD, DTS)	1.22 $\times$	73%	0.57
504M	20	0	1.03 $\times$	90%	0.66

Table 4: Self-speculative decoding speedup across model scales using Draft & Verify ( Zhang et al., 2024a ) on XSUM ( Narayan et al., 2018 ) .

\alpha\in[0,1]

is acceptance rate for draft length

\gamma=5

T_{\text{Draft}}

is time to decode a single token for the selected subset of layers, and

T_{\text{Target}}

for the full model.

Figure 8: Validation loss across TPP for ILD with DTS, showing competitive performance with dense baselines as tokens-per-parameter increase.

Model Size	1.8B		3.9B		8.2B
TPP	20	20	20	20	20
Max Dropout	0	0.6	0	0.8	0.99
Layer Dist.	–	Inc.	–	Inc.	Inc.
Time Sched.	–	Dec.	–	Dec.	Dec.
FLOPs Savings $\uparrow$	0%	15%	0%	20%	25%
Val. Loss $\downarrow$	1.849	1.836	1.732	1.745	1.663

Table 5: Training larger models with aggressive dropout preserves accuracy while accelerating training and inference.

Panel of Figure 9 — Figure 9: Validation loss versus training FLOPs for larger model sizes at aggressive maximum dropout rates for 20 TPP. At the same training FLOPs, training with layer dropout leads to better loss.

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Model	Hidden Dim. $D$	Layers $L$	Heads	Head Size	FFN Mult.
271M	640	13	10	64	8
504M	896	17	14	64	8
906M	1152	23	9	128	8
1.8B	1536	30	12	128	8
3.9B	2048	40	16	128	8
8.2B	2688	52	21	128	8

Appendix

Table A.1: Model architectures used in the experiments.

Parameterization	Base Training	Scaled Training
Width	$d_{\text{base}}$	$d_{\text{base}}\cdot\color[rgb]{1,0.5,0}m_{d}$
Depth	$L_{\text{base}}$	$L_{\text{base}}\cdot\color[rgb]{0.132,0.545,0.132}{m_{L}}$
Dataset Size	$D_{\text{base}}$	$D_{\text{base}}\cdot\color[rgb]{0.75,0.5,0.25}{m_{D}}$
Layer Density	$1$	$\color[rgb]{0.75,0,0.25}{\rho}$
Model Size	$N_{\text{base}}=g(d_{\text{base}},L_{\text{base}})$	$N=g(d_{\text{base}}\cdot\color[rgb]{1,0.5,0}m_{d},L_{\text{base}}\cdot\color[rgb]{0.132,0.545,0.132}{m_{L}})$
Tokens per Parameter	$\text{TPP}_{\text{base}}=D_{\text{base}}/N_{\text{base}}$	$\text{TPP}=D_{\text{base}}\cdot\color[rgb]{0.75,0.5,0.25}{m_{D}}/N$

Appendix

Table A.2 : Summary of SP, µP, and CompleteP with layer dropout rate for a transformer model. Terms related to width (introduced by

\mu

P Yang & Hu (2021) ), depth (introduced by CompleteP Dey et al. (2025) ), data size (introduced by Bergsma et al. (2025a) ), and dropout (introduced in this paper) controls are highlighted in orange , green , brown , and purple respectively. Additional tunable parameters are highlighted in blue . Hidden refers to all linear layers in the transformer backbone. Layer density,

\rho

is the complement of layer dropout,

p

such that

\rho=1-p

Panel of Figure 17 — Figure A.1 : Coordinate Check passing for non-uniform distribution.

		Training Config.		Downstream Tasks $\uparrow$
Size	TPP	Layer Dropout	FLOPs Sav. $\uparrow$	BBH	PIQA	SIQA	Hella.	Wino.	ARC-c	ARC-e	OBQA	Lamb.	COPA	RACE
271M	20	0	0%	0.083	0.593	0.358	0.293	0.500	0.224	0.415	0.274	0.226	0.580	0.265
271M	20	0.2 (ILD)	10%	0.095	0.579	0.349	0.285	0.505	0.224	0.399	0.262	0.218	0.580	0.273
271M	20	0.4 (ILD, DTS)	10%	0.106	0.584	0.358	0.290	0.523	0.241	0.409	0.274	0.229	0.600	0.269
271M	20	0.4 (ILD)	20%	0.128	0.584	0.345	0.283	0.499	0.234	0.400	0.266	0.223	0.610	0.253
503M	20	0	0%	0.178	0.594	0.360	0.316	0.515	0.246	0.441	0.282	0.280	0.560	0.279

Appendix

Table A.3 : Downstream task performance benchmarks.

Panel of Figure 19 — Figure A.2 : Comparison of early-exit validation losses for models trained with different dropout configurations. All models trained with 20 TPP.

Panel of Figure 20 — Figure A.3 : Comparison of intermediate layer skipping validation losses for models trained with different dropout configurations. All models trained with 20 TPP.

LLM Training

Large language model training studies the data, objectives, optimizers, and systems used to build and adapt language models.

See all 332 papers

OPTIMA: Optimal One-shot Pruning for LLMs via Quadratic Programming Reconstruction
Dec 15, 2025Mohammad Mozaffari, Samuel Kushnir, Maryam Mehri Dehnavi +1LLM Pruning Neural Network Pruning
University of Toronto · Google DeepMind
Locality-Aware Redundancy Pruning for LLM Depth Compression
May 27, 2026Vincent-Daniel Yun, Youngrae Kim, Woosang Lim +3LLM Pruning LLM Compression
University of Southern California, United States · Neural Superintelligence Lab, MODULABS, Republic of Korea · Seoul National University, Republic of Korea +1
Adaptive Depth Sparse Framework: Similarity-Driven Resource Allocation for Pre-Trained LLMs
Jul 23, 2026Yidu Wu, Xiang Wang, Kejie Zhao +3LLM Pruning LLM Inference Acceleration
Department of Electronic and Electrical Engineering, Southern University of Science and Technology, Shenzhen, China · ACS Lab, Huawei Technologies Co., Ltd., Shenzhen, China · Department of Computer Science and Engineering, Southern University of Science and Technology, Shenzhen, China

Dec 15, 2025J/K move · Enter open · S save

OPTIMA: Optimal One-shot Pruning for LLMs via Quadratic Programming Reconstruction

Mohammad Mozaffari, Samuel Kushnir, Maryam Mehri Dehnavi, Amir Yazdanbakhsh

University of Toronto · Google DeepMind

LLM Pruning Neural Network Pruning Transformer Pruning Quadratic Programming

Open paper AskPDF arXiv

Post-training pruning cuts the inference cost of large language models, yet it faces a trade-off: simple heuristics that zero weights are fast but degrade accuracy, while principled joint optimization methods recover accuracy but are computationally infeasible at modern scale. One-shot methods such as SparseGPT sit in between and apply efficient but approximate weight updates. To close this gap, we introduce OPTIMA, a practical one-shot post-training pruning method that balances accuracy and scalability. OPTIMA casts layer-wise weight reconstruction after mask selection as independent, column-wise quadratic programs (QPs) that share one Hessian per weight matrix. Solving these QPs yields, at solver convergence, the per-column minimizer of the layer-wise reconstruction objective under the shared Hessian, and the shared Hessian makes the QPs easy to batch on accelerators. We implement an accelerator-friendly QP solver that accumulates one Hessian per weight matrix and solves many small QPs in parallel with a convergence certificate per column, enabling one-shot post-training pruning at scale on a single accelerator without fine-tuning. OPTIMA integrates with existing mask selectors and improves average zero-shot accuracy in 40 of 49 (model, mask selector, sparsity) settings across the LLaMA and Gemma families, with mean gains of +0.61, +0.56, and +1.01 points at 50% unstructured, 2:4, and 60% unstructured sparsity, and up to +3.97 points at 60% sparsity, while matching ADMM-based exact reconstruction on its own masks. On an NVIDIA H100, OPTIMA prunes an 8B-parameter transformer end-to-end in 40 hours and 60 GB of peak GPU memory. This shows that principled column-wise QP reconstruction is a practical alternative to heuristic weight updates in one-shot post-training pruning.

Abstract

LLM Training

Explore similar work

OPTIMA: Optimal One-shot Pruning for LLMs via Quadratic Programming Reconstruction

Locality-Aware Redundancy Pruning for LLM Depth Compression

Adaptive Depth Sparse Framework: Similarity-Driven Resource Allocation for Pre-Trained LLMs