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

한 줄 요약

레이어 드롭아웃을 최적화하여 LLM의 학습 및 추론 효율성을 동시에 향상시킨다.

핵심 기여도

최적의 레이어 분포, 시간 스케줄, 옵티마이저 하이퍼파라미터를 통해 레이어 드롭아웃이 동일한 FLOPs에서 낮은 손실을 달성함.
2400개 이상의 실험을 통해 271M~8.2B 파라미터 모델에서 25%의 학습 FLOPs 절감과 1.5× 추론 가속을 달성함.
레이어 드롭아웃이 추론 단계에서 이른 종료, 중간 레이어 건너뛰기, 자기 추론 해독 등에 효과적임을 입증함.
레이어 드롭아웃을 배치 단위와 시퀀스 단위로 비교하여 후자가 더 낮은 손실을 보임.

핵심 아이디어

기존 연구에서는 레이어 드롭아웃이 대규모 언어 모델에서 성능 저하를 유발한다고 보고되었으나, 본 연구는 이를 하이퍼파라미터 및 스케줄링을 통해 완화할 수 있음을 입증한다. 특히, 레이어 드롭아웃은 전체 트랜스포머 블록을 건너뛰는 구조적 희소성을 유도하여 FLOPs를 줄이며, 추론 시에도 유연한 실행이 가능하다는 점에서 기존의 활성화 드롭아웃과 구분된다. 연구는 레이어 드롭아웃을 **배치 단위**(per-batch)와 **시퀀스 단위**(per-sequence)로 적용하는 방식을 비교했으며, 후자가 더 낮은 손실을 보였다. 또한, 레이어 드롭아웃의 분포를 **균일**(uniform), **선형 증가**(ILD), **대체 적용**(ALD) 방식으로 정의하고, 이 중 선형 증가 방식이 추론 효율성과 성능 사이의 균형을 잘 유지함을 밝혔다.

기술적 접근법

**모델 아키텍처**: 디코더-온리 트랜스포머, ALiBi 포지셔널 임베딩, 제곱 ReLU 활성화, Llama3 어휘 사용.
**하이퍼파라미터 최적화**: 학습률, 배치 크기, 가중치 감소를 레이어 드롭아웃 비율에 따라 조정.
**레이어 드롭아웃 분포**: uniform, ILD, ALD 방식 비교.
**드롭아웃 스케줄**: 학습 단계에 따라 점진적으로 증가하는 분포와 감소하는 스케줄 적용.
**실험 환경**: Cerebras CS-3 시스템에서 2400개 이상의 실험 수행, 최대 160B 토큰 데이터셋 사용.

주요 결과

**학습 효율성**: 동일 FLOPs 조건에서 레이어 드롭아웃이 기존 밀집 모델보다 낮은 손실을 보임.
**FLOPs 절감**: 주어진 학습 단계에서 25%의 FLOPs 절감을 달성함.
**추론 가속**: 이른 종료, 중간 레이어 건너뛰기, 자기 추론 해독을 통해 1.5× 추론 속도 향상.
**모델 규모**: 271M~8.2B 파라미터 모델에서 일관된 성능 개선 관찰됨.

의의 및 한계

**의의**: 레이어 드롭아웃이 대규모 언어 모델에서 학습 및 추론 효율성을 동시에 향상시킬 수 있음을 입증.
**한계**: 드롭아웃 비율이 모델 크기와 상호작용하는 방식은 명확히 규명되지 않았으며, 세부적인 하이퍼파라미터 조합에 따라 결과가 달라질 수 있음.
**추론 유연성**: 레이어 드롭아웃은 추론 시 다양한 최적화 기법과 호환되지만, 실제 배포 환경에서의 성능은 추가 연구가 필요함.

실용적 활용

대규모 언어 모델의 학습 및 추론 과정에서 컴퓨팅 자원을 절감할 수 있어, 클라우드 기반 모델 훈련 및 실시간 추론 시스템에 적용 가능.
특히, 다양한 계산 예산에 따라 유연하게 실행 가능한 모델을 요구하는 산업 분야(예: 모바일, 임베디드)에 유용함.
추론 단계에서의 이른 종료 및 레이어 건너뛰기는 실시간 응답이 필요한 어플리케이션(예: 챗봇, 번역)에 적합함.