论文

SNLP:通过结构化牛顿校正进行层并行推理

SNLP: Layer-Parallel Inference via Structured Newton Corrections

模型推理批处理与并行

摘要

自回归语言模型按顺序执行 Transformer 层,从而产生传统张量或管道并行性无法消除的延迟瓶颈。我们研究是否可以通过将跨层的隐藏状态迹视为非线性残差方程的解并使用并行牛顿式更新对其进行求解来放松这种分层依赖性。虽然这种观点是有原则的,但精确的牛顿校正需要昂贵的雅可比向量乘积,并且朴素的定点迭代在经过训练的 Transformer 上是不稳定的。我们引入了结构化牛顿层并行(SNLP),这是一种训练和推理框架,用廉价的架构引发的代理动力学取代了精确的层雅可比行列式。在残差 Transformer 中,这会产生恒等牛顿 (IDN),其中校正减少为类似前缀和的更新;在 mHC 式架构中,HC Newton (HCN) 使用模型的残差混合矩阵。我们还研究 SNLP 感知训练,包括预训练正则化和直接 SNLP 前向 SFT。 Nanochat 规模 Transformer 上的实验表明,SNLP 暴露了一个实用的速度质量前沿:在 0.5B 模型上,它达到了高达 2.58 倍的挂钟加速,而不太激进的配置则在不增加 PPL 的情况下达到了 1.40 倍的加速。有用的权衡来自于 IDN/HCN 引起的有偏差的有限迭代计算,而不是顺序跟踪的精确恢复。我们进一步表明,SNLP-forward SFT 可以保持下游任务的准确性,并且 SNLP 可以充当 self-推测解码 的起草者,而顺序验证器可以保持输出的正确性。