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

2609.05275 ↗ · 2026-09-04 · cs.AI

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

This is a slightly extended version of the paper published at ICML 2026

The finding

Headline claim

Properly configured layer dropout (stochastic depth) should be restored to state-of-the-art LLM pretraining, because it cuts training FLOPs (up to 25%) at equal or better validation loss and enables post-training inference speedups (up to 1.5x) with negligible accuracy loss.

The test

More than 2400 pretraining runs on Cerebras CS-3 systems, decoder-only Celerity-style transformers, spanning 271M to 8.2B parameters (systematic sweeps up to 3.9B/116B tokens, larger confirmatory runs at 8.2B/160B tokens), at compute-optimal 20 tokens-per-parameter budgets, varying dropout rate, granularity, depth-wise distribution, temporal schedule, and jointly tuned optimizer hyperparameters (learning rate, batch size, weight decay) via muP/CompleteP transfer.

Effect size & scope

At equal training FLOPs, optimal layer dropout yields lower validation loss than dense baselines; for a fixed number of steps, up to 25% training FLOP savings at similar or lower loss. Inference: up to 1.5x speedup via early exit, intermediate-layer skipping, and self-speculative decoding with negligible accuracy loss. The key enabling detail is the scaling factor r_train = 1/rho, which allows one set of optimizer hyperparameters to transfer across dropout rates. The largest fully controlled sweeps are at a few billion parameters and ~100B tokens, below frontier scale.

Limitations

The authors concede that prior degradations may reflect suboptimal schedules/hyperparameters rather than fundamentals, and their own earlier work reported accuracy degradation with naive layer dropout. Unstated: all experiments use one architecture family (ALiBi, squared ReLU), one dataset mix, one hardware vendor (Cerebras), and single-epoch 20 TPP regimes; the largest aggressive-dropout demonstration (8.2B) is small by frontier standards, and wall-clock speedups depend on hardware that can exploit per-sequence layer skipping.

Counterpoint

Results are loss-based on an in-house model family and corpus rather than frontier benchmarks or third-party replication; the headline 25% FLOP savings and 1.5x inference speedup are best-case numbers from tuned configurations, not average gains; 271M to 8.2B is far from frontier scale so extrapolation to trillion-parameter training is untested; and Cerebras, which sells the hardware this method favors, published it, so incentive alignment warrants caution.

Economy implication

If the results transfer to frontier scale, training-time compute per capability unit falls and a single pretrained checkpoint can serve multiple latency/price tiers, modestly reinforcing the trend of falling inference cost per unit of capability and softening compute-demand growth. But at the scale tested this is a training-recipe efficiency result with no demonstrated frontier deployment yet, so its near-term economic significance is uncertain.

Who this matters for

Technical Capability: Practitioners planning pretraining budgets get a concrete recipe (increasing depth distribution, decreasing schedule, 1/rho scaling) that claims real FLOP savings and depth-elastic deployment from one checkpoint.

Market & Valuation: Investors in compute-demand theses should note that training-recipe efficiency gains like this are one route by which capability per FLOP rises faster than raw compute scaling, pressuring capex-justification narratives at the margin.

Bears on (advisory · papers never write evidence)

[6.2] Algorithmic / training-efficiency gains reduce the compute (and chip-access) needed for frontier capability.

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.