Microsoftが差分アテンション第2世代を発表——推論速度を標準Transformerと同等に
MicrosoftのUnifiedLanguage Modelチーム(Tianzhu Ye、Li Dong、Yutao Sun、Furu Wei)は、大規模言語モデル(LLM)のアテンション機構を改良したDifferential Transformer V2(DIFF V2)を発表した。
Differential Transformerとは
差分アテンション(Differential Attention)は、Microsoftが提唱したアテンション機構で、「2つのSoftmaxアテンションの差分を取る」という独自のアプローチでノイズを打ち消し、重要なコンテキストへの集中力を高める技術だ。第1世代(DIFF V1)は精度面で優れた結果を示していたが、デコード速度の低下とカスタムアテンションカーネルの必要性という2つの課題を抱えていた。
DIFF V2の核心——KVヘッドを増やさずクエリヘッドを2倍に
DIFF V2が採用した設計の要点は、クエリヘッド数を2倍にしつつ、キーバリュー(KV)ヘッド数はそのまま維持するという点にある。
LLMの推論(デコード)フェーズはメモリ帯域幅に律速されるため、KVキャッシュのサイズが速度を左右する。DIFF V2はKVヘッドを増やさないことで、標準Transformerと同等のデコード速度を実現した。DIFF V1ではバリューキャッシュを2回ロードする必要があり速度低下が生じていたが、DIFF V2ではこの問題が解消されている。
差分演算では、同じGQA(Grouped Query Attention)グループに属する2つのクエリヘッドの出力から差し引く仕組みを採用。このグループ内でキーとバリューが共有されることが、性能面で重要な役割を果たしている。
カスタムカーネル不要でFlashAttentionと完全互換
DIFF V1では差分アテンションの実装にカスタムCUDAカーネルが必要だったが、DIFF V2ではクエリ・キー・バリューのヘッド次元が揃っているため、標準的なFlashAttentionカーネルをそのまま利用できる。NVIDIAのHシリーズ・Bシリーズ(Hopper/Blackwell世代)GPUでの事前学習においても、スループット低下はほぼ無視できる水準に抑えられているという。
また、DIFF V1で採用されていたRMSNormによる後処理や複雑なλ計算を廃し、シグモイド関数を使ったシンプルなλ推定(トークンごと・ヘッドごとに入力Xから動的に射影)に置き換えた。これにより実装が大幅に簡素化された。
長文脈処理にはYOCOとの組み合わせを推奨
長シーケンスのプリフィリング(初期コンテキスト処理)については、YOCO(You Only Cache Once)と組み合わせることを研究チームは推奨している。YOCOはGemma 3nでも採用されており、シーケンス長に対するプリフィリング計算量を線形に抑える技術だ。
実装・コードの公開
コードはMicrosoftのGitHubリポジトリ(microsoft/unilm)で公開されており、Hugging Face Blogでも詳細な解説と数式が読める。差分アテンションの改良版として、実際のLLM開発への応用が現実的な選択肢となってきた。日本のAI研究者・エンジニアにとっても、標準的な開発環境でそのまま試せる実用的な技術として注目に値する。