MoD(Mixture-of-Depths): Transformer 기반 언어 모델 연산 최적화를 위한 접근법, 그리고 MoDE(MoD+MoE)

9bow · 4월 7, 2024, 5:19오전

PyTorchKR

Jamba나 Qwen1.5-MoE, 또는 DBRX 등과 같이 여러 '전문가' 블록들 중 하나 이상을 선택하여 입력 token의 처리를 진행하는 MoE(Mixture-of-Experts) 방법이 최근들어 많이 사용되고 있습니다. 이러한 MoE 방식은 방대한 LLM 중 일부를 동적으로 조절하여 연산 자원을 효율적으로 할당하면서 더 나은 성능을 보이고 있습니다.

MoD(Mixture-of-Depths): Transformer 기반 언어 모델 연산 최적화를 위한 접근법

이렇게 MoE가 여러 전문가들 중 일부를 선택하는 방식으로 '모델의 폭(width)'을 동적으로 조절했다면, 이번에 소개하는 MoD(Mixture-of-Depth)는 '모델의 깊이(depth)'를 동적으로 조절하는 접근 방식을 취하고 있습니다. 즉, 특정 토큰에 대해서 모델의 전체 레이어에 대해서 연산을 하지 않고, 일부 레이어들을 건너뛰는 방식으로 동적으로 연산을 하는 것인데요, 함께 살펴보시죠.

MoD(Mixture-of-Depths): Transformer 기반 언어 모델 연산 최적화를 위한 접근법 (Mixture-of-Depths: Dynamically allocating compute in transformer-based language models)

소개

전통적으로, 트랜스포머 모델은 입력 시퀀스의 모든 부분에 동일한 계산 자원을 할당합니다. 하지만 이번에 소개하는 MoD(Mixture-of-Depth) Transformer 논문은 Trnasformer의 특정 시퀀스 위치에 컴퓨트 자원을 동적으로 할당하고, 다른 레이어를 통해 할당을 변화시키는 시도를 합니다. 이러한 Mixture-of-Depths(MoD) 전략을 사용하기 위해 저자들은 어떤 토큰을 처리할지 결정하기 위해 top-k 라우팅 메커니즘을 사용하여 전체 연산량을 관리합니다. 이 기술은 기존 모델의 성능에 필적하거나 뛰어나면서도, 훨씬 적은 FLOPs를 요구하며, 학습 후 샘플링 동안 최대 50%의 속도 향상을 제공합니다.

Attention is All You Need 논문에 소개된 Transformer 구조

최근의 트랜스포머 모델과 LLM(Large Language Model)은 그 규모와 복잡성이 지속적으로 증가하고 있습니다. 이러한 모델들은 인간 수준의 언어 이해와 생성 능력에 접근하기 위해 수십억 개의 파라미터를 활용하며, 그 결과로 놀라운 성능을 보여주고 있습니다. 그러나, 이러한 발전에는 상당한 계산 비용과 에너지 소비가 수반되며, 모델 규모의 증가는 학습과 추론 시간을 급격히 증가시킵니다.

이는 Transformer가 입력 시퀀스의 모든 토큰에 대해서 동일한 양의 연산을 수행하고 있기 때문입니다. 하지만, 어떤 토큰은 문맥 이해에 있어 핵심적인 역할을 하지만 어떤 토큰은 상대적으로 덜 중요하는 등, 모든 토큰들이 동등한 양의 정보를 제공하지 않을 수 있다는 아이디어에서 출발을 합니다.

MoD Transformer는 각 토큰에 대해 동적으로 계산을 할당함으로써, 모델이 필요한 곳에 자원을 집중시킬 수 있게 합니다. 이는 모델의 전체적인 연산 비용을 줄이는 동시에, 중요한 토큰에 더 많은 주의를 기울임으로써 정보 처리의 효율성을 높일 수 있게 됩니다. 특히, LLM과 같은 대규모 모델에서 MoD의 접근 방식은 학습과 추론 과정에서의 시간과 자원을 절약할 수 있는 중요한 방법이 될 수 있습니다.

MoD(Mixture-of-Depths) 기법 소개

MoD 기법은 트랜스포머 레이어의 각각에서 어떤 토큰을 처리할지 동적으로 선택하는 top-k 라우팅 메커니즘을 통합합니다. 이를 통해 선택된 일부 토큰들만 전체 연산 과정을 거치며, 다른 토큰들은 레이어(들)을 건너뛰게 됩니다.

