# 🧠 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）

```python
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)