Title: Efficient Large Language Model Inference with Neural Block Linearization
Venue: NeurIPS 2025
Accepted: yes
Keywords: Inference

Transformer

LLM

CCA

Lightweight Architectures
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.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nNeural Block Linearization (NBL) is a method that replaces blocks of a neural network with linear functions, specifically analytically derived Linear Minimum Mean Squared Error (LMMSE) estimators. The proposed selection criteria of layers for pruning is based on Canonical Correlation Analysis (CCA). NBL is best used to replace the attention layers in Transformer LLMs, and ablations are performed using it to replace entire Transformer blocks (attention + feedforward layers), speeding up the autoregressive generation process while maintaining reasonable performance compared to other structural pruning methods. The results demonstrate that certain attention layers in some models can be replaced with analytically computed linear functions with minimal decreases in performance and some improvement to generation speed.\n\nSTRENGTHS AND WEAKNESSES:\nThe submission is technically sound and clearly written.  \nThe claims of increased inference speed with minimal loss of accuracy for the method are well supported within ablations on number and type of layers replaced.  \nHowever, the evaluation is limited as some structural pruning baselines are missing (such as PruneNet, ICLR 2025), and some models have no pruning baselines at all (DeepSeek-R1-Distill-Llama-8B and Llama-3.1-70B).  \nThe levels of sparsity are unclear for the given methods in the tables besides SliceGPT.\n\nWithout fair comparison to pruning baselines, the significance of the results is diminished.  \n\nThe results from the ablation on CCA-Bound criterion vs cosine distance, as well as for greedy selection are overclaimed: Tables 17-19 lack uncertainty estimates and do not show meaningful difference between CCA and the cosine/greedy criteria.\n\nQUESTIONS:\nCan you report evaluations of SliceGPT and SLEB for Tables 4-5, as well as PruneNet (or other structured pruning baselines that you think are more suitable) for all performance evaluations?  \nIncluding more comparisons will make me more inclined to raise the score.\n\nCan you report the level of sparsity for SLEB/Block/Attn Drop/NBL-m, similar to SliceGPT or other structured pruning methods?  \nThis will make the sparsity comparison between methods more clear.\n\nCan you add ablations on replacing the feedforward layers? Or, explain why the only ablations are on dropping attention layers or Transformer blocks (equivalent to dropping attention+feedforward layers)?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 4, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-06-26"}

comments: SUMMARY:
This paper introduces Neural Block Linearization (NBL), transforming the self-attention layers into linear approximations by Linear Minimum Mean Squared Error estimators. After transforming the attention layers from the existing LLM into linear layers, NBL effectively balances inference speed-up with accuracy.

STRENGTHS AND WEAKNESSES:
Strengths:
1. This paper is well-written and easy to follow. 
2. NBL uses the Canonical Correlation Analysis (CCA) to guide the transformation of attention layers, building on a good theoretical foundation. 
3. Experimental results show the effectiveness of NBL on reasoning tasks.

Weaknesses:
1. Computing the linear approximations requires a longer time than other pruning methods (e.g., SliceGPT).
2. Some typos: 
- Inconsistent method name between “Attn SLEB-m” in Line 218 and ‘’SLEB-m’’ in Table 2, 3.
- As the shape of $X$ is $\mathbb{R}^{(s \cdot t)} \times h_{in}$ in equation(1), equation (2) should be $\hat{Y} = X W + b$.

QUESTIONS:
1. The capability of linear layers may be limited by the context length of the calibration set. Can using a calibration set with relatively short context length still be good at handling long context?
2. For tasks that require long-distance dependencies, will this linearization cause a "cliff-like" drop in model performance?

LIMITATIONS:
yes
TITLE: Official Review
CLARITY: 3
IMPACT: 3
ORIGINALITY: 3
DATE: 2025-06-29

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work proposes Neural Block Linearization (NBL) as a method of replacing non-linear or complex layers / blocks in a deep learning architecture by the optimal linearization obtained via the SVD of the input-output data matrix over a calibration dataset. This method provides clearer error bounds compared to pruning or other methods for speeding up inference. In experiments on NBL replacing Self-Attention blocks in LLMs, the author can improve inference speeds at only slight performance degradation.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n- Theoretical derivation, analysis and bounds\n- No training needed\n- Identification of target layers with minimal performance impact after linearization\n\nWeaknesses:\n- There is no history for linearized attention blocks, outputs depending on the past are approximated by a linear average over typical outputs for this input in the calibration dataset\n- When put to extreme (i.e. replacing all self-attention layers) this method should evolve into a bi-gram-like model, that only takes into account the previous token\n- Calibration dataset needs to cover the whole inference distribution - acts similar as a training dataset\n- missing limitations section\n\nQUESTIONS:\n- Why is the limitation that applying NBL to attention blocks removes the history not covered in your paper?\n\n- Is there a possiblity to extend this to general random input vectors by just analyzing the QKV matrices and how they apply to random vectors? This might reduce the reliance on the calibration dataset and improve the bounds to general vectors beyond calibration \"in-distribution\".\n\n- It would be interesting to analyze the vector norms of the layers in comparison to the residual vector norms. This could explain why later layers are more likely to be chosen.\nCould you investigate those?\n\nLIMITATIONS:\nThe paper is missing a limitations section.\n\nThe paper does not cover the limitation of its method that for linearizing self-attention to a linear relationship between per-token input and output is reduced to a 1-convolution.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-06-30"}

comments: SUMMARY:
The authors introduce a fine-tuning-free method that uses Canonical Correlation Analysis (CCA) to identify and replace attention blocks in pretrained LLMs with their best linear approximations. This approach significantly reduces computational costs while largely preserving the model's original performance. The key innovation is identifying which attention blocks behave most like linear layers and can therefore be simplified without extensive retraining.

STRENGTHS AND WEAKNESSES:
Strengths: 

1. The paper is generally well-written and the motivations and methods are clearly stated. Necessary theoretical analysis is provided. 
2. The authors conduct rather complete evaluation with ablation study, demonstrates the strong performance of the proposed method. 

Weaknesses: 

1. I'm wondering what is the purpose of using the upper bound of Theorem 3.2 in the algorithm rather than directly evaluating the NMSE by definition. To obtain the sum of $\rho_i$, you need to compute $C_W$ already (or the trace of $C_W$). By a simple expansion of the MSE, it is not hard to see that 
$$\mathrm{MSE} = \mathrm{Tr}[C_{YY} - C_{YX} C_{XX}^{-1} C_{YX}^\top].$$
It seems calculating both objectives have roughly the same cost. I don't see the necessity of using the approximation. I'm happy for further discussion on this point. 
2. Is it a typo? In the algorithm output line 11: "Compressed LLM with linearized attention layers." -> linear layer actually.

QUESTIONS:
Already listed in the Weaknesses

LIMITATIONS:
N/A
TITLE: Official Review
CLARITY: 3
IMPACT: 3
ORIGINALITY: 3
DATE: 2025-07-03