Show HN: Proxy Block-CAGE、新しいスパースブロックアテンション
本記事では、長いコンテキストのプリフィル向けスパースアテンションプロトタイプ「Proxy Block-CAGE」を紹介します。Q/K/Vをブロックに分割し、プールされた要約を使用してキーブロックを選択し、選択したブロック内でのみ正確なアテンションを計算します。小さなTransformerで16Kトークンにおいて、密なアテンションと比較して2.61倍の高速化を達成し、密度はわずか0.39%でありながら最終トークンの一致を100%維持しました。遺伝的アルゴリズムから数学的最適化への進化について詳述しています。
人工知能研究において、長いシーケンスを扱うTransformerモデルは、アテンション機構の計算量が系列長の二乗に比例して増加するという課題を抱えています。最近、George Mason大学の博士課程学生Iqbal Addou氏がGitHub上で「Proxy Block-CAGE」というスパースアテンションのプロトタイプを公開しました。これは、長いコンテキストのプリフィル段階における計算オーバーヘッドを削減することを目的としており、ChatGPTを研究パートナーとして活用した実験でもあります。
Proxy Block-CAGEの核となるアイデアは、クエリ（Q）、キー（K）、バリュー（V）を固定サイズのブロックに分割し、各ブロックの平均プール要約を用いてブロック間の代理スコアを計算することです。このスコアに基づき、各クエリブロックは最もスコアの高い少数のキーブロックを選択し、選択されたブロック内でのみ正確なアテンションを実行します。これにより、アテンションの品質を維持しながら、計算するクエリ-キーペアの数を大幅に削減します。
著者はNVIDIA RTX 2080 Max-Q GPUを搭載したノートPC上で、2層の文字レベルTransformerをテストしました。その結果、16,384トークンのコンテキスト長において、Proxy Block-CAGEは密なアテンションに比べて2.61倍の高速化を達成し、アテンション密度はわずか0.39%でありながら、最終トークンのtop-1およびtop-5精度は100%を維持しました。テスト規模は小さいものの、この手法の潜在能力を示しています。
記事では、プロジェクトの進化が詳細に説明されています。当初、著者は遺伝的アルゴリズム（GA）を使って各クエリトークンにスパースなアテンションマスク（「CAGE」と呼称）を選択しようとしましたが、トークンレベルの検索は非常に時間がかかりました。次に、検索単位をトークンからブロックに昇格させることで検索回数を減らしましたが、それでもGAの計算コストは高すぎました。最終的に、単純なlog-sum-exp目的関数に対しては、最適解が単純なtop-k選択と等価であることが判明し、GAはより効率的な順序選択に置き換えられました。この発見がProxy Block-CAGEの最終形態に繋がり、プール要約による代理スコアのtop-kブロック選択が採用されました。
数学的には、Proxy Block-CAGEの計算複雑性はおおよそO(BHd(n²/(b_q b_k) + n m b_k))です（b_qとb_kはブロックサイズ、mは各クエリブロックが選択するキーブロック数）。典型的な設定（b_q=b_k=32, m=2）では、16Kコンテキスト時の計算位置数は密なアテンションの約1/100になります。エンジニアリングの最適化度合いが異なるため実速度向上は理論値には及びませんが、長いコンテキストにおけるスパースアテンションの利点を示すには十分です。
著者は、これはあくまで再現可能な研究プロトタイプであり、プロダクション対応ではないと強調しています。このプロジェクトを公開することでコミュニティからのフィードバックを得て、スパースアテンションにおける数学的・アルゴリズム的革新をさらに探求したいと考えています。