ループ型 Transformer のスケーリング則

ロール仲の皆さん、こんにちは。rumor です。

Transformer が台頭し始めた頃、あるアイデアが提案されていました。パラメータを増やさず、活性化値を同じ層グループ内で複数回ループさせるだけで、計算量を増やせるのではないか?というものです。

このアイデアはループ型 Transformer(ループ構造)、あるいは再帰的深さモデル(Recurrent Depth Model: RDM)と呼ばれます。理論的には非常にエレガントです。推論時にループ回数を柔軟に調整でき、メモリ使用量は固定であり、テスト時スケーリングを本質的にサポートします。

しかし実際には、こうしたモデルの訓練は極めて不安定でした。残差爆発(residual explosion)や損失のスパイクが発生し、ハイパーパラメータの選択に対して非常に敏感だったのです。

UCSD と Together AI による今回の論文「Parcae: Scaling Laws For Stable Looped Language Models」[1]は、制御理論の観点からこの不安定性の根源を解明し、安定版である Parcae を提案しました。

不安定性はどこから来るのか

Parcae の著者らは、非常に巧妙な視点を用いました。ループ型 Transformer の各ステップの更新を、制御システムにおける力学系として捉えたのです。

ループ型 Transformer の力学系としての模式図

ループごとに、モデルの隠れ状態 h は Transformer 層の計算を経て新たな h へと更新されます。このプロセスは以下の 3 つに分解できます。

  1. 前回の隠れ状態 h_t が行列によって変換され、次回へ伝達される。
  2. 初期入力である埋め込み e が、モデルが逸脱しないよう各回の計算に継続的に注入される。
  3. 残りは Transformer 内のアテンションや MLP といった非線形計算部分である。

この過程を数式で表すと以下のようになります。

ここで、

  • は t 回目のループにおける隠れ状態
  • は前置モジュール P が出力する入力埋め込み
  • は状態遷移行列。前ステップの隠れ状態の伝達を制御する
  • は入力注入行列。入力 e が当前の状態に与える影響を制御する
  • は Transformer モジュール(アテンション+MLP)の非線形部分

ここで最も重要なのが、前回の状態伝達を制御する行列です。最も分かりやすく例えるなら、このループ過程は雪だるまを転がすようなものです。

  • 隠れ状態 h は、あなたの手の中の雪だるま。
  • ループごとの処理は、雪だるまを雪原で一回転させる行為。
  • 行列は、雪だるまが一回転するごとの拡大係数。

もしの拡大係数が 1 より大きければどうなるでしょうか。雪だるまは一回転するごとに大きくなり、1 回目は拳大から椀大へ、2 回目で洗面器大へ、そして十数回も回せば小山ほどに膨れ上がり、最後には「爆発」します。つまり数値オーバーフローを起こし、モデルは発散してしまうのです。

では、どうすればこの雪だるまを爆発させずに安定して転がせるのでしょうか。

古典的な制御理論ですでに答えが出ています。こうしたループ型の線形システムが安定であるためには、行列のスペクトル半径でなければなりません。

ここでのスペクトル半径とは、この行列の固有値の絶対値の最大値と理解できます。この値が 1 より小さければ、雪だるまは一回転するごとに最大でも元の大きさを保つか、あるいは縮小するだけで、決して巨大化することはなく、自然と爆発もしないのです。

そこで従来のループ構造を振り返ると、問題が一気に露呈します。

  • 加算によって入力を注入するモデルでは、が単位行列 I そのものとなり、スペクトル半径は 1 となります。これは「限界安定」の状態であり、少しのきっかけで容易に爆発してしまいます。
  • 連結と射影によって入力を注入するモデルでは、に全く制約がありません。訓練中にモデルがスペクトル半径が 1 より大きい行列を学習してしまい、即座に発散してしまいます。

著者らは実験でもこの結論を証明しました。訓練が発散したモデルはすべて、訓練過程でスペクトル半径が 1 以上となるを学習していました。一方、安定して収束したモデルは、常にスペクトル半径が 1 より小さいままでした。

これでループ構造を長年悩ませてきた安定性の謎は、完全に解明されたのです。

Parcae による对症下药(病に合わせた治療)

病根が見つかった以上、解決策は自明です。不安定性がのスペクトル半径の制御不能に起因するのなら、これを厳しく制約し、スペクトル半径を常に 1 より小さく保てばよいのです。

