[서론]
문제 정의
트랜스포머 모델은 모든 토큰에 동일한 연산을 수행하며, 이는 불필요한 계산 자원을 소모합니다. 이상적으로는 토큰의 중요도에 따라 연산 자원을 동적으로 할당하여 총 연산 비용을 줄여야 합니다. 이 문제를 해결하기 위해 조건부 계산(Conditional Computation) 기법이 제안되었습니다. 그러나 기존 방법들은 동적 계산 그래프를 도입하여 하드웨어 제약을 초래합니다. 본 논문에서는 고정된 계산 예산 하에서 토큰별로 동적으로 연산 자원을 할당하는 방법을 제안합니다.
선행 연구
조건부 계산 기법은 Bengio et al. (2016)에서 제안되었으며, 이후 다양한 연구들이 이를 발전시켰습니다. Ainslie et al. (2023), Bapna et al. (2020), Fedus et al. (2022) 등은 토큰의 중요도에 따라 연산을 조절하는 방법을 제안했습니다. 그러나 이들 방법은 기존 하드웨어와의 호환성 문제를 야기할 수 있습니다.
본 논문의 기여
본 논문에서는 고정된 계산 예산을 사용하여 토큰별로 연산 자원을 동적으로 할당하는 Mixture-of-Depths (MoD) 기법을 제안합니다. MoD는 각 레이어에서 토큰별로 필요한 연산만 수행하여 총 연산 비용을 줄입니다. 이를 통해 성능을 유지하면서도 연산 비용을 절감할 수 있습니다.
[방법]
데이터셋 및 벤치마크
본 논문에서는 다양한 언어 모델 데이터셋과 벤치마크를 사용하여 제안한 기법의 성능을 평가합니다. 실험에 사용된 주요 데이터셋은 다음과 같습니다.
MoD 개요
MoD의 기본 개념는 Mixture of Experts(MoE)에 기반을 두고 있습니다. 이 방식은 각 토큰에 대해 표준 트랜스포머 블록의 계산을 수행하거나, 계산을 생략하고 잔차 연결을 통해 토큰을 전달하는 두 가지 선택지를 제공합니다. 이는 다음과 같은 수학적 표현으로 나타낼 수 있습니다.
\[\text{if } r_i > \beta(R) \text{ then } x_{i+1} = f(x_i) + x_i \text{ else } x_{i+1} = x_i\]수식에서 \(r_i\)는 라우터가 토큰 \(i\)에 할당한 가중치, \(\beta(R)\)는 설정된 임계값, \(f(x_i)\)는 토큰에 적용된 변환 함수입니다. 이 방법은 불필요한 계산을 줄이면서도 네트워크가 주어진 컴퓨팅 예산 내에서 최적의 성능을 달성하도록 합니다.
[방법]
트랜스포머
트랜스포머 모델에서 모든 토큰이 동일한 연산을 수행하면 불필요한 계산 자원이 낭비될 수 있습니다. 이를 해결하기 위해 특정 토큰들만 선택적으로 연산에 참여시킴으로써 연산량을 줄입니다. 즉, '용량(Capacity)' 개념을 도입하여 각 레이어에서 처리할 토큰의 수를 제한합니다.
용량의 정의 용량 $C$는 각 레이어에서 연산에 참여할 수 있는 최대 토큰 수를 의미합니다.
\[C = \text{total number of tokens participating in computation}\]라우팅 기법
라우팅 기법은 각 토큰의 중요도를 평가하여, 중요도가 높은 토큰만 연산에 참여시키는 방법입니다. 이를 위해 각 토큰에 대해 스칼라 값을 계산하고, 이 값을 기반으로 연산 참여 여부를 결정합니다.
라우터 스칼라 값 계산
각 토큰 $x_i$는 라우터 파라미터 $\theta$와의 내적을 통해 스칼라 값 $r_i$를 얻습니다.
\[r_i = \mathbf{w}^T \mathbf{\theta} x_i\]상위 $k$ 토큰 선택
각 레이어에서 상위 $k$개의 스칼라 값을 가진 토큰만을 선택하여 연산에 참여시키고, 상위 $k$ 값의 기준인 퍼센타일 $P_\beta$를 계산합니다.
\[P_\beta(R) = \beta\text{-th percentile of } R\]라우팅 결정
각 토큰의 스칼라 값 $r_i$와 퍼센타일 값 $P_\beta(R)$를 비교하여 연산 참여 여부를 결정합니다.
\[x_i^{l+1} = \begin{cases} f(x_i^l) + x_i^l & \text{if } r_i > P_\beta(R) \\ x_i^l & \text{if } r_i \leq P_\beta(R) \end{cases}\]연산 경로 선택
선택된 토큰은 셀프 어텐션 및 MLP 연산을 통해 업데이트되며, 선택되지 않은 토큰은 입력 값 그대로 유지됩니다. 이를 통해, 전체 연산 비용을 줄이면서도 필요한 연산만 수행하게 됩니다.
[라우터의 학습]
라우터는 학습 가능한 파라미터 $\theta$를 통해 각 토큰의 중요도를 학습합니다. 이를 위해 다음과 같은 손실 함수를 사용합니다.
\[\mathcal{L} = \mathcal{L}_{LM} + \lambda \mathcal{L}_{aux}\]라우터의 학습을 통해, 모델은 각 토큰의 중요도를 정확히 평가하고, 불필요한 연산을 줄이면서도 성능을 유지할 수 있다고 언급합니다.
학습 방법 라우팅 기법을 학습하기 위해 다음과 같은 손실 함수를 사용합니다.
\[\mathcal{L} = \mathcal{L}_{LM} + \lambda \mathcal{L}_{aux}\]수식에서 \(\mathcal{L}_{LM}\)은 언어 모델링 손실, \(\mathcal{L}_{aux}\)는 라우터의 예측 정확도를 높이기 위한 보조 손실입니다.
[실험 및 결과]
훈련 및 isoFLOP 비교
제안한 MoD 기법을 사용하여 다양한 모델을 훈련하고 isoFLOP 분석을 수행했습니다. MoD 트랜스포머는 벤치마크 데이터셋에서 기존 트랜스포머 모델 대비 연산 비용을 줄이면서도 성능을 유지하거나 향상시켰습니다.
토큰 라우팅 분석
토큰별 라우팅 결과를 분석한 결과, MoD 트랜스포머는 특정 토큰에 대해 더 많은 연산을 할당하여 예측 정확도를 높였습니다.
자가회귀 샘플링 평가
자가회귀 샘플링 동안에도 MoD 기법은 성능 저하 없이 연산 비용을 절감할 수 있음을 확인했습니다.
[결론]
본 논문에서는 고정된 계산 예산 하에서 토큰별로 연산 자원을 동적으로 할당하는 Mixture-of-Depths (MoD) 기법을 제안했습니다. MoD 기법은 기존 트랜스포머 모델 대비 성능을 유지하거나 향상시키면서도 연산 비용을 절감할 수 있음을 실험적으로 확인했습니다.