1 Introduction↩︎

Transformer language models [1] are sequential in two distinct senses. Token generation is autoregressive [2], [3], but even for a fixed token prefix, the hidden state must normally pass through the network one layer at a time. Tensor parallelism [4], pipeline parallelism [5], kernel fusion [6], batching, KV caching [7], and speculative decoding [8], [9] improve the efficiency of each layer or token step, but they do not remove the layer dependency chain. As models become deeper [10], [11] and decoding remains latency-sensitive, this depthwise dependency becomes a natural target for algorithmic parallelism.

A principled way to expose such parallelism is to view the entire sequence of hidden states across layers as the solution of a nonlinear residual equation. This is analogous to DEER-style [12] parallelization of nonlinear recurrences, where Newton iterations solve for all states in a chain jointly rather than executing the chain strictly left-to-right [13]. Applied along the depth axis, this perspective suggests that many Transformer layer states could be updated in parallel. However, exact Newton updates require the Jacobian of each full layer with respect to its input. For language-model hidden states, these Jacobians are too large to materialize, and even Jacobian-vector or finite-difference approximations can consume the latency budget that layer parallelism is meant to save. Cheap fixed-point or Jacobi iterations avoid this cost [14], [15], but are often unstable or slow on trained residual networks.

We introduce Structured Newton Layer Parallelism (SNLP), a training and inference framework that makes this Newton view practical by replacing exact layer Jacobians with cheap structured surrogates. In residual Transformers, the identity residual path gives the simplest surrogate, yielding Identity Newton (IDN): the correction reduces to additive prefix-style propagation over depth. Diagonal Newton (DiagN) connects SNLP to quasi-DEER and scan-based linear recurrences [16]. For HC/mHC-style models [17], [18], the architecture exposes a learned residual mixing matrix, yielding HC Newton (HCN). In all cases, the expensive nonlinear layer or chunk forwards are parallelizable, while the Newton correction is a lightweight structured recurrence.

The second ingredient is training co-design. A pretrained sequential model need not be compatible with a cheap surrogate Jacobian, so we introduce SNLP-aware regularization: during training, we ask one or a few structured Newton iterations over a suffix of layers to match the ordinary sequential hidden state. This regularizer encourages suffix dynamics that are easier to solve with the chosen surrogate. Empirically, it also improves the standard sequential model in several trained-from-scratch Nanochat settings [19], suggesting that it acts as a useful regularizer on layer dynamics rather than merely an inference-time approximation loss.

Our experiments show that layer-parallel inference can be useful in practice, but not as a universal post-training acceleration trick. On trained-from-scratch Nanochat-scale models, SNLP-aware regularization improves sequential PPL by 4.7%–23.4%. At the 0.5B scale, SNLP inference with chunkwise layer fusion reaches up to \(2.3\times\) speedup with comparable or lower PPL than the model’s own sequential forward. For 3B models, we observe lower-PPL SNLP configurations but do not yet realize wall-clock speedups with our current PyTorch-level implementation, likely because the wider sequential blocks already saturate the H100 more effectively.

These lower-PPL cases should not be interpreted as monotonic inference-time scaling. Exact convergence of the Newton formulation recovers the sequential trace. The improvement arises because practical SNLP uses approximate structured corrections, finite iteration counts, chunking, fusion, and initialization choices; together these define a distinct inference computation. We therefore interpret SNLP as a form of solver-induced inference bias: an approximate solver over depth can sometimes produce a better computation path than strict sequential execution, while still retaining enough structure to be accelerated. Our contributions are:

• We formulate layer-parallel language-model inference as structured surrogate Newton solving over the hidden-state trace, instantiated as IDN, DiagN, and HCN.

• We introduce SNLP-aware regularization, which improves layer-parallel compatibility and can also improve sequential perplexity.

• We introduce chunkwise layer fusion, which groups multiple depthwise-parallel layers into wider executable chunks before applying the structured Newton correction.

• We analyze the resulting solver-induced inference bias through correction ordering, propagation, variance-reduction, and layer-coupling ablations.

2 Related Work↩︎

Parallel nonlinear solvers. SNLP builds on the view that a sequential computation can be solved as a coupled nonlinear system. DEER applies Newton’s method to nonlinear recurrences and uses parallel scan to solve the resulting linearized dynamics [12]; later work extends this perspective to MCMC chains [20] and improves stability and scalability with quasi-Newton and Kalman-style approximations [16]. Song et al. [14] frame feedforward computation as parallel nonlinear equation solving, and Jacobi decoding applies fixed-point iteration to parallelize autoregressive translation [15]. Deep Equilibrium Models [21] take a complementary view, finding fixed points of weight-tied infinite-depth networks via root-finding. Our work rotates this line of work from sequence length to Transformer depth, and focuses on structured surrogates that avoid full layer Jacobians.

Associative scans and structured recurrences. Parallel prefix scan is a classical primitive [22] that has become central to efficient recurrent and state-space models. Linear recurrent networks can be parallelized over sequence length with scan [23]; structured state-space models such as S4 [24] and Mamba [25] use related hardware-aware recurrent kernels and scan-style algorithms. SNLP uses the same computational principle for depthwise correction: when the surrogate is identity, diagonal, or a small matrix, the Newton correction becomes a cheap structured recurrence.

Depth mixing and residual architectures. Residual connections [26] are central to deep Transformer training, and several architectures modify how information flows across depth. Hyper-Connections and mHC introduce learned residual-stream mixing and stabilization mechanisms [17], [18]; AttnRes replaces fixed residual accumulation with learned attention over previous layer outputs [27]. Value residual learning and x0-style residual connections also alter how features persist through depth [28], [29]. Weight-tied and looped architectures, including Universal Transformers [30], ALBERT [31], recurrent-depth models [32], and Hyperloop Transformers [33], reuse layers across depth. SNLP is complementary: rather than only changing the forward architecture, it asks whether the resulting depth dynamics expose a cheap surrogate for Newton-style layer-parallel inference.

