일반적인 트랜스포머 (Transformer)는 모든 토큰을 모든 레이어(layer)를 통과하도록 강제하므로, "the"라는 단어는 "photosynthesis"만큼의 연산 비용을 정확히 소모합니다. 이는 명백히 낭비이며, Mixture-of-Depths (DeepMind, 2024)는 이러한 낭비를 마침내 해결 가능한 문제로 만든 해결책입니다. 이 방식은 각 블록(block)에 아주 작은 라우터 (router)를 추가하여 모든 토큰의 점수를 매기고, 상위 k개(top-k) 토큰만을 유지하여 전체 어텐션(attention)+MLP를 실행합니다. 나머지 토큰들은 블록을 완전히 건너뛰어 약 0 FLOPs의 비용으로 잔차 연결 (residual connection)을 타고 넘어갑니다. 깊이 (depth)는 동일한 총 예산 내에서 토큰별, 레이어별로 동적으로 할당됩니다. 저는 모든 라우터 점수와 FLOP 계산량이 실시간으로 계산되는 8개 레이어, 12개 토큰 데모를 구축했습니다. 작동 방식은 다음과 같습니다.

우리가 제거하려는 낭비

일반적인 트랜스포머는 모든 T개 토큰에 대해 L개의 블록을 실행합니다. 즉, 정보 밀도가 높은 단어와 채우기용 단어가 동일한 예산을 사용하는 L·T번의 블록 통과 (block-passes)가 발생합니다.

const L = 8;
const tokens = embed(sentence);    // T개 토큰 벡터
// 밀집 순전파 (dense forward) = L * T 블록 통과.
...

라우터는 하나의 아주 작은 선형 레이어입니다

각 블록에 라우터를 부여합니다. 이는 각 토큰의 벡터를 _하나의 스칼라 (scalar)_로 변환하는 단일 가중치 벡터입니다. 즉, "이 토큰이 이 블록의 연산을 얼마나 원하는가?"를 나타냅니다. 모든 토큰에 대해 실행되지만, 각각 하나의 내적 (dot product)만 수행하므로 블록 연산의 0.1% 미만입니다.

function router(x, Wr){
  return x.map(tok => dot(tok, Wr));   // 토큰당 하나의 "연산 수요" 수치
}

Top-k는 고정된 예산입니다 — 그리고 그것이 핵심 비결입니다

용량 (capacity, 예: 50%)을 정하고 점수에 따라 상위 k = capacity·T개의 토큰만을 유지합니다. 결정적으로 k는 임계값 (threshold)이 아니라 매 통과 시마다 동일한 개수를 유지하는 고정된 값입니다. 임계값을 사용하면 입력마다 처리되는 토큰 수가 달라져 GPU가 싫어하는 동적 연산 그래프 (dynamic compute graph)가 생성됩니다. Top-k는 그래프를 **정적 (static)**으로 유지합니다. 즉, FLOPs를 알 수 있고 깔끔한 배치 (batching)가 가능합니다.

function topK(weights, capacity){
  const k = Math.round(capacity * weights.length);   // 고정된 k = 예산
  const order = weights.map((w,i)=>[w,i]).sort((a,b)=>b[0]-a[0]);
...

선택된 토큰은 처리하고, 나머지는 잔차를 타고 넘어갑니다

이것이 핵심입니다. 선택된 토큰은 전체 Attention + MLP 과정을 거치지만, 나머지 모든 토큰은 변화 없이 잔차 연결 (residual connection)을 타고 블록을 통과하여 약 0 FLOPs의 비용으로 넘어갑니다. 어떤 토큰도 버려지는 것이 아니라, 단지 여기서는 변환되지 않을 뿐입니다.

return x.map((tok,i) =>
  keep.has(i)
    ? tok + sigmoid(w[i]) * attnMlp(tok, x)   // 선택됨: 라우터 점수에 의해 스케일링된 전체 블록 처리
...

(Sigmoid를 통해) 라우터 점수로 출력을 스케일링하는 것이 이 방식을 학습 가능하게 (trainable) 만듭니다. 하드 Top-k 방식은 그래디언트 (gradient)가 없지만, 연산 경로에 점수를 배치함으로써 손실 (loss)이 라우터로 역전파될 수 있게 하여, 별도의 레이블 없이도 내용어 (content words)는 깊이가 필요하고 "the"와 같은 단어는 필요하지 않다는 것을 스스로 학습하게 합니다.

절감 효과 계산하기

각 레이어는 독립적으로 라우팅하므로, 특정 토큰은 레이어 2에서 처리되고, 3을 건너뛴 뒤, 4에서 다시 실행될 수 있습니다. Dense 모델은 L·T 번의 블록 통과를 수행하지만, MoD는 L·k 번의 통과와 아주 작은 라우터 연산만을 수행하므로, 절감량은 거의 정확히 1 − capacity와 같습니다.

const modCost = L * Math.round(CAPACITY * T) + L * T * routerFLOPs;
const saved   = 1 - modCost / (L * T);   // ~= 1 - CAPACITY  -> 약 50%의 FLOPs 절감

논문에서는 Dense 모델과 동일한 성능을 내면서 순방향 전파 (forward pass) 당 FLOPs를 약 50% 수준으로 맞추었다고 보고합니다. 추론 (inference) 시 한 가지 주의할 점은, Top-k를 결정하려면 전체 시퀀스가 필요하며 이는 인과적 (causal)이지 않다는 것입니다. 따라서 MoD는 토큰 자체만 보고 "이 토큰을 처리할 것인가?"를 결정하는 작은 예측기 (predictor)를 학습시켜, 자기회귀적 (autoregressive) 특성과 KV-캐시 (KV-cache) 친화성을 유지하면서도 고정된 예산을 보존합니다.

MoD vs Mixture-of-Experts

두 방식 모두 토큰을 "라우팅"하지만, 절감하는 요소가 다릅니다. MoE는 모든 토큰을 서로 다른 전문가 (expert)에게 라우팅합니다. 이는 모델을 더 넓게 (wider) 만들어 파라미터 (parameter) 수는 늘리되 토큰당 연산량은 동일하게 유지합니다. 반면 MoD는 상위 k개 토큰만을 블록을 통과하도록 라우팅합니다. 즉, 파라미터 수는 동일하지만 FLOPs는 더 적습니다. MoE는 너비 (width)를 교환하고, MoD는 깊이 (depth)를 교환합니다. 이 둘은 직교(orthogonal)하므로, 전문가에게 라우팅하면서 동시에 토큰을 건너뛰는 "MoDE" 블록처럼 두 방식을 모두 사용할 수도 있습니다.

용량(capacity)을 설정하고 스택이 레이어별로 토큰을 라우팅하는 과정을 지켜보세요:
https://dev48v.infy.uk/ai/days/day47-mixture-of-depths.html