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

Layer dropout enables 25% training FLOPs savings and 1.5× inference speedups without sacrificing LLM accuracy.

How can we optimize layer dropout (stochastic depth) to enable efficient LLM training and inference without sacrificing model accuracy?

Modern Large Language Models (LLMs) are typically trained as dense, static architectures, wasting compute on redundant layers and lacking the flexibility to adapt to varying latency requirements at inference time. The authors reintroduce layer dropout—stochastically skipping entire transformer blocks during training—using a specific configuration of increasing dropout across layers and decreasing dropout over time. This approach maintains or improves validation loss while enabling zero-shot inference optimizations like early exit and self-speculative decoding.

Paper Primer

The core mechanism hinges on two design choices: applying dropout at the transformer-layer level rather than sub-layers, and using a scaling factor of $1/\rho$ (where $\rho$ is layer density) to ensure stable hyperparameter transfer across different dropout rates. This is like a structural stress test: by forcing the model to learn with varying subsets of its layers, it develops an inherent elasticity that allows it to function reliably even when layers are skipped during deployment.

Layer dropout significantly improves compute-accuracy trade-offs.

Across 2400+ experiments, models trained with the recommended Increasing Layer Distribution (ILD) and Decreasing Time Schedule (DTS) achieved validation loss competitive with or superior to dense baselines. Up to 25% reduction in training FLOPs and up to 1.55× inference speedup.

Larger models exhibit higher robustness to aggressive structural sparsity.

Models up to 8.2B parameters converged successfully even with maximum dropout rates ($p_{\max}$) as high as 0.99, where the final layer was skipped 80% of the time during early training. Aggressive dropout rates enable deeper inference elasticity without degrading baseline generalization.

Why has dropout been abandoned in modern LLM pre-training?

Empirical evidence suggested that activation-level dropout degrades performance in single-epoch, large-data regimes, leading to the prevailing view that it is unnecessary or harmful at scale.

How does this method differ from existing depth-aware training techniques?

Unlike methods that introduce auxiliary losses, routers, or architectural changes, layer dropout induces robustness directly during pre-training without modifying the model architecture or adding parameters.

Introduction

Layer dropout can cut compute while preserving model quality, but its use in LLMs is unclear.

Layer dropout (stochastic depth) can skip entire transformer blocks, offering a direct path to lower compute. Yet, as LLMs grew to billions of parameters, practitioners abandoned it, assuming it harms accuracy or offers no wall‑clock benefit.

Layer dropout skips residual blocks during training, trimming active FLOPs while preserving the model’s representational power if the dropout pattern is chosen wisely.

**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.

Layer dropout enables efficient training and inference.

Related Work

We situate layer dropout among prior dropout variants and depth‑elastic training methods.

Dropout techniques differ in granularity: activation‑level dropout masks individual neuron outputs, weight‑level dropout (e.g., DropConnect) masks individual parameters, and structured dropout operates on groups such as channels, layers, or blocks.

This work concentrates on structured, depth‑wise dropout—commonly called layer dropout or stochastic depth—which randomly removes entire transformer blocks during training.

In the era of billion‑parameter language models, activation dropout has largely vanished from pretraining pipelines; early decoder‑only models such as GPT‑3 and OPT kept the original dropout settings, whereas later systems like PaLM and LLaMA omit dropout entirely.

Empirical studies show that in single‑epoch, massive‑data regimes activation dropout can hurt performance, while in multi‑epoch or data‑constrained regimes it still provides regularization benefits.

Weight decay, another ubiquitous regularizer, primarily shapes optimization dynamics rather than classic generalization in large‑scale pretraining.

Layer dropout was first introduced to stabilize training of very deep residual networks and later became a staple in large‑scale vision models; however, its optimal strength appears to shrink as dataset size grows, as observed in ConvNeXt models trained on ImageNet‑22K versus ImageNet‑1K.

Techniques such as Progressive Layer Dropout and LayerDrop demonstrated improved convergence and robustness in BERT‑era, multi‑epoch settings, yet recent applications of layer dropout to decoder‑only LLMs trained on massive token budgets have reported non‑negligible accuracy loss.

Beyond dropout, a broader class of training‑aware depth‑elastic methods—including quantization‑aware training, early‑exit models, routing‑based skipping, and elastic architectures like Once‑for‑All—enable inference‑time efficiency but typically require architectural changes, extra parameters, or auxiliary objectives.