Efficient language-model inference. Most efficient LLM inference work accelerates token-level decoding through batching, KV caching [7], quantization, memory-aware execution [34], kernel engineering [6], speculative decoding [8], [9], early exit [35], or serving systems [36], [37]. These techniques improve the execution of the standard sequential layer stack, whereas SNLP targets a different bottleneck: the dependency chain across layers for a fixed token prefix. Our experiments use Nanochat as a compact from-scratch training and evaluation harness [19]; we also run preliminary post-hoc and finetuning experiments on representative open-weight decoder-only models, including Qwen2.5, TinyLlama, and Gemma [38]–[40]. The gap between trained-from-scratch and off-the-shelf results suggests that layer-parallel inference benefits from training/inference co-design, leaving stronger pretrained-model adaptation to future work.

3 Methods↩︎

3.1 Background↩︎

Layer traces as residual equations. Consider a depth-\(L\) model with hidden states \(h_0,\ldots,h_L\) and layer maps \[h_l = f_l(h_{l-1}), \qquad l=1,\ldots,L .\] Here \(l\) indexes depth, while superscripts such as \((k)\) will index iterative solver steps. Rather than viewing the forward pass only as a sequential program, we can view the entire hidden-state trace \(\mathbf{h}=(h_1,\ldots,h_L)\) as the solution of a nonlinear residual equation. Define \[\label{eq:residual} G_l(\mathbf{h}) = h_l - f_l(h_{l-1}), \qquad G(\mathbf{h}) = (G_1(\mathbf{h}),\ldots,G_L(\mathbf{h})).\tag{1}\] The usual sequential forward pass is exactly the zero-residual trace \(G(\mathbf{h})=0\). This formulation exposes a different source of parallelism: instead of computing layers one after another, one may iteratively solve for all layer states jointly.

Newton-style updates over depth. DEER applies Newton’s method to nonlinear recurrences by linearizing the transition at the current iterate and solving the resulting linear recurrence in parallel [12], [20]. Rotating this view by \(90^\circ\), the depth axis of any block-sequential model–a Transformer, CNN, or recurrent stack–can be treated as the recurrence axis. At solver iteration \(k\), the exact Newton update over layers can be written as \[h_l^{(k+1)} = f_l\!\left(h_{l-1}^{(k)}\right) + J_l^{(k)} \left(h_{l-1}^{(k+1)} - h_{l-1}^{(k)}\right), \qquad J_l^{(k)} = \tfrac{\partial f_l}{\partial h_{l-1}}\!\left(h_{l-1}^{(k)}\right).\] This recurrence is equivalent to applying Newton’s method to the stacked residual system in Eqn. 1 because the residual Jacobian is block lower-bidiagonal; we refer readers to prior derivations of this equivalence in DEER-style solvers [16], [20]. The challenge is that \(J_l^{(k)}\) is the Jacobian of an entire layer or block output with respect to its input. For language-model hidden states, materializing this operator is infeasible, and even Jacobian-vector products or finite-difference approximations can consume the latency budget that layer parallelism is meant to save. Naive fixed-point updates avoid this cost but are often unstable on trained residual networks. The practical question is therefore whether we can replace the exact layer Jacobian with a cheap structured surrogate that preserves enough of the Newton correction to make finite-iteration, layer-parallel inference useful.

3.2 Structured Newton Layer Parallelism↩︎

SNLP replaces the exact layer Jacobian in Eqn. 1 with a cheap structured surrogate. Let the first \(S\) layers be evaluated sequentially, producing a prefix state \(h_S\). The remaining suffix \(\{S+1,\ldots,L\}\) is solved by iterative correction. At iteration \(k\), each suffix layer is first evaluated using the current estimate of its input, \[\widetilde{h}_l^{(k)} = f_l\!\left(h_{l-1}^{(k)}\right), \qquad l=S+1,\ldots,L .\] These evaluations are independent across \(l\) and can be batched or fused. SNLP then applies the structured Newton correction \[\label{eq:snlp-update} h_l^{(k+1)} = \widetilde{h}_l^{(k)} + A_l^{(k)} \left(h_{l-1}^{(k+1)} - h_{l-1}^{(k)}\right), \qquad h_S^{(k+1)}=h_S ,\tag{2}\] where \(A_l^{(k)}\) is a surrogate for the exact block Jacobian \(J_l^{(k)}\). If \(A_l^{(k)}=J_l^{(k)}\), this recovers the exact DEER/Newton update over depth. SNLP instead chooses \(A_l^{(k)}\) so that the correction is much cheaper than evaluating or materializing the true Jacobian, while still propagating information from earlier corrected layer states to later ones.

The update in Eqn. 2 separates the two costs that matter for inference. The nonlinear layer evaluations \(\widetilde{h}_l^{(k)}\) are parallel across the suffix and dominate GPU work. The Newton correction still propagates through depth, but because \(A_l^{(k)}\) is either trivial to compute or directly available from the architecture, this sequential correction is cheap relative to a Transformer block. Thus SNLP realizes speedup by parallelizing the expensive block forwards while keeping only a lightweight structured recurrence on the critical path. After \(K\) iterations, the model projects the final corrected state \(h_L^{(K)}\) to logits.

