在科技的浩瀚星空中,有两位闪耀的新星——金天,这位麻省理工学院(MIT)计算机科学与人工智能实验室(CSAIL)博士五年级的学生,如同一位在知识海洋中执着探索的航海家,师从 Michael Carbin 和 Jonathan Ragan - Kelley,专注于机器学习与编程系统的结合。他曾在 IBM Research 主导实现深度神经网络在 IBM 主机上的推理部署,本科毕业于 Haverford College,拥有计算机科学与数学双学位,知识的光芒在他身上交织闪耀。
还有鄭鈺熹,麻省理工学院 CSAIL 博士三年级的学生,像是在科技森林中寻找独特花朵的寻芳者,师从 Michael Carbin,在编程语言与机器学习的交叉领域展开深入研究。
大语言模型(LLM)的发展,就像一部波澜壮阔的史诗,正从传统的「单人书写」时代迈向「分身协作」的新纪元。传统的自回归解码方式,如同一位一丝不苟的工匠,按顺序生成内容,而新兴的异步生成范式,则如同一个高效的团队,通过识别语义独立的内容块,实现并行生成。
看那对比图,传统方法按部就班地生成所有内容,而异步生成则同时处理多个互不依赖的内容块,宛如一场速度与激情的较量。在 AlpacaEval 长度控制评测中,异步生成实现了 1.21 - 1.93× 的几何平均提速,这速度的提升,就像给科技的列车装上了强大的引擎,不过生成质量也有 +2.2% 至 -7.1% 的变化,如同在追求速度的道路上,偶尔也会遇到一些小颠簸。
MIT 与谷歌研究团队在最新研究 PASTA(PArallel STructure Annotation)中,如同勇敢的开拓者,首次从策略学习(policy learning)角度探索异步生成范式的可能。
- 论文标题:Learning to Keep a Promise: Scaling Language Model Decoding Parallelism with Learned Asynchronous Decoding
- 论文地址:https://arxiv.org/abs/2502.11517
研究团队没有依赖人工设计规则来识别异步生成机会,而是让模型像一个聪明的学习者,通过策略学习自主发现并标注这些机会,系统地优化质量与速度的平衡。这就好比在生活中,我们不断学习,找到工作与生活的最佳平衡点。这种方法让 LLM 根据内容特点自适应地确定最佳异步生成策略,为生成效率优化开创了一条学习驱动的全新路径,仿佛为科技的发展打开了一扇崭新的大门。
PASTA - LANG:划分独立内容的标记语言
研究人员开发的新标记语言 PASTA - LANG,就像一把神奇的钥匙,专为异步生成而设计。大模型使用它在生成过程中标记语义独立块,指示并行生成机会。这种语言包含三种核心标记,它们共同构成了一种「承诺 - 履行」的生成模式。
- :标记语义独立的内容块,通过 topic 属性总结内容主题,大模型用它表明「这部分将会由一个独立子线程异步生成」,就像在生活中许下一个承诺。
- :在
- 后标识对应的异步生成的内容,表示这一部分由独立子线程负责生成,如同去履行许下的承诺。
- :在主线程标记同步点,表明后续内容生成将会依赖于之前的异步生成,主线程需等待所有异步线程完成后才能继续,这就像生活中等待团队成员完成各自任务后再一起前进。
这些标记共同构建了一个有序的世界,大模型首先通过标记「承诺」生成某些内容,推理系统再创建异步线程来「履行」这些承诺,最后在处将异步内容组合,就像一场精心编排的舞蹈,每个环节都紧密相连。
看那线段长度计算案例,面对计算任务,大模型就像一个智慧的指挥官,首先识别出「坐标提取」和「长度公式」两个可并行的独立部分,生成相应的标记,随后用标记表明需要等待这些内容完成。图中红色和绿色区域显示了两个异步线程并行生成的内容,最终组合成完整解答,这一过程就像团队协作完成一个复杂的项目,每个成员各司其职,最终达成目标。
这个新的标记语言简单而又具有强大的可扩展性,它开启了新的未来研究范式,就像一颗种子,在科技的土壤中种下,未来可能会长成参天大树。
PASTA 训练:从标注到优化的双阶段学习
PASTA 系统采用双阶段训练流程,让大模型像一个不断成长的学生,自主学习使用上述标记语言,完成异步生成。
第一阶段:监督微调。研究团队选取 SlimOrca 指令跟随数据集,用 Gemini 1.5 Flash 为 100K 条样本添加 PASTA - LANG 标记,在样本回答中插入标记,创建 PASTA 微调数据集。然后对 Gemma 7B 进行监督微调,得到能插入 PASTA - LANG 标记的 PASTA - SFT 模型,这就像老师指导学生完成基础的学习任务。
第二阶段:偏好优化。为了优化标注策略,团队设计了策略学习方案。对每个样本从 PASTA - SFT 模型采样多种标注方案,然后基于理论加速比和内容质量(由 Gemini 1.5 Pro 评估)两项指标评估这些方案。根据评估结果,构建「拒绝采样数据集」,该数据集包含每个输入的最佳和最差标注方案。最后,用 BoNBoN 算法对 PASTA - SFT 模型进行偏好优化,得到最终的 PASTA 模型,这就像学生在不断地自我反思和改进中成长。
PASTA 推理系统:并行生成与缓存管理
推理系统设计难点。异步并行生成的主要挑战在于如何协调多个线程高效协作,这就像组织一场大型活动,需要协调各方人员的行动。传统方法通常需要为每个线程创建独立的 KV 缓存池,创建新线程时必须复制主线程的前缀内容到子线程缓存池,完成后再复制结果回主线程,这两次大规模矩阵复制操作严重限制了系统性能,使理论加速难以转化为实际收益,就像在前进的道路上遇到了绊脚石。
KV 缓存的存储布局。PASTA 设计了交错式 KV 缓存布局,所有线程共享单一连续内存池。系统初始以连续方式存储用户输入,在推理过程中动态将不同线程在同一时间点生成的 token 交错存储在相邻位置,这就像合理规划资源,让每个参与者都能高效地利用空间。
注意力控制与位置编码。PASTA 通过两个机制确保大模型正确理解多线程交错存储的 KV 缓存:
- 注意力掩码控制:限制子线程只能访问与自己相关的内容,在
- 后通过移除掩码使主线程能访问所有子线程生成的内容,这就像在团队中,每个人专注于自己的任务,到了关键节点再共享信息。
- 位置编码调整:每个线程都使用独立且连续的位置编码,使线程处理自己的内容时,将交错存储的内容视为逻辑上连续的序列,确保模型能正确理解上下文,就像给每个成员一个清晰的任务指南,让他们能准确地完成工作。
这些设计共同确保 PASTA 能在提高速度的同时保持输出质量,就像在追求速度的同时,也不忘记保证工作的质量。
实验结果:Pareto 最优与可扩展性
PASTA 在性能与质量的平衡上取得了突破性成果,这就像在生活中找到了事业与家庭的完美平衡点。实验结果表明它不仅实现了显著加速,还在某些情况下提高了输出质量。研究团队在 AlpacaEval 基准上进行了全面评估,该基准包含 805 个具有代表性的指令跟随任务。
质量 - 速度平衡的 Pareto 前沿。如图所示,PASTA 通过调节质量权重参数生成了一系列的模型。在不同的生成质量的情况下,PASTA 均能提供非常可观的加速。结果显示,即使最注重质量的 PASTA 模型也能提供显著加速,而最快的模型则以一定的质量牺牲换取接近 2 倍的速度提升。与基于手动设计的异步生成方案(Skeleton - of - Thought, APAR)相比,PASTA 模型展现出全面优势,这就像在竞争中脱颖而出的强者。
可扩展性。研究结果展示了 PASTA 方法出色的可扩展性,随着偏好优化不断推进,PASTA 模型的性能持续提升。图中清晰展示了从第一轮开始到第一轮结束,再到第二轮开始和第二轮后半程的整个优化过程,质量 - 速度的 Pareto 前沿大体持续向右上方推进。这种稳定的改进趋势表明,PASTA 方法具有良好的可扩展性,随着投入更多计算资源,仍未饱和。与传统依赖固定规则的异步解码方法不同,PASTA 通过策略学习驱动的训练算法提供了可持续的优化路径,能够有效地将额外计算资源转化为更高的推理效率,这就像一个不断成长的企业,有无限的发展潜力。
总结与展望
PASTA 首次证明,通过策略学习让 LLM 自主优化生成策略,能够突破传统自回归和基于规则的异步生成的效率极限。这一工作不仅为实时大模型应用提供了实用加速方案,更印证了未来 LLM 可能具备推理时自我优化能力的发展方向。在科技的长河中,PASTA 就像一颗闪耀的星星,照亮了未来大语言模型发展的道路,让我们对未来充满了无限的期待。