SNLP：基于结构化牛顿修正的层级并行推理 SNLP: Layer-Parallel Inference via Structured Newton Corrections
把 Transformer 层间隐藏状态视为非线性残差方程，用廉价结构化替代 Jacobian 做并行求解，最快实现 2.58× 加速。
前置知识
Transformer 层序依赖
Transformer 语言模型的自回归推理存在两重顺序性：token 自回归生成以及对固定 token 前缀，隐藏状态必须逐层执行 $h_l = f_l(h_{l-1})$。即便有了张量并行、流水并行、KV cache、推测解码等优化，固定前缀下的深度依赖链依然存在，这是本文想要打破的核心瓶颈。
理解这点是看懂 SNLP 切入点的关键：之前的所有推理加速都在'每层更快'或'每 token 更便宜'两个轴上做文章，SNLP 则挑战'是否真的必须一层一层算'。
DEER 风格并行非线性求解
DEER 将非线性递推的整条轨迹视为一个非线性方程组的解，用 Newton 法把每步线性化后用 parallel scan 沿递推轴并行求解。SNLP 把这个视角从序列长度轴旋转 90° 到 Transformer 深度轴，对隐藏状态轨迹 $(h_1,\dots,h_L)$ 建立块下双对角残差 Jacobian 并做 Newton 求解。
这是 SNLP 理论根基；理解 DEER 才能理解为什么 Newton 修正天然适合沿深度轴传播，也才能理解 SNLP 的精确 Newton 更新为什么在语言模型上不可行（Jacobian 太大）。
Jacobian 替代/结构化近似
精确 Newton 需要每层 $J_l = \partial f_l/\partial h_{l-1}$，对 LM 隐藏状态而言不可能材料化。SNLP 的核心思想是换成廉价的结构化替代 $A_l$：残差 Transformer 上 $A_l=I$（即 Identity Newton，简称 IDN），mHC 架构上 $A_l=H^{mlp}_{res,l}H^{attn}_{res,l}$（HC Newton，HCN），或者只取 Jacobian 的对角线（DiagN）。
这是 SNLP 唯一的关键工程 trick；不同替代决定训练目标、可达速度-质量权衡，是理解全部实验配置差异的钥匙。
自推测解码 (Self-Speculative Decoding)
用同一个模型不同推理路径作为草稿/验证组合：先用廉价的近似前向（如 SNLP）一次性生成 $B$ 个候选 token，再用一次顺序前向并行验证，按接受率 $\alpha$ 决定接受几个。理想端到端加速约为 $(1-\alpha^{B+1})/(1-\alpha) / (B/s + 1)$。
SNLP 用作 drafter 时即便自身输出有偏也能保证最终结果正确，是 SNLP 实用性最直接的体现，也是 Table 3 的核心指标解释前提。
研究动机
现代 Transformer 自回归推理有两个独立但叠加的顺序来源：一是 token 自回归，二是对固定前缀，隐藏状态必须沿 $h_l=f_l(h_{l-1})$ 逐层执行。张量并行把单层算力拆到多卡，流水并行按层切分流水线，KV cache 复用历史 token，kernel fusion 与量化压低单层开销，推测解码则把'生成多个 token 再验证'作为突破口——但所有这些都没拆掉那条对固定前缀的层间依赖链。问题随着模型越来越深（例如几十层甚至上百层）越发尖锐：单 token 解码时深度方向的延迟天然成为瓶颈，且对 batch size=1 的在线推理场景尤其致命。在 Nanochat-3B 这种 32 层结构上，单 token prefill/decode 的 wall-clock 时间几乎由这 32 次顺序层调用主导。
本文的目标是论文提出一个训练-推理协同设计的框架 Structured Newton Layer Parallelism (SNLP)，目标是在保持模型质量的代价内，把对固定前缀的层序依赖松弛掉。具体说，作者希望用结构化 Newton 修正让多个 suffix 层在一次或少数几次迭代里同步推进，从而把延迟从 $O(L)$ 的层数依赖降到接近 $O(1)$（按 chunk 数）级别，并在 0.5B Nanochat 模型上证明至少 1.40× 无损加速以及最高 2.58× 有损加速的可行性。
与已有工作不同的是，已有思路大致分两类：(1) DEER/Jacobi decoding 等把'非线性链'作为耦合方程组求解，但 DEER 原版要求全 Jacobian，Jacobi decoding 又在训练好的 Transformer 上发散或收敛太慢；(2) Hyper-Connections/AttnRes 等改造残差流结构以暴露更多深度耦合，但只在架构层改变，没有回答'即便改了架构是否真的能并行推理'。SNLP 的独特切入是把两者缝合：先用架构诱导的廉价 Jacobian 替代（残差路径的 $I$、mHC 的残差混合矩阵）让 Newton 修正在推理侧跑得起，再通过一个把有限迭代 SNLP 状态匹配到顺序轨迹的辅助损失，把模型训练成'层动力学适合这套廉价替代'。这种 training/inference co-design 是论文区别于单纯推理 trick 的核心。
核心方法
直觉上，SNLP 想把 Transformer 前向传播的隐藏状态轨迹 $(h_1,\dots,h_L)$ 看作一个非线性残差方程 $G(h)=0$ 的解，然后像解非线性方程组那样，用 Newton 法线性化后并行更新所有层。直接照搬 Newton 会有两个问题：(a) 每层 $J_l=\partial f_l/\partial h_{l-1}$ 对 LM 隐藏状态维度是天文数字，材料化、JVP、有限差分都会吃掉本想省的延迟；(b) Jacobi/不动点迭代在训练好的残差网络上要么发散要么太慢。SNLP 的策略是'架构诱导的结构化替代'：在残差 Transformer 上 $f_l(x)=x+g_l(x)$，自然有 $J_l\approx I$，于是 Newton 修正退化成前缀式传播 $h^{(k+1)}_l = \tilde h^{(k)}_l + h^{(k+1)}_{l-1} - h^{(k)}_{l-1}$，计算量几乎为零，这就是 Identity Newton (IDN)；在 mHC 上有显式的残差混合矩阵 $H_{res}$，于是用 $A_l=H^{mlp}_{res,l}H^{attn}_{res,l}$，称为 HC Newton (HCN)。这样一次迭代内，昂贵的前向计算（残差分支 $g_l$）可以在 suffix 层间并行跑，Newton 修正只是轻量的逐层递归。推理侧还引入层融合（把多层的 QKV 拼接成宽矩阵）和 chunkwise 策略（按 chunk 迭代），把并行结构变成对 GPU 友好的大 GEMM。训练侧引入一个匹配有限迭代 SNLP 输出与顺序轨迹的辅助正则损失，把模型训练成'与这套廉价替代兼容'。
核心创新是用'架构诱导的廉价结构化替代'取代精确 Newton 步里的层 Jacobian，并在训练阶段同步优化模型以让这套替代在有限迭代下可用。这与已有工作的本质区别在于：(1) 与 DEER/quasi-DEER/ELK 不同，SNLP 不去逼近真 Jacobian，而是直接把残差路径的 $I$（或 mHC 的混合矩阵）当成 Jacobian 用，再用训练把残差分支的灵敏度压下来，使 $J_{f_l}=I+J_{g_l}\approx I$；(2) 与 Jacobi decoding/deep equilibrium 等不动点类方法不同，SNLP 不是依赖收缩映射的迭代收敛，而是显式做 Newton 修正并允许有限迭代残差，把这种'求解器诱导的偏置'作为一种可调的速度-质量权衡；(3) 与 Hyper-Connections 等纯架构改进不同，SNLP 把架构信号和训练目标耦合成同一套'廉价 Newton 求解'机制。
方法步骤详情
方法实现分为推理侧与训练侧两条流水线。推理侧（Algorithm 2）：第一步对 prefix 部分 $l=1,\dots,S$ 顺序执行，得到 $h_S$；第二步对 suffix 部分初始化所有 $h^{(0)}_{S+j}\gets h_S$（h0 初始化）或做一次单次 batched 前向（fwd 初始化）；第三步进入迭代 $k=0,\dots,K-1$，其中 (a) 对 suffix 内所有 $l=S+1,\dots,L$ 并行计算 $\tilde h^{(k)}_l = f_l(h^{(k)}_{l-1})$；(b) 对 $l=S+1,\dots,L$ 顺序做 Newton 修正 $h^{(k+1)}_l = \tilde h^{(k)}_l + A^{(k)}_l(h^{(k+1)}_{l-1} - h^{(k)}_{l-1})$，其中 $A^{(k)}_l=I$（IDN）或 $A_l=H^{mlp}_{res,l}H^{attn}_{res,l}$（HCN）；(c) DiagN 时 $A^{(k)}_l=\mathrm{diag}(J^{(k)}_l)$，可用 Hutchinson 有限差分/JVP/VJP 估计，再用 associative scan 沿深度高效求解。chunkwise 模式下把 suffix 切成 $C$ 个 chunk，每个 chunk 内的层用融合 QKV/MLP 宽矩阵一次计算，Newton 修正在 chunk 之间进行；同一 chunk 内多个并行层共享输入 $h$ 并通过相加注意力输出和 MLP 分支实现隐式跨层耦合。最后一步把 $h^{(K)}_L$ 投到 logits。训练侧有两类：(a) 预训练时在 CE 损失基础上加辅助正则 $\mathcal{L}=\mathcal{L}_{CE}+\lambda \sum_{N\in\mathcal{S}} \sum_{l\in T_N} \|h^{SNLP}_l(N,K;A) - h^{seq}_l\|^2 / (\|h^{seq}_l\|^2+\epsilon)$，其中 $K=1$、$T_N$ 是 stride 选出的监督层集合；(b) 任务适应阶段直接用 SNLP 前向做 SFT，让模型在 SNLP 计算图上学任务。DiagN 的特殊处理是先用 Algorithm 3 的 FD/JVP/VJP 估计对角 $\mathrm{diag}(J^{(k)}_{S+j})$，再做 affine prefix scan。
技术新颖性
技术新颖性体现在三处。第一处是把 Newton 求解的'架构先验'用足：以往工作要么做不动点（Jacobi/DEQ），要么逼近 Jacobian（quasi-DEER/ELK），SNLP 直接用残差流本身作为 Jacobian 的一阶项，再训练模型让二阶项可忽略，这样修正项几乎免费。第二处是引入了 solver-induced inference bias 这一新概念：精确 Newton 在收敛极限下恢复顺序前向，但 SNLP 用的是有限迭代 + 廉价替代 + chunk 初始化 + 融合卷积 的组合，这些引入的偏置不再是噪声而是可以被训练侧吸收的结构性偏差，从而把'加速'和'质量'解耦成可调旋钮。第三处是层融合 + chunkwise 策略：把多层的 QKV 矩阵沿输出维拼接、MLP 沿输入维拼接，让一次宽矩阵乘覆盖多个并行层；以及把 DeltaNet 风格的 chunkwise 并行思想搬到深度方向，使 chunk 内 attention 的输出相加引入隐式跨层耦合，论文在 Section C.3 给出耦合项的显式表达并指出'过于激进'的融合会改变模型行为。
实验结果
实验在 Nanochat 框架的两个尺度（0.5B 与 3B，均 32 层）以及若干变体（standard、带 x0/VE、不带 x0/VE 的标准残差、mHC）上跑 pretraining 评估、主 PPL、下游任务（ARC-Easy/ARC-Challenge/MMLU）和自推测解码。Table 1 是主表，对所有模型 No Reg.、IDN Reg.、DiagN Reg. 各自报告 2-3 个 SNLP 配置的 PPL、∆PPL、加速比、Top-1、LogitSim、EmbSim、AR Match。关键发现：(1) 在 0.5B w/o x0/VE IDN Reg. 上 12xF2-h0-K1 达到 2.58× 加速 + 17.0% PPL 上升 (20.55 vs 17.57)，同时 1xF12-h0-K1 达到 1.40× 加速且 PPL 不增 (17.56 vs 17.57)；这是最干净的速度-质量前沿点。(2) 在 0.5B standard IDN Reg. 上 12xF1-fwd-K2 给出 1.14× 加速 + 仅 1.5% PPL 上升 (15.59 vs 15.36)，24xF1-h0-K2 在 1.88× 加速下 PPL 仅升 17.8%。(3) 3B 模型上 SNLP-aware 训练仍能给出有用的前沿点：IDN Reg. 8xF1-h0-K1 给出 1.20× 加速 + 3.9% PPL 上升 (10.46 vs 10.07)，但更激进的配置在当前 PyTorch 实现下无法继续放大加速比，因为顺序 Transformer block 已经较好地打满了 H100。(4) mHC 模型上 HCN Reg. 16xF1-h0-K1 给出 1.60× 加速 + 9.1% PPL 上升，HCN 保持了很高的 greedy 接受率。Table 2 显示下游任务上 IDN Reg. 模型做 Sequential SFT + IDN inference 在 3B 上达到 57.03% 平均精度，接近 Sequential SFT + Sequential 的 57.85%，并远好于 No Reg. 模型上 Sequential SFT + IDN inference 的 54.60%，说明 IDN 正则化让模型在 SFT 后仍保有 SNLP 兼容行为。Table 3 把 SNLP 用作 self-speculative decoding 的 drafter，最佳配置 12xF2-K2-h0 B=4 α=0.958 给出估计端到端 1.67× 加速，HCN 因保持 token 排序能力在高 α 上尤其受用。Table 4 的诊断显示 IDN Reg. 在 0.5B 末尾 8 层的 Jacobian 谱估计降约 12×、Frobenius 估计降约 12×，相对替换误差 $|\epsilon_l|$ 控制在 0.03%–0.15% 区间（vs No Reg. 的 2%–24%），与 Lipschitz 隐式正则化解释一致。Table 5 的修正顺序实验显示顺序前向在 N≤12 时对随机置换几乎不变（shuffle std≤0.06），而 IDN 修正则更依赖前向顺序（K=1 下 forward 几乎总是最优）。Table 6 反驳了'输入近似不变的层即可丢掉'的直觉：跳过 8–16 层 PPL 立刻爆到 50–1700+。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| Nanochat-0.5B w/o x0/VE PPL（长度 2048） | Perplexity 与加速比 | IDN Reg. 12xF2-h0-K1: PPL 20.55 (+17.0%) 加速 2.58×；IDN Reg. 1xF12-h0-K1: PPL 17.56 (-0.1%) 加速 1.40× | No Reg. 顺序 PPL = 17.65 | 无损配置加速 1.40×；激进配置加速 2.58× 但 PPL +17% |
| Nanochat-0.5B standard PPL（长度 2048） | Perplexity 与加速比 | IDN Reg. 24xF1-h0-K2: PPL 18.09 (+17.8%) 加速 1.88×；12xF1-fwd-K2: PPL 15.59 (+1.5%) 加速 1.14× | No Reg. 顺序 PPL = 15.21 | 保守配置仅 +1.5% PPL 换 1.14×；激进 1.88× 换 +17.8% PPL |
| Nanochat-3B standard PPL（长度 2048） | Perplexity 与加速比 | IDN Reg. 8xF1-h0-K1: PPL 10.46 (+3.9%) 加速 1.20×；8xF1-h0-K4: PPL 10.09 (+0.2%) 加速 0.95× | No Reg. 顺序 PPL = 10.10 | 3B 上仅能给出 1.20× 无大损加速点；更大并行度当前 PyTorch 实现打不满 H100 |
| Nanochat-0.5B-mHC PPL（长度 2048） | Perplexity 与加速比 | HCN Reg. 16xF1-h0-K1: PPL 16.93 (+9.1%) 加速 1.60×；8xF1-h0-K4: PPL 15.60 (+0.5%) 加速 1.00× | No Reg. 顺序 PPL = 15.16 | HCN 在 K=4 下能接近顺序 PPL 但加速 1.00×；K=1 给 1.60× 但 +9.1% PPL |
| 下游多选题准确率（3B IDN Reg.） | ARC-E / ARC-C / MMLU 平均准确率 | Sequential SFT + IDN-8xF1-K1 推理: 57.03% (74.39 / 56.86 / 39.84) | No Reg. Sequential SFT + Sequential 推理: 57.85% (74.91 / 58.86 / 39.78) | 仅落后 0.82 个百分点但推理用 IDN 并行；远好于 No Reg. Sequential SFT + IDN 推理的 54.60% |
| Self-Speculative Decoding（SNLP 作 drafter） | 估计端到端加速比 | 12xF2-K2-h0 B=4 α=0.958: 估计加速 1.67×（理想值 1.604×，drafter 加速 2.15×） | 顺序推理基准 1.00× | 在保证 greedy 输出与顺序完全一致的前提下最多 1.67× 加速 |
局限与改进
作者明确承认的局限包括：(1) 速度-质量权衡而非保证：SNLP 给出的是可调前沿，不是'无损'方案；3B 模型只能给出 1.20× 加速且伴随 3.9% PPL 上升，激进配置在 3B 上因 PyTorch 级实现的层融合无法克服宽块顺序执行的开销。(2) 离训练好的大模型不易加装：Qwen2.5-0.5B-Instruct 和 TinyLlama-1.1B-Chat-v1.0 上 SNLP 后处理匹配顺序 PPL 要 8 次迭代却更慢（Table 13），finetune 注入 SNLP 兼容性也不容易（Table 16：更高 IDN 权重直接拉到 PPL 26+）。(3) DiagN 路径代价高：Table 15 显示 JVP 比 plain forward 慢 6-8×，VJP 慢约 3×，且 VJP 在 mHC 这类非对称路由下估计 $J^\top v$ 而非 $Jv$ 导致对角近似偏离实际方向，常引起发散。(4) 顺序偏差：Table 6 表明 $J\approx I$ 不等于'层可丢掉'，训练后非残差分支的灵敏度被压小但仍必须存在，跳过 8-16 层 PPL 立刻爆到 50-1700+。我自己的观察：(a) 实验仅在 batch=1 H100 上时延，缺大 batch 与多卡扩展性数据，推理系统的实际可行性还待验证；(b) 主表全部基于 Nanochat 自训练 32 层模型，是否能泛化到 70+ 层甚至 100+ 层的现代大模型仍是开放问题；(c) 顺序 SFT 后切换到 IDN 推理的精度损失 0.82 个百分点看似不大，但若迁移到更敏感的下游任务（如生成式 QA、代码补全）影响可能放大，论文没给出这类评估；(d) IDN 顺序偏差主要靠训练吸收，那么一旦后续在 SNLP 模型上继续 SFT/DPO/RLHF，是否会破坏已经学到的'输入近似不变'的层动力学也没讨论。
独立分析的弱点
独立审视弱点有三类。第一类是'层序偏差只能被吸收不能被消除'：SNLP 用 $A_l=I$ 等廉价替代，本质是用前向输出差异当残差，只要 $g_l(h_S)\neq g_l(h^{seq}_{l-1})$ 就会引入系统性偏差；论文通过训练正则化把 $J_{g_l}$ 压小从而缩小偏差，但这反过来限制了模型表达能力（残差分支必须输入近似不变），在需要高度顺序依赖的任务（如长链推理、多步算术）上可能掉点明显。第二类是系统实现层：当前 chunkwise 融合在 PyTorch 级别，3B 模型上 fused 块没能超越顺序执行，主要瓶颈是宽矩阵乘没把 H100 喂满；缺乏定制 CUDA/Triton kernel 也意味着难以扩展到 7B/13B/70B 规模。第三类是修正传播的鲁棒性：Table 5 显示 IDN 修正在 K=1 下前向顺序几乎总是最优，其他顺序随机置换会出现 PPL 数百到数千甚至发散，这意味着算法对'深度方向因果性'的依赖非常重，一旦未来架构用更多非因果耦合（如跨层 skip、双向残差），前向顺序的优越性可能消失，需要重新设计 $A_l$。
未来方向
作者提出的方向包括：(1) 更系统的 SNLP-aware 后训练和 SFT 研究，特别是给现成大模型做 SNLP 兼容适配；(2) 更大模型的算法-硬件协同设计，如定制 fused kernel、compute-in-memory 执行；(3) 让 DiagN 路径在 Jacobian 估计上更便宜的实现工作；(4) 把 SNLP 推广到 Jacobi decoding、扩散语言模型、Homomorphic Encryption 等特殊执行约束场景。基于成果可延伸的方向还有：(a) 把 SNLP 与 early exit 结合：早期层做 SNLP 解、中后期层恢复顺序，看能否进一步提升 wall-clock；(b) 把 HCN 推广到其他显式残差混合架构（如 AttnRes、Value Residual Learning），看是否能复用流维小矩阵替代；(c) 自适应迭代数 K：在每次迭代后用 cheap proxy（如 hidden state 变化幅度）判定收敛，提前终止；(d) 与 MoE 结合：把专家路由也并入 chunk，并在 Newton 修正中考虑专家负载均衡；(e) 训练侧探索非 $L_2$ 的匹配损失（如 spectral norm 或 Sinkhorn 距离），可能更直接压住 Jacobian 谱。
复现评估
代码已在 GitHub 开源（https://github.com/phymhan/nanochat-snlp），主要实验基于公开的 Nanochat 训练流水线、ClimbMix（FineWeb-Edu 子集）、ARC、MMLU 与 SmolTalk 数据集，所以数据与训练代码可完全复现。模型架构细节（nembd=2048 / 640、head 数、32 层、rotary、x0 残差、value embeddings）在 Table 8 与正文 §5 给出，正则化超参（N、λ、stride、detach）在 Table 9 详尽列出。算力需求方面：Nanochat-3B 标准模型 9600 步、0.5B 模型 4800 步、每个规模需训 No Reg./IDN Reg./DiagN Reg. 三个变体（外加 0.5B w/o x0/VE 与 mHC 变体），单卡 H100 训练 0.5B 量级数小时到一天，3B 需要更长时间，整体可在 8×H100 节点上几天内完成。复现难度评估为中等偏上：硬件门槛低（H100 即够），但 (1) PyTorch 级 fused/chunkwise 实现细节、kernel launch overhead、warmup 协议会显著影响加速比，需对照原 timing 协议（batch=1、50 warmup、200 measured runs）严格对齐；(2) 多个训练随机种子下加速-质量前沿点的稳定性未充分报告；(3) SNLP-aware 正则损失中 stride 与 λ 的选择敏感（Table 9 中 stride 2 直接把 0.5B standard PPL 拉到 21.36），需要细致复现消融；(4) self-speculative decoding 的接受率只在 128 个 prompt 上跑且仅 greedy，对采样场景的结论外推需谨慎。
论文图表
列出 4 类模型的训练步数（9600 for 3B, 4800 for 0.5B 等）、正则化方法、suffix 长度 N、λ、stride、detach 标志。
是 Table 1 数字背后的完整训练配方，对复现至关重要。
对 4 类模型系统扫描正则超参：3B 上 IDN N=8 λ=0.5 stride=0 给出最优 -0.3% PPL，0.5B standard 敏感于 λ 与 stride（stride=2 直接 PPL 21.36），w/o x0/VE 偏好 stride=6。
为'为何这套正则化有效'与'超参如何选'提供完整的实证地图。
Qwen2.5-0.5B-Instruct 和 TinyLlama-1.1B-Chat-v1.0 在 K=4 时 PPL 上升 8-10% 且速度 0.74-0.75×；K=8 恢复到 0-1.7% PPL 上升但速度仅 0.63-0.66×，没有加速。
说明 SNLP 作为通用后处理加速技巧的局限性，强化'training/inference co-design'的核心论点。
对 Nanochat-3B、Nanochat-0.5B、Qwen2.5-0.5B 在 batch=1 不同序列长度下比较 plain Forward、FD、JVP、VJP 的单层耗时。JVP 比 plain 慢 6-8×、VJP 慢约 3×、FD 慢约 2×。
为 DiagN 失败提供具体代价证据，是理解'DiagN 路径不实用'结论的关键基准。
对 TinyLlama-1.1B-Chat-v1.0 跑 8 个 IDN 正则微调配置：baseline PPL 17.82，IDN 权重 0.5 时 PPL 17.91，权重 5-10 时 PPL 跳到 26+，长程训练 (idn1_long, idn5_long_npar4) 卡在 ~17.8 即崩。
直接展示 SNLP 后处理适配离训练模型的难度，与 Table 13 一起支撑'需要 co-design'的结论。
在序列长度 2048 下逐步激活 suffix 的 last 1..8 层，比较有/无 Newton 修正的 PPL。有修正时所有 active 层从 K=1 起就影响 PPL（last 1..8 列分别为 38.70→17.19），无修正时呈阶梯形——K=1 所有列都是 37.41，K=2 之后 last 2..last 8 才开始一致，体现'无修正每次只前进 1 层'。
直接验证 Section 3.2 中关于'Newton 修正让信息在一次迭代内传播到整个 suffix'的 claim。