RBF-Attentionを試した実験が示す、attentionを入れ替える難しさ

r/MachineLearningで165ポイント、23件のコメントを集めたこの投稿は、Transformer の中心にある scaled dot-product attention を distance-based RBF-Attention に置き換えたらどうなるかを記録した実験レポートだ。出発点は明快で、dot product はベクトルの向きだけでなく大きさにも強く影響されるため、norm の大きい key が softmax を押し切ってしまう場合がある、という問題意識だ。そこで著者は Euclidean distance に基づく類似度を使い、query と key が本当に近いときだけ高いスコアを与える構成を試した。

面白いのはアイデアそのものより、その実装過程だ。単純に torch.cdist を使うと巨大な距離行列が materialize されてすぐ OOM になるため、著者は式を変形し、softmax の shift-invariance を使って query norm を落とし、実質的に 2(Q·K)-||K||^2 に近い形で処理できると説明する。それでも PyTorch の fused SDPA には key norm penalty を差し込めないため、最終的には Triton で custom kernel を書く必要があった。

• distance ベースにすると通常の attention sink が失われるため、register tokens を追加して余った attention mass の受け皿を作った。
• RoPE は回転ベースの幾何なので Euclidean distance と相性が悪いとして外し、代わりに additive な SuSiE embedding を採用した。
• TinyStories 上の小さな causal model では baseline よりやや速い収束が見えたが、著者自身も FlashAttention をすぐ置き換える話ではないと述べている。

コメント欄でも、kernelized attention との関連を指摘する声や、最終的には hardware lottery が採用を決めるという指摘が出ていた。つまりこの投稿の価値は、次の標準 attention を宣言したことではなく、attention の根本仮定を変えると LLM スタック全体がどこで壊れるかを具体的に見せた点にある。

参照先は Reddit スレッド、技術ブログ、コードリポジトリ。