Mixture of Depths (MoD)

깊이 혼합(MoD)이 토큰을 동적으로 라우팅하여 AI 효율성을 최적화하는 방법을 살펴봅니다. 이 기법이 Ultralytics YOLO26 및 LLM에서 FLOPs를 줄이는 방법을 알아봅니다.

딥러닝 아키텍처에서는 특히 긴 시퀀스나 고해상도 입력을 처리할 때 계산 효율성이 가장 중요합니다. 새로운 접근 방식은 네트워크가 입력의 어느 부분에 전체 처리가 필요한지, 어느 부분이 특정 레이어를 안전하게 건너뛸 수 있는지를 결정하도록 하여 계산 리소스를 동적으로 할당합니다. 이러한 동적 라우팅 전략은 모델의 예측 성능이나 정확도를 저하시키지 않으면서 전체 계산 복잡도를 줄입니다.

개념 이해하기#

깊이 혼합 (MoD)은 주로 Transformer 아키텍처에 적용되는 아키텍처 기법으로, 모델이 여러 레이어에서 특정 토큰의 계산을 동적으로 건너뛰도록 학습합니다. 기존 Transformer는 중요한 정보인지 단순한 채움 콘텐츠인지와 관계없이 모든 토큰을 모든 레이어에서 처리합니다. 반면 MoD 모델은 라우터 메커니즘을 사용하여 토큰을 평가하고 점수를 할당합니다. 사전에 정의된 용량 한도까지 가장 높은 점수를 받은 토큰만 어텐션 메커니즘이나 밀집 피드포워드 레이어와 같은 고비용 계산 블록을 통과합니다. 나머지 토큰은 잔차 연결을 통해 해당 블록을 우회하여, 토큰마다 서로 다른 수준의 처리 깊이를 경험하는 "깊이 혼합"을 효과적으로 구현합니다.

최근 DeepMind 연구를 통해 널리 알려지고 arXiv 저장소에 광범위하게 문서화된 이 방법은 학습과 추론 과정에서 필요한 전체 부동 소수점 연산 (FLOPs) 수를 크게 줄입니다.

전문가 혼합 (MoE)과의 차이#

이 개념은 전문가 혼합 (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와 같은 엔터프라이즈 생태계와 원활하게 통합할 수 있습니다.