SNLP: Inferência Paralela por Camadas via Correções de Newton Estruturadas

Resumo

Modelos de linguagem autorregressivos executam camadas Transformer sequencialmente, criando um gargalo de latência que não é removido pelo paralelismo convencional de tensor ou pipeline. Estudamos se essa dependência entre camadas pode ser relaxada ao tratar o traço de estado oculto entre camadas como a solução de uma equação residual não linear e resolvê-la com atualizações paralelas no estilo Newton. Embora essa visão seja fundamentada, correções exatas de Newton exigem produtos Jacobiano-vetor dispendiosos, e iterações ingênuas de ponto fixo são instáveis em Transformers treinados. Apresentamos o Paralelismo de Camada Newton Estruturado (SNLP, na sigla em inglês), uma estrutura de treinamento e inferência que substitui Jacobianos exatos de camada por dinâmicas substitutas baratas induzidas pela arquitetura. Em Transformers residuais, isso resulta no Newton Identidade (IDN), onde a correção se reduz a uma atualização do tipo soma de prefixo; em arquiteturas estilo mHC, o Newton HC (HCN) utiliza a matriz de mistura residual do modelo. Introduzimos ainda uma regularização consciente do SNLP, que treina modelos para que uma ou poucas iterações Newton estruturadas aproximem com precisão o forward sequencial. Experimentos em Transformers de escala nanochat mostram que a regularização SNLP melhora a compatibilidade de paralelismo de camada e também pode melhorar a perplexidade sequencial padrão, reduzindo a PPL de base em 4,7%–23,4%. Em tempo de inferência, o SNLP combinado com fusão de camadas e decomposição por blocos atinge acelerações práticas de tempo real: em um modelo Nanochat de 0,5B, obtém aceleração de 2,3x enquanto ainda melhora a PPL em 6,1%. Esses resultados sugerem que a inferência paralela entre camadas não é meramente uma aproximação numérica da execução sequencial, mas pode atuar como um viés de inferência útil induzido pelo solver. Também caracterizamos limitações: modelos pré-treinados prontos para uso são menos adequados a esse procedimento, e a convergência exata recupera a computação sequencial em vez de fornecer escalonamento monotônico de inferência.

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.