問題
知識蒸留は、大きなモデルの知識を小さなモデルに転移する標準的手法です。教師モデルの出力分布を生徒モデルに模倣させることで、小さなモデルでも高い性能を引き出せます。
再帰型Transformer(同一ブロックを深度方向に繰り返し適用するアーキテクチャ)に蒸留を適用する場合、通常は全ての反復深度で均一に教師信号を与えます。
しかし、ここに問題があります。
再帰型Transformerの内部状態は、浅い反復では活発に変化しますが、深い反復では収束して変化しなくなります。前回の記事(Ghost Detectionによる内部状態収束解析)で、この収束現象を定量的に確認しました。
収束済みの深い反復に教師信号を与え続けることは、安定した内部表現に不要な擾乱を与えることになります。では、収束度合いに応じて教師信号を制御したらどうなるか。これが本記事の実験です。
アイデア
Ghost Detectionの結果から、各反復深度の収束状況が分かっています。
反復深度 収束状況
1〜5 活性状態(出力が大きく変化)
6〜7 移行領域(一部の入力で収束開始)
8〜12 収束状態(出力がほぼ変化しない)
この情報を蒸留に反映します。3つの方式を比較しました。
固定方式(従来): 全ての反復で同じ強さの教師信号を与える
連続制御方式: 収束状況に応じて蒸留の適用方法を段階的に変える
二値制御方式: 収束状況に応じて蒸留の適用方法を切り替える
実験
722Mパラメータの再帰深度Transformerを、3条件で同一の100Mトークン(日本語Wikipedia)により学習しました。
- 教師モデル: llm-jp-4-8b-base(4bit量子化)
- GPU: RTX 3090 × 1
- 最大反復深度: 12
- 乱数シード: 42
評価は51,200トークン(50バッチ)で、各反復深度におけるCross-Entropy(CE)を測定しました。
結果
深度別CE
| 反復深度 | 固定(従来) | 二値制御 | 連続制御 |
|---|---|---|---|
| n=1 | 6.910 | 8.469 | 8.628 |
| n=3 | 6.377 | 7.508 | 8.415 |
| n=6 | 6.276 | 7.325 | 8.042 |
| n=12 | 5.997 | 7.042 | 7.653 |
| 単調性 | ✓ | ✓ | ✓ |
3条件すべてで深度単調性を達成しました(反復を増やすほどCEが改善する)。
追加評価: 浅い深度での比較
50バッチ再評価において、連続制御の反復2回が固定方式の反復12回に相当する性能を示しました。反復回数にして約83%の削減に相当します。
考察
本実験では、内部状態の情報を蒸留に活用することで、蒸留の「適用場所」についての制御が可能であることを確認しました。
3条件すべてで深度単調性が達成されたことは、内部状態に応じた制御が深度利用構造を破壊しないことを意味しています。また、連続制御のn=2が固定方式のn=12に相当する性能を示したことは、浅い深度での表現形成が促進されていることを示唆しています。
一方で、CE絶対値は固定方式が最良であり、内部状態のみに基づく単純な制御では性能向上には至りませんでした。今回用いた収束情報だけでは、入力ごとの最適な反復深度までは判断できないためと考えられます。
つまり、内部状態の収束情報は蒸留制御の判断材料として成立するが、それ単体では最適な制御には不十分であり、入力依存の制御やGate Decay機構との組み合わせが必要です。
3段階構造
連続制御方式のCE推移には、3段階の構造が観測されました。
- 第1段階(反復1→2): 急激な改善。初期状態から有効な内部表現への遷移
- 第2段階(反復3→6): 漸進的改善。追加反復による表現の精緻化
- 第3段階(反復7→12): CE飽和。内部状態の収束
この3段階構造は、Ghost Detectionで得られた深度プロファイル(浅層Active→中層移行→深層Ghost)と対応しています。
応用の可能性
この「内部状態の収束度に応じて処理を変える」という考え方は、蒸留制御に限らず幅広い応用が考えられます。
推論時の早期停止: 収束を検出した時点で残りの反復を省略すれば、推論コストを大幅に削減できます。本実験のn=2≈n=12の結果は、入力によっては反復2回で十分な場合があることを示しています。エッジデバイスや低レイテンシが求められる環境で特に有効です。
計算資源の動的配分: バッチ内の各入力に対して、収束が早い入力は浅い深度で打ち切り、難しい入力だけ深い深度まで計算する。入力の「難しさ」に応じて計算量を動的に配分するアプローチです。同一バッチ内でもサンプルごとに異なる反復回数を適用できれば、GPUの計算効率が改善します。
モデル圧縮との組み合わせ: 収束済み深度では内部状態がほとんど変化しないことを利用して、深い反復段階だけ量子化や低ランク近似を適用するという戦略が考えられます。状態が安定している領域であれば、精度を落としても出力への影響は小さいはずです。
学習カリキュラムの自動設計: 学習の進行に伴って内部状態の収束パターンは変化します。学習初期はどの深度も活性状態ですが、学習が進むと深い深度から順に収束していきます。この変化を追跡すれば、学習段階に応じた蒸留の自動調整が可能になります。
まとめ
Ghost Detectionで得られる内部状態の情報が、知識蒸留にも活用できる可能性を確認しました。CE絶対値の改善には至りませんでしたが、蒸留の適用方法に内部状態の情報を活用できること、浅い深度で深い深度に相当する性能が得られることを確認しました。
Ghost Detectionは、推論時の停止制御(ALTH-009で報告済み)だけでなく、学習時の蒸留制御にも応用可能な、再帰Transformerの汎用的な内部状態監視手法であると言えます。今後は入力依存制御やGate Decayとの組み合わせにより、再帰Transformer全体の深度制御へ発展させる予定です。
実験環境
- GPU: NVIDIA RTX 3090 (24GB) × 1
- フレームワーク: PyTorch 2.6.0+cu124
- 教師モデル: llm-jp-4-8b-base(4bit量子化、BitsAndBytes)
- 学習データ: 日本語Wikipedia 100Mトークン
本記事の内容に関連する技術は特許出願済みです。
関連記事
- Gate Decayは何を制御しているのか(ALTH-006: τの設計思想)
- 深度依存ゲート減衰がTransformerの深度利用を変える(ALTH-008: 4条件アブレーション)
- Ghost Detectionによる内部状態収束解析(ALTH-009: 収束判定)


