論文の概要: StreamKL: Fast and Memory-Efficient KL Divergence for Boosting Attention Distillation
- arxiv url: http://arxiv.org/abs/2606.20005v1
- Date: Thu, 18 Jun 2026 09:40:38 GMT
- ステータス: 翻訳完了
- システム内更新日: 2026-06-19 18:23:39.772918
- Title: StreamKL: Fast and Memory-Efficient KL Divergence for Boosting Attention Distillation
- Title(参考訳): StreamKL: 注意蒸留を増強するための高速かつメモリ効率の良いKL分岐
- Authors: Guangda Liu, Yiquan Wang, Chengwei Li, Wenhao Chen, Jing Lin, Yiwu Yao, Danning Ke, Wenchao Ding, Jieru Zhao,
- Abstract要約: 留置蒸留は、KL(Kulback-Leibler)の分散を最小化して、互いに一致するように注意を向ける。
既存のアプローチはKLの削減を計算する前に注意を喚起し、O(N_QN_K)$メモリとIOコストが長期にわたって禁止される。
本稿では,この二次的な物質化を排除した最初の融合GPUプリミティブであるStreamKLを紹介する。
- 参考スコア(独自算出の注目度): 9.82608878263804
- License: http://creativecommons.org/licenses/by/4.0/
- Abstract: 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.
- Abstract(参考訳): Kullback-Leibler (KL) の発散を最小限に抑えて, 注意分布を訓練する注意蒸留は, 知識蒸留, モデル圧縮, 連続学習, 疎注意LDM訓練に広く用いられている。
しかし、既存のアプローチはKLの削減を計算する前に注意分布を具体化しており、長いコンテキストで禁止となるメモリとIOコストが$O(N_QN_K)となる。
本稿では,この二次的な物質化を排除した最初の融合GPUプリミティブであるStreamKLを紹介する。
StreamKLは、結合された2分散KL削減のための新しいオンライン定式化を導き、オンチップSRAMを介してクエリキータイルをストリームする単一のワンパスフォワードカーネルを実現する。
後方通過のために、StreamKLは2次中間体の保存を回避し、注意確率をタイル単位で再計算する。
我々は、専用最適化による効率的なGPUカーネルの設計と実装をさらに進める。
実験によると、StreamKLは、それぞれ前方パスと後方パスのベースラインメソッドよりも最大43\times$と14\times$のスピードアップを提供する。
最も重要なことは、StreamKLは、注意蒸留のHBMフットプリントを$O(N_QN_K)$から$O(1)$に減らし、単一のGPU上での長いコンテキスト蒸留を可能にすることである。
関連論文リスト
- AdaSplash-2: Faster Differentiable Sparse Attention [14.338708749838881]
AdaSplash-2を導入し、$を1--2に計算するのに必要なイテレーション数を削減します。
AdaSplash-2はブロック間隔が適度に高い場合、FlashAttention-2と比較してステップごとのトレーニング時間を一致または改善する。
ダウンストリームタスクでは、効率的な$-entmaxアテンションでトレーニングされたモデルは、短いコンテキスト長でソフトマックスベースラインと一致し、長いコンテキスト設定で大幅に向上する。
論文 参考訳(メタデータ) (2026-04-16T16:03:13Z) - FlashSinkhorn: IO-Aware Entropic Optimal Transport [12.298153797924074]
シンクホーン反復によるエントロピック最適輸送(EOT)は、現代の機械学習で広く使われているが、解法は大規模に非効率である。
正方形ユークリッドコストに対するIO対応EOTソルバである textbfFlashSinkhorn を提案する。
A100$では、FlashSinkhornは、ポイントクラウドOT上の最先端のオンラインベースラインよりも、32倍のフォワードパスと161倍のエンドツーエンドのスピードアップを実現している。
論文 参考訳(メタデータ) (2026-02-03T03:52:20Z) - Spava: Accelerating Long-Video Understanding via Sequence-Parallelism-aware Approximate Attention [63.69228529380251]
Spavaはシーケンス並列フレームワークで、ロングビデオ推論に最適化されている。
Spavaは、FlashAttn、ZigZagRing、APBで12.72x、1.70x、1.18xのスピードアップを提供する。
論文 参考訳(メタデータ) (2026-01-29T09:23:13Z) - Efficient Low Rank Attention for Long-Context Inference in Large Language Models [41.24530756499533]
低ランククエリとキーアテンション(LRQK)は、プリフィル段階で、完全精度クエリとキー行列をコンパクトなランク-(r)要素に分解するフレームワークである。
トップ(k)トークンと、最近のトークンの小さな固定セットだけを選択することで、LRQKは、完全に精度の低いKVペアだけを転送するヒットアンドミス機構を備えた混合GPU-CPUキャッシュを使用する。
論文 参考訳(メタデータ) (2025-10-25T11:43:27Z) - SLA: Beyond Sparsity in Diffusion Transformers via Fine-Tunable Sparse-Linear Attention [88.47701139980636]
Diffusion Transformer(DiT)モデルでは、特にビデオ生成において、注意遅延が大きなボトルネックとなっている。
注目重量は2つの部分に分けられる: 高いランクの大型重量のごく一部と、非常に低いランクの残りの重量の2つである。
本稿では,拡散モデルを高速化するために,疎度と直線的注意を融合させる訓練可能な注意法SLAを提案する。
論文 参考訳(メタデータ) (2025-09-28T17:58:59Z) - TPLA: Tensor Parallel Latent Attention for Efficient Disaggregated Prefill and Decode Inference [48.40143137402824]
MLA(Multi-Head Latent Attention)は、キー値の状態を低ランクの潜在ベクトルに圧縮し、このベクトルだけをキャッシュしてメモリを減少させる。
しかし、テンソル並列性(TP)では、アテンションヘッドは複数のデバイスにまたがって計算され、各デバイスはフルキャッシュをロードしなければならない。
本稿では,潜在表現と各頭部の入力次元をデバイス間で分割し,シャード毎に独立して注目を行い,結果を全再現と組み合わせる方式であるTPLAを提案する。
論文 参考訳(メタデータ) (2025-08-21T15:25:40Z) - LeanK: Learnable K Cache Channel Pruning for Efficient Decoding [12.370497592637179]
大きな言語モデル(LLM)は、長いコンテキストタスクを可能にするが、キーバリュー(KV)キャッシュの増加による効率上の問題に直面している。
本稿では,静的チャネル空間を利用して重要でないキー(K)キャッシュチャネルを創り出す学習ベースの手法であるLeanKを提案する。
論文 参考訳(メタデータ) (2025-08-04T09:08:43Z) - Decoupled Kullback-Leibler Divergence Loss [90.54331083430597]
我々は、クルバック・リブラー(KL)の除算損失がデカップリングカルバック・リブラー(DKL)の除算損失と等価であることを証明した。
我々はKL/DKLにクラスワイドなグローバル情報を導入し、個々のサンプルからバイアスを取ります。
提案手法は,新たな最先端の対人ロバスト性を公衆のリーダーボード上で実現する。
論文 参考訳(メタデータ) (2023-05-23T11:17:45Z)
関連論文リストは本サイト内にある論文のタイトル・アブストラクトから自動的に作成しています。
指定された論文の情報です。
本サイトの運営者は本サイト(すべての情報・翻訳含む)の品質を保証せず、本サイト(すべての情報・翻訳含む)を使用して発生したあらゆる結果について一切の責任を負いません。