王全东-SUBLLM新架构文本下采样机制革新大语言模型效率

长文本模型的技术挑战

长文本模型在多人会议摘要、行业报告、新闻摘要等领域需求旺盛，但面临训练成本高、推理速度快等挑战。传统Decoder only Transformer模型（如Llama）受限于attention的平方复杂度，导致训练和推理成本随窗长增长显著。现有扩展方法如数据工程、位置编码调整虽能提升窗长，但训练成本仍高。此外，长文本模型训练和推理的infra开发（如DeepSpeed Ulysses、Ring-attention）存在并行度限制和兼容性问题，推理成本随窗长指数级增加。

SUBLLM架构

为优化资源使用，小米大模型团队提出SUBLLM新架构，目标在保持模型能力不变的前提下降低训练和推理成本。该架构受语音领域下采样启发，通过区分重要token和不重要token，将主要算力分配给前者。SUBLLM包含：

• Learnable Subsampling Module：通过去除不重要token缩短序列，使用Score层衡量token重要性，实现序列采样。
• Upsampling Module：恢复序列长度，使token选择可导。
• Bypass Module：对下采样前和上采样后序列加权求和，增强训练收敛性和稳定性。

主要实验结果

在SlimPajama数据集上，SUBLLM 1.3B模型仅增加8192个参数，实现：

• 训练加速比26%，最大加速比31%，显存减少10%。
• 推理加速比37%，窗长越大加速越明显（8k窗长加速50%）。
• 预训练valid loss持平，Few-shot分数持平。

分析与讨论

• 预训练加速：SUBLLM有效降低训练成本，且窗长越大加速越明显。
• 推理加速：推理加速比与窗长正相关，8k窗长加速50%。
• 模型结构：适用于不同优化器，最优下采样次数为2次，保留比例75%时valid loss最低。
• 下采样有效性：下采样前后的attention和保留index分布接近，证明下采样有效性。

总结与展望

SUBLLM贡献：

1. 提出结合下采样、上采样和旁路模块的新架构，动态分配资源给重要token。
2. 提出token序列子采样方法，通过Score层测量重要性并控制分数值分布。
3. 相比Llama，SUBLLM训练加速34%，推理加速50%，显存降低，模型能力保持。

未来研究方向：

1. 不同tokenizer的压缩率差异，确保SUBLLM结构的通用性。
2. 在200K窗长等长文本场景的应用，探索最低保留比例。
3. 应用于多模态模型中的视觉标记冗余处理，提升VLM效率。