Mixture of Depths (MoD)

딥러닝 아키텍처에서 특히 긴 시퀀스나 고해상도 입력을 처리할 때 계산 효율성은 무엇보다 중요합니다. 새로운 접근 방식인 이 기술은 네트워크가 입력의 어느 부분이 전체 처리를 필요로 하고 어느 부분이 특정 계층을 안전하게 우회할 수 있는지 결정하도록 하여 컴퓨팅 자원을 동적으로 할당합니다. 이러한 동적 라우팅 전략은 모델의 예측 성능이나 정확도를 희생하지 않으면서도 전체적인 computational complexity를 줄여줍니다.

Link to this section개념 이해#

Mixture of Depths (MoD) is an architectural technique primarily applied to Transformer architectures where the model learns to dynamically skip computation for specific tokens at various layers. Traditional transformers process every token through every layer, whether it is a crucial piece of information or filler content. In contrast, an MoD model uses a router mechanism to evaluate tokens and assigns them a score. Only the top-scoring tokens—up to a predefined capacity limit—are passed through the heavy computation blocks, such as attention mechanisms or dense feed-forward layers. The remaining tokens bypass the block via residual connections, effectively creating a "mixture of depths" where different tokens experience varying levels of processing depth.

최근 DeepMind research를 통해 대중화되고 arXiv repository에 광범위하게 문서화된 이 방법은 학습 및 추론 과정에서 필요한 floating-point operations (FLOPs)의 총 개수를 획기적으로 줄여줍니다.

Link to this sectionMixture of Experts (MoE)와의 차이점#

이 개념을 Mixture of Experts (MoE)와 혼동하기 쉽습니다. 두 기술 모두 라우팅 메커니즘을 사용하지만, 서로 다른 문제를 해결합니다.

• MoE는 토큰을 계층 내의 서로 다른 서브 네트워크(전문가)로 라우팅합니다. 모든 토큰에 대해 계산 깊이는 동일하게 유지되지만 모델의 파라미터 수는 증가합니다.
• MoD는 토큰을 계산 블록이나 스킵 연결 중 하나로 라우팅합니다. 파라미터 수는 엄격하게 일정하게 유지되지만 중요도가 낮은 토큰에 대한 계산 깊이는 감소하여 inference latency를 직접적으로 개선합니다.

Link to this section실제 애플리케이션 사례#

컴퓨팅 자원을 동적으로 관리하는 능력 덕분에 이 기술은 computer vision 및 자연어 처리의 여러 영역에서 매우 가치 있게 활용됩니다.

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)

By leveraging frameworks like the PyTorch framework or TensorFlow, developers can integrate these custom model optimization blocks. Furthermore, tools like the Ultralytics Platform help teams manage the training data required to accurately train these routers, alongside integrating seamlessly with enterprise ecosystems like Google Cloud AI.