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

AI summary

Overview

Research area: Efficient large language model pre-training and inference — specifically structured, depth-wise regularization (layer dropout / stochastic depth) and its interaction with optimizer hyperparameters, sparsity distributions, and inference-time depth elasticity.

Technical level: Intermediate. The paper assumes familiarity with transformer residual blocks, learning-rate/batch-size scaling rules, and standard pre-training loss comparisons, but its central arguments are expressed in accessible terms.

Scope: A large-scale empirical study (2,400+ training runs on Cerebras CS-3 systems) establishing how to configure layer dropout so that LLM pre-training saves compute without losing accuracy, and so that the resulting models support depth-reduced inference.

What This Paper Is About

Layer dropout — randomly skipping entire transformer blocks during training — was standard in vision transformers and early language models, but has largely vanished from modern LLM pre-training recipes because prior work reported accuracy degradation at scale. The authors argue that those degradations come from poorly chosen configurations rather than from layer dropout itself. The goal is to identify the optimizer settings, dropout distributions, and time schedules that make layer dropout a net win for both training efficiency and inference.

Key Contributions

1. Improved compute–accuracy trade-offs. Properly configured layer dropout reduces training FLOPs while achieving validation loss competitive with, and in several cases superior to, dense baselines at scale.
2. Joint optimization framework. The paper identifies key interactions between dropout configurations, time schedules, and optimizer hyperparameters that mitigate the degradations reported in prior work.
3. Depth-elastic inference. The average training dropout rate predicts zero-shot robustness to early exit and layer skipping without any retraining.
4. Scaling analysis and best practices. Across model and data scales, the authors recommend a progressively increasing distribution across depth paired with a decreasing schedule across steps, yielding up to 25% training FLOPs savings and up to 1.5× inference speedup.

Main Findings

• Layer dropout beats sub-layer dropout. Dropping whole transformer blocks per sample yields higher accuracy than dropping the attention and FFN sub-blocks independently. On the 271M model at dropout rate 0.1, layer dropout produced a 3.95% training-loss delta and 3.18% validation-loss delta, versus 4.72% and 3.61% for sub-layer dropout. On the 503M model, layer dropout gave 2.60% training and 2.84% validation deltas, versus 3.94% and 3.03% for sub-layer. The authors note that dropout-induced degradation shrinks as model size grows.

• Per-sequence masking beats per-batch masking. For any given dropout rate, sampling the Bernoulli mask independently per sequence produces lower loss than drawing one mask per layer per step. The authors note a middle ground — dropping different batches on different devices during distributed training, or a different mini-batch per gradient-accumulation step — as unexplored future work.

• Non-uniform depth distributions beat uniform ones at fixed FLOPs. Comparing uniform, increasing layer dropout (ILD, linearly ramping from 0 to p_max across depth), and alternating layer dropout (ALD, applied every other layer), non-uniform distributions consistently outperform uniform at matched average dropout. ALD was best at the smallest model size but its advantage diminished with scale, while ILD's advantage over uniform widened. On the 906M model at 20% FLOPs savings, ILD reached 1.9983 validation loss (2.34% delta) versus 2.0078 for uniform (2.82%) and 2.0045 for ALD (2.66%).

• Decreasing time schedules dominate. Across 271M, 503M, and 906M models, decaying the dropout rate from p_max to zero over training consistently outperforms constant and increasing schedules at the same total FLOPs budget. Increasing schedules degraded loss substantially — for example, at 20% savings with p_max = 0.8, the increasing/increasing configuration reached 7.04% validation delta on the 906M model versus 1.55% for the increasing/decreasing combination.

• Layer dropout can beat the dense baseline with fewer FLOPs. At 5% training FLOPs savings, combining an increasing depth distribution with a decreasing time schedule achieved lower validation loss than the dense baseline on the 503M model (2.1103 versus 2.1096 baseline, a 0.03% delta) and on the 906M model (1.9513 versus 1.9526 baseline, a −0.06% delta). The authors describe this as the first demonstration that the dense baseline can be beaten with fewer training FLOPs.

• r_train = 1/ρ is critical for hyperparameter transfer. Different codebases scale layer dropout output differently. The authors show empirically that scaling training-time output by the inverse layer density ρ = 1 − p satisfies the Maximal Residual Stream Update desideratum and keeps optimal learning rate, batch size, and weight decay roughly constant across dropout rates, allowing hyperparameters to be tuned once and transferred. Setting r_eval = 1 ensures the eval-time activation equals the expected training-time activation.

• Recommended recipe. Increasing dropout across layers combined with decreasing dropout across time, with per-sequence masking on whole transformer blocks, and r_train = 1/ρ.

