GDN, KDA の整理

Qwen3.8-Flash-Next と Kimi K3 が話題になっていたものの、具体的にどういうアーキテクチャなのかを理解していなかったので、論文を読んで理解の確認も兼ねて blog にしています。まずは Attention 周りから始めて気持ちが続けば residual、Engram 周りも出すかもです(未定)。執筆・実装には Codex と Claude Code を使っています。

今回は Gated DeltaNet (GDN) と Kimi Delta Attention (KDA) の更新式を整理したうえで、学習済みの Olmo Hybrid を使い、状態に対する局所的な回帰の振る舞いと、状態の更新が後の回答にどう影響するかを調べます。

Attention

causal softmax attention の 1 head の出力は、位置 $t$ の query に対して

と書けます。過去の token ごとに KV を保存しておくので、後から来た query に応じて、その時点で必要な位置を参照できます。その代わり、系列長を $N$ とすると KV cache は $N$ に比例して増えていきます。

GDN と KDA は、過去の情報を固定サイズの状態 $S_t\in\mathbb{R}^{d_k\times d_v}$ に保持します。token ごとに $S_t$ を更新し、$o_t=S_t^\top q_t$ として読み出します。つまり、後からどんな query が来るかを知らないまま、有限の状態に何を残すかを決めなければなりません。

図 1. 共通の token 列と、各方式が保存するもの。softmax attention は token ごとの KV を保持し、GDN / KDA は固定サイズの状態を更新する。上下の対応は token の位置を表す (projection は方式ごとに異なる)。

Qwen3.8-Flash-Next は GDN と global softmax attention を、Kimi K3 は KDA と Gated Multi-head Latent Attention (Gated MLA) を組み合わせています（Qwen3.8-Flash-Next は full attention で学習を進め、continued pretraining の段階で full attention の層を Qwen Sparse Attention (QSA) に置き換えます。その際、まず full attention の attention distribution を QSA の indexer に蒸留してから、sparse attention での学習に移ります）。再帰的な層の状態は固定サイズですが、softmax attention の層では系列長に応じて KV cache が増えます (Qwen Team, 2026; Kimi Team, 2026)。

Gated DeltaNet と Kimi Delta Attention

Gated DeltaNet

Additive update では、associative memory を

として更新します。この場合、同じ key への書き込みも足し合わされ続けます。これに対して delta rule は、現在の状態から読み出した value と新しい value の差を使って補正します (Schlag et al., 2021)。

GDN は、この補正の前にスカラー gate $\alpha_t$ による減衰を加えます (Yang et al., 2025)。

$\alpha_t,\beta_t\in(0,1)$ は、head ごとに入力から決まる gate です。減衰は状態全体にかかり、delta rule は現在の key の方向だけを補正します。GDN はこの 2 つを組み合わせることで、不要になった文脈を消しつつ、個々の key–value の対応を書き換えられるように設計されています (Yang et al., 2025)。

図 2. GDN の 1 token 分の更新。$\widetilde S_t$ から residual と delta update を求め、最後に $\widetilde S_t$ を加えて $S_t$ を得る。枝分かれした線は、最後の加算に使う $\widetilde S_t$ を渡す経路を表す。

Qwen の GDN では、正規化した residual stream の入力を $x_t\in\mathbb{R}^d$ とし、学習される線形射影 $W_q,W_k,W_v$ と causal ShortConv から各 head の $q_t,k_t,v_t$ を作ります。

ここでは 1 head 分だけを書いており、$W_q,W_k\in\mathbb{R}^{d_k\times d}$、$W_v\in\mathbb{R}^{d_v\times d}$ です。ShortConv は channel ごとに独立にかかる短い causal convolution で、添字 $t$ はその出力の位置を表します。また Qwen の GDN では、各 head から読み出した出力に RMSNorm と sigmoid output gate をかけます (Qwen Team, 2026)。以下では、この出力側の処理より前の状態更新だけを扱います。

Kimi Delta Attention

KDA は、GDN のスカラーによる減衰を key channel ごとの減衰に拡張します (Kimi Team, 2025)。$D_t=\operatorname{Diag}(\boldsymbol\alpha_t)$ とすると、更新は

となります。$D_t=\alpha_t I$ とすれば GDN と一致します。

一般には $(I-\beta_t k_tk_t^\top)D_t\ne D_t(I-\beta_t k_tk_t^\top)$ なので、減衰と delta update の順序は交換できません。また、$k_tk_t^\top$ を通して channel 間の結合が生じます。

各 channel の decay logit を $z$、head ごとの learned log-scale を $A_h$ とすると、log-decay $g=\log\alpha$ の parameterization は次のように変わります。

