SNLP: Inferencia Paralela por Capas mediante Correcciones Newton Estructuradas

Resumen

Los modelos de lenguaje autorregresivos ejecutan las capas Transformer de forma secuencial, creando un cuello de botella de latencia que no se elimina con el paralelismo convencional de tensores o de pipelines. Estudiamos si esta dependencia entre capas puede relajarse tratando la traza de estados ocultos a través de las capas como la solución de una ecuación residual no lineal y resolviéndola con actualizaciones paralelas de estilo Newton. Si bien este enfoque es sólido desde un punto de vista teórico, las correcciones exactas de Newton requieren costosos productos Jacobiano-vector y las iteraciones de punto fijo ingenuas son inestables en Transformers entrenados. Presentamos el Paralelismo de Capas con Newton Estructurado (SNLP, por sus siglas en inglés), un marco de entrenamiento e inferencia que reemplaza los Jacobianos exactos de las capas por dinámicas sustitutas baratas inducidas por la arquitectura. En Transformers residuales, esto da lugar a Newton Identidad (IDN), donde la corrección se reduce a una actualización similar a suma de prefijos; en arquitecturas de estilo mHC, Newton HC (HCN) utiliza la matriz de mezcla residual del modelo. Además, introducimos una regularización consciente de SNLP, que entrena a los modelos para que una o unas pocas iteraciones de Newton estructurado aproximen con precisión el forward secuencial. Experimentos en Transformers a escala nanochat muestran que la regularización SNLP mejora la compatibilidad con el paralelismo entre capas y también puede mejorar la perplejidad secuencial estándar, reduciendo la PPL de referencia entre un 4.7% y un 23.4%. En tiempo de inferencia, SNLP combinado con fusión de capas y descomposición por fragmentos logra aceleraciones prácticas en tiempo real: en un modelo Nanochat de 0.5B parámetros, alcanza una aceleración de 2.3x mientras sigue mejorando la PPL en un 6.1%. Estos resultados sugieren que la inferencia con paralelismo entre capas no es simplemente una aproximación numérica a la ejecución secuencial, sino que puede actuar como un sesgo de inferencia útil inducido por el solucionador. También caracterizamos las limitaciones: los modelos preentrenados estándar son menos adecuados para este procedimiento, y la convergencia exacta recupera el cómputo secuencial en lugar de proporcionar un escalado monótono en tiempo de inferencia.

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.