그림1 / 좌측: MoD(Mixture-of-Depths) Transformer 개요: MoE(Mixture-of-Experts) 구조와 유사하게 라우터를 통해 연산 경로를 선택합니다. 이 때, MoE와는 다르게 표준 블럭(Self-Attention 및 MLP)을 연산할지 말지를 결정하게 됩니다. 이 때 입력 토큰(X_i)이 라우터의 결정에 따라 전체 연산을 활성화하지 않는 경우, 해당 레이어의 연산을 거치지 않게 됩니다.
그림 1 / 우측: 64개 토큰으로 축소된 짧은 시퀀스에 대한 라우팅 결정 결과입니다. (X축: 시퀀스, Y축: 레이어) 보라색은 해당 레이어의 연산을 모두 진행하였으며, 주황색은 해당 레이어의 연산을 건너뛴 것입니다.

위 그림의 설명과 같이, MoD는 트랜스포머 아키텍처 내에 예측 라우터를 구현하여 각 토큰마다 레이어 연산의 진행 여부를 결정합니다. 이를 통해 연산 자원을 상대적으로 "중요한" 토큰에 집중할 수 있습니다.

MoD Transformer의 학습

MoD 트랜스포머의 학습 과정은 기본적으로 표준 트랜스포머와 유사합니다. 입력 시퀀스는 모델에 공급되고, 각 레이어를 통해 전달되면서 필요한 계산이 수행됩니다. 그러나 MoD 트랜스포머는 라우터를 통해 각 토큰이 어떤 계산 경로를 따를지 결정하는 추가적인 단계를 포함합니다.

라우터 가중치 계산: 라우터는 모델의 각 레이어에서 각 토큰에 대한 가중치를 계산합니다. 이 가중치는 해당 토큰이 받아야 할 계산의 양을 나타내며, 토큰의 중요성에 기반합니다.
상위 𝑘 토큰 선택: 라우터는 계산된 가중치를 기반으로, 각 레이어에 대해 상위 𝑘개의 토큰을 선택합니다. 이 토큰들은 전체 계산을 수행하는 경로로 안내됩니다.
연산 경로 할당: 선택된 토큰은 표준 계산(Self-Attention 및 MLP)을 수행하는 경로로 할당됩니다. 나머지 토큰은 잔여 연결을 통해 전달되어, 추가 계산 없이 다음 레이어로 넘어갑니다.
손실 계산 및 역전파: 모델의 출력이 최종 손실 함수를 통해 평가되고, 그래디언트가 모델을 통해 역전파되어 가중치를 업데이트합니다. 이 과정에서 라우터도 학습되어, 시간이 지남에 따라 더 효율적인 라우팅 결정을 내릴 수 있습니다.

MoD Transformer의 라우팅 방식 (Routing Schemes)

라우팅 구현은 MoD Transformer의 핵심적인 부분입니다. 저자들은 3가지 라우팅 방식을 비교하여 Expert-Choice MoD 방식의 토큰 라우팅이 가장 나은 성능을 보이는 것을 확인하였습니다. 소개된 다른 라우팅 방식들과의 비교를 해보겠습니다:

토큰 선택(Token-choice) 라우팅

토큰 선택 라우팅에서는 각 토큰이 자신이 거칠 경로를 스스로 선택합니다. 모델은 각 토큰에 대해 다양한 계산 경로(예: 다른 레이어 또는 처리 블록)에 대한 확률 분포를 생성하고, 이 분포에 따라 토큰을 해당 경로로 전달합니다.

토큰 선택 라우팅은 각 토큰이 가장 적합한 계산 경로를 선택할 수 있어, 매우 유연한 라우팅이 가능합니다. 또한, 모델이 각 토큰의 중요도나 문맥을 더 정밀하게 고려할 수 있어, 개별 토큰에 최적화가 가능합니다.

전문가 선택(Expert-choice) 라우팅

전문가 선택 라우팅에서는 모델이 각 경로(또는 "전문가")가 처리할 토큰을 직접 선택합니다. 이는 각 토큰의 라우터 가중치를 기반으로 상위 𝑘개의 토큰을 선별하는 top-𝑘 메커니즘을 사용하여 이루어집니다.

이렇게 전문가 선택 라우팅을 하게 되면 토큰이 각 경로에 균등하게 분배될 수 있어, 연산 부하가 균형있게 분산되므로 모델의 연산 자원을 보다 효율적으로 활용할 수 있습니다. 또한, 상위 𝑘개의 토큰이 명확하게 결정되므로, 모델의 예측 가능성과 안정성이 향상됩니다.

전문가 선택 MoD(Expert-choice MoD)

전문가 선택 MoD는 전문가 선택 라우팅의 개념을 MoD 트랜스포머의 구조에 적용한 것입니다. 이 방식에서는 모델이 각 레이어 또는 처리 단계에 대해 상위 𝑘개의 토큰을 선택하고, 선택된 토큰만이 계산을 수행하는 경로를 따라 처리됩니다. 나머지 토큰은 잔여 연결을 통해 비용이 적게 드는 경로를 따릅니다.

