Mixture of Depths (MoD)
探索深度混合 (MoD) 如何通过动态路由 Token 来优化 AI 效率。了解该技术如何减少 Ultralytics YOLO26 和 LLM 中的 FLOPs。
在深度学习架构中,计算效率至关重要,特别是在处理长序列或高分辨率输入时。一种新颖的方法通过允许网络决定输入中哪些部分需要完全处理,哪些可以安全地跳过某些层,从而动态分配计算资源。这种动态路由策略降低了整体 计算复杂度,且不会牺牲模型的预测能力或准确性。
Link to this section概念理解#
Mixture of Depths (MoD) 是一种主要应用于 Transformer 架构 的架构技术,模型在其中学习在不同层动态跳过特定 Token 的计算。传统的 Transformer 会将每个 Token 通过每一层进行处理,无论它是关键信息还是填充内容。相比之下,MoD 模型使用路由机制来评估 Token 并为其分配分数。只有得分最高的 Token(达到预定义的容量限制)才会通过繁重的计算模块,例如 注意力机制 或全连接前馈层。其余的 Token 通过 残差连接 绕过该块,有效地创造了一种“深度混合”,使不同的 Token 经历不同程度的处理深度。
这种方法由最近的 DeepMind 研究 普及,并在 arXiv 存储库 中有详尽记录,它极大地减少了在训练和推理过程中所需的 浮点运算 (FLOPs) 总数。
Link to this section与 Mixture of Experts (MoE) 的区别#
很容易将此概念与 Mixture of Experts (MoE) 混淆。虽然两者都使用路由机制,但它们解决的问题不同:
- MoE 将 Token 路由到层内的不同 子网络(专家)。所有 Token 的计算深度保持不变,但模型的参数数量会增加。
- MoD 将 Token 路由到计算块或跳过连接。参数数量严格保持不变,但对于不太重要的 Token,计算深度会降低,从而直接改善 推理延迟。
Link to this section实际应用#
动态分配计算预算的能力使该技术在 计算机视觉 和自然语言处理的多个领域中都极具价值。
-
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.
-
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 代码片段,演示了基本的路由机制如何跳过部分输入 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 等企业生态系统无缝集成。