Layer dropout occupies a distinct niche: it induces robustness to depth‑wise inference optimizations directly during pretraining without altering the model architecture or adding auxiliary losses.

To date, the literature lacks a controlled, large‑scale evaluation that systematically varies dropout configurations and optimizer settings, leaving the trade‑off between sparsity and accuracy at modern LLM scales unresolved.

Theoretical Formulation

We formalize layer dropout, detail its training and inference forms, and illustrate its transformer implementation.

Training large language models incurs huge compute; skipping whole residual blocks can cut FLOPs but risks destabilizing the forward signal.

Instead of always applying every layer, we randomly drop the entire residual branch during training, letting the model learn to propagate information through the shortcut path.

Sample $M_{\ell,t}\sim\text{Bernoulli}(0.6)$ → suppose $M_{\ell,t}=0$ (drop).

Output becomes $H_{\ell+1,t}=H_{\ell,t}+0\cdot1.67\cdot2 = H_{\ell,t}$ (identity shortcut).

Next training step, sample $M_{\ell,t}=1$ (keep).

Output becomes $H_{\ell+1,t}=H_{\ell,t}+1\cdot1.67\cdot2 = H_{\ell,t}+3.34$.

Across many steps the expected added value is $0.6\cdot1.67\cdot2 \approx 2$, matching the full‑layer contribution.

The scaling $r_{\ell,t}^{\text{train}}$ restores the average magnitude of the residual, so the model sees a consistent signal despite stochastic skipping.

Because the mask is applied per‑sample, only the active subset needs the expensive $f_{\ell}$ computation, yielding a FLOP reduction proportional to the dropout rate $p_{\ell,t}$.

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

How does Layer Dropout differ from standard neuron‑wise dropout?

Standard dropout randomly zeroes individual neurons inside a layer, preserving the layer’s overall computation. Layer Dropout instead drops the entire residual branch, skipping the whole $f_{\ell}$ sub‑network for a given sample, which yields a larger compute saving and forces the model to rely on the shortcut path.

Hyperparameter Optimization

We tune key hyperparameters and set the training scaling to $1/\rho$ for stable transfer.

To avoid the “hyperparameter lottery” and fairly compare dropout configurations, we first optimize learning rate $\eta$, batch size $B$, and weight decay $\lambda$ for each dropout rate before any evaluation.

Because dropout changes the effective depth of a model, scaling the training activations by the inverse layer density ($1/\rho$) keeps each residual block’s contribution proportional to its depth, allowing hyperparameters to transfer across dropout rates.

Compute layer density: $\rho = 1 - p = 0.8$.

For $r_{\text{train}} = 1$, activations of the remaining $80\%$ of layers are unchanged.

For $r_{\text{train}} = 1.25$, each surviving layer’s activation is multiplied by $1.25$, restoring the overall signal magnitude.

Run a single training step with learning rate $\eta = 0.001$; the update magnitude with $r_{\text{train}} = 1$ is $0.8 \times 0.001 = 0.0008$, whereas with $r_{\text{train}} = 1.25$ it is $0.8 \times 1.25 \times 0.001 = 0.001$, matching the full‑depth case.

Thus the optimal $\eta$, $\lambda$, and $B$ found for $p=0$ also work for $p=0.2$ when using $r_{\text{train}} = 1/\rho$.

Scaling by $1/\rho$ equalizes the effective update strength across dropout rates, eliminating the need to retune hyperparameters for each rate.

How does this scaling differ from the standard dropout scaling used in regular (non‑layer) dropout?

Regular dropout typically scales activations by $1/\rho$ during training and leaves them unchanged at evaluation. For layer dropout, many prior implementations kept $r_{\text{train}} = 1$, which under‑scales the remaining layers. Our method restores the $1/\rho$ factor specifically for layer dropout, ensuring that the residual contribution matches the full‑depth network.

Dropout Granularity

Choosing the right granularity for dropout preserves accuracy while cutting compute.

Applying dropout at the wrong granularity can cripple a transformer’s ability to learn joint attention‑FFN dynamics, causing unnecessary loss spikes.

Instead of masking individual sub‑layers (attention or FFN) independently, the model can drop the entire transformer layer as a single unit, preserving the joint computation of attention + FFN.

Step 1: Sample $M_{1,t}^{\text{attn}}=1$, $M_{1,t}^{\text{ffn}}=0$ → only the FFN is dropped.