Effect of the correction. The correction in Eqn. 2 is what moves information across the whole suffix within a single solver iteration. Once the layer outputs \(\widetilde{h}_l^{(k)}\) are computed, the corrected prefix state propagates from layer \(S\) to layer \(L\) through the structured recurrence, so \(h_L^{(k+1)}\) depends on the corrections from all layers \(S+1,\ldots,L\). Without this correction, a naive parallel fixed-point update only advances information by one layer per iteration: after \(K\) iterations, the effect of the prefix can reach only the next \(K\) layers of the suffix. We verify this propagation effect empirically in 5.3 and 10.8.

3.3 Structured Surrogates↩︎

Identity Newton (IDN). For residual Transformer blocks, \(f_l(x)\) contains an explicit identity path. SNLP uses the architecture-induced surrogate \[A_l^{(k)} = I .\] The correction becomes \[h_l^{(k+1)} = \widetilde{h}_l^{(k)} + h_{l-1}^{(k+1)} - h_{l-1}^{(k)} ,\] which reduces the Newton correction to additive propagation of the previous-layer correction. This is our main residual-Transformer instantiation because it requires no Jacobian estimation and makes the correction essentially a prefix-sum over depth. We refer to this variant as Identity Newton (IDN).

Diagonal Newton (DiagN). A closer approximation to the exact Newton step uses only the diagonal of the layer Jacobian, \[A_l^{(k)} = \operatorname{diag}\!\left(J_l^{(k)}\right).\] This connects SNLP to quasi-DEER and ELK-style approximations [16]. With a diagonal surrogate, the correction in Eqn. 2 becomes an elementwise affine recurrence over depth and can be evaluated efficiently by an associative prefix scan [22], [23], [25]. In our implementation, the diagonal can be estimated by a Hutchinson-style finite-difference or VJP estimator [20], [41], [42], optionally only on a subset of layers.

HC Newton (HCN). For hyper-connection and mHC-style models [17], [18], the architecture exposes an explicit residual mixing matrix over streams. If a block applies residual mixing matrices \(H^{\mathrm{res}}_{\mathrm{attn},l}\) and \(H^{\mathrm{res}}_{\mathrm{mlp},l}\), we use \[A_l = H^{\mathrm{res}}_{\mathrm{mlp},l} H^{\mathrm{res}}_{\mathrm{attn},l}. \label{eq:hcn-A}\tag{3}\] This surrogate is small: it acts on the stream dimension rather than on the full hidden dimension. The mHC case demonstrates that SNLP is not tied to the identity residual path; any architecture with a cheap structured approximation to inter-layer sensitivity can define an SNLP correction.

3.4 SNLP-Aware Training↩︎

Off-the-shelf sequential models need not have layer dynamics that match a cheap surrogate. We therefore add an auxiliary loss that makes a finite SNLP solve match the sequential trace. For each suffix length \(N\in\mathcal{S}\), let \(\mathcal{T}_N\) be the stride-selected supervised layers in that suffix, and let \(\widehat h_l^{\mathrm{SNLP}}(N,K;A)\) be the SNLP state at layer \(l\) after \(K\) iterations with surrogate family \(A\). We optimize \[\mathcal{L} = \mathcal{L}_{\mathrm{CE}} + \lambda \sum_{N\in\mathcal{S}} \sum_{l\in\mathcal{T}_N} \tfrac{ \left\|\widehat h_l^{\mathrm{SNLP}}(N,K;A)-h_l^{\mathrm{seq}}\right\|_2 }{ \left\|h_l^{\mathrm{seq}}\right\|_2+\epsilon } . \label{eq:snlp-training-loss}\tag{4}\] In our runs, \(K=1\) during training and \(\mathcal{S}\) contains one or more configured suffix lengths. The set \(\mathcal{T}_N\) controls where the matching loss is applied: stride 0 uses only the final layer, \(\mathcal{T}_N=\{L\}\), while positive strides add sparse intermediate layers and always include \(L\) to reduce memory cost; see 11 for ablations. The surrogate \(A\) is identity for IDN, diagonal for DiagN, and the stream-mixing matrix for HCN. This objective does not make layers removable; rather, it makes the chosen structured correction a better finite-iteration solver for the sequential trace.

3.5 Inference With Fusion and Chunking↩︎

At inference time, SNLP runs a sequential prefix and applies Eqn. 2 to a suffix of \(N=L-S\) layers. The suffix hidden states can be initialized from the prefix state \(h_S\), from a one-shot batched forward, or from a lightweight predictor; our main evaluations focus on simple prefix-state and batched-forward initializations. The number of iterations \(K\) controls the quality-cost tradeoff.

Layer fusion. Wall-clock speedups require more than replacing the Jacobian. We therefore combine SNLP correction with GPU-oriented execution of the suffix. In the batched form, per-layer weights are stacked so all suffix layers evaluate in one grouped operation. In the fused form, several layers that read the same input are combined into one wider layer: the attention \(Q,K,V\) projections and MLP expansion matrices are concatenated along their output dimension, while the attention output projection and MLP down-projection are concatenated along their input dimension. Equivalently, the fused layer computes all branch outputs in one wide matmul and performs the required sum-reductions after the attention output projection and after the MLP projection. This converts layer-parallel algorithmic structure into larger GPU-efficient matrix multiplies.

Chunkwise strategy. For more aggressive parallelization, we split the suffix into multiple fused chunks, inspired by DeltaNet-style chunkwise parallelization [43]. Each chunk is treated as a wide layer as above, and all chunk forwards are parallelizable because they use the current chunk-input estimates from iteration \(k\). SNLP then applies the structured Newton correction between chunk outputs rather than between individual layers: \[h_c^{(k+1)} = \widetilde{h}_c^{(k)} + A_c^{(k)} \left(h_{c-1}^{(k+1)} - h_{c-1}^{(k)}\right), \label{eq:snlp-iterative}\tag{5}\] where \(c\) indexes chunks and \(A_c^{(k)}\) is the corresponding identity or architecture-induced chunk surrogate. Chunking trades a coarser solver approximation for better hardware utilization, since the expensive work is executed as a small number of wide parallel chunk forwards followed by a cheap correction across chunks. These fusion choices change the finite-iteration computation, so the resulting model should be understood as practical SNLP inference rather than exact recovery of the sequential forward.

