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

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

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

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

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

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

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

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

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

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

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

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

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

라우팅 구현 방식

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

1. 토큰 임베딩과 라우터 가중치 결정: 모델의 각 레이어에 대해, 주어진 시퀀스의 토큰 임베딩 집합이 있습니다. 각 토큰 임베딩에 대해서, 라우터는 선형 변환을 통해 해당 토큰이 연산에 참여할 가능성을 나타내는 스칼라 가중치를 계산합니다.

2. 상위 ? 가중치 결정: 토큰의 라우터 가중치 집합 R에서 β번째 가중치의 백분위 확률 Pβ을 찾습니다. 이 때 β는 모델이 각 레이어에서 처리할 토큰의 비율을 결정합니다. 따라서 라우터 가중치가 Pβ보다 큰 토큰이 연산에 참여하게 됩니다.

3. 연산 경로 선택: 각 토큰은 라우터 가중치 r과 비교하여 Pβ보다 크거나 같으면 해당 레이어의 연산을 처리하며, 작으면 다음 레이어로 직접 전달됩니다. 이를 통해 선택된 토큰에 대해서만 계산을 수행하고, 나머지 토큰은 계산 비용을 발생시키지 않는 경로를 따라 처리할 수 있게 됩니다.

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

🌑 가중치에 기반한 동적 결정: 라우팅 과정은 각 토큰의 가중치를 기반으로 동적으로 결정되며, 이는 토큰의 중요도를 반영합니다. 중요한 토큰은 더 많은 계산을 받아 모델의 성능을 최적화하는 반면, 덜 중요한 토큰은 잔여 경로를 통해 비용을 절감합니다.

🌑 계산 효율성과 성능 최적화: 이러한 라우팅 메커니즘을 통해 MoD 트랜스포머는 계산 자원을 효율적으로 사용하면서도, 모델의 성능을 유지하거나 향상시킬 수 있습니다. 특히, 계산 비용이 높은 레이어에서 중요한 토큰만을 처리함으로써 전체적인 FLOP 사용량을 줄일 수 있습니다.

🌑 구현의 단순성과 효율성: 라우팅은 선형 변환과 백분위수 계산을 통해 구현되므로, 상대적으로 단순하며 효율적입니다. 이는 모델 학습과 추론 과정에서 계산 부담을 추가하지 않으면서도, 계산 자원의 동적 할당을 가능하게 합니다.

MoD Transformer의 동작 예시

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

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

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

MoD Transformer의 성능 비교

학습 성능