基于结构牛顿校正的层并行推理: SNLP
机器学习
2026-05-28 v2
摘要
自回归语言模型的 Transformer 层按顺序执行, 形成延迟瓶颈, 常规张量或流水线并行无法消除。我们研究是否可通过将跨层隐藏状态轨迹视为非线性残差方程的解, 并采用并行牛顿式更新来放宽层间依赖。虽然这种观点具有原则性, 但精确牛顿校正需要昂贵的 Jacobian-Vector 乘积, 而朴素的固定点迭代在训练好的 Transformer 上不稳定。我们引入结构化牛顿层并行 (SNLP), 一个训练和推理框架, 用低成本架构诱导的类比动力学取代精确层 Jacobian。在残差 Transformer 中, 这导致恒等牛顿 (IDN), 其中校正简化为前缀求和式更新;在 mHC 类架构中, HC 牛顿 (HCN) 使用模型的残差混合矩阵。我们还研究 SNLP-aware 训练, 包括预训练正则化和直接 SNLP 前向 SFT。在 Nanochat 规模的 Transformer 上实验表明, SNLP 显示出实用的速度-质量前沿:在 0.5B 模型上,其达到最高 2.58 倍 wall-clock 加速,较激进配置可达 1.40 倍加速且无需增加 PPL。有用权衡源于 IDN/HCN 引起的偏有有限迭代计算,而非对顺序轨迹的精确恢复。我们进一步表明, SNLP 前向 SFT 可保持下游任务准确性, SNLP 可作为自推测解码的草案,而顺序验证器保持输出正确性。
关键词
引用
@article{arxiv.2605.17842,
title = {SNLP: Layer-Parallel Inference via Structured Newton Corrections},
author = {Ligong Han and Kai Xu and Hao Wang and Akash Srivastava},
journal= {arXiv preprint arXiv:2605.17842},
year = {2026}
}
备注
Project webpage: https://github.com/phymhan/nanochat-snlp