PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2609.05275v1cs.AI

TL;DR

Modern LLM pretraining lacks comprehensive evidence on whether layer dropout can preserve accuracy while improving training and inference efficiency. This paper systematically studies dropout distributions, schedules, and optimizer settings, finding that properly configured layer dropout reduces training cost and enables depth-elastic inference across large-scale regimes.

  • Problem

    Layer dropout’s role in modern large-scale LLM pretraining remained unresolved because prior evidence was fragmented and reported degradations could reflect suboptimal configurations.

  • Method

    The paper systematically evaluates optimizer hyperparameters, depth-wise dropout distributions, temporal schedules, and post-training inference optimizations across large LLM training experiments.

  • Results

    Properly configured layer dropout achieves competitive or lower validation loss while saving up to 25% of training FLOPs and enabling inference speedups up to 1.55×.

  • Takeaways & Limitations

    Increasing dropout across depth with a decreasing schedule over time provides a non-invasive route to efficient training and zero-shot inference elasticity.

  • Takeaways & Limitations

    Hyperparameter transfer was primarily validated for constant dropout schedules, and other dropout granularities were not studied.

Abstract

from arXiv · show

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.

1 Introduction

Layer dropout offers structured training sparsity and depth-wise inference flexibility, but its value and configuration in modern large-scale LLM pretraining remained unresolved. This study evaluates the design space systematically and reports efficiency gains across model and data scales.

  • Motivation: Layer dropout skips entire transformer blocks, reducing active training FLOPs and supporting zero-shot depth-wise inference optimizations.Unlike fine-grained dropout, block skipping provides structured sparsity that can translate into wall-clock and deployment benefits.
  • Research gap: The study addresses fragmented evidence on whether layer dropout preserves accuracy in modern LLM pretraining and how it should be configured.Prior degradations may reflect suboptimal schedules or hyperparameters rather than fundamental limitations.
  • Study design: Across 2400+ runs, the authors vary optimizer hyperparameters, depth-wise sparsity distributions, and temporal dropout schedules.Experiments span 271M to 3.9B parameters and up to 116B tokens in the cited study scope.
  • Findings: Properly configured layer dropout achieves validation loss competitive with or superior to dense baselines while reducing training FLOPs by up to 25%.The recommended configuration uses an increasing distribution across depth paired with a decreasing schedule over training.
  • Findings: Layer dropout also supports depth-elastic inference, with the average training dropout rate predicting robustness to early exit and layer skipping.The paper positions this as a training-aware route to inference flexibility without changing the pretrained architecture.

4 Preliminary

The preliminary formulation applies stochastic block execution to transformer residual layers during training and disables dropout during inference. Efficient implementation executes each layer only on sequences whose Bernoulli mask is active, saving a portion of layer FLOPs.

  • Formulation: Layer dropout applies a stochastic mask to transformer residual-layer computation at each training step.The formulation covers layer index ℓ and training step t, with dropout applied to transformer operations.
  • Notation: The activation tensor H has dimensions B×S×d, representing batch size, sequence length, and hidden dimension.The feed-forward network is denoted FFN in the transformer-layer formulation.
  • Training operation: A Bernoulli mask Mℓ,t retains each sequence with probability 1−pℓ,t, while rtrain rescales the retained computation during training.The scaling factor is chosen separately from the mask and is discussed in the hyperparameter analysis.
  • Implementation: An efficient implementation evaluates the layer only for sequences with active masks rather than computing all sequences and masking afterward.This implementation saves a portion p of the layer’s training FLOPs.
  • Inference operation: During inference, dropout is disabled and a distinct inference scaling factor is used.The training and inference operations therefore use separate scaling choices.

5 Hyperparameters

