Mixture of Depths (MoD)

Mixture of Depths(MoD)는 학습된 라우터가 블록마다 어떤 토큰을 처리할지, 어떤 토큰이 잔차 연결을 통해 블록을 건너뛸지 결정하여 토큰마다 서로 다른 계산량을 할당하는 Transformer 기법입니다. Google DeepMind의 David Raposo, Adam Santoro와 동료들은 2024년 4월 Mixture-of-Depths: Transformer 기반 언어 모델의 동적 연산 할당에서 이 기법을 소개했습니다. 고정된 용량 k는 라우팅된 각 블록에서 처리할 토큰 수를 정하므로, 처리할 토큰은 입력에 따라 달라지더라도 연산 예산과 텐서 형태를 사전에 알 수 있습니다. MoD 대규모 언어 모델(LLM)은 학습 FLOPs와 실제 실행 시간이 동일할 때 기준 모델과 동등한 품질을 달성하면서 순전파당 FLOPs의 일부만 사용했고, 샘플링 속도는 50% 이상 빨라졌습니다. 고정된 학습 예산에서 최상의 MoD 모델은 최상의 기준 모델보다 더 컸으며 더 낮은 손실을 달성했습니다.

MoD 라우팅 작동 방식#

라우팅된 블록 앞에서 선형 라우터가 각 토큰에 스칼라 점수를 부여합니다. 점수가 가장 높은 k개 토큰은 블록의 셀프 어텐션과 MLP를 거치며, 그 출력에 라우터 점수를 곱한 뒤 잔차 스트림에 더합니다. 점수를 곱하면 라우터가 그래디언트 경로에 포함되므로 역전파를 통해 어떤 토큰이 유용한지 학습할 수 있습니다. 나머지 토큰은 변경 없이 통과하며 해당 블록의 어텐션에서 제외됩니다. 따라서 이후 토큰은 그 블록에서 제외된 토큰에 주의를 기울일 수 없으며 어텐션 비용은 용량의 제곱에 비례해 감소합니다. 토큰 수를 절반으로 줄이면 쿼리-키 곱셈 비용은 4분의 1이 됩니다. 논문에서 가장 좋은 설정은 시퀀스의 12.5% 용량으로 격 블록을 라우팅하는 것이었습니다. MLP만 사용하는 라우팅 블록의 최소 PyTorch 버전은 다음과 같습니다.

import torch
from torch import nn

torch.manual_seed(0)
batch, seq_len, dim = 2, 64, 128
capacity = int(0.125 * seq_len)  # 12.5% of tokens, the paper's best setting
x = torch.randn(batch, seq_len, dim)
router = nn.Linear(dim, 1)
block = nn.Sequential(nn.Linear(dim, 4 * dim), nn.GELU(), nn.Linear(4 * dim, dim))

# The router scores every token, and only the top-k tokens enter the block
weights, idx = router(x).squeeze(-1).topk(capacity, dim=1)
idx = idx[..., None].expand(-1, -1, dim)
update = block(x.gather(1, idx)) * weights[..., None]  # router weight keeps the router trainable

# Selected tokens get the residual update; the rest skip the block unchanged
out = x.scatter_add(1, idx, update)
unchanged = (out == x).all(dim=-1).sum().item()
print(f"Block ran on {batch * capacity} of {batch * seq_len} tokens; {unchanged} skipped it unchanged")

Block ran on 16 of 128 tokens; 112 skipped it unchanged을 출력합니다. 블록의 MLP는 평소 작업량의 8분의 1만 수행했으며, 다음 블록을 위해 출력 텐서는 전체 시퀀스 형태를 유지합니다.

생성 중 라우팅#

Top-k 선택은 비인과적입니다. 시퀀스에서 토큰이 상위 k개에 포함되는지는 그 뒤에 나오는 토큰의 점수에 따라 달라지지만, 모델이 한 번에 토큰 하나씩 텍스트를 생성할 때는 뒤에 나오는 토큰이 아직 존재하지 않기 때문입니다. 저자들은 두 가지 방법으로 라우팅을 인과적으로 만들었습니다. 첫 번째는 라우터 출력이 top-k 포함 여부를 예측하도록 가르치는 보조 이진 교차 엔트로피 손실로, 언어 모델링 손실이 0.2~0.3%만 변했습니다. 두 번째는 그래디언트를 차단한 입력을 받는 별도의 소형 MLP 예측기로, 빠르게 97~99%의 정확도에 도달했습니다. 두 방법 모두 자기회귀 샘플링에서 품질 저하는 크지 않았습니다. 학습된 라우터에서도 특정 패턴이 나타납니다. 일부 토큰은 모델 깊이에 걸쳐 모든 라우팅 블록을 활성화하는 반면, 다른 토큰은 가능할 때마다 블록을 우회합니다.

MoD와 관련 기법 비교#

• 전문가 혼합(MoE): 각 레이어에서 모든 토큰을 하나 또는 소수의 전문가 MLP로 라우팅합니다. 파라미터 수는 늘리면서 토큰당 계산량은 일정하게 유지합니다. 논문에서는 두 기법을 MoDE로 결합하기도 했습니다. 이 방식에서는 토큰을 전문가 또는 무연산 경로로 라우팅할 수 있으며, 두 기법의 효과가 더해졌습니다.
• 조기 종료: 토큰이 특정 레이어 이후 처리를 멈추고 나머지 레이어를 모두 건너뜁니다. MoD는 한 블록을 건너뛰더라도 이후 블록을 계속 사용할 수 있습니다.
• 토큰 병합(ToMe): 비전 Transformer에서 유사한 토큰을 병합해 이후 레이어가 더 적은 토큰을 처리하도록 하며, 학습된 라우터는 사용하지 않습니다.
• 모델 가지치기: 가중치를 영구적으로 제거하므로 모든 입력에 동일하게 축소된 네트워크를 적용합니다.
• 모델 라우팅: 전체 요청을 여러 모델 중 하나로 전달하는 방식으로, 완전한 모델 단위에서 동일한 개념을 적용합니다.

대부분의 네트워크는 구축 시 깊이가 고정됩니다. Ultralytics YOLO26은 모델 스케일마다 깊이를 설정합니다. 모델 YAML의 깊이 승수는 n, s, m 스케일에서 0.50이고 l, x 스케일에서 1.00입니다. 따라서 반복 횟수가 2로 지정된 C3k2 단계는 YOLO26n에서 1개 블록을, YOLO26x에서 2개 블록을 실행하며 모든 이미지가 해당 블록을 전부 통과합니다.