🤖 machine learning

Kaczmarz Linear Attention

本論文は、経験的に学習された更新係数を理論的に導出されたキーノルム正規化カチャマツステップサイズに置き換えた修正版ゲートドデルタネットであるカチャマツ線形アテンション（KLA）を導入し、モデルのアーキテクチャや状態形状を変更することなく、優れたパープレキシティ、長文脈安定性、およびデコーディング効率を実現する。

原著者： Jiaxuan Zou, Ruifeng Ren, Yong Liu

公開日 2026-05-12

📖 1 分で読めます☕ さくっと読める

CC BY 4.0

原著者： Jiaxuan Zou, Ruifeng Ren, Yong Liu

原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 ✨ これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。免責事項の全文を読む

ロボットに非常に長い本を読ませることを想像してください。ロボットは現在の文を理解するために、以前読んだ内容を記憶しておく必要があります。

問題：「二次関数的」なボトルネック
従来の AI モデル（トランスフォーマー）は、新しい単語を読むたびに、すべての過去の単語とのつながりを確認するために、全体の本を振り返って調べる学生のように動作します。本が短ければ問題ありません。しかし、本が 10 万ページもある場合、学生は単語一つ一つに対して膨大な作業を強いられます。これはあまりにも遅く、高価になり、実質的に拡張することが不可能になります。

解決策：「再帰的状態」
新しいモデルは、この問題を修正するために、固定サイズの小さなノートを持つ学生のように動作しようとします。本全体を振り返る代わりに、読むにつれてノートを更新します。最も重要な部分を書き留め、それ以外は忘れ、動き続けます。これは高速（線形時間）ですが、正しく行うのは困難です：何を書き留めるべきか？どれくらい消すべきか？そして、同じトピックが再び現れた場合、どのようにノートを更新すべきか？

以前の試み：ゲート付きデルタネット（GDN）
ゲート付きデルタネット（GDN）と呼ばれる人気のあるモデルは、「ノート」アプローチを使用します。新しい情報を見ると、自分が知っていると思っていることと、実際に見ていることとの差を計算し、その差をノートに書き込みます。

しかし、GDN には欠点があります。それは「学習された推測」（トレーニング中に導き出される数値）を使用して、どの程度の大きさの変更を行うかを決定することです。まるで学生が、「うーん、この内容は 5 号のマーカーで書き留めるべきかな」と推測しているようなものです。時にはページを滲ませるほどマーカーが大きすぎたり、逆に文字が薄くて見失われるほど小さすぎたりします。この推測は、モデルが学習した習慣に過ぎず、数学的な規則ではありません。

新しいアイデア：カチャマールツ線形アテンション（KLA）
この論文の著者、Jiaxuan Zou とその同僚たちは、「推測を止めて、変更の大きさを正確に決定するために数学を使えないか？」と問いかけました。

彼らは、カチャマールツ射影と呼ばれる古い数学的手法に注目しました。

アナロジー： 紙の上に特定の点を通る線を引こうとしていると想像してください。あなたは定規（現在の状態）を持っています。定規がその点に当たらない場合、定規を少しずらす必要があります。
洞察： カチャマールツ法によれば、定規をずらす最良の方法は、その点の「強さ」や「大きさ」を測定することです。点が非常に「大きい」（強い信号）場合、その点に合わせるためにはわずかなずらしで済みます。点が「小さい」（弱い信号）場合、大きなずらしが必要です。

論文の用語では、「キー」（信号）を見て、その強さ（ノルム）を測定します。そして、正確なステップサイズを計算します：

ステップサイズ = （学習率）/（信号の強さ）

これがカチャマールツ係数です。

何が変わったのか？
著者たちは新しいロボットや新しいノートを作ったわけではありません。ハードウェアも変更していません。単に、GDN モデル内の「推測」の数値を、この正確で数学的に導き出された数値に置き換えただけです。

旧方式： 「トレーニングデータがそう指示しているので、この内容は 0.5 のマーカーサイズで書き留めます」
新方式（KLA）： 「この信号の大きさで割った 0.5 のマーカーサイズで書き留めます」

