0
0

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?

Kimi K3の1M文脈を支えるKimi Delta Attentionという線形注意

0
Posted at

Kimi K3の1M文脈を支えるKimi Delta Attentionという線形注意

1Mトークンの文脈を「現実的な速度とメモリで」回す。この一句を額面どおり実現するのは、実はモデルを大きくするより難しい。文脈が伸びるほど注意機構のコストが効いてきて、モデルの賢さより先に速度とメモリのほうが破綻するからだ。Moonshot AIのKimi K3は2.8兆パラメータ・896エキスパート(トークンあたり16個だけ活性化)・1Mトークン文脈という数字で話題をさらったが、実務エンジニアとして本当に見るべきは総パラメータ数ではない。その1M文脈を破綻させずに回すために、Moonshotが注意機構そのものを作り替えている点だ。正体は Kimi Delta Attention(KDA) で、独立した論文とオープンな実装まで出ている。ここを解説する。

なぜ長い文脈は「思い出すたびに」高くつくのか

普通のTransformerの自己注意(self-attention)は、生成中のトークンが過去の全トークンを毎回参照する。そのために過去の鍵と値を保存しておくのが KVキャッシュ だ。文脈が長くなるほどこのキャッシュは線形に膨らみ、1Mトークンともなると、メモリを食うだけでなく、1トークン出すたびに巨大なキャッシュを舐め直すので生成が遅くなる。エージェント用途のように、ツール呼び出しの履歴や長いコードベースを丸ごと文脈に積む使い方では、この「思い出すコスト」がそのまま推論コストになる。

対極にあるのが 線形注意(linear attention) だ。過去を全部保持する代わりに、RNNのように固定サイズの状態へ情報を畳み込んでいく。メモリは文脈長に依存せず一定、生成も速い。ただし固定状態に押し込む以上、遠くの細かい情報を正確に「思い出す」能力(想起)は落ちやすい。速さと想起はトレードオフというのが、これまでの相場だった。

KDAは「忘れ方」を細かくして想起を取り戻す

KDAはこの線形注意側の弱点に手を入れたものだ。論文 Kimi Linear: An Expressive, Efficient Attention Architecture(Kimi Team, arXiv:2510.26692)は、KDAを次のように定義する。

an expressive linear attention module that extends Gated DeltaNet with a finer-grained gating mechanism

土台のGated DeltaNetは、線形注意の状態を「デルタ則(delta rule)」で更新する系譜だ。デルタ則はもともと、新しい入力と現在の記憶のズレ(誤差)分だけ状態を書き換えるという、古典的な連想記憶の更新規則を指す。KDAはここに より細かい粒度のゲーティング を足した。ゲートとは、状態のどこをどれだけ忘れ・上書きするかを制御する弁のことで、これを次元ごとに細かく効かせることで、限られた固定状態メモリをより無駄なく使い切る、という発想である。

実装面では、状態遷移行列に汎用のDPLR(対角+低ランク)ではなく、その特殊化した形を使うと論文は述べている。

a specialized variant of the Diagonal-Plus-Low-Rank (DPLR) transition matrices, which substantially reduces computation compared to the general DPLR formulation

つまり表現力は残しつつ、行列計算そのものを軽くしているわけだ。「ゲートを細かくすると重くなるのでは」という当然の懸念を、遷移行列側を削ることで相殺している。

全部を線形注意に置き換えない割り切り

面白いのは、KDAで全層を置き換えていない点だ。Kimi Linearは KDAと通常のフル注意(MLA, Multi-head Latent Attention)を3:1の比率で層ごとに混ぜる ハイブリッド構成をとる。ざっくり言えば、3層ぶんは安くて速いKDAに任せ、1層ぶんだけ想起の効くフル注意を残す、という配分だ。

この割り切りが効く。公式リポジトリと論文が挙げる数字は次の通り。

指標 Kimi Linearのハイブリッド構成
KVキャッシュ使用量 最大75%削減
1M文脈でのデコード throughput 最大6倍
MLA比のTPOT(1トークンあたり生成時間) 6.3倍高速
品質(全評価タスク) フルMLAを上回る

論文自身が「Kimi Linear outperforms full MLA with a sizeable margin across all evaluated tasks」と書いている点は強調しておきたい。線形注意系はこれまで「速いが少し弱い」妥協案として語られがちだったが、少なくともこの構成では、フル注意より軽く、かつ品質でも負けていないと主張している。速さのために品質を差し出す話ではない、というのがこれまでとの一番の差分だ。

なお、ここまでの数字はKDAを検証した研究モデル「Kimi Linear」(総48B・活性化3BのMoE)のものだ。K3(2.8兆パラメータ)そのものの内部ベンチマークではない点は区別しておく。

K3はこの土台の上に載っている

Kimi K3は、このKDAに加えて Attention Residuals(AttnRes) という別の工夫を重ねている。MarkTechPostの解説によれば、AttnResは情報を層を追うごとに一様に積み上げるのではなく、深さ方向で表現を選択的に取り出す仕組みで、追加コスト2%未満で訓練効率を約25%高めるという。KDAが「文脈長(横方向)」のコストに効くのに対し、AttnResは「層の深さ(縦方向)」の情報の流れに効く、と対比で理解すると整理しやすい。

K3の公開そのものにも注意点がある。7月16日に推論APIとアプリで先行公開されたが、複数の報道はフルの重みが後日(7月27日と報じられている)オープンウェイトとして出る予定としている。この重み公開日はMoonshotの一次発表で私が直接確認できたものではなく、二次報道ベースなので、自己ホストを計画するなら公式リポジトリでの確定を待つのが安全だ。2.8兆パラメータのMoEを動かすには相応のハードも要る。

重すぎるK3の代わりに48BのKimi Linearを動かす

K3本体はまだ重すぎるが、KDAの挙動そのものは研究モデルのKimi Linearで今すぐ触れる。公式リポジトリ MoonshotAI/Kimi-Linear が手順を公開している(以下は同リポジトリで確認したもの)。

必要要件はPython 3.10以上、PyTorch 2.6以上。KDAのカーネルはFlash Linear Attention由来のパッケージで入る。

pip install -U fla-core

Hugging Faceからの読み込みは、カスタム注意機構を含むため trust_remote_code=True が要る。

from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "moonshotai/Kimi-Linear-48B-A3B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    trust_remote_code=True,
    device_map="auto",
)

サーブするならvLLMがそのまま使える。

vllm serve moonshotai/Kimi-Linear-48B-A3B-Instruct --port 8000

チェックポイントは Kimi-Linear-48B-A3B-Base(素の事前学習版)と -Instruct(指示調整版)の2種がHugging Faceに出ている。

派手な数字より地味なハイブリッド設計を見る

長文脈やエージェントを実運用に載せようとすると、真っ先にぶつかるのはモデルの賢さではなくKVキャッシュのコストだ。純粋な線形注意で丸ごと置き換えると想起が落ち、フル注意のままだと1M文脈が高すぎる。KDAが示したのは、線形注意の「忘れ方」を細かく作り込み、フル注意を少量だけ混ぜるという中間解で、コストと品質の両取りが現実になりつつある、という一点だ。K3という派手な数字の裏で、実際に効いているのはこの地味なハイブリッド設計の方だと私は見ている。総パラメータ数より、注意機構の作りを見るべき局面に来ている。

0
0
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
0
0

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?