Step 2: Sample $M_{1,t}=0$ → the entire layer (both attention and FFN) is skipped.

Step 3: Forward pass computes loss; the Layer‑Dropout run preserves the attention‑FFN coupling, avoiding a mismatch that can increase loss.

When masks differ across sub‑layers, the model sees an incomplete computation graph, which can destabilize training; a single mask keeps the graph intact.

How does Layer Dropout differ from the classic stochastic depth used in CNNs?

Classic stochastic depth drops whole residual blocks but assumes each block is a single convolutional unit. In transformers, a “block” contains two distinct sub‑layers (attention + FFN). Layer Dropout treats the whole transformer layer as the unit, preserving the internal attention‑FFN partnership, whereas Sub‑Layer Dropout would break that partnership by dropping only one half.

Applying the same dropout mask to every sequence in a batch (per‑batch) is cheaper, but letting each sequence draw its own mask (per‑sequence) yields finer‑grain stochasticity and better loss.

Step 1: Per‑batch mask $M_{1,t}=1$ → both sequences see the layer.

Step 2: Per‑sequence masks $M_{1,t}[1]=0$, $M_{1,t}[2]=1$ → sequence 1 skips the layer, sequence 2 processes it.

Step 3: The loss computed over the batch reflects both a dropped and a kept instance, providing richer gradient signals.

Independent masks break the batch‑wise coupling, preventing the model from over‑fitting to a single deterministic pattern of dropped layers.

Why doesn’t per‑batch dropout always win when compute is the bottleneck?

Per‑batch dropout saves weight‑loading cost, but if the training pipeline is already compute‑bound (large batch size or long context), both per‑batch and per‑sequence drop the same number of FLOPs. The extra stochasticity of per‑sequence then yields a measurable loss improvement without extra compute.

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

Dropout Distributions

How to vary layer dropout across a model to keep training cheap without hurting quality.

Applying the same dropout rate to every layer wastes compute because early layers receive unnecessary noise. Varying dropout across depth lets the model keep cheap early representations while still regularizing deeper parts.

Instead of dropping the same fraction of units everywhere, we assign a layer‑specific dropout rate so that shallow layers stay mostly intact while deeper layers are more aggressively thinned.

Layer 0 receives rate $0/3 \times 0.2 = 0.0$.

Layer 1 receives rate $1/3 \times 0.2 \approx 0.067$.

Layer 2 receives rate $2/3 \times 0.2 \approx 0.133$.

Layer 3 receives rate $3/3 \times 0.2 = 0.200$.

The mean dropout is $(0 + 0.067 + 0.133 + 0.200)/4 \approx 0.100$, i.e. 10 % of units are dropped on average.

ILD concentrates regularization where representations are most abstract, letting early layers stay fully expressive while still achieving the same overall FLOPs reduction.

How does ILD differ from simply increasing the global dropout rate?

Increasing the global rate would thin every layer equally, harming early feature extraction. ILD keeps early layers untouched (rate 0) and only ramps dropout later, so the same average sparsity yields less degradation of low‑level information.

Dropout Scheduling

Defines how dropout rates vary over training steps and why the chosen schedule works.

Prior work reported that decreasing dropout schedules help large‑language‑model pre‑training, while increasing schedules aid fine‑tuning; recent studies claim both fail under modern training regimes.

The schedule maps each training step to a dropout probability, gradually changing regularization pressure much like dimming a light—early training sees full signal, later training sees more dropout.

Step $t=0$: $p_{\ell,0}^{\text{constant}} = 0.2$.

Step $t=1$: $p_{\ell,1}^{\text{constant}} = 0.2$.

Step $t=2$: $p_{\ell,2}^{\text{constant}} = 0.2$.

Step $t=3$: $p_{\ell,3}^{\text{constant}} = 0.2$.

Step $t=4$: $p_{\ell,4}^{\text{constant}} = 0.2$.

With a constant schedule the regularization strength never changes, so the model sees the same amount of stochastic depth throughout training.

**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.

How does an increasing dropout schedule differ from a typical learning‑rate schedule?

A learning‑rate schedule adjusts the step size of parameter updates, while an increasing dropout schedule raises the probability of dropping entire residual blocks, progressively reducing model capacity as training proceeds.

Inference Optimizations

Inference optimizations exploit layer‑dropout robustness to enable faster, flexible inference.

