Mixture of Depths (MoD)

探索深度混合（MoD）如何通过动态路由 token 优化 AI 效率。了解该技术如何减少 Ultralytics YOLO26 和 LLMs 中的 FLOPs。

在深度学习架构中，计算效率至关重要，尤其是在处理长序列或高分辨率输入时。一种新颖的方法是动态分配计算资源，让网络决定输入的哪些部分需要完整处理，哪些部分可以安全地跳过某些层。这种动态路由策略在不牺牲模型预测能力或准确性的情况下，降低了整体计算复杂度。

理解这一概念#

深度混合（MoD）是一种架构技术，主要应用于Transformer 架构，模型会学习在不同层动态跳过对特定 token 的计算。传统 Transformer 会让每个 token 经过每一层，无论它是关键信息还是填充内容。相比之下，MoD 模型使用路由器机制评估 token 并为其分配分数。只有得分最高的 token（最多达到预先设定的容量上限）才会通过计算量较大的计算模块，例如注意力机制或稠密前馈层。其余 token 则通过残差连接绕过该模块，从而有效形成一种“深度混合”，让不同 token 经历不同程度的处理深度。

这种方法因近期的DeepMind 研究而广为人知，并在arXiv 论文库中得到广泛记录，大幅减少了训练和推理期间所需的浮点运算（FLOPs）总数。

与专家混合（MoE）的区别#

这个概念很容易与专家混合（MoE）混淆。虽然两者都使用路由机制，但它们解决的是不同的问题：

MoE 将 token 路由到某一层中的不同子网络（专家）。所有 token 的计算深度保持不变，但模型的参数量会增加。
MoD 将 token 路由到计算模块或跳过连接。参数量严格保持不变，但不太重要的 token 的计算深度会降低，从而直接改善推理延迟。

实际应用#

动态分配计算预算的能力，使这项技术在计算机视觉和自然语言处理等多个领域都极具价值。

**语言模型中的上下文优化：**来自OpenAI和Anthropic等组织的现代大语言模型（LLMs）会处理超大上下文窗口。通过采用动态深度路由，这些模型可以跳过结构性或重复性的填充词，将深度计算留给复杂的推理步骤和事实提取。
**高分辨率视觉 AI：**在Ultralytics YOLO26等先进视觉系统中，为目标检测和图像分割处理大图像需要大量内存。深度路由使网络能够跳过对均匀背景（如空旷的天空或空白墙壁）的特征提取，将计算能力集中于复杂的前景目标。这对于将模型部署到由CUDA 优化库优化的资源受限边缘 AI硬件上至关重要。

实现示例#

下面是一段概念性的 PyTorch 代码片段，用于演示基本路由机制如何跳过部分输入 token 的计算，以模拟深度路由行为。

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等企业生态系统实现无缝集成。