SNLP: Параллельная по слоям инференция через структурированные поправки Ньютона

Аннотация

Авторегрессионные языковые модели выполняют слои Трансформера последовательно, что создаёт узкое место по задержке, которое не устраняется обычным тензорным или конвейерным параллелизмом. Мы исследуем, можно ли ослабить эту послойную зависимость, рассматривая след скрытых состояний по слоям как решение нелинейного остаточного уравнения и решая его с помощью параллельных ньютоновских обновлений. Хотя такой подход является теоретически обоснованным, точные ньютоновские коррекции требуют дорогостоящих произведений Якобиана на вектор, а наивные итерации неподвижной точки оказываются нестабильными на обученных Трансформерах. Мы предлагаем структурированный ньютоновский параллелизм слоёв (Structured Newton Layer Parallelism, SNLP) — фреймворк для обучения и инференса, который заменяет точные якобианы слоёв дешёвыми суррогатными динамиками, порождёнными архитектурой. В остаточных Трансформерах это даёт тождественный ньютоновский метод (Identity Newton, IDN), где коррекция сводится к обновлению, подобному префиксной сумме; в архитектурах типа mHC — ньютоновский метод HC (HC Newton, HCN), использующий матрицу остаточного смешивания модели. Мы также вводим регуляризацию, учитывающую SNLP, которая обучает модели так, чтобы одно или несколько структурированных ньютоновских итераций точно аппроксимировали последовательный прямой проход. Эксперименты на Трансформерах масштаба nanochat показывают, что регуляризация SNLP улучшает совместимость с параллелизмом по слоям и может также улучшать стандартную последовательную перплексию, снижая базовую PPL на 4,7%–23,4%. Во время инференса SNLP в сочетании с объединением слоёв и поблочным разложением даёт практическое ускорение по астрономическому времени: на модели Nanochat с 0,5 млрд параметров достигается ускорение в 2,3 раза при одновременном улучшении PPL на 6,1%. Эти результаты позволяют предположить, что параллельный по слоям инференс — это не просто численная аппроксимация последовательного выполнения, но и полезное смещение, обусловленное решателем. Мы также характеризуем ограничения: готовые предобученные модели менее поддаются этой процедуре, а точная сходимость восстанавливает последовательные вычисления, а не обеспечивает монотонное масштабирование во время инференса.

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.