Kimi K3 の log-decay は $-5<g<0$、retention factor は $e^{-5}<\alpha<1$ です。そのため 16 token の区間では累積 log-decay が $(-80,0)$ に収まり、chunkwise 計算で使う累積減衰の逆数は $e^{80}$ 未満になります (Kimi Team, 2026)。

図 3. $A_h=0$ での log-decay $g=\log\alpha$。破線は $g=-5$。Kimi K3 では log-decay に下限があり、retention factor は $\alpha>e^{-5}$ となる。学習後の logit の分布ではなく、同じ $z$ を入れたときの関数の形を比べている。

channel ごとに減衰率を持てるからといって、各 channel が特定の事実に対応するとは限りません。減衰の自由度が増えることと、意味的に独立した情報を選んで消せることは別の話です。

Online learning

GDN を提案した Yang et al. (2025) は、Section 3.1 の Table 1 で、fast weight matrix の更新を正則化付きの online optimization として定式化しています。Kimi Linear の Section 2.2 でも、DeltaNet を online gradient descent として、GDN をそれに weight decay を加えたものとして明示的に導出しています (Kimi Team, 2025)。

現在の token から作った $k_t,v_t$ を固定し、$S$ に対する損失を

と置くと、勾配は

です。GDN の delta update は

となり、線形回帰の online gradient descent に一致します。$S$ が回帰モデルの重み、$k_t$ が入力、$v_t$ が regression target、$\beta_t$ が learning rate に当たります。なお、ここでは $\alpha_t$ による減衰を先にかけ、減衰後の $\widetilde S_t$ で勾配を評価しています。減衰前の $S_{t-1}$ で勾配を評価してから weight decay を加える更新とは、順序が異なります。

KDA でも inner loop の回帰問題と delta update は同じで、$S_t=\widetilde S_t-\beta_t\nabla_S\ell_t(\widetilde S_t)$ と書けます。違うのは、勾配 step の前にかける減衰が、GDN ではスカラーの $\alpha_t I$、KDA では channel ごとの $D_t$ である点だけです。

ここでいう online learning は、系列を token の順に読みながら $S$ を更新していくことを指します。この状態更新は、事前学習でも推論でも forward の中で行われます。事前学習では、next-token prediction の損失の勾配を状態更新を通して流し、projection や gate などのパラメータを学習します。推論ではそれらのパラメータを固定したまま、入力に応じて $S$ だけを更新します。$v_t$ は現在の入力から作られるので、この更新に正解の次 token は要りません。

表 1. GDN / KDA における outer loop と inner loop。

このように、状態そのものをモデルの重みとみなし、forward の中で自己教師あり損失を使って更新する構成は、Test-Time Training (TTT) と共通しています。TTT では、この内側のモデルを線形モデルや MLP とし、inner loop の入力と target を作る射影も outer loop で学習します (Sun et al., 2024)。

局所損失

更新式から予測できる振る舞いを整理しておきます。$\|k_t\|_2=1$ のとき、更新直後の残差は

なので、同じ回帰損失は

になります。$\beta_t=1$ という極限では、その key に対する読み出しは $v_t$ に一致します。

また、delta update の大きさ (Frobenius norm) と、固定した任意の query $q$ への影響は

です。$k_t$ と直交する query には delta update の影響が及ばず、$k_t$ と重なる方向にだけ補正が伝わります。そのため、ある token での局所損失の変化からは、後で別の query で読み出したときに何が得られるかは分かりません。

ただし、これらはあくまで潜在表現の上での回帰の話です。同じ事実を言い換えて与えても、学習された key / value が同じ関係を表すとは限りません。回帰の残差が減ることが文章中の情報の更新に対応しているかどうかは、実際のモデルで確かめる必要があります。

Olmo Hybrid を使った観察

ここからは学習済みモデルで、前半で整理した状態更新が実際にどう使われているかを調べます。問いは 4 つです。1 つ目は、state が前半で見た key から value への回帰として働いているか。2 つ目は、その回帰の損失に passage の内容がどれだけ表れるか。3 つ目は、passage の書き込み、とくに答えに当たる token の書き込みを止めたときに、回答の尤度がどれだけ変わるか。4 つ目は、答えの情報が GDN の state と softmax attention のどちらの経路で、どれだけ運ばれているかです。

対象とする allenai/Olmo-Hybrid-Instruct-SFT-7B は、GDN と softmax attention を 3:1 の比率で積んだモデルです (Qwen3.8-Flash-Next の GDN と full attention、Kimi K3 の KDA と Gated MLA も同じ 3:1 です)。以下では、そのうち 24 層の GDN を計測します (Merrill et al., 2026; Qwen Team, 2026; Kimi Team, 2026)。