The hyperparameter study treats dropout as a change in effective depth and jointly tunes optimization and scaling choices to enable stable comparisons across dropout rates. It identifies rtrain = 1/ρ as the key scaling rule for hyperparameter transfer.

  • Hyperparameter search: The authors jointly optimize learning rate, batch size, and weight decay for each dropout rate before comparing configurations.This procedure is intended to avoid the “hyperparameter lottery” and establish a strong baseline.
  • Scaling choices: Layer density is defined as ρ = 1−p, and the study evaluates training and evaluation scaling factors across prior conventions.Different dropout implementations use different rtrain and reval choices, motivating an explicit comparison.
  • Scaling analysis: The coordinate-check analysis finds that rtrain = 1 fails stable scaling across dropout rates, whereas rtrain = 1/ρ largely satisfies it.This scaling compensates for reduced effective depth during training.
  • Hyperparameter transfer: With rtrain = 1/ρ, optimal learning rate, weight decay, and batch size remain constant across dropout rates in the transfer test.The authors therefore adopt this scaling and transfer rule for subsequent experiments.

6 Dropout Granularity

The study compares dropout granularity in transformers and finds that dropping whole layers and sampling masks per sequence provide better accuracy than finer or batch-level alternatives.

  • Model granularity: Layer Dropout drops entire transformer layers, whereas Sub-Layer Dropout samples attention and FFN masks independently.These correspond to different mask granularities within each transformer layer.
  • Model granularity: Whole-layer dropout achieves higher accuracy than separately dropping attention and FFN sub-blocks.The comparison evaluates layer and sub-layer dropout across model sizes at 20 tokens per parameter.
  • Tensor granularity: Per-sequence dropout can provide compute speedups similar to per-batch dropout when training is compute-bound.Both methods save the same compute FLOPs, although per-batch dropout avoids loading skipped-layer weights.
  • Tensor granularity: Per-sequence dropout produces lower loss than per-batch dropout at the same layer dropout rate.Per-batch masks are shared across sequences, while per-sequence masks are sampled independently for each sequence.

7 Dropout Configurations

Layer dropout is evaluated through depth-wise distributions and time schedules under matched training FLOPs. Non-uniform distributions and decreasing schedules provide the strongest accuracy–compute trade-offs, with ILD increasingly favored as models scale.

  • Dropout distributions: Increasing Layer Distribution (ILD) starts dropout at zero in the first layer and increases linearly to pmax at the last layer.The study compares ILD with uniform and alternating layer distributions.
  • Dropout distributions: Non-uniform layer distributions consistently outperform uniform dropout at matched average dropout and training FLOPs.ALD is strongest at the smallest scale, but its advantage diminishes with scale while ILD’s improvement over uniform widens.
  • Dropout schedules: The mean training dropout averages sparsity across depth and time and represents total active training FLOPs savings.For the listed schedules, ¯P_uniform,ITS = 0.5pmax and ¯P_ILD,ITS = ¯P_ILD,DTS = 0.25pmax.
  • Dropout schedules: 5% FLOPs savings: ILD combined with a decreasing schedule achieves lower validation loss than the dense baseline for 503M and 906M models.This result demonstrates that fewer training FLOPs can outperform the dense baseline in those settings.
  • Dropout schedules: A decreasing schedule from the maximum rate to zero consistently achieves the highest accuracy across scales, outperforming constant and increasing schedules.The schedules are compared at fixed total active non-embedding FLOPs savings.
  • Recommended configuration: The recommended configuration increases dropout across layers and decreases it across time for the best accuracy at a given training FLOPs budget.Other schedules, including mid-training dropout, supervised fine-tuning, and continual pre-training, are left for future work.

8 Inference Optimizations

