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

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 LLM 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, LLM can achieve lower or similar validation loss while saving upto 20\% 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.4x inference speedup with negligible accuracy loss. All pre-training experiments were run on Cerebras CS-3 systems.

Layer dropout (i.e., making skipping layers randomly during training) was common in training neural networks and transformers. However, in the age of large language models (LLMs) it disappeared from pretraining recipes, as it was shown to harm their accuracies. In this paper, we show that layer dropout could work and we find configurations that minimize - if not eradicate - their accuracy degradations. Not only that, we show that it saves FLOPs during training, leading up to 20% reduction in training FLOPs, and make the model robust to layer-wise compression techniques during inference, leading to upto 1.43x speedup.