表 2. Checkpoint と実行条件 (config)。

Olmo Hybrid Instruct SFT 7B の GDN は負の固有値 (negative eigenvalues) を許す拡張を使っていて (上で紹介した GDN では $\alpha_t,\beta_t\in(0,1)$)、$\beta_t\in(0,2)$ です。$\beta_t>1$ では、更新後の回帰残差の符号が反転することがあります。また、上の「局所損失」では $\|k_t\|_2=1$ を仮定しましたが、測定ではこの checkpoint が実際に出力する gate の値 $\alpha_t,\beta_t$ と key $k_t$ をそのまま使います。

文書 QA

まず、文書を読んで答える QA の基本性能を測るため、Stanford Question Answering Dataset (SQuAD) v1.1 の dev set (以下、SQuAD dev) 全 10,570 問を評価します。データは Wikipedia の 48 記事から取った passage と、それぞれに対応する質問からなります。入力では passage の後に質問を置き、batch size 32 の greedy decoding で最大 1,024 token まで生成しました。SQuAD 自体には文書を与えない評価条件はないので、比較のために passage を取り除いた「文書なし」の条件を独自に加えています (Rajpurkar et al., 2016)。

採点は SQuAD 公式の exact match (EM) と token F1 に従います。小文字にそろえ、句読点と英語の冠詞を取り除き、空白を整えたうえで、EM では回答全体が一致するか、token F1 では空白で区切った単語がどれだけ重なるかを測ります。正答の annotation が複数ある設問では、それぞれに対するスコアの最大値を取ります。例えば正答 thousands に対する thousands of years は EM = 0、token F1 = 0.5 です。この採点では、同義語や単数形・複数形は別物として扱われます (公式の採点コード)。

表 3. SQuAD dev 全 10,570 問の結果。回答全体を採点し、EM・平均 token F1 は % で示す。「上限到達」は、EOS を出す前に 1,024 token に達した件数。

文書なしの 7 問は、質問や The short answer is: などを繰り返し続けて生成上限に達したもので、表 3 では保存された出力をそのまま採点しています。文書なしの条件でも「与えた文書だけを使う」という指示は残しているので、答えられなくても、モデルの知識不足とは限りません。

State 解析の設問

state の解析には、SQuAD dev から、1 つの passage と 3 つの質問を組み合わせたものを規則で選びました。以下では、この 3 問のまとまりを triplet と呼びます。各 triplet は、同じ事実を言い換えて問う 2 問 (Q1、Q2) と、同じ passage の別の事実を問う 1 問 (Q3) からなります。以下では、passage の一部の書き込みを止めたり差し替えたりして、正答の尤度がどう変わるかを見ます。Q1 と Q2 で同じように変われば、その変化は特定の質問文によるものではなく、言い換えても共通する事実に関わるものだと考えられます。Q3 は、変化が passage 全体に及ぶのか、対象の事実に限られるのかを比べるための対照です。

Q1 と Q2 では、すべての正答 annotation が passage 中の同じ 1 箇所を指しています。その正答に当たる token の位置の集合を $\mathcal{A}$ とし、区間 $\mathcal{A}$ と呼びます。

後半の activation patching では、元の passage $P$ に加えて、区間 $\mathcal{A}$ を別の語に置き換えた passage $P'$ を使います。そのため、置き換えができる triplet だけを選びます。置き換えは Longpre et al. (2021) の corpus substitution に倣い、SQuAD dev のほかの記事の質問の正答から、同じ種類の語を選びます。ただし patching では、$P$ と $P'$ で token の位置がそろっている必要があるので、token の数が同じ語に限ります。

表 4. triplet の選び方と、各段階で残った triplet の数。

固有表現抽出には spaCy の en_core_web_trf を使い、種類は Longpre et al. と同じ 5 種類にまとめました。置き換える語は、段階 5 の候補から乱数の種を固定して選びます。人物の姓のように、区間 $\mathcal{A}$ と同じものを指す語が passage のほかの場所にも出てくる場合は、同じ語で置き換えます。段階 7 は $P'$ を読んで判定し、patching の結果を見る前に triplet を固定しました。

最終的な 31 triplet (93 問) の答えの種類は、人物 15、数 7、日付 6、組織 3 です。場所を答えとする triplet は段階 7 で残りませんでした (国の位置に都市が入るなど、置き換えで場所の粒度がずれたため)。記事は 9 本で、そのうち Super_Bowl_50 が 17 triplet を占めます。