Layer dropout improves robustness to static early exit and intermediate-layer skipping, while enabling weight-frozen post-training methods and faster self-speculative decoding. The benefits depend on dropout configuration, with distinct trade-offs between early-exit and skip robustness.

  • Early Exit: Dropout-trained models retain lower loss under static early exit than models trained without dropout, including with decreasing schedules.Higher dropout improves early-exit loss, while uniform and ILD improve early exit but ALD does not.
  • Early Exit: 10% FLOPs savings: dropout configurations are compared against no-dropout models for early-exit validation loss across model sizes.The comparison uses models trained at 20 tokens per parameter.
  • Intermediate Layer Skipping: 20% training compute reduction: ALD maintains lower validation loss than ILD under intermediate-layer skipping.ALD also provides superior retention for non-contiguous skipping, whereas ILD offers stronger early-exit robustness.
  • Configuration Trade-offs: ALD excels at skip robustness, while ILD provides better base accuracy and early-exit robustness.The paper recommends choosing between them according to deployment needs.
  • Zero-Shot Elasticity: Layer dropout creates depth elasticity, allowing larger models to degrade gracefully toward smaller-model performance when layers are skipped.Average training dropout rate predicts zero-shot robustness to early exit and layer skipping without retraining.
  • Post-Training Methods: Weight-frozen post-training methods preserve the pretrained backbone while adding adapters or changing decoding to improve depth-wise inference flexibility.The study covers early-exit adapters and self-speculative decoding, while leaving routers for future work.
  • Early-Exit Adapters: Layer-dropout pretraining consistently improves post-training early-exit adapter losses over dense baselines across 270M, 503M, and 906M models.Higher dropout rates give adapters a stronger head-start at earlier layers without increasing pretraining memory or computational overhead.
  • Self-Speculative Decoding: Higher layer dropout produces higher self-speculative decoding speedups when the model searches for efficient draft-layer subsets.The study compares Bayesian optimization, genetic algorithms, hill climbing, and simulated annealing.

9 Scaling Analysis

The study extends its optimized layer-dropout recipe to higher tokens-per-parameter regimes to test whether its training benefits persist as models receive more data. ILD+DTS remains close to dense-baseline loss, indicating stable scaling behavior.

  • Scaling Motivation: The authors frame high-TPP scaling as a test of whether the optimized recipe generalizes beyond compute-optimal pretraining.Earlier compute-optimal experiments found that ILD+DTS could surpass dense-baseline validation accuracy.
  • Scaling Motivation: The scaling analysis targets how layer dropout’s regularization and structural benefits evolve as models exhaust inherent redundancy through prolonged training.It extends established model-size and data-budget scaling frameworks to high-TPP regimes.
  • Scaling Results: ILD+DTS loss degradation typically remains within ≈0.50% of the dense baseline at high tokens per parameter.The configuration maintains competitive performance as tokens per parameter increase.

10 Large-Scale Runs

Large-scale experiments apply aggressive layer dropout to models above 1B parameters. The results report competitive accuracy, lower loss at matched FLOPs, substantial training savings, and faster self-speculative inference.

  • Scaling Evaluation: Figure 8 summarizes validation loss across tokens per parameter for ILD with DTS against dense baselines.The figure is used to assess scaling performance across training-data regimes.
  • Training Efficiency: At matched training FLOPs, aggressive dropout achieves lower validation loss than dense training through most of training.The comparison is reported across larger model sizes at 20 tokens per parameter.
  • Inference Efficiency: 1.54× speedup: the 3.9B model improves self-speculative decoding, while the dense baseline produces a 1.02× regression.The result demonstrates an inference benefit at multi-billion-parameter scale.
  • Large-Scale Results: pmax = 0.99: aggressive layer dropout enables up to 25% cumulative training-FLOPs savings while largely preserving baseline generalization and enabling up to 1.55× zero-shot inference speedups.This is the paper’s reported multi-billion-parameter takeaway.

11 Limitations

