推論系統
RMM、入力に応じて行列乗算の縮約軸を枝刈り――Llama 3.1 8Bの長文コンテキストで最大1.40倍高速化
Reduced Matrix Multiplicationは重みを変更せず、現在のactivationの列ノルムに基づき、行列乗算でスコアが最も高い共有次元のみを計算する。A100でのテストでは、シーケンスが長くなるほど効果が増した一方、短いコンテキストでは選択処理とkernel起動のコストが計算削減分をほぼ相殺した。
新たに公開されたReduced Matrix Multiplication（RMM）の論文は、Transformer推論における近似の単位を、token、ニューロン、固定サブネットワークから、個々の行列乗算の「縮約軸」へと移した。Y=ABの場合、システムは現在のactivation Aの各列についてL2ノルムを計算し、比率ρで指定されたTop-K次元を残したうえで、縮小されたA[:,I]B[I,:]を実行する。同じルールは、QKᵀのhead次元、attention重みとVの乗算におけるtoken次元、線形射影にも適用できる。選択内容は入力、層、head、デコードステップに応じて変化し、再学習や重みの恒久的な削除は不要だ。
著者らは、1Bから70Bまでの言語モデルとQwen2.5-VL-7Bでテストを実施した。Llama 3.1 8B、A100、batch 1、ρ=0.8を例にすると、エンドツーエンドのレイテンシは1,024 tokenで109.39ミリ秒から103.91ミリ秒へ短縮したにすぎず、高速化は1.05倍だった。2,048 tokenと4,096 tokenで初めて、それぞれ1.27倍、1.40倍に達した。70Bでは2,048 tokenで1.41倍となった。4,096 tokenではdense版がメモリ不足に陥った一方、RMMは推論を完了したが、これはモデルの重みが圧縮されたことを意味しない。
品質は枝刈りを行う位置に大きく左右される。8Bモデルでは、ρ=0.8におけるCNN/DailyMailのROUGEはほぼ変化しなかった。しかしρ=0.5まで下げると、5つのQAタスクの平均精度は完全なモデルの69.8から59.8へ低下した。ARC-Easyのρ=0.7の実験では、attention側のみを枝刈りした場合の低下は3.52ポイントだったのに対し、MLP全体では18.78ポイント低下しており、コンポーネントごとに比率を設定する必要性を裏付けている。デプロイ前の評価には、なお慎重さが求められる。レイテンシはわずか10回の平均であり、denseベースラインはHugging Face SDPA、RMMはカスタムTriton kernelを使用している。OpenReview版では初期の手法を追跡できるものの、論文に記載されたGitHubリポジトリは原稿締め切り時点でも404を返していた。本当に注目すべきなのは、batch 1の長いprefillだけで成立するかではなく、vLLM、量子化kernel、continuous batching、高並行性サービスへ統合できるかどうかだ。