Back to AI Research

AI Research

Don't Drop Dropout: Optimizing Layer Sparsity f... | AI Research

Key Takeaways

  • Don't Drop Dropout: Optimizing Layer Sparsity for Efficient LLM Training and Inference This research investigates the role of "layer dropout"—a technique tha...
  • 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.
Paper AbstractExpand

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.

Don't Drop Dropout: Optimizing Layer Sparsity for Efficient LLM Training and Inference
This research investigates the role of "layer dropout"—a technique that randomly skips entire transformer blocks during training—to improve the efficiency of Large Language Models (LLMs). While dropout was common in early neural networks, it has largely been abandoned in modern LLM pre-training due to concerns that it might degrade model accuracy. This study challenges that trend, providing a comprehensive framework to show that when configured correctly, layer dropout can significantly reduce training costs and enable faster inference without sacrificing performance. The same large language models question is explored in When Does Bigger Help? A Controlled..., which adds a research perspective.

Rethinking Dropout for Modern LLMs

The researchers argue that the perceived failure of dropout in modern LLMs is often due to suboptimal implementation rather than a fundamental flaw. By systematically testing over 2,400 training experiments, the study identifies that the key to success lies in how the model is scaled during training. Specifically, they found that using a scaling factor of $1/\rho$ (where $\rho$ is the density of the layers) ensures stable training and allows for consistent hyperparameter settings across different model sizes. This approach effectively resolves the "hyperparameter lottery," where researchers previously struggled to find the right balance between dropout rates and learning stability.

Improving Efficiency and Performance

The study demonstrates that layer dropout offers two major advantages. First, it improves training efficiency: by skipping layers during the training process, models can achieve similar or even better validation loss while saving up to 25% of the total computational effort (FLOPs). Second, it creates "depth-elastic" models. Because the model learns to function effectively even when some layers are skipped during training, it becomes naturally robust to similar shortcuts during inference. This allows for post-training optimizations like early exiting or intermediate-layer skipping, which can speed up text generation by up to 1.5x with negligible impact on accuracy. The same ai systems question is explored in Efficient Test-Time Adaptation through Human-AI Interaction, which adds a research perspective.

Best Practices for Implementation

The researchers provide clear guidelines for applying these techniques. They found that "Layer Dropout"—dropping the entire transformer block—consistently outperforms "Sub-Layer Dropout," where attention and feed-forward networks are dropped independently. Furthermore, they recommend a strategy of using a progressively increasing dropout distribution across the depth of the model, paired with a decreasing schedule over the course of training steps. These findings suggest that layer dropout is a powerful, underutilized tool that can make the development and deployment of large-scale language models significantly more sustainable and cost-effective. The same large language models question is explored in Rethinking On-Policy Distillation of Large Language..., which adds a research perspective. as detailed in the full paper on Arxiv

More from Franklin

Rethinking On-Policy Distillation of Large Language...researchWhen Does Bigger Help? A Controlled...researchClean Engineering, Unstable MeasurementresearchEfficient Test-Time Adaptation through Human-AI InteractionresearchCalibrated Enough to Know, Not Calibrated...researchThe Dually Flat Geometry of Planning...research

Comments (0)

No comments yet

Be the first to share your thoughts!