4 Analysis↩︎

Exact convergence of Newton’s method on Eqn. 1 recovers the sequential forward pass, so lower-PPL SNLP configurations should not be interpreted as monotonic inference-time scaling. Practical SNLP uses approximate surrogates, finite iterations, initialization, fusion, and chunking; together these define a solver-induced inference bias. We summarize the main mechanisms here and defer derivations to the appendix.

Training-side effects. SNLP-aware training makes a cheap structured correction match the sequential final state. For residual blocks \(f_l(x)=x+g_l(x)\), IDN training encourages \(g_l(h_S)\approx g_l(h_l^{\mathrm{seq}})\) over the suffix, putting implicit Lipschitz pressure on the non-residual branch: smaller \(J_{g_l}\) makes \(J_{f_l}=I+J_{g_l}\) closer to the IDN surrogate. This can improve gradient flow and encourage capacity partitioning, with prefix layers handling more input-dependent processing and suffix layers acting more like stable feature-correction modules. The suffix is still useful, but its features become less sensitive to the exact state where they are evaluated.

Inference-side effects. IDN and fused SNLP evaluate many suffix contributions at a common or chunk-level input instead of along the fully accumulated sequential chain. This removes part of the variance from layerwise error compounding, at the cost of bias from evaluating \(g_l(h_S)\) rather than \(g_l(h_l^{\mathrm{seq}})\). When SNLP training makes this bias small, the variance reduction can dominate. Fusion and chunking add another bias: fused branches see summed cross-layer signals, which can create useful feature conjunctions but can also hurt when chunks are too aggressive.

Layer coupling. The results suggest that SNLP benefits from structured depth coupling rather than removing depth interactions entirely. HC and mHC parameterize residual-stream mixing [17], [18], while AttnRes learns attention over depth [27]. SNLP is complementary: IDN uses identity coupling, HCN uses the learned residual mixing matrix, and fused SNLP induces implicit cross-layer coupling inside each chunk.

5 Experiments↩︎

From-scratch models. Our main results use Nanochat [19] models trained from scratch at two 32-layer scales: a 3B model with \(n_{\mathrm{embd}}=2048\) and 16 heads, and a 0.5B model with \(n_{\mathrm{embd}}=640\) and 5 heads. Forward Nanochat uses rotary position embeddings [44], an \(x_0\) residual connection [29], and value embeddings [28]. For both standard scales, we train No Reg., IDN Reg., and DiagN Reg. variants. We also train 0.5B variants without \(x_0\)/VE to isolate SNLP from Nanochat-specific features, and an mHC model [17], [18], where HCN uses the learned matrix surrogate.

Off-the-shelf models. We also evaluate SNLP post hoc on Qwen2.5-0.5B-Instruct [38], TinyLlama-1.1B-Chat-v1.0 [39], and Gemma-3-1B-it [40]. Their best SNLP configurations can match, but do not improve, sequential perplexity and are slower than sequential execution; see 16. We additionally finetune TinyLlama to test whether SNLP compatibility can be introduced after pretraining.

Inference configurations. We evaluate chunkwise SNLP with IDN correction for residual models and HCN correction for mHC models. A configuration NxFM-init denotes \(N\) parallel chunks, each fusing \(M\) layers, with initialization h0 from the prefix state \(h_S\) or fwd from a one-shot parallel forward \(f_l(h_S)\). For example, 8xF2-fwd uses 8 fused 2-layer chunks with one-shot initialization. Preheat initialization is deferred to the appendix.

Metrics and protocol. All models are implemented in PyTorch [45] using HuggingFace Transformers [46]. Training and evaluation use ClimbMix, a FineWeb-Edu subset [47]; PPL is measured on a fixed 1M-token validation split. Timing uses batch size 1 on H100 GPUs with 50 warmup and 200 measured runs. Top-1 and LogitSim are prefill metrics against the original sequential model. AR Match is token-level greedy-generation agreement over 32 tokens; for fused-weight variants, it compares to sequential execution of the same fused model. EmbSim compares generated text to original sequential generation using BGE-small-en-v1.5 embeddings [48].

5.1 Main results↩︎

5 reports two representative SNLP configurations per model: one speed-oriented and one quality-oriented; checkpoint settings are listed in 10. All quality-oriented configurations, and most speed-oriented configurations, achieve lower PPL than the corresponding sequential forward pass. On 0.5B models, SNLP reaches up to \(2.3\times\) speedup while maintaining comparable or lower PPL; the fastest configurations use aggressive 24-layer parallelization, either 12xF2 or 4xF6, covering 75% of the model depth. Across both 3B and 0.5B models, SNLP-aware regularization also improves the sequential model itself: IDN/HCN and DiagN regularization reduce baseline PPL by 4.7%–23.4%. The 3B models also obtain lower PPL but not speedup in our implementation. At this width, sequential Transformer blocks already saturate the H100 more effectively, so PyTorch-level layer fusion does not overcome the overheads. Custom fused kernels or software-hardware co-design, such as compute-in-memory-style execution [49], may be needed to realize the algorithmic parallelism at larger scale.

The diagonal variant is useful as an ablation because it relaxes the strong identity assumption, but it adds extra block evaluations or autodiff work. 6 shows that, unlike IDN inference, most quality-oriented diagonal-correction configurations only recover comparable PPL to sequential inference. Compared with IDN, DiagN introduces a different solver-induced bias that is numerically closer to the sequential computation, but this bias is not necessarily beneficial for reducing PPL.

