기술적 요약: 트랜스포머에서 학습된 토큰 라우팅을 통한 적응적 계산 깊이

문제 제기

표준 트랜스포머 아키텍처는 시퀀스 내의 모든 토큰에 대해 고정된 수의 레이어를 적용합니다. 이는 해당 토큰의 맥락적 난이도나 정보 밀도와 무관하게 수행됩니다. 이러한 설계는 토큰별 적응성보다 아키텍처의 단순성을 우선시하여, 상당한 계산 비효율성을 초래합니다. 실제로 예측 가능한 맥락에서 자주 등장하는 토큰은 새로운 구조에서 드문 토큰보다 훨씬 적은 처리가 필요함에도 불구하고, 모든 순전파 (forward pass) 동안 동일한 계산 자원을 할당받습니다. 이러한 비효율성은 추론 규모에서 특히 비용이 큽니다. 왜냐하면 주된 비용은 모든 토큰에 대한 모든 레이어를 통한 완전한 순전파이기 때문입니다. RNN 을 위한 적응적 계산 시간 (ACT), 유니버설 트랜스포머, 깊이 혼합 (MoD), 조기 종료 (Early Exit) 메커니즘과 같은 기존 솔루션들은 이를 해결하려 시도하지만, 종종 완전히 미분 가능하거나 최적이지 않을 수 있는 하드 선택, 누적 예산, 또는 사후 신뢰도 임계값에 의존합니다.

방법론: 토큰 선택적 어텐션 (TSA)

본 논문은 연속적인 트랜스포머 블록 간의 잔차 업데이트에 토큰별 학습된 게이트를 도입하는 메커니즘인 **토큰 선택적 어텐션 (Token-Selective Attention, TSA)**을 제안합니다.

아키텍처 및 라우팅

TSA 는 프리-노름 (pre-norm) 디코더 전용 트랜스포머에서 각 블록 $f_l$ ($l = 0, \dots, L-2$) 이후에 경량 라우터 $r_l$을 삽입합니다. 스템 블록 $f_0$은 항상 무조건 실행됩니다.

• 라우터 설계: 각 라우터는 시그모이드 출력을 가진 2 층 멀티 레이어 퍼셉트론 (MLP) 으로, 각 토큰에 대해 연속적인 정지 확률 $p_l \in (0, 1)$을 생성합니다. 은닉 차원은 $d/4$ (최소 16 으로 내림) 이며, 기본 모델에 약 1.7% 의 파라미터 오버헤드를 추가합니다.
• 게이트된 업데이트: 후속 블록 $f_{l+1}$에 대한 잔차 업데이트는 $(1 - p_l)$로 스케일링됩니다.
  • $p_l = 0$인 경우, 업데이트는 표준 방식입니다.
  • $p_l = 1$인 경우, 상태는 변경되지 않아 블록이 효과적으로 건너뜁니다.
  • 보간은 매끄러우므로, 스트레이트-쓰루 추정기, Gumbel 샘플링, 또는 강화 학습 없이도 메커니즘이 완전히 미분 가능하도록 보장합니다.

학습 및 정규화

학습 중에는 모든 레이어가 완전히 실행됩니다. 게이트는 잔차 업데이트를 스케일링하지만 계산을 건너뛰지는 않습니다. 총 손실은 $L = L_{task} + L_{depth}$로 정의됩니다.

• 깊이 정규화: 조기 정지를 부드럽게 장려하기 위해 항 $L_{depth} = \lambda \cdot \frac{1}{L-1} \sum (1 - p_l)$이 추가됩니다.
• 창발적 행동: 주목할 점은, $\lambda = 0$(명시적 깊이 압력 없음) 일 때도 라우팅 메커니즘이 작업 손실 기울기만으로 작동하여 연산을 건너뛰는 것을 학습한다는 것입니다.

추론

추론 시 메커니즘은 "희소 (sparse)" 모드로 전환됩니다. $p_l > 0.5$인 토큰은 gather/scatter 연산을 통해 피드포워드 네트워크 (FFN) 에서 제외되지만, 정확한 키 - 값 의미론을 보존하기 위해 어텐션은 밀집 (dense) 상태로 유지됩니다. 이는 실제 FLOPs 절감과 월클럭 속도 향상을 가져옵니다.

주요 기여

