Don't Drop Dropout: 优化层稀疏性以实现高效 LLM 训练与推理

TL;DR 系统研究 layer dropout 在 LLM 预训练中的最优配置，证明同 FLOPs 下 loss 更低或节省 25% 训练 FLOPs，并支持 early exit 等推理优化，实现 1.5 倍加速。

问题

问题背景

大模型预训练的核心诉求是在不损失精度前提下降低训练与推理成本，条件计算与结构化稀疏成为当前热点方向。

现有方法局限

Layer dropout（随机深度）早在视觉 Transformer 中验证过加速训练与提升鲁棒性，但在 LLM 预训练配方中几乎消失。已有工作仅零散报告“dropout 会损害精度”，缺乏对 dropout 分布、时间调度、优化器超参 的系统量化。此外，没有覆盖 271M 到 8.2B 参数、160B token 量级的缩放分析，导致训练侧收益无法预测；训练时感知深度弹性的方法又往往改动架构或增加额外模块，难以直接复用标准 Transformer 训练栈。

为什么难/重要

随机深度会改变每步有效网络深度，造成梯度与优化器状态统计偏移，简单套用固定分布或静态调度会在大规模训练中放大损失。同时，为获得推理加速（早退、跳层、自投机解码）需要模型具备零样本深度弹性，这要求训练阶段精心设计稀疏模式。业界对每 FLOP 效率极致关注，任何可节省 25% 训练 FLOPs 且保持 loss 的方法都有直接落地价值。

行业类比

类似 MoE 中通过路由实现条件计算以摊薄每 token 算力，layer dropout 试图用随机层稀疏换计算效率，但挑战在于保证稠密主干的知识密度不下降。

核心洞察

• 层 dropout 在大规模 LLM 预训练中可以恢复为有效的正则化与算力节省手段，而非必需丢弃的过时技术。已有工作报道 dropout 会损害精度，但本研究表明通过系统调优层丢弃分布、时间调度和优化器超参数，在相同训练 FLOPs 下层 dropout 能获得更低验证 loss，或在相同训练步数下节省最多 25% FLOPs 且 loss 相当或更低。这推翻了“大模型预训练不应使用 dropout”的默认假设，并为实践者提供了可复现的配置方案。
• 预训练阶段使用层 dropout 隐式地让模型获得深度弹性，可直接转化为推理加速，无需额外的训练或蒸馏流程。由于训练时随机跳过中间层，每个子网络都被优化，因此训练后能支持零样本 early exit、中间层跳过，配合轻量 early exit adapters 或 self-speculative decoding 实现最高 1.5 倍推理加速，准确率损失可忽略。这与专为深度弹性设计的训练方法形成对比，后者通常需要修改训练目标或额外蒸馏阶段，而本文方法复用标准预训练流程即可获得此类能力。

方法

方法概览

输入：Transformer 架构的 LLM 预训练任务，包括模型参数、训练数据集、优化器状态与超参数配置。

关键模块

1. 层 Dropout 分布优化
   通过模型粒度与张量粒度的分析，确定各层被随机跳过的概率分布（例如前层、中层、后层的不同丢弃率），使训练信号在深度上更有效地传递。

2. Dropout 调度策略
   设计训练过程中的 dropout 概率动态调整方案（如 warmup、线性衰减或分段常数），平衡训练早期的探索能力与后期的收敛稳定性。

3. 优化器超参数调整
   针对部分层被跳过导致的梯度路径变化，调整学习率、动量等超参数，并引入 scale transfer 机制，确保剩余层能充分补偿被丢弃层的信息。

4. 推理加速模块
   利用训练得到的层弹性，在推理时直接应用 early exit、intermediate-layer skipping 和 self-speculative decoding，无需额外训练即可获得加速。

输出：在相同训练 FLOPs 下达到更低验证损失，或在相同训练步数下节省最多 25% 的 FLOPs；同时模型保留零样本推理加速能力（最高 1.5×），精度损失可忽略。

与同类方法的差异：以往工作通常将 layer dropout 视为 LLM 预训练中的精度退化因素而弃用，或仅作为简单正则化手段；本研究首次系统优化 dropout 的分布、调度与优化器配置，使其在大规模训练中同时提升效率与精度，并直接解锁推理阶段的深度弹性加速。

实验

实验设计