5.2 Training effect↩︎

SNLP-aware regularization improves the sequential standard models: on 3B, IDN and DiagN regularization reduce PPL from 37.16 to 35.31 and 35.41; on 0.5B, they reduce PPL from 69.54 to 53.25 and 63.08. This suggests that the loss changes learned layer dynamics rather than only fitting the parallel inference path. 7 measures the Jacobian of the non-residual branch \(g_l\) for the last 8 layers of the 0.5B standard models. For layers 24–30, IDN Reg. reduces spectral estimates by roughly \(12\times\) and Hutchinson Frobenius estimates by roughly \(12\times\), supporting the implicit Lipschitz-regularization interpretation. [tab:amplification] shows the same reduction in full-layer amplification \(\|J_{f_l}v\|/\|v\|\), including removal of the previous layer-31 outlier. For diagonal-Jacobian inference, VJP-based estimates may introduce additional bias, especially in HC/mHC-style models where asymmetric routing makes \(J^\top v\) deviate from the required forward sensitivity.

5.3 Inference effect↩︎

Correction ordering. 8 tests whether correction order matters on the IDN Reg. model; full tables are in 10.7. Sequentially executing the permuted suffix is almost invariant to layer order up to \(N=12\), with very low shuffle variance, suggesting that IDN regularization makes these suffix layers effectively parallelizable. SNLP with IDN correction is more order-sensitive: the best \(K=1\) PPL in the summary is always achieved by forward order, possibly because the correction recurrence preserves the causal direction of the original depth computation. At \(N=16\), however, sequential permutation reaches 47.0 PPL with the best shuffle; and in an earlier, less parallelized IDN \(\lambda=0.5\), stride-0 checkpoint, 23 shows that a random shuffle can outperform forward order. This suggests that single-iteration SNLP may benefit from searching over correction orders.

Correction propagation. Layer-activation ablations in 10.8 confirm the propagation pattern from 3: with correction, every active suffix layer can affect the final output in one iteration; without correction, influence moves only one layer per iteration, producing a staircase pattern.

Variance reduction. For each suffix layer \(l\), write \(f_l(h)=x_{\mathrm{in},l}(h)+g_l(h)\), where \(g_l\) is the non-residual attention/MLP update. Let \(h_S\) be the clean prefix state and \(h_l^{\mathrm{seq}}\) the sequential input to layer \(l\). We define the substitution error \(\epsilon_l=g_l(h_l^{\mathrm{seq}})-g_l(h_S)\) and report: \[|\epsilon_l| = \tfrac{ \mathbb{E}_{b}\left[\|\epsilon_{l,b}\|_2\right] }{ \mathbb{E}_{b}\left[\|h_{S,b}\|_2\right] + \epsilon }, \qquad C_\epsilon(l) = \tfrac{ \mathbb{E}_{b}\left[ \left\|\sum_{r=S+1}^{l}\epsilon_{r,b}\right\|_2 \right] }{ \mathbb{E}_{b}\left[\|h_{S,b}\|_2\right] + \epsilon }, \qquad \Delta_g(l) = \tfrac{ \mathbb{E}_{b}\left[ \left\|g_l(h_{l,b}^{\mathrm{seq}})-g_l(h_{S,b})\right\|_2 \right] }{ \mathbb{E}_{b}\left[\left\|g_l(h_{S,b})\right\|_2\right] + \epsilon }.\] Here \(|\epsilon_l|\) measures the per-layer substitution error relative to the prefix-state scale, \(C_\epsilon(l)\) measures how these errors accumulate through the suffix, and \(\Delta_g(l)\) measures sensitivity relative to the layer update magnitude. [tab:variance95reduction] shows that IDN Reg. keeps relative per-layer error in the 0.03%–0.15% range of \(\|h_S\|\), while No Reg. ranges from 2% to 24%; IDN Reg. also has lower \(\Delta_g\) than No Reg. at every reported layer.

What \(J\approx I\) really means. Input-invariant suffix layers are not droppable. In an earlier d32 IDN model, early exiting before the last seven layers increases PPL by 52%, while evaluating those layers on the same prefix input preserves the sequential output; see [tab:not95droppable]. The same holds for the highly parallel IDN Reg. checkpoint used in our main 0.5B results and most ablations: skipping 8–16 suffix layers greatly increases PPL ([tab:early95exit95h5v3]). Thus \(J\approx I\) means that the non-residual features are nearly input-invariant, not unnecessary.

6 Conclusion↩︎

We introduced Structured Newton Layer Parallelism, a training and inference framework for relaxing the strict layerwise dependency in Transformer inference. SNLP treats the hidden-state trace across depth as a residual equation, but replaces exact layer Jacobians with cheap structured surrogates induced by the architecture or training objective. In residual Transformers this yields IDN, while mHC-style models use HCN with a learned matrix surrogate. With SNLP-aware regularization and chunkwise layer fusion, trained-from-scratch Nanochat models can execute groups of layers in parallel and, in 0.5B settings, reach up to \(2.3\times\) speedup with comparable or lower perplexity. These results should be interpreted through the finite-iteration solver rather than as monotonic inference-time scaling. Exact convergence of the residual formulation recovers the sequential trace, while practical SNLP uses structured surrogates, initialization choices, and chunkwise fusion to define a related but distinct computation. When training makes the surrogate accurate enough, this solver-induced inference bias can preserve or improve quality while exposing useful layer parallelism.

6.0.0.1 Limitations.

