Differential Transformer

Differential Transformer（DIFF Transformer）は、Transformerアーキテクチャの一種です。アテンションヘッドが2つのソフトマックスアテンションマップを計算し、一方からもう一方を差し引くことで、無関係なコンテキストに両方のマップが割り当てる重みを相殺します。Microsoft Researchと清華大学のTianzhu Ye、Li Dong、Furu Weiらが2024年10月に考案し、ICLR 2025の口頭発表論文Differential Transformerで発表しました。この設計は、2つの信号を差し引いて共通するノイズを取り除く差動増幅器やノイズキャンセリングヘッドホンに着想を得ています。言語モデリングでは、標準的なTransformerと同じ損失を達成するのに必要なパラメーター数または学習トークン数は約65%で済みました。68億パラメーターのモデルが、110億パラメーターのTransformerと同等の性能を達成しました。

Differential Attentionの仕組み#

標準的なself-attentionでは、すべてのトークンに非負の重みが割り当てられるため、重要でないトークンにも必ず重みが割り当てられます。著者らはこれをattention noiseと呼んでいます。Differential attentionでは、各ヘッドのクエリ射影とキー射影を2つに分け、各ペアからsoftmaxマップを計算し、値を(map1 − λ × map2) × Vとして組み合わせます。どちらのマップも同じ無関係なトークンに重みを割り当てるため、差し引くことでその重みは相殺され、map1だけが割り当てた重みは残ります。各マップの合計は1なので、結果の各行の合計は1 − λとなり、個々の重みは負になることがあります:

import torch

torch.manual_seed(0)
n, d, lam = 6, 8, 0.8  # 6 tokens, head size 8, lambda near the paper's 0.8 initialization
q1, k1, q2, k2 = torch.randn(4, n, d)
v = torch.randn(n, 2 * d)  # values are twice as wide as each query-key pair

# Two ordinary softmax attention maps from separate query and key projections
map1 = torch.softmax(q1 @ k1.T / d**0.5, dim=-1)
map2 = torch.softmax(q2 @ k2.T / d**0.5, dim=-1)

# Subtracting cancels weight that both maps place on the same tokens
diff = map1 - lam * map2
out = diff @ v
print(f"Each row sums to {diff.sum(dim=-1)[0]:.2f}; {(diff < 0).sum().item()} of {n * n} weights are negative")

これはEach row sums to 0.20; 13 of 36 weights are negativeを出力します。softmax attentionマップでは持ち得ない2つの性質です。実際のモデルではλは4つの小さなベクトルを通じて学習され、層lに対して0.8 − 0.6 × exp(−0.3 × (l − 1))で初期化されるため、深い層ほど0.8に近い値から始まります。各ヘッドの出力は個別に正規化され、論文ではGroupNormと表記されています。その後、標準的なTransformerと同じ勾配の大きさを保つため、(初期化時の1 − λ)を乗じます。モデルはヘッド数を半分にし、各ヘッドの値の幅を2倍にすることで、ベースラインと同じパラメーター数にしています。各マップは通常のsoftmax attentionなので、実装ではflash attentionカーネルを再利用できます。また、H100 GPUでの学習スループットは、同規模のTransformerより6%から12%低くなりました。

報告された利点#

6本のneedleと2つの質問を含む4Kトークンのコンテキストウィンドウを使ったneedle-in-a-haystackテストで、DIFF Transformerは0.85を記録し、標準的なTransformerの0.55を上回りました。論文のattention分析では、この差は焦点の当て方によるものとされ、DIFF Transformerは回答にattentionの0.27から0.40を割り当てたのに対し、Transformerは0.03から0.09でした。また、テキスト要約と質問応答でハルシネーションが少なく、few-shotのin-context learningでは精度が高く、例の順序への依存も小さくなりました。

活性化の外れ値も小さくなりました。最大attention logitは標準的なTransformerの318.0に対して38.8、最大hidden-state活性化は3,609に対して1,688でした。attention logitを4ビットに量子化した場合、標準的なTransformerを6ビットで使った場合と同等の精度を達成し、4ビットのTransformerより精度が約25%向上しました。

この結果はモデル量子化にとって重要です。少数の極端な活性化によって、ほかのすべての値が共有しなければならない範囲が決まるためです。同じことはビジョンモデルにも当てはまります。Ultralytics YOLO26をquantize=8でエクスポートすると、エクスポーターはdataデータセットの画像で活性化の範囲をキャリブレーションし、少なくとも300枚の使用を推奨します。一方、quantize="w8a16"では、LiteRT、Core ML、QNNへのエクスポート時に、重みを8ビット、活性化を16ビットに維持します。

DIFF Transformerと関連用語#

• Diffusion Transformer (DiT)：画像や動画の生成を行う拡散モデルで、U-Netを置き換えるTransformerです。
• 差分プライバシー：学習データに対する数学的なプライバシー保証です。
• Sparse attention：固定された一部のトークンペアのみをスコアリングし、それ以外の重みは事前にすべて0に設定します。
• 線形可変差動変圧器：線形変位を測定する電気センサーです。