SUBLLM新架构：文本下采样机制革新大语言模型效率_44页_4mb

2024 AI+研发数字峰会总结

核心内容

本次2024 AI+研发数字峰会上，小米大模型团队的王全东介绍了其提出的SUBLLM新架构，该架构旨在优化大语言模型（LLM）的资源使用效率，同时保持模型性能不变。SUBLLM通过引入下采样（Subsampling）、**上采样（Upsampling）和旁路（Bypass）**模块，实现对长文本处理的高效优化，适用于训练和推理场景，显著降低计算和显存成本。

主要观点

1. 长文本模型的技术挑战

• 长文本需求旺盛：包括多人会议摘要、行业报告分析、学术论文处理、长文写作和翻译等。
• 模型结构限制：传统的Decoder only Transformer结构（如LLaMA）由于attention的平方复杂度，导致训练和推理成本高。
• 训练成本主要集中在8k预训练阶段：随着模型窗长增加，训练和推理资源消耗呈指数级增长。
• 现有方法的局限性：如投机解码、动态稀疏attention、Kv cache压缩等方法在推理加速方面表现良好，但训练加速效果有限。

2. SUBLLM架构设计

• 架构目标：
  • 优化资源使用，保持模型能力不变。
  • 区分重要token与不重要token，将主要算力分配给重要token。
  • 兼容主流attention-based模型生态，便于广泛应用。
• 设计灵感：来自语音识别领域的下采样技术，如Zipformer模型，通过减少冗余信息保留关键内容。
• 模块组成：
  • 下采样模块：通过Score层评估token重要性，使用TopK选择保留token。
  • 上采样模块：通过加权求和恢复原始序列长度，确保token判别能力。
  • 旁路模块：在下采样前和上采样后加权求和，增强训练的收敛性和稳定性。
• 关键创新点：
  • 通过减法操作使token选择可导。
  • 通过位置编码下采样解决训练和推理不一致性问题。
  • 兼容性高，支持主流attention机制（如Llama、Flash Attention等）。

3. 实验结果

• 训练加速：SUBLLM在不同窗长下（2k、4k、8k）均实现了34%的训练加速，显存减少10GB/GPU。
• 推理加速：在8k窗长下，推理加速比达到50%，显存减少1.35GB。
• 性能保持：在few-shot学习任务中，SUBLLM保持与LLaMA相当的性能，同时有效降低资源消耗。

4. 与其他方法的对比

• SUBLLM优势：
  • 同时支持训练和推理加速。
  • 降低显存成本。
  • 兼容主流模型结构。
• 对比Mixture-of-Depths (MoD)：
  • MoD通过间隔block筛选token，实现**50%**加速。
  • SUBLLM在训练和推理中分别实现**34%和50%**加速，性能保持良好。

关键信息

• SUBLLM结构：采用嵌套式模块设计，包括下采样、上采样和旁路模块。
• 参数量：SUBLLM仅比LLaMA多8192个参数，实现显著性能提升。
• 下采样保留比例：通过实验确定最优保留比例为75%，在训练和推理中保持有效性能。
• 应用前景：SUBLLM适用于长文本场景（如200K窗长），也具备在多模态模型（如VLM）中应用的潜力。

未来研究方向

• 通用性提升：探索不同tokenizer对压缩率的影响，确保SUBLLM结构的通用性。
• 长文本场景优化：研究在200K窗长下的最小保留比例，进一步优化资源使用。
• 多模态扩展：探索SUBLLM在视觉标记处理中的应用，提升高分辨率图像处理效率。

总结

SUBLLM通过创新性的下采样、上采样和旁路模块，显著提升了大语言模型在长文本处理中的资源使用效率，同时保持了模型性能。其在训练和推理中的加速效果分别为34%和50%，并有效降低了显存成本。未来，该架构有望在更广泛的场景中应用，包括长文本处理和多模态模型优化，进一步推动AI研发的降本增效。