Mixture-of-Depths: Dynamically allocating compute in transformer-based language models

Abstract

특정 계층에서 계산에 참여하는 토큰 수를 제한하여 FLOPs를 동적으로 할당하는 MoD(Mixture-of-Depths) 제안

 

[arXiv](2024/04/02 version v1)

 

 

 

Implementing Mixture-of-Depths Transformers

(1) Self-attention & MLP, (2) Residual connection (계산 우회) 중 하나로 토큰을 라우팅하는 설정을 고려한다.

---

Routing schemes

라우팅 방식에는 2가지가 있다.

(왼쪽) Token-choice routing은 각 토큰에 대해 전문가 선호도 분포를 생성하여 라우팅하는 것으로 불균형 할당 문제가 있다.

(중간) Expert-choice routing은 반대로 각 전문가가 토큰을 선택하는 것이다. 완벽한 load balancing을 보장하지만 일부 토큰이 과잉, 과소 처리되는 경우가 생김.

 

Expert-choice routing을 선택한 이유:

1. Load balancing loss가 필요하지 않음.
2. 중요한 토큰을 반드시 계산에 포함시킬 수 있다.
3. 단일 top-k 연산을 통해 효율적으로 집합을 분리할 수 있다.

---

Routing implementation

라우터 가중치 top-k에 들지 못하는 토큰 x는 block을 우회, 속하는 토큰은 self-attention & MLP를 통과한다.

 

추가로 라우터 가중치 r을 출력에 곱하여 gradient에 속하도록 한다.

---

Sampling

한 가지 문제가 있는데 top-k 작업은 비인과적(non-causal)으로, autoregressive sampling 시에는 미래 토큰에 접근할 수 없기 때문에 적용할 수 없다.

 

두 가지 해결 방법:

1. 라우터의 출력 logits에 추가적인 binary cross-entropy loss를 적용한다. (top-k에 속할지 말지)
2. 라우터와 동일한 입력을 받는 추가 MLP를 사용하여 해당 토큰이 top-k에 포함될 지 여부를 예측한다.

 

 

 

Results

IsoFLOP 비교.

 

 

Autoregressive model에서 성능 저하 좀 있음.

 

 

MoE와 통합 가능. (단계적 라우팅 또는 통합 라우팅)