🧠 Efficient Large Language Model Inference with Neural Block Linearization (NBL)

Paper: Mete Erdogan, Francesco Tonin, Volkan Cevher, EPFL, 2025
arXiv: 2505.21077v1
Goal: 无需训练的 LLM 推理加速框架，通过线性近似替换部分 Attention 层。

🚀 1. Motivation

Transformer 的 Multi-Head Self-Attention (MHSA) 在推理阶段具有二次复杂度： [ O(n^2 d) ] 其中 (n) 为上下文长度，(d) 为隐藏维度。

NBL 的核心思想：

某些注意力层在输入输出关系上高度线性，可用一次线性变换近似。
因此可用一个线性层（矩阵乘 + 偏置加）替代整个 attention。

🧩 2. Method Overview

Step 1. 核心替换公式

对原 Attention 层： [ Y = \text{softmax}\left(\frac{QK^\top}{\sqrt{d_k}}\right)V ] 其中 ( Q = XW_Q, K = XW_K, V = XW_V )。

NBL 用线性层替代： [ \hat{Y} = W X + b ] 其中 (W, b) 通过 线性最小均方误差 (LMMSE) 闭式求解。

Step 2. 数据采样（Calibration Data）

采样少量无标注文本序列（如来自 C4 或 WikiText-2）；
通常取 256 条文本，长度约 128–512；
对每层记录输入输出激活： [ X_k^{(i)}, Y_k^{(i)} \quad \text{for } i=1,\dots,s ]
叠加形成： [ X \in \mathbb{R}^{(s\cdot t)\times d_{in}}, \quad Y \in \mathbb{R}^{(s\cdot t)\times d_{out}} ]

Step 3. LMMSE 求解线性映射

最小化均方误差： [ \min_{W,b} \mathbb{E}|Y - WX - b|_2^2 ]

闭式解为： [ W = C_{YX} C_{XX}^{-1}, \quad b = \mathbb{E}[Y] - W\mathbb{E}[X] ]

其中：

(C_{XX})：输入协方差；
(C_{YX})：输入输出互协方差。

Step 4. 层选择标准（CCA-Based Criterion）

不线性化所有层，而是选择“最线性”的层。

Canonical Correlation Analysis (CCA)

对每层输入输出计算典型相关系数： [ \rho_1, \rho_2, ..., \rho_r ] 衡量输入输出子空间线性对齐程度。

线性化误差上界（定理 3.2）

[ \text{NMSE}(Y, \hat{Y}) \le (h_{out}-r) + \sum_{i=1}^r (1 - \rho_i^2) ]

因此：

(\rho_i \approx 1) → 该层近似线性；
NMSE 越小 → 越适合替换。

选择 NMSE 最低的 m 层作为替换目标。

Step 5. 替换算法（Algorithm 1）

for each attention layer Ak:
    collect (Xk, Yk) from calibration data
    compute canonical correlations rho_k
    estimate NMSE_k = sum(1 - rho_k^2)

select m layers with lowest NMSE_k

for each selected layer Aj:
    compute Wj = CYX * inv(CXX)
    compute bj = mean(Y) - Wj * mean(X)
    replace Aj with Linear(Wj, bj)