[논문 리뷰] Mixture-of-Depths: Dynamically allocating compute in transformer-based language models

David Raposo, Sam Ritter외 4명|arXiv (Cornell University)|2024. 04. 02.

Natural Language Processing Techniques인용 수 4

한 줄 요약

이 논문은 Mixture-of-Depths (MoD)를 소개하며, 토큰이 학습된 top-k 라우팅 메커니즘을 통해 레이어를 건너뛸 수 있도록 허용함으로써 트랜스포머 언어 모델에서 계산을 동적으로 할당하고, 각 레이어의 총 FLOPs를 제한하는 방법을 제안한다. MoD 모델는 정상적인 트랜스포머보다 최대 50% 적은 FLOPs로 전방향 계산을 수행하고 최대 50% 빠른 추론을 달성하면서도, 정적 계산 그래프를 통해 학습 효율성과 하드웨어 호환성을 유지한다.

ABSTRACT

Transformer-based language models spread FLOPs uniformly across input sequences. In this work we demonstrate that transformers can instead learn to dynamically allocate FLOPs (or compute) to specific positions in a sequence, optimising the allocation along the sequence for different layers across the model depth. Our method enforces a total compute budget by capping the number of tokens ($k$) that can participate in the self-attention and MLP computations at a given layer. The tokens to be processed are determined by the network using a top-$k$ routing mechanism. Since $k$ is defined a priori, this simple procedure uses a static computation graph with known tensor sizes, unlike other conditional computation techniques. Nevertheless, since the identities of the $k$ tokens are fluid, this method can expend FLOPs non-uniformly across the time and model depth dimensions. Thus, compute expenditure is entirely predictable in sum total, but dynamic and context-sensitive at the token-level. Not only do models trained in this way learn to dynamically allocate compute, they do so efficiently. These models match baseline performance for equivalent FLOPS and wall-clock times to train, but require a fraction of the FLOPs per forward pass, and can be upwards of 50\% faster to step during post-training sampling.

연구 동기 및 목표

• 모든 토큰이 동일한 FLOPs를 받는 트랜스포머의 계산 분포의 비효율성 문제를 해결하기 위해.
• 동적 계산 그래프를 도입하지 않고도 토큰 및 레이어 차원에서 맥락에 민감한 계산 할당을 가능하게 하기 위해.
• 스마트한 라우팅을 통해 모델의 다양한 깊이를 통과하는 토큰을 라우팅함으로써 전방향 계산 시 FLOPs를 줄이고, 학습 속도와 하드웨어 효율성을 유지하기 위해.
• 학습된 라우팅이 학습 시간이나 정확도를 희생시키지 않고 모델의 효율성과 성능을 향상시킬 수 있는지 탐색하기 위해.
• 추론 시 성능 저하 없이 계산을 절약하여 더 빠른 추론 또는 모델 확장이 가능함을 보여주기 위해.

제안 방법

• 각 레이어에서 라우터 헤드를 기반으로 top-k 라우팅 메커니즘이 토큰이 자기주도 주의 및 MLP 계산에 참여할지 여부를 선택한다. k는 사전에 고정된다.
• 라우터에 의해 선택되지 않은 토큰들은 잔차 연결을 통해 전달되어, 실제로 트랜스포머 블록 계산을 건너뛰며 FLOPs를 절약한다.
• 계산 그래프가 정적이고 알려진 텐서 크기를 가지므로 하드웨어 효율성과 기존 가속기와의 호환성이 보장된다.
• 라우팅 결정은 학습 기간 동안 종단 간 라우팅으로 학습되며, 보조 균형 손실이 필요로 하지 않는다. 상대적으로 안정적인 top-k 라우팅을 활용한다.
• 자기주도적 샘플링 중 라우터는 미래 토큰 정보 없이도 학습 시 라우팅 결정을 모방할 수 있도록 단순한 보조 분류기 또는 손실을 통해 정제된다.
• 각 레이어에서의 계산 예산(k)을 조정함으로써 속도와 성능 간의 트레이드오���을 가능하게 하여 더 빠른 추론 또는 모델 확장이 가능하다.

Figure 1: Mixture-of-Depths Transformer. As in mixture-of-experts (MoE) transformers we use a router to choose among potential computational paths. But unlike in MoE transformers the possible choices are a standard block’s computation (i.e., self-attention and MLP) or a residual connection. Since so

실험 결과

연구 질문

• RQ1 트랜스포머가 정적 계산 그래프를 유지하면서도 토큰 및 레이어 차원에서 계산을 동적으로 할당할 수 있는가?
• RQ2 라우팅을 통한 동적 계산 할당이 성능 저하 없이 상당한 FLOP 절감을 이끌 수 있는가?
• RQ3 동일한 학습 FLOP 예산 하에서 MoD 모델는 표준 트랜스포머보다 더 높은 성능 또는 더 빠른 추론을 달성할 수 있는가?
• RQ4 추론 중 향후 토큰 정보가 없을 경우 라우팅 성능에 어떤 영향을 미치며, 이를 어떻게 완화할 수 있는가?
• RQ5 라우팅 메커니즘을 표준 블록과 잔차 연결 외의 다양한 유형의 계산(예: 메모리 검색, 도구 사용 등)으로 확장할 수 있는가?

주요 결과

• MoD 모델는 동일한 FLOPs와 월클럭 시간으로 훈련되었을 때 기준 트랜스포머와 동일하거나 더 높은 성능을 달성하며, 계산 효율성이 향상됨을 보여준다.
• MoD 모델는 표준 트랜스포머 대비 전방향 계산 시 최대 50%까지 FLOPs를 절감하여 성능 저하 없이 더 빠른 추론을 가능하게 한다.
• 훈련 FLOPs와 시간을 동일하게 유지한 상태에서 후기 훈련 샘플링 중 최대 50% 빠른 추론 단계를 달성할 수 있다.
• 학습 중의 top-k 라우팅 메커니즘이 자기주도적 추론으로 일반화되며, 미래 토큰의 정체성이 가용하지 않더라도 성능 저하가 최소한도로 유지된다.
• MoD 모델는 지능적으로 라우팅을 학습하여 간단한 예측에 대해 불필요한 계산을 건너뛰는 반면, 전략적인 주의 라우팅을 통해 장거리 의존성을 유지한다.
• 이 프레임워크는 확장 가능하며, 다양한 계산 유형(예: 메모리 검색, 도구 사용 등)으로 라우팅을 지원할 수 있어 제어 가능한 FLOP 비용을 가진 새로운 아키텍처 설계를 가능하게 한다.

Figure 2: Routing schemes. Tokens are funnelled to the computational path of their choice when using token-choice routing (left). If a given path exceeds its capacity (e.g., more than two tokens in this example) then surplus tokens must be dropped (purple token). The exact token that is ultimately d