Layer dropout skips residual blocks during training, which makes the model tolerant to missing layers at inference.

Decreasing dropout schedules achieve the lowest validation loss for a fixed training FLOPs budget, beating dense baselines at 5 % FLOPs savings.

Table 3 shows that ILD + DTS at 5 % FLOPs savings for 503 M and 906 M models yields lower loss than the dense baseline.

Because layer dropout forces the model to operate correctly even when some layers are missing, we can safely skip layers or exit early at inference time without retraining.

How does the early‑exit mechanism here differ from a naïve early‑exit that simply truncates a dense model?

In a naïve truncation, the model has never seen missing layers, so the representations become unstable and loss spikes. With layer dropout pre‑training, the model has been exposed to random layer omissions, so the hidden states remain well‑behaved up to the chosen exit layer, yielding a smooth loss curve.

**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.

**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.

Layer dropout enables efficient early‑exit inference without sacrificing accuracy.

Scaling Analysis

Scaling analysis quantifies how layer dropout performance scales with compute and tokens‑per‑parameter.

We now examine whether the benefits of Improved Layer Dropout (ILD) persist when models are scaled to higher tokens‑per‑parameter (TPP) regimes and when used for post‑training inference tricks.

Validation loss stays within 0.5 % of dense baselines across model sizes up to 906 M parameters at high TPP.

Across 270 M, 504 M, and 906 M models, ILD + DTS maintains loss within 0.5 % of the dense baseline even as TPP increases.

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

The figure displays two line charts comparing Validation Loss against Training FLOPs for models of different sizes (1.8B and 3.9B parameters) under varying maximum dropout rates.

ILD maintains performance across varying compute constraints.

Large-Scale Evaluation

Scaling aggressive layer dropout to billion‑parameter models reveals speedups and FLOP savings.

Aggressive layer dropout (pmax = 0.99) yields up to 25 % training FLOP savings while preserving validation loss at the billion‑parameter scale.

Table 5 shows the 8.2B model with pmax = 0.99 matches the dense baseline’s validation loss despite using 25 % fewer FLOPs.

Across all three scales, higher dropout rates improve robustness: larger models tolerate pmax = 0.99 with minimal loss degradation, and the resulting sparsity translates into both training‑time FLOP reductions and inference‑time speedups.

**Table 1.** Impact of Dropout Configuration on Self-Speculative Decoding performance across different model sizes.

**Table 1.** Hyperparameters and results for different model sizes (1.8B, 3.9B, 8.2B).

Additional Results

Additional ablations quantify how each dropout component impacts validation loss and downstream performance.

This section reports a suite of ablations that answer the core question “does removing or altering this component hurt?” for each dropout design choice.

Increasing the maximum dropout to 0.4 in the Linear Distribution scheme reduces early‑exit validation loss.

Figure 10 a shows the 0.4 (ILD) curve achieving the lowest loss across all exit layers for 271 M, 503 M, and 906 M models.

At a fixed maximum dropout of 0.2, the Improved Layer Dropout (ILD) configuration yields the smallest validation loss.

Figure 11 b reports the 0.2 (ILD) line consistently below the Alternating and DTS variants.

Using an alternating dropout schedule with a higher maximum dropout (0.4) attains lower loss than a constant schedule for the same memory budget.

Figure 13 a shows alternating points (circles) dominate constant points (squares) in the loss‑vs‑memory scatter.

When focusing on configurations that skip layers, alternating dropout still outperforms constant dropout.

Figure 14 b highlights the alternating cluster achieving lower loss at comparable memory footprints.

For a maximum dropout of 0.2, the ILD configuration again provides the best validation loss among all alternatives.

Figure 15 c places the ILD (circle) curve below the constant and other variants across model sizes.

Scaling up to 8.2 B parameters while applying aggressive dropout (≈0.99) continues the trend of lower validation loss.

Figure 16 d shows the 0.99 (ILD, DTS) line achieving the lowest loss among all plotted configurations.

Models trained with 20 TPP and Layer Dropout retain or improve downstream task scores despite up to 30 % FLOPs savings.

Table A.3 reports comparable or higher scores on BBH, PIQA, and ARC‑e while the “Sav.” column shows up to 0.30 × reduction.

**Figure A.1.** Coordinate Check passing for non-uniform distribution

**Figure a.** Ablating different maximum dropout values for Linear Distribution.

