StreamKL让注意力蒸馏快14倍且内存降为常数级
StreamKL: Fast and Memory-Efficient KL Divergence for Boosting Attention Distillation
- 用单次遍历流式计算注意力分布的KL散度，避免存储全部中间结果
- 前向加速43倍，反向加速14倍，内存占用从二次方降为常数
- 适合长序列注意力蒸馏，可在单卡上运行超长上下文模型
注意力蒸馏通过最小化两个注意力分布之间的Kullback-Leibler（KL）散度来训练模型，广泛应用于知识蒸馏、模型压缩、持续学习和稀疏注意力大模型训练。然而，现有方法在计算KL散度前需完整存储两个注意力分布，导致$O(N_QN_K)$的内存和IO开销，在长上下文场景下难以承受。本文提出StreamKL，首个用于注意力KL散度计算的融合GPU原语，消除了这一二次方的存储需求。StreamKL推导出一种新型在线公式，实现一次遍历的前向核，将查询-键块流式传输至片上SRAM。反向传播时，按块重新计算注意力概率，避免存储二次方中间结果。我们进一步设计并实现了高效的GPU内核及专用优化。实验表明，StreamKL在前向和反向传递中分别实现最高43倍和14倍的加速。更重要的是，它将注意力蒸馏的额外显存占用从$O(N_QN_K)$降至$O(1)$，使单卡支持长上下文蒸馏成为可能。
原文摘要 · Abstract (English)
Attention distillation, which trains one attention distribution to match another by minimizing their Kullback-Leibler (KL) divergence, is widely used in knowledge distillation, model compression, continual learning, and sparse-attention LLM training. However, existing approaches materialize both attention distributions before computing the KL reduction, incurring $O(N_QN_K)$ memory and IO costs that become prohibitive at long context lengths. We present StreamKL, the first fused GPU primitive for attention KL divergence that eliminates this quadratic materialization. StreamKL derives a novel online formulation for the coupled two-distribution KL reduction, enabling a single one-pass forward kernel that streams query-key tiles through on-chip SRAM. For the backward pass, StreamKL recomputes attention probabilities tile-by-tile, avoiding storage of quadratic intermediates. We further design and implement efficient GPU kernels with dedicated optimizations. Experiments show StreamKL delivers up to $43\times$ and $14\times$ speedups over baseline methods in the forward and backward passes, respectively. Most importantly, StreamKL reduces the extra HBM footprint of attention distillation from $O(N_QN_K)$ to $O(1)$, enabling long-context distillation on a single GPU.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。