作者在 Cerebras CS-3 上开展 2400+ 组预训练实验，覆盖模型规模 271M 至 8.2B 参数、数据量最高 160B tokens。系统调查 layer dropout 的多个关键维度：dropout 粒度（模型级 / 张量级）、dropout 分布、时间调度，以及优化器超参数（如 lr 与 batch size 的协同）。通过控制变量与缩放实验，确立可迁移至大规模训练的最佳配置。

关键发现

• 在 相同训练 FLOPs 下，采用最优 layer dropout 可获得更低验证损失。
• 对给定训练步数，模型在节省最高 25% 训练 FLOPs 的同时，保持损失不变或更低。
• 训练得到的深度弹性网络支持 零样本 early exit、中间层跳过 与 self-speculative decoding，推理加速最高 1.5x，精度损失可忽略。

与基线对比

传统 LLM 预训练普遍弃用 dropout，往往因早期工作报告其有害。本文证明：只要对 dropout 分布、调度与优化器做系统调优，layer dropout 不仅不损害性能，反而在算力预算约束下带来实质收益。与无 dropout 或随机配置 dropout 的基线相比，深度弹性带来的推理加速是纯训练侧方法无法提供的额外优势，为 训练-推理协同优化 提供了新范式。

行业影响

落地场景

Layer dropout 可直接用于 LLM 预训练与推理优化。预训练服务商（云厂商、AI 实验室）可将它作为默认训练技巧，在同等训练步数下节省最多 25% 的训练 FLOPs；推理部署方（如 LLM API、对话机器人、边缘端侧服务）可利用训练得到的深度弹性，通过 early exit 或 intermediate-layer skipping 将解码速度提升至 1.5 倍，且精度损失可忽略。

商业价值

• 降本：训练环节减少 FLOPs 意味着更低的 GPU 租赁或采购成本；推理加速可提升单卡吞吐，降低每次请求的算力开销。
• 体验提升：更低的延迟直接改善用户交互感受，有助于提升留存与转化；同时，同一模型可动态调整深度以适应不同硬件等级，扩大部署范围。
• 质量保障：论文证实层丢弃在最优配置下不损害验证损失，消除了此前对 dropout 降低精度的顾虑，使工程团队更有信心采用。

与现有产品/工作流的接口

• 训练集成：需在优化器中加入针对 dropout 的适配（如调整 beta 参数），并配置 dropout 分布与时间表；可封装为 PyTorch/JAX 插件，接入 Hugging Face Trainer 或 Megatron-LM 等框架。
• 推理集成：early exit 需在特定层添加轻量 exit head，并在推理引擎（如 vLLM、TensorRT-LLM）中实现条件退出逻辑；intermediate-layer skipping 可作为模型变换插件动态加载。
• 与自推测解码结合：训练时采用 layer dropout 的模型可直接作为 draft model 进行 self-speculative decoding，无需额外蒸馏或训练。

具体落地案例

• 电商智能客服：大型电商平台使用 LLM 提供实时对话服务。采用 layer dropout 训练后，在推理时启用 early exit，简单查询在浅层退出，延迟从 500ms 降至 300ms，提升用户体验的同时降低 GPU 集群负载。
• 内容平台自动摘要：视频或文章平台每天处理海量内容生成摘要。训练时使用 layer dropout 削减 25% 训练成本；推理时采用 intermediate-layer skipping 提高吞吐，在不损失摘要质量的前提下，单卡吞吐提升 50%。

局限

• **验证规模有限**：论文最大模型为 8.2B 参数、数据集 160B tokens，虽声称可扩展到更大训练体制，但未在数十 B 或百 B 级模型、数 T tokens 上验证。当前主流 LLM（如 70B、405B）的超参数（层 dropout 分布、时间表、优化器设置）可能需要重新搜索，结论的普适性仍有待确认。
• **超参数搜索成本高**：论文通过超过 2400 次实验找到最优层 dropout 配置，但搜索过程依赖特定模型架构与数据集。实际应用中，工程师若为不同模型重新搜索，会带来显著额外计算开销，且论文未提供自动搜索策略或通用启发式规则，限制了方法的可移植性。
• **推理优化依赖训练阶段改动**：early exit、intermediate-layer skipping、self-speculative decoding 等收益均要求训练时应用层 dropout，已部署模型无法直接受益。此外，early exit adapters 等后训练微调会引入额外延迟和内存开销，可能部分抵消推理加速效果。