カードの 3 問は、それぞれ独立した入力として評価します。どの forward も同じ初期 state から始め、passage → question → answer の順に処理します。Q1 の回答や state を Q2 に引き継ぐことはありません。3 問とも passage は共通なので、区間 $\mathcal{A}$ を読み終えるまでの入力と state の更新は同じです。Other fact でも区間 $\mathcal{A}$ はそのままにして、別の正答を問います。

図 4. 3 問の評価。初期 state $S_0$ から区間 $\mathcal{A}$ を含む passage を読んだ state $S_P$ は 3 問で共通で、そこから質問ごとに分かれて回答する。各質問は独立に評価し、ほかの質問の回答や state は引き継がない。色は図 6・図 9 の系列に対応する。

State と局所回帰損失の測り方

各設問について、全 24 層の GDN・全 30 head で、モデルが次の token に渡す recurrent state を取り出しました。layer の番号は 0 始まりで、3、7、11、15、19、23、27、31 は softmax attention の層です。計測用の処理を入れても、最終 logits は入れない場合と全 token で bitwise に一致しました。

局所回帰損失は、delta update 前の state で測ります。質問の token の位置の集合を $\mathcal{Q}$ とし、質問の各 token を読む時点で、同じ $k_t,v_t$ に対する $\ell_t(\widetilde S_t)$ を比べます。比べる相手の state を $\widetilde S'_t$ として、設問・layer ごとに $\mathcal{Q}$ と全 head について損失を足し合わせてから比を取ります。

$R<1$ なら、通常の state のほうが比べる相手よりも現在の $v_t$ をよく再現しています。比べる相手には次の 4 つを使います。括弧内は図 5 での表記です。

• 何も書き込まれていない state ($S=0$)：$\ell_t(0)=\frac12\lVert v_t\rVert_2^2$ です。
• 質問より前の書き込みを除いた state (prefix removed)：その layer だけで、質問より前のすべての token について $\beta_t=0$ とします。passage のほかに、chat template や Question: のような定型の token も含みます。passage に限らず、質問より前の文脈全体がどれだけ寄与しているかの上限を見るための比較です。
• passage の書き込みを除いた state (passage removed)：その layer だけで、passage の token について $\beta_t=0$ とします。
• 区間 $\mathcal{A}$ の書き込みを除いた state ($\mathcal{A}$ removed)：その layer だけで、$t\in\mathcal{A}$ について $\beta_t=0$ とします。

$\beta_t=0$ にしても $\alpha_t$ による decay は残し、質問の token 自身の更新もそのままにします。layer $L$ の $k_t,v_t,\alpha_t,\beta_t$ はそれより下の層だけで決まるので、layer $L$ だけで書き込みを止めても、layer $L$ の入力は変わりません。実際に、書き込みを止めた forward で得た state と、通常の forward で得た $k_t,v_t,\alpha_t,\beta_t$ から更新式を計算し直した state は、全条件で相対誤差 $6.3\times10^{-8}$ 以内で一致しました。

State は回帰として働いているか

$S=0$ と比べると、93 問 × 24 層の 2,232 条件のうち 2,226 条件で $R<1$ でした。中央値は 0.48、10–90 パーセンタイルは 0.17–0.76 です。学習済みモデルでも、state の読み出し $\widetilde S_t^\top k_t$ はゼロベクトルよりも現在の $v_t$ に近く、中央値では残差の大きさ $\lVert v_t-\widetilde S_t^\top k_t\rVert_2^2$ が $\lVert v_t\rVert_2^2$ の半分ほどまで小さくなっています。前半で見た「状態更新は key から value への回帰」という読み方は、学習済みモデルでも成り立っています。

ただし、$S=0$ は何も覚えていない state との比較です。この結果だけでは、state のどの書き込みがこの改善を生んでいるのかは分かりません。

局所回帰損失は passage を反映しているか

比べる相手を、書き込みを除いた state にすると、$R$ は 1 に近づきます。

表 5. 局所回帰損失の比 $R$ (31 triplet・93 問 × 24 層の 2,232 条件)。

図 5. layer ごとの $R$ の中央値 (31 triplet・93 問)。比べる相手ごとに 1 本の線で示す。

$S=0$ との差の大部分は、質問より前の書き込みを除いても残ります。passage の書き込みを除いたときの変化はさらに小さく、区間 $\mathcal{A}$ の書き込みを除いたときの $R$ はすべて 0.987–1.013 の範囲に収まります。質問の token での回帰の改善は、主に質問の token 自身がそれまでに書き込んだ内容によるもので、passage、とくに答えに当たる token の書き込みはほとんど寄与していません。

この損失は、現在の token から作った latent value をどれだけ再現できるかを測るもので、質問への正答を直接測るものではありません。以下では、回答の尤度で書き込みの寄与を調べます。