Methodology in Plain English

The authors train decoder-only transformers following the Celerity architecture: ALiBi position embeddings, squared ReLU activations, and the Llama3 vocabulary, on a corpus of natural language text and code. Following Hoffmann et al. (2022), most experiments use a compute-optimal budget of 20 tokens per parameter at each model size.

Rather than comparing dropout rates using a single fixed set of hyperparameters — which the authors call the "hyperparameter lottery" — they first tune learning rate, batch size, and weight decay on a small base model, then scale those settings using μP, CompleteP, and Power Lines scaling rules. They verify that with r_train = 1/ρ, the optimal hyperparameters stay roughly constant across dropout rates, so later experiments can hold them fixed.

They then run controlled ablations in sequence: first whether to drop whole blocks or sub-blocks, then whether to sample masks per batch or per sequence, then which distribution to use across depth (uniform, increasing, or alternating), and finally which schedule to use across training time (constant, increasing, or decreasing). To compare schedules fairly, they define the mean training dropout rate P̄ as the average rate across depth and time, which equals total training FLOPs savings, and group configurations by matched P̄. Larger-scale runs with aggressive dropout rates follow, along with evaluations of depth-wise inference optimizations. All pre-training experiments were run on Cerebras CS-3 systems.

Why This Matters

Impact on research. The paper pushes back on the prevailing assumption that dropout is useless or harmful at LLM scale, showing that the earlier negative results were configuration artifacts. It provides a quantitative framework — hyperparameter transfer via r_train = 1/ρ, matched-FLOPs comparisons across distributions and schedules — that makes future dropout studies comparable rather than anecdotal. It also connects depth-aware pre-training to the broader literature on training-aware efficiency methods like quantization-aware training.

Real-world applications:

• Reducing pre-training cost. Up to 25% of training FLOPs can be saved at comparable or better validation loss, which translates directly into reduced GPU-hours and energy consumption for large runs.
• Latency-adaptive serving. A single trained model exposes "elastic depth," letting a deployment serve the same weights at different latency or compute budgets via early exit or intermediate-layer skipping.
• Speculative decoding. Depth-robust models support self-speculative decoding, where a shallower sub-network proposes tokens and the full model verifies them, yielding up to 1.5× inference speedup with negligible accuracy loss per the abstract.
• Edge and on-device deployment. Zero-shot robustness to layer skipping may benefit settings where the full-depth model cannot meet a memory or latency envelope.

Industry relevance. The findings are directly actionable for teams training or serving LLMs: the recommended configuration is a small change to an existing training loop rather than an architectural redesign, and the hyperparameter-transfer property means it can be adopted without re-running expensive tuning sweeps. For inference providers, depth elasticity offers a way to get multiple serving profiles out of one pre-trained checkpoint.

Future Directions

• Why does sub-layer dropout underperform layer dropout? The authors leave this unexplained, speculating that attention and FFN may work in tandem, and also leave open whether dropping only attention or only FFN would help.
• Finer or interleaved granularity. The authors propose combining the weight-loading advantage of per-batch dropout with the fine-grained sparsity of per-sequence dropout via per-device or per-gradient-accumulation-step masking, and leave per-token and per-neuron dropout unexplored.
• Schedules beyond pre-training. Other temporal schedules and applying dropout to mid-training, supervised fine-tuning, or continual pre-training are explicitly listed as future work.
• Higher dropout at larger scale. Because dropout-induced degradation shrank with model size, and because ALD trailed off at larger scales while ILD improved, the authors call for further investigation at larger scales and with more aggressive dropout rates; the truncated content does not report the final results of that large-scale aggressive-dropout study.

The paper reports two different figures for the upper end of its scaling range: the abstract describes experiments spanning 271M to 8.2B parameters and datasets up to 160B tokens across more than 2,400 training experiments, while the introduction describes 2,400+ runs spanning 271M to 3.9B parameters and up to 116B tokens. The detailed tables reproduced in the available content cover model sizes of 271M, 503M, and 906M. Section 8 (Inference Optimizations) and the scaling analysis of Section 9 are truncated in the available content, so the specific mechanisms behind the 1.5× speedup and the large-scale results are not reported here.

Target Audience

Researchers and engineers working on LLM pre-training efficiency, training-recipe design, and inference optimization. It is most useful to practitioners who already run or tune large transformer pre-training jobs and want a concrete, empirically validated configuration change, and to researchers studying regularization, structured sparsity, or depth-elastic architectures. Readers need comfort with transformer internals and hyperparameter scaling terminology; those without it can still follow the high-level findings, which are stated in plain loss-comparison terms.

Authors’ 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.