弾力的な潜在再帰（ELR）：1ループなら同サイズの通常モデルと同等で、ループを増やすほど精度が上がる1つのチェックポイント

「潜在再帰型の言語モデル」の考え方は単純です。数十の別々の層を積み上げる代わりに、小さな層のブロックを何度も繰り返して使い、ループが1回増えるごとに推論の計算量を1単位ずつ余計に使います。すると1つのチェックポイントで、速度と精度を調整できます。ループが少なければ速く、多ければより正確です。ただ、長く気まずい問題がありました。1ループでは、そうしたモデルは通常、同サイズの普通の非再帰モデルにはっきり劣ります。論文2610.12582はこの差を「1ループ税」と呼び、遅延が重要なときは別に非再帰モデルを選ぶほかなく、「1つのチェックポイントで多用途」という狙いが失われると指摘します。提案する弾力的潜在再帰（ELR）の目標は、1ループでは同サイズの非再帰モデルと同じ提供コストで、同じデータで訓練して同程度の精度に達し、その後ループを増やすたびに、追加の推論計算量でより高い精度を買えるようにすることです。 著者は3つの工夫が効いたと述べます。第一に、学習可能なステップ幅を持つゲート付きの状態更新で、各パスが既存の状態に対して制御された変更を行います。第二に、同じ最適化更新の中でループ数を混ぜます。マイクロバッチごとにループ数を一様にサンプルし、ブロックのすべての適用を通して誤差を逆伝播します。1ループの精度を失わないよう、データ並列のランクの一部を順番に強制的に1ループにし、1ループの確率を約3分の1に保ち、1ループと複数ループの勾配を同じ更新の中で平均します。第三に、各パスにループ内での進み具合と全体の予算を条件づけ、同じブロックがパスごとに違うことをできるようにします。アブレーションでは、他の一般的なやり方は1ループの精度を保てないことが示されます。最初の2つが保ち、3つ目が追加のループごとの損失を下げます。 論文はコストにも目を向けています。まず、混合深度の訓練は、事前学習の最後の4分の1で導入でき、基準より訓練FLOPsは推定で約16%増えるだけで、切り替え時にゼロ初期化の条件を加えれば、最初から混合深度で訓練するより損失が低くなります。次に、専門家混合モデルでは、後のループが最初のループの専門家の割り当てを再利用することで、混合深度の更新時間が11.7%短くなり、精度の測定可能な変化はありません。さらに、標準の推論ではループごとに別の注意のKVキャッシュを持つため、メモリは深さとともに増えます。後訓練したKoi-4Bでは、6つの注意層のうち3つが再帰ブロックにあるので、4ループは1ループの2.5倍のKVメモリ、つまり128Kトークンの要求あたりBF16で7.5 GiB対3.0 GiBが必要です。最後のループのKVキャッシュだけを保持すれば、再訓練は不要で、注意のKVメモリはどのループ数でも1ループの大きさに保たれ、4ループの得点低下はループごとのキャッシュと比べて0〜3%です。

主な結果は次のとおりです。4B密モデルと30B-A3B混合モデルで、1ループのELRは10タスク平均で対応する基準より0.52点低く、1.04点高く、4ループでそれぞれ7.25点と4.75点増えます。4Bモデルのタスクごとの最大の向上はGSM8Kの11.8点とMATH-500の11.3点、30B-A3BではMMLU-Proの7.5点とGSM8Kの6.5点です。代償は遅延に現れます。1ループの順伝播は各層をちょうど1回実行するので、最初のトークンまでの時間とトークンごとのデコード遅延は非再帰の基準と同じです。ループを1回足すごとに再帰ブロックだけが繰り返され、4Bでは4ループで最初のトークンまでの時間が基準の2.33倍、デコード遅延が2.06倍になります。3ループなら、1.87倍で1から4ループの向上の96%を得られます。この作り方を後訓練まで持ち込むと、1つのチェックポイントKoi-4Bが得られ、人工分析のインテリジェンス指数v4.1の重み付けで、1ループでは推定約18でQwen3.5-4Bと同等、4ループでは約29です。

どう読むべきでしょうか。第一に、数字はすべて論文自身の実験によります。著者は、Koi-4Bと比較モデルの得点は人工分析のものではなく自分たちの環境のもので、公表値と異なる可能性があり、Koi-4Bには対応する非再帰の基準がないと述べています。第二に、「ループを増やして精度を買う」ことは無料ではありません。遅延、訓練の複雑さ、提供側の対応が関わり、4ループでは最初のトークンまでの時間が2倍以上になります。第三に、アブレーションは主に小さめのモデルで行われており、大規模へ広げられるかは今後の検証が必要です。それでも、低遅延の即答と、難問での数ラウンドの思考の両方が要る製品にとって、「1つのチェックポイントで深さを調整できる」という発想は、2つのモデルを維持するより魅力的です。