Skip a Layer or Loop it? Test-Time Depth Adaptation of Pretrained LLMs

著者 Ziyue Li, Yang Li, Tianyi Zhou
所属 University of Maryland, College Park
投稿日 2025年7月11日
カテゴリ cs.LG, cs.CL
arXiv 2507.07996v1 ↗

Skip a Layer or Loop it? Test-Time Depth Adaptation of Pretrained LLMs

基本情報

  • arXiv ID: 2507.07996v1 (https://arxiv.org/abs/2507.07996)
  • 著者: Ziyue Li, Yang Li, Tianyi Zhou
  • 所属: University of Maryland, College Park
  • 投稿日: 2025年7月11日
  • カテゴリ: cs.LG, cs.CL

簡単に説明すると

この論文は、事前学習済みLLMの各層を再利用可能なモジュールとして扱い、入力ごとに最適な層の組み合わせを動的に選択するChain-of-Layers（CoLa）を提案している。スキップ、繰り返し、並び替えの3つの操作が可能である。モンテカルロ木探索（MCTS）を使用して追加学習なしでテスト時に最適な層の実行パスを発見する。例えばLLaMA-3-3BのARC-Easyでの精度は27.8%から95.8%に向上し、計算効率も改善された。

関連リンクは以下の通りである。

  • GitHub: （論文中に記載なし）

1. 研究概要

1.1 背景と動機

現代のLLMは、入力の難易度に関わらず同じ固定アーキテクチャを使用している。しかし、簡単なタスクで全ての層が必要なのか、難しいタスクで固定の深さで十分なのかという根本的な疑問がある。人間の認知は「速い思考」と「遅い思考」を使い分ける。一方でLLMは常に同じ計算パスを使用する。

従来の適応的計算手法（早期終了やlayer-wiseの削除）は追加学習を必要とし、サブネットワークやハードウェア最適化、データ圧縮技術もそれぞれ独自の学習コストがかかる。一方、CoLaは事前学習済みモデルの各層が独立したモジュールとして機能できることを示し、追加学習なしでテスト時に動的な深さ適応を実現する。

1.2 主要な貢献

CoLaの主要な貢献は以下の通りである。

  • 事前学習済みLLMの層を追加学習なしでモジュラーに操作できることを実証
  • モンテカルロ木探索を用いた効率的な層パス探索手法の開発
  • 精度向上と計算効率化の同時実現（例：エラー訂正時により少ない層で正解を達成）
  • LLMの「過剰思考」現象の発見（正しい予測でも不要な計算を行っている）

2. 提案手法

2.1 手法の概要

Chain-of-Layers（CoLa）は、事前学習済みLLMの各層を以下の3つの方法で操作する。

  • スキップ（Skip）：特定の層を飛ばして計算を省略
  • 繰り返し（Recurrence）：同じ層を複数回実行してRNNのような深い処理を実現
  • 並び替え（Reorder）：層の実行順序を変更（実験では効果が限定的）

例えば、元のモデルが[L₁, L₂, ..., Lₙ]という層構成の場合を考える。CoLaは入力ごとに[L₁, L₁, L₃, L₅, L₅, L₅, L₇, ...]のような新しい実行パスを構築する。

2.2 技術的詳細

モンテカルロ木探索（MCTS）による最適化：
組み合わせ爆発的な探索空間を効率的に探索するため、MCTSを使用する。

  1. 状態空間：部分的または完全な層シーケンス（初期状態は元の順伝播パス）
  2. アクション：
    • k層をスキップ（k ∈ {1,2,3,4}）
    • k層をr回繰り返し（k,r ∈ {1,2,3,4}）
  3. 報酬設計：正解時に1、不正解時に0、オプションで長さペナルティを追加
  4. UCBスコア：
    UCB(P) = Q(P)/v(P) + c√(ln V/v(P)) - λ||P||/N
    ここで、Q(P)は累積報酬、v(P)は訪問回数、Vは総訪問回数、||P||はパス長

実装の詳細：

  • 各入力に対して50-100回のシミュレーションを実行
  • 精度と効率性のパレート最適解を返す
  • キャッシュを使用して同じ層シーケンスの再計算を回避

2.3 新規性

CoLaの新規性は以下の点にある。

  • 事前学習済みモデルの層をテスト時に動的に再構成する初の手法
  • 追加学習なしで大幅な性能向上を実現
  • 層のスキップと繰り返しを組み合わせた柔軟な実行パス構築
  • MCTSを用いた効率的な探索による実用的な推論時間の実現

3. 実験結果

3.1 実験設定

評価は以下の設定で実施された。

  • モデル：LLaMA-3-3B、LLaMA-3-8B、Mistral-3-7B
  • モデルタイプ：ベースモデルとinstruction-tunedモデル
  • データセット：ARC-Easy、GSM8K、DART（難易度1-5）、Math
  • 評価指標：精度、平均パス長、エラー訂正率
  • MCTSパラメータ：50-100シミュレーション、c=2、λ=0.1

3.2 主要な結果

精度の大幅な向上：

  • LLaMA-3-3B on ARC-Easy: 27.8% → 95.8%（+68.0%）
  • LLaMA-3-3B on GSM8K: 14.8% → 29.2%（+14.4%）
  • 全てのモデルとデータセットで一貫した改善

計算効率の改善：

  • 正しく予測されたサンプルの75%以上がより短いパスで実現可能
  • エラー訂正時は元の正解よりも短いパスを使用する傾向
  • 平均的に10-20%の層削減を達成しながら精度向上

スキップvs繰り返しvs組み合わせ：

  • スキップのみ：簡単なタスクで効果的、計算削減に優れる
  • 繰り返しのみ：中程度の複雑さのタスクで有効
  • 組み合わせ（CoLa）：特に難しいタスクで劇的な性能向上

3.3 既存手法との比較

CoLaは以下の点で既存手法を上回る。

  • 早期終了手法：固定的な終了条件vs動的な層構成
  • Layer dropping：学習が必要vs学習不要
  • Mixture of Depths：特定の位置での層選択vs任意の層構成
  • ハードウェア最適化：物理的な制約vs柔軟なソフトウェア実装

4. 実用性評価

4.1 実装の容易性

CoLaの実装は比較的シンプルである。

  • 既存のLLM実装に対して層単位の実行制御を追加するのみ
  • MCTSは標準的なアルゴリズムで実装が容易
  • キャッシュ機構により計算の重複を回避
  • バッチ処理との互換性は要検討

4.2 計算効率

計算効率の観点では以下の特徴がある。

  • MCTS探索のオーバーヘッド：入力あたり50-100回のシミュレーション
  • キャッシュにより同じパスの再計算を回避
  • 最終的な推論は元のモデルより高速（層数削減のため）
  • メモリ使用量は元のモデルと同等

4.3 応用可能性

CoLaは以下の応用分野で活用可能である。

  • エッジデバイスでのLLM推論：動的に計算量を調整。バッテリー寿命を延長
  • リアルタイムシステム：レイテンシ要件に応じた適応的な計算
  • 大規模推論システム：負荷に応じた動的なリソース割り当て
  • モデル解釈性研究：層の重要度と役割の理解

5. まとめと所感

5.1 論文の意義

CoLaはLLMの推論に関する根本的な仮定に挑戦する重要な研究である。固定アーキテクチャが必須ではないことを示した。事前学習済みの層が驚くほどモジュラーであることを実証した。

この研究の意義は以下の点にある。

  1. 新しい一般化の次元：テスト時のアーキテクチャ適応という未探索の領域を開拓
  2. 効率性と精度の両立：従来のトレードオフを打破し、両方を改善
  3. LLMの内部動作の理解：層の使用パターンから新しい知見を獲得

5.2 今後の展望

CoLaの今後の発展可能性として以下が考えられる。

  1. 自動パス学習：MCTSの結果を蒸留して高速な推論を実現
  2. マルチモーダルモデルへの拡張：視覚や音声を含むモデルでの適用
  3. ハードウェア最適化：動的実行パスに最適化されたアクセラレータの開発
  4. 継続学習との統合：新しいタスクに対する層パスの効率的な発見

総じてCoLaはLLMの推論効率化に新しい道を開く革新的な研究である。実用的な応用と理論的な洞察の両方を提供している。