Efficient Large Language Model Inference with Neural Block Linearization
Mete Erdogan, Francesco Tonin, Volkan Cevher
TL;DR
This work tackles the high inference cost of transformer-based LLMs by introducing Neural Block Linearization (NBL), which replaces selected self-attention layers with closed-form linear estimators learned via Linear Minimum Mean Squared Error (LMMSE). A Canonical Correlation Analysis (CCA) based bound quantifies the potential accuracy loss from linearization and guides layer substitution by ranking layers with the smallest bound, enabling calibration-free compression of pre-trained models. Empirically, NBL yields significant speedups (e.g., up to 32% in certain configurations) while maintaining competitive accuracy across multiple models and reasoning benchmarks, and it remains effective when combined with post-training quantization (AWQ) and speculative decoding. The approach provides a scalable, interpretable path toward deploying large LLMs in resource-constrained environments, with robust ablations and extensions to larger models and hardware-aware settings.
Abstract
The high inference demands of transformer-based Large Language Models (LLMs) pose substantial challenges in their deployment. To this end, we introduce Neural Block Linearization (NBL), a novel framework for accelerating transformer model inference by replacing self-attention layers with linear approximations derived from Linear Minimum Mean Squared Error estimators. NBL leverages Canonical Correlation Analysis to compute a theoretical upper bound on the approximation error. Then, we use this bound as a criterion for substitution, selecting the LLM layers with the lowest linearization error. NBL can be efficiently applied to pre-trained LLMs without the need for fine-tuning. In experiments, NBL achieves notable computational speed-ups while preserving competitive accuracy on multiple reasoning benchmarks. For instance, applying NBL to 12 self-attention layers in DeepSeek-R1-Distill-Llama-8B increases the inference speed by 32% with less than 1% accuracy trade-off, making it a flexible and promising solution to improve the inference efficiency of LLMs. The implementation is available at: https://github.com/LIONS-EPFL/NBL.