はじめに
再帰ブロックを深度方向に繰り返し適用するアーキテクチャでは、「何回繰り返すか」が推論品質と計算コストの両方に影響します。前回の記事で報告した「有効深度帯」の存在は、深すぎる反復は品質を悪化させることを示していました。
では、推論時にその「止めどき」を自動的に検出できないか？
本記事では、隣接する深度間の出力分布差異（KLダイバージェンス等）を監視することで、追加学習なしに有効深度帯の変化点を検出し、不要な演算を省略する手法について報告します。
要点
- 隣接深度間のKLダイバージェンスは非単調なプロファイルを示す
- 変化の極小付近（安定領域）から増大に転じる位置が有効深度帯の上限に対応
- 連続増加パターンの検出により、閾値パラメータなしで停止可能
- 計算量を約50%削減しつつ、出力品質は維持または向上
- 追加学習は不要。既存の学習済みモデルにそのまま適用可能
背景: 従来の停止制御手法
Adaptive Computation Time (ACT)
ACT（Graves, 2016）は各反復で累積するhalting probabilityが閾値を超えた時点で停止する手法です。しかし停止判断のための追加パラメータ（halting unit）の学習を必要とし、停止位置が有効深度帯の上限と直接対応しない場合があります。
Early Stopping
機械学習における学習の早期停止は検証損失の増大を検出して学習を打ち切る手法ですが、これは学習プロセスの制御であり、推論時の計算深度制御とは本質的に異なる問題です。
本手法の位置づけ
本手法は、モデル自身の出力分布の変化を監視することで、推論時に有効深度帯の変化点を検出します。追加パラメータの学習は不要で、既存の学習済みモデルに対してそのまま適用可能です。
手法: 出力分布差異の監視による停止制御
基本的な考え方
再帰ブロックを深度方向に適用するモデルでは、各深度dで中間的な出力分布P_dが得られます。隣接する深度間の出力分布差異を測定すると、深度方向に特徴的なパターンが現れます。
安定領域（極小付近）から増大に転じる位置が有効深度帯の変化点に対応します。
安定領域では分布間差異が減少し、有効深度帯の上限を超えると急激に増大します。
停止判定の2つの方式
方式A: 固定閾値方式
所定の初期深度d_min以降で、分布間差異が閾値Tを超えた場合に、その直前の深度の出力を最終出力として採用します。
方式B: 連続増加パターン方式
分布間差異が連続して増加する傾向（連続的な増加傾向）を検出した場合に、増加開始の深度で停止します。閾値パラメータの設定が不要です。
初期深度d_minの必要性
再帰ブロックの初期反復では、状態ベクトルがランダム初期化値から安定状態へ遷移する過渡応答期間が存在します。この期間では分布間差異が大きな値を示しますが、これは有効深度帯の上限を示すものではありません。そのため、初期深度d_minを設定し、この過渡応答期間を停止判断の対象から除外します。
実験結果
実験環境
| 項目 | 値 |
|---|---|
| モデル | 2.1Bパラメータ 再帰深度Transformer |
| データ | 日本語Wikipedia（トークン化済み） |
| GPU | NVIDIA RTX 3090（24GB）×1 |
出力分布差異のプロファイル
日本語Wikipediaの実データを用いて、各深度でのKLダイバージェンスを測定しました。
定性的な結果:
- 浅い深度ではKLが高い（初期過渡）
- 深度が進むにつれKLが減少し、安定領域を形成
- 安定領域の底を過ぎると、KLが連続的に増大
この非単調な変化パターンは、有効深度帯の変化点に対応する位置でKLが増大に転じることを示しています。
方式Aと方式Bの比較
| 方式 | 停止深度 | 計算量削減 |
|---|---|---|
| 固定深度（最大深度まで計算） | 24 | 0% |
| 方式A（固定閾値、d_min設定） | 有効深度帯付近 | 約42〜50% |
| 方式B（連続増加パターン） | 有効深度帯付近 | 約54% |
方式Bは方式Aと比較して:
- 閾値パラメータの設定が不要
- 分布間差異のスケールに依存しない
- 有効深度帯の最適深度により近い位置で停止
複数の分布間距離指標で同じパターン
KLダイバージェンスだけでなく、Jensen-Shannonダイバージェンス、コサイン距離、L2ノルムでも同様の非単調な変化パターンが観測されました。
このことから、非単調なプロファイルは特定の距離指標に依存する現象ではなく、再帰系の状態遷移に起因する構造的特性であると考えられます。
停止判定の安定性
異なるサンプル集合を用いた複数回の評価において、連続増加パターンによる停止深度のばらつきを測定しました。
- 30バッチ平均での停止深度: 安定（標準偏差≈0）
- 1バッチでも停止深度の変動は±1以内
- 異なる学習段階のモデル（学習途中/追加学習後）でも成立
停止判定はモデルの内部状態構造に由来するため、入力サンプルの選択に対して高いロバスト性を示します。
入力長に依存しない
系列長を変えて実験した結果、有効深度帯の位置（≒停止深度）は入力系列長に依存しないことが確認されました。
ACTとの差異
| 観点 | ACT | 本手法 |
|---|---|---|
| 停止判断の根拠 | 学習されたhalting unit | モデル自身の出力分布変化 |
| 追加パラメータ | 必要 | 不要 |
| 追加学習 | 必要 | 不要 |
| 既存モデルへの適用 | 再学習が必要 | 推論時にそのまま適用可能 |
| 停止位置の意味 | 入力依存の計算量調整 | 有効深度帯の変化点検出 |
今後の展望
- 分布間差異の非単調プロファイルの理論的解釈（収縮写像→発散遷移の位相境界として）
- より大規模なモデル・異なるアーキテクチャでの再現性検証
- 実運用環境でのレイテンシ・スループットへの影響測定
実験環境
- GPU: NVIDIA RTX 3090 (24GB) ×1
- フレームワーク: PyTorch
- 学習データ: 日本語Wikipedia（HuggingFace datasets形式）
本記事の内容に関連する技術は特許出願済みです。
BathysRDTプロジェクトの詳細については @nakatada-lab の他の記事もご参照ください。