The study identifies several boundaries for applying layer dropout, including incomplete hyperparameter transfer, untested granularities and architectures, and missing scaling analyses for dropout and inference benefits.

  • Hyperparameter Transfer: Hyperparameter transfer to aggressively high dropout rates remains incomplete, especially for decreasing schedules and high training-period budgets.Optimal learning rate and weight decay decrease at high dropout; transfer analysis mainly covered constant schedules.
  • Alternative Granularities: The study tests only transformer-level dropout, leaving attention-head and neuron-level dropout as open alternatives for structural resilience.
  • Learned Depth Optimization: The work does not compare stochastic layer dropout with learned depth mechanisms such as Mixture-of-Depths.
  • Cross-Architecture Generalization: Results have not established how aggressive layer dropout interacts with Mixture-of-Experts or non-transformer architectures.
  • Scaling Laws for Maximum Dropout: Comprehensive scaling laws for the maximum dropout rate that avoids accuracy degradation relative to dense baselines remain undeveloped.
  • Scaling Analysis for Inference Benefits: The study does not quantify how early-exit, layer-skipping, adapter, and self-speculative-decoding benefits vary as training-period dropout increases.

12 Conclusion

The paper presents layer dropout as a scalable way to improve training efficiency while preserving accuracy and enabling flexible inference. Its recommended increasing layer distribution and decreasing time schedule support zero-shot speedups that grow with model scale, while several broader scaling questions remain open.

  • Conclusion: An increasing dropout distribution with a decreasing schedule can maintain or surpass dense-baseline accuracy while reducing training FLOPs and enabling zero-shot inference elasticity.
  • Conclusion: Up to 1.55× inference speedups emerge in 1.8B, 3.9B, and 8.2B hero runs using the recommended recipe.
  • Future Directions: The findings suggest progressively increasing effective model capacity during training as a generalized curriculum beyond depth.
  • Future Directions: Future work should identify scale- and data-specific maximum dropout rates and investigate learned mechanisms that select skipping paths dynamically.

Appendix

The appendix defines the paper’s layer-dropout terminology, distributions, and temporal schedules. It distinguishes structural layer skipping from sub-layer dropout and describes decreasing schedules as a way to stabilize early training before final dense convergence.

  • Layer Distributions: Alternating Layer Distribution applies dropout to alternating layers while leaving the first layer undropped.
  • Temporal Schedules: Decreasing Time Schedule begins at maximum dropout and decays linearly to zero by the final training step.The paper states that this stabilizes early training while allowing final convergence in a dense state.
  • Layer Distributions: Increasing Layer Distribution starts at zero dropout in the first layer and increases linearly across layers to a maximum rate.
  • Temporal Schedules: Increasing Time Schedule starts at zero dropout and increases linearly to its maximum by the final training step.
  • Definitions: Layer Dropout, or Stochastic Depth, randomly skips entire transformer blocks during training, reducing effective network depth per forward pass.
  • Definitions: Sub-Layer Dropout applies layer dropout separately to attention and FFN sub-blocks, while FFN denotes feed-forward networks.
  • Definitions: Uniform Distribution sets the dropout rate equally across all layers, whereas a constant time schedule keeps dropout unchanged over time.

A Experimental Settings

The experimental appendix documents the model architectures, scaling and parameterization controls, activation-coordinate checks, and validation-loss comparisons used across dropout configurations.

  • Architectures: Table A.1 lists the model architectures used in the experiments.
  • Parameterization: Table A.2 summarizes SP, µP, and CompleteP controls together with layer-dropout rates and tunable width, depth, data-size, and dropout parameters.
  • Coordinate Checks: Coordinate-check plots compare activation Frobenius norms after merged attention and FFN residual streams in a 40-layer model after 10 training steps.The analysis uses CompleteP scaling for residuals.
  • Coordinate Checks: Figure A.1 presents coordinate checks for non-uniform layer-dropout distributions.
  • Downstream Evaluation: Table A.3 reports downstream task performance benchmarks.
  • Inference Evaluations: Figure A.2 compares early-exit validation losses across dropout configurations, with all models trained using 20 TPP.
  • Inference Evaluations: Figure A.3 compares intermediate-layer-skipping validation losses across dropout configurations, also using 20 TPP.
Loading 2609.05275v1…