이러한 방식은 선택된 토큰만이 계산을 수행하므로, 전체 모델의 FLOP 사용량이 크게 감소합니다. 또한, 중요한 토큰에 대해서는 전체 레이어에 대한 연산을 보장하면서도, 불필요한 연산은 최소화하여 모델의 성능을 유지하거나 향상시킬 수 있습니다. 마지막으로 라우터를 통한 동적 라우팅을 통해 모델이 입력 시퀀스의 특성에 따라 계산 자원을 유동적으로 할당할 수 있는 장점이 있습니다.

라우팅 구현 방식

라우팅의 구현은 다음 3가지 주요 단게를 통해 가능합니다:

토큰 임베딩과 라우터 가중치 결정: 모델의 각 레이어에 대해, 주어진 시퀀스의 토큰 임베딩 집합 X^i 이 있습니다. 각 토큰 임베딩 x^l_i 에 대해서, 라우터는 선형 변환을 통해 해당 토큰이 연산에 참여할 가능성을 나타내는 스칼라 가중치 r^l_i 를 계산합니다.
상위 𝑘 가중치 결정: 토큰의 라우터 가중치 집합 R^l 에서 \beta 번째 가중치의 백분위 확률 P_\beta(R^l) 을 찾습니다. 이 때 \beta 는 모델이 각 레이어에서 처리할 토큰의 비율을 결정합니다. 따라서 라우터 가중치가 P_\beta(R^l) 보다 큰 토큰이 연산에 참여하게 됩니다.
연산 경로 선택: 각 토큰은 라우터 가중치 r^l_i 과 비교하여 P_\beta(R^l) 보다 크거나 같으면 해당 레이어의 연산을 처리하며, 작으면 다음 레이어로 직접 전달됩니다. 이를 통해 선택된 토큰에 대해서만 계산을 수행하고, 나머지 토큰은 계산 비용을 발생시키지 않는 경로를 따라 처리할 수 있게 됩니다.

이러한 라우팅 구현 시 다음과 같은 내용을 핵심적으로 고려해야 합니다:

가중치에 기반한 동적 결정: 라우팅 과정은 각 토큰의 가중치를 기반으로 동적으로 결정되며, 이는 토큰의 중요도를 반영합니다. 중요한 토큰은 더 많은 계산을 받아 모델의 성능을 최적화하는 반면, 덜 중요한 토큰은 잔여 경로를 통해 비용을 절감합니다.
계산 효율성과 성능 최적화: 이러한 라우팅 메커니즘을 통해 MoD 트랜스포머는 계산 자원을 효율적으로 사용하면서도, 모델의 성능을 유지하거나 향상시킬 수 있습니다. 특히, 계산 비용이 높은 레이어에서 중요한 토큰만을 처리함으로써 전체적인 FLOP 사용량을 줄일 수 있습니다.
구현의 단순성과 효율성: 라우팅은 선형 변환과 백분위수 계산을 통해 구현되므로, 상대적으로 단순하며 효율적입니다. 이는 모델 학습과 추론 과정에서 계산 부담을 추가하지 않으면서도, 계산 자원의 동적 할당을 가능하게 합니다.

MoD Transformer의 동작 예시

지금까지 살펴본 내용에 따라 MoD Transformer의 동작을 예시를 통해 설명하겠습니다.

예를 들어, 입력 시퀀스에 100개의 토큰이 있고, 라우터가 𝑘를 10으로 설정했다고 가정해 봅시다. 이 경우, 각 레이어에서 라우터는 가장 중요한 10개의 토큰을 선택하고, 이 토큰들만이 표준 계산 경로를 따라 처리됩니다.

나머지 90개의 토큰은 계산을 건너뛰고, 잔여 연결을 통해 다음 레이어로 직접 전달됩니다. 이 과정을 통해, MoD 트랜스포머는 전체 FLOP 사용량을 줄이면서도 중요한 토큰에 대해 필요한 계산을 보장합니다.

MoD Transformer의 성능 비교

학습 성능

하이퍼파라미터 최적화: 12.5% 용량 라우팅 블록을 전체 주의(attention) 블록과 교차하여 배치한 MoD 트랜스포머 변형을 학습시키는 실험을 통해, 최적의 하이퍼파라미터 설정을 찾았습니다. 이 설정은 isoFLOP 분석에서 최적 기준선보다 더 나은 성능을 제공했습니다.

계산 효율성: 특정 MoD 변형은 기존 모델 대비 더 적은 FLOP으로 동일한 성능을 달성하면서, 66% 더 빠른 처리 속도를 보였습니다.

