LLM 推論系統
RMM、入力に応じて行列乗算の縮約軸を枝刈り――Llama 3.1 8Bの長文コンテキストで最大1.40倍高速化
Reduced Matrix Multiplicationは重みを変更せず、現在のactivationの列ノルムに基づき、行列乗算でスコアが最も高い共有次元のみを計算する。A100でのテストでは、シーケンスが長くなるほど効果が増した一方、短いコンテキストでは選択処理とkernel起動のコストが計算削減分をほぼ相殺した。
新たに公開された[Reduced Matrix Multiplication（RMM）の論文](https://arxiv.org/abs/2608.13426)は、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版](https://openreview.net/forum?id=2uxuiykvA4)では初期の手法を追跡できるものの、論文に記載されたGitHubリポジトリは原稿締め切り時点でも404を返していた。本当に注目すべきなのは、batch 1の長いprefillだけで成立するかではなく、vLLM、量子化kernel、continuous batching、高並行性サービスへ統合できるかどうかだ。