回答は passage の書き込みに依存しているか

指定した SQuAD の正答 token 列を $a_r=(a_{r,1},\ldots,a_{r,M})$ とし、teacher forcing で

を計算します。各時点の最終 logits に log_softmax をかけ、正答 token の成分を取り出して足し合わせます。書き込みを止めた条件と通常の条件の差を $\Delta\log p$ とします。$\exp(\Delta\log p)$ は正答列の尤度比なので、例えば $\Delta\log p=-\log 2$ は尤度が半分になったことを表します。正答列は teacher forcing で与えていて、書き込みを止めた条件で自己回帰生成はしていません。そのため $\Delta\log p$ はモデルが正答列をどれだけ支持するかの変化を表しますが、生成される回答文字列が変わるかどうかまでは分かりません。

24 層の GDN すべてで passage の token の書き込みを止める ($\beta_t=0$) と、31 triplet 93 問のすべてで正答列の尤度が下がりました。$\Delta\log p$ の中央値は −7.24 nats で、最も小さい変化でも −0.73 nats です。

図 6. 24 層の GDN すべてで passage の書き込みを止めたときの、質問ごとの正答列の $\Delta\log p$。

このモデルの softmax attention の 8 層は、passage の token を KV cache から直接読めます。それでも GDN への passage の書き込みを止めると、回答の尤度は大きく下がります。上で見たとおり、passage の書き込みは質問の token での局所回帰損失にはあまり表れません。局所回帰損失の変化が小さくても、その書き込みが後で使われていないとは限らないようです。

ただし、passage 全体の書き込みを止めると、GDN の state は通常の入力では現れない状態になります。尤度の低下が「passage の情報がなくなったこと」によるのか、「見たことのない state でモデル全体の振る舞いが崩れたこと」によるのかは、この操作だけでは区別できません。別の事実を問う Q3 でも尤度は同じように下がっていて (中央値 −7.55 nats)、この操作は特定の事実に絞ったものにはなっていません。

答えの書き込みを止める

ここで確かめたい仮説は、passage を読むときの区間 $\mathcal{A}$ の delta update が対象の事実を state に書き込み、その情報が同じ入力の後半にある質問への回答に使われる、というものです。この仮説が正しければ、区間 $\mathcal{A}$ の書き込みを止めると、独立に評価した 2 つの質問でどちらも同じ正答の尤度が下がり、別の事実を問う質問への影響は相対的に小さくなるはずです。

図 7. 仮説と、書き込みを止める操作。各 token は delta update で固定サイズの state $S$ に書き込み、同時に $S$ から読み出す ($o_t=S_t^\top q_t$)。図では passage の書き込みと質問の読み出しだけを描いている。区間 $\mathcal{A}$ の token だけ $\beta=0$ として書き込みを止めると、同じ事実を問う 2 問では正答列の尤度が下がり、別の事実を問う 1 問ではほとんど変わらないと予想される。

24 層の GDN のすべての head について、$t\in\mathcal{A}$ のときだけ $\beta_t=0$ にします。

このとき区間 $\mathcal{A}$ では $S_t^{(\ell)}=\widetilde S_t^{(\ell)}=\alpha_t S_{t-1}^{(\ell)}$ となり、state には decay だけがかかります。同じ事実を問う 2 問の $\Delta\log p$ を $\Delta_1,\Delta_2$ とし、別の事実を問う質問の $\Delta\log p$ を $\Delta_c$ として、triplet ごとに

を計算します。$D>0$ なら、対象の事実の正答列の尤度のほうが、別の事実の正答列の尤度より下がったことになります。書き込みを止めること自体の影響と区別するため、各 passage で区間 $\mathcal{A}$ と同じ長さで、答えと関係のない区間 $\mathcal{C}$ を 1 つ選び、その token についても同じように $\beta_t=0$ にして書き込みを止めます。区間 $\mathcal{A}$ の書き込みを止めたときの $D$ を $D_{\mathcal{A}}$、区間 $\mathcal{C}$ の書き込みを止めたときの $D$ を $D_{\mathcal{C}}$ と書きます。区間 $\mathcal{C}$ は区間 $\mathcal{A}$ から 4 token 以上離し、3 問のどの正答 annotation とも重ならない位置にしました。

31 triplet のすべてで $D_{\mathcal{A}}$ は正で、$D_{\mathcal{C}}$ を上回りました。$D_{\mathcal{A}}$ の中央値は 1.30 nats (10–90 パーセンタイルは 0.11–7.80)、$D_{\mathcal{C}}$ は 0.0006 nats (−0.15–0.10) です。区間 $\mathcal{A}$ を止めたときの $\Delta\log p$ の中央値は、Q1 が −1.12、Q2 が −1.26、Q3 が +0.006 nats でした。