SNLP does not yet provide universal wall-clock acceleration: our 3B models improve PPL but do not speed up with the current PyTorch implementation, likely because the larger sequential blocks already saturate the H100. More efficient kernels, runtime support, or software-hardware co-design may be needed at larger scale. We also only observe lower-PPL SNLP inference on models trained from scratch with SNLP-compatible objectives; on off-the-shelf models such as Qwen2.5, SNLP can match but does not improve sequential perplexity.

Appendix↩︎

[1.4em]

[3.2em]

7 Algorithm↩︎

The main SNLP update is given in Eqn. 2 . 2 shows the standard layer-by-layer execution, while 3 shows the batched structured-correction implementation used by the NxF1-h0 configurations. For HCN, \(A_l\) acts only on the HC/mHC stream dimension rather than the full hidden dimension.

In 4, FD, JVP, and VJP change only how the diagonal estimator is obtained; once \(a_j=\operatorname{diag}(J_{S+j}^{(k)})\) is available, all variants use the same affine scan recurrence.

8 Analysis Details↩︎

8.1 Variance Reduction Derivation↩︎

Consider a residual suffix beginning at prefix state \(h_S\), with \[f_l(h)=h+g_l(h), \qquad l=S+1,\ldots,L .\] The sequential suffix computes \[h_L^{\mathrm{seq}} = h_S + \sum_{l=S+1}^{L} g_l(h_{l-1}^{\mathrm{seq}}),\] where each branch \(g_l\) is evaluated at a different accumulated hidden state. In contrast, one-step IDN initialized from \(h_S\) computes \[h_L^{\mathrm{idn}} = h_S + \sum_{l=S+1}^{L} g_l(h_S).\] Thus the difference is entirely due to evaluation points: \[h_L^{\mathrm{seq}} - h_L^{\mathrm{idn}} = \sum_{l=S+1}^{L} \left[g_l(h_{l-1}^{\mathrm{seq}})-g_l(h_S)\right].\] First-order expansion around \(h_S\) gives \[g_l(h_{l-1}^{\mathrm{seq}})-g_l(h_S) \approx J_{g_l}(h_S)\left(h_{l-1}^{\mathrm{seq}}-h_S\right) = J_{g_l}(h_S) \sum_{r=S+1}^{l-1} g_r(h_{r-1}^{\mathrm{seq}}).\] Sequential execution therefore couples every later layer deviation to all earlier deviations through the branch Jacobians. If the branch contributions have covariance scale \(\sigma^2\) and \(\|J_{g_l}\|\le \rho\) over the suffix, this term contributes a variance scale of order \(\rho^2\sum_l(l-S)\sigma^2\) under a first-order independence approximation. IDN removes this compounding term by evaluating all branches at the same prefix state. The price is bias: \[\operatorname{Bias}_{\mathrm{idn}} = \sum_{l=S+1}^{L} \left[g_l(h_S)-g_l(h_{l-1}^{\mathrm{seq}})\right].\] SNLP-aware training directly reduces this bias by making the branch functions less sensitive over the suffix trajectory. When the induced bias is small, the variance reduction from avoiding chain-wise error compounding can improve the finite-iteration inference path.

8.2 Connection to Hyper-Connections and mHC↩︎