これが論文で提案された Parcae 構造の中核的な設計思想です。派手なトリックは一切なく、すべての工程が痛烈な急所を突き、かつ追加されるパラメータはごく僅かです。まさに「パラメータを増やさずにスケーリングする」ことを実現しました。

1. 状態遷移行列に枷をはめ、根源から安定を保証する

著者らは行列に対し、特殊なパラメータ化形式を設計しました。まず連続領域の行列 A を負対角行列とし、それを制御理論における標準的なゼロ次ホールド(ZOH)法を用いて、離散ループ内のへと変換します。

この設計の妙は、負対角行列のすべての固有値が負の数であるため、ZOH 離散化を経て得られたのスペクトル半径は必然的に 1 より小さくなり、数学的にシステムの安定性が保証され、雪だるまが巨大化する事態が二度と起こらなくなる点にあります。

2. 入力に正規化を加え、訓練後期の損失スパイクを解消する

を制約するだけでは不十分でした。著者らが発見したところ、大規模モデルの訓練後期には依然として損失のスパイクが散発することがあるというのです。調査の結果、原因は入力注入される e であることが判明しました。入力 e の数値が大きすぎると、隠れ状態が突然爆発してしまうのです。

解決策は簡単です。入力 e がループに入る前に正規化レイヤーを一段追加し、その数値範囲を厳密に制御するだけです。この小さな変更一つで、訓練後期の損失スパイクは見事に平滑化されました。

3. 訓練サンプリング方式を最適化し、モデルにより安定して学習させる

従来のループモデルでは、訓練時にバッチ単位でループ回数をサンプリングしていました。つまり、同一バッチ内のすべての文が同じループ回数を用いていたのです。しかしこれでは、モデルがループ回数の期待値を正確に推定できず、訓練中の損失変動が大きくなる原因となっていました。

Parcae ではこれをシーケンスごとの深さサンプリング(Per-Seq. Sampling)へ変更しました。同一バッチ内でも、各文が独立して自身のループ回数をサンプリングするようにしたのです。これにより、モデルは異なるループ回数の分布をより正確に学習できるようになり、訓練は安定し損失変動は消滅。さらに、異なるループ回数に対する汎化能力も向上しました。

これらの手を尽くした結果、効果はいかに凄まじいものでしょうか。

異なる学習率における収束状況の比較グラフ

著者らは 2e-4 から 1e-3 までの 5 種類の学習率で実験を行いました。

  • 元の RDM モデルは、最小の 2e-4 でのみ収束。
  • 残差正規化を追加した RDM でさえ、4e-4 以下でのみ収束。
  • 一方 Parcae は、5 組すべての学習率で安定して収束しました。

これでループモデルを長年悩ませてきたハイパーパラメータ過敏という難問が見事に解決されたのです。

効果

安定性の問題が解決され、ループ構造の真価が完全に解放されました。著者らは 140M から 1.3B パラメータまでを網羅する大規模実験を実施。同パラメータ数の標準 Transformer との比較結果は以下の通りです。770M パラメータの Parcae は Core スコア 25.07 を記録。1.3B パラメータの標準 Transformer(25.45)とほぼ互角です。パラメータ数を半分に削減し、メモリ使用量は据え置きのまま、性能差はわずか 0.38 ポイントに抑え込みました。

Parcae と標準 Transformer の性能比較グラフ

(注)本論文では Parcae と Transformer の推論効率の比較は行われていません。例えば 770M の Parcae(8 回ループ)は 1.3B の Transformer と同等の効果ですが、前者は 8 回のループを要します。トークンあたりの推論遅延、スループット、メモリ帯域幅の占有状況などの差異については、今後の検証が待たれます。

既存最高性能のループモデル RDM と比較しても、Parcae は下流タスクの平均スコアで最大 1.8 ポイントの向上を達成しました。

RDM と Parcae の下流タスク性能比較

さらに著者らは、前述の 3 つの最適化ポイントに関するアブレーション実験も実施。制約付き A により T=4/8 での発散を防止し、シーケンスごとのサンプリングで損失スパイクの分散を低減。プレリュード正規化が全体品質の向上と後期スパイクの解消に寄与したことが確認されました。

各コンポーネントのアブレーション実験結果

ループモデルのスケーリング則

安定性の問題解決がループ構造の経絡を整えることであれば、本論文がもたらす長期的価値として最も大きいのは、ループ構造に対するスケーリング則(Scaling Laws)を初めて体系的に導出した点にあります。

