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

小米大模型团队演讲摘要

1. 报告主题

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

2. 主讲人

王全东

• 背景：小米大模型高级算法工程师、中科院声学所博士、美国访问学者
• 贡献：主导SUBLLM架构，获2024 CCF创新奖
• 学术成果：顶会论文十余篇、多次竞赛获奖

3. 技术背景

3.1 长文本模型痛点

• 结构：Decoder-only Transformer，attention复杂度平方增长
• 训练成本：窗口扩展时token需求激增（8k→1M对应2~15T训练tokens）
• 推理成本：长文本导致attention计算量级急剧上升
• 特殊结构方案：Infini-Transformer、MEGALODON等新型架构探索

3.2 优化方向

• 数据层：Data Engineering（Meta文献）
• 位置编码：YaRN（NTK-by-parts+温度控制）|PoSE（位置跳过训练）
• 训练基础设施：DeepSpeed-Ulysses | Ring-attention

4. SUBLLM架构

4.1 设计灵感

• 受语音识别领域“降采样”启发（如Conformer中的4-16倍帧率压缩）
• 类比思路：文本序列虽无语音时间属性，但文本内仍存在冗余信息

4.2 核心结构

• 采样模块（Learnable Subsampling）

  • 特征：score-based token重要性评估+动态保留率控制
  • 优势：解决位置编码不对齐问题，保持token分布一致性
  • 公式：Balancer（正分值token数量调控）

• 上采样模块（Upsampling）

  • 方法：合并采样token序列与原始序列
  • 创新：反向梯度使score层具有防过拟合能力

• 旁路模块

  • 功能：残差连接提高训练稳定性（跨通道加权）

4.3 技术对比

• vs MoD（Mixture-of-Depths）：
  • 实现方式不同（MoD静态剔除vs动态评分）
  • 资源分配差异（令牌级动态分时vs层级分配）

5. 实验结果

5.1 性能提升

• 原型对比：SUBLLM1.3B vs LLaMA1.3B
• 资源消耗：
  • 显存节约10+GB/GPU
  • 训练加速34%
  • 推理加速最高50%

5.2 能力保持

• 预训练：valid loss持平，达100%基线水平
• 少样本测试：分数完全一致
• 窗长适用性：8k→200k文本均有效

6. 未来方向

• 多tokenzer适配：不同压缩率下的保留策略优化
• 超长文本专项：200K输入下的最小保留比例模型
• 多模态扩展：VLM中视觉token冗余处理的应用