結果
この新しい規則は、メモリ更新という特定の作業に対して数学的に完璧であるため、モデルのパフォーマンスが向上しました：

より賢い： 前の文の次の単語を、以前の最高性能モデルよりも正確に予測します（「パープレキシティ」が低い）。
より長い記憶： 以前のモデルが混乱し始めたり忘れ始めたりするのに対し、65,000 語までの非常に長いコンテキストを処理できます。
タスクへの優位性： 巨大な「干し草の山」の中から特定の「針」を見つけるテストでは、KLA は 100% 正解し、他のモデルは見逃しました。
同等の速度： ノートの構造ではなく、更新の数学式のみを変更したため、モデルは以前と同じ速度で動作します。実際、長い長さでのデコード（テキスト生成）は 2.1 倍高速です。

まとめ
この論文は、KLAを紹介しています。これは、先行モデルと同じ高速で効率的な構造を維持しつつ、「推測」に基づく更新規則を、正確で数学的に導き出されたものに置き換えたモデルです。これは、すでに良く走る車を運転手の推測から、完璧な GPS 航法システムに乗り換えるようなものです。車は同じですが、目的地への到達がより正確かつ効率的になります。

技術概要：カチャルツ線形アテンション（KLA）

問題定義

長文脈言語モデル化は、Transformer における標準的な Softmax アテンションの二次的な計算コストとメモリコストによって阻害されている。線形時間の再帰モデルは、文脈を固定サイズの状態に圧縮することでこれを解決するが、状態更新則の設計が決定的なボトルネックとなる。Gated DeltaNet（GDN）などの既存のアプローチは、「デルタ則」を用いて目標値と現在の状態予測との間の残差を書き込むことで、既存の関連性を単に蓄積するのではなく編集する。

しかし、GDN は忘却と更新のバランスを取るために学習されたスカラー係数に依存している。著者らは、この係数が理論的に導出されたステップサイズではなく、経験的な設計選択であると主張する。不適切な係数は、サブオプティマルな更新の大きさをもたらす可能性がある。すなわち、ノルムが大きいキーに対して過剰に修正し、ノルムが小さいキーに対して修正不足となり、結果として状態の破損とサブオプティマルな検索性能を引き起こす。

手法：カチャルツ線形アテンション

本論文は、学習されたスカラーをカチャルツ射影法に基づいて理論的に導出されたステップサイズに置き換える GDN の改良版である**カチャルツ線形アテンション（KLA）**を提案する。

理論的導出

著者らは、キー・バリューの書き込み操作を制約付き最適化問題として定式化する。トークン $t$ において、目的は、線形制約 $S_t^T k_t = v_t$ （ここで $k_t$ はキー、 $v_t$ はバリュー）を満たす、減衰した前状態 $\tilde{S}_{t-1}$ に最も近い状態 $S_t$ を見つけることである。

最適化目的: $\min_S \frac{1}{2}\|S - \tilde{S}_{t-1}\|_F^2$ 制約条件 $S^T k_t = v_t$ 。
解: 最適解は、制約超平面への $\tilde{S}_{t-1}$ の直交射影である。これにより、ステップサイズ（係数） $\beta_t = 1/\|k_t\|_2^2$ が得られる。

KLA 更新則

KLA は、この正確な射影を、学習可能な緩和ゲート $\eta_t$ と数値安定化項 $\epsilon$ を含む実用的な更新則に一般化する。
$\beta_t = \frac{\eta_t}{\|k_t\|_2^2 + \epsilon}$
状態更新は以下のようになる。
$S_t = \tilde{S}_{t-1} + \beta_t k_t e_t^T$
ここで、 $e_t = v_t - \tilde{S}_{t-1}^T k_t$ は残差誤差である。

主要な特性

本論文は、この係数が（ $\eta_t=1, \epsilon=0$ の緩和されていない場合において）同時に 3 つの望ましい特性を満たすことを立証している。

直交射影: 制約を満たすために状態に対して最小ノルムの変更を行う。
正確な線探索: 更新方向に沿ったトークンごとの損失の正確な最小化器である。
正規化勾配ステップ: 正規化された確率的勾配降下（SGD）ステップとして機能する。

