기술 요약: 심층 디코더 전용 트랜스포머를 위한 다중 해상도 잔차 라우팅 (Multi-Resolution Residual Routing)

문제 정의

표준 디코더 전용 트랜스포머는 서브레이어 업데이트가 고정된 단위 가중치와 함께 집계되는 PreNorm 잔차 연결(residual connections)에 의존합니다. 이러한 구조는 안정적인 학습을 가능하게 하지만, 네트워크 깊이가 깊어짐에 따라 개별 레이어의 기여도를 희석시킬 수 있는 경직된 누적 규칙을 부과합니다. 최근의 발전인 Attention Residuals 및 **Block Attention Residuals (Block AttnRes)**는 고정된 누적 대신 이전 출력들에 대한 학습된 콘텐츠 의존적 라우팅을 사용하여 이 문제를 해결합니다. 특히 Block AttnRes는 여러 레이어를 블록 수준의 잔차 요약($C_b$)으로 압축함으로써 효율성을 개선합니다.

그러나 저자들은 Block AttnRes의 결정적인 한계를 지적합니다. 단일 블록 요약은 블록 내 잔차 궤적(residual trajectory)의 저주파(DC) 성분만을 나타낸다는 점입니다. 업데이트의 합($\sum u_{b,i}$)만을 저장함으로써, 이 메커니즘은 다음과 같은 귀중한 **방향성 구조(directional structure)**를 유실합니다:

1. 서브레이어 간 불균형: 해당 블록이 Attention 업데이트에 의해 지배되는지 또는 MLP 업데이트에 의해 지배되는지 여부.
2. 시간적 역학(Temporal dynamics): 블록 내 초기 업데이트가 후기 업데이트와 다른 방향을 가리키는지 여부.

저자들은 잔차 궤적의 길이가 길어질수록, 이 유실된 방향성 정보가 효과적인 깊이별 라우팅(depth-wise routing)을 위해 점점 더 중요해진다고 가설을 세웁니다.

방법론: WAV v1

저자들은 Block AttnRes의 경량화된 다중 해상도 확장판인 WAV v1(Waveform Attention v1)을 제안합니다. WAV v1은 블록을 오직 총합으로만 표현하는 대신, 기존의 Attention 및 MLP 모듈을 유지하면서 모든 블록 표현에 두 개의 **제로 합(zero-sum) 방향성 디테일 기저(detail bases)**를 추가로 부여합니다.

1. 다중 해상도 표현

$m$개의 서브레이어 업데이트 $\{u_{b,i}\}_{i=1}^m$를 포함하는 블록 $b$에 대해, WAV v1은 다음을 계산합니다:

• 상태 기저 (State Basis, $C_b$): 저주파 상태를 나타내는 표준 블록 합입니다.
  $$C_b = \sum_{i=1}^m u_{b,i}$$
• 위상 기저 (Phase Basis, $D^{phase}_b$): Attention과 MLP 업데이트를 대조하는 디테일 소스입니다. Attention 업데이트에는 $+1$, MLP 업데이트에는 $-1$의 부호를 할당합니다.
  $$D^{phase}_b = \sum_{i=1}^m a_i u_{b,i}, \quad \text{단, } a_i \in \{+1, -1\}$$
• 분할 기저 (Split Basis, $D^{split}_b$): 블록의 전반부와 후반부를 대조하는 디테일 소스입니다. 전반부에는 $+1$, 후반부에는 $-1$을 할당합니다.
  $$D^{split}_b = \sum_{i=1}^m r_i u_{b,i}, \quad \text{단, } r_i \in \{+1, -1\}$$

이 기저들은 업데이트가 생성될 때 온라인으로 누적되며, 계산 오버헤드는 미미합니다.

2. 라우팅 메커니즘

깊이별 소프트맥스 믹서(depth-wise softmax mixer)는 확장된 소스 풀을 전달받습니다. Block AttnRes의 경우 풀은 $\{e, C_0, \dots, C_{b-1}, PC\}$이지만, WAV v1은 각 완료된 블록과 현재 진행 중인 부분 블록의 디테일 기저를 포함하도록 확장됩니다:
$$S_{WAV} = \{e, C_0, \tilde{D}^{phase}_0, \tilde{D}^{split}_0, \dots, C_{b-1}, \tilde{D}^{phase}_{b-1}, \tilde{D}^{split}_{b-1}, PC, \tilde{P}^{phase}, \tilde{P}^{split}\}$$

3. 안정화 메커니즘

학습 초기 단계에서의 불안정성을 방지하기 위해, WAV v1은 두 가지 보수적인 전략을 채택합니다:

