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

Don't Drop Dropout. Optimizing Layer Sparsity for Efficient LLM Training and Inference. By Mostafa Elhoushi and eight others, from Cerebras Systems and MBZUAI. Published on the fourth of September. No code repository is linked to it.

Layer dropout, also known as stochastic depth, has been shown to do three good things in both language and vision transformers: faster training, higher accuracy, and robustness to zero-shot layer pruning. And yet, as models and datasets scaled, dropout of every kind and layer dropout in particular largely disappeared from large language model pre-training recipes. Some prior work reported that dropout can degrade accuracy, but no comprehensive study had quantified that effect, let alone mitigated it. The claim made here is that layer dropout should be used in state-of-the-art training.

The result arrives with its conditions attached, which is what makes it a recipe rather than an observation. With an optimal layer distribution, time schedule and optimizer hyperparameters, layer dropout leads to a lower loss at the same training budget. Turned around, for a given number of training steps a model reaches lower or similar validation loss while saving up to twenty-five percent of its training compute. The benefits continue after training. Layer dropout enables early exit, intermediate-layer skipping and self-speculative decoding, together yielding up to a one point five times inference speedup with negligible accuracy loss.

The scale of the study is the argument for trusting it. More than two thousand four hundred training experiments, spanning models from two hundred and seventy-one million to eight point two billion parameters, on datasets of up to one hundred and sixty billion tokens. That range is what lets the authors say the findings extend reliably to large-scale training regimes, rather than holding only at the size where the effect was first noticed, which is the usual way a dropout result comes apart. All pre-training experiments were run on Cerebras CS-3 systems.

The framing is a correction rather than a discovery. Layer dropout was known to work, and then quietly left the recipes as scale increased, on a belief about accuracy that had never been carefully measured. What the study supplies is that measurement, together with the settings which make the technique pay: the distribution across layers, the schedule over time, and the optimizer hyperparameters that go with them. The gains are stated under those settings, not in general.

Don't Drop Dropout. Optimizing Layer Sparsity for Efficient LLM Training and Inference. By Mostafa Elhoushi and eight others, from Cerebras Systems and MBZUAI. Published on the fourth of September. This is the long version: the whole paper rather than the abstract.

Dropout used to be everywhere. It was standard in convolutional networks and in early transformers, and then, as language models scaled to billions of parameters and trillions of tokens, it quietly left the recipes. The reasoning was straightforward. A model trained for a single epoch over an enormous dataset has little opportunity for classical overfitting, so a regularizer aimed at overfitting has nothing to do, and empirical evidence suggested activation dropout actively degraded performance in exactly those conditions. The prevailing view became that dropout is unnecessary or harmful at scale.

The paper's opening move is to separate two things that got abandoned together. Layer dropout, also called stochastic depth, skips entire transformer blocks rather than individual activations, and that difference is not cosmetic. Activation dropout and unstructured sparsity typically do not produce wall-clock speedups, because sparse kernels carry overheads that eat the savings. Skipping whole blocks is structured sparsity, and it reduces active training compute almost linearly with the dropout rate. It also does something the regularization framing misses entirely: it makes the model robust to running at reduced depth later, so one pretrained model can adapt to different latency budgets without retraining.

There is a further conceptual permission the authors take, and it is the one that makes the study thinkable. Techniques introduced as regularizers can persist for reasons unrelated to preventing overfitting. Weight decay is the precedent they cite: it has been shown to influence optimization dynamics in large-scale pretraining rather than classical generalization. So dropout can be re-examined without assuming its value must come from regularization. The study then varies three things across fixed architecture and data: optimizer hyperparameters, the distribution and granularity of sparsity across depth, and the schedule of dropout over training steps.

Several findings come out of more than two thousand four hundred training runs. For any given rate, applying dropout per sequence produces lower loss than applying it per batch, which is a granularity choice most implementations make without comment. The recommended configuration pairs a distribution that increases progressively across depth with a schedule that decreases across training steps, and it yields up to twenty-five percent savings in training compute at competitive loss, with up to a one and a half times inference speedup afterwards through early exit, layer skipping and self-speculative decoding. The cleanest result is predictive: the average training dropout rate predicts how robust the resulting model will be to early exit and layer skipping, with no retraining required.

Two things deserve noting about scope. The study is explicitly restricted to structured depth-wise dropout, with neuron-level and weight-level variants left alone because they interact differently with hardware efficiency, so nothing here rehabilitates dropout in general. And all pre-training experiments ran on one vendor's hardware, which matters for a result partly about wall-clock efficiency. There is also an internal inconsistency in how the scale is reported: the abstract gives a range up to eight point two billion parameters and a hundred and sixty billion tokens, while the introduction describes the same study as spanning up to three point nine billion parameters and a hundred and sixteen billion tokens.

The shape of the argument is a correction rather than a discovery, and corrections of this kind are rarer than they should be. A technique was known to work, left the recipes as the field's regime changed, and stayed out on the strength of a belief about why it had worked in the first place. Re-testing it required suspending that belief, holding compute fixed, and running enough configurations to tell a bad setting from a bad idea. The finding that the average dropout rate predicts depth elasticity is the part with the longest reach: it turns a training knob into a statement about what the deployed model will tolerate.