1. 미분 가능한 토큰 라우팅: 기본 트랜스포머의 아키텍처 변경 없이 토큰별, 레이어별 라우팅을 가능하게 하는 잔차 업데이트에 대한 단순하고 연속적인 소프트 게이트.
2. 기울기 주도 창발: 작업 손실 기울기만으로 난이도에 비례하는 라우팅이 자연스럽게 창발한다는 증거. 깊이 정규화 ($\lambda=0$) 가 없더라도 라우터는 토큰 - 레이어 연산의 약 20% 를 건너뛰는 것을 학습합니다.
3. 효율성 - 품질 트레이드오프: 문자 단위 언어 모델링 검증에서 토큰 - 레이어 연산 (TLOps) 을 14~23% 절감하면서도 0.5% 미만의 품질 손실.
4. 조기 종료 대비 우월성: 일치된 효율성에서 TSA 는 조기 종료 베이스라인보다 낮은 검증 손실을 달성하여, 사후 신뢰도 임계값 적용보다 엔드 - 투 - 엔드 학습된 라우팅 결정이 우월함을 시사합니다.
5. 실제 속도 향상: 상용 하드웨어 (Apple M1 Pro) 에서 희소 추론을 통해 최대 2.3% 의 실제 월클럭 속도 향상 입증.

실험 결과

합성 작업

길이 10 시퀀스의 복사 및 정렬 작업에서:

• 복사 작업: 라우터가 작업이 항등 매핑임을 학습하여 정확도 손실 없이 TLOps 의 54.9% ($\alpha = 0.341$) 를 건너뜁니다.
• 정렬 작업: 라우터가 더 많은 계산을 할당하여 TLOps 의 22.5% ($\alpha = 0.730$) 를 건너뛰며, 이는 더 높은 난이도를 반영합니다.

문자 단위 언어 모델링

• Tiny-Shakespeare: TSA 는 검증 손실 증가가 0.006 nats(+0.4%) 에 불과하면서 TLOps 의 22.8% ($\alpha = 0.726$) 를 절감했습니다.
• enwik8: 더 다양한 코퍼스 (위키백과) 에서 TSA 는 품질 저하 없이 (검증 손실이 약간 더 좋음) TLOps 의 13.9% ($\alpha = 0.833$) 를 절감했습니다.
• 강건성: 메커니즘은 두 개의 차수 ($[0, 0.1]$) 에 걸쳐 정규화 하이퍼파라미터 $\lambda$에 대해 강건합니다. $\lambda=0.05$에서 50.4% 의 TLOps 절감을 달성했으며 품질 비용은 무시할 수 있었습니다.

조기 종료와의 비교

일치된 활성 분율 ($\alpha \approx 0.726$) 에서 TSA 는 조기 종료 베이스라인보다 0.71% 낮은 검증 손실을 달성했습니다. 이는 추론 시 신뢰도 임계값을 적용하는 것보다 엔드 - 투 - 엔드 라우팅 결정을 학습하는 것이 더 나은 결정임을 나타냅니다.

월클럭 처리량

• Soft-TSA(학습): $\alpha$와 무관하게 약 1% 의 오버헤드가 추가되었습니다.
• Sparse-TSA(추론): Tiny-Shakespeare($\alpha=0.726$) 에서 2.3% 속도 향상을 달성했으며 enwik8($\alpha=0.833$) 에서는 브레이크이븐 성능을 보였습니다. 속도 향상은 CPU-GPU 동기화 오버헤드를 극복하기 위해 배치 크기가 $\ge 64$여야 했습니다.

중요성 및 주장

본 논문은 TSA 가 고정 깊이 트랜스포머의 비효율성에 대한 최소적이고 아키텍처적으로 단순한 해결책을 제공한다고 주장합니다. 그 주요 중요성은 난이도에 대한 명시적 감독이나 복잡한 강화 학습 신호 없이 적응적 계산 깊이를 작업 목적에서 직접 학습할 수 있음을 입증하는 데 있습니다.

저자들은 라우터가 잔차 기여가 중복될 때 업데이트를 감쇠시키는 "학습된 노이즈 게이트"로 작용한다고 강조합니다. 이를 통해 모델은 토큰 난이도에 따라 동적으로 계산을 할당하여 합성 작업에서는 최대 55%, 언어 모델링에서는 14~23% 의 상당한 효율성 향상을 달성하면서도 최소한의 품질 트레이드오프를 보입니다. 이 연구는 이산적 선택이나 사후 임계값 적용보다 엔드 - 투 - 엔드 미분 가능 라우팅이 적응적 깊이를 위한 더 효과적인 전략임을 시사하며, 상용 하드웨어에서 더 효율적인 추론을 위한 길을 열고 있습니다.

논문의 규모에 대해서는 겸손하게, 모든 실험이 약 500 만600 만 파라미터 모델을 대상으로 수행되었으며, 더 큰 모델 (1000 만1 억+) 로의 확장 행동은 여전히 열린 질문이라고 명시하고 있습니다. 또한 현재 속도 향상은 배치 크기에 의존하며, 어텐션 희소화는 아직 구현되지 않았습니다.