表 6. 記事ごとの $D_{\mathcal{A}}$ と $D_{\mathcal{C}}$ の中央値 (nats)。どちらも 24 層の GDN すべてで書き込みを止めた。

$D_{\mathcal{C}}$ はどの記事でもほぼ 0 で、答えと関係のない区間の書き込みを止めても、対象の事実の答えだけが下がることはありません。$D_{\mathcal{A}}$ は triplet によって 0.02 から 12.7 nats まで大きく異なりますが、どの記事でも $D_{\mathcal{C}}$ より大きくなっています。

ただし、この操作も書き込みをゼロにするもので、state を通常の入力では現れない状態にします。また、区間 $\mathcal{A}$ の情報は softmax attention の層からも読めます。GDN の state と softmax attention のどちらの経路で、答えの情報がどれだけ運ばれているかは、この操作からは分かりません。続いて、この 2 点を activation patching で調べます。

答えの情報はどの経路で運ばれるか

activation patching では、2 つの入力でモデルを走らせ、一方の途中の値をもう一方の値に差し替えて、出力がどれだけ変わるかを見ます (Heimersheim & Nanda, 2024)。値をゼロにする代わりに、実際の入力から得た値に差し替えるので、差し替えた後の state は通常の入力で現れる state のままです (Zhang & Nanda, 2023)。

ここでは、元の passage $P$ と、区間 $\mathcal{A}$ を置き換えた passage $P'$ を使います。$P$ での正答を $r$、置き換えた語を $r'$ とし、2 つの答えの対数尤度の差 (log-probability difference) を

と書き、teacher forcing で計算します。以下では LD と呼びます。同じ事実を問う Q1・Q2 の 62 問はすべて、$P$ では $\mathrm{LD}>0$、$P'$ では $\mathrm{LD}<0$ でした。モデルは passage に書かれた答えに従っています。

差し替える経路は 2 つです。

• GDN の state の経路：置き換えた token の区間を読み終えた直後に、24 層の GDN の state と conv cache を、もう一方の入力の値に差し替えます。conv cache は、短い畳み込みが直前 3 token 分の入力を保持しておくキャッシュです (実装では conv state と呼ばれます)。畳み込みは kernel 幅 4 なので、区間の直後の 3 token にも区間の情報が入るためです。
• softmax attention の経路：8 層の softmax attention の KV cache のうち、置き換えた token の位置の key と value を差し替えます。

図 8. activation patching で差し替える 2 つの経路。$P$ と $P'$ の 2 つの入力でモデルを走らせ、区間の直後の GDN の state と conv cache、または区間の位置の key と value だけを、一方の run からもう一方の run へ差し替える。denoising は $P'$ の run に $P$ の値を、noising は $P$ の run に $P'$ の値を入れる。区間の位置の residual stream は差し替えない。

どちらの場合も、置き換えた token の位置の residual stream は差し替えません。このため、区間の情報が後ろの token に届く経路は、GDN の state と conv cache と KV cache の 3 つだけになります。3 つをすべて差し替えると、出力は差し替え元の入力と bitwise に一致します。特定の経路だけを差し替えるこの方法は、path patching と呼ばれるものの一種です (Goldowsky-Dill et al., 2023)。

差し替えの向きは 2 つ測ります。$P'$ で走らせながら $P$ の値に戻す向き (denoising) と、$P$ で走らせながら $P'$ の値を入れる向き (noising) です。前者はその経路だけで答えを戻せるか、後者はその経路が欠けると答えが崩れるかを表します (Heimersheim & Nanda, 2024)。それぞれ、$\mathrm{LD}$ がどれだけ元に戻ったか (崩れたか) を、$P$ と $P'$ の $\mathrm{LD}$ の差で割った割合で表します。

表 7. 経路ごとの回復の割合 (Q1・Q2 の 62 問の中央値)。1 なら、答えが差し替え元のものに完全に入れ替わる。

図 9. 問ごとの回復の割合。左は GDN の state の経路と softmax attention の経路を、denoising と noising のそれぞれで示す (縦線は中央値)。右は denoising での 2 つの経路の回復の割合を問ごとに並べたもので、斜めの線は和が 1 になる位置。