라우팅 분석

라우팅 결정의 희소성: 12.5% 용량(capacity) 라우팅을 사용함으로써, 대부분의 토큰은 블록을 우회하고 소수의 중요한 토큰만이 블록을 통과했습니다. 이는 라우터 가중치 분포에서도 확인할 수 있으며, 네트워크는 깊이에 따라 특정 토큰을 선호하여 라우팅하기도 했습니다.

추론 성능

자동 회귀 추론: 학습 시 비인과적(non-causal) top-𝑘 라우팅 스킴에서 추론 시 인과적(causal) 예측 기반 접근법으로 전환하였을 때, 성능 저하가 거의 없었습니다. 이는 해당 예측 문제의 학습이 쉽고, 학습 초기부터 97% 이상의 정확도를 달성했기 때문으로 보입니다.

Mixture-of-Depths-and-Experts(MoDE) = MoD + MoE

MoDE(Mixture-of-Depths-and-Experts)는 MoD(Mixture-of-Depths) 기법과 MoE(Mixture of Experts) 기법의 통합을 의미합니다. 이 방식은 두 가지 접근 방식의 장점을 결합하여, 트랜스포머 모델의 계산 효율성과 성능을 동시에 최적화하는 데 목적이 있습니다. MoD 방식이 각 토큰에 대해 더 깊이 또는 덜 깊이 계산을 수행하도록 선택하는 반면, MoE 방식은 다양한 "전문가" 중에서 가장 적합한 전문가를 통해 토큰을 처리합니다. MoDE는 이 두 가지 접근 방식을 통합하여, 토큰이 특정 블록을 건너뛰거나 선택적으로 더 깊은 계산을 수행할 수 있게 함과 동시에, 다양한 전문가 중에서 최적의 처리 방식을 선택할 수 있도록 합니다.

MoDE의 구현 방식

MoDE는 두 가지 주요 구현 방식을 가집니다: 단계적(Staged) MoDE와 통합된(Integrated) MoDE.

Staged MoDE(단계적 MoDE): 단계적 MoDE는 먼저 MoD 기법을 적용하여 각 토큰의 계산 경로를 결정하고, 그 다음에 MoE 기법을 적용하여 선택된 경로에서 처리될 토큰을 각각의 전문가 블록으로 라우팅합니다. 이렇게 하면, 모델은 먼저 계산의 깊이를 조절한 후, 각 단계에서 최적의 전문가를 선택할 수 있습니다.
Integrated MoDE(통합된 MoDE): 통합된 MoDE는 MoD와 MoE 기법을 하나의 라우팅 단계에서 통합하여 구현합니다. 즉, 토큰을 라우팅할 때 한 번의 결정으로 계산 경로의 깊이와 처리할 전문가를 동시에 선택합니다. 이 방식은 구현이 더 단순하고 효율적일 수 있으며, 토큰마다 가장 적합한 계산 경로와 전문가를 동시에 결정할 수 있는 장점이 있습니다.

MoDE의 장점

계산 효율성: MoDE는 토큰마다 필요한 계산의 양과 처리할 전문가를 동적으로 결정함으로써, 모델의 전체적인 계산 비용을 효과적으로 관리합니다.
성능 최적화: 각 토큰을 가장 적합한 계산 경로와 전문가에 할당함으로써, MoDE는 모델의 성능을 최대화합니다. 이는 특히 복잡한 문제를 해결하거나 다양한 종류의 데이터를 처리할 때 유용할 수 있습니다.
유연성과 확장성: MoDE는 다양한 모델 구조와 태스크에 적용될 수 있는 유연성을 제공하며, 모델의 계산 효율성과 성능을 동시에 향상시킬 수 있는 확장 가능한 프레임워크를 제공합니다.

MoDE는 트랜스포머 모델의 미래 발전 방향을 제시하며, 계산 비용을 최소화하면서도 모델의 성능을 극대화하는 방법을 탐색하는 연구에 중요한 기여를 합니다.

더 읽어보기

이 글은 GPT 모델로 정리한 글을 바탕으로 한 것으로, 원문의 내용 또는 의도와 다르게 정리된 내용이 있을 수 있습니다. 관심있는 내용이시라면 원문도 함께 참고해주세요! 읽으시면서 어색하거나 잘못된 내용을 발견하시면 덧글로 알려주시기를 부탁드립니다.

파이토치 한국 사용자 모임이 정리한 이 글이 유용하셨나요? 회원으로 가입하시면 주요 글들을 이메일로 보내드립니다! (기본은 Weekly지만 Daily로 변경도 가능합니다.)

아래쪽에 좋아요를 눌러주시면 새로운 소식들을 정리하고 공유하는데 힘이 됩니다~