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

SUBLLM新架构：文本下采样机制革新大语言模型效率王全东 小米大模型团队演讲嘉宾王全东小米大模型团队 大模型高级算法工程师中国科学院声学研究所博士、美国佐治亚理工访问学者、中科院认证高级工程师，长期从事大语言模型、多模态、语音识别等领域研究，曾获多项顶会竞赛冠亚军奖项，已发表顶会论文十余篇，拥有专列多项。深度参与了小米自研大模型从0到1的研发过程，荣获2024年度CCF计算机应用创新技术一等奖。近期和Daniel Povey等提出SUBLLM新架构，被量子位等科技媒体报道。目 录CONTENTS1.长文本模型的技术挑战2.SUBLLM架构3.主要实验结果4.分析与讨论5.总结与展望长文本模型的技术挑战PART 01长文本模型的技术挑战长文本需求旺盛多人会议摘要行业报告新闻摘要学术论文分析长文写作长篇翻译…长文本模型的技术挑战长文本模型结构：Decoder only Transformer结构：Llama 类似结构，attention的平方复杂度, by Meta, 2023长文本模型的技术挑战8k窗长32K窗长200K窗长1M窗长2~15T tokens5~10B tokens5~10B tokens5~10B tokens训练成本高：attention的平方复杂度主要优化点长文本模型的技术挑战模型窗长扩展方法：数据方向，训练成本不高1. Data Engineering for Scaling Language Models to 128K Context长文本模型的技术挑战模型窗长扩展方法：位置编码方向，训练成本不高2. YaRN: Efficient Context Window Extension of Large Language ModelsNTK-by-parts + 温度控制长文本模型的技术挑战模型窗长扩展方法：位置编码方向，训练成本不高3. PoSE: Efficient Context Window Extension of LLMs via Positional Skip-wise Training长文本模型的技术挑战8k窗长32K窗长200K窗长1M窗长2~15T tokens5~10B tokens5~10B tokens5~10B tokens训练成本主要在8k预训练阶段：attention的平方复杂度主要优化点长文本模型的技术挑战Decoder only Transformer长文本模型训练infra开发1. DeepSpeed Ulysses：优势：对Attention 的实现不敏感，适合各种attention方法劣势：序列并行度不能超过头数长文本模型的技术挑战Decoder only Transformer长文本模型训练infra开发2. Ring-attention：“大号”的flash attention优势：并行度的扩展性较好劣势：对Attention 变种不友好，eg. Sparse Attention长文本模型的技术挑战推理成本高：attention的平方复杂度推理速度角度，200k 比 20k 贵28倍，比2k 贵112倍 长文本模型的技术挑战其他长文本模型结构：1. Infini-Transformer：长期压缩记忆和局部因果注意力attentionInfini-Transformer 模型结构 by Google, 2024长文本模型的技术挑战其他长文本模型结构：2. MEGALODON：继承MEGA(带有门控注意力的指数移动平均) 并改进MEGALODON 模型结构 by Meta, 2024长文本模型的技术挑战加速方法很多 备受关注 SUBLLM应运而生SUBLLM架构PART 02SUBLLM架构目标:1. 开发一种优化资源使用的架构，同时保持模型能力不变。2. 区分重要token和不重要token，重要token占主要算力3. 兼容现有attention based模型生态，模型广泛应用的关键SUBLLM架构受语音领域启发：语音信号下采样减少冗余 保留必要信息语音识别中语音分帧后降采样4倍Conformer结构by Google, 2020SUBLLM架构受语音领域启发：语音信号下采样减少冗余 保留必要信息语音识别Zipformer，最高降采样16倍Zipformer 模型结构 by Daniel Povey, 2024SUBLLM架构受语音领域启发：语音信号下采样减少冗余 保留必要信息文本序列是否存在冗余？Yes!Funnel-Transformer 模型结构 by Google, 2020SUBLLM架构受语音领域启发：语音信号下采样减少冗余 保留必要信息文本序列是否存在冗余？Yes!Fourier Transformer 模型结构, 2023SUBLLM架构受语音领域启发：语音信号下采样减少冗余 保留必要信息文本序列中的token是否同等重要？No!CoLT5 模型结构 by Google, 2023SUBLLM架构受语音领域启发：语音信号下采样减少冗余 保留必要信息SUBLLM：采样(Sub-sampling)-上采样(Up-sampling)-旁路(Bypass) LLMSUBLLM 模型结构 by 小米AI实验室https://arxiv.org/pdf/2406.06571，accepted by ECAI 2024SUBLLM架构受语音领域启发：语音信号下采样减少冗余 保留必要信息SUBLLM：采样(Sub-sampling)-上采样(Up-sampling)-旁路(Bypass) LLM新模块插入位置示例，嵌套结构，逐渐下采样+逐渐上采样SUBLLM架构Learnable Subsampling Module1. 通过去除不太重要的tokens来缩短序列长度2. 使用Score层来衡量token的重要性序列采样，index选择预设采样保留比例dScore层给每个token打分再得到weight保留index，舍弃indexSUBLLM架构Learnable Subsampling Module妙处1：解决训练和推理的不一致性位置编码下采样Inference Mode：推理时看不到未来token，设定阈值为0，正数分值的token保留Balancer：通过对梯度进行惩罚来限制score的正数比例=预设下采样保留比例只保留正数分值的tokenSUBLLM架构Upsampling Module通过将采样后的tokens与原始序列合并来恢复序列长度妙处2：用减法使得token的完全选择可导，使score层有token判别能力实际基本为0，不影响weight大小，但是有梯度SUBLLM架构Subsampling and Upsampling ModuleSUBLLM架构Bypass Module：1. 下采样前和上采样后的序列加权求和，per channel2. 增强训练的收敛性和稳定性。SUBLLM架构对比Google: Mixture-of-Depths (MoD) ：1. 给重要token分配更多算力2. MoD 每间隔一个block做一次topk的token筛选3. 加速50%Mixture-of-Dep