Efficient Large Language Model Inference with Neural Block Linearization
Authors: Mete Erdogan
Francesco Tonin
Volkan Cevher
Venue: NeurIPS 2025 poster ()
Keywords: Inference
Transformer
LLM
CCA
Lightweight Architectures
Primary area: deep_learning
TL;DR: We propose to replace self-attention layers with linear estimators through the derived CCA error bound, achieving inference speedups with favorable accuracy trade-off.

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.

Efficient Large Language Model Inference with Neural Block Linearization
Mete Erdogan, Francesco Tonin, Volkan Cevher
Laboratory for Information and Inference Systems École Polytechnique Fédérale de Lausanne (EPFL), Switzerland
Abstract
The high inference demands of transformer-basee Large Language Models (LLMs) pose substantial challenges in their deployment. To this end, we introduce Neural Block Linearization (NBL), 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 theoretical upper bound 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 multiple reasoning benchmarks. For instance, applying NBL to 12 self-attention layers in MiSSeekkRR1-Disil12-
increases the inference speed by 32% with less than 1% accuracy trade-off, making a flexible and promising solution to improve the inference efficiency of LLMs. The implementation is available bttps://githwe Btw..///S//SSSEELL///I
Introduction
Transformer-based models have become foundational in machine learning, with wide applications in NLP and language modeling [Vaswani al., 2017, Brown et al., 2020, Jiang et al., 2023, Achiam et al., 2023, Dubey et al., 2024]. Due to their ability to learn longerange dependencies and capture complex patterns, transformer models have achieved state-of-the-ar performance in tasks like language modeling, text generation, and translation. However, their growing size and complexity strongly limit the widespread adoption of Large Language Models (LLMs) real-world applications, especially in resource cost constrained scenarios.
therefore crucial to develop methods to reduce LLM inference costs. Proposed techniques mainly include weight pruning (Kusupati et al., 2020, Hoefler et al., 2021], low rank approximations [Yuan et al., 2023, Wang et al., 2024a], quantization [Lin al., 2024, Saha et al., 2024], speculative decoding [Leviathan et al., 2023, Cai al., 2024], distillation [Jiao et al., 2020, Liu et al., 2024] and subquadratic attention [Wang et al., 2024b, Zhang et al., 2025]. In this work, we focus on pruming, though our method can be integrated on top of other techniques such quantization.
Structured pruning methods eliminate specific layers, attention heads, or hidden dimensions to accelerate LLM inference [Voita et al., 2019, Ma et al., 2023, Xia et al., 2024a, Muralidharan et al., 2024, Men et al., 2024, Ashkboos et al., 2024, Song et al., 2024]. In fact, it is well-established that attention mechanisms exhibit redundancy [Voita et al., 2019, He et al., 2024]. However, many existing methods show substantial performance degradation as removing specific layers without
Correspondence to: Mete Erdogan <merdogan @ stanford.edus, Francesco Tonin <francesco.tomi eptl.ch>.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

L.M
Transformer Block
Transformer Block
MILP
MLP
LaverNorm
LayerNerm
NBL
Transformer
Compressed Block
Linear Layer W Gyfai A EPI WEIT
Atteation
Trausformer Block
Transformer
LaperNuem
Pootinnal Freading
Excadin
l'mbedding
Iepur Kinhwldlig
Figure 1: Illustration of Neural Block Linearization (NBL), which replaces multi-hean attention layer with an efficient linear layer using the closed-form LMMSE estimator.
properly replacing them lead to substantial accuracy drop, highlighting the need for more reliable strategies to identify and replace network blocks effectively.
In this paper, introduce Neural Block Linearization (NBL), novel method for compressing transformer-basee models by replacing self-attentions (costly network blocks) with efficient linear layers derived using Linear Minimum Mean Squared Error (LMMSE) estimators [Kay, 1993, Kailath et al., 2000]. without the need for computationally expensive gradient- baase training. While applicable to any network block, this paper mainly focuses substituting the attention layers rather than entire transformer blocks, as experimental results show that this selective approach effectively balances inference speed-up with accuracy.
Unlike methods that entirely remove specific network components, for example Ashkboos al. [2024], Xia et al. [2024a]. Song et al. [2024], He et al. [2024], NBL maintains consistency by substituting these layers with their linear approximations. NBL further integrates theoretical lin- carization eтror quantification using Canonical Correlation Analysis (CCA) [Hotelling, 1992]. By deriving an error bound and employing as a criterion for layer substitution, our approach provides a theoretically grounded framework to quantify redundancy, ensuring optimal compatibility with the linearization process. This enables significant inference speed-ups while maintaining performance, positioning NBL effective solution for optimizing large language model inference, enhancing both efficiency and scalability in applications.
Contributions. Our contributions include:
A principled and efficient substitution of attention layers with linear transformations using the closed form solution of the Linear Minimum Mean Squared Error (LMMSH)) estimators;
A theoretical error bound derived using Canonical Correlation Analysis (CCA), used to quantify approximation error and serve the layer substitution criterion;
Accelerating pre-trained transformer model inference, demonstrated through empirical results preserve competitive performance on reasoning tasks.
2 Motivation
The computational efficiency of transformer-basee models has become a critical concern as their use expands across NLP tasks. The attention mechanism, a cornerstone of transformer architecture, particularly resource-intensive due its quadratic complexity with respect sequence length. Some methods address this issue with subquadratic linear attention mechanisms [Katharopoulos al., 2020), e-g., [Mercat al., 2024, Zhang et al., 2025]. These methods fundamentally change the architecture, often requiring significant retraining and leading to potential trade-offs in performance. Moreover, methods such as SLEB [Song et al., 2024], DROP [He et al., 2024] rely on redundancy quantification metrics to remove specific transformer blocks attention layers, to accelerate inference. However, these approaches often lead to performance degradation due to the abrupt removal of critical components, which disrupt the model's ability to capture essential patterns and dependencies.

---

In this work, we propose fundamentally different approach. Rather than completely removing attention layers, we approximate those that exhibit redundancy by replacing them with linear approx- imations based on their input-output relationships. To achieve this, we use Canonical Correlation Analysis (CCA) [Hotelling, 1992] as a theoretical foundation to identify and quantify redundancy in a layer considering its inputs and outputs.
CCA is statistical method that measures the linear relationship between two random vectors and Y. It achieves this by finding pairs of canonical directions-linens projections of X and Y-that are maximally correlated. Mathematically, CCA solves the following optimization problem:
a Cyxb
max P LuTCxxaVbTCy 0.6
where Cxx and Cyy are the covariance matrices of X and Y, respectively, and Cyx is their cross-covariance matrix. To compute and the corresponding canonical directions and b, we first standardize X and Y by normalizing their variances. This leads to the construction of the standardized cross-correlation matrix:
Cy= --1 Cyx
The canonical correlations A are obtained singular values of Cw through Singular Value Decom- position (SVD):
Cw = UEVT,
where Σ is a diagonal matrix containing the singular values pi and U and V are the canonical direction matrices for X and Y, respectively. The strength of these correlations P: provides clear indication of how well the components of can be linearly predicted from X. If the canonical correlations are high, it suggests that significant portion of the oulput can be captured using a linear transformation, without incurring substantial loss in predictive fidelity.
We use CCA to guide the Linear Minimum Mean Squared Error (LMMSSE) [Kay, 1993, Kailath al., 2000] approximation of the attention layers, allowing us quantify the approximation error introduced during linearization. By identifying and targeting the layers where the redundancy is high (i.e., canonical correlations are close 1), our method compresses transformer models a principled manner while preserving their functionality with minimal compromise.
Most notably, approach replaces attention layers with linear approximations without any fine- tuning or gradient- optimization, either during or after substitution, making it highly efficient for compressing pre- pre-trained models. Despite this, NBL achieves strong accuracy, demonstrating that attention layers can be effectively approximated without retraining. This direct substitution significantly reduces computational overhead while preserving model functionality. Our method not only enables efficient transformer compression but also provides new insights into model component analysis and redundancy quantification, aligning with the growing demand for scalable NLP solutions.
3 Methods
We propose method that replaces selected blocks, particularly attention layers, with linear approxi- mations based on their estimation task, where the goal is to approximate the output of an attention layer using its input through the well-known Linear Minimum Mean Squared Error (LMMSE) formulation.
relationships. This approach frames the problem as a channel
The LMMSE framework offers a principled way to minimize the mean squared error (MSE) between the true output and the estimated output of an attention layer. Using the covariance structure of the inputs and outputs.
derive the optimal linear approximation of the attention layer's behavior. Our process ensures that the replaced layer closely approximates the behavior of the original while significantly reducing computational complexity.
The following sections detail the theoretical formulation, algorithmic implementation, and practical considerations of the proposed Neural Block Linearization (NBL) method. In Section 3.1, we introduce NBL primarily in the context of replacing self-attention layers. However, NBL is flexible framework that can be applied to any neural network block. In Section (3.2), we derive error bound for the linearization the NBL method, which serves as a valuable tool for redundancy characterization within the NBL framework.
3

---

