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
ICML 2026 Proceedings of the
cs.AI
2026-09-04
在 2400 多次实验里,按深度递增、训练过程递减来丢整层,同位 FLOPs 可以更低 loss。8.2B 上省 25% 训练计算,自投机解码加速 1.55 倍。
Layer dropout(也叫 stochastic depth)在视觉和早期 BERT 里能加快训练、提高正确率,还让模型对少层推理不那么脆。LLM 预训练把 dropout 几乎删光了:数据按单 epoch 喂,激活 dropout 还被报告会伤效果。layer dropout 和激活 dropout 不是一回事,整层跳过能换成接近线性的 FLOPs 下降,也能在推理时做早退、跳层、自投机解码。缺的是一张在现代单 epoch、大 token 预算下怎么配才不掉点的地图。
Cerebras 在 CS-3 上跑了 2400 多次预训练,模型从 271M 到 8.2B,数据到 160B token,把这张地图补上。
Decoder-only,Celerity 结构:ALiBi、squared ReLU、Llama3 词表。Layer dropout 按序列独立采样整层(注意力和 FFN 共用同一掩码),被跳过的序列不跑该层,训练 FLOPs 近似按丢弃率下降。
关键缩放是训练时残差乘 1/ρ(ρ 是层密度),推理时乘 1。这个选择让学习率、batch、权重衰减能在不同丢弃率之间迁移;乘 1 则过不了坐标检查。超参加完再比粒度、深度分布和时间表,计算最优预算按 20 tokens-per-parameter。
深度分布比了均匀、隔层(ALD)和从浅到深线性增加(ILD)。时间表比了恒定、递增和递减(DTS)。推荐组合是 ILD+DTS:浅层少丢、深层多丢,训练开始噪声大、结束时回到满深度。平均训练丢弃率 P̄ 就是总 FLOPs 节省;ILD+DTS 下 P̄ = 0.25 pmax。
同样 5% FLOPs 节省时,ILD+DTS 在 503M 和 906M 上验证 loss 低于或持平稠密基线(906M: 1.951 vs 1.953)。递增时间表明显更差。整层丢优于把注意力和 FFN 分开丢;按序列丢优于按 batch 丢。
大模型可以更狠。1.8B 用 pmax=0.6(省 15%)验证 loss 1.836,对稠密 1.849;3.9B 用 pmax=0.8(省 20%)是 1.745,对稠密 1.732,略差一截;8.2B 用 pmax=0.99(省 25%)验证 loss 1.663,表里没有同尺寸稠密对照。3.9B 训练开头有效深度只有约 0.6L,最后一层 80% 的时间被跳过,仍能收敛。同位 FLOPs 下,带 dropout 的曲线在大部分训练里更低。
推理侧,平均训练丢弃率能预测零样本早退和跳层的稳健程度。3.9B 自投机解码从稠密的 1.02× 升到 1.54×;8.2B 到 1.55×。隔层跳过时,3.9B 的 loss 从稠密的 6.446 降到 2.129。早退适配器(Balcony,骨干冻结)加在 dropout 预训练模型上,各尺寸的早退 loss 都低于稠密骨干。
高 TPP 下 ILD+DTS 相对稠密的 loss 劣化大约稳住在 0.50% 量级,作者据此认为这条配方能跟到更长的预训练。
这是在说 dropout 不该当过拟合工具来用,而该当深度课程:先逼模型在更浅的有效网络里学,再把深度还回去。对预训练,25% 的非嵌入 FLOPs 不是边角料;对部署,同一套权重能早退、能跳层、能当自己的草稿模型,不必再训一套弹性架构。ILD 更利于正确率和早退,ALD 更利于非连续跳层,按线上需求选。
攻击性丢弃率下最优学习率和权重衰减会下降,超参迁移没有完全跟上。没有和 Mixture-of-Depths 这类学出来的跳层比,也没有做到 MoE 或非 Transformer。pmax 随规模怎么取还没有定律;推理收益随 TPP 怎么变也没量化。8.2B 的 25% 节省没有同表稠密对照,3.9B 在 20% 节省时验证 loss 还略高于稠密。实验全在 Cerebras 上,按序列丢层在 GPU 上能省多少墙钟,要另测。