SNLP : Inférence parallèle par couches via des corrections de Newton structurées

Résumé

Les modèles de langage autorégressifs exécutent les couches Transformer séquentiellement, créant un goulot d'étranglement de latence que le parallélisme tensoriel ou pipeline conventionnel ne supprime pas. Nous étudions si cette dépendance inter-couche peut être assouplie en traitant la trace des états cachés à travers les couches comme la solution d'une équation résiduelle non linéaire et en la résolvant avec des mises à jour parallèles de type Newton. Bien que cette vision soit fondée, les corrections de Newton exactes nécessitent des produits jacobien-vecteur coûteux et les itérations naïves de point fixe sont instables sur les Transformers entraînés. Nous introduisons le Parallélisme de Couche Newton Structuré (SNLP), un cadre d'entraînement et d'inférence qui remplace les jacobiens de couche exacts par une dynamique surrogate peu coûteuse induite par l'architecture. Dans les Transformers résiduels, cela donne l'Identity Newton (IDN), où la correction se réduit à une mise à jour de type somme préfixe ; dans les architectures de type mHC, le HC Newton (HCN) utilise la matrice de mélange résiduelle du modèle. Nous introduisons également une régularisation consciente du SNLP, qui entraîne les modèles à effectuer une ou quelques itérations de Newton structurées pour approximer précisément le passage avant séquentiel. Les expériences sur des Transformers à l'échelle nanochat montrent que la régularisation SNLP améliore la compatibilité du parallélisme de couche et peut également améliorer la perplexité séquentielle standard, réduisant la perplexité de référence de 4,7 % à 23,4 %. En temps d'inférence, le SNLP combiné à la fusion de couches et à la décomposition par morceaux atteint des accélérations pratiques en temps réel : sur un modèle Nanochat de 0,5B, il atteint un facteur d'accélération de 2,3x tout en améliorant la perplexité de 6,1 %. Ces résultats suggèrent que l'inférence parallèle inter-couche n'est pas simplement une approximation numérique de l'exécution séquentielle, mais peut agir comme un biais d'inférence utile induit par le solveur. Nous caractérisons également les limitations : les modèles pré-entraînés prêts à l'emploi sont moins adaptés à cette procédure, et la convergence exacte retrouve le calcul séquentiel plutôt que de fournir un passage à l'échelle monotone en temps d'inférence.

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.