CCA Iound
5 8
Mistral-71
UCA Itond
Liama 3.100
Index
Index
Mannd--f 75
Lana Lana 1-UU Eastre
(a)
(h)
Figure 2: Illustration of layer selection in the NBL method, guided by the CCA-based bound from the Theorem 3.2, as applied to (a) Asstral-7B and (b) L1 alama-3 1-8B models.
3.1 Neural Block Lincarization (NBL)
To determine the linear weights, utilize a calibration dataset D (5(0)) )== with input se- quences, each having a context length of t. Each sequence S = { [w] wh w (0) is processed through multiple transformer blocks before reaching the k-th block. The input representations to the k-th attention layer are obtained by applying the first (A: 1) transformer blocks sequentially:
x
fk-1 fk-20.... f(( E gdax1,
where each function (.) represents the mapping applied by the j-th transformer block, including its self-attention and MLP layers. The corresponding output representations are then extracted directly from the self-attention layer of the k-th transformer block:
(x. IR
where Ax(.) denotes the softmax self-attention transformation with input and output activation sizes as Iin and hoa, respectively, inside the k-th transformer block. Then, to construct the dataset for learning the linear mapping, stack the boken-wise representations across all sequences as:
X = [x(] (1) x(2)
x(()) ((xxt)
Y = [y(]]
(x)
E R ((sxt(
Here, x and y( are the matrices containing the representations of all tokens in sequence i, and stacking them along the batch dimension ensures that each column X and corresponds to a single token's s representation across all sequences. Then, X and Y can be treated as vector-valued random variables with s realizations. Then, our goal is to minimize the MSE between the actual attention output Y and the estimated output y, formulated as:
MSE(Y, V) = E ||Y - 813
where [[[-] represents the expectation, and || ||2 denotes the euclidean norm of a vector. The linear estimator takes the form:
y = W X + b.
(2)
Then, the optimal linear estimator minimizing the MSE in Expression (I) is derived presented in the following proposition.
Proposition 3.1. [Kay, 1993]. The Linear Minimum Mean Squared Error (LMMSE) estimator defines the optimal linear relationship between the vector- valued random variables X and Y by the following weight W and bias b:
W = =yxx
(3)
b E[Y WE[X].
(4)
where Cxx is the covariance of X and Cy is the cross-covariance belween Y and X.
This formulation demonstrates that the linearized block operates as a single linear layer with the weight W and bias b. Figure illustrates the process of using NBL to replace an attention layer of an L.M with its linear approximation. For the derivation of Proposition 3.1 and an in-depth discussion of MMSE estimators, refer to Appendix A.

---

3.2 Canonical Correlation Analysis (CCA) based error bound and redundancy analysis
key aspect the NBL is understanding and quantifying the error introduced by replacing a network block with its linear approximation. While the LMMSE estimator provides a principled approach to minimize the error between the true outputs and their linear estimates, is essential bound this emor to ensure acceptable approximation and preserve the model performance.
To this end, present upper bound on the normalized mean squared error (NMSE) of the linear approximation in the following theorem. Leveraging Canonical Correlation Analysis (CCA), this theorem directly links the approximation error the CCA singular values, which quantify the alignment between input and output spaces. These singular values provide valuable insights into the conditions under which linearization most effective.
Theorem 3.2. For a given self-attention layer with input X and output Y, the Normalized Mean MSE(Y,Y)
Squared Error (NMSE) of the LMMSE estimator Y is defined as: NMSE(Y, TKCrv)' where Cyy is the auto-covariance matrix of Y. Then, the NMSE satisfies the following upper-bound:
NMSE(Y, 8) < (hout r) Σ (1-P))
(5)
where min(haas hin), and Pi are the canonical correlations between X and Y.
This theorem highlights that the NMSE is
Algorithm 1 Neural Block Linearization (NBL)
bounded by summation term dependent on
1: Input: Attention layer inputs X and outputs Y from calibration dataset D. number of layers to linearize m. 2: Initialization: Load the pre-trained LLM,, extract all
1 1-pī. where represents the canonical cor-
relations between the input and output Y.
Intuitively, the larger the canonical correla-
attention layers A = {A1, A2, Ak}.
tions (p; close to 1), the better the alignment 3: for cach attention layer A E AA do
(linear relationship) between the input and out-
4: Collect the input-output pairs (Xe, Yx) from D cor- responding the attention layer As..
put spaces, and thus the lower the approxima-
tion error. The result also implies that lay- ers with weaker correlations (pt closer to 0) are less suitable for linearization, as the er- ror bound becomes larger. This insight en- ables the systematic identification of blocks where linearization will have less impact on overall model performance, providing quan- titative basis for deciding which attention lay- blocks to replace. The additive term corresponding to the undetermined mappings,
5: Compute the bound on the linearization NMSE (Theorem-322) for the attention layer Ar, using the canonical correlations Pki between Y, and X s.t. h hom hin:
NMSE, MSE(Y,, Vi). Tr(Crr)
A < >EO (1 -Pa)).
6: end for
7: Select m layers with the lowest NMSE bounds, form-
ing Anm {A AJz...,, Ajm }
A.A.
(how -)) is zero when howt < Nim condition typically satisfied in most LLM layers such as self-attention.
8: for cach layer A5 E Aim do
9: Replace A₁ with its linear approximation (linear layer) obtained by solving the LMMSE problem (Proposition 3.1), with weight and bias terms as
We propose using this bound to rank the lay- ers of a pre-trained transformer model by their suitability for linearization, as illustrated in Figures 2a and 2b. By targeting the layers with high canonical correlations, significant computational savings can be achieved with minimal approximation error. Although one could directly compute the NMSE, we adopt its CCA-based upper bound, which provides variance- aggossic, modewise indicator of linear approximability with more stable em- pirical behavior (see Appendix C.I).
W; = Cryx, Xxxxx,XXxxx CxJx, by = E[Y] ] WgE[X,]
10: end for 11: Output: Compressed LLM with m lincarized atten- tion layers.
Table I: Calibration runtime scaling with model size.
L Model
8B
70B
405B
Hidden size (d) Hidden layers
4092
8192
16384
32
80
126
Runtime layer Total
26.04
78.91 s
371,79
3.3 Algorithm
00.3 hr
~ 1.75 hr
~33.1 hr
Our methodology, detailed in Algorithm 1, compresses an LLM by replacing m of the computationally intensive attention layers with efficient linear estimators. Using a calibration dataset, extract
5

---

Table 2: Performance of Mistral-7B on reasoning benchmarks across baselines, PruneNet [Sengupta et al., 2025], SliceGPT [Ashkboos et al., 2024], SLEB [Song et al., 2024], DROP [He et al., 2024], and ours (NBL). Prefill and throughput speeds are reported relative to the baseline original model.
Method
Sparsity (56
ARC- (norm)
ARC-c (norm)
Bool)
HellaSwag (norm)
MMLU (5-shot)
OBQA
PIQA (nom))
Wino- Grande
Avg. (1)
Prefil (1)
Through-
(nomn)
put (f)
Baseline
0
79.4
54.2
63.8
81.1
65.3
43.8
82.1
74,2
70.2
PruneNer-15%
15
72.1
46.1
68.7
75.0
50.2
39.8
78.8
69.7
62.5
1.11
1.05
PruneNes-25%
64.3
39.9
62.7
67.1
31.6
35.0
73.1
66.0
1.09
PraneNer335%
35
26.3
27.3
44.5
26.8
24.0
25.0
50.5
49.5
34.2
1.32
1.13
Slice 15%
15
67.9
41.9
82.2
73.8
52.0
40.2
79.1
69.9
63.4
1.14
1.01
55.9
33.5
77.6
35.9
34.8
65.1
55.0
1.08
SliceGPT-35%
35
43.2
27.6
67.7
48.8
27.0
29.8
66.2
60.7
46.4
1.44
1.14
SL HH-4
12.5
70.5
43.2
77.4
72.5
38.9
41.6
79.1
66.8
61.2
1.14
1.16
SLEB-S
58.9
51.2
28.6
35.4
72.7
51.8
1.35
SLEB-12
37.5
43.1
30.1
49.4
50.4
24.4
30.8
66.4
53.0
43.5
1.59
1.50
Block DROP-4
12.5
70.8
47.2
80.4
75.4
61.4
40.2
77.9
71.0
65.5
1.14
1.16
Block DROP-8
52.9
37.4
71.6
59.8
31.2
67.6
56.2
1.33
1.34
Block DROP-12
37.5
31.7
31.6
42.8
27.0
24.2
30.0
54.1
57.1
37.3
1.59
1.50
Block NBL.44
11.5
72.0
47.1
82.1
74.4
61.5
40.4
78.2
73.0
66.1
1.14
1.14
Block NBL-8
23.1
58.8
38.8
82.2
60.6
58.6
36.2
71.4
69.0
59.4
1.32
1.31
Block NBI.-12
34.6
42.0
32.7
62.9
43.7
55.9
32.0
60.3
65.8
49.4
1.47
Attn DROP-4
2.4
79.9
53.5
83.4
80.8
62.4
44.6
$2.1
73.8
20.0
1.08
1.11
Anm DROP-8
4.8
79.6
52.3
82.6
80.2
62.0
44.2
72.9
69.4
1.18
1.22
DROP-12
76.3
76.7
77.7
41.8
72.9
1.29
Arm DROP-16
9.6
57.8
41.4
49.8
67.9
24.9
38.6
73.6
69.4
52.9
1.44
1.41
Anm NBL-4
1.5
80.2
53.9
83.5
80.6
62.4
44.0
$1.9
74.0
70.1
1.08
1.10
NBL-8
62.4
70.0
1.20
Atm NBL-12
4.3
77.3
50.5
82.6
76.9
62.3
43.0
80.9
73.2
68.3
1.28
1.27
Attn NBL.-1
5.8
62.4
42.8
76.7
69.4
33.0
39.8
76.4
70.2
58.8
1,40
1.37
inputs and outputs from each attention layer and compute the weight matrices (Wa) and bias terms (bç) via I.MSSE estimation. Layers are ranked by CCA-derived NMSE bounds, and the m layers with the lowest approximation error are selected for linearization. These layers are then replaced with their linear approximations, reducing complexity while preserving accuracy. The key computational steps of NBL have an overall complexity of Odd s- -s-t.dd). where d is the embedding dimension of the attention layers, 8 is the number of sequences, and is the context length of calibration samples. Since calibration is layer-wise, runtime scales linearly with the number of layers. We demonstrate that 8B-scale models can be compressed in under an hour, with scalability to larger models (see Table 1). While NBL effectively replaces attention layers, it can also linearize any network block, including transformer blocks, as also demonstrated next. Refer to Appendix B for CCA derivations, and Appendix D for NBL implementation and calibration details corresponding to Table 1.
4 Experiments
In this section, we assess the impact of Neural Block Linearization (NBL) on LLM inference performance and efficiency. Evaluation benchmarks include 5-shot performance on the MMLU task [Hendrycks et al., 2020] and O-shot performance [Clark et al., 2018], BoolQ [Clark et al., 2019], HellaSwag [Zellers et al., 2019], OBQA [Mihaylov et al., 2018], PIQA [Bisk al., 2020] and WinoGrande [Sakaguchi et al., 2021], following a similar evaluation as Zhang et al. [2025]. We implemented and evaluated NBL on an NVIDIA A100 GPU (80GB) using Py PyTorch [Paszke al., 2019] and HuggingFace Transformers [Wolf, 2019]. Evaluation is carried out with the default parameters from the Evaluation Harness framework [Gao et al., 2024].
ARC-casy (ARC-e), ARC-challenge (ARC-c)
We compare NBL. with several baseline methods, including SLEB [Song et al., 2024], SliceGPT [Ashkboos et al., 2024], PruneNet [Sengupta et al., 2025] and DROP [He et al., 2024], evaluating their performance on reasoning tasks and their improvements in latency and throughput. In the calibration of all methods, we used 256 samples from the C4 dataset [Raffel et al., 2020]. In the tables, "Attn NBL-m" denotes the NBL applied to attention layers, while "Attn DROP-m" removes m attention layers with the DROP method. Similarly, "SLEB-m" refers to dropping m transformer blocks based on the SLEB algorithm, and "SliceGPT-d%" or "PruneNet-d%% indicates pruning d% of the model parameters using SliceGPT or PruneNet methods. We note that PruneNet required tuning in our evaluations: the default 5e 4 learning rate led to poor convergence, while l -4 performed better.
Table shows that Attn NBL outperforms other methods on the 32-layer Mistral-78 model [Jiang et al., 2023]. Replacing 8 layers (Attn NBL-8) retains the original performance (70.0), exceeding Attn DROP-8 (69.4), SliceGPT-25% (55.0), Pru5eNet-25% (55.0), and SLEB-8 (51.8). At higher compression (c.g., Attn NBL-16), the gap widens, where NBL scores 58.8 compared to 52.9

---

Table 3: Performance of 1lama-3.1-8B various methods and reasoning benchmarks.
Method
Spursity (%)
ARC-e (norm)
ARC-c (norm)
BoolQ
HellaSwag (norm)
MMLU (S-shet)
OBOA (norm)
PIOA (noеm)
Wino- Grande
Arg. (f)
Prdill
Through
put
Baseline
81.2
33.6
81.9
78.9
65.1
45.0
81.4
74.2
70.2
PraneNet-15%
15.0
66.8
42.06
61.8
66.2
42.9
38.2
74,9
67.4
58.9
1.08
1.05
PruneNec-2544
55.7
25.5
1.09
PruneNes-55 35%
35.0
37.5
25.9
60.0
37.7
25.2
28.0
60.6
54.8
41.2
1.32
1.13
SliceGP--1SS
15.0
66.6
40.8
77.9
68.3
40.8
36.4
75.2
64.7
58.8
1.13
1.05
SP
1.09
SliceGPT-35%
33.0
43.1
26.1
61.8
44.9
26.3
26.2
65.8
52.7
43.4
1.41
1.12
SLEB-4
12.5
73.7
44.3
67.3
71.0
34.6
40.2
77.8
68.9
59.7
1.14
1.16
XI.B--
25.0
59.4
32.7
38.7
54.2
1.32
SLEB-12
37.5
47.1
28.2
46.8
46.0
21.6
27.8
67.3
52.5
42.5
1.54
1.56
Block DROP-4
12.5
71.4
47.6
70.6
73.7
61.7
41.4
77.5
70.2
64.3
1.13
1.16
Bloek DROP-8
1.32
Block DROP-12
37.5
39.5
35.1
29.8
Block NBL-4
11.5
77.1
19.0
81.4
73.3
6/1.3
41.4
78.6
72.5
67.2
1.12
1.15
Block NBL.8
23.1
66.2
41.9
62.5
62.6
62.6
35.6
73.1
70.2
59.3
1.29
1.27
48.7
1.50
Attn DROP-
2.4
81.4
54.2
81.9
78.4
65.0
45.4
80.6
74.4
70.2
1.08
1.12
Atin DROP-8
1.8
80.9
53.2
81.1
78.0
65.0
45.2
$1.1
73.5
69.8
1.18
1.19
DROP-12
76.3
51.3
79.7
1.29
Atm DROP-16
9.6
55.2
39.1
80.6
63.4
27.2
36.6
69.8
69.6
55.2
1.43
1.42
Arn NBI
1.5
81.9
$4.0
82.2
78.1
65.0
45.8
81.1
73.4
70.2
1,08
1.11
NBL-S
82.1
Atm NBL--12
4.3
79.1
52.2
82.3
75.2
64.8
45.2
79.9
74.0
69.1
1.24
1.25
Attn NBL-16
5.8
71.8
46.8
81.6
69.0
39.1
41.8
77.0
73.1
62.5
1.39
1.37
Table 4: Performance of -aapsSeek-RR--Diii----LL-a-
across baselines and reasoning benchmarks. Prefill and throughput speeds are reported relative to the baseline original model.
Sparsity (%)
ARC-e (поrm)
ARC- (norm)
BoolQ
HellaSwag (norm)
OBQA (norm)
PIQA (corm)
Wino- Grande
Avg. (?)
Prefll (?)
Through-
(5-sbot)
put (1)
Buseline
66.1
42.3
82.9
74.3
55.7
41.8
77.8
67.7
63.6
PruneNet-155
15%
15.0
55.1
35.3
66.7
63.7
40.6
34,8
69.9
63.5
53.7
1.08
1.08
25
25.0
44.3
30.5
62.4
52.6
27.7
30.2
65.6
58.2
46.4
1.16
1.15
PrumeNet-35%
38.0
32.3
24.6
62.0
37.6
25.5
25.2
$8.4
53.1
39.9
1.22
1.19
SliceGPT-15%
15.0
51.3
34.1
73.6
62.1
31.0
31.4
69
56.4
51.2
1.09
1.09
SliceGPT-254 25%
25.0
45,3
29.7
68.8
$1.4
25.6
29.2
65.5
55.0
46.3
1.15
1.13
Sl1ceGPT:355%
35.0
37.2
25.5
60.9
$1.8
40.6
1.21
1:20
SLEB-4
12.5
37.7
60.8
74,9
67.3
36,9
36,0
74.4
61.2
56.2
1.09
1.14
SLEB-8
25.0
33.6
52.4
58.6
55.5
25.9
30,8
68.9
55.7
47.7
1.17
1.29
SLEB-12
1.53
Attn DROP-S
4.8
66.9
43.5
83.6
73.0
55.8
40.6
76.0
68.3
63.5
1.16
1.19
Arnn DROP-12
7.2
60.7
41.6
82.7
68.1
55.6
37.8
73.9
66.0
60.7
1:27
1.34
DROP-16
Atin NBL-8
2.9
65.5
41.7
83.8
72.3
55.2
43.8
76.3
66.8
63.2
1.15
1.19
Artn NBL-12
4.3
63.6
41.2
84.3
70.0
55.3
42.0
74.9
67.7
62.4
1.24
1.32
73.3
(DROP-16), 46.4 (SliceGPT-35%). 34.2 (PruneNet-35%) and 43.5 (SLEB-16). NBL also improves inference speed (up to 1.27x) while keeping performance loss under 2%. Similar trends appear for Llama-3. 1-BB [Dubey et al., 2024] (Table 3) and DeepSeek- -R--Dsttiiiii 1atilll--amma-88 [Guo et al., 2025] (Table 4), where NBL outperforms other methods, particularly with 12-16 layer replacements, confirming its advantages improving accuracy-efficiency trade-off across models.
We report sparsity comparisons in Tables 2 to alongside efficiency results. Sparsity is defined the fraction of dropped replaced blocks, excluding the always-preserved embedding and projection layers (1.81,, 6.54%, and 2.98% of parameters in Mistral-7B, Llama-3 -. 1-8B, and Llama-3, 3. 70B respectively) While most parameters lie MLPs 80%). attention layers dominate FLOPs and runtime, making them the most effective targets. In this view, Block NBL achieves the best accuracy-spanse trade-off: for instance; on Llama-3 3. 1--8B Block NBL-12 attains 50.4 50.4 accuracy at 34.6% sparsity, surpassing Block DROP-8 (40.2), SLEB-8 (46.5), and MliceGPT-2594 (49.9) each with % sparsity. Atn NBL, while inducing less sparsity, yields the strongest sccuracy-throughee gains, improving throughput by 32% on DeepSeek, with under 1% accuracy loss.
4.1 Latency and throughput analysis
In addition to reasoning performance, we evaluate the inference speed-up L.LMs in two stages: prompt processing (prefill) and token generation (throughput). Prompt processing builds the
7

---

Table 5: Larger 70B models. Performance of Llama-3.1-708 across various methods and reasoning benchmarks. Prefill and throughput speeds are reported relative to the baseline quantized model.
Method
Sparsity (%)
ARC- (norm)
ARC-c (norm)
BoolQ
HellaSwag (norT)
MMLU (5-sbot)
OBQA (corm)
PIQA (norm))
Wine- Grande
Avg. (?)
Prefil
Through-
<T)
put (
Baseline (quant.)
63.7
84.5
84.6
77.7
47.8
84.1
79.2
76.0
SLEB-16
20.0
78.7
55.0
78.3
77.5
71.2
41.6
81.1
77.1
70.4
1.23
1.24
SLEB-32
40.0
60.7
37.2
64.6
63.6
44.6
35.8
73.2
67.3
55.9
1.65
1.61
Arm DROP-32
7.1
85.5
64.1
84.3
83.7
77.3
48.0
$4.3
79.0
75.7
1.20
1.30
DROP-48
10.1
77.2
56.2
76.3
70.3
45.6
79.4
75.1
70.4
1,42
Arm DROP-Se
11.9
43.4
31.4
65.5
52.2
36.6
29.4
67.0
61.0
48.3
1.34
1.53
Attn NBL-32
3,9
86.5
66.4
85,0
83.2
77.6
16.8
83.7
78.8
76.0
1,16
1.21
NBL-4S
76.2
$1.9
71.3
1.35
Atm NBL-54
6.5
63.0
48.2
82.3
71.6
65.3
43.4
78.3
71.7
63.4
1.30
1.44
key-value (KV) cache for 2048-token input, while token generation autoregressively produces 2048 tokens with a batch size of 1, following He et al. [2024]. We measure the prefill speed, calculated as the number of context tokens divided by the time generate the first token. Whereas, the throughput is defined as the median of the number of tokens generated per second, across time intervals. Results are presented in Tables 2 and for Llama-3. 1-8B and Mistral-78 models respectively. which show the speed-ups achieved by methods such as SliceGPT, PruneNet, SLEB, DROP, and NBL, normalized against the baseline model. Our findings show that NBL significantly improves computational efficiency, delivering consistent prefill and throughput gains across various compression levels with significantly less performance degradation than the layer-dropping methods.
4.2 Inference complexity and KV-cache usage
Bascline (Norm.)
Attm NIL-4
Let K be the total number of attention layers in an LLM, with m layers modified by NBL. The prefill speed for full attention is limited by complexity O(K-ndd), where n is the context length and the model dimension. Replacing m attention layers reduces their complexity to Omm nd), yielding an overall complexity of O(K m) n2d + m. nd). Thus, prefill speed improves as m increases, particularly for large n, where quadratic terms dominate. Correspondingly, Figure 3 shows that applying NBL more layers results in even higher prefill speed-ups at longer contexts. Furthermore, applying NBL reduces the usage of KV cache by a factor of K-m as only the K - m layers require the storage of key and value tensors. For details on the complexity, refer Appendix H.
Prefill Spood-up
15
14
25000
Context Length
Figure 3: Prefill speed-up of Llama-3.1 1--8B with varying context lengths. NBL. values are normal- ized by the baseline.
4.3 Application of NBL to larger models
We extend the NBL to a larger model, namely Llama-3, 1-70B [Dubey et al., 2024]. To fit the model in an NVIDIA A100 (80GB), we apply 4-bit post-training quantization using Activation-awarc Weight Quantization (AWQ) [Lin et al., 2024] with default settings and 128 Pile samples [Gao et al., 2020] for calibration. Table 5 presents comparison between Attn NBL and Attn DROP, with the 4-bit quantized Llama-3. 1-B0B model as baseline. For this model, PruneNet and SliceGPT were excluded from evaluation, as their public repositories currently lack support for post-quantization pruning. Both Attn NBL-32 and Attn NBL-48 outperform SLEB-16 in accuracy while matching or exceeding its inference throughput, with significantly lower KV-cache usage due to greater number of pruned attention layers. NBL demonstrates a notable advantage by preserving model accuracy greater extent while achieving significant improvements in inference speed relative baseline. For more details on quantization and implementation, see Appendix F.4.
4.4 Ablation studies
Replacing the Attention layer vs Transformer block: Tables 2 and present the results for substiquting m transformer blocks with NBL (Block NBL-m) and dropping m blocks using the DROP method (Block DROP-m). Although modifying transformer blocks achieves significant inference speed-ups, it results in much greater performance degradation compared to adjustments
8

---

made exclusively to the attention layer. Importantly, Block NBL demonstrates greater control and better accuracy preservation compared to Block DROP, as well as SliceGPT and SLEB.
Dependency on the calibration dataset: We examine the impact of the calibration dataset 'D. During calibration, 256 samples drawn from either the C4 or WikiText-2 training sets. The pruned models are then evaluated on both C4 and WikiText-2 validation datasets, with results provided in Tables 13 and 14, in Appendix F.I. Our comparison includes several methods, including SLEB, Attn DROP, Attn NBL (each with 8 layers dropped or linearized), and SliceGPT (with 25% sparsity). NBL performs favorably compared to other methods, where SliceGPT exhibits the greatest sensitivity.
LoRA fine-tuning of NBL-linearized layers:
We also explore the effect of applying LoRA [Hu et al., 2022] to the linear layers produced by NBL. Experiments show that LoRA fine-tuning provides only marginal accuracy improvements over NBL alone, reinforcing the strength and generality of NBL without requiring additional tuning. Full results and setup details are provided in Appendix F.2.
CCA-bound vs cosine distance criterion: We further examine the impact of the layer selection criterion on NBL, comparing the CCA-bound based criterion with the cosine distance criterion originally used in the DROP method. Our findings suggest, the CCA-bound criterion provides more reliable results, likely due to its better suitability for assessing linearization error (see Appendix F.4).
Related work
The computational and memory demands of transformer-based LLMs have inspired numerous techniques to reduce the complexity of inference while maintaining the performance:
Structured and unstructured pruning: Techniques such as weight pruning. structured pruning, and layer removal have been widely explored to reduce model size and computational requirements. Unstructured pruning focuses on sparsifying model weights [LeCun et al., 1989, Hassibi et al., 1993, Han et al., 2015, Kusupati et al., 2020, Hoefler et al., 2021]. but managing sparse matrices modern hardware introduces significant challenges [Wang et al., 2024c]. Structured pruning methods [ [Voita et al., 2019, Ma et al., 2023, Xia et al., 2024a, Muralidharan et al., 2024, Men et al., 2024, Gromov et al., 2024, Song et al., 2024, He et al., 2024, Sengupta et al., 2025] target specific layers, blocks, or dimensions for removal, offering improved efficiency. However, these methods often rely on empirical insights and lack a robust thcoretical framework for quantifying approximation errors.
Methods leveraging low-rank decompositions [Hsu et al., 2022, Kaushal et al., 2023, Yuan et al., 2023, Wang et al., 2024a], quantization techniques [Li et al., 2024, Zhang et al., 2024a, Tseng et al., 2024], and their combinations [Yao et al., 2023, Zhang et al., 2024b, Saha et al., 2024] have been explored to simplify model inference. While NBL introduces architectural modifications by linearizing specific layers, it remains straightforward to implement on pre-trained models. It eliminates the need for data-type type changes or additional hardware compatibility, streamlining its application. Also, in Section 4.3, NBL was succesfully applied to Llama-3. 1-70B, quantized with the AWQ [Lin et al., 2024]. This highlights the potential for deeper algorithmic integration between NBL and quantization techniques for greater efficiency.
Low-rank approximations and quantization:
Linear Attention mechanisms: Linear attention mechanisms aimed at achieving sub-quadratic complexity have been proposed (Katharopoulos et al., 2020, Choromanski et al., 2020, Peng et al., 2021, Shen al., 2021, Mercat et al., 2024, Zhang et al., 2025]. While these methods reduce inference complexity, they often demand extensive architectural modifications, retraining, or distillation, leading to trade-offs between performance and flexibility. In contrast, NBL utilizes the closed-form LMMSE solution to compute linear weights, avoiding these challenges.
Speculative decoding: Speculative decoding meth- ods [Leviathan et al., 2023, Cai et al., 2024, Xia et al., 2024b] accelerate autoregressive generation via draft- and-verify strategies. These approaches are orthogonal to model compression techniques like NBL, as they target decoding strategy rather than model structure. We show that NBL can be combined with speculative decoding (e.g., the recent EAGLE-3 [Li et al., 2025]) to achieve compounding speed-ups, up to 4.07x (Table 6) (Please refer Appendix E.5 for more details).
Table 6: Speculative decoding + NBL speed-ups on PrepSSek-R1-Distill---e 8B (MT-bench [Zheng et al., 2023], A100).
Model Configuration
Speedup
EAGLE-3 Alone
3.23 x
Attn NBL-4 + EAGLE-3
3.51 x
Attn NBL-8 + EAGLE-3
3.89 x
Atn NBI- EAGLE-3
4.07 x

---

Canonical Correlation Analysis (CCA) as an analysis tool: CCA emerged as tool to analyze and interpret neural networks. For instance. SVCCA was introduced to compare representations across layers and models [Raghu et al., 2017], and CCA was applied to investigate robustness and generalization [Morcos al., 2018). Additionally, a similarity index for comparing representations between different architectures was proposed [Kornblith et al., 2019], and representation alignment in wide networks was studied [Nguyen et al., 2020J. These works highlight the effectiveness of CCA in examining neural network structure and functionality. Similarly, we use CCA to analyze neural activations and assess the suitability of attention layers for linear approximations.
Further discussions, possible limitations, and future work
Empirical results show the effectiveness of our approach across several challenging reasoning benchmarks (e.g.. MMLU, HellaSwag, ARC). Nonetheless, it is important acknowledge certain limitations. In general, certain nonlinear transformations in large language models may not admit accurate linear approximations. This is an inherent property of deep architectures. NBL addresses this by ranking layers based on linear predictability via CCA, and substituting only those with low approximation error. The method is explicitly designed for such selective substitution; applying it indiscriminately to all attention layers would overly simplify the model, leading to shallow and history-agnostic behavior, which our approach does not advocate.
The choice of calibration dataset can affect performance on new inputs, where a form of calibration history averaging could lead to hallucination-like effects. These effects could be possible especially in domain shift or insufficient calibration coverage cases. Our ablations show that NBL demonstrates superior robustness than SliceGPT and that NBL's performance is remarkably stable even in the case of two very different calibration distributions (C4 and WikiTeex-22. Despite being calibrated on C4, the model was tested on reasoning tasks that likely exhibit substantially different data distributions. yet it achieved consistently strong performance (Tables 2-5).
We believe NBL offers significant promise for future research on optimizing LLM inference. Although NBL shows only mild sensitivity to calibration data, this can be further reduced using calibration data mixture techniques or adaptive strategies. A promising direction is to combine theoretical analysis of QKV structure with data-driven calibration to improve generalization and reduce reliance task-specific data. More principled selection of layers could also be achieved by analysing operator norms of attention layers, or examining the internal behavior of attention matrices. Investigating the relative norms of attention outputs and residual connections may further complement the CCA based criterion. Finally, hybrid designs that retain critical attention heads within a layer present promising path for balancing inference efficiency with contextual reasoning capacity.
7 Conclusion
This work proposes Neural Block Linearization (NBL), a novel framework for the compression of transformer-basee LLMs by selectively replacing attention layers other -nntenssvie network blocks with linear approximations derived using Linear Minimum Mean Squared Error estimators. NBL effectively integrates theoretical rigor into practical implementation, resorting to Canonical Correlation Analysis to bound approximation errors and guide layer substitution. Experimental results demonstrate that NBL achieves significant computational efficiency while maintaining competitive performance across various reasoning benchmarks, outperforming other recent pruning techniques in most cases, without resorting additional retraining, and adapts smoothly to pre-trained models. These directions position NBL as a scalable and adaptable solution for compressing and deploying LLMS in resource-constrained environments.
Acknowledgements
This work was supported by Hasler Foundation Program: Hasler Responsible Al (project number 21043). This research was sponsored by the Army Research Office and was accomplished under Grant Number 19 Foundation (SNSF) under grant number 200021 205011.
This work was supported by the Swiss National Science
10

---

References
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
Albert Q Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, et al. Mistral 7b. arXiv preprint arXiv:2310.0682. 2023.
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint arXivv2303.08887 230300877, 2023.
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, et al. The Ilama 3 herd of models. arXiv preprint arXiv:240.221783, 2024.
Aditya Kusupati, Vivek Ramanujan, Raghav Somani, Mitchell Wortsman, Prateek Jain, Sham Kakade, and Ali Farhadi. Soft threshold weight reparameterization for learnable sparsity. In International Conference on Machine Learning, pages 3544-5555 PMLR, 2020.
Torsten Hoefler, Dan Alistarh, Tal Ben-Nun, Nikoli Dryden, and Alexandra Peste. Sparsity in deep learning: Pruning and growth for efficient inference and training in neural networks. Journal of Machine Learning Research, 22(241):1-124 2021.
Zhihang Yuan, Yuzhang Shang, Yue Song, Qiang Wu, Yan Yan, and Guangyu Sun. Asvd: Activation- aware singular value decomposition for compressing large language models. arXiv preprint arXiv:2312.058211 2023.
Xin Wang, Yu Zheng, Zhongwei Wan, and Mi Zhang. Svd-Ilm: Truncation- singular value decomposition for large language model compression. arXiv preprint arXiv:2403.07378, 2024a.
Ji Lin, Jiaming Tang, Haotian Tang, Shang Yang, Wei-Ming Chen, Wei-Chen Wang, Guangxuan Xiao, Xingyu Dang, Chuang Gan, and Song Han. Awq: Activation-awar weight quantization for on-device Ilm compression and acceleration. Proceedings of Machine Learning and Systems, 87-100, 2024.
Rajarshi Saha, Naomi Sagan, Varun Srivastava, Andrea J Goldsmith, and Mert Pilanci. Compress- ing large language models using low rank and low precision decomposition. arXiv preprint arXiv:2405. 18886, 2024.
Yaniv Leviathan, Matan Kalman, and Yossi Matias. Fast inference from transformers via speculative decoding. In International Conference on Machine Learning, pages 19274-19288 PMLR, 2023.
Tianle Cai, Yuhong Li, Zhengyang Geng, Hongwu Peng, Jason D Lee, Deming Chen, and Tri Dao. Medusa: Simple I1m inference acceleration framework with multiple decoding heads. In International Conference Machine Learning, pages 5209-5235. PMLR, 2024.
Xiaoqi Jiao, Yichun Yin, Lifeng Shang, Xin Jiang, Xiao Chen, Linlin Li, Fang Wang, and Qun Liu. TinyBERT: Distilling BERT for natural language understanding. In Trevor Cohn, Yulan He, and Yang Liu, editors, Finulings of the Association for Computational Linguistics: EMNLP 2020, pages 4173-4174, Association for Computational Linguistics, 2020. doi: 10. 1/22000 findings-emmlp.372. URL attps://achi haca:///ol/attooogg..oooooo 2rg/2020..o. findingss-nnnn eetpsseemll.3333 372.
Jiaheng Liu, Chenchen Zhang, Jinyang Guo, Yuanxing Zhang, Haoran Que, Ken Deng, ZhiqiBai, Jie Liu, Ge Zhang, Jiakai Wang, Yanan Wu, Congnan Liu, Jiamang Wang, Lin Qu, Wenbo Su, and Bo Zheng. DDK: Distilling domain knowledge for efficient large language models. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024. URL https:/ htaps:://pppereee http///pppreeeeeeeeeee/ee nnt/ httpe//oommniddiiiiii

---

Junxiong Wang, Daniele Paliotta, Avner May, Alexander M Rush, and Tri Dao. The mamba in the aama:: Distilling and accelerating hybrid models. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024b. URL httpthttt///peeeeeee jAZHODjALU.
www/eeetuurrne id=
Michael Zhang, Simran Arora, Rahul Chalamala, Benjamin Frederick Spector, Alan Wu, Krithik Ramesh, Aaryan Singhal, and Christopher Re. Lolcats: On low-rank linearizing of large language models. In The Thirteenth International Conference on Learning Representations, 2025. URL
https://openrevieww http://parrrvvveeee..eerrrrrraaa meee//forummidd8ttteery...
Elena Voita, David Talbot, Fedor Moiseev, Rico Sennrich, and Ivan Titov. Analyzing multi-head self-attention: Specialized heads do the heavy lifting, the rest can be pruned. arXiv preprint arXiv:1905.00418. 2019.
Xinyin Ma, Gongfan Fang, and Xinchao Wang. LIm- Aim-prune:: On the structural pruning of large language models. Advances in neural information processing systems, 36:21702-21722 2023.
Mengzhou Xia, Tianyu Gao, Zhiyuan Zeng, and Danqi Chen. Sheared LLaMA: Accelerating language model pre-training via structured pruning. In The Twelfth International Conference on Learning Representations, 2024a.
Saurav Muralidharan, Sharath Turuvekere Sreenivas, Raviraj Joshi, Marein Chochowski, Mostofa Patwary, Mohammad Shoeybi, Bryan Catanzaro, Jan Kautz, and Pavlo Molchanov. Compact language models via pruning and knowledge distillation. arXiv preprint (iti::2007.166799 2024.
Xin Men, Mingyu Xu, Qingyu Zhang, Bingning Wang, Hongyu Lin, Yaojie Lu, Xianpei Han, and Weipeng Chen. Shortgpt: Layers in large language models are more redundant than you expect. arXiv preprint vrXiv:2403.03853 2024.
Salch Ashkboos, Maximilian Croci, Marcelo Gennari do Nascimento, Torsten Hoefler, and James Hensman. Slicegpt: Compress large language models by deleting rows and columns. arXiv preprint arXiv:2401.1500 2024.
Jiwon Song, Kyungseok Oh, Taesu Kim, Hyungjun Kim, Yulhwa Kim, et al. Sleb: Streamlining 1lms through redundancy verification and elimination of transformer blocks. In Forty-first International Conference on Machine Learning, 2024.
Shwai He, Guoheng Sun, Zheyu Shen, and Ang Li. What matters in transformers? not all attention is needed, 2024. URL /ttps:///aaaaie / http:/aawvvv..oraaaaa/aa/a
15786.
Steven M Kay, Fundamentals of statistical signal processing: Estimation theory, 1993.
Thomas Kailath, Ali H Sayed, and Babak Hassibi. Linear estimation. Prentice Hall, 2000.
Harold Hotelling. Relations between two sets of variates. In Breakthroughs in statistics: methodology and distribution, pages 162-190. Springer, 1992.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are mnns: Fast autoregressive transformers with linear attention. In International conference on machine learning, pages 5156-5166 PMLR, 2020.
Jean Mercat, Igor Vasiljevic, Sedrick Keh, Kushal Arora, Achal Dave, Adrien Gaidon, and Thomas Kollar. Linearizing large language models. arXiv preprint arXiv:2405.008 06640, 2024.
Ayan Sengupta, Siddhant Chaudhary, and Tanmoy Chakraborty. You only prune once: Designing calibration-free model compression with policy learning. arXiv preprint arXiv:2501.15296 2025.
Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazcika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding. arXiv preprint arXiv:2009.033000 2020.
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprint arXiv: 1003::8880000455 2018.
12

---

Christopher Clark, Kenton Lee, Ming- Toutanova. Boolq: Exploring the surprising difficulty of natural yes/no questions. arXiv preprint arXiv: 10044, 2019.
Ming-Wei Wei Chang, Tom Kwiatkowski, Michael Collins, and Kristina
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can machine really finish your sentence? arXiv preprint arXiv: vrxxiv:/05.000 190.007800, 2019.
Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. Can suit of armor conduct electricity? a new dataset for open book question answering. arXiv preprint arXiv: 180900289, 2018.
Yonatan Bisk, Rowan Zellers, Jianfeng Gao, Yejin Choi, al. Piqa: Reasoning about physical commonsense in natural language. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pages 7432-7439, 2020.
Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adversarial winograd schema challenge at scale. Communications of the ACM, 64(9))99-106,
2021.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32, 2019.
T Wolf. Huggingface's 's transformers: State-of-the--e natural language processing. arXiv preprint arXiv:1910.03771, 2019.
Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding. Jeffrey Hsu, Alain Le Noac' h. Haonan Li, Kyle McDonell, Niklas Muennighoff, Chris Ociepa, Jason Phang, Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang Sutawika, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. framework for few-shot language model evaluation, 2024. URL htttss://zenod:: www.ooogg/oooo.ooeooooooooooooeo seccords/12608600
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of machine learning research, 21(140): -6,, 2020.
Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang, Xiao Bi, et al. Deepseek-rl: Incentivizing reasoning capability in llms via reinforcement learning. arXiv preprint arX1v:2501.12948, 2025.
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, et al. The pile: An 800gb dataset of diverse text for language modeling. arXiv preprint arXiv:2101.0000 2020.
Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang. Lu Wang, Weizhu Chen, et al. Lora: Low-rank adaptation of large language models. ICLR, 1(2):3, 2022.
Yann LeCun, John Denker, and Sara Solla. Optimal brain damage. Advances in neural information processing systems, 2, 1989.
Babak Hassibi, David G Stork, and Gregory J Wolff. Optimal brain surgeon and general network pruning. In IEEE international conference neural networks, pages 293-299. IEEE, 1993.
Song Han, Jeff Pool, John Tran, and William Dally. Learning both weights and connections for efficient neural network. Advances in neural information processing systems, 28, 2015.
Wenxiao Wang, Wei Chen, Yicong Luo, Yongliu Long, Zhengkai Lin, Liye Zhang, Binbin Lin, Deng Cai, and Xiaofei He. Model compression and efficient inference for large language models: A survey. arXiv preprint 4rX 2024c.
Andrey Gromov, Kushal Tirumala, Hassan Shapourian, Paolo Glorioso, and Daniel Roberts. The unreasonable ineffectiveness of the deeper layers. arXiv preprint arXiv:2403.17888 2024.
13

---

Yen-Chang Hsu, Ting Hua, Sungen Chang, Qian Lou, Yilin Shen, and Hongxia Jin. Language model compression with weighted low-rank factorization. arXiv preprint arXiv:2207.00112, 2022.
Ayush Kaushal, Tejas Vaidhya, and Irina Rish. Lord: Low rank decomposition of monolingual code Ilms for one-shot compression. arXiv preprint arXiv:2309.14022 2023.
Yixiao Li, Yifan Yu, Chen Liang, Nikos Karamparziakis, Pengcheng He, Weizhu Chen, and Tuo Zhao. Loftq: Lora-tine-tuning toneeettninn--wve aware quantization for large language models. In The Twelfth International Conference on Learning Representations, 2024.
Yang Zhang, Yawei Li, Xinpeng Wang, Qianli Shen, Barbara Plank, Bernd Bischl, Mina Rezaei, and Kenji Kawaguchi. Finercut: Finer-grained interpretable layer pruning for large language models. arXiv preprint arXiv:2405. 182 18, 2024a.
Albert Tseng, Jerry Chee, Qingyao Sun, Volodymyr Kuleshov, and Christopher De Sa. QuIPSW#S: Even better LLM quantization with hadamard incoherence and lattice codebooks. In Forty-first International Conference on Machine Learning, 2024.
Zhewei Yao, Xiaoxia Wu, Cheng Li, Stephen Youn, and Yuxiong He. Zeroquant-v2: Exploring post-training quantization in Iims from comprehensive study to low rank compensation. arXiv preprint 2.0.333030008802,, 2023.
Cheng Zhang, Jianyi Cheng, George A Constantinides, and Yiren Zhao. Lqer: Low-rank quantization error reconstruction for Iims. arXiv preprint arXiv:2402.02446, 2024b.
Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. arXiv preprint (xXi::2009.14794 2020.
Hao Peng, Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah A Smith, and Lingpeng Kong, Random feature attention. arXiv preprint arXiv:2103.02144, 2021.
Zhuoran Shen, Mingyuan Zhang, Haiyu Zhao, Shuai Yi, and Hongsheng Li. Efficient attention: Atten- tion with linear complexities. In Proceedings of the IEEE/CVF winter conference on applications of computer vision, pages 3531-3539, 2021.
Lianmin Zheng, Wei-Lin Chiang. Ying Sheng, Siyuan Zhuang, Zhanghao Wu, Yonghao Zhuang. Zi Lin, Zhuohan Li, Dacheng Li. Eric Xing, et al. Judging ilm-as-a-jude juddge with mt-bench and chatbot arena. Advances in Neural Information Processing Systems, 36:46595546623, 2023.
Heming Xia, Zhe Yang, Qingxiu Dong, Peiyi Wang, Yongqi Li, Tao Ge. Tianyu Liu, Wenjie Li, and Zhifang Sui. Unlocking efficiency in large language model inference: A comprehensive survey of speculative decoding. arXiv preprint arXiv:2401.07851, 2024b.
Yuhui Li, Fangyun Wei, Chao Zhang, and Hongyang Zhang. Eagle-3: Scaling up inference ac- celeration of large language models via training-time test. arXiv preprint arXiv:2503.014 01840, 2025.
Maithra Raghu, Justin Gilmer, Jason Yosinski, and Jascha Sohl-Dickstein. Sveca: Singular vector canonical correlation analysis for deep learning dynamics and interpretability. Advances in neural information processing systems, 30, 2017.
Ari Morcos, Maithra Raghu, and Samy Bengio. Insights representational similarity in neural networks with canonical correlation. Advances in neural information processing systems, 31, 2018.
Simon Kornblith, Mohammad Norouzi, Honglak Lee, and Geoffrey Hinton. Similarity of neural network representations revisited. In International conference on machine learning, pages 3519- 3529. PMLR, 2019.
Thao Nguyen, Maithra Raghu, and Simon Kornblith. Do wide and deep networks learn the same things? uncovering how neural network representations vary with width and depth. arXiv preprint arXiv:2010. 15327, 2020.
David Hardoon, Sandor Szedmak, and John Shawe-Taylon Canonical correlation analysis: An overview with application to learning methods. Neural computation, 16(12):2639-266 2004.
14

---

Pauli Virtanen, Ralf Gommers, Travis E. Oliphant, Matt Haberland. Tyler Reddy, David Couraaeea,, Evgeni Burovski, Pearu Peterson, Warren Weckesser, Jonathan Bright, Stéfan J. van der Walt, Matthew Brett, Joshua Wilson, K. Jarrod Millman, Nikolay Mayorov, Andrew R. J. Nelson, Eric Jones, Robert Kern, Eric Larson, C J Carey, Ilhan Polat, Yu Feng. Eric W. Moore, Jake VanderPlas, Denis Laxalde, Josef Perktold, Robert Cimrman, Ian Henriksen, E. A. Quintero, Charles R. Harris, Anne M. Archibald, Antônio H. Ribeiro, Fabian Pedregosa, Paul van Mulbregt, and SciPy Contributors, SciPy 1.0: Fundamental Algorithms for Scientific Computing in Python. Nature Methods, 17:261-272, 2020, doi: 12.1038/441522-09--00
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, al. Training verifiers to solve math word problems. arXiv preprint arXiv:21 10. 14168, 2021.
Daria Soboleva, Faisal Al-Khateeb, Robert Myers. Jacob R Steeves, Joel Hestness, and Nolan Dey. SlimPajama: A 627B token cleaned and deduplicated version RedPajama, 2023. URL https://huugiingface co/ datasetss tetaaeetseeeeeeeeeeeeeeeeeeeeeee StweerrasSSimPaaammm 627B.
Cheng-Ping Hsieh, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, and Boris Ginsburg. RULER: What's s the real context size of your long-context language models? In First Conference on Language Modeling, 2024. URL https: / /twp////ep//eeee////////e/e/eeee www.eeettuunneee id= kIoBbc76Sy.
15

---

Appendix
Table of contents
Overview of Linear Minimum Mean Squared Error (LMMSE) estimation
17
A.1 General MMSE estimation
17
A.2 Linear MMSE (LMMSE) estimation
17
A.3 Derivation of the LMMSE estimator (Proposition 3.1)
18
Overview of Canonical Correlation Analysis (CCA)
20
B.1 Problem formulation
20
B.2 Canonical Correlation matrix
20
B.3 Performing CCA using the Singular Value Decomposition (SVD)
20
C Proof of Theorem 3.2
21
C.1 Rationale for the CCA-based upper bound in layer substitutability
23
D Algorithmic details for NBL
24
D.1 Computational cost of the calibration
25
D.2 Calibration runtime
25
D.3 Implementation details
26
Experiments
26
E.1
Performance efficiency 2D plots
26
E.2
Accuracy and speed-up comparisons with intervals
28
E.3 Qualitative analysis of LLM outputs after NBL
29
Quantization Llama-3, 1-70B using the AWQ algorithm
................................
30
E.5
Combining NBL with speculative decoding for faster inference
31
F
Ablation studies
32
El
Dependency on the calibration dataset
32
F2
LoRA fine-tuning of NBL-linearized layers
32
F3
Attention heavy experiments with NVIDIA RULER benchmark
33
E4
CCA-Bound criterion vs cosine distance criterion
33
F5
Greedy selection
35
G Selected Transformer layers
36
Inference complexity analysis and KV-cache usage details
37
H.l Figure analysis: prefill speed-up context length
37
H.2 KV-cache calculation for Grouped Query Attention based models.
.e..ss..........................
37
I Limitations and broader impacts
38
16

---

A Overview of Linear Minimum Mean Squared Error (LMMSE) estimation
Minimum Mean Squared Error (MMSE) estimation a fundamental statistical method for estimating unknown quantity from noisy observations. The goal of MMSE estimation is minimize the mean squared error (MSE) between the actual value and its estimate. While the general MMSE approach does not impose restrictions on the form of the estimator, the Linear Minimum Mean Squared Error (LMMSE) estimation specifically considers linear estimators, making it computationally efficient and widely applicable. This section provides concise overview of LMMSE estimation, including its derivation and its connection to our proposed method. For more detailed discussion on LMMSE estimation, refer the textbooks by Kailath et al. [2000] and Kay [1993].
A.1 General MMSE estimation
In the general MMSE framework, given observations X E R and a quantity to be estimated Y E R", the MMSE estimator is defined as the conditional expectation of Y given X:
(6)
This estimator minimizes the MSE, defined as:
MSE(Y, 8) = | - 813
The MMSE estimator is optimal in terms of minimizing the expected squared error and does not assume a specific functional form for the relationship between and Y. However, the computation of the conditional expectation high-dimensional or non-Gaussin Gaussian distributions. This motivates the use of LMMSE, which simplifies the problem by restricting the estimator linear form.
be intractable in many practical scenarios, especially for
A.2 Linear MMSE (LMMSE) estimation
In the LMMSE, the estimator y is constrained to be linear function of the observations x, given by:
W X + b
(8)
where W Raxm is the weight matrix, and Rn is the bias vector. The goal is determine W and such that the MSE is minimized. The MSE is expressed as:
MSE(Y, P)
E || 81
Tr (E [Y WX - D)(Y - WX - -b)
A.2.1 Properties of LMMSE
The LMMSE estimator has several key properties. It minimizes the MSE among all linear estimators and coincides with the general MMSE estimator when the joint distribution of X and Y is Gaussian. An important property of the LMMSE estimator is the orthogonality of the estimation error Y y to the observations X, such that E [E(X- - -[X])=] o. This orthogonality implies that no further reduction in error is possible through linear adjustments.
A.2.2 Common Use Cases of LMMSE
LMMSE estimation is widely applied in various fields. In signal processing, it is used for tasks such as noise reduction, channel equalization, and adaptive filtering. In communications, serves as foundational method for channel estimation and interference cancellation. In sensor networks, it utilized fuse noisy measurements from multiple sensors for more accurate state estimation. Control systems leverage LMMSE for state estimation, with Kalman and Wiener filtering methods being prominent examples.
17

---

In machine learning, LMMSE underpins fundamental techniques such as linear regression, where it estimates optimal weights by solving peast-squares problem to establish relationships between imput features and outputs. Beyond its computational advantages, IMMSE provides a robust framework for analyzing noise-affected data. addition, fields such as medical imaging and remote sensing, LMMSE is used to denoise and reconstruct images, delivering high-quality results from noisy or incomplete measurements.
A.2.3 Advantages and limitations of LMMSE
The LMMSE is computationally efficient due its closed-form solution, making it suitable for real- time applications. It provides interpretable results when the covariance structures are well understood and is robust under Gaussian noise assumptions. However, reliance on accurate covariance estimation and the linearity assumption can limit its effectiveness in highly non-linear systems. Furthermore, may be sensitive outliers and model specifications. Despite these challenges, LMMSE re- mains an important framework in estimation theory, valued for its theoretical rigor and practical utility.
The advantages of using LMMSE in this work include its ability to maintain computational efficiency, the closed-form solution avoids the need to retrain the model after layer replacement. Additionally, LMMSE ensures that the approximation error is minimized in a principled way, leveraging the statistical structure of the data. This property is critical for preserving the model's performance on downstream tasks after compression.
However, the linearity assumption inherent LMMSE introduces limitations. While the approxima- tion is effective for layers with strong linear correlations between inputs and outputs, it may not capture the complex, nonlinear interactions in some attention layers. This limitation is addressed in the paper by applying LMMSE selectively layers where canonical correlations suggest that linearization will have minimal impact on performance.
A.3 Derivation of the LMMSE estimator (Proposition 3.1)
Proof. To derive the Linear Minimum Mean Squared Error (LMMSS) estimator, we first solve for the bias vector b:
MSE(Y, Tr TE [(Y WX () - wx -x) Tr (E(Y*) - EY(WX + b)r] R[(WX + D)
Expanding the quadratic term E(WXX X + (W b)7):
E[WX + () X + +b E W 2 + +b"..
The derivative of the trace of matrix Tr(X) with respect matrix X is simply the identity matrix, i.e,
TTr(X)
I,
ax
where the trace operator sums the diagonal elements of X. This result is used in conjunction with the chain rule in the derivation to handle the differentiation of MSE that includes trace terms with respect W and Then, the derivative of the MSE with respect b is:
aMSE(Y,Y) ab
2E[Y] + 26 + 2WE
@MSELY )) O, solve for b:
Setting 0h
b = E - WE[X]
18

---

Next, substituting b back into the MSE expression, we focus on the centered terms X E[X] and Y- E|y]:
MSE(Y, 9) Tr ( (
E [(Y (
-
-E - - (WV(X-(())(WCC -E(X)))]
The derivative of the MSE with respect to W is:
ÖMSE(Y,Y) Y)
( - - 2wE(( [[ E[X])"].
aW
Setting
(MSE(Y.)) OW
0, we find:
WE - - (X
Rewriting in terms of covariance matrices:
WCxx 1Cxx=Co CCx
where Cxx and C x are:
Cxx = E [(X - ((X] B(x))
Crr [Y -[-]) [V])(--([X))
The optimal W is then:
W = CyxXxx
Substituting W and into the linear model, the LMMSE estimator becomes:
= B[|] Xx -X
This form highlights how the LMMSE estimator utilizes the mean and covariance of X and Y to minimize the MSE.
19

---

B Overview of Canonical Correlation Analysis (CCA)
Canonical Correlation Analysis (CCA) is a statistical technique that identifies linear transformations of two random vectors, X E R" and Y IR, such that the transformed variables maximally correlated. The goal to find pairs of canonical variables u, = a. and e = by Y for i E (1...., 9) where a, Iº and b; E IR,, with the correlation Pt between u, and vs being maximized.
B.1 Problem formulation
CCA aims to solve the following optimization problem:
C'Cxxb
max 2,
V Cxxax 12 Cyyb'
where Cxx, Cyy, Cxy are the autocovariance matrices of X, Y, and their cross-covariance, respec- tively.
B.2
Canonical Correlation matrix
To simplify the problem, change of basis applied normalize the covariance matrices Cxx and Cyy. Define:
= = xxa.
where the transformed variables 31 and Ú1 have unit variance. Substituting these into the optimization problem and using the Schwwrz inequality, the problem reduces to an eigenvalue problem:
f arg max
1
The solution involves solving the eigenvalue problem:
Cx Cr Ad,
where the cigenvalues A, are the squared canonical correlations př. and the corresponding eigenvectors T; determine the canonical weights in the whitened space.
B.3 Performing CCA using the Singular Value Decomposition (SVD)
The eigenvalue problem in (9) can alternatively be solved by performing the full SVD on the matrix
A =
UEV?.
where U E RRxP and E Rox4 are orthogonal matrices, and Σ E RRx4 is a diagonal matrix containing the singular values. From this decomposition, the relationship AAT U is obtained. Then, the matrix in the eigenvalue problem in (9), given by Cx is symmetric and therefore orthogonally diagonalizable as:
C where Q E pPxP is an orthogonal matrix and A E RExD is a diagonal matrix containing the eigenvalues. Noting that AAT = CY we can equate:
QAQT,
QAQ -T.UT.
Since Q and are orthogonal matrices and diagonalize the same symmetric matrix, follows that A Σ Thus, performing SVD on Cx Cxy Cyÿ provides the canonical correlations PL.P2.... ... as the singular values of Σ, and the results for CCA can be summarized as follows:
U RR Orthogonal matrix containing the left singular vectors, which correspond the canonical weights for X.
20

---

V E Rix;; Orthogonal matrix containing the right singular vectors, which correspond to the canonical weights for Y.
E E R Diagonal matrix containing the canonical correlations Pl P2, ... Pmiin(p.))
which are the square roots of the eigenvalues A A.
Then, the canonical variables can also be computed as:
u. at X (C
(10)
1/- b = (Cyy V,)Y,
(11)
where U; and V; are the i-th columns of U and V, respectively. Furthermore, the singular values in are the canonical correlations p;. If P; (p: 0), the i-th canonical pair is perfectly correlated (uncorrelated).
C Proof of Theorem 3.2
Proof. First, we put the description of the linear estimator Y as in (1) into the MSE formula in (2)
we have:
MSE(Y, Y) E | 813
Tr (T(B (Y -Y -)/1)
Tr (E [( WX -b)(Y WX b)
= Tr (E [(Y - WX -E(Y]+ WE(X))( - WX - + WE(X))
Tr (E [((Y- E(Y]) - W(X - -[XJ))(Y - E[y])
W(X - E[X))"],,
where Tr(-) is the trace operator defined on matrices. Then, opening the outer-product, we end up with the following expression:
MSE(Y, 8) = Tr (E [(Y (Y - B[
E [(Y - (X(T) - (
E [W(X E[X])(Y E[Y
[W(X [[ B[
As the weight matrices W are deterministic, we can take them out of the expectation, and simplify the covariance expressions using C, and using the fact CfTy = Cyx x in the second line below:
MSE(Y, F) = Tr [Cry - CyxWT - WCXY WCxxW")
Tr [Crr "xx(Cxxxxx))" - Cxx + Cxx xx
Tr [Cry Cyx Cxv Cxx
[Cry - Oxx - Cxx CxxC
= Tr [Crr -C Cr Cxx
(12)
Then, by using the matrix square root of Cry, and using the fact that for A E RKL. and B R Tr( (AB Tr(BA), we can further simplify the expression:
MSE(Y, 8) Tr [Cyr Cxx
ctrl = =Tr T C cnll
= =Tr c 1-
[ [ ond
= =Tr
T - c ( crr]
= T
21

---

Here, the canonical correlation matrix is corr(Y, X) C CopCrxOxx Then if we apply full Singular Value Decomposition on this matrix as UEVT = corr(Y, such that U TEAuu hhea R hu xh and V Rh x Also, using the fact that U and V are orthogonal matrices such that
UU = U and VVT = VVTV = II we have:
MS 8) = [ [I- UEV" (UEVT)* cry] Cyrl
= Tr -UCU"-UEE'UUN crv]
Tr [U [I- ---E-11 U Cyr]
-Tr [[T W Cy-UU.
We have ΣV" = diag(p? p2... Pron where P: = 0
for for > min(hbaa hm). Then, we can
write the MSE expression as:
MSE(Y, = Tr - ....., P U CrrU]
Σ (1 (2)(UUCyyU))
Then the product CrrCru can be written the following by using the symmetric ma- trix singular value decomposition of the covariance matrix Crr = Uv and Ey diag(oy,,0y,2 Nyhh ):
CrrUa= Σ Uy
Then, we have the MSE expression as
MSE(Y, 8) Σ (1 - (U)(UTCTrUU)
1 p7) Σ Σoyn(UUUuy)2
2 pi) Σor,,
E (CyT)(Crv),
where the inequality arises from the fact that UT product of orthogonal matrices, ensuring that its elements have magnitudes less than or equal to one. Also, the last equality comes by the fact that the sum the singular values of matrix equal to the trace the same matrix. Then we can write the bound, the MSE normalized by the trace of the auto-correlation
NMSE(Y, Ý) MSE(Y, Ý)
p?
Tr(Cry)
1=1
(hout r) +
Σ
(1 p?))
7m
such that min (in) using the fact P. 0 for i r.
22

---

C.1 Rationale for the CCA-based upper bound in layer substitutability
Let Y E IR denote layer's output and Y its linear approximation from the corresponding imput X.
Writing the error as E Y, the normalized mean-squared error (NMSE) is
Y 813 E Y 12
Tr(Cyr - -CyxCxX
NMSE(Y,Y) Y)
(13)
Tr[Cyr]
Our analysis employs a CCA-based upper bound that evaluates a whitened error:
NMSE(Y, P) S (hout r) + -0-
(14)
1=1
where P1 2 2 Pr are the canonical correlations between X and Y (with == rank). This
connects per-mode approximation quality directly to canonical alignment.
The choice of CCA-based bound over the scalar NMSE in (13) is motivated by the following considerations:
1. Mode-wise, variance-agnostic assessment. The scalar NMSE aggregates error with variance weighting through Tr[Cyv]; high-variance directions can dominate, potentially obscuring failure low-variance (but semantically important) modes. In contrast, (14) decomposes eπror along canonical directions and evaluates per-mode alignment via 1 pf after
whiteming by Cyy This avoids skew from marginal variances and directly targets linear predictability of Y from X, which is the relevant notion for safe linear substitution.
2. Comparable computational profile with improved numerical behavior Both approaches require second-order statistics, i.e., estimates of Cxx, Crv, and Cxy. Forming linear predictors and their induced error covariances for NMSE involves matrix inver- sions/multiplication:: the CCA route performs whitening and an SVD of the whitened cross-covariance. The big-O complexity is comparable, while the explicit whitening in CCA often yields more stable behavior in low-sample or near-singunt singllrn regimes by regularizing anisotropy the covariance estimates.
3. Empirical evidence for superior layer selection. On MISTRAL--B, selecting the 12 most "substitutable" layers by the CCA-based criterion (smallest average p?) preserves reasoning accuracy significantly better than selecting by direct NMSE:
Criterion
Baseline (%)
After replacement (%)
CCA-based
70.2
68.3
Direct NMSE
70.2
39.1
The -CA-based selection aligns with the goal of identifying layers whose outputs are well explained by linear projections of their inputs, thereby enabling safer linear substitution.
While NMSE standard scalar figure of merit estimation theory [Kay, 1993, Kailath et al., 2000], it can be biased toward high-variance directions and provides limited diagnostic insight. The CCA-based bound (14) delivers principled, variance-agnostic, and mode- mose wise indicator tightly coupled to linear substitutability. This theoretical alignment, together with comparable computational cost and favorable empirical behavior, motivates its use as the guiding criterion for layer replacement.
CCA is classical tool for quantifying linear relations between multivariate random vectors and has been widely used to probe representational alignment and predictability [Hotelling, 1992, Hardoon al., 2004]. Its canonical-me mode decomposition naturally supports analyses like the Equation (14), that evaluate approximation quality per direction.
23

---

D Algorithmic details for NBL
Algorithm 2 Pseudocode for NBI. Weight/Bias and CCA Bound Computation for given Self-
Attention Layer
Input: X (attention layer input), Y (attention layer output).
2: Output: Weight matrix W and bias b, CCA-bound.
3: Y4 + + [Residual connection output.)
4:
5: Mean Computations
6: E[X] Mean(X)
7: Y f MY)
8: E[Y+] Mean(Y+)
9:
10:
11: Cyx < Cov(Y, X)
(Unbiased estimator for cross-covariance) }
12: Cx,X + Cou(Y,,
13:
Covariance Matrices
15: Cxx + Cov(X,X) (Unbiased estimator for auto-covariance.) CY (Y),
17:
Digen-ecommpsittion
19: (14 k + (igh(Cy
(Symmetric matrix eigenvalue decomposition, )
20:
(Axx, Vxx) f- Figh(Cxxx)
21:
Inverse Square Roots
23:
Cy 1/2
+ Viy diag( (>
1/2 Y1
VT
1444
(Matrix square roots (Cy,y, = YLY, CLYD
C2/2 C2/2
24: CXV + Vxx diag Axxx
1/2
1/2
Vxx
25:
26: Correlation Matrix
27: Cw +
+ Cx,Y, Cy,x OXX
1/2
(Calculate the standardized cross-correlation matrix.)
28:
29: Singular Value Decomposition (SVD)
30: VU) SVD(Cw) (Perform SVD to compute the CCA singular values.)
31: (21,P2...., pr) &- (Sa (r is the embedding dimension.)
32: CCA-bound t Ei=(1 pt)
[Calculate CCA Bound based Theorem-(3.2)
33:
Weight Matrix and Bias
35: W +
Xx -Cxx
(NBL weight.)
E[Y] WE[X] (NBL bias.)
Algorithm provides detailed procedure for NBL parameter estimation and CCA-bound computation. This approach approximates self-attention layers by analytically deriving linear weights and biases while quantifying redundancy through Canonical Correlation Analysis (CCA). takes the input X and output Y of any given self-attention layer. To capture the full behavior of the outputs, the residual connection is added as Y, = Y +X. While the CCA-bound calculated based on the output of the residual connection, the lincarization weights and biases computed directly from the attention output. As a result, the residual connection is retained in the compressed layer. Although it is possible to calculate the weights directly from the residual connection output and remove the residual connection in the compressed model, its inclusion adds minimal complexity to the model and helps preserve the input information more effectively.
First, the algorithm calculates F[X], F|Y) and F, then unbiased estimates of cross-covariances matrices, Cyx and Cr, x, and also auto-covariances, Cxx and Cr,y,- y, Further, based on these covariance matrices, cigen-decompossiton is applied, which gives eigenvalues and cigenvectors for
24

---

calculating inverse square roots of covariance matrices C,,Y, and CXX² The CCA standardized corss-conrllation matrix (Ow) is then computed using the cross-covariance matrix Cy, x and with the inverse square roots of both input and output covariance matrices. Thereafter, pi. the CCA singular values from SVD over Cty computed. The CCA-bound quantifies the redundancy present the attention layer by computing the total of 1 - P summed over all dimensions of the embeddings.
At last, the weight matrix and bias of the NBL are computed. The weight matrix computed from the cross-covariance (Cyx) and inverse (Cxx), whereas the bias can be induced from the expression (b= = _[ - W E[X]. This method realizes not only a linear approximation of the attention layer but also rigorous quantification of the redundancy, hence this turns out to be strong approach toward the compression and optimization of layers.
D.I Computational cost of the calibration
The proposed algorithm for NBL parameter estimation and CCA-bound computation is designed to be computationally efficient, making it feasible for large-scale models such Mistral-7B. Llama-3, 3, 1-8B and Llama-3. 3. 1B. The primary computational steps involve mean and covariance computation, symmetric matrix sigen-decomposition, and singular value decomposition (SVD). These operations dominate the overall runtime complexity, with eigen-decomposinn and SVD scaling (d²), where d is the embedding dimension of the attention layers.
The algorithm first computes the mean of input and output activations, requiring O(s operations, where 8 is the number of sequences and t the context length. Covariance matrices, including cross-covariances,
estimated with complexity of O(s t d²). The cigen-decomposition of covariance matrices and the SVD of the standardized correlation matrix both scale as O(d²). Matrix inversion for computing the weight matrix also has complexity of Odd)), while bias computation is relatively inexpensive at (dP).
d.). Overall, the total complexity of the algorithm is O t d²). The method is efficient due to optimized linear algebra routines and direct computations, Additionally, the algorithm requires no hyperparameter tuning or fine-tuning beyond specifying the number of calibration samples.
D.2 Calibration runtime
The main computational overhead the calibration process arises from Algorithm which performs Canonical Correlation Analysis (CCA) and computes the associated linear parameters. To empirically assess the runtime, we run the algorithm using float32 precision on randomly initialized matrices that simulate the attention representations of different model scales in the Llama-3.1 family, as presented in Tables and Each experiment used 256 calibration samples with a context length of 2048. Since the calibration is applied independently each attention layer, the total runtime for model can be estimated by multiplying the per-layer runtime by the number of layers.
For the GPU-based implementation, the tensors required for the Llama-3 -. 1-00BB scale-with the specified precision, sample size, and context length-do do not fit within the memory of single NVIDIA A100 (80 GB). Therefore, we report GPU results Tables and 7 using two A100 GPUs running in parallel across all three model scales, whereas all other runtimes in this paper are calculated with a single GPU. As shown in these analyses, although calibration can be computationally demanding at large scales on CPU, it remains tractable with GPU ac- celeration and becomes more efficient when the cost amortized across layers and repeated use cases.
In our implementation, compressing the Llama 3. 1-8B and Mistral-78 models-each containing 32 attention layers takes under 30 minutes using an NVIDIA A100 GPU (80 GB) for model inference and activation extraction, computing CCA bounds and NBL. weights and biases. This runtime includes all necessary steps, including covariance estimation, SVD, and linear parameter computation. These
25

---

Table 7: Calibration Runtime for Different Model Scales (256 samples, 2048 context length)
Model
Llama-8B
Llama-70B
Llama-405B
Hidden size (d)
4092
8192
16384
Hidden layers
32
80
126
Runtime per layer (CPU)
62.16 sec
324.79 sec
1839.64 sec
Total (CPU)
00.5 hr
~7.22 hr
~64.38 hr
Runtime per layer (GPU)
26.04 sec
78,91 sec
371.79 sec
Total (GPU)
00.24 hr
11.75 hr
~11300 hr
results demonstrate the practicality of our method for compressing large-scale models with minimal overhead. Moreover, the implementation can be further accelerated through additional optimization and system-level improvements.
D.3 Implementation details
The algorithm is implemented using PyTorch [Paszke et al., 2019] and SciPy [Virtanen al., 2020] for tensor operations, linear algebra routines, and digen-decomposition and HuggingFace ormers [Wolf, 2019] for loading and managing the pretrained models. These libraries provide optimized and robust functionality for large-scale matrix operations, enabling efficient computation on standard hardware. Additionally, the implementation builds the code repository provided by He et al. [2024]]
After passing the calibration data through the model and computing the activations, typically on the GPU, the computations in Algorithm 2 can performed on either the CPU or GPU, depending on the available compute resources. We provide efficient implementations for both options. Running the algorithm on the CPU helps reduce GPU memory overhead and ensures broader hardware compatibility, while using the GPU offers faster runtime for large-scale modcls. PyTorch and SciPy's s linear algebra routines ensure efficient execution of the covariance, eigen-decomposstton, and SVD steps efficiently either backend, enabling smooth integration with diverse workflows.
E
Experiments
E.1 Performance - effficiecc 2D plots
We have included performance-efffciency plots with the pooled standard error (SE) intervals, where visualize accuracy against both throughput and KV-cache gains, for the two best performing methods, Attention NBL and Attention DROP. The plots for Llama-3 3. 1-8B, Mistral-Pn 7B, DeepSeek-Distiee Llama-8B and Llama-3. 1-8B can be found in Figure 4. To evaluate the statistical significance of the differences, used the pooled standard error (Pooled SE) aggregate the uncertainty across tasks when reporting average performance. The Pooled SE is computed as:
Pooled_SE
Σ SE
where SE, denotes the standard error of task and is the total number of tasks. We have n=8 benchmarks similar to the setting in Tables 2, 3, 4. The plots 4a, 4b, 4c indicate that, particularly at higher compression rates, NBL demonstrates a more favorable trade-off between KV-cache usage, accuracy, and throughput, suggesting statistically significant Pareto optimality.
'hanss/gittibbccm//AAEEELLL LL--LLLLLL--L-M-MMA
26

---

Avg. Accuracy
Avg. Acturacy
65.0
65.0
69.5
62.5
57.5
57.5
25.0
Attn DROP
35.0
Attn DROP NBL
52.5
Attn NBL
52.5
% KV-Cache Savings
Throughput Speed-up
(a)
68
68-
Avg. Accuracy
06
Avg. Accuracy
68-
64
64-
62
62
56
Attn DROP
56-
DROP
Attn NBL
Attn NBL
54
56
10
20
30
40
50
1.0
1.1
12
13
14
KV-Cache Savings
Throughput Speed-up
Avg. Accurscy
Avg. Accuracy
60
60-
00
50
55-
Attn DROP
Attn DROP
50
50
% KV-Cache Savings
Throughput Speed-up
(c)
Figure Accuracy versus (left) KV cache savings (right) throughput speed-up across varying compression levels for NBL and DROP, on (a) Mistral-7B (b) Llama-3.1-8B (c) DeepSeek-R1- Distill-Llama-BB models with the corresponding Standard Error intervals.
27

---

E.2 Accuracy and speed-up comparisons with intervals
In this section, we present Tables 9, 10, and 11 with the average accuracy, prefill and throughput speed-ups ups with standard error and standard deviation intervals for our Attention NBL method and the baseline Attention DROP, corresponding to the results in the main text.
Table 8: Mistral-7B: Accuracy and Speedup (Table 2).
Method
Avg. Ace. SE (T)
Prefill Speed-up (1)
Throughput Speed-up (1)
Base
70.1 0.41
Drop-4
70.0 0.41
1.08 # 0.03
1.10
Drop-8
69.4 0.41
1.18 # 0.05
1.22 0.01
Drop-12 12
66.5 : 0,42
1.29 : 0.04
1.29 0.03
Drop-16
52.9 0.42
1.44 0.05
1.41
NBL-4
70.1 # 0.41
1.08 # 0.04
1.00
NBL-8
70.0 # 0.41
1.17 = 0.03
1.20
NBL-12
68.3 # 0.42
1.28 = 0.05
1.27 0.01
NBL-16
58.8 0.42
1.40 0.04
1.37 0.02
Table Llama-3.1-8B: Accuracy and Speedup (Table 3).
Method
Avg. Acc. # SE (1)
Prefill Speed-up (1)
Throughput Speed-up (t)
Base
70.1 0.41
Drop-4
1.08 + 0.035
1.11 0.02
Drop-8
69.8 0.41
1.18 0.048
1.01
Drop-12 12
67.9 0.42
1.033
1.29
1rop-16
55.2 0.43
1.02
1.42
NBL-4
70.2
0.040
1.11
NBL-8
69.8 0.42
1.06
1.17
NBL-12
69.1 + 0.42
1.24
1,25 0.02
NBL-16
62.5 0.42
1.39 # 0.043
1.37 # 0.01
Table 10: DeepSeekkRR--Diiill--aaaaaa
Accuracy and Speedup (Tables 4).
Method
Avg. Ace. = SE (t)
Prefill Speed-up (1) Throughput Speed-up (1)
Base
63.6 0.42
Drop-4
63.9 0.42
1.08 0.04
1.09 0.01
Drop-8
63.5 0.42
1.16 0.04
1.01
Drop-12 12
60.7 0.42
1.28 # 0.05
0.002
1rop-16
50.2 0.42
1.37 # 0.04
1.07
NBL-4
63.6 0.42
1.08
1.08
NBL-8
63.4
1.15
1.09
NBL-12
62.4 0.42
1.24
1.32
NBL-16
56.8 0.42
1.35 # 0.05
1.42 = 0.01
Table 11: Llama-3.1-70B (quant.): Accuracy and Speedup (Table 5).
Method
Avg. Acc. SE (t)
Prefill Speed-up (1)
Throughput Speed-up (t)
Base
76.0 0.40
Drop-32
75.7 0.40
1.20 # 0.033
1.30 + 0.01
Drop-48
70.4 0.40
1.29 # 0.052
1.42 = 0.02
Drop-54
48.3 + 0.42
1.34 0.059
1.53 0.02
NBL-32
76.0 0.41
1.16 0.039
0.00
NBL-48
71.3 + 0.41
1.25 0.044
1.35 0.01
NBL-54
65.4 0.43
1.30 0.021
1.44 1.400.00.
28

---

E.3 Qualitative analysis of LLM outputs after NBL
Here we present the output of Llama- -1ama-3-880 --B-Instuct" model on sample from the test set of GSM8K [Cobbe et al., 2021] as the input in Table 12. Here, the NBL gives the correct answer until 12 layer modification, whereas DROP fails to give the correct answer in similar compression.
Table 12:
Correct answers are marked with the symbol "V "v", and incorrect ones with "X" All outputs are in response to math word problem; the correct answer extracted with the token "#"""""
Question
Mark has garden with flowers. Hc planted plants of three different colors in it. Ten of them are yellow, and there are 80% more of those in purple. There are only 25% as many green flowers as there are yellow and purple flowers. How many flowers does Mark have garden?
Reference Answer
There are 80/100 10 = <80/1000 10=8*8 more purple flowers than yellow flowers. \n So Mark's ' garden, there are + 8 = 10+8=18+18 purple flowers. in Purple and yellow flowers sum up to 10 + 18 flowers. in That means Mark's garden Mark has
=25/100=28=7w7 green flowers, <28+7=35=35 plants his garden. ####
there are 25/100
Method
Output
Baseline
There are 10 ycllow flowers. 80% more of those are purplc. So 10 + (10 x 0.8) = 18 purple flowers. 10 + 18 28 yellow and purple flowers. 25% of that is 28 x 0.25 = green flowers. /
10 + 18 7 35 flowers. The answer 35. \n NMN! 35
Attn NBL-8
18 purple flowers. There are 18 yellow and purple flowers. There are 25% as many green as yellow and The
NBL-8 Ans: There are 10 yellow flowers. There are 10
purple, there are 0.25 28 green. 7
18
35.
in #### 35
Åttn NBI.-12
The number of yellow flowers is 10. The number of the other type of flower is 10 1.8 18. The number of green ones is 10 + 18 28, and 0,25 28 = 7. So the total number of flowers + 18
35. ln
35
Atta NBL-16
The number of the first group (the one with the highest number) is 100% and the other two are 1.5 and 0.75. So the number of the first is 100% and the other two are 1,5 and 0.5. So the total number the first 100% and the other two are and 1.5.
Attn DROP-8
There are yellow flowers. There are 80% more of those in purple, so there are 10 18 purple flowers. total number yellow and purple flowers There 25% many green flowers there are yellow and purple flowers, so there
0.25 28 7 green flowers. The total number of flowers is 7 28 35. \n #e## 35
Attn DROP-12
There 10 yellow flowers. 80% of 10 is there
more, 18 in total. 25% of 18 is 4, so there are 4 green. 10 + 8 = 22. So there are 22 of the non-green flowers. 4 is the green. So the total is 26. The answer 26. X
Attn DROP-16
We'r given that the number of fish are is the main character, the most important thing is the the the problem. The problem states 1. the (the phrase and if, the (the phrase a, The (the " or (a, (the. The the (the, the (The the. The the. Let's analyze the first line, the. is and, the, is and, the. and, and. The a, the. and, the, is and, the.
29

---

E.4 Quantization of L.ama-3. using the AWQ algorithm
The Activation-awarc Weight Quantization (AWQ) algorithm Lin al. [2024] a hardware-friendly method designed for low-bit weight-only quantization of large language models (LLMs). AWQ focuses on preserving model performance while significantly reducing memory and computation costs. Below, we outline the mathematical framework and process to quantize the Llama-3. 10ma-3.1-8B.
Weight quantization:
Quantization maps floating-poin weights W to low-bit integers Q(W). For N-bit quantization, the function defined as:
W
max ( (W))
Q(W) = Round
2
where is the quantization scale, and (ound(.) maps values to the nearest integers. This reduces precision but can incur noticeable error if all weights are treated equally.
Preserving salient weights:
A key observation is that not all weights contribute equally to model performance. AWQ identifies the most important weights, termed salient weights, using activation magnitudes. Channels corresponding to large activation values are more likely to process critical features. Instead of mixed-precision quantization, AWQ applies scaling factor error. For weight W, and input vectOR X, the scaled quantization process is:
to these weights to reduce relative quantization
X = A'or Round Q(W s)
W
s A'
where A' is the updated quantization scale. The ratio of quantization errors before and after scaling given by:
A'
~ 1 for A' ~ A.
Error Ratio
A
Thus, scaling effectively reduces the relative error for salient weights.
Optimal scaling factor search:
The AWQ algorithm optimizes the scaling factors by minimizing the output difference between quantized and original weights:
L(s) Q W diag(s)) (diag(s)) WX
where s is per-channel scaling vector, and X represents cached activation features from small calibration dataset. To simplify the search, AWQ uses the average activation magnitude per channel, sx, proxy for saliency: s = x a
arg min L(s).
which allows an efficient search over single exponent a.
30

---

AWQ and NBL:
The AWQ algorithm's per-channel scaling and activation-awen aware strategy enable efficient compression without requiring retraining or large calibration datasets. AWQ is also designed support layer fusion to optimize hardware efficiency. By fusing linear layers with adjacent operations, AWQ reduces intermediate memory access and kernel overhead, enhancing its performance even further. However, in the implementation used for the Llama-3 1-70B quantization in the NBL framework, we opted to use unfused version for simplicity and flexibility during experimentation. Additionally, the linear weights calculated by NBL were also quantized by the AWQ method for compatibility. The speed-up measurements were done by generating 2048 tokens with a context length of 2048, and a batch size of 1, similar to the experiments with smaller models. We used the publicly available AutoAWQ repository for the implementation of the AWQ on top of NBL..
Future work can address these limitations by integrating quantization-awal techniques directly into the NBL pipeline. For instance, fully quantized NBL framework could adopt fine- see-ggriied scaling for linear layers while retaining the ability to fuse operations dynamically, thereby combining flexibility with optimal hardware utilization. Such advancements could further reduce memory overhead while maintaining model accuracy.
E.5 Combining NBL with speculative decoding for faster inference
We evaluate the integration of speculative decoding and Neural Block Linearization (NBL) using the DeepSeek-R1-DD -i-isttil-iiiill---- Masttll-llaaaae -L1ama-88 model [Guo et al., 2025]. Experiments conducted on the MT-bench [Zheng et al., 2023] benchmark using an NVIDIA A100 GPU. For speculative decoding, adopt the EAGLE-3 method [Li et al., 2025], which extends the traditional sraft-and-veriine framework by enabling the generation of multiple tokens per iteration with a low-cost draft model. In each decoding step, the draft model proposes sequence of candidate tokens, which then passed to a verifier model. The verifier checks these tokens in parallel, accepting the longest prefix that aligns with its output. This process amortizes the cost of expensive autoregressive generation and significantly accelerates decoding.
To evaluate compatibility with NBL, we replace the standard verifier model in EAGLE-3 with its NBL-compressed counterpart. Since NBL preserves the input-output behavior of selected attention layers through linear approximations, the verifier remains functionally consistent while being more efficient. This substitution does not require changes to the speculative decoding logic genera- tion pipeline. For the draft model, we use the publicly available weights from the EAGLE repository 4.
Our results show that combining NBL with EAGLE-3 leads to compounding speed-ups, achieving up to 4.07x acceleration (see Table 6). This demonstrates that NBL. is not only orthogonal to decoding- level acceleration strategies but can also enhance them with negligible integration overhead. These findings highlight NBL's practicality for real-world deployment scenarios, where both model size and inference latency are critical.
31

---

Ablation studies
F.
Dependency on the calibration dataset
Tables 13 and 14 present validation results on and WikiText-2 for pruned or linearized Llama-3.1- 8B and Mistral-7B models, comparing SLEB, Attn DROP, Attn NBL (each applied to 8 attention layers), and SliceGPT with 25% sparsity. These results are included to assess each method's dependeney the choice of calibration data used to prune linearize LLMs.
Table 13: Perplexity Results on C4.
Table 14: Perplexity Results on WikiText-2.
Calibration (D):
Wiki-2
C4
Wiki-2
C4
Calibration (D):
Wiki-2
C4
Wiki-2
C4
Method
Llams-3 3. 1-88
Mistral 778
Method
Llama 3. 1-88
Mistral 1-lall178 7B
SliceGPT-25%%
123.29
28.28
13.07
SliceGPT-25%%
14.27
74.33
12.11
SLEB-8
20.53
20.54
15.07
15.49
SLEB-8
16.70
16.67
10.23
11.13
Attn DROP-8
11.40
11.35
9.20
9.20
Attn DROP-8
7.39
7.48
6.81
6.81
Attn NBL-8
12.22
11.37
9.85
9.08
Attn NBL-8
7.18
9.57
5.59
7.33
F.2 LoRA fine-tuning of NBL-linearized layers
We further investigate the potential of fine-tuning the linear layers derived by our NBL method using the Low-Rank Adaptation (LoRA) framework [Hu al., 2022]. Specifically, apply LoRA to the linearized layers on the same C4 dataset used for calculating the NBL parameters. This dataset closely resembles the original pretraining data, ensuring consistency in domain and distribution.
Our experiments use the DeepSeek-R1- R1--isttiDllil 1aw.i11----aa--------------- model, where 12 and 16 attention layers are replaced with their NBL-linearized counterparts (denoted Attn NBL-12 and Attn NBL- 16). We then fine-tune these lincarized layers using LoRA with a rank of 32, a = 64, and a dropout of 0.1. Training is performed in bfloat16 for 3 epochs with learning rate of 1e-4, effective batch size of 16, and context length of 1024 tokens using a 5000-sample subset of the C4 validation split under causal language modeling objective.
To further assess generalization beyond the calibration data, we also apply LoRA fine-tuning using SlimPajama [Soboleva et al., 2023] dataset-2 large-scale corpus with different statistical properties than C4. This setup allows us to explore LoRA's effectiveness under a moimatemismmtched regime.
These results in Table 15 suggest that LoRA-based refinement offers only marginal improvements over NBL alone-e.g. 62.5% vs. 62.4% for NBL-12 and 58.2% vs. 56.8% for NBL-16. indicating that most of the performance gains stem from the underlying NBL mechanism itself. Furthermore, the consistent behavior across both matched and mismatched data regimes highlights the robustness and generality of NBL, where it provides competitive reasoning performance even without additional parameter-efficient fine-tuning, making a lightweight and effective tool for compressing large pretrained language models.
Table 15: Average reasoning accuracy of NBL applied DeepSeek- R1-Distill--Lamaaeee model before and after LoRA fine-tuning on C4 and SlimPajama pretraining datasets.
Model Variant
Average Accuracy (%)
Baseline
63.6 0.42
NBL-12
62.4 + 0.42 62.5 0.42
NBL-12 + LoRA (C4)
NBL-12 + LoRA (SlimPajama)
62.6 0.42
NBL-16
56.8 + 0.42
NBL-16 + LoRA (C4)
58.2 + 0.41
NBL-16 + LoRA SlimPajama)
58.1 0.43
32

---

F.3 Attention heavy experiments with NVIDIA RULER benchmark
We report additional results on the RULER benchmark [Hsieh et al., 2024], which designed to evaluate models on iniention-intensivy tasks requiring long-context reasoning and dependency tracking. Using the Attn NBL retains over 95% of baseline accuracy even when 2-4 attention layers are linearized, while SLEB exhibits much sharper degradation. These findings highlight that NBL. can selectively reduce computational cost in attention-heavy settings while preserving substantially better performance than recent baselines on long-context reasoning tasks.
model, the results in Table 16 show that
Table 16: Average accuracy (%) of DeepSeck-R--Distille RI-Distill-LLama-- on NVIDIA RULER benchmark [Hsieh et al., 2024] at 4K context length).
Method
Average Accuracy (%)
Baseline
90.02
SLEB-2
69.78
SLEB-3
61.29
Attn DROP-2
87.78
Attn DROP-3
86.66
Attn DROP-4
83.15
Attn NBL-2 (ours)
88.59
Attn NBL-3 (ours)
88.21
Attn NBL-4 (ours)
86.19
F.4 CCA-Bound criterion vs cosine distance criterion
We analyze the performance of NBL under two different layer selection criteria: the CCA-bbund criterion introduced in Theorem 3.2, and the cosine distance criterion originally employed in the DROP method of He et al. [2024]. Results for the cosine distance criterion are reported in Tables 17 and 18, with the last column ("CCA Avg.") showing the averages obtained from the CCA-bound criterion (Tables 2 and 3).
For the Mistral-78 model (Table 17), the cosine-based criterion performs comparably at smaller layer intervals (e.g., Attn NBL-4 and Attn NBL-8), but its performance degrades sharply for larger intervals (e.g., Attn NBL-12, Attn NBL-16), where both the average accuracy and CCA scores decline. Similar trends are observed for the L1ama-3. 3. 1-8B model (Table 18), where performance drops from 70.2 for Attn NBL-4 to 58.0 for Attn NBL-16. These results suggest that the cosine distance criterion is less stable when approximating transformations across broader sections of the network, whereas the CCA-bound criterion provides greater reliability.
While the performance gap between CCA and cosine distance may appear less pronounced in some benchmarks, choice of CCA is motivated by its stronger theoretical grounding. CCA explicitly measures the shared subspace between input and output representations, capturing structural alignments that preserve information flow. In contrast, cosine similarity only compares Fuclidean angles, potentially overlooking linear dependencies. To illustrate this, consider a two-dimensional example with input and output tensors each consisting three samples:
0
X
0
Y =
1
0
0
There exists transformation matrix A where:
1 1
A
such that = X A. Although this mapping linear, the cosine similarity between each corresponding row of X and Y is zero, since the vectors are orthogonal. CCA, in contrast, identifies the full-rank
33

---

alignment of the shared subspace and assigns a maximal similarity score, reflecting the underlying linear connection. While this example is simplistic compared the highly structured transformations in L.Mss, it highlights fundamental limitation of cosine-based criteria in detecting linearly aligned signal transformations.
These challenges become particularly relevant in attention-heavy heavy tasks, where subtle but important features may be selectively amplified in sparse or nonlinear ways. Consequently, the difference between CCA and cosine distance is most pronounced in more demanding evaluations where errors in layer selection carry greater cost.
To validate this, we evaluated both criteria on the attention heavy NVIDIA RULER benchmark (Hsich al., 2024] under the same setting in Table 16 using the DeepSeek- -R---iiiiiitl--- L1ama-88 and Llama-3 1-8B models in Tables 19 and 20 respectively, observing a consistent trend.
Across both models, the CCA-based criterion consistently outperforms the cosine-based criterion as the number of substituted layers increases, highlighting its advantage guiding structure aware substitutions and maintaining robustness in tasks that emphasize long-range reasoning. This empirical evidence further supports Theorem 3.2, which establishes the CCA-bound as a principled measure of the linear structure preserved between layers, and justifies its use as the preferred selection criterion for layer linearization with NBL.
Table 17: Performance of Mistral-7B with NBL using CCA vs cosine based criterion.
Method
ARC- (aorm)
ARC-c (norm)
BoolQ
HellaSwag (norm)
MMLU (5-shot)
OBQA (norm)
PIQA (norm)
Wino- Grande
Cosine- Busine-Based
CCA-Based
Avg. (t)
Avg.
Atnn NBL-4
80.0
53.7
83.4
80.6
62.4
44.6
81.9
74.0
70.1 + 0.40
70.1
Atm NBL-8
80.4
52.8
83.6
79.8
62.1
44.0
81.8
74.4
69.9 0.40
70.0 0.41
Atnn NBL-12 12
76.6
49.1
83.5
76.5
60.4
42.6
79.8
73.3
67.7 0.42
0.42
Attn NBL-16
62.4
42.8
76.7
69.4
33.0
39.8
76.4
70.2
58.8 + 0.40
58.8 0.42
Table 18: Performance Llama-3.1-8B with NBL using CCA vs cosine based criterion.
Method
ARC-e (norm)
ARC-c (norm)
BoolQ
HellaSwag (norn)
MMLU (S-shot)
OBQA (norm)
PIQA (norm)
Wino- Grande
Cosine-Basod Avg. (f)
CCA-Based
Avg. (t)
Altn NBL-4
81.9
$4.0
82.2
78.1
65.0
45.8
81.1
73.4
70.2 2 0.43
70.2:041
NBL-8
77.2
64.0
45.4
73.3
Atnn NBI-12 12
79.1
52.2
82.3
75.2
64.8
45.2
79.9
74.0
69.0 + 0.41
69.1 0.42
Attn NBL-16
69.0
39.1
41.8
73.1
58.0
62.5+043
Table 19: Performance of DeepSoek-RR-DDsttil-Llamm criterion on NVIDIA RULER benchmark [Hsich et al., 2024].
with NBL using CCA cosine based
Method Variant
Cosine-Basee Based (%)
CCA-Basee (%)
Bascline
90.02
90.02
Attn NBI-2
88.50
88.59
Attn NBL-3
87.38
88,21
Altn NBL-4
85.62
86.19
Table 20: Performance Llama-3.1-8h with NBL using CCA vs cosine based criterion on NVIDIA RULER benchmark [Hsieh et al., 2024].
Method Variant
Cosine-Based (%)
-CA-Based (%)
Baseline
93.28
93.28
Attn NBL-2
91.51
91.89
Attn NBL-3
90.36
91.56
Attn NBL-4
89.68
90.85
34

---

E55 Greedy selection
In this section, we perform an ablation with greedy selection, which iteratively checks for the differences in the bound scores and compresses the model incrementally in multiple iterations. Results are shown in Table 21 and 22, where the last column shows our NBL results from main body. Our results show that our NBI. with CCA criterion outperforms greedy selection. Greedy linearization alters the activation distribution leading to inconsistencies in layer ranking. We also noticed that greedy selection chooses earlier layers, which are usually more less amenable to linearization, as shown in the main body.
Table 21: Performance of Mistral-7B with greedy selection.
Method
ARC-e
ARC-e
BoolQ
HellaSwag (norm)
MMLU (5-shat)
OBQA (norm)
PIQA (norm)
Wino- Grande
Greedy Avg. (?)
NBL
(norm)
(norm)
Avg. (1)
Greedy-4
79,9
53.4
83.8
80.6
44.2
81.9
73.9
70.0 0.40
70.1
Greedy-8
79.3
51.7
82.7
79.1
62.2
43.6
81.0
73.6
69.2 + 0.44
70.0 0.41
Gready- 12
69.2
46.1
81.0
74.1
46.3
40.2
79.2
72.9
63.6 1: 0.43
68.3 a 0.42
Greedy-16
60.6
40.8
76.1
68.5
36.8
67.8
70.2
57.7
58.8 0.42
Table 22: Performance of Llama-3.1-8B with greedy selection.
Method
ARC-c (norm)
ARC-E
BoolQ
HellaSwag (norm)
MMLU (5-shot)
OBQA (norm)
PIQA (norm)
Wino- Grande
Greedy Avg.
NBL
(nor))
Avg. (t)
Greedy-4
82.2
$4.4
82.3
78.4
64.7
46.0
80.9
73.1
70.2 1 0.42
70.2 0.41
Greedy-8
81.7
53.8
82.2
77.2
64.7
45.2
80.7
74.4
70.0
69.8
Greedy-12
79.2
52.1
82.5
74.9
45.6
80.2
73.0
66.9 0.40
69.1 0.44
Greedy-16
69.2
39.8
71.3
68.1
26.1
42.2
77.6
66.6
4.49
0.45
35

---

G Selected Transformer layers
Table 23 presents the sorted importance rankings of attention layers selected by the Attn DROP and Attn NBL methods across various models and calibration datasets. In all configurations, the methods consistently prioritize dropping or linearizing the higher-indexed layers, particularly those the end of the model, suggesting that the later layers may be more redundant or compressible in terms of their contribution to overall performance. This pattern is prominent in all models, where the top-ranked layers for pruning are overwhelmingly concentrated toward the final half of the network.
Table 23: Sorted attention layer rankings selected by Attn DROP and Attn NBL across different models and calibration datasets. Lower ranks indicate higher importance.
Method & Model
Sorted Layer IDs (Most to Least Important)
25, 26, 27, 24, 22, 28, 23, 30, 31, 21, 29, 20, 19, 18, 17, 16 6, 9, 12, 14, 13, 11, 15, 4, 10, 7, 5, 0,
Atn DROP Mistral-78 (C4)
24, 25, 22, 23, 27, 26, 20, 28, 19, 29, 21, 18, 30, 17, 16, 15
Attn DROP - (C4)
31. 11, 10, 13, 12,
24, 23, 22, 25, 26, 20, 27, 19, 28, 18, 21, 29, 17, 30, 16, 15
Attn DROP - -D-DP----istill---- (C4)
31, 11, 14, 10, 13, 4, 9, 2, 6, 3, 12, 7, 1,0
25, 24, 22, 23, 26, 28, 27, 20, 21, 18, 17, 16, 31 4,
Atta DROP Ll -lama-3-83 Mamma---BBIIntteeee Instruct (C4)
15, 11, 14, 10, 13,
62, 65, 59, 61, 63, 46, 50, 58, 48, 51, 53, 57, 60. 64, 66
Attn DROP Llama -. 1-70B (quant.) (C4)
67, 68, 49, 69, 55, 56, 47, 44, 52, 42, 45, 70, 43, 40, 71,
72, 77, 79, 73, 76, 74, 78, 39, 38, 75, 36, 25, 32, 7,
22, 23, 26, 24, 5, 30, 28, 35, 33, 4, 20, 10
2. 31, 12, 14. 18, 15, 16, 19, 17, 13.
23,24, 24, 22, 25, 20, 19, 28, 26, 21, 18, 20, 17, 16, 15
Attn DROP 11l.am3...1.. 331-BB (WikiText-2)
31, 11, O, 14, 13, 12, 9, 8, 6, 4. 7, 1,0
27, 24, 25. 21, 26. 28, 20, 29, 30, 19, 18. 17. 16. 31
Atta DROP Mistral-7B (WikiText-2)
14, 15, 13, 9, 12, 8, 6,
1,0
25, 26, 27, 24, 23. 22, 28. 30. 31, 21, 20, 19. 18, 17, 16
Attn NBL - Mit.-MMMstra-7BBB (C4)
6, 8,8, 9, 12, 14, 13, 11, 15, 7,
25, 25, 26, 29, 28, 27, 24. 31, 22, 20, 21, 19, 18, 17,
Attn NBL Llama- 1-8B (C4)
11, 10, 15, 14, 12, 6, 8, 13, 3,0,
25, 26, 29, 27, 24, 20, 23, 28, 22, 19, 21, 31, 30, 18, 17, 16
Attn NBL. - DHt.--DD-Detill1--- (C4)
15, 11, 14, 10, 5, 13, 12, 6, 7, 0,
25, 22, 23, 26, 28, 27, 20, 19, 29, 18. 30, 17. 16, 31
Atta NBL Llama- --n-trutttIttttttt (C4)
11, 14, 10, 13, 8, 1,,
62, 65, 59, 66, 63, 61, 46 58, 55, 54, 50, 48, 53, 51, 57, 60
Attn NBL Llama 1-70B (quant.) (C4)
67, 69, 64, 68, 49, 47, 79, 70, 56, 52, 45, 72, 43, 44, 71, 42
76, 73, 74, 78, 77, 41, 75, 38, 39, 26, 25, 24, 32, 36,
12. 28, 11, 30, 22, 23, 29, 9, 8. 27, 10,
34, 35, 21, 14, 16, 15, 3, 13, 19, 18, 17, 0, 2,
25, 25, 26, 23, 24, 29, 27, 28, 20, 22, 31, 21, 19, 30, 18, 17, 16 4. 15, 10, 14, 12, 6, 13, 3,
Atta NBL Llama-3, 1-8B (WikiText-2)
27, 25, 26, 24, 22, 23, 28, 31, 30, 21, 29, 20, 19, 18, 17, 16
Attn NBL - MBs-Miit (WikiText-2)
6,8, 14, 4, 13, 12, 15, 11, 10, 7, 5, 0, 2,
36

---

H Inference complexity analysis and KV-cache usage details
This section provides a detailed explanation of the relationship between context length, prefill speed, and KV-cache usage in transformer models with NBL-applied attention layers.
H.1 Figure analysis: prefill speed-up context length
Figure illustrates how prefill speed-up improves as more attention layers are modified by NBL.. The baseline, represented by the black line, is normalized to a value of across all context lengths. Each other line corresponds to different NBL configurations, where m attention layers have been replaced with linear layers. The speed-up becomes more pronounced at longer context lengths, as the quadratic complexity (n²)) of attention begins to dominate, while the linear layers operate at a lower complexity of (nd)). At shorter contexts, the differences between configurations are minimal since the computational overhead of attention is less significant. However, as the context grows, models with more linear layers (e.g., NBL-16) 16) maintain higher prefill speeds due to the reduced cost of quadratic operations. This behavior aligns with the theoretical complexity expression:
O((K m) n² + nd)
where is the total number of attention layers, is the number of self-attention layers replaced by linear approximations through NBL, n is the context length, and d is the model embedding dimension. In this particular experiment setting to generate the figure, we used NVIDIA A100 (80GB) GPU's, and batch size of 16.
H.2 V-cache calculation for Grouped Query Attention based models
The KV-cache stores key and value tensors for each token processed by attention layers, enabling efficient incremental decoding, In models like Llama and Mistral, which use grouped-query attention, the KV-cache requirements depend on both the number of active attention layers and the context length. Grouped- attention is designed to improve the efficiency of multi-ha head attention by reducing the redundancy in key-value storage across attention heads. In standard multi-head attention. each head independently maintains its own key and value tensors, leading to a cache size proportional to the total number of attention heads. Specifically, the cache size scales as Batch Size n where n is the context length, d is the hidden dimension of the model, and the factor of 2 accounts for both keys and values.
In grouped-query attention, multiple heads are organized into groups, with each group sharing single key-value cache. Let g denote the number of groups and / the total number of attention heads and bs denote the batch size. Instead of each head having separate key-value storage, all h, h/g heads within a group use the same set of keys and values. As a result, the cache size is reduced by factor of g/ yielding the following expression for grouped-query cache size:
KV-cache size = 2 bs n d
This optimization is particularly beneficial for large-scale models, where the number of attention heads can be substantial. By reducing redundant storage, grouped-query attention helps manage memory usage, especially for longer context lengths and large batch sizes. The application of NBL further optimizes the KV-cache by modifying subset of the attention layers. When m out of K layers are replaced by linear layers, these layers longer require key-value storage. Consequently, the KV cache size is reduced to:
g m K
KV-ccche size with NBL = 2- bs d.
h
This dual optimization-grroneee grouped-query attention and NBL substantially lowers the memory requirements during inference, allowing models like Llama and Mistral to handle longer sequences more efficiently. For example, the KV cache sizes in Table 24 demonstrate the impact of these optimizations. With a batch size of 64 and grouped-query attention (32 total heads divided into groups), the cache size at a context length of 512 is 4 GB in the original configuration, reduced
37

---

to 2.5 GB with 12 layers modified by NBL. Similarly, at context length of 4096, the cache size decreases from 32 GB to 20 GB. This reduction follows the expected scaling behavior, highlighting the effectiveness of NBL and grouped-query attention in optimizing memory usage.
The table values reflect the KV cache size scaling with context length. For instance, at a context length of 512, the original cache size is GB, reduced to 2.5 GB with Attn NBL-12. At context length of 4096, the cache size decreases from 32 GB to 20 GB. These reductions demonstrate how NBL reduces both computational complexity and cache requirements in large-scale transformers.
Table 24:
KV-Cache sizes of both Llama-3: -31ma-3.1--8 and Mistral-78 models for different context lengths with batch size of 64 and half precision.
Context Len. Original (GB)
Attn NBL-4
Attn NBL-8
Attn NBL-12
Att NBL-16
(GB)
(GB)
(GB)
(GB)
512
3.5
3.0
2.5
2.0
1024
8
7.0
6.0
5.0
4.0
2048
16
14.0
12.0
10.0
8.0
4096
32
28.0
24.0
20.0
16.0
128000
1000
875.0
750.0
625.0
500.0
I Limitations and broader impacts
In this work, we present Neural Block Linearization (NBL), a novel and theoretically grounded approach for compressing LLMs without retraining. By identifying and substituting layers exhibiting high linear redundancy, NBL achieves significant inference speedups and memory savings across challenging reasoning benchmarks, including MMLU, HellaSwag, and ARC-Challenge. We derive the CCA-based bound to quantify input-output redundancy. One limitation is that certain nonlinear transformations in L.MSs may not admit low-error linear approximations. This is intrinsic property of deep architectures. Our framework explicitly accounts for such cases by ranking layers based on their approximation impact and applying substitutions only where low-error can be achieved. We further demonstrate that even local approximations result in strong end-to-end performance. The calibration data selection and local linearization considerations balance tractability and performance: we show empirically that NBL remains stable across diverse calibration sets and architectures.
Broader impacts. Looking forward, NBL offers several promising directions for extension. Apply- ing NBL to MLP sub-blocks, cooss-attention layers, or even finer-grained structures like individual attention heads may unlock additional efficiency without sacrificing quality. Moreover, NBL is particularly well-suited for long-contex, inference, where its linearized attention yields increasing benefits in both speed and memory efficiency as sequence lengths grow. From broader perspective, NBL contributes to the democratization and sustainability of AI by enabling faster, cheaper, and more private inference on edge devices and under-resourced environments. As compression becomes central to deploying foundation models at scale, techniques like NBL offer reliable and interpretable path forward, supporting both practical utility and responsible innovation.
38

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The abstract and introduction clearly state that NBL a post-training method for compressing LLMs via linear substitutions of attention layers, guided by a CCA-derived bound. This supported by the theoretical analysis and extensive empirical results across several models and benchmarks.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: Limitations discussed in Appendix I. The theoretical assumptions also clearly stated and are supported by the strong performance of our algorithm in numerical simulations. We conducted extensive experiments with different setups, constrained only by our available computational resources.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
39

---

3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification: The main theoretical result (Theorem 3.2) is formally stated and proved in Appendix C, with all assumptions explicitly listed and the derivation clearly shown.
Guidelines:
The answer NA means that the paper does not include theoretical results,
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification;
The paper provides model names, compression ratios, dataset descriptions, and detailed algorithm steps in the Experiments Section 4, and Appendix D and E. Open source datasets are used, and the code associated with the paper is also shared in the supplementary files with proper annotations.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access to hosted model (e.g., in the of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily the architecture clearly and fully.
model architecture, the paper should describe
(c) the contribution a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
40

---

(d)
We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: We provide the implementation of the proposed method in the supplementary material, accompanied by README file that outlines the structure and usage instructions.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.. cc/
public/guides/
cesiinnooliioe icyy for more details.
While we encourage the release of code and data, we understand that this might be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed run reproduce the results. See the NeurlPS code and data submission guidelines (https: / cttpss.cccpbbbiccc/cccc wetisc/ggiee//eeee http://gggiddee/eeeuuubmiiiiii ottpe///oooomammsssoooooooooeee ssioooolccy) for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We describe model selection, number of layers substituted, calibration sizes, and evaluation benchmarks. Appendix D details complexity and runtime considerations. Appendix lists the layers selected by the NBL. for replacement.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
41

---

Justification: We present the 2-D performance -etermaannn-fffiinnyy plots for the results presented in Tables 2, 3 and 4, with Standard Error intervals in Appendix E.I, supporting the statistical significance of the proposed method. We further present the numerical version of these standard error, together with the standard deviation for prefill and throughput speed-ups in Appendix E.2.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the enror bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall
with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Details are provided in the Experiments Section 4, and Appendix D and including the details on compute and runtimes.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurt httpa//ppp.p.pppppaaccacccaaaaaa
Answer: [Yes]
Justification:
We have reviewed the NeurIPS Code of Ethics and our article conforms to
them.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
42

---

The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: We discuss some of the broader impacts in Appendix Note that these impacts are not just with NBL, but generally applicable for most LLM inference optimization methods such as quantization and speculative decoding.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation. generating fake profiles, surveillance), fairesss considerations (e.g.. deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
there are negative socictal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? Answer:[NA]
Question:
Justification: No new datasets are introduced in this work, other than publicly available and open-sourced ones.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
43

---

Answer: [Yes]
Justification: All asset creators and original rights holders referenced in this paper have been appropriately acknowledged, with their respective licenses and terms of use fully adhered to. Further details are provided in Section 4 and Appendix E.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, paperswi thcoee com comddataeess has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-pane se-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is available online, the authors encouraged to reach out the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification:
Our code is the only new asset which we include in the supplemental material.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: the paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
Institutional review board (IRB) approvals or equivalent for research with human subjects
44

---

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: This paper does not involve crowd-sourcing research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with
human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that the LLM used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to our L.LM policy (https: http::///erripp.cc/ooffeeeeeeete /ot.e////eeeeeeee///25////////// for what should or should not be described.
45


Reviews:
[{"review_id": "ovtMZMraKP", "reviewer": "Reviewer_mjoZ", "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.", "questions": "Already listed in the Weaknesses", "limitations": "N/A", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths: \n\n1. The paper is generally well-written and the motivations and methods are clearly stated. Necessary theoretical analysis is provided. \n2. The authors conduct rather complete evaluation with ablation study, demonstrates the strong performance of the proposed method. \n\nWeaknesses: \n\n1. 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 \n$$\\mathrm{MSE} = \\mathrm{Tr}[C_{YY} - C_{YX} C_{XX}^{-1} C_{YX}^\\top].$$\nIt 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. \n2. Is it a typo? In the algorithm output line 11: \"Compressed LLM with linearized attention layers.\" -> linear layer actually.", "was_revised": true, "final_justification": "The CCA based linearization method itself is of interest and the authors provide rather comprehensive experimental studies.", "rebuttal": "We thank the reviewer for the feedback. Below we address the points raised.\n\n---\n## Strengths:\n> The paper is generally well-written and the motivations and methods are clearly stated. Necessary theoretical analysis is provided.\n\n> The authors conduct rather complete evaluation with ablation study, demonstrates the strong performance of the proposed method.\n\nWe appreciate the positive comments by the reviewer.\n\n---\n## Weaknesses:\n\n>  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, 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.\n\nWe appreciate the reviewer’s thoughtful question and are happy to clarify our rationale. While both the NMSE and the upper bound in Theorem 3.2 rely on second-order statistics (i.e., covariance matrices), we intentionally use the CCA-based bound instead of the direct NMSE expression for both theoretical and practical reasons.\n\nThe scalar NMSE, defined as:\n- $\\mathrm{NMSE} = \\frac{\\mathrm{Tr}[C_{YY} - C_{YX} C_{XX}^{-1} C_{YX}^\\top]}{\\mathrm{Tr}[C_{YY}]}$, \n\nwhich aggregates the approximation quality into a single number, obscures the contribution of individual directions in the output space. In contrast, the bound in Theorem 3.2 decomposes the error along canonical directions via \n- $\\mathrm{NMSE} \\leq (h_{\\text{out}} - r) + \\sum_{i=1}^r (1 - \\rho_i^2)$, \n\nwhere $\\rho_i$ are the canonical correlations.\n\nThis decomposition allows us to have a **whitened** version of the NMSE enabling the understanding of *how many* directions are well-approximated and *how strongly* each output mode aligns with the input, which NMSE does not offer. Therefore, although MSE and NMSE is a commonly used criterion in the estimation theory literature [1,2] (*and we aimed to preserve this connection through our bound*), directly using the NMSE can be biased by high-variance directions, masking poor approximations. \n\nBecause the NMSE normalizes by $\\mathrm{Tr}[C_{YY}]$, it can be dominated by a small number of high-variance output directions. Even if the approximation fails in lower-variance (but semantically important) dimensions, NMSE can still appear deceptively low. In contrast, the CCA-based bound evaluates approximation error uniformly across canonical modes, without being skewed by marginal variances.\n\nTherefore, our method is not just trying to minimize average error, but to **identify layers whose outputs are largely determined by their inputs in a linear fashion**. Canonical correlations directly quantify this, where they are the optimal linear correlations between input and output projections. This makes them an ideal indicator for substitutability. \n\nTo assess our claim, we performed an ablation where we calculated the layer rankings using both the direct NMSE and the CCA based criterion on the Mistral-7B model as the following:\n\n- NMSE values sorted from low to high: `[2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 25, 22, 26, 24, 13, 27, 14, 23, 28, 21, 16, 20, 15, 30, 31, 17, 19, 18, 29, 1, 0]`\n\n- CCA criterion values sorted from low to high: `[25, 26, 27, 24, 22, 28, 23, 30, 31, 21, 29, 20, 19, 18, 17, 16, 6, 8, 9, 12, 14, 13, 11, 15, 4, 10, 7, 5, 3, 0, 2, 1]`\n\nThen, replacing 12 attention layers with the CCA criterion in the Mistral-7B model drops the average reasoning accuracy of the model **70.2% $\\rightarrow$ 68.3%** as also indicated in Table-2 of our paper, whereas **the replacement of 12 layers with the direct NMSE causes a much larger performance drop of 70.2% $\\rightarrow$ 39.1%**, suggesting the better performance of the CCA based criterion.\n\nFurthermore, we also would like to note that while both approaches (direct NMSE vs CCA based bound) involve similar matrix operations (inversions and matrix multiplications vs. whitening and SVD), we found the CCA route more stable in practice, especially in low-data or near-singular regimes. The whitening step in CCA ensures that numerical instabilities in covariance estimation do not overly bias the error estimate.\n\nTo sum up, our framework is built around Canonical Correlation Analysis as a network analysis tool [3,4] to understand and exploit representational redundancy. \n\nReferences:\n- [1]. Steven M Kay. Fundamentals of statistical signal processing: Estimation theory, 1993.\n- [2]. Thomas Kailath, Ali H Sayed, and Babak Hassibi. Linear estimation. Prentice Hall, 2000.\n- [3]. Maithra Raghu et al. \"Svcca: Singular vector canonical correlation analysis for deep learning dynamics and interpretability\", Neurips, 2017.\n- [4]. Ari Morcos et al. \"Insights on representational similarity in neural networks with canonical correlation\", Neurips, 2018.\n\n---\n\n> Is it a typo? In the algorithm output line 11: \"Compressed LLM with linearized attention layers.\" -> linear layer actually.\n\nThank you for pointing this out. We agree that the phrasing could be improved. Our intention was to convey that the selected attention layers are replaced with their corresponding optimal linear approximations, each implemented as a single linear layer. We will revise the wording in Algorithm 1, line 11 to more clearly state: *\"Compressed LLM with selected attention layers replaced by linear layers.\"*\n\n---\nWe hope that through our comments, we have been able to address the questions raised by the reviewer. In case the reviewer has any remaining questions, we are happy to address them.\n\nBest,"}, {"review_id": "3LKwmxMoWU", "reviewer": "Reviewer_an8B", "summary": "This 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.", "questions": "- 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?", "limitations": "The 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.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths:\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", "was_revised": true, "final_justification": "The experimental results show an improvement over competitive methods, and together with the included limitations section (that was missing before to fulfill all mandatory submission requirements) and potential future work, I want to recommend this work for acceptance.", "rebuttal": "We thank the reviewer for the feedback. Below we address the points raised.\n\n## Weaknesses:\n\n> W1. 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\nNBL does not replace all attention layers. It only replaces those that can be well-approximated by a learned linear mapping, while keeping the rest—especially the ones important for modeling complex or long-range dependencies—unchanged. This allows the model to retain much of its ability to handle context and history.\n\nThe replaced layers use a token-wise linear transformation learned from real activation data. This transformation is applied independently to each token, so it doesn't average over history or remove contextual information. **Many original attention layers still remain and perform dynamic computations**.\n\nOur results (Tables 2–5) show that even after replacing multiple layers, the model still performs well on tasks that rely on history. This confirms that the remaining attention layers and the learned linear mappings together preserve important contextual reasoning.\n\n> W2. 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\nWhile it is true that full replacement of all attention layers would limit the model’s expressiveness, this scenario is **not aligned with the goals or design of our methodology**. NBL is explicitly built around **selective replacement**, using canonical correlation scores to retain layers that are not linearly predictable and likely essential for modeling long-range or nonlinear dependencies.\n\nMoreover, even the substituted layers do not behave like static averaging as they apply token-wise learned linear transformations, fitted to real activation data. These retain per-token contextual information without collapsing the output into a local or averaged form.\n\nTherefore, *although full replacement is theoretically possible, it is neither practiced nor advocated in our method*. Our empirical results demonstrate that **selectively replacing even a few layers yields efficiency gains without collapsing into shallow behavior**, preserving reasoning ability across a wide range of tasks as evidenced by our strong empirical results, consistently outperforming recent pruning baselines in maintaining performance.\n\n> W3. Calibration dataset needs to cover the whole inference distribution - acts similar as a training dataset\n\nThank you for the comment. We would like to clarify that the calibration dataset does not act like a training dataset. **Our method does not involve gradient computation, parameter updates, or task-specific supervision**. Instead, our method performs a lightweight, unsupervised forward pass over 256 samples to estimate CCA-based linear approximability. This process is scale-invariant, task-agnostic, and computationally efficient.\n\nAs shown in our experiment results in Tables 2–5, we calibrate using a generic corpus (C4 dataset) yet evaluate on reasoning-heavy tasks like ARC and MMLU, which have distinct distributional properties. The strong performance in these settings indicates strong robustness to calibration–inference mismatch.\n\nTo further analyze this, we conducted an explicit ablation (Appendix F.1) where we calibrate on either C4 or WikiText-2, then evaluate on both datasets. Below, we summarize the perplexity results from Tables 14 and 15:\n\n|Calibration Dataset|Method|C4↓(Llama-3.1-8B)|C4↓(Mistral-7B)|WikiText-2↓(Llama-3.1-8B)|WikiText-2↓(Mistral-7B)|\n|-|-|-|-|-|-|\n|WikiText-2|SliceGPT-25%|123.29|28.28|14.27|7.61|\n||SLEB-8|20.53|15.07|16.70|10.23|\n||Attn DROP-8|11.40|9.20|7.39|6.81|\n||**Attn NBL-8**|**12.22**|**9.85**|**7.18**|**5.59**|\n|C4|SliceGPT-25%|25.41|13.07|74.33|12.11|\n||SLEB-8|20.54|15.49|16.67|11.13|\n||Attn DROP-8|11.35|9.20|7.48|6.81|\n||**Attn NBL-8**|**11.37**|**9.08**|**9.57**|**7.33**|\n\nAs shown above, Attn NBL exhibits minimal degradation across calibration sources (e.g., C4 eval: 12.22 → 11.37), in contrast to SliceGPT, which suffers from large mismatches (e.g., WikiText-2 eval: 14.27 → 74.33). This **confirms that NBL is robust to moderate calibration–inference distribution shifts**. In short, calibration is computationally efficient, task-agnostic, and significantly lighter than training, both in cost and in assumptions.\n\n> W4. missing limitations section.\n\nWe would like to clarify that we **do include a limitations section in Appendix-I** of the paper. There, we explicitly discuss several key assumptions of our method. We will make this section more detailed and more visible in the revised version by moving it to the main text.\n\n---\n## Questions:\n> Q1. Why is the limitation that applying NBL to attention blocks removes the history not covered in your paper?\n\nWe respectfully disagree that NBL removes access to history, as our method does not replace all attention layers. Instead, it applies selective substitution, where only the layers whose outputs are well-approximated linearly (as measured by CCA scores) are replaced. Layers that are critical for modeling long-range dependencies or complex reasoning are retained. This data-driven approach ensures that the model preserves its ability to attend over extended context when necessary.\n\nThe strong performance of our method on tasks requiring multi-hop or long-context reasoning (e.g., ARC, MMLU) in Tables 2–5 supports this point. Nonetheless, we agree that if NBL were applied indiscriminately or too aggressively — especially replacing early attention layers — it could, in the limit, reduce the model's effective receptive field, resembling a bi-gram-like architecture that lacks full history access.\n\nWhile we already mention a related limitation in Appendix-I (*\"One limitation is that certain nonlinear transformations in LLMs may not admit low-error linear approximations.\"*), we agree this point deserves more emphasis. We will update the limitations and also add a discussion in the conclusion of the main text to clarify that:\n\n- *\"With NBL, we currently rely on cautious and selective layer replacement based on linear approximability. However, aggressive substitution across all attention layers may reduce the model’s ability to capture the history and long-range dependencies. Future work will explore ways to mitigate this, such as applying NBL at the head level or retaining critical heads to preserve contextual reasoning, thereby achieving a finer tradeoff between efficiency and capability.\"*\n\n> Q2. 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.\n\nThis is an insightful direction. Analyzing the QKV structure directly, for example, through probing or synthetic inputs, could provide a layerwise estimate of sensitivity or redundancy. However, **such methods may miss important correlations or positional biases present in real data**. Our Canonical Correlation Anslysis (CCA)-based method captures data informed structure, which is critical to preserving fidelity. This situation is existent in may prior works, such as SLEB, SliceGPT and DROP which utilizes calibration dataset to accurately model the data distribution while compressing LLMs.\n\nThat said, we agree that hybrid approaches combining theoretical analysis of QKV structure with data-driven calibration could further improve robustness and generalization beyond the calibration distribution. We consider this a promising area for future work. We will add this point to the conclusion section as a potential direction for future work.\n\n- *\"An interesting direction for future work is to combine theoretical analysis of QKV structure with data-driven calibration, enabling layer substitution decisions that generalize beyond the calibration distribution and reduce reliance on task-specific data.\"*\n\n> Q3. 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.\n\nThank you for the suggestion. While norm-based diagnostics may offer interpretability, our method is not driven by absolute magnitudes. Instead, we use **Normalized Mean Squared Error (NMSE)** to evaluate how well the output of each attention layer can be linearly approximated **- independent of scale**. Therefore, raw output or residual norms alone do not determine which layers are substituted. Importantly, **our focus is on preserving functional computation paths, not on the magnitude of activations or weights.**\n\nBut, we investigated the ratio $\\|Y\\| / \\|X + Y\\|$ of attention layer output's ($Y$) norm to the residual's ($X + Y$) norm, in DeepSeek-8B. However, this analysis did not reveal any meaningful correlation with the layers chosen by our method. The top-4 layers ranked by $\\|Y\\| / \\|X + Y\\|$ are:\n\n- Layer indices with highest $\\|Y\\| / \\|X + Y\\|$: `[0, 1, 31, 30]`  \n\nThis further supports our choice to use **NMSE**, which is scale-invariant and more directly reflects functional compressibility. Parameter norms can also be analyzed, but they are entangled with output norms, being hard to interpret.  That said, **if the reviewer has further details for a more specific diagnostic or hypothesis in mind regarding the role of vector norms, we would be happy to explore it during the discussion period.**\n\n---\n## Limitations\n> L1. The paper is missing a limitations section ... the paper does not cover the limitation of its method ...\n\nWe already included a limiations section in Appendix-I. However, we will extend this and move it to the main text as pointed by the reviewer. Please see Weakness-4 and Question-1 for further details.\n\n---\nWe hope our responses address the points raised by the reviewer and are happy to clarify any remaining.\n\nBest,"}, {"review_id": "jOCMbRcfF8", "reviewer": "Reviewer_8B3g", "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.", "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?\n2. For tasks that require long-distance dependencies, will this linearization cause a \"cliff-like\" drop in model performance?", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n1. This paper is well-written and easy to follow. \n2. NBL uses the Canonical Correlation Analysis (CCA) to guide the transformation of attention layers, building on a good theoretical foundation. \n3. Experimental results show the effectiveness of NBL on reasoning tasks.\n\nWeaknesses:\n1. Computing the linear approximations requires a longer time than other pruning methods (e.g., SliceGPT).\n2. Some typos: \n- Inconsistent method name between “Attn SLEB-m” in Line 218 and ‘’SLEB-m’’ in Table 2, 3.\n- 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$.", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the feedback. Below we address the points raised.\n\n---\n## Strengths:\n\n> This paper is well-written and easy to follow.\n\n> NBL uses the Canonical Correlation Analysis (CCA) to guide the transformation of attention layers, building on a good theoretical foundation.\n\n> Experimental results show the effectiveness of NBL on reasoning tasks.\n\nWe appreciate the positive comments by the reviewer.\n\n---\n## Weaknesses:\n\n> W1. Computing the linear approximations requires a longer time than other pruning methods (e.g., SliceGPT).\n\nThank you for raising this point. It is true that our calibration step is somewhat more computationally intensive than simpler methods like DROP, primarily because it involves computing second-order statistics (covariance matrices) and performing SVD. However, our method takes comparable time compared to SliceGPT. As shown in our Table 1, pruning an 8B Llama-3 model takes under 30 minutes on an A100 GPU. In comparison, SliceGPT reports a runtime of 44 minutes to compress a Llama-2 7B model on a more powerful H100 GPU (Table 3 of their paper).\n\nHowever, we emphasize that this step is performed **only once, offline**, and **does not require gradient-based optimization or fine-tuning**. As shown in Table 1, the wall-clock time remains practical even for 405B sized models (e.g., under 30 minutes for 8B models, in around 13 hours for 405B models). \n\nImportantly, our method yields more robust accuracy retention across tasks (see Tables 2–5), particularly in the mid-to-high compression regimes where pruning-based approaches like SliceGPT degrade more rapidly. We believe this slight increase in calibration cost is justified by the **better trade-off between efficiency and accuracy**, and by the fact that NBL can be applied plug-and-play to any pre-trained model without retraining.\n\n---\n> W2. Some typos: \n> - Inconsistent method name between “Attn SLEB-m” in Line 218 and “SLEB-m” in Table 2, 3. \n> - As the shape of $X$ is $\\mathbb{R}^{(s \\cdot t)} \\times h_{\\text{in}}$ in Equation (1), Equation (2) should be $\\hat{Y} = XW + b$.\n\nThank you for pointing out these. There was a typo in Line 218 as SLEB acts on te whole transformer block, therefore all the methods that refers to SLEB are the same and will be changed to “SLEB-m”. Furthermore, we will revise the definition of $X$ and $Y$ in Line 127 to be the transpose of the current version, and become $h_{\\text{in}} \\times \\mathbb{R}^{(s \\cdot t)}$ and $h_{\\text{out}} \\times \\mathbb{R}^{(s \\cdot t)}$ respectively. This way, the Equation (2) will correctly match the matrix dimensions, and be consistent with the rest of the paper.\n\n---\n## Questions:\n\n> Q1. 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?\n\nThank you for the insightful question. It's important to emphasize that the purpose of the calibration set in NBL is not to emulate full-context inference, but to assess the **local linearity and redundancy structure** of each attention layer. Since our method substitutes layers based on the canonical correlation between inputs and outputs, rather than their absolute sequence-level behavior, very long contexts were not required during calibration.\n\nIn practice, we use sequences of 2048 tokens, which already cover a wide range of attention behaviors for linear substitution. The selected calibration length aligns with the default context range for the evaluated models (e.g., Mistral, LLaMA-3.1, DeepSeek), and suffices to capture the essential token-wise interaction patterns for reasoning tasks. \n\n---\n> Q2. For tasks that require long-distance dependencies, will this linearization cause a \"cliff-like\" drop in model performance?\n\nThank you for the question. In our evaluations, we did not observe a cliff-like collapse in performance. Instead, the degradation is generally gradual and varies by task complexity. Some benchmarks, particularly those relying on shallow reasoning or factual recall (e.g., BoolQ, HellaSwag), are more resilient to linearization. Others, such as ARC-e or multi-hop MMLU categories, can be more sensitive depending on which layers are replaced. This variability reinforces the value of our CCA-based selection criterion, which helps identify layers whose output structure is already close to linear and less critical for modeling long-range dependencies. \n\nThat said, as context length increases, **the computational cost of attention dominates**, and even **replacing a small number of attention layers** yields substantial gains. This is evidenced by Figure 3, which shows growing speedups at longer contexts even when only a few layers are substituted, and Table-21 that suggests huge KV-cache savings in longer contexts. Thus, for tasks that require long-range reasoning, **a conservative application of NBL still provides meaningful efficiency improvements without sacrificing core performance**.\n\nTo evaluate robustness under such conditions, we report results on the NVIDIA RULER benchmark (Hsieh et al., 2024), which consists of attention-heavy tasks requiring reasoning, using DeepSeek-R1-Distill-Llama-8B model. RULER is specifically designed to test a model’s ability to track and integrate long-range dependencies, making it a suitable stress test for compression methods like NBL. Here our average accuracies over the RULER benchmarks:\n\n| Method   | Avg. Accuracy (%) |\n| --------------- | ------------ |\n| **Uncompressed Baseline**  | 90.02        |\n| **SLEB-2**      | 69.78        |\n| **SLEB-3**      | 61.29        |\n| **Attn DROP-2** | 87.78        |\n| **Attn DROP-3**| 86.66        |\n| **Attn DROP-4** | 83.15        |\n| **Attn NBL-2 (ours)**  | **88.59**    |\n| **Attn NBL-3 (ours)**  | **88.21**    |\n| **Attn NBL-4 (ours)**  | **86.19**    |\n\nThese results show that Attn NBL retains over 95% of full accuracy even with 2–4 attention layers linearized, whereas SLEB degrades much more rapidly. This demonstrates that NBL can selectively reduce computational cost in attention-heavy settings with better  performance compared to recent baselines on tasks requiring long-context understanding.\n\n---\nWe hope that through our comments, we have been able to address the questions raised by the reviewer. In case the reviewer has any remaining, we are happy to address them.\n\nBest,"}, {"review_id": "x0u9N8UHIB", "reviewer": "Reviewer_ZrVu", "summary": "Neural 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.", "questions": "Can 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)?", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 4, "significance": 2, "originality": 2, "strengths_and_weaknesses": "The 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.", "was_revised": true, "final_justification": "I feel that NBL may be interesting to the community, but am unconvinced by the argument for CCA over cosine distance.\n\nRESOLVED: \n* including additional baselines\n* including explicit sparsity levels\n\nUNRESOLVED: \n* evidence showing superiority of CCA over cosine distance", "rebuttal": "We thank the reviewer for the feedback. Below we address the points raised.\n\n---\n## Strengths:\n\n> This paper is well-written and easy to follow.\n\n> NBL uses CCA to guide the transformation of attention layers, building on a good theoretical foundation.\n\n> Experimental results show the effectiveness of NBL on reasoning tasks.\n\nWe appreciate the positive comments by the reviewer.\n\n---\n## Weaknesses:\n> W1. The evaluation is limited as some structural pruning baselines are missing.\n\nThank you for the suggestion. We would like to clarify that Tables 2–5 already include evaluations against a strong structural pruning baseline, DROP (He et al. \"What Matters in Transformers? Not All Attention is Needed\", 2024) which is one of the most recent and competitive pruning methods. And for Mistral-7B and Llama-3.1-8B, also SliceGPT and SLEB. These represent some of the most recent and competitive attention-pruning techniques.\n\nPreviously:\n- Tables 2 and 3 evaluate NBL against DROP, SliceGPT, and SLEB on Mistral-7B and Llama-3.1-8B, demonstrating NBL's favorable trade-off between accuracy and speed.\n-  Table 4 extends this to DeepSeek-R1-Distill-Llama-8B, a newly released model, again comparing with DROP. \n-  Table 5 explores scalability by applying NBL to a quantized Llama-3.1-70B model, where we aimed to show feasibility on even larger models.\n\nHowever, we appreciate this observation and agree that broader baseline coverage is useful as time permitted. In response to this feedback, we extended Table 4 to include PruneNET, SliceGPT, and SLEB, along with existing DROP and NBL entries:\n\n### Extended Table-4: Deepseek-8B:\n|Method|Avg.Acc.(↑)|Prefill(↑)|Throughput(↑)|\n|-|-|-|-|\n|Baseline|63.6|1.00|1.00|\n|PruneNET-15%|53.7|1.08|1.08|\n|PruneNET-25%|46.4|1.16|1.15|\n|PruneNET-35%|40.8|1.22|1.19|\n|SliceGPT-15%|51.2|1.09|1.09|\n|SliceGPT-25%|46.3|1.15|1.13|\n|SliceGPT-35%|40.6|1.21|1.20|\n|SLEB-4(12.5%)|56.2|1.09|1.14|\n|SLEB-8(25%)|47.7|1.17|1.29|\n|SLEB-12(37.5%)|42.4|1.26|1.53|\n|Attn DROP-12|60.7|1.27|1.34|\n|Attn DROP-16|50.2|1.37|1.47|\n|Attn NBL-12|__62.4__|1.24|1.32|\n|Attn NBL-16|56.8|1.35|1.42|\n\nWe note that PruneNET required parameter tuning for DeepSeek. The default learning rate (5e-4) for the action model training mentioned in manuscript led to poor convergence, and sweeping between 1e-5 and 1e-3, we found that 1e-4 yielded better results, but performing marginally better than SliceGPT.\n\nThrough these results, we see that Attn NBL gives the least drop in accuracy, while significantly increasing the inference speed, suggesting a better trade-off compared to baselines. In the revised manuscript, we will also extend Tables 2 and 3 with the PruneNET evaluation for consistency.\n\nFor Table 5, which evaluates Llama-3.1-70B quantized to 4-bit using AWQ, we could not include PruneNET or SliceGPT as **these methods do not currently support post-quantization pruning**. However, we have now included SLEB in addition to DROP and our method NBL:\n\n### Extended Table-5: Llama-3.1-70B (quant.):\n|Method|Avg.Acc.(↑)|Prefill(↑)|Throughput(↑)|\n|-|-|-|-|\n|Baseline(quant.)|76.0|1.00|1.00|\n|SLEB-16|70.9|1.23|1.24|\n|SLEB-32|55.9|1.63|1.61|\n|Attn DROP-32|75.7|1.20|1.30|\n|Attn DROP-48|70.4|1.29|1.42|\n|Attn DROP-54|48.3|1.34|1.53|\n|Attn NBL-32|76.0|1.16|1.24|\n|Attn NBL-48|71.3|1.25|1.35|\n|Attn NBL-54|65.4|1.30|1.44|\n\nAgain, both Attn NBL-32 and Attn NBL-48 outperform SLEB-16 in accuracy while matching or exceeding its inference throughput, with significantly lower KV-cache usage due to more attentions removed. These results further highlight that NBL consistently achieves a better accuracy/efficiency trade-off, even under aggressive compression settings.\n\n> W2. The levels of sparsity are unclear for the given methods in the tables besides SliceGPT.\n\nThank you for this suggestion. While we agree that sparsity is an important comparative metric, our primary focus in this work is not to maximize sparsity, but rather to improve inference efficiency through the selective substitution of attention layers. Therefore, our main criteria while designing our method is to achieve a better \"*accuracy / inference-speed*\" or \"*acuracy / KV-cache-reduction*\" trade-offs.\n\nThat said, we understand that some readers may be interested in comparing methods from the perspective of accuracy vs sparsity (**even though this is not the main objective of our paper**). In our main text, we already report the number of layers replaced or dropped, which corresponds directly to sparsity in block-based SLEB and Block DROP methods (e.g., dropping 8 blocks out of 32 gives 25% sparsity). For Attn NBL, we replace attention blocks with lightweight linear layers. These replacements preserve the interface but eliminate the computation of QKV projections and attention maps, significantly reducing FLOPs and runtime.\n\nIn response to the reviewer’s comment, we have added explicit sparsity ratios for all methods in Table 2,3,4, along with their average reasoning accuracy, together with a wider discussion of the advantages and disadvantages of each method. Note that in all methods, the embedding and final projection layers remain intact, which comprise 1.81% and 6.54% of total parameters in Mistral-7B and Llama-3.1-8B, respectively. Therefore, we report sparsity only over the remaining layers. \n\nFor Tables 2 (Mistral-7B) and 4 (DeepSeek-8B) the sparsity ratios are the same as Llama-3.1-8B, since the only differences lie in the size of the embedding and final projection layers.\n\n### Table-3 extended with sparsities: Llama-3.1-8B:\n|**Method**|**Sparsity(%)**|**Avg. Accuracy(↑)**|\n|-|-|-|\n|**Baseline**|0|70.2|\n|SliceGPT-15%|15.0|58.8|\n|SliceGPT-25%|25.0|49.9|\n|SliceGPT-35%|35.0|43.4|\n|SLEB-4|12.5|59.7|\n|SLEB-8|25.0|46.5|\n|SLEB-12|37.5|42.5|\n|Block DROP-4|12.5|64.3|\n|Block DROP-8|25.0|40.2|\n|Block DROP-12|37.5|43.9|\n|**Block NBL-4**|11.5|67.2|\n|**Block NBL-8**|23.1|59.3|\n|**Block NBL-12**|34.6|50.4|\n|Attn DROP-4|2.4|__70.2__|\n|Attn DROP-8|4.8|__69.8__|\n|Attn DROP-12|7.2|67.9|\n|Attn DROP-16|9.6|55.2|\n|**Attn NBL-4**|1.5|__70.2__|\n|**Attn NBL-8**|2.9|__69.8__|\n|**Attn NBL-12**|4.3|__69.1__|\n|**Attn NBL-16**|5.8|62.5|\n\nAs shown in table, operating only on attention layers naturally result in lower sparsity as the majority of parameters (~80%) reside in the MLP layers. However, attention layers are significantly more computationally expensive due to their $O(N^2)$ complexity, making them ideal targets for improving inference efficiency in practice.\n \nHowever, our proposed **Block NBL method achieves a notably strong accuracy/sparsity trade-off**. For instance, **Block NBL-12 yields 50.4 accuracy with 34.6% sparsity, whereas methods like Block DROP-8 gives 40.2%, SLEB-8 has 46.5% and SliceGPT-25% has 49.9% accuracies each with a lower 25% sparsity**. This highlights the effectiveness of our selective substitution strategy, particularly in maintaining accuracy\n\nMoreover, although inducing less sparsity by acting exclusively on attention layers, Attn NBL provides the best accuracy–throughput trade-off: achieving less than 1% accuracy drop while improving throughput by over 32% on DeepSeek, also suggesting pareto-optimality (Appendix-E.3). This demonstrates that even minimal modifications, when strategically applied to computational bottlenecks, can lead to substantial efficiency gains without sacrificing accuracy.\n\n> W3. The results from the ablation on CCA-Bound criterion vs cosine distance, as well as for greedy selection are ... no uncertainty estimates.\n\nWe include the uncertainty estimations for the ablation with cosine distance criterion, as in Tables 17 and 18:\n\nMistral-7B:\n|Method|Cosine Dist.Avg.Acc.±SE (%)|CCA Avg.Acc.±SE(%)|\n|-|-|-|\n|Attn NBL-4|70.1±0.40|70.1±0.41|\n|Attn NBL-8|69.9±0.40|70.0±0.41|\n|Attn NBL-12|67.7±0.42|68.3±0.42|\n|Attn NBL-16|58.8±0.40|58.8±0.42|\n\nLlama-3.1-8B:\n|Method|Cosine Dist.Avg.Acc.±SE (%)|CCA Avg.Acc.±SE(%)|\n|-|-|-|\n|Attn NBL-4|70.2±0.43|70.2±0.41|\n|Attn NBL-8|70.0±0.41|69.8±0.41|\n|Attn NBL-12|69.0±0.41|69.1±0.42|\n|Attn NBL-16|58.0±0.42|62.5±0.43|\n\nWhile the results are mostly similar, we observe statistically significant improvements in specific settings. For example, on Mistral-7B, Attn NBL-12 with the CCA criterion outperforms cosine distance by 1.5× the standard error. On LLaMA-3.1-8B, the gap is even more pronounced for Attn NBL-16, where CCA yields a 10× SE improvement. \n\nThese results suggest that the CCA-based selection is more stable and effective, particularly in deeper linearization settings. Furthermore, although the accuracy results of NBL versus the greedy selection are similar, NBL is preferred as it is computationally more efficient by eliminating the iterative procedure.\n\n---\n## Questions:\n\n> Q1. Can you report evaluations of SliceGPT and SLEB for Tables 4-5, as well as PruneNet? More comparisons will make me more inclined to raise the score.\n\nPlease refer Weakness-1, for this matter.\n\n> Q2. Can you report the level of sparsity for SLEB/Block/Attn Drop/NBL-m, similar to SliceGPT?\n\nPlease refer Weakness-2 for this matter.\n\n> Ablations on replacing the feedforward layers? Or, explain why the only ablations are on dropping attention layers or Transformer blocks?\n\nWe chose not to include ablations on dropping only the feedforward (MLP) layers because prior work, specifically the DROP paper has shown that modifying or removing MLP layers not only leads to a significant degradation in reasoning performance but also yields negligible improvements in inference speed, particularly in terms of KV-cache efficiency. In contrast, targeting the $O(N^2)$ attention layers—either individually (Attention-NBL) or jointly with MLPs via full block removal (Block-NBL)—has demonstrated clear and consistent benefits across both accuracy and speed. Therefore, we focused our experiments on these interventions that we believe to be more impactful.\n\n---\nWe hope our responses address the reviewer’s questions and are happy to clarify any remaining points.\n\nBest,"}]

Author rebuttal:

Meta-review: This paper proposes a technique for linearizing self-attention layers in Transformers to achieve faster inference, using a linear minimum mean squared error estimator. The reviewers felt that the approach was novel, with strong theoretical backing and demonstrated wall-clock speedup, and solid experiments. During the rebuttal period, the authors responded well to reviewer questions and reported additional experiments that made a stronger case for the method.

I recommend accept as a poster.
Decision: Accept (poster)