Hyper-Connections expand the residual stream into \(M\) streams and mix them through learned matrices [17]. A simplified HC sublayer can be written as \[x_{l+1} = H^{\mathrm{res}}_l x_l + H^{\mathrm{post}}_l F_l\!\left(H^{\mathrm{pre}}_l x_l\right),\] where the \(H\) matrices act on stream dimension and \(F_l\) is the nonlinear branch. For an mHC Transformer block, the attention and MLP wrappers each have their own residual mixing: \[\begin{align} x'_l &= H^{\mathrm{res}}_{\mathrm{attn},l} x_l + H^{\mathrm{post}}_{\mathrm{attn},l} \operatorname{Attn}_l\!\left(H^{\mathrm{pre}}_{\mathrm{attn},l}x_l\right),\\ x_{l+1} &= H^{\mathrm{res}}_{\mathrm{mlp},l} x'_l + H^{\mathrm{post}}_{\mathrm{mlp},l} \operatorname{MLP}_l\!\left(H^{\mathrm{pre}}_{\mathrm{mlp},l}x'_l\right). \end{align}\] The exact sublayer Jacobians contain both the residual mixing and the nonlinear branch sensitivity: \[\begin{align} J_{\mathrm{attn},l} &= H^{\mathrm{res}}_{\mathrm{attn},l} + H^{\mathrm{post}}_{\mathrm{attn},l} J_{\operatorname{Attn}_l} H^{\mathrm{pre}}_{\mathrm{attn},l},\\ J_{\mathrm{mlp},l} &= H^{\mathrm{res}}_{\mathrm{mlp},l} + H^{\mathrm{post}}_{\mathrm{mlp},l} J_{\operatorname{MLP}_l} H^{\mathrm{pre}}_{\mathrm{mlp},l}. \end{align}\] If training makes the nonlinear branches locally input-invariant, the branch Jacobian terms become small and the block Jacobian is approximated by \[J_{\mathrm{block},l} \approx H^{\mathrm{res}}_{\mathrm{mlp},l} H^{\mathrm{res}}_{\mathrm{attn},l}.\] This is exactly the HCN surrogate used in Eqn. 2 . It is available from the architecture and acts only on stream dimension, avoiding full hidden-state Jacobian estimation.

IDN is the degenerate residual case. For \(f_l(x)=x+g_l(x)\), the known residual Jacobian is \(I\). IDN uses \(A_l=I\) and trains the branch sensitivity \(J_{g_l}\) to be small enough that \(J_{f_l}=I+J_{g_l}\approx I\). In this sense, IDN and HCN follow the same template: keep the cheap architecture-induced residual transition, and train the nonlinear branch to be compatible with it.

8.3 Fused Cross-Layer Coupling↩︎

Layer fusion combines several parallel layers that read the same input \(h\) into one wide layer. For a chunk \(\mathcal{C}\), the fused attention computes the branch projections for all \(l\in\mathcal{C}\) and sums the output projections: \[a_{\mathcal{C}}(h) = \sum_{l\in\mathcal{C}} O_l\,\operatorname{Attn}_l(Q_lh,K_lh,V_lh).\] The fused MLP then receives the shared post-attention state \[u_{\mathcal{C}}(h)=h+a_{\mathcal{C}}(h),\] and applies the concatenated expansion and down-projection weights, equivalently summing the per-layer MLP branches: \[m_{\mathcal{C}}(h) = \sum_{l\in\mathcal{C}} P_l\,\phi\!\left(W_l\,\operatorname{Norm}(u_{\mathcal{C}}(h))\right),\] where \(\phi\) is the MLP nonlinearity. The fused chunk output is \[F_{\mathcal{C}}^{\mathrm{fused}}(h) = u_{\mathcal{C}}(h)+m_{\mathcal{C}}(h).\] This is not identical to independently evaluating and summing each layer branch, because every MLP branch sees the aggregate attention state \(h+\sum_{r\in\mathcal{C}}a_r(h)\) rather than only its own \(h+a_l(h)\). The cross-layer coupling term is \[\sum_{l\in\mathcal{C}} P_l \left[ \phi\!\left(W_l\operatorname{Norm}\!\left(h+\sum_{r\in\mathcal{C}}a_r(h)\right)\right) - \phi\!\left(W_l\operatorname{Norm}\!\left(h+a_l(h)\right)\right) \right].\] This term captures the change induced by evaluating every MLP branch on the aggregate post-attention state rather than on its own per-layer post-attention state. It is generally nonzero: layer normalization, nonlinear MLP activations, and non-canceling projections can all make a branch react to attention evidence produced by other layers in the same chunk. This implicit cross-layer coupling connects fused SNLP to the broader observation that richer depth mixing can be beneficial, as in HC/mHC and AttnRes [17], [18], [27], while also explaining why overly aggressive fusion can change model behavior.

9 Training Configuration and Ablations↩︎

9 summarizes an early training-side ablation that motivated SNLP-aware regularization. Vanilla Jacobi, which initializes every layer state from the prefix state and repeatedly applies \(h_l^{(k+1)}=f_l(h_{l-1}^{(k)})\), diverged on trained models: measured layer Jacobian norms were far from contractive. Spectral regularization reduced this norm substantially, but it only modestly reduced the number of Newton iterations and did not address the dominant cost of Jacobian estimation. This motivated the IDN/HCN direction: instead of making exact Newton cheaper, train the model so a cheap structured correction is useful.

10 lists the training settings used for the trained-from-scratch models in 5. 11 shows that the best loss configuration depends on architecture and scale: the 0.5B standard model benefits strongly from stride-3 IDN regularization, while detaching the MSE target hurts; the no-x0/VE model works best with stride 6; the 3B model prefers stride 0 over stride 3. For HCN on the mHC model, not detaching hurts PPL, so our default is no detach for IDN and detach for DiagN/HCN.

10 Additional Inference Ablations↩︎

10.1 Interpreting AR Match Rate↩︎

12 gives a rough conversion between the observed greedy autoregressive match rate and an effective local per-token agreement rate. For AR Match, each sample is evaluated over \(T=32\) generated tokens. The reported AR match rate \(\alpha\) is token-level agreement against the sequential baseline over \(N\) samples \(\alpha = \tfrac{\#\text{matched generated tokens}}{N T}\). Because generation is autoregressive, a single mismatch changes the prefix for all later positions. Under a simple absorbing-divergence model, let \(\beta\) be the probability of matching the next token conditioned on all previous generated tokens matching. Then token \(i\) matches with probability \(\beta^i\), so \[\alpha = \tfrac{1}{T}\sum_{i=1}^{T}\beta^i = \tfrac{\beta(1-\beta^T)}{T(1-\beta)} .\] Thus an apparently modest AR match can still imply high local agreement before prefix divergence; for \(T=32\), \(\alpha=0.30\) corresponds to \(\widehat\beta\approx0.905\).

10.2 Coupling SNLP with Jacobi Decoding↩︎

SNLP parallelizes over the layer axis, while Jacobi decoding (JD) parallelizes over future token positions by iteratively refining a block of draft tokens [15]. Speculative decoding and speculative Jacobi decoding (SJD) further combine draft proposals with verification [8], [9], [50]. A natural extension is to couple the two axes and solve over a layer-token lattice \[h_{\ell,t}, \qquad \ell=0,\ldots,L,\quad t=1,\ldots,T .\] In principle, one could update hidden states and token guesses jointly, \[h_{\ell,t}^{(r+1)} = \mathcal{F}_{\ell,t}\!\left(h^{(r)}, x^{(r)}\right), \qquad x_t^{(r+1)} = \operatorname{JDUpdate}\!\left(\operatorname{LMHead}(h_{L,t}^{(r+1)})\right),\] rather than using a nested loop, \[x^{(r+1)} \leftarrow \operatorname{JDUpdate}\!\left( \operatorname{SNLPForward}(x^{(r)};K) \right).\] This view is appealing because a successful 2D solver could avoid paying a full inner SNLP solve for every token-level Jacobi iteration. In practice, the main difficulty is initialization. JD changes the draft tokens between iterations; reinitializing all layer states from the new prefix state \(h_0\) reduces to the nested baseline, while reusing hidden states from the previous JD iteration carries features computed for old draft tokens.

13 evaluates this design on the 0.5B IDN Reg. model. The naive variants run a fresh SNLP forward inside each JD iteration. The h0-JD variants recompute the first parallel input but warm-start deeper parallel states from the previous JD iteration. Naive \(K=1\) already obtains 100% token match across all tested configurations, leaving no inner-iteration gap for coupling to close. Increasing to \(K=2\) adds cost without improving match. By contrast, h0-JD degrades match to roughly 21–37%, because stale hidden states encode features for old draft tokens and are propagated by the IDN correction. SJD-style variable-length acceptance makes direct hidden-state reuse even less straightforward, so the tested SJD path effectively reduces to naive composition. We therefore leave useful 2D layer-token coupling to future work; it likely requires a better transport or reinitialization rule when draft tokens change.

10.3 ELK Tempering and Preheat Initialization↩︎

We also explored two auxiliary inference knobs that are not included in the main configuration search. The first is an ELK-style tempering of the Newton correction [16], which scales the correction term without changing the number of block forwards and therefore should not affect speed. The second is preheat initialization: offline calibration fits a low-rank affine predictor for each layer output, \[\widehat h_l(x_0) = (x_0 V_l) U_l^\top + b_l ,\] where the basis \(V_l\) is obtained from a truncated SVD of calibration embeddings and \(U_l,b_l\) are fit by linear regression to sequential hidden states. At inference time, \(\widehat h_l(x_0)\) initializes the parallel suffix. In the table below, preheat was calibrated on random tokens; later validation-set calibration did not consistently improve results.

14 shows that ELK tempering can substantially reduce PPL for some configurations, but the effect is not uniform. Preheat is also inconsistent: it can be close to h0, but it can also be much worse. To reduce search space, our main experiments do not tune ELK or preheat; better quality may be achievable at little or no additional runtime cost with a more systematic search.

10.4 Off-the-Shelf Model Results↩︎

16 summarizes the best post-hoc SNLP configurations for off-the-shelf models. We select a fast configuration as the fastest run within \(\pm 8\%\) PPL of sequential, and a quality configuration as the lowest-PPL run, breaking ties within 1% PPL by speed. Matching sequential perplexity requires multiple Newton iterations and does not produce speedup, supporting the need for SNLP-aware training. Gemma’s absolute PPL is unusually high on this WikiText-style evaluation: on short raw text, Gemma-3-1B is comparable to Qwen2.5-0.5B and TinyLlama-1.1B, but its WikiText PPL is much higher, likely because its instruction tuning is less compatible with the article formatting and structure in this benchmark.

10.5 Diagonal-Jacobian Correction Results↩︎

17 benchmarks the local cost of Jacobian-vector primitives used by the diagonal-correction sweep in 6. Finite-difference (FD) estimates \(Jv\) by evaluating \(f(x+\epsilon v)-f(x)\), exact forward-mode JVP returns both \(f(x)\) and \(Jv\) in one call, and VJP uses reverse-mode autodiff to compute \(J^\top v\). FD is efficient in our solver because \(f(x)\) is already computed by the parallel block forward, so the finite difference only adds the extra \(f(x+\epsilon v)\) evaluation, similar to the cost of forward initialization. However, FD often diverges in our sweep, plausibly because the correction stacks multiple approximations: replacing the full Jacobian by its diagonal, estimating that diagonal with a Hutchinson-style probe, and then using a noisy finite difference. JVP and VJP are more stable but expensive: 17 shows roughly \(3\times\) overhead for VJP and \(6\)–\(8\times\) for JVP over a plain forward. Exact JVP also lacks fused-SDPA support, so the diagonal experiments use FD or VJP estimators.

For an earlier mHC baseline, no diagonal Newton configuration reaches the 10% fast threshold relative to sequential PPL 88.67: the best diagonal result reaches 104.2 PPL, or +17.5% above sequential. With the x0+VE mHC models in 6, diagonal correction can match PPL but remains much slower. This failure mode is expected: mHC routing through \(H_{\mathrm{res}}\), \(H_{\mathrm{pre}}\), and \(H_{\mathrm{post}}\) is asymmetric, so a VJP estimates \(J^\top v\) rather than the required \(Jv\) direction. The resulting diagonal approximation is poor, causing slow convergence or divergence.

10.6 Pretrained TinyLlama Fine-Tuning↩︎

We finetune TinyLlama-1.1B-Chat-v1.0 [39] with IDN regularization to test whether layer-parallel compatibility can be retrofitted onto an off-the-shelf pretrained model. TinyLlama is a standard LLaMA-style model with 22 layers, hidden dimension 2048, 32 attention heads, and 4 KV heads. We finetune and evaluate on ClimbMix [47] using AdamW with cosine learning-rate decay. 18 summarizes the finetuning grid, and 19 reports representative IDN inference results.

The finetuning results support the co-design interpretation. Mild IDN regularization preserves base quality but leaves large \(K=1\) gaps, while aggressive regularization reduces the IDN loss but degrades sequential PPL. The best runs can bring \(K=2\) one-shot-forward SNLP close to sequential perplexity, but they do not produce the lower-PPL behavior observed in models trained from scratch. This suggests that pretrained layer Jacobians are difficult to reshape late in training without damaging the base model.

10.7 Correction Ordering↩︎

20 21 22 report the correction-ordering ablations used to construct the summary in 8. 23 reports an earlier IDN \(\lambda=0.5\), stride-0 checkpoint with explicit permutations. AR-ness is the average of local- and global-AR-ness scores as defined in DiffuCoder [51]. At \(K=1\), several non-forward orders remain usable; after repeated correction at \(K=4\), only forward ordering remains stable.

10.8 Correction Propagation↩︎

24 compares subset ablations with and without Newton correction. The identical columns in the no-correction block show that information advances only one layer per iteration, while the correction block shows that all active suffix layers can influence the output immediately.

References↩︎