• 음의 디테일 편향 (Negative Detail Bias): 디테일 소스는 강한 음의 편향($\beta_D = -2.0$)으로 초기화되어, 모델이 초기에는 표준 Block AttnRes처럼 동작하고 유익할 때만 디테일을 활용하도록 보장합니다.
• 분리된 RMS 매칭 (Detached RMS Matching): 디테일 텐서는 최대 스케일 인자 $\rho$로 클리핑된, 대응하는 블록 요약($C$)에 대한 분리된(detached) RMS 비율을 사용하여 스케일링됩니다. 이는 디테일이 단순히 더 큰 원시 크기 때문에 지배적인 영향을 미치는 것을 방지합니다.
• 보수적 리드아웃 (Conservative Readout): 기본적으로 최종 예측 믹서는 디테일 소스가 아닌 임베딩과 $C$ 소스만을 읽어 들여 안정성을 더욱 확보합니다.

주요 기여

1. 정식화: 본 논문은 블록 잔차 라우팅을 저주파 상태 정보($C_b$)와 블록 내 궤적 정보(방향성 디테일)를 구분하는 다중 해상도 표현 문제로 프레임화합니다.
2. 알고리즘: 핵심 Attention 또는 MLP 함수를 수정하지 않고도 블록 요약에 위상 및 분할 디테일 기저를 추가하는 드롭인(drop-in) 확장 방식인 WAV v1을 도입합니다.
3. 실증적 증거: 다중 해상도 라우팅의 효용성이 깊이에 의존적이며, 더 깊은 아키텍처에서만 유의미하게 이득이 된다는 것을 보여주는 예비 스케일링 실험을 제시합니다.

실험 결과

저자들은 캐릭터 수준 데이터셋인 TinyStories와 Text8에서 12, 24, 48 레이어 GPT 디코더 전용 모델을 평가했습니다.

• 깊이 의존적 경향:
  • 12 레이어: WAV v1이 Block AttnRes보다 성능이 낮게 나타나며, 이는 방향성 디테일이 얕은 네트워크에서는 중복되거나 노이즈가 될 수 있음을 시사합니다.
  • 24 레이어: WAV v1이 경쟁력을 갖추기 시작하며, TinyStories에서 약간의 개선을 보이고 Text8에서는 대등한 성능을 보입니다.
  • 48 레이어: WAV v1이 두 데이터셋 모두에서 가장 좋은 검증 손실(validation loss)을 달이트하며, Block AttnRes, ReZero, LayerScale을 크게 상회합니다.
• 정량적 이득 (48 레이어):
  • TinyStories: 검증 손실을 0.4960 (Block AttnRes)에서 0.4738 (WAV v1)로 감소시켜, 2.19%의 PPL 감소를 달성했습니다.
  • Text8: 검증 손실을 0.9363에서 0.9305로 감소시켜, 0.57%의 PPL 감소를 달성했습니다.
• 학습 역학: 48 레이어에서 WAV v1은 학습 초기에 Block AttnRes와 격차를 벌리며 지속적인 우위를 유지하는데, 이는 이 개선이 후기 단계의 일시적 현상이 아니라 근본적인 최적화 궤적의 변화임을 나타냅니다.

의의 및 주장

본 논문은 깊은 잔차 라우팅이 단순히 어떤 블록을 읽을 것인가뿐만 아니라, 각 블록의 내부 방향성 구조를 라우팅하는 것에서도 이득을 얻는다고 주장합니다. 결과는 잔차 궤적이 길어질수록 궤적의 "형태"(예: Attention 지배 여부 또는 초기 대비 후기 역학)가 깊이별 믹서에 유익한 신호가 됨을 시사합니다.

저자들은 WAV v1을 기존 방식들과 상호 보완적인 접근법으로 위치시킵니다:

• 업데이트의 크기를 조절하는 ReZero 또는 LayerScale과 달리, WAV v1은 라우터가 사용할 수 있는 표현 자체를 변경합니다.
• 개별 델타를 라우팅하는 Delta Attention Residuals(높은 비용 발생)와 달리, WAV v1은 효율성과 구조적 풍부함 사이의 균형을 맞춘 블록 수준의 다중 해상도 근사치를 제공합니다.

저자들은 향후 잔차 라우팅 방법의 평가 시 깊이를 고려해야 한다고 결론짓습니다. 이는 다중 해상도 정보의 이점이 잔차 궤적의 길이 및 구조와 연관될 수 있기 때문입니다. 또한 시드(seed) 간의 표준 편차 부족, 소규모 캐릭터 수준 모델 사용, 수기로 설계된 디테일 기저 사용 등의 현재 한계를 인정하며, 향후 연구에서 학습 가능한 기저와 더 큰 규모의 토큰 수준 모델을 탐구할 것을 제안합니다.