区間 $\mathcal{A}$ の中身は、多くが softmax attention の経路で答えに届き (中央値 0.84)、GDN の state の経路を通るのは 2 割弱です (中央値 0.16)。GDN の経路の割合は問によって大きく違い、−0.03 から 0.49 まで広がっています。2 つの経路の回復の割合を問ごとに足すと、denoising で 0.80–1.19、noising で 0.81–1.20 になり、denoising と noising の値もほぼ同じです。片方の経路だけを壊しても、もう片方が補って答えが保たれる、という関係にはなっていません。2 つの経路は、それぞれの分だけ答えに寄与しています。conv cache を含めると GDN の経路の割合が増える (0.14 → 0.16、0.12 → 0.16) ので、短い畳み込みも区間の情報を後ろに運んでいます。

別の事実を問う Q3 では、どの差し替えでも正答の $\log p$ はほとんど変わりませんでした (変化の絶対値の中央値 0.01 nats、最大 0.33 nats)。

1 つの層だけを差し替えると、GDN の層ごとの回復の割合は、中央値ではどれも小さくなりました。最も大きい layer 26 でも 0.013、次の layer 29 で 0.011 です。ただ、これは問ごとのばらつきをならした値です。GDN の経路での回復がある程度ある 50 問 (24 層の和が 0.02 を超える問) では、その大部分 (中央値で 95%) が layer 16 以上の層から来ています。最も効く 1 層は layer 26 (62 問中 28 問) と layer 29 (13 問) が多いものの、残りの 3 分の 1 ほどの問ではそれ以外の層でした。24 層それぞれの回復の割合を問ごとに足すと中央値は 0.12 で、全層をまとめて差し替えたときの 0.16 に近い値です。GDN の各層の寄与は、ほぼ足し算で合わさっています。

softmax attention では、layer 23 と layer 27 がそれぞれ 0.23 と 0.22、layer 19 が 0.08 で、この 3 層が大部分を担っています。8 層それぞれの回復の割合を問ごとに足した中央値は 0.54 で、全層をまとめて差し替えたときの 0.84 を下回ります。softmax attention の層は、まとめて差し替えると 1 層ずつの和より大きく効くので、層どうしが関係し合いながら答えを運んでいるようです。

GDN の層は softmax attention の層の 3 倍あるので、層ごとの値を比べるときは層の数の違いを考える必要があります。1 層ずつの和を層の数で割ると、1 層あたりの回復の割合は GDN が 0.005 (0.12/24)、softmax attention が 0.068 (0.54/8) で、softmax attention の層のほうが 10 倍以上大きくなります。1 層あたりの寄与が同じなら GDN の経路が全体の 4 分の 3 を担うはずですが、実際には 2 割弱です。層の数の違いを考えに入れても、答えの情報を主に運んでいるのは softmax attention の層です。

図 10. layer ごとの回復の割合。点は 62 問の中央値、縦の帯は 10–90 パーセンタイル。GDN の層では state と conv cache を、softmax attention の層では key と value を、その 1 層だけ差し替えた。

上で、区間 $\mathcal{A}$ の書き込みを止めると答えの尤度が特異的に下がることを見ました。activation patching の結果と合わせると、GDN の state は区間 $\mathcal{A}$ の中身を実際に後ろの token へ運んでいますが、その量は softmax attention の経路より小さい、ということになります。

Limitations

• データの範囲が狭く、偏っている
  • 使ったのは SQuAD dev だけです。英語の Wikipedia の段落を読み、答えが passage 中の語句そのものである抽出型の QA に限られます。ほかの形式の文章や、答えを組み立てる必要がある質問に同じ結果が当てはまるかは分かりません。
  • triplet は 31 個 (93 問) で、記事は 9 本しかなく、半分以上が Super_Bowl_50 です。答えの種類も人物、数、日付、組織だけで、場所は含まれていません。数値は、この狭い範囲での中央値として読む必要があります。
  • triplet は規則と、置き換えた passage を読んでの判定で選びました。判定は結果を見る前に固定しましたが、自然に置き換えられる答えだけが残っているので、置き換えにくい答えでの振る舞いは測れていません。
• 何を測っているかに制約がある
  • $\mathcal{A}$ は正答に当たる token だけです。答えの文字列そのものを読むことと、事実を覚えておくことは、この設計では区別できません。
  • $P'$ の置き換え先は実在の語なので、モデルがもともと持つ知識と食い違います。モデルは 62 問すべてで passage の答えに従いましたが、この食い違いが経路の割合に影響している可能性はあります。
  • 書き込みを止める操作 ($\beta_t=0$) は、通常の入力では現れない state を作ります。activation patching はこの問題を避けますが、差し替えたのは置き換えた token の位置から後ろへ情報を運ぶ経路だけで、同じ位置の上の層に届く情報は差し替えていません。