ループ回数は、パラメータ数・データ量に並ぶ、3 つ目の独立かつ予測可能なスケーリング次元なのです。

訓練時の計算リソースはどう配分するのが最も効率的か?

著者らは大量の isoFLOP(総計算量固定)実験を実施。その結果、固定された計算量予算のもとでの最適訓練戦略は、すべての計算リソースをデータ量に投じるのではなく、ループ回数と訓練データ量の双方を同期して増加させることだと判明しました。

isoFLOP 実験における最適配分の可視化

両者とも総計算量の増加に伴い、厳密な冪乗則に従います。

ループ回数とデータ量の冪乗則を示すグラフ

つまり、訓練の総計算量が 2 倍になるたび、最適なループ回数は約 32%増加し、最適な訓練データ量は約 72%増加させる必要があります。この組み合わせが最も優れたモデル性能をもたらすのです。

実験でも確認されています。同じ計算量予算であれば、「ループ回数を増やしデータ量を減らす」構成の方が、「低ループ回数でデータ量を最大化する」構成よりも損失が低く、性能も向上します。

異なる配分戦略における損失比較

これは大規模モデル訓練に全新的な道筋を開くものです。計算リソースに限りがあり、より大きなモデルやより多くのデータを準備できない場合でも、ループ回数を増やすことで、より低コストでより高い性能を達成可能なのです。

推論時は何回ループさせるのが最も効率的か?

訓練に加え、著者らは推論時におけるループ回数の増加による性能向上が飽和指数減衰の法則に従うことも発見しました。

推論時のループ回数と性能向上の関係を示すグラフ

この数式が示すのは、初期段階ではループ回数を増やすことで性能が顕著に向上するものの、回数が増えるにつれて限界効用は急速に低下し、最終的には最低損失の下界に収束し、それ以上は向上しないという事実です。

そしてこの下界は、訓練時のループ回数によって決定されます。訓練時のループ回数が多ければ多いほど、この下界は低くなり、推論時に到達可能な性能上限も高まります。訓練時の性能上限を、推論時の無限ループで突破することは不可能なのです。

さらに素晴らしいことに、著者らは訓練と推論のスケーリング則を統合した統一数式を構築。異なる計算量・異なるループ回数におけるモデル性能を精密に予測可能とし、その予測誤差はわずか 0.85%〜1.31%に収まっています。今後はループモデルを訓練する際、やみくもにハイパーパラメータを試行錯誤する必要はなく、この数式を用いて最適解を算出できるようになります。

訓練・推論を統合したスケーリング則の予測精度グラフ

(注)本論文の全実験は、最大でも 1.3B パラメータ・104B トークン規模に留まっており、産業界で主流の大規模モデルのサイズには遥かに及びません。より大規模なモデルにおいても、安定性やパラメータ効率の優位性が維持されるかどうかは、今後の検証が待たれます。

まとめ

総じて、本論文はループ構造研究における記念碑的な成果です。理論的根源からループ型 Transformer の長年の訓練不安定性問題を解決しただけでなく、完全なスケーリング則を提示。「パラメータを増やさずにモデル性能を向上させる」という方向性に、全新的な想像空間を切り開きました。

ただし、本論文の中核的な安定性に関する結論は線形近似に基づく導出であり、理論的な前提境界が存在します。アテンションや活性化関数を含む完全な非線形システムにおいては、スペクトル半径<1 は安定のための必要条件ではあるものの、十分条件ではありません。また性能面でも、Transformer モデルとの事後訓練(Post-train)後の比較は行われていません。

現時点では未だ埋めるべき課題は多く、大規模な産業応用までは距離がありますが、パラメータ・データ量・ループ回数という 3 つのスケーリング軸を同時に解放したとき、最適解はいったいどこにあるのか――これこそが真に興味深い問いであり、Parcae はまさにその扉をほんの少し開けただけに過ぎないのです。

参考資料

[1] Parcae: Scaling Laws For Stable Looped Language Models: https://arxiv.org/abs/2604.12946

執筆者プロフィール画像


パンクでギークな AI アルゴリズム研究員の rumor です。

大規模言語モデルのアルゴリズム研究に従事。Google 開発者エキスパート。

私をフォローして、一緒に学び、一緒に頑張りましょう。

人工知能という時代を、共に踊り、跳ね、瞬きながら駆け抜けましょう。

「ループに称賛を」称賛ボタンまたはアイコン

関連記事

分享網址
AINews·AI 新聞聚合平台
© 2026 AINews. All rights reserved.