- decoder-only Transformer は LLM の標準アーキテクチャになったが、長いコンテキストでは重要情報を見落とし、無関係なトークンに注意を割り当てる アテンションノイズ が性能を揺るがす
- differential attention は query と key を 2 つのグループに分けて 2 つの softmax attention map を作り、その後、学習可能な λ を適用した 2 番目のマップを差し引いて最終的な attention score を計算する
- DIFF Transformer は、モデルサイズ、学習トークン数、コンテキスト長を拡大する実験で Transformer より良い結果を示し、同等の性能に必要な モデルサイズまたは学習トークン数は約 65% 水準だった
- 長いコンテキストのモデリング、重要情報検索、幻覚の緩和、in-context learning、数理推論、activation outlier の低減で利点を示し、QA と要約では無関係なコンテキストの影響が減少する
- 全体レイアウトは Transformer と同じまま保ち、softmax attention だけを置き換えるため、パラメータ数と計算量 を揃えながら FlashAttention も再利用できる
Transformer のアテンションノイズ問題
- decoder-only Transformer は LLM の事実上の標準アーキテクチャであり、その中核にはシーケンス内のトークン重要度を softmax で重み付けする attention mechanism がある
- LLM はコンテキストから重要情報を正確に見つけ出すことに苦労しており、特に無関係なコンテキストが多いほど正答の手がかりが埋もれやすい
- 文書の山の途中に挿入された正答を見つける必要がある例では、Transformer は正答には小さな attention score しか与えず、無関係なコンテキストには過剰な score を割り当てる傾向を示す
- このように無関係なコンテキストへ割り当てられた無視しにくい attention score が attention noise として作用する
- Figure 1 の Multi-Needle Retrieval の例では、Transformer と Differential Transformer の精度がそれぞれ 55% と 85% と示されている
Differential attention の動作方式
- DIFF Transformer は sequence modeling と LLM のための基盤アーキテクチャであり、既存 Transformer の macro layout は維持しつつ、通常の softmax attention を differential attention に置き換える
- 入力
X から query、key、value を射影するが、query と key は Q1, Q2, K1, K2 の 2 グループに分け、value は V とする
- attention 出力は 2 つの softmax attention map の差として計算される
DiffAttn(X) = (softmax(Q1K1^T / √d) − λ softmax(Q2K2^T / √d))V- 1 つ目の attention map から 2 つ目の attention map を引くことで 共通ノイズ を除去する構造である
λ は学習可能な scalar であり、学習ダイナミクスを合わせるため次のように再パラメータ化される
λ = exp(λq1 · λk1) − exp(λq2 · λk2) + λinit- 基本実験では
λinit = 0.8 − 0.6 × exp(−0.3 · (l − 1)) を使用する
- すべての layer に同じ
λinit、たとえば 0.8 を使う方法も検討され、ablation では初期化戦略の違いに比較的頑健な性能を示した
- この方式は、2 つの信号の差によって common-mode noise を除去する differential amplifier や、ノイズキャンセリングヘッドホンの発想に似ている
- Naderi et al. は、differential attention が attention matrix の spectral distribution をよりバランス良くし、rank collapse を効果的に解決することを証明した
Multi-head 構造と全体アーキテクチャ
- multi-head differential attention は、複数の head ごとに異なる projection matrix を使用し、同じ layer 内では scalar λ を head 間で共有 する
- 各 head の出力には独立に RMSNorm を適用した後、
(1 − λinit) を掛け、head を channel dimension で concatenate してから出力 projection WO を通す
- Figure 2 の GroupNorm 表記は、各 head に独立して normalization を適用する点を強調している
- differential attention はより sparse なパターンを持つ傾向があるため、head 間の統計情報がより多様になる
- head ごとの normalization は concatenate 前に各 head を正規化し、gradient statistics を改善する
- 全体の DIFF Transformer layer は 2 つのモジュールで構成される
MultiHead(LN(Xl)) + XlSwiGLU(LN(Yl)) + Yl
- 構造上 pre-RMSNorm と SwiGLU を使用しており、これは LLaMA 系の改善に従ったものである
効率性と学習安定性
- differential attention は FlashAttention をそのまま再利用でき、モデル効率を大きく改善できる
- head 数は
h = dmodel / 2d に設定し、ここで d は Transformer の head dimension と同じである
- この設定は Transformer と パラメータ数および計算複雑性 を揃えるためのものである
- head normalization の後に固定 multiplier
(1 − λinit) を使って Transformer と gradient flow を合わせる
- Appendix G では、全体の gradient flow が Transformer と同様に維持されることが示されており、これにより近い hyperparameter を継承しつつ学習安定性を確保できる
実験結果と適用効果
- 言語モデリング実験では、parameter count、training tokens、context length を拡大する方向で DIFF Transformer をスケールさせた
- scaling curve では、DIFF Transformer は Transformer と同程度の言語モデリング性能を達成するのに必要な モデルサイズまたは学習トークン数が約 65% の水準だった
- downstream task でも Transformer より高い性能を示し、long-sequence 評価ではコンテキストが増えるほどそれを効果的に活用した
- 重要情報検索では、正答 span により高い attention score を与え、無関係なコンテキストにはより低い score を与えるパターンを示した
- QA と text summarization では、無関係なコンテキストに気を取られにくくなり、hallucination を緩和する
- in-context learning では精度を高めると同時に、慢性的な頑健性問題とみなされてきた order permutation に対してもより頑健である
- activation outlier を減らす結果は、quantization に向けた新たな機会を提供する
1件のコメント
Hacker News のコメント
ここで肝心な直感を見落としている気がします。通常のソフトマックス・アテンションでは、無関係なものにほぼ 0 の注意を向けるのが難しいという問題は理解できますし、減算構造があれば外れ値的な活性値なしでも、正確に、あるいはほぼ 0 の注意重みを作れるという点も理解できます。
ただしこの構造は負のアテンション重みも簡単に作ってしまいそうで、それは値ベクトルの否定に正の注意を向けるのと同じに見えます。直感的には、関心のないものをすべて 0 付近に保つバランス調整が難しそうです。それでも Figure 1 はうまく動作することを示しているので可能性は疑っていませんが、ネットワークが具体的にどうやってそれを実現しているのか、うまくイメージできません。
exp()/1+∑exp()であるべきです。分母に 1 が追加される点が核心です。負の極限でソフトマックスは何らかのイプシロンではなく 0 になれます。
xに追加の 0 値を入れても同じ効果を得られます。欠点は、これを直すにはモデルを最初から再学習する必要があることです。1 つ目のアテンション重みの図を見ると、ノイズ区間には実際に負のスコアがあります。それでもその区間への注意はどのみち非常に小さいです。2 つ目のアテンションマップは 1 つ目のノイズを予測するだけでよく、1 つ目の入力に完全にアクセスできるので、かなり正確にできる作業です。
論文の現実世界の比喩に戻ると、ノイズキャンセリングヘッドホンはマイクを通じて耳に聞こえる音にアクセスできるため、正確な相殺信号を出せます。同様に、2 つ目のアテンションマップは 1 つ目に何が入力されているかを知っているので、対応する相殺信号を出せます。完璧ではありませんが、ノイズキャンセリングヘッドホンも完璧ではなく、それでも 99% までは到達して性能向上には十分です。
枝刈りはパラメータ数を驚くほどよく減らす方法として、すでに文献でかなり確立されており、おおよそ 40% まで減らせます。実際のモデルが正確にそう動くわけではないでしょうが、結局は普通の Transformer を近似する形であっても驚きません。
とても賢いです。こういう細かな仕事は好きですし、変更も小さいので他の人が簡単に適用できそうです。素晴らしいです。
ただし「2 Differential Transformer」セクション導入部の最後の文が少し気になります。以前の論文の改善を使うと言っていますが、文法上の文脈では、その改善が通常の Transformer と diff Transformer の両方に入っているのか不明です。そうでなければ比較が曖昧になります。直前の文の「main difference」という表現で警戒しました。
もちろん善意の研究者ならこれを理解していて、あえて明確にする必要を感じなかった可能性もあります。しかし、この分野の一部の出版研究では、いくら慎重でも慎重すぎることはありません。
MoE 構造を思い出します。あの世界では、推論作業の一部または全部を処理する最適な小さなモデルを選びます。Transformer が代替可能性を区別するよう強制されることで、MoE も同様の利得を得たのか気になります。
いずれにせよ数値が維持されるなら、広く採用されそうです。先ほど言ったように、実質的に欠点がなさそうで、再現も簡単そうです。
機械学習というこの新しい世界にある大半のものと同じく、なぜこれが動くのか本当に混乱します
ノイズキャンセリングヘッドホンの比喩は役に立ちますが、その場合は何が信号で何がノイズかを明確に分かっています。ここでもそれが分かっているなら、そもそもなぜノイズ除去作業をしなければならないのか分かりません
さらに悪いことに、低いアテンション値の勾配は非常に小さくなるため、そのような誤りを元に戻すには多くの重み更新が必要になります。一方、2つのソフトマックスの出力を引けば、モデルは一部の値について正確に0である重みを予測でき、妥当な勾配の流れも維持できます
つまりモデルはすでに何がノイズかを知っていますが、単一のソフトマックスがそれを排除しにくくしています。また単一のソフトマックスでは、すべてのヘッドの出力が値ベクトルの凸包の中に留まるよう強制されますが、この変形では各ヘッドが独自のlambdaを選び、出力範囲を値があらかじめ定める凸包の外へ移せます。そのためモデル全体の表現力が大きくなります
よりよい例は、プロ向けオーディオやEthernet、HDMI、USBのような多くのデジタル信号プロトコルで使われる差動信号です。接地を基準にした1本の線を使う代わりに、2本の線の差で信号を送ります。2本の線は同じ信号を逆極性で運び、互いに並行して通るので、外部ノイズは両方に同じように加わります
電圧は変わるでしょうが、2本の線の間の電圧差はそのままです。受信側で2つの電圧を引くと、ノイズは単に相殺されます
これを行う単純な方法はソフトマックスをなくすかsigmoidを使うことでしょうが、実際にはソフトマックスの方がうまく動くようです
簡単に言うと、RoPEはアテンションを行う際にクエリとキーがどれだけ離れているかの情報をモデルに与える現代的な戦略です。現時点で最善の戦略ですが、遠く離れたトークン間の一部の結合を、望ましいよりもはるかに強くしてしまう大きな欠点があります。Xpos (https://arxiv.org/pdf/2212.10554)もMicrosoftがRoPEの問題を扱った論文で、4ページのFigure 1を見ると、正弦波状のアテンション強度の視覚的解釈が見られます。本来は滑らかであってほしいものです
Differential Transformerが特に長いシーケンスでうまく動く大きな理由は、
q1とq2があるトークンとどちらも一致しない場合でも、RoPEの相対強度は同じ値を持つため、ノイズが相殺されるからだと見ています。意図した一致だけが残りますが、RoPEが本来もたらしていた値がある程度弱まるというコストはありますもちろん仮説にすぎません。どちらもalibi attention (https://arxiv.org/pdf/2108.12409)を使うベースラインと比較して実験すれば、簡単に検証できます。alibiにはこの方式では緩和できない別のトレードオフがありますが、それでも本当に興味深い結果です
なぜ動くのかについての直感は、勾配降下の地形を少し扱いやすく変え、小さなステップで学習しやすくすることにあるようです。今では、ネットワーク自体を「最初は予測に多くの誤りを出し、時間とともに改善する」というアイデアに合わせて明示的に設計するからです
「Differential attention takes the difference between two softmax attention functions to eliminate attention noise」を正しく理解しているなら、この構造はアテンションメモリを2倍使う代わりに、より高品質なモデル、または同程度の品質でより少ないパラメータを得るトレードオフに見えます
「6.8B-size DIFF Transformer achieves a validation loss comparable to 11B-size Transformer, requiring only 62.2% of parameters」という部分を見ると、いくつか疑問が湧きます。パラメータが60%しかないなら、アテンション空間が2倍である点を相殺して従来のTransformerと似たメモリ特性になるのか、そしてそのトレードオフが学習と推論の間で目に見えて異なるのかが気になります
論文には「We set the number of heads h = dmodel/2d, where d is equal to the head dimension of Transformer. So we can align the parameter counts and computational complexity.」とあります。言い換えると、層ごとのアテンションヘッド数を半分に減らして、これを補っています
「We empirically find that the setting λᵢₙᵢₜ = 0.8 − 0.6 × exp(−0.3 · (l − 1)) works well in practice」という式の裏に、どんな経緯があったのか気になる
lで負のアテンション項を後方の層より小さくするように見える。もっともらしい。実際に見るべきいくつかの位置だと結論づける前は、あらゆるものに少しずつ注意を向けたいはずだからただ、論文で著者がこの選択を個別に議論してはいないように思う
最初に理解できなかった核心は、2つのアテンショングループが同じものを学習したらどうなるのかという点だった。アテンションマスクを互いに引き算するので、両方が似た値を出力すると全体のアテンションが0に落ちて損失が大きくなる
したがって損失を減らす唯一の方法は、互いに異なるものに注意を向けるよう学習することになる。彼らが学べる最も単純な戦略の1つは、論文が主張するように、一方のグループは関連する文脈に集中し、もう一方のグループは無関係な文脈に集中することだ。そうして一方のグループはノイズを、もう一方のグループは信号を学ぶことになる。実際にはここまできれいに分かれるわけではないが、理解のための単純化としては有用だ
2つのコピーが同一ならソフトマックス出力も同一になり、差があらゆる場所で0になることを考えると筋が通る。しかしスケールされたコピーを引くと、差を正規化する過程が信号値をノイズより大きく際立たせ、正規化前より信号が目立つように見える
あるアテンションヘッドが別のヘッドの学んだことに驚けば重みを上げ、両方が同じものを見つけたらそれほど驚くべきではないと見て重みを下げる、という感じだ
認めると、「驚き」は私の知識基盤の中でかなり大きな領域を占めている[1][2][3]。主観的な感情であり、心の適応機能として、私たちが知る最も複雑な適応システムの1つだ
[1] https://plus.maths.org/content/information-surprise
[2] https://blakeelias.name/papers/Multi-Agent-Cooperation-Intri...
[3] https://complexity.simplecast.com/episodes/81/transcript
ここで何を失うのかが気になる。明らかにトレードオフはありそうだ
創造性や概念間の補間能力に影響があるのかも気になる。ハルシネーションと創造性はかなり関係があるように見える。ハルシネーションは、人間が適切だと感じる補間空間から外れたものだと理解している
もちろん、失敗がときどき着想を与えることはあるが、創造性は失敗よりはるかに多くのものを含む
この種の言語モデルは次トークン予測器だ。次のトークンは、モデルが出力した確率空間からサンプリングして予測される。そのサンプリング過程は非決定的になり得る
ハルシネーションとは、そのサンプリング結果として、偽または意図しない文を作るトークンが出てきた場合だ。モデルが出力するすべてをハルシネーションと見なすこともできるが、私たちはモデルが自分たちの望むものをハルシネートする確率がより高い空間を出力するよう学習させている。そうでなければ意味のないノイズだけを出す
「ハルシネーション」は、説明しようとしている対象に対して本当にひどい言葉だ
ここでの価値のどれだけが、RoPEが生む位置ノイズを相殺することから来ているのか気になる。ここにあるRoPEモデルだけでなく、alibi版とalibiベースラインを比較した表も見てみたい
それでも大きな改善であり、研究者たちに祝意を送りたい
ここで起きていることは、ソフトマックスは値を0に押しやることはできないが、2つのソフトマックスマップを引けば0出力を作れる、ということなのか
解くべき良い問題だが、アプローチは間違っていると思う
何に注意を向けたのかと全体の文脈を知るには、階層的な方法でやる必要がある。差分ベクトルがアテンションベクトルと同じ入力から計算されるなら、アテンションベクトルをどう正しく修正すべきかをどうやって知るのか分からない