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, Joel Hestness
ICML 2026 Proceedings of the
cs.AI
2026-09-04
Across 2,400+ runs, ILD+DTS layer dropout matches dense loss at equal FLOPs. An 8.2B model saves 25% training compute and gets 1.55x speculative decoding.
Layer dropout, also called stochastic depth, sped up training and helped accuracy in vision and early BERT, and it made reduced-depth inference less brittle. LLM pretraining mostly deleted dropout: one epoch over huge data, and activation dropout has been reported to hurt. Layer dropout is a different tool. Skipping a whole block turns into nearly linear FLOP savings and, at inference, into early exit, layer skipping, and self-speculative decoding. What was missing was a map of how to set it in the modern single-epoch, large-token regime without giving up loss.
Cerebras ran 2,400-plus pretraining jobs on CS-3, from 271M to 8.2B parameters and up to 160B tokens, and drew that map.
The models are decoder-only Celerity stacks: ALiBi, squared ReLU, Llama3 vocabulary. Layer dropout samples a whole transformer block per sequence (attention and FFN share the mask). Skipped sequences do not run that block, so training FLOPs fall roughly with the drop rate.
The scaling that matters is a train-time residual multiplier of 1/ρ, where ρ is layer density, and an eval multiplier of 1. That choice lets learning rate, batch size, and weight decay transfer across drop rates; a multiplier of 1 fails the coordinate check. Hyperparameters are tuned first; then granularity, depth distribution, and time schedule are compared at a compute-optimal 20 tokens-per-parameter.
Depth distributions: uniform, alternating layers (ALD), and linearly increasing with depth (ILD). Time schedules: constant, increasing, and decreasing (DTS). The recommended pair is ILD+DTS: drop little near the bottom, more near the top, noisy at the start of training, full depth at the end. Mean training drop rate P̄ is total FLOP savings; under ILD+DTS, P̄ = 0.25 pmax.
At a matched 5% FLOP saving, ILD+DTS matches or beats the dense baseline at 503M and 906M (906M: 1.951 vs 1.953). An increasing schedule is clearly worse. Whole-layer dropout beats dropping attention and FFN independently; per-sequence masks beat per-batch masks.
Larger models tolerate harsher rates. A 1.8B model at pmax=0.6 (15% saving) reaches validation loss 1.836 against dense 1.849. A 3.9B model at pmax=0.8 (20%) reaches 1.745 against dense 1.732, a small loss. An 8.2B model at pmax=0.99 (25%) reaches 1.663; the table has no dense 8.2B control. The 3.9B run starts at about 0.6L effective depth, with the last layer skipped 80% of the time, and still converges. At matched training FLOPs, the dropout curves sit lower for most of training.
At inference, mean training drop rate predicts zero-shot early-exit and skip robustness. Self-speculative decoding on the 3.9B model moves from 1.02× dense to 1.54×; the 8.2B run hits 1.55×. Skipping alternate layers, 3.9B loss falls from 6.446 dense to 2.129. Balcony early-exit adapters on a frozen backbone also land lower exit loss when the backbone was pretrained with dropout.
At high tokens-per-parameter, ILD+DTS stays within about 0.50% of dense loss. The authors read that as a path that continues to longer pretraining.
Dropout here is a depth curriculum, not an overfitting hammer: force a shallower effective net early, then give depth back. For pretraining, 25% of non-embedding FLOPs is real money. For serving, one checkpoint can early-exit, skip, or draft for itself without a second elastic architecture. ILD is the better default for accuracy and early exit; ALD is better for non-contiguous skipping. Pick from the deployment constraint.
At aggressive drop rates the best learning rate and weight decay fall, so hyperparameter transfer is incomplete. There is no comparison with learned routers such as Mixture-of-Depths, and no MoE or non-transformer run. How pmax should scale is still uncharted, as is how inference gains move with tokens-per-parameter. The 8.2B 25% saving has no dense control in the same table, and the 3.9B 20% run is slightly worse on validation loss. Every job ran on Cerebras; wall-clock savings of per-sequence layer drops on GPUs still need a separate measurement.