SNLP: Laag-parallelle Inferentie via Gestructureerde Newton-correcties

Samenvatting

Autoregressieve taalmodellen voeren Transformerlagen sequentieel uit, wat een latentieknelpunt creëert dat niet wordt weggenomen door conventionele tensor- of pijplijnparallelisme. We onderzoeken of deze laagafhankelijkheid kan worden versoepeld door de verborgen-toestandtrace over lagen te beschouwen als de oplossing van een niet-lineaire residuvergelijking en deze op te lossen met parallelle Newton-achtige updates. Hoewel deze benadering principieel is, vereisen exacte Newton-correcties dure Jacobiaan-vectorproducten en zijn naïeve vastpuntiteraties instabiel op getrainde Transformers. We introduceren Structured Newton Layer Parallelism (SNLP), een trainings- en inferentieframework dat exacte laag-Jacobianen vervangt door goedkope architectuur-geïnduceerde surrogaatdynamieken. In residu-Transformers levert dit Identity Newton (IDN) op, waarbij de correctie reduceert tot een prefix-som-achtige update; in mHC-achtige architecturen gebruikt HC Newton (HCN) de residu-mengmatrix van het model. We introduceren verder SNLP-bewuste regularisatie, die modellen traint om met één of enkele gestructureerde Newton-iteraties de sequentiële voorwaartse doorgang nauwkeurig te benaderen. Experimenten op nanochat-schaal Transformers tonen aan dat SNLP-regularisatie de laag-parallelle compatibiliteit verbetert en ook de standaard sequentiële perplexiteit kan verbeteren, met een reductie van de baseline PPL met 4,7%-23,4%. Bij inferentie bereikt SNLP gecombineerd met laagfusie en chunksgewijze decompositie praktische wandklokversnellingen: op een 0,5B Nanochat-model wordt een 2,3x versnelling bereikt terwijl de PPL nog steeds met 6,1% verbetert. Deze resultaten suggereren dat laag-parallelle inferentie niet slechts een numerieke benadering van sequentiële uitvoering is, maar kan fungeren als een nuttige door de oplosser geïnduceerde inferentiebias. We karakteriseren ook beperkingen: kant-en-klare voorgetrainde modellen zijn minder vatbaar voor deze procedure, en exacte convergentie herstelt de sequentiële berekening in plaats van monotone inferentie-tijdschaling te bieden.

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.