File size: 2,653 Bytes
d73500e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 | # 🧠 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)
|