arXiv:2605.17842cs.LG2026-05被引 1

通过结构化牛顿法实现分层并行推理,显著提升大模型生成速度。

SNLP: Layer-Parallel Inference via Structured Newton Corrections

论文配图:SNLP: Layer-Parallel Inference via Structured Newton Corrections
图 1 · 摘自论文原文
  • 将层间状态视为非线性方程,用近似牛顿迭代并行求解
  • 0.5B模型下最快提速2.58倍,且不增加困惑度
  • 适合追求推理加速的部署场景,尤其支持自推测解码

自回归语言模型按层顺序执行Transformer,形成延迟瓶颈,传统张量或流水线并行无法解决。本文将层间隐藏状态序列视为非线性残差方程的解,采用并行牛顿式更新求解。虽理论严谨,但精确牛顿修正需昂贵的雅可比-向量乘积,而朴素固定点迭代在训练好的Transformer上不稳定。为此提出结构化牛顿分层并行(SNLP),以廉价的架构诱导替代动态取代精确层雅可比。在残差Transformer中,得到恒等牛顿(IDN),修正简化为前缀和式更新;在mHC类架构中,使用残差混合矩阵的HC牛顿(HCN)。还研究了适配SNLP的训练策略,包括预训练正则化与直接SNLP前向微调。在纳米聊天级模型上实验表明,SNLP揭示了实用的速度-质量权衡:0.5B模型最高达2.58倍时钟速度提升,较温和配置达1.40倍提速且无困惑度上升。有效权衡源于IDN/HCN引入的有偏有限迭代计算,而非精确恢复串行轨迹。进一步证明SNLP前向微调可保持下游任务精度,且SNLP可作为自推测解码的草稿器,由串行验证器保障输出正确性。

原文摘要 · Abstract (English)

Autoregressive language models execute Transformer layers sequentially, creating a latency bottleneck that is not removed by conventional tensor or pipeline parallelism. We study whether this layerwise dependency can be relaxed by treating the hidden-state trace across layers as the solution of a nonlinear residual equation and solving it with parallel Newton-style updates. While this view is principled, exact Newton corrections require expensive Jacobian-vector products and naive fixed-point iterations are unstable on trained Transformers. We introduce Structured Newton Layer Parallelism (SNLP), a training and inference framework that replaces exact layer Jacobians with cheap architecture-induced surrogate dynamics. In residual Transformers, this yields Identity Newton (IDN), where the correction reduces to a prefix-sum-like update; in mHC-style architectures, HC Newton (HCN) uses the model's residual mixing matrix. We also study SNLP-aware training, including pretraining regularization and direct SNLP-forward SFT. Experiments on Nanochat-scale Transformers show that SNLP exposes a practical speed-quality frontier: on 0.5B models, it reaches up to 2.58x wall-clock speedup, and a less aggressive configuration reaches 1.40x speedup without increasing PPL. The useful tradeoff comes from the biased finite-iteration computation induced by IDN/HCN rather than exact recovery of the sequential trace. We further show that SNLP-forward SFT can preserve downstream task accuracy, and that SNLP can serve as a drafter for self-speculative decoding while a sequential verifier preserves output correctness.

模型加速并行推理牛顿方法自回归生成

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。