AI-Papers
- 約2.1Mパラメータの軽量予測ネットワークが各入力に応じてレイヤーのスキップ・反復を決定し、ベースモデルの再学習なしに動的な推論経路を実現
- 正解できる入力の75.5%は標準の全レイヤー通過より短いプログラムで解決でき、推論オーバーヘッドはわずか0.8%（3.05ミリ秒）
- DART-Mathの数学推論ベンチマークで精度向上と実行レイヤー数削減を同時達成し、分布外テストでも改善効果が持続
研究の背景
大規模言語モデル（LLM）は、入力テキストに対してすべてのレイヤーを固定された順番で通過させる推論を行います。簡単な質問であっても複雑な数学問題であっても、同じ計算量が費やされる構造です。
この設計には2つの課題があります。一つは計算の非効率さで、単純な入力でも不要なレイヤーが実行されます。もう一つは表現力の制限で、特定の問題では一部のレイヤーを繰り返し使用した方が有益なケースがあっても、固定順序の実行では対応できません。
動的な深さを持つ先行研究も存在しますが、レイヤーを途中で打ち切るだけで「反復」を考慮したものは少なく、各入力に最適な実行順序を探索する余地が残されていました。
提案手法PoLarの仕組み
ICML 2026に採択された本研究「Program-of-Layers（PoLar）」は、事前学習済みLLMのレイヤーを「関数ライブラリ」として捉え、各入力に固有の実行プログラムを動的に生成します。プログラムは標準の全レイヤー通過とは異なり、一部のレイヤーをスキップしたり同じレイヤー群を繰り返したりする経路を取ります。ベースモデルの再学習は一切不要です。
プログラムの表現には2種類の構造が使われます。まず、隣接するレイヤーを「モジュール」として束ねる区切りマスク（最大4レイヤー単位）。次に、各モジュールに対して「スキップ」「キープ」「リピート」の3操作のいずれかを割り当てる操作ラベルです。
予測ネットワークの設計
入力ごとの最適プログラムを予測するため、約2.1Mパラメータ（ベースモデルの0.01〜0.06%）の軽量ネットワークが設計されました。入力テキストはQwen3-Embedding-0.6Bで埋め込まれ、クロスアテンション層とクロスレイヤーエンコーダを経て、区切りマスクと操作ラベルを同時に出力します。
学習用データはMonte Carlo Tree Search（MCTS、モンテカルロ木探索）でオフラインに収集されます。MCTSがプログラム空間の木構造を探索して有効な実行経路を発見し、その結果で予測ネットワークを訓練することで、推論時に高コストな探索を行わずに済みます。実測オーバーヘッドは全フォワードパスの0.8%（3.05ミリ秒）に収まります。
実験結果と性能
DART-Math（数学推論の難易度別ベンチマーク）を用いた実験では、LLaMA-3.2-3B-InstructのPass@5精度が難易度DM-1で47.6%から68.4%（+20.8ポイント）に向上しました。Qwen2.5-3B-InstuctやQwen3-8Bなど複数のモデルでも同様の改善が確認されています。
計算効率の面では、正解できる入力の75.5%が標準の全レイヤー通過より短いプログラムで解決できることが判明しました。最も難易度が低い問題（DM-1）では、実行時間がベースラインの0.83倍に収まりつつ精度が5.8ポイント向上するケースもあります。ASDiv、MAWPSといった学習分布外のベンチマークでも改善効果が持続することも確認されています。
主な発見
実験を通じて、LLMの潜在的な計算構造についていくつかの知見が得られました。
- スキップのみ、リピートのみよりも両者の組み合わせが最も精度が高い
- 難易度が上がるほど、スキップやリピートへの依存割合が増加する
- 有効なプログラムの54.5%が単一レイヤーセグメントで構成され、複雑な制御フローは不要なケースが多い
- リピートは各モジュールにつき1回で十分な効果を発揮する
数学推論のさらなる高精度化に取り組んだ研究として、集団探索と生成的検証器を組み合わせてIMO・USAMOレベルの数学証明を達成したMaxProofも注目されていますが、PoLarは既存モデルの推論パスそのものを最適化するアプローチで異なる方向性を持ちます。
まとめと今後の展望
PoLarは、LLMの固定深度実行が潜在的な推論能力の一部しか引き出せていないことを実証した研究です。軽量な予測ネットワークで入力ごとの実行プログラムを動的生成することで、ベースモデルの再学習なしに精度と効率の両立を実現しました。
現状の課題として、評価が数学推論タスク中心であること、モジュールサイズの上限が4レイヤーに固定されていること、MCTSによるオフライン収集コストが必要なことが挙げられます。数学以外の領域への拡張や、より柔軟なモジュール構造への対応が今後の研究課題となるでしょう。
論文情報: "Skip a Layer or Loop It? Learning Program-of-Layers in LLMs"（Ziyue Li, Yang Li, Tianyi Zhou, 2026） arXiv:2606.06574, CC BY 4.0
本記事の図（図1〜図2）は上記論文より引用しています（縮小・形式変換あり）。
