기술 요약: 멀티 게이트 잔차 (MGR)

1. 문제 제기

심층 신경망, 특히 트랜스포머 아키텍처는 소실된 그래디언트를 완화하고 깊은 층의 안정적인 훈련을 가능하게 하기 위해 잔차 연결에 의존합니다. 그러나 PreNorm 패러다임 하의 표준 가법 잔차 업데이트 ($x_{l+1} = x_l + F(LN(x_l))$) 는 깊이가 증가함에 따라 두 가지 치명적인 한계를 겪습니다:

1. 정보 희석: 고정된 집계 메커니즘은 깊이를 균일한 누적 경로로 취급하여, 이전 층들의 정보가 점진적으로 희석되게 만듭니다. 이로 인해 깊은 모듈이 이전 표현들을 선택적으로 주시하기 어려워집니다.
2. 경계 없는 활성화 성장: 잔차 스트림의 누적적 특성으로 인해 "거대한 활성화 (massive activation)" 현상이 발생하여 활성화의 크기가 경계 없이 drifting 됩니다. 이는 수치적 불안정성을 초래하며, 깊은 층들이 영향력을 유지하기 위해 불균형하게 큰 출력을 생성하도록 강요합니다. 이를 PreNorm 희석 문제라고 합니다.

최근 어텐션 잔차 (AttnRes) 와 하이퍼-커넥션 (HC) 과 같은 해결책들이 이러한 문제들을 다루었지만, 심각한 실용적 병목 현상을 도입합니다. AttnRes 는 모든 선행 층의 출력을 집계해야 하므로 심각한 통신 및 메모리 오버헤드를 발생시킵니다. HC 는 은닉 상태를 병렬 스트림으로 확장하지만 복잡한 연산과 메모리 접근 패턴을 도입하여 대규모 분산 훈련에서의 확장성을 저해합니다.

2. 방법론: 멀티 게이트 잔차 (MGR)

저자들은 전체 어텐션의 통신 부담이나 밀집 다중 스트림 토폴로지의 복잡성 없이 활성화 스케일을 안정화하는 새로운 아키텍처인 멀티 게이트 잔차 (MGR) 를 제안합니다. MGR 은 AttnPool(집계기) 과 게이팅 메커니즘(믹서) 을 포함하는 간소화된 설계를 통해 AttnRes 와 HC 의 강점을 통합합니다.

2.1 아키텍처 구성 요소

• 멀티-스트림 잔차: 단일 잔차 스트림 대신 MGR 은 $n$개의 병렬 은닉 벡터 (스트림) 를 유지합니다.
• AttnPool(집계기): AttnRes 에서 차용된 이 모듈은 $n$개의 스트림에서 잠재적 특징들을 집계하여 현재 층의 처리를 위한 컴팩트한 표현을 생성합니다. 층별 학습 가능한 쿼리 벡터 $w^{(\alpha)}_l$ 과 RMSNorm 을 사용하여 Softmax 가중 분포 $\alpha_{i \to l}$ 을 계산함으로써 관련 스트림 상태를 효과적으로 선택합니다.
• 게이팅 메커니즘 (믹서): 스트림을 업데이트하기 위해 MGR 은 현재 스트림 상태 $s_i$ 와 층 출력 $F_l(h_l)$ 을 보간하는 경량 게이팅 메커니즘을 사용합니다. 두 가지 변형이 제안됩니다:
  1. 독립 게이팅: 각 스트림에 대해 Sigmoid 활성화 함수를 사용하여 보간 계수 $\beta_{i \leftarrow l}$ 을 독립적으로 생성합니다.
  2. 경쟁 게이팅: 스트림에 대한 Softmax(기억 게이트 편향 추가) 를 사용하여 정규화된 가중치 분포를 보장하고 엄격한 볼록 결합을 강제합니다.
     업데이트 규칙은 다음과 같습니다: $s'_i = (1 - \beta_i) \odot s_i + \beta_i \odot F_l(h_l)$.

2.2 안정성 분석

이 논문은 MGR 의 안정성에 대한 이론적 근거를 제공합니다. 무한한 누적을 허용하는 표준 ResNet 과 달리, 게이팅 메커니즘은 볼록 결합을 강제합니다.

• 유계 노름: 업데이트를 볼록 합으로 제한함으로써 연산자 노름이 유계로 유지됩니다: $\|x_{l+1}\| \leq \max(\|x_l\|, \|F_l(x_l)\|)$.
• 전역 보장: $L$층 스택의 경우, 활성화 크기는 모든 층의 누적 효과가 아니라 가장 극단적인 단일 층에 의해 제한됩니다. 이는 깊이에 기인한 불안정성을 곱셈적 누적 문제에서 유계 선택 문제로 변환하여 수학적으로 "거대한 활성화"를 방지합니다.

2.3 게이트 편향 초기화

신호 분산이 깊이와 함께 선형적으로 ($O(L)$) 폭발하는 것을 방지하기 위해, 저자들은 게이트 편향 $b_l$ 에 대한 초기화 전략을 유도했습니다. 총 누적 기여도가 $O(1)$ 로 유지되어야 한다는 요구에 기반하여, 게이트 활성화 $g_{init}$ 은 $1/\sqrt{L}$ 로 스케일링되어야 합니다. 결과적으로 편향은 다음과 같이 초기화됩니다:
$$b^{(L)}_{init} \approx \ln \left( \sqrt{\frac{L}{L_{base}}} \left( e^{-b^{(base)}_{init}} + 1 \right) - n \right)$$
이는 훈련 시작 시 게이팅 행동이 정보 전달을 용이하게 하는 "캐리 (carry)" 상태를 유지하면서도 모델이 필요에 따라 스트림을 업데이트하도록 학습할 수 있도록 보장합니다.

