- Google DeepMind의 Mixture-of-Depths(MoD) 는 트랜스포머 언어 모델이 모든 토큰에 같은 FLOPs를 쓰지 않고, 레이어마다 필요한 토큰만 attention과 MLP 계산에 참여시키는 방식임
- 레이어별 라우터가 토큰마다 스칼라 가중치를 만들고, 미리 정한 top-k 용량 안의 토큰만 블록을 통과하며 나머지는 residual connection으로 우회함
- k를 사전에 고정해 계산 그래프와 텐서 크기는 정적으로 유지하면서도, 문맥에 따라 토큰별 계산 배분은 동적으로 달라짐
- 실험에서 MoD는 isoFLOP 기준 기본 transformer보다 같은 학습 시간에 더 낮은 loss를 내거나, 같은 성능에서 forward pass당 FLOPs를 줄여 학습·샘플링 step을 빠르게 만들 수 있었음
- 가장 좋은 설정은 12.5% 용량 블록을 격블록(every other block)에 적용하는 쪽이었고, auto-regressive sampling에서는 predictor 기반 라우팅으로 top-k의 비인과성 문제를 줄임
MoD가 겨냥한 계산 낭비
- 일반적인 transformer 언어 모델은 forward pass에서 모든 토큰에 동일한 계산량을 사용함
- MoD는 모든 토큰과 시퀀스가 예측에 같은 시간과 노력을 요구하지 않는다는 전제에서, 전체 계산 예산을 줄이면서 필요한 토큰에 계산을 배분함
- 기존 조건부 계산(conditional computation)은 동적 계산 그래프를 만들 수 있어 현재 하드웨어가 선호하는 정적 계산 그래프와 맞지 않을 수 있음
- 사용자가 학습 전에 총 계산 예산을 정하면, 모델은 그 안에서 토큰별·레이어별로 계산을 어디에 쓸지 학습함
- 계산량, 메모리 사용량, forward pass당 FLOPs를 사전에 예측할 수 있음
- 라우팅 대상 토큰은 입력 문맥에 따라 달라짐
라우팅 구조와 계산 절감 방식
- MoD는 MoE transformer처럼 라우터를 쓰지만, 여러 expert 중 하나를 고르는 대신 토큰을 두 경로 중 하나로 보냄
- attention과 MLP를 포함한 표준 블록 계산
- 값을 그대로 넘기는 residual connection
- residual connection은 계산 비용이 낮고, 블록 출력은 입력값으로 결정됨
- 블록의 capacity를 전체 토큰 수 T보다 작게 설정하면 기본 transformer보다 forward pass당 FLOPs가 줄어듦
- 예를 들어 capacity를 T/2로 줄이면 attention의 query-key 행렬 곱은 T² 대신 (T/2)²가 되어 FLOPs가 25% 수준이 됨
- capacity를 0으로 만들면 매우 빠르지만 대부분의 transformer 파라미터를 쓰지 않아 성능이 나빠지고, capacity를 T로 두면 기본 transformer와 같아짐
- 목표는 이 두 극단 사이에서 기본 transformer와 같거나 더 나은 성능을 내면서 step 속도를 높일 수 있는 지점을 찾는 것임
top-k expert-choice 라우팅
- 각 레이어 l에서 라우터가 토큰 임베딩마다 스칼라 가중치를 생성함
- 사용자가 정한 capacity C에 따라, 시퀀스 토큰 중 top-k에 해당하는 토큰만 attention과 MLP 계산에 참여함
- 선택되지 않은 토큰은 블록 계산을 건너뛰고 residual connection으로 다음 레이어에 전달됨
- 이 방식은 expert-choice routing에 기반함
- token-choice routing은 토큰이 경로를 고르므로 load balancing 문제가 생길 수 있음
- expert-choice routing은 각 경로가 top-k 토큰을 고르므로 capacity가 정확히 채워짐
- MoD에서는 경로가 사실상 “블록 계산”과 “우회” 두 가지라, 단일 top-k 연산으로 토큰을 두 집합으로 나눌 수 있음
- 라우터 가중치는 블록 출력에 곱해져 gradient path 위에 놓이고, 언어 모델링 학습 과정에서 경사하강법의 영향을 받음
샘플링에서 생기는 비인과성 문제
- expert-choice routing의 top-k 연산은 비인과적(non-causal) 임
- 어떤 토큰이 top-k에 들어가는지는 이후 토큰의 라우터 가중치까지 알아야 결정됨
- auto-regressive sampling에서는 미래 토큰 정보를 사용할 수 없음
- 이를 피하기 위해 두 가지 방법을 실험함
- 라우터 출력의 sigmoid가 top-k 선택 여부를 예측하도록 하는 binary cross-entropy 보조 손실
- 라우터와 같은 입력을 받되 stop-gradient를 적용한 작은 보조 MLP predictor
- 첫 번째 방법은 주 언어 모델링 objective에 약 0.2~0.3% 영향을 줌
- 두 번째 predictor 기반 방법은 언어 모델링 objective에 영향을 주지 않았고, step speed에도 유의미한 영향을 주지 않음
- 보조 예측 과제는 실험적으로 빠르게 99% 정확도에 도달함
실험 결과와 관찰
- 초기 실험은 총 6e18 FLOPs 예산에서 MoD 하이퍼파라미터를 찾는 방식으로 진행됨
- 최적 MoD 설정은 격블록마다 라우팅을 적용하고, top-k를 256으로 둔 경우였음
- 시퀀스 길이 2048 기준 256개 토큰, 즉 12.5% 만 attention과 MLP를 통과함
- 나머지 1792개 토큰, 즉 87.5% 는 블록을 우회함
- 220M 파라미터 MoD 변형은 isoFLOP 최적 baseline보다 약간 더 좋은 성능을 내면서 학습 중 step 속도가 60% 이상 빨랐고, 동일 하드웨어에서 전체 학습 wall-clock time은 거의 같았음
- 6e18, 2e19, 1e20 FLOPs의 isoFLOP 분석에서도 MoD의 FLOP-optimal 모델은 baseline보다 더 많은 파라미터를 가지는 경향을 보임
- 같은 성능 또는 더 낮은 loss를 내면서 forward pass당 FLOPs가 더 적은 MoD 변형들이 존재함
- step 속도 향상은 두 가지 요인에서 나옴
- 일부 토큰이 블록을 우회하므로 같은 모델 크기에서 FLOPs per parameter가 줄어듦
- baseline과 같은 성능을 내는 더 작은 MoD 변형을 고를 수 있음
- stochastic routing을 쓴 MoD는 기본 transformer와 일반 MoD보다 크게 낮은 성능을 보여, 학습된 라우팅이 중요했음
- auto-regressive 평가에서는 256,000개 시퀀스, 500M 토큰의 held-out data를 사용했고, predictor 기반 라우팅으로 바꿔도 성능 저하는 작았음
- MoD는 MoE와 결합해 Mixture-of-Depths-and-Experts(MoDE) 로도 구성할 수 있음
- staged MoDE는 attention 전에 토큰을 블록으로 보낼지 우회할지 결정함
- integrated MoDE는 기존 MLP expert 사이에 “no-op” expert를 넣어 MoD 라우팅을 통합함
- 라우팅 분석에서는 일부 토큰이 여러 블록을 자주 통과하고, 다른 토큰은 가능할 때마다 우회하는 패턴이 관찰됨
- 블록을 자주 통과하는 토큰은 출력 예측의 entropy가 더 높은 경우와 상관되는 예비 분석 결과가 있었고, 이는 예측이 더 어려운 토큰일 가능성과 연결됨