Mixture of Depths (MoD)

ディープラーニングのアーキテクチャにおいて、特に長いシーケンスや高解像度の入力を処理する際には、計算効率が極めて重要です。新しいアプローチとして、ネットワークが入力のどの部分に完全な処理が必要で、どの部分が特定のレイヤーを安全にスキップできるかを判断することで、計算リソースを動的に割り当てる手法があります。この動的なルーティング戦略により、モデルの予測能力や精度を犠牲にすることなく、全体的な計算複雑性を低減します。

Link to this sectionコンセプトの理解#

Mixture of Depths (MoD) は、主にTransformerアーキテクチャに適用されるアーキテクチャ技術であり、モデルがさまざまなレイヤーで特定のトークンに対する計算を動的にスキップすることを学習します。従来のTransformerは、重要な情報であるかフィラーコンテンツであるかに関わらず、すべてのトークンをすべてのレイヤーで処理します。これに対し、MoDモデルはルーターメカニズムを使用してトークンを評価し、スコアを割り当てます。定義された容量制限までの上位スコアのトークンのみが、アテンションメカニズムや密なフィードフォワード層などの重い計算ブロックを通過します。残りのトークンは残差接続を介してブロックをバイパスし、トークンごとに処理の深さが異なる「深さの混合（Mixture of Depths）」を効果的に作り出します。

この手法は、最近のDeepMindの研究によって広まり、arXivリポジトリで広く文書化されており、トレーニングと推論の両方で必要な浮動小数点演算数 (FLOPs)の総数を劇的に削減します。

Link to this sectionMixture of Experts (MoE) との違い#

このコンセプトは、Mixture of Experts (MoE) と混同されがちです。どちらもルーティングメカニズムを使用しますが、解決する問題は異なります。

• MoEは、レイヤー内の異なるサブネットワーク（エキスパート）にトークンをルーティングします。すべてのトークンに対する計算の深さは変わりませんが、モデルのパラメータ数が増加します。
• MoDは、計算ブロックまたはスキップ接続のいずれかにトークンをルーティングします。パラメータ数は厳密に一定ですが、重要度の低いトークンに対する計算の深さが減少し、推論レイテンシが直接的に向上します。

Link to this section実社会での応用#

計算リソースを動的に割り当てる能力により、この技術はコンピュータビジョンや自然言語処理の複数の領域において非常に価値の高いものとなっています。

1. Context Optimization in Language Models: Modern Large Language Models (LLMs) from organizations like OpenAI and Anthropic process massive context windows. By employing dynamic depth routing, these models can skip structural or repetitive filler words, reserving deep computation for complex reasoning steps and factual extraction.

2. High-Resolution Vision AI: In advanced vision systems like the Ultralytics YOLO26 model, processing large images for object detection and image segmentation requires immense memory. Depth routing allows the network to bypass feature extraction on uniform backgrounds (like empty skies or blank walls), focusing computational power on intricate foreground objects. This is crucial for deploying models to resource-constrained edge AI hardware optimized by CUDA optimization libraries.

Link to this section実装例#

以下は、基本的なルーティングメカニズムがどのように入力トークンの一部に対する計算をスキップし、深さルーティングの挙動をシミュレートするかを示す概念的な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 framework や TensorFlow といったフレームワークを活用することで、開発者はこれらのカスタムモデル最適化ブロックを統合できます。さらに、Ultralytics Platform のようなツールは、これらのルーターを正確にトレーニングするために必要なトレーニングデータをチームが管理するのを助け、Google Cloud AI といったエンタープライズエコシステムとのシームレスな統合を可能にします。