Reproduce “SNLP: Layer-Parallel Inference via Structured Newton Corrections” on OpenResearch

About this paper

Ligong Han, Kai Xu, Hao Wang, Akash Srivastava

A new method called Structured Newton Layer Parallelism (SNLP) enables parallel inference across Transformer layers by framing hidden state computation as solving a coupled nonlinear residual equation. This approach achieved up to 2.58x wall-clock speedup on 0.5B models, with configurations maintaining near-sequential perplexity at 1.40x speedup.

Key insights

• The prohibitive cost of exact Jacobian computation in Newton's method can be overcome by using cheap, structured Jacobian surrogates derived from Transformer architecture, enabling practical layer-parallel updates.
• Co-designing training with SNLP inference is critical, as sequential models are not inherently compatible with approximate parallel execution; SNLP-aware training methods mitigate "solver-induced inference bias" and improve compatibility.
• The efficacy of SNLP stems from exposing a practical speed-quality frontier by leveraging controlled inference bias, rather than perfectly recovering the exact sequential forward pass.

Summary by alphaXiv.