Kimi Delta Attentionはあなたも考案できたはず

数式表記 ⟨k|q⟩ kᵀq この記事では、量子計算にインスパイアされた考え方から、導出過程における形状を明確にするために、ブレイケット記法をデフォルトで使用します。上記の数式表記の切り替えは、すべての等式を従来の太字ベクトルと明示的な転置を用いて書き換えます。ブレイケット記法では、∣q⟩は列ベクトル、⟨k∣は行ベクトル、⟨k∣q⟩はスカラー、∣v⟩⟨k∣は行列です。デフォルトではベクトルは右向きですが、線形アテンションの状態に書き込む際にはキーは左向きになります。ここでは、1つの因果的アテンションヘッドと実数値ベクトルを扱い、DeltaNetのキーは正規化されていると仮定し、状態がキー空間から値空間へのマッピングを行うとします。 現代の線形アテンションのバリアントは複雑であり、初見ではそれらが何を達成しようとしているのかを容易に理解することはできません。参考までに、Kimi Delta Attention (KDA) の状態更新方程式を以下に示します。 St−1=St−1Diag⁡(αt) ∣v^t⟩=St−1Diag⁡(αt)∣kt⟩ ∣et⟩=βt(∣vt⟩−∣v^t⟩) St=St−1+∣et⟩⟨kt∣ ∣ot⟩=St(dk−1/2∣qt⟩) それらが理解しにくい理由は、過去数年間に開発されてきた線形アテンションのバリアント群の最新版であり、その複雑さが必然的に膨れ上がっているため、外部からは最新のバリアントにアクセスするのが困難に見えるからです。この記事では、QwenとKimiの最新モデルファミリーで使用されている2つのバリアントを含む、DeltaNetファミリーの線形アテンションのバリアントを順を追って説明し、隠れ状態について単純なことを仮定することで、どのようにして同じ等式に至ることができるかを示します。それが私たちが取る道筋です。 softmaxアテンション → 線形アテンション → DeltaNet → Gated DeltaNet → KDA KDAを導出した後でなければ、それを実行する再帰的およびチャンクワイズのTritonプログラムには触れません。 1. 二次アテンションから始める トークンtのクエリに対して、通常の因果的ソフトマックスアテンションは以下のようになります。 ati=exp⁡(s⟨ki∣qt⟩)∑j≤texp⁡(s⟨kj∣qt⟩),s=dk−1/2, ∣ot⟩=∑i≤tati∣vi⟩. すべての<0xE3><0x82><0x83>テンション重みはスカラーです。それは1つのキーと1つのクエリの類似性を測定し、その後ソフトマックスがそのクエリのすべてのスコアを分布に変換します。出力は値ベクトルの加重合計です。シーケンス長TTの場合、T²個のキー・クエリペアが存在します。自己回帰推論中にキーと値を再計算せずにキャッシュすることはできますが、キャッシュはシーケンスとともに成長し、新しいクエリは常に履歴全体を検査する必要があります。この計算の並べ替えを妨げるのはソフトマックスです。その分母は、現在のクエリと以前のすべてのキーに共同で依存します。したがって、いったんそれを削除します。 1.1 ソフトマックスの削除 明確にするために、定数スケールsをクエリに吸収します。アテンションの意図的に裸のバージョンは次のようになります。 ∣ot⟩=∑i≤t⟨ki∣qt⟩∣vi⟩. 内積は右に移動できます。 ∣ot⟩=∑i≤t∣vi⟩⟨ki∣qt⟩=(∑i≤t∣vi⟩⟨ki∣)∣qt⟩. 過去に依存するすべてを、固定サイズのV×K行列に収集できます。 St=∑i≤t∣vi⟩⟨ki∣ そしてアテンションは、再帰的な書き込みとそれに続く読み取りになります。 St=St−1+∣vt⟩⟨kt∣, ∣ot⟩=St∣qt⟩. 恒等式(∣v⟩⟨k∣)∣q⟩=⟨k∣q⟩∣v⟩がトリック全体です。外積は行列であり、内積はスカラーです。もはや過去のすべてのキーと値を保存しません。それらの合計外積を固定サイズの状態Stに保存します。これはシーケンス長に対して線形であり、二次ではありません。トークンを一度スキャンし、各ステップで同じdv×dk状態を更新します。その効率のために、ソフトマックスの正規化と選択性を破棄しました。より洗練された線形アテンション手法は、特徴マップと正規化を使用しますが、この飾り気のない形式は、DeltaNetを動機付けるメモリの問題を露呈します。 1.2 加算は代入ではない ペア∣vt⟩⟨kt∣を書き込み、その同じキーで新しい状態をすぐにクエリすると仮定します。 St∣kt⟩=(St−1+∣vt⟩⟨kt∣)∣kt⟩=St−1∣kt⟩+∣vt⟩⟨kt∣kt⟩⏟1=St−1∣kt⟩+∣vt⟩. 書き込みはメモリに∣vt⟩を返しません。それは、メモリが既に返したものに∣vt⟩を追加します。古い状態が既に正しい値を返していた場合、加算書き込みは新しい状態がその値の2倍を返すようにします。より一般的には、キーは互いに直交しないため、各書き込みは以前の書き込みに干渉する可能性があります。線形アテンションは、コンパクトな連想メモリを提供してくれましたが、その更新は、私たちが望むものよりも=に近い動作をします。 2. DeltaNet: 値ではなくエラーを書き込む DeltaNetは、無条件の線形アテンション書き込みをデルタ規則の修正に置き換えます。それを導出する2つの有用な方法があります。 2.1 導出1: 書き込みが読み取れるように要求する トークンtを書き込む前に、メモリに新しいキーと現在関連付けられているものを尋ねます。 ∣v^t⟩=St−1∣kt⟩. メモリに∣vt⟩を返させたい場合、値全体を追加するのではなく、差分のみを追加する必要があります。 ∣vt⟩−∣v^t⟩. 学習可能な書き込み強度βt∈[0,1]を導入し、次のように定義します。 ∣et⟩=βt(∣vt⟩−St−1∣kt⟩). 次に、このエラーを現在のキーに書き込みます。 St=St−1+∣et⟩⟨kt∣. これで、同じキーをすぐに読み取ります。 St∣kt⟩=St−1∣kt⟩+∣et⟩⟨kt∣kt⟩=(1−βt)St−1∣kt⟩+βt∣vt⟩. βt=1の場合、結果は正確に∣vt⟩になります。より小さいβtは、古い予測部分をターゲットに近づけます。修正はキー空間でも局所的です。現在のキーに直交する任意のクエリ∣x⟩に対して、⟨kt∣x⟩=0の場合、(St−St−1)∣x⟩=∣et⟩⟨kt∣x⟩⏟0=0となります。