実装

KLA は GDN のためのドロップイン代替として設計されている。以下を維持する。

状態の形状（ $S \in \mathbb{R}^{d_k \times d_v}$ ）。
ゲーティング機構（減衰ゲートと更新ゲート）。
線形再帰構造。
チャンク単位並列トレーニングアルゴリズム。

構造的な変更は、チャンク単位ソルバ内での対角係数行列 $B = \text{Diag}(\beta_1, \dots, \beta_C)$ の計算のみである。新しいハードウェアカーネルや状態レイアウトは不要である。

主要な貢献

自己正規化デルタ則: カチャルツ射影からの KLA 係数の導出。それがトークンごとの損失に対する直交射影、正確な線探索最小化器、および正規化された勾配ステップであることを証明。
ドロップイン再帰アルゴリズム: KLA が GDN の再帰とチャンク単位ソルバを維持し、対角係数行列の変更のみを必要とすることを示す。
経験的優位性: 構造の拡張なしに、GDN や他の線形時間ベースラインよりも低いパープレキシティと優れた状態追跡を達成することを検証。
維持された効率性: KLA が $O(n)$ の時間/メモリ複雑性を維持し、GDN のプレフィル遅延と同等でありながらデコードスループットを改善することを確認。

実験結果

実験は、SlimPajama-672B データセット上で、0.4B パラメータ規模、1B トークン予算で行われた。

言語モデル化

検証パープレキシティ（PPL）: KLA は8.09 PPLを達成し、GDN（8.50）、DeltaNet（8.26）、GLA（8.42）、Longhorn（8.79）、Mamba2（9.41）を上回った。
長文脈外挿: 長文脈微調整なしで、KLA は 1K から 65K トークンまでの文脈長全体で最低のパープレキシティを維持し、GDN に対して一貫して 3.7–4.1 PPL の優位性を示した。

合成タスク

シングル・ニードル・イン・ア・ヘイスタック（S-NIAH）: KLA はテストされたすべての文脈長（1K–8K）で100% の精度に達したのに対し、GDN はわずかに低下し、Mamba2 は著しく失敗した。
マルチクエリ連想想起（MQAR）: KLA は、GDN に対して 8 倍の長さ外挿精度を7.03 ポイント改善した（73.84% 対 66.81%）。
スタックと回文: KLA は状態追跡および反転タスクにおいて競争力のある、あるいは優れた性能を示した。

効率性

プレフィル遅延: KLA は GDN のプレフィル遅延と一致した（例：131K 文脈で約 1220ms）。これにより、カチャルツ係数の計算が支配的なチャンク単位カーネルコストに影響を与えないことが確認された。
デコードスループット: 32K 文脈において、KLA は3.54 トークン/秒を達成し、GDN（1.65 トークン/秒）の2.1 倍、Mamba2 の 1.7 倍であった。

除去実験

著者らは正規化戦略、ゲーティング機構、状態拡張を除去実験した。結果は、キーノルム正規化が安定性にとって決定的であることを示した。これを除去するか、学習されたスカラーのみを使用すると、事前学習の安定性（PPL）と合成タスクの性能が著しく低下した。

意義と主張

本論文は、キーノルム正規化されたカチャルツ係数がデルタ則シーケンシャルモデルのための「第一序の設計軸」であると主張する。書き込み操作を制約付き射影として扱うことで、KLA は経験的に学習されたスカラーに内在するサブオプティマルな更新の大きさの問題を解決する。

その意義は、これらの改善が再帰状態、ハードウェアカーネル、または構造的複雑性を変更することなく達成された点にある。著者らは、状態容量やゲーティング構造ではなく、係数そのものが、デルタ則書き込みが局所回帰問題を成功裡に解決するかどうかを決定する主要因であると仮説を立てている。これは、将来の線形時間モデルが、この理論的に裏付けられたステップサイズを採用するだけで、精度、外挿、デコード効率において大幅な改善を達成できることを示唆している。