**Figure b.** Ablating different dropout configurations for the same maximum dropout of 0.2.

**Figure A.2.** Comparison of early-exit validation losses for models trained with different dropout configurations. All models trained with 20 TPP. (c) Larger model sizes with aggressive maximum dropout rates.

**Figure a.** Ablating different maximum dropout values for Alternate Distribution.

**Figure b.** Ablating different maximum dropout values for Alternate Distribution, zooming on skipping configurations with lower losses.

**Figure c.** Ablating different dropout configurations for the same maximum dropout of 0.2.

**Figure d.** Larger model sizes with aggressive maximum dropout rates.

**Table A.3.** Downstream task performance benchmarks.

Limitations and Conclusion

We outline remaining challenges and summarize the broader implications of layer dropout.

Our hyperparameter transfer analysis shows that learning rate $\eta$ and weight decay $\lambda$ stay stable for modest dropout levels but drop for aggressive rates, and it only covers constant dropout schedules.

We limited our study to transformer‑level layer dropout; it remains an open question whether attention‑head or neuron‑level dropout can provide comparable depth‑wise robustness.

We did not compare our structured dropout against learned depth‑aware methods such as Mixture‑of‑Depths, leaving the trade‑offs between stochastic removal and dynamic routing for future investigation.

Our experiments reach 8.2 B parameters, but we have not examined how aggressive dropout interacts with alternative architectures like Mixture‑of‑Experts or non‑transformer models.

We lack comprehensive scaling laws that predict the maximum dropout rate $p_{\text{max}}$ achievable without harming accuracy relative to a dense baseline.

Quantifying inference benefits—early‑exit loss, intermediate‑layer skipping, adapter loss, and self‑speculative decoding speedup—as tokens‑per‑parameter (TPP) grows is another open direction.

Overall, we demonstrate that layer dropout is a practical, non‑invasive tool for large‑scale pre‑training, offering flexibility without sacrificing performance.

By systematically varying dropout distributions and employing an increasing‑distribution across layers together with a decreasing schedule, we preserve or improve dense‑baseline accuracy while cutting training FLOPs and enabling zero‑shot inference elasticity.

Our “Hero Runs” on 1.8 B, 3.9 B, and 8.2 B models show that larger architectures naturally tolerate higher dropout, achieving up to 1.55× speedup in zero‑shot inference.

The training curriculum of progressively increasing effective capacity can be extended to other dimensions—model width, quantization bit‑widths, or unstructured sparsity—using analogous spatial distributions and temporal schedules.

Future work should derive optimal maximum dropout rates for given model scales and data budgets, and explore learned depth‑aware mechanisms that dynamically select skip paths.

We note that while the work advances machine‑learning capabilities, we do not highlight specific societal impacts.

Glossary

Key dropout terms and schedules used throughout the paper.

Alternating Layer Distribution (ALD) – dropout is applied on alternating layers, mathematically $p_{\ell,t}^{\text{ALD}} = p_{\text{max}} \cdot \mathbf{1}_{\ell \equiv 1 \pmod{2}}$.

Constant Time Schedule – the dropout rate stays fixed during training, $p_{\ell,t}^{\text{const}} = p_{\ell}^{\text{dist}}$.

Decreasing Time Schedule (DTS) – dropout starts maximal and linearly decays to zero, $p_{\ell,t}^{\text{DTS}} = p_{\ell}^{\text{dist}} \left(1 - \frac{t}{T}\right)$.

Feed‑Forward Networks (FFN) – the sub‑layers within transformer blocks that apply two linear transforms with a non‑linearity in between.

Increasing Layer Distribution (ILD) – dropout probability grows linearly across layers, $p_{\ell,t}^{\text{ILD}} = \frac{\ell}{L-1} \cdot p_{\text{max}}$.

Increasing Time Schedule (ITS) – dropout rate ramps up linearly over training steps, $p_{\ell,t}^{\text{ITS}} = p_{\ell}^{\text{dist}} \frac{t}{T}$.

Layer Dropout – entire transformer blocks are randomly skipped during a forward pass (see §2).

Stochastic Depth – another name for Layer Dropout, used interchangeably in the literature.

Sub‑Layer Dropout – dropout applied independently to the attention and FFN sub‑layers within a transformer block.

Read the original paper

Open the simplified reader on Paperglide

Browse all simplified papers