Mixture of Depths (MoD)
Mixture of Depths(MoD)がトークンを動的にルーティングしてAIの効率を最適化する方法を解説します。この手法がUltralytics YOLO26とLLMのFLOPsを削減する仕組みを学びます。
ディープラーニングアーキテクチャでは、特に長いシーケンスや高解像度の入力を処理する場合、計算効率が最も重要です。ネットワークが入力のどの部分に完全な処理が必要で、どの部分が特定のレイヤーを安全にスキップできるかを判断できるようにすることで、計算リソースを動的に割り当てる新しいアプローチがあります。この動的ルーティング戦略により、モデルの予測能力や精度を損なうことなく、全体的な計算量を削減できます。
概念の理解#
深度の混合(MoD)は、主にTransformerアーキテクチャに適用されるアーキテクチャ技術で、さまざまなレイヤーで特定のトークンの計算を動的にスキップする方法をモデルが学習します。従来のTransformerは、それが重要な情報であるか単なる埋め草のコンテンツであるかにかかわらず、すべてのトークンをすべてのレイヤーで処理します。これに対して、MoDモデルはルーターメカニズムを使用してトークンを評価し、スコアを割り当てます。スコア上位のトークンのみが、あらかじめ定義された容量上限まで、アテンションメカニズムや高密度フィードフォワードレイヤーなどの負荷の大きい計算ブロックを通過します。残りのトークンは残差接続を介してブロックをバイパスします。これにより、トークンごとに処理の深さが異なる、効果的な「深度の混合」が実現します。
この手法は、近年のDeepMindの研究によって広く知られるようになり、arXivリポジトリにも詳しく記録されています。トレーニングと推論の両方で必要となる浮動小数点演算(FLOPs)の総数を大幅に削減できます。
Mixture of Experts(MoE)との違い#
この概念はMixture of Experts(MoE)と混同しやすいものです。どちらもルーティングメカニズムを使用しますが、解決する問題は異なります。
- MoEは、レイヤー内の異なるサブネットワーク(エキスパート)にトークンをルーティングします。すべてのトークンで計算の深さは同じですが、モデルのパラメーター数は増加します。
- MoDは、トークンを計算ブロックまたはスキップ接続のいずれかにルーティングします。パラメーター数は厳密に一定のままですが、重要度の低いトークンでは計算の深さが減少し、推論レイテンシが直接改善されます。
実世界での利用例#
計算予算を動的に割り当てられるこの手法は、コンピュータービジョンや自然言語処理の複数の分野で非常に有用です。
-
言語モデルにおけるコンテキストの最適化: OpenAIやAnthropicなどの組織が開発する最新の大規模言語モデル(LLMs)は、巨大なコンテキストウィンドウを処理します。動的な深度ルーティングを採用することで、これらのモデルは構造上の単語や繰り返しの埋め草となる単語をスキップし、複雑な推論ステップや事実の抽出に深い計算を割り当てられます。
-
高解像度ビジョンAI: Ultralytics YOLO26モデルのような高度なビジョンシステムでは、物体検出や画像セグメンテーションのために大きな画像を処理する際、膨大なメモリが必要です。深度ルーティングにより、ネットワークは均一な背景(何もない空や無地の壁など)での特徴抽出をバイパスし、複雑な前景オブジェクトに計算能力を集中できます。これは、CUDA最適化ライブラリによって最適化された、リソースに制約のあるエッジAIハードウェアにモデルをデプロイするうえで重要です。
実装例#
以下は、入力トークンの一部の計算をスキップする基本的なルーティングメカニズムを示し、深度ルーティングの動作をシミュレーションする概念的なPyTorchスニペットです。
import torch
import torch.nn as nn
class MixtureOfDepthsBlock(nn.Module):
def __init__(self, d_model, capacity_factor=0.5):
super().__init__()
self.capacity_factor = capacity_factor
self.router = nn.Linear(d_model, 1)
self.heavy_compute = nn.Sequential(nn.Linear(d_model, d_model * 4), nn.GELU(), nn.Linear(d_model * 4, d_model))
def forward(self, x):
# x shape: (batch_size, seq_len, d_model)
seq_len = x.size(1)
capacity = int(seq_len * self.capacity_factor)
# 1. Compute routing scores
scores = self.router(x).squeeze(-1) # Shape: (batch_size, seq_len)
# 2. Identify top-k tokens to process
topk_indices = torch.topk(scores, capacity, dim=1).indices
# 3. Create an output tensor mirroring the input (residual baseline)
output = x.clone()
# 4. Apply heavy computation only to the selected tokens
for b in range(x.size(0)):
selected_tokens = x[b, topk_indices[b]]
processed_tokens = self.heavy_compute(selected_tokens)
output[b, topk_indices[b]] += processed_tokens
return output
# Example usage
dummy_input = torch.randn(2, 64, 128) # Batch=2, Seq=64, Dim=128
mod_block = MixtureOfDepthsBlock(d_model=128, capacity_factor=0.5)
output = mod_block(dummy_input)
print(f"Output shape: {output.shape}") # Expect (2, 64, 128)PyTorchフレームワークやTensorFlowなどのフレームワークを活用することで、開発者はこれらのカスタムモデル最適化ブロックを統合できます。さらに、Ultralytics Platformなどのツールを使用すると、これらのルーターを正確にトレーニングするために必要なトレーニングデータをチームで管理できるほか、Google Cloud AIなどのエンタープライズエコシステムともシームレスに統合できます。






