QKV Projections Require a Fraction of Their Memory¶

会议: ICLR 2026
arXiv: 2506.02939
代码: 无
领域: 模型压缩
关键词: 训练内存压缩, 注意力机制, 矩阵乘法近似, 激活压缩, LLM训练

一句话总结¶

提出 PAMM（Point-Approximate Matrix Multiplication），一种激活压缩技术，通过随机选取少量代表性 token 来近似 QKV 投影层激活，实现高达 512× 压缩率且不影响模型性能。

LLM 训练中，注意力层的 QKV 投影占用大量内存：输入 \(X\) 需要在前向过程中保存以用于反向传播（计算 \(\nabla W = X^\top \cdot \nabla Z\)）。这部分内存可占注意力块总峰值 GPU 内存的 20%。

现有内存优化方法的不足： - 高效注意力（FlashAttention 等）：优化缩放点积本身，未涉及线性投影 - 低秩方法（CompAct 等）：沿隐藏维度压缩，但序列维度的冗余更大 - 优化器状态压缩：不随 batch size 和序列长度扩展

核心洞察：序列维度存在巨大冗余。训练 batch 中的 token 数量 \(b = BL\)（如 16384）远大于隐藏维度 \(n\)（如 2048），\(\text{rank}(X) \leq n\)，理论上仅需 \(n\) 个基向量即可表示 \(X\)，压缩比可达 8×。

PAMM 分两阶段工作：(1) 前向时将 \(X\) 压缩为少量生成点和辅助信息；(2) 反向时用压缩表示近似计算梯度 \(\nabla W\)。

激活压缩 (Compression Stage)：
- 从 \(X \in \mathbb{R}^{b \times n}\) 中随机采样 \(k = r \cdot b\) 行作为生成点 \(C \in \mathbb{R}^{k \times n}\)
- 对每个点 \(A_i\)，选择最佳生成点：\(f(i) = \arg\max_j |\text{csim}(A_i, C_j)|\)（Lemma 1）
- 计算缩放系数：\(\tilde{A}_i = \alpha(i, f(i)) \cdot C_{f(i)}\)，其中 \(\alpha = \frac{\langle A_i, C_j \rangle}{\|C_j\|_2^2}\)
- 邻域条件：\(\|A_i - \tilde{A}_i\|_2 \leq \varepsilon \|A_i\|_2\)，不满足则丢弃
近似矩阵乘法 (Approximate Multiplication)：
- 不重建完整 \(\tilde{A}\)，而是先聚合 \(\tilde{B}_j = \sum_{i:f(i)=j} \alpha_i B_i\)
- 计算 \(\tilde{O} = C^\top \tilde{B}\)，维度从 \(b \times n\) 降为 \(k \times n\)
- 引入归一化因子 \(\beta = \frac{b}{b-\eta}\) 保证无偏估计 \(\mathbb{E}[\tilde{O}] = O\)
理论保证：
- Lemma 2（\(k\) 的充分条件）：\(k > \frac{b}{n_{\min}} \ln(\frac{b}{\delta})\)，仅需对数级别的生成点
- 近似误差上界：\(\|O - \tilde{O}\|_F^2 \leq \|B\|_2^2 (\varepsilon^2 \|A_\mathcal{I}\|_F^2 + \|A_{\bar{\mathcal{I}}}\|_F^2)\)
- 实践中 \(\varepsilon \to \infty\)（不使用邻域约束）效果最好

模型	PAMM r	验证 PPL	QKV 内存 (MB)	内存减少
LLaMA-60M	无 PAMM	31.8	432	-
LLaMA-60M	1/512	31.6	0.85	>99%
LLaMA-350M	无 PAMM	18.7	1,296	-
LLaMA-350M	1/512	18.5	2.53	>99%
LLaMA-1B	无 PAMM	15.1	2,592	-
LLaMA-1B	1/512	15.0	5.06	>99%

方法	QKV 内存 (MB)	GLUE 平均	内存减少
Full Fine-Tuning	288	86.28	-
PAMM r=1/128	6.75	86.11	97.7%
PAMM r=1/256	3.37	86.18	98.8%

阶段	基线 (tok/s)	PAMM (tok/s)	吞吐量降低
前向	247.6K	235.4K	4.92%
反向	141.9K	138.3K	2.53%
总计	88.4K	85.2K	3.61%