2.4 엔지니어링 효율성

여러 스트림으로 인한 메모리 오버헤드를 해결하기 위해 MGR 은 다음을 활용합니다:

• 커널 퓨전: 메모리 트래픽을 최소화하기 위해 lerp(보간) 와 attnpool 연산을 단일 커스텀 커널로 융합합니다 (I/O 를 약 $3nC$읽기와$1nC$ 쓰기로 감소).
• 폴백 역변환 (재계산): 역전파를 위해 모든 중간 스트림 상태를 저장하는 대신, 선형 보간의 가역성 ($\beta$가 작을 때) 을 활용합니다. 역방향 통과 시 입력 스트림을 출력에서 재계산합니다. $\beta$가 큰 이상치를 처리하기 위해 "폴백 역변환" 전략은 상위 $p$개 가장 큰 게이트 값만 저장하여 최소한의 메모리 풋프린트로 정확한 재구성을 보장합니다.

3. 주요 기여

1. 새로운 아키텍처: 전체 깊이 어텐션의 통신 오버헤드 없이 심층 신경망을 안정화하기 위해 멀티-스트림 잔차, 어텐션 풀링, 게이팅을 결합한 MGR 도입.
2. 이론적 안정성: 볼록 결합 게이팅 메커니즘이 활성화 크기를 제한하여 PreNorm 아키텍처에 내재된 "거대한 활성화" 문제를 직접 해결한다는 공식적 증명.
3. 확장 가능한 구현: AttnRes 의 통신 병목 현상이나 밀집 다중 스트림 설계의 메모리 비효율성 없이 MGR 이 대규모 모델로 확장될 수 있도록 하는 엔지니어링 전략 (커널 퓨전 및 선택적 재계산) 개발.
4. 초기화 전략: 안정적인 훈련 역학을 보장하기 위해 모델 깊이에 따라 스케일링되는 유도된 편향 초기화 체계.

4. 실험 결과

저자들은 FineWeb-10BT 데이터셋을 사용하여 세 가지 규모 (S: 0.12B, M: 0.35B, L: 0.77B 파라미터) 와 다양한 스트림 수 ($n=4, 8$) 에서 처음부터 훈련된 언어 모델에 대해 MGR 을 평가했습니다.

• 성능: MGR 은 PreNorm, mHC-lite, Block AttnRes 를 포함한 베이스라인보다 일관되게 우수한 성능을 보였습니다.
  • $n=4$에서 독립 및 경쟁 MGR 변형 모두 최상의 Full AttnRes 베이스라인을 약간 능가했습니다.
  • $n=8$에서 성능 격차는 더욱 커졌으며, 경쟁 MGR 은 모든 규모에서 가장 낮은 훈련 및 검증 손실을 달성했습니다.
• 활성화 안정성: 진단 결과 PreNorm 은 초기 층부터 거대한 활성화 성장을 보인 반면, MGR 은 Full AttnRes 와 비교 가능한 안정적인 활성화 크기를 유지하여 "거대한 활성화" 현상을 효과적으로 제거했습니다.
• 특징 다양성: 각도 거리 분석 결과, MGR 스트림은 이질적인 의미적 경계를 발전시켜 PreNorm 에서 관찰된 표현의 동질화를 피했습니다.
• 중복성: 블록 가지치기 실험은 MGR 이 기능적 참여를 깊이 전반에 걸쳐 균일하게 분배하는 반면, PreNorm 은 초기 블록이 중요하고 깊은 블록은 중복된다는 것을 보여주었습니다.

5. 중요성 및 주장

이 논문은 MGR 이 복잡한 수작업 프로젝트 연산자에 의존하는 현대 아키텍처에서 철학적 전환을 나타낸다고 주장합니다. 대신 MGR 은 슬롯 어텐션 (Slot Attention) 원칙과 일치하여, 단순한 스코어링과 게이팅에 의해 업데이트되는 병렬 스트림 상태를 동등한 메모리로 취급합니다.

• 실용성: 저자들은 MGR 이 대규모 훈련 및 배포에 매우 실용적이라고 주장합니다. 이는 AttnRes 의 통신 오버헤드와 하이퍼-커넥션의 구현 복잡성을 피합니다.
• 효율성: 기존 아키텍처보다 성능 향상을 달성하면서 훈련 오버헤드는 미미하고 추론 오버헤드는 최소한으로 유지함으로써 엔지니어링 확장성을 위한 실질적인 개선을 제공합니다.
• 한계 및 향후 작업: 저자들은 겸손하게 MGR 이 Full AttnRes 를 능가하지만, 후자가 임의의 역사적 층에 접근할 수 있는 능력이 특정 특징 중요성 작업에서 여전히 이점을 제공할 수 있다고 지적합니다. 이 가설에 대한 엄격한 검증과 더 큰 규모에서의 이점 확인은 향후 작업에 맡깁니다.

요약하자면, MGR 은 단순하지만 효과적인 게이팅 메커니즘을 통해 심층 신경망의 정보 희석과 활성화 폭발이라는 이중 문제를 해결하는 견고하고 효율적이며 이론적으로 근거 있는 표준 잔차 연결의 대안을 제공합니다.