CORAL-KV、層間予測で動画拡散モデルのKVキャッシュを圧縮

導入

自己回帰型の動画拡散モデルは、生成済みフレームの情報をアテンションのキー・バリュー（KV）キャッシュに保持します。動画が長くなるほどキャッシュは増え続け、GPUメモリの大きな負担になります。圧縮すれば使用量は抑えられますが、復元誤差が後続フレームのアテンション計算に何度も影響し、画質低下や時間方向の不連続につながる可能性があります。

arXivで公開されたCORAL-KVは、この問題に対し、キャッシュの層間に存在する依存関係を圧縮へ取り込む手法です。

手法のポイント

従来のクラスタリング型量子化は、主に同一層のキーとバリューをグループ化し、残差を量子化します。CORAL-KVでは、ある層のキャッシュ状態から後続層のキーとバリューを予測します。その後、元のベクトル全体ではなく、予測しきれなかった残差を量子化します。予測が十分に正確なら、保存すべき情報量を減らせます。

予測器には、ブロック対角行列と低ランク項を組み合わせたblock-DPLR構造を採用しています。

• ブロック対角成分：各アテンションヘッド内のチャネル間関係を表現します。
• グローバル低ランク成分：異なるヘッド間にある依存関係をまとめて捉えます。

また、精度とメモリ使用量のバランスを調整できる点も特徴です。予測器を動画間で再利用すれば適応処理を抑えられます。一方、現在の動画に合わせて軽量なブロック対角予測器を再フィットすれば、より高い復元精度を狙えます。

実験結果

HY-WorldPlayとMAGI-1を用いた評価では、再利用可能な構成がQVG-4に対してPSNRをそれぞれ0.32 dB、4.95 dB向上させました。同時に、キャッシュメモリ使用量は48%、37%削減されています。動画適応型の構成では、MAGI-1においてQVG-ProよりPSNRが3.77 dB高く、メモリ使用量も33.7%少ないと報告されています。

さらに、長時間の生成でも復元品質を維持しやすく、時間方向の不連続を抑えられたとされます。生成レイテンシは比較対象と同程度でした。ただし、これらは論文の実験条件に基づく結果であり、実際の効果はモデル構造、量子化設定、動画の長さによって変わる可能性があります。

意義と今後の論点

CORAL-KVの重要な点は、KVキャッシュを層ごとに独立した圧縮対象として扱わず、動画モデル内部の層間構造を利用したことです。層間の相関を安定して再利用できれば、より小さなメモリ予算で長い動画を生成しながら、圧縮誤差の蓄積も抑えられる可能性があります。

実運用では、動画間で予測器を再利用するか、動画ごとに調整するかが選択になります。後者は品質向上が期待できる一方、適応処理が必要です。予測器のフィッティングと残差量子化が各種ハードウェア上でどれだけの負荷になるかも、今後検証すべき課題です。

出典：arXiv