Mixture of Depths (MoD)

Mixture of Depths（MoD）は、学習されたルーターがブロックごとにどのトークンを処理し、どのトークンを残差接続を通じてブロックをスキップさせるか決定するTransformerの手法です。これにより、トークンごとに異なる量の計算を適用できます。Google DeepMindのDavid Raposo、Adam Santoroらは、2024年4月にMixture-of-Depths: Dynamically Allocating Compute in Transformer-Based Language Modelsでこの手法を発表しました。固定容量kによって、各ルーティング対象ブロックが処理するトークン数が決まるため、トークンの選択が入力に依存していても、計算予算とテンソル形状は事前に分かります。MoDの大規模言語モデル（LLM）は、訓練時のFLOPsと実時間が同じ場合、ベースラインと同等の品質を達成しながら、順伝播あたりのFLOPsを大幅に削減し、サンプリング時のステップ速度を50%超向上させました。訓練予算を固定した場合、最良のMoDモデルは最良のベースラインより大規模で、より低い損失を達成しました。

MoDルーティングの仕組み#

ルーティング対象のブロックの前で、線形ルーターが各トークンにスカラー値のスコアを付与します。スコアが上位k個のトークンはブロックの自己注意とMLPを通り、その出力にルーターのスコアを掛けて残差ストリームに加算します。スコアを掛けることでルーターが勾配経路に含まれるため、誤差逆伝播によって、どのトークンが有益かを学習できます。残りのトークンは変更されずに通過し、そのブロックの注意計算から除外されるため、そのブロックでは後続のトークンから注意を向けられません。注意コストは容量の二乗に比例して小さくなり、トークン数を半分にするとクエリーとキーの積の計算コストは4分の1になります。論文で最良だった設定は、系列の12.5%を容量として、ブロックを1つおきにルーティングするものでした。PyTorchで、MLPのみを使ったルーティングブロックの最小実装を示します。

import torch
from torch import nn

torch.manual_seed(0)
batch, seq_len, dim = 2, 64, 128
capacity = int(0.125 * seq_len)  # 12.5% of tokens, the paper's best setting
x = torch.randn(batch, seq_len, dim)
router = nn.Linear(dim, 1)
block = nn.Sequential(nn.Linear(dim, 4 * dim), nn.GELU(), nn.Linear(4 * dim, dim))

# The router scores every token, and only the top-k tokens enter the block
weights, idx = router(x).squeeze(-1).topk(capacity, dim=1)
idx = idx[..., None].expand(-1, -1, dim)
update = block(x.gather(1, idx)) * weights[..., None]  # router weight keeps the router trainable

# Selected tokens get the residual update; the rest skip the block unchanged
out = x.scatter_add(1, idx, update)
unchanged = (out == x).all(dim=-1).sum().item()
print(f"Block ran on {batch * capacity} of {batch * seq_len} tokens; {unchanged} skipped it unchanged")

Block ran on 16 of 128 tokens; 112 skipped it unchangedを出力します。ブロックのMLPは通常の8分の1の処理を行い、出力テンソルは次のブロックに渡すため、系列全体の形状を維持します。

生成時のルーティング#

Top-k選択は因果的ではありません。トークンが系列内で上位k個に入るかどうかは、その後に続くトークンのスコアにも依存しますが、モデルがテキストを生成して1トークンずつ出力するとき、後続のトークンはまだ存在しません。著者らは2つの方法でルーティングを因果的にしました。1つは、ルーターの出力がtop-kに含まれるかを予測するよう訓練する補助的な二値交差エントロピー損失で、言語モデリングの損失への影響はわずか0.2～0.3%でした。もう1つは、勾配を停止した入力に対する小規模な独立MLP予測器で、すぐに97～99%の精度を達成しました。どちらの方法でも、自己回帰サンプリングの品質低下はわずかでした。訓練済みのルーターにはパターンも見られます。あるトークンはモデルの深さに沿ってルーティング対象のすべてのブロックを利用する一方、別のトークンは可能な限りブロックを迂回します。

MoDと関連手法#

• エキスパート混合（MoE）: 各層内で、すべてのトークンを1つまたは少数のエキスパートMLPに振り分けます。パラメーター数を増やしつつ、トークンごとの計算量は一定に保ちます。論文ではMoDEとして2つを組み合わせた手法も提案しています。トークンはエキスパートまたは何もしない処理に振り分けられ、両者の効果が相乗的に高まりました。
• 早期終了: トークンがある層で処理を終え、残りのすべての層をスキップします。MoDでは、1つのブロックをスキップしても後続のブロックを利用できます。
• トークンマージ（ToMe）: ビジョントランスフォーマー内で類似するトークンを統合し、後続の層で処理するトークン数を減らします。学習されたルーターは使用しません。
• モデルプルーニング: 重みを恒久的に削除するため、すべての入力が同じ縮小ネットワークを使います。
• モデルルーティング: リクエスト全体を複数のモデルのいずれかに振り分けます。完全なモデル単位で同じ考え方を適用します。

ほとんどのネットワークでは、構築時に深さが固定されます。Ultralytics YOLO26では、モデルのスケールごとに深さを設定します。モデルYAML内の深さ倍率は、n、s、mスケールで0.50、lとxで1.00です。そのため、繰り返し回数が2と記載されたC3k2ステージは、YOLO26nでは1ブロック、YOLO26xでは2ブロック実行し、すべての画像がそれらすべてを通過します。