SNLP: Schichtparallele Inferenz mittels strukturierter Newton-Korrekturen

Zusammenfassung

Autoregressive Sprachmodelle führen Transformer-Schichten sequentiell aus, was einen Latenz-Engpass erzeugt, der durch herkömmlichen Tensor- oder Pipeline-Parallelismus nicht beseitigt wird. Wir untersuchen, ob diese schichtweise Abhängigkeit gelockert werden kann, indem die Hidden-State-Spur über Schichten hinweg als Lösung einer nichtlinearen Residuumsgleichung behandelt und mit parallelen Newton-ähnlichen Aktualisierungen gelöst wird. Obwohl diese Sichtweise prinzipienbasiert ist, erfordern exakte Newton-Korrekturen teure Jacobi-Vektor-Produkte, und naive Fixpunktiterationen sind bei trainierten Transformatoren instabil. Wir führen den Strukturierten Newton-Schicht-Parallelismus (SNLP) ein, ein Trainings- und Inferenzframework, das exakte Schicht-Jacobis durch kostengünstige architekturinduzierte Ersatzdynamiken ersetzt. In Residual-Transformatoren ergibt dies Identity Newton (IDN), bei dem die Korrektur auf eine präfixsummenähnliche Aktualisierung reduziert wird; in mHC-artigen Architekturen verwendet HC Newton (HCN) die residuelle Mischmatrix des Modells. Wir führen ferner eine SNLP-bewusste Regularisierung ein, die Modelle darauf trainiert, mit einer oder wenigen strukturierten Newton-Iterationen den sequentiellen Vorwärtsdurchlauf genau zu approximieren. Experimente mit NanoChat-großen Transformatoren zeigen, dass SNLP-Regularisierung die Schichtparallel-Kompatibilität verbessert und auch die standardmäßige sequentielle Perplexität verbessern kann, wobei die Baseline-PPL um 4,7 %–23,4 % reduziert wird. Zur Inferenzzeit erzielt SNLP in Kombination mit Schichtfusion und blockweiser Zerlegung praktische Wanduhr-Beschleunigungen: Bei einem 0,5B NanoChat-Modell wird eine 2,3-fache Beschleunigung erreicht, während die PPL dennoch um 6,1 % verbessert wird. Diese Ergebnisse deuten darauf hin, dass schichtparallele Inferenz nicht nur eine numerische Approximation der sequentiellen Ausführung ist, sondern als nützlicher, lösungsinduzierter Inferenz-Bias wirken kann. Wir charakterisieren auch Einschränkungen: Fertig vortrainierte Modelle sind für dieses Verfahren weniger geeignet, und exakte Konvergenz führt zur Wiederherstellung der sequentiellen Berechnung und nicht zu einer monotonen Inferenzzeitskalierung.

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 further introduce SNLP-aware regularization, which trains models to make one or a few structured Newton iterations accurately approximate the sequential forward. Experiments on nanochat-scale Transformers show that SNLP regularization improves layer-parallel compatibility and can also improve standard sequential perplexity, reducing baseline PPL by 4.7%-23.4%. At inference time, SNLP combined with layer fusion and chunkwise decomposition achieves practical wall-clock speedups: on a 0.5B Nanochat model, it reaches 2.3x speedup while still improving PPL by 6.1%. These results suggest that layer-parallel inference is not merely a numerical approximation to sequential execution, but can act as a useful solver-induced inference bias. We also characterize limitations: off-the-shelf pretrained models are less amenable to this procedure, and exact convergence recovers the sequential computation rather than providing monotonic inference-time scaling.