• 一般化できる範囲が限られる
  • 調べたのは Olmo Hybrid 1 つだけです。KDA を使う Kimi K3 や、softmax attention を QSA に置き換えた Qwen3.8-Flash-Next で同じ結果になるかは分かりません。
  • triplet どうしや記事どうしの相関があるので、統計的な検定や信頼区間は付けていません。すべての値を点で示しています。

まとめ

今回、GDN / KDA を式のレベルまで整理したことで、アーキテクチャのお気持ちが少し分かった気がします。最初は softmax attention の key / value と同じような役割だと思っていましたが、token ごとに KV を残すのではなく、key / value を固定サイズの state への書き込みに使うという捉え方の違いがあることが分かり納得 (ほんとか?)。

state の更新が online での SGD の 1 step に対応していることにも、今回初めて気づきました。単なる recurrent update として眺めていましたが、forward の中で小さな回帰問題を token ごとに解いていると捉えると見え方はおもしろいと思いました。

関連研究

Sun et al. (2024) の Figure 4 は、125M の学習済みモデルについて、初期状態 $W_0$・更新前 $W_{t-1}$・更新後 $W_t$ の TTT 損失を、token の進行に沿って比べています。全 test sequence の平均を示した図で、今回の $\ell_t(0)$ と $\ell_t(\widetilde S_t)$ の比較に近いものです。

recurrent state への書き込みを直接扱った研究に、Young (2026) の WriteSAE があります。状態への行列更新を rank-one の atom で表し、実際の書き込みを atom に置き換えたときの出力を評価しています。対象は Qwen3.5-0.8B などの学習済みモデルで、複数の layer・head で置き換えを検証しています。今回の操作は特定の区間で更新を止めるだけで、WriteSAE のように更新を意味のある成分に分解して置き換えるものではありません。

activation patching の方法は Zhang & Nanda (2023) と Heimersheim & Nanda (2024) に倣いました。前者は、ノイズを加えたり値をゼロにしたりする操作が分布の外の状態を作ることを示し、同じ長さの別の語に置き換えた入力を使うことを勧めています。後者は、denoising と noising がそれぞれ何を示すかを整理しています。特定の経路だけを差し替える方法は Goldowsky-Dill et al. (2023) の path patching にあたります。

再帰的な状態を持つモデルに activation patching を使った例として、Sen Sharma et al. (2024) は Mamba で事実の想起を調べています。この研究では状態の行列そのものではなく、token の位置ごとの出力を差し替えています。今回は GDN の state と conv cache を直接差し替えました。

文脈に書かれた情報とモデルがもともと持つ知識のどちらが使われるかを、activation patching で調べた研究に Yu et al. (2023) があります。state space model や hybrid model の内部を、値を書き換えて調べた研究には Wang et al. (2025) があり、文脈内学習に関わる成分が attention と Mamba の層にあることを報告しています。

実装

実装と計測結果は整理してから公開する予定です。

References

• Kimi Team. (2026). Kimi K3: Open Frontier Intelligence.
• Qwen Team. (2026). On the Design of Qwen3.8-Next Architecture: Evaluation, Efficiency, and Training Stability.
• Goldowsky-Dill, N., et al. (2023). Localizing Model Behavior with Path Patching.
• Heimersheim, S., & Nanda, N. (2024). How to use and interpret activation patching.
• Kimi Team. (2025). Kimi Linear: An Expressive, Efficient Attention Architecture.
• Longpre, S., et al. (2021). Entity-Based Knowledge Conflicts in Question Answering. EMNLP.
• Merrill, W., et al. (2026). Olmo Hybrid: From Theory to Practice and Back.
• Rajpurkar, P., et al. (2016). SQuAD: 100,000+ Questions for Machine Comprehension of Text.
• Schlag, I., et al. (2021). Linear Transformers Are Secretly Fast Weight Programmers. ICML.
• Sen Sharma, A., et al. (2024). Locating and Editing Factual Associations in Mamba.
• Sun, Y., et al. (2024). Learning to (Learn at Test Time): RNNs with Expressive Hidden States.
• Wang, S., et al. (2025). Understanding In-Context Learning Beyond Transformers: An Investigation of State Space and Hybrid Architectures.
• Yang, S., et al. (2025). Gated Delta Networks: Improving Mamba2 with Delta Rule. ICLR.
• Young, J. (2026). WriteSAE: Sparse Autoencoders for Recurrent State. arXiv preprint, v4.
• Yu, Q., et al. (2023). Characterizing Mechanisms for Factual Recall in Language Models. EMNLP.
• Zhang, F., & Nanda, N. (2023). Towards Best Practices of Activation Patching in Language Models: Metrics and Methods. ICLR 2024.