Unlocking Feature Learning in Gated Delta Networks at Scale

기술 요약: Gated Delta Network의 대규모 기능 학습(Feature Learning) 구현

문제 정의
대규모 언어 모델(LLM)의 스케일링은 두 가지 주요 병목 현상에 직면해 있습니다: 더 큰 모델을 훈련하는 데 드는 막대한 계산 비용과 표준 트랜스포머(Transformer) 셀프 어텐션의 이차 복잡도(quadratic complexity)입니다. 선형 순환 구조(Linear recurrent architectures)인 Gated Delta Network(GDN)는 효율적인 대안으로 등장하였으며, 장거리 의존성을 포착하면서도 선형 시간 추론을 유지합니다. 그러나 이러한 모델을 효과적으로 스케일링하는 데에는 중요한 격차가 존재합니다. Maximal Update Parametrization(µP)가 표준 트랜스포머 및 일부 선형 모델에서 제로샷 하이퍼파라미터 전이(특히 학습률)를 성공적으로 가능하게 했음에도 불구하고, GDN에 대한 µP의 적용은 아직 탐구되지 않았습니다. 핵심 과제는 이 아키텍처의 순환 상태 역학(recurrent state dynamics)이 기존 µP 유도 방식에서 가정하는 스칼라 순환이나 표준 피드포워드 구조가 아닌, 데이터 의존적 게이팅 메커즘과 풀 매트릭스 델타 규칙(full-matrix delta rules)을 통해 업데이트된다는 점에 있습니다. Mamba와 같은 구조적 상태 공간 모델(SSM)에 대한 선행 연구는 바닐라 µP가 대각 SSM(diagonal SSMs)에서 실패함을 보여주었으나, GDN의 풀 매트릭스 상태 업데이트는 이러한 특정 교정법을 적용할 수 없게 만듭니다.

방법론
저자들은 GDN을 위한 µP 공식을 엄밀하게 유도함으로써 이 격차를 해결합니다. 방법론은 전체 순전파(forward pass) 과정, 즉 고유한 게이팅 메커니즘과 순환 상태 전이를 포함하여 좌표 크기 추정치(coordinate-size estimates)를 전파하여 필요한 초기화, 순전파 승수(forward multipliers) 및 학습률을 결정하는 과정을 포함합니다.

주요 분석 단계는 다음과 같습니다:

1. 순전파 분석: 저자들은 투영된 특징(projected features), 잠재 상태 $S_t$, 그리고 리드아웃(readout)의 좌표 크기를 유도합니다. 이들은 쿼리(query)와 키(key)의 L2 정규화가 랭크-원 쓰기 업데이트(rank-one write updates)와 결합되어 특정 분산 스케일링을 초래함을 식별합니다. 결정적으로, 저자들은 RMSNorm 레이어 앞에 $\sqrt{d}$ 승수를 삽입하여 정규화 레이어의 입력이 $\Theta(1)$의 좌표 크기를 유지하도록 함으로써, 기능 학습에 필요한 그래디언트 스케일링이 방해받지 않도록 제안합니다.
2. 게이팅 메커니즘 분석: 논문은 비선형 게이팅 파라미터 $\alpha_t$ (감쇠)와 $\beta_t$ (쓰기 강도)를 분석합니다. 이들은 일반적인 가중치와 달리, 훈련 가능한 가중치 벡터($W_\alpha, W_\beta$)와 스칼라 파라미터($a_{log}, b$)를 포함하는 스칼라 프리-액티베이션(pre-activations)으로부터 유도됩니다. 분석 결과, 이러한 게이팅 구성 요소들은 상태 업데이트 방정식에서의 특정한 역할 때문에 표준 µP 스케일링 법칙을 따르지 않음을 밝혀냈습니다.
3. 역전파 및 학습률 유도:
   • AdamW: AdamW는 그래디언트를 두 번째 모멘트로 정규화하기 때문에, 유효 업데이트 크기가 가중치 클래스 간에 일관되게 유지되어 학습률 요구 사항을 단순화합니다.
   • SGD: 순수 SGD 환경에서는 원시 그래디언트 크기가 가중치 클래스마다 크게 다릅니다. 저자들은 게이팅 가중치 행렬($W_\alpha, W_\beta$)이 $\Theta(1/\sqrt{d})$의 학습률 스케일링을 필요로 하고, 스칼라 게이팅 파라미터($a_{log}, b$)는 $\Theta(\sqrt{d})$의 스케일링을 필요로 한다는 것을 유도했습니다. 이러한 스케일링은 학습률이 일반적으로 균일하거나 단순한 $1/n$ 규칙을 따르는 표준 µP 설정에서 벗어납니다.

주요 기여

1. 이론적 유도: 본 논문은 Gated Delta Network를 위한 최초의 완전한 µP 공식을 제공합니다. 저자들은 전체 순전파를 통해 좌표 크기 추정치를 유도하고, 모든 가중치 클래스에 대한 원칙적인 초기화 분산과 순전파 승수를 확립했습니다.
2. 새로운 스케일링 규칙 발견: 저자들은 SGD 하에서 게이팅 구성 요소가 비표준 학습률 스케일링을 필요로 한다는 것을 발견했습니다:
   • 게이팅 가중치 행렬($W_\alpha, W_\beta$)은 $\Theta(1/\sqrt{d})$ 스케일링이 필요합니다.
   • 스칼라 게이팅 파라미터($a_{log}, b$)는 $\Theta(\sqrt{d})$ 스케일링이 필요합니다.
3. 경험적 검증: 저자들은 다양한 너비(숨겨진 차원 256에서 1563까지)에 걸쳐 AdamW와 SGD 옵티마이저를 사용하여 GDN 언어 모델을 사전 훈련했습니다.
   • 결과: 제안된 µP 구성 하에서, 최적의 학습률은 두 옵티마이저 모두에서 모든 모델 너비에 대해 제로샷으로 전이되었습니다. 반면, 표준 파라미터화(Standard Parametrization, SP)와 기존 µP 구성은 전이에 실패했으며, 모델 너비가 증가함에 따라 최적 학습률이 크게 변화했습니다.
   • 안정성: 제안된 구성은 안정적인 학습률 전이를 가능하게 하여 이론적 유도의 타당성을 입증했습니다.

의의 및 주장
본 논문은 자신의 분석이 Gated Delta Network를 규모에 맞게 파라미터화하여 기능 학습을 수행하는 미해결 문제를 해결한다고 주장합니다. 제안된 µP 공식이 제로샷 하이퍼파라미터 전이를 허용함을 보여줌으로써, 이 연구는 이론적 유도의 정확성과 실용적 유용성을 검증합니다. 저자들은 이것이 Gated Delta Network에 특화된 µP 일관적 파라미터화를 유도하고 검증한 최초의 연구임을 강조하며, 이는 대각 SSM에 관한 기존 연구와 구별됩니다. 연구 결과는 적절한 파라미터화가 선형 순환 모델의 스케일링에 필수적이며, GDN의 특정 게이팅 메커니즘은 표준 트랜스포머나 다른 선형 아키텍처와는 다른 고유한 스케일링 규칙을 필요로 한다는 점을 시사합니다. 저자들은 이러한 유도가 향후 다른 선형 또는 하이브리드 아키텍처의 스케일링 법칙에 대한 연구를 안내하기를 기대합니다.