기술 요약: 디커플된 DiLoCo 를 위한 코사인 게이트된 Adam-Decay (CGAD)

1. 문제 제기

비동기 디커플된 DiLoCo(DDiLoCo) 시스템에서 작업자들은 전역 파라미터 서버와 의사-기울기를 교환하기 전에 여러 번의 로컬 최적화 단계를 수행합니다. 이 아키텍처의 핵심적인 과제는 **노후화 (staleness)**입니다. 외부 라운드 $t$에서 적용되는 의사-기울기는 네트워크 지연, 지체 작업자, 또는 비동기성 예산으로 인해 $\tau$ 라운드 전에 계산되었을 수 있습니다.

DiLoCo 의 표준 외부 옵티마이저는 네스테로프 모멘텀을 사용한 SGD 입니다. 이 방법은 의사-기울기의 나이 ($\tau$) 에 관계없이 모든 들어오는 의사-기울기에 균일한 업데이트 단계를 적용합니다. $\tau$가 증가함에 따라 모멘텀 버퍼는 손실 지형의 노후화된 뷰를 축적하게 됩니다. 결과적으로 옵티마이저는 더 이상 현재 손실을 감소시키지 않는 방향으로 과도하게 이동하는 단계를 취하게 되어, 특히 모델 규모가 커질수록 불안정성이나 발산이 발생합니다. 기존 비동기 모멘텀-SGD 에 대한 이론적 분석에 따르면, 수렴 속도는 최대 지연 $\tau_{max}$의 제곱에 비례하는 인자로 인해 저하됩니다.

2. 방법론: 코사인 게이트된 Adam-Decay (CGAD)

저자들은 의사-기울기의 나이에 명시적으로 조건을 부여한 업데이트를 수행하는 네스테로프 외부 옵티마이저를 대체할 수 있는 **코사인 게이트된 Adam-Decay (CGAD)**를 제안합니다.

핵심 메커니즘

CGAD 는 표준 Adam 업데이트를 수정하여, Adam 의 1 차 및 2 차 모멘트 버퍼에 들어가기 전에 조각 $f$에 대한 각 들어오는 의사-기울기 $g_f$를 인자 $\sigma(\tau)$로 스케일링합니다:
$$\sigma(\tau) = \gamma(\tau) e^{-\alpha \tau}$$

여기서:

• 지수 감쇠 ($e^{-\alpha \tau}$): 노후화된 기울기의 정보 감쇠를 모델링합니다. $\tau$가 커질수록 모멘텀 버퍼에 대한 기울기의 기여도는 지수적으로 감소합니다.
• 코사인 게이트 ($\gamma(\tau)$): 다음과 같이 정의된 부드러운 컷오프 함수입니다:
  $$\gamma(\tau) = \frac{1}{2} \left(1 + \cos\left(\frac{\pi \tau}{\tau_{cut}}\right)\right) \cdot \mathbb{I}[\tau \le \tau_{cut}]$$
  이 게이트는 신선한 기울기에 대해서는 1 이며, 선택된 컷오프 $\tau_{cut}$에서 0 으로 부드럽게 감쇠하며, $\tau > \tau_{cut}$인 경우 기여도를 엄격히 0 으로 만듭니다.

주요 설계 선택 사항

1. 버퍼 전 게이팅: 스케일링 $\sigma(\tau)$는 1 차 모멘트 ($m$) 및 2 차 모멘트 ($v$) 버퍼를 업데이트하기 전에 기울기에 적용됩니다. 이는 노후화된 기울기가 2 차 모멘트 ($v$) 를 팽창시켜 이후의 신선한 업데이트를 약화시키는 것을 방지합니다.
2. 조각별 변형 (PA-CGAD): 부분 동기화 스케줄러 (예: Streaming DiLoCo) 의 경우, 이 방법은 개별 조각의 나이 ($a_f$) 를 추적하고 $\max(\tau, a_f)$를 기반으로 업데이트를 게이팅하여 네트워크 지연과 스케줄링 지연에 대한 통합된 노후화 모델을 가능하게 합니다.
3. 하이퍼파라미터: 이 방법은 두 가지 하이퍼파라미터인 $\alpha$(감쇠율) 와 $\tau_{cut}$(컷오프) 을 도입하며, 모델 규모 간에 전이되는 기본값 ($\alpha=0.2, \tau_{cut}=32$) 을 제공합니다.

3. 이론적 기여

이 논문은 매끄러운 비볼록 목적 함수에 대한 이상화된 게이트된 적응형 업데이트에 대한 비점근적 수렴 경계를 제공합니다.

• 노후화 편향 독립성: 수렴 경계의 노후화 편향 항은 실현된 최대 지연 $\tau_{max}$가 아닌 게이팅 하이퍼파라미터 $\alpha$(구체적으로 $O(1/e^\alpha)$) 에 의존합니다.
• 표준 분석과의 대비: 비동기 모멘텀-SGD 에 대한 표준 분석은 $\tau_{max}^2$에 비례하는 복잡도 항을 산출합니다. 게이팅 메커니즘은 이를 $\alpha$에 의존하는 항으로 대체함으로써, $\tau_{cut} \ge 1/\alpha$가 보장되는 한 수렴 페널티를 최악의 경우 지연 분포에서 이론적으로 분리합니다.
• 범위: 저자들은 이것이 경험적 Adam 에 대한 완전한 수렴 정리가 아니라 게이팅 메커니즘에 대한 동기 부여임을 지적합니다. 이는 경험적 통계에서 유도된 것이 아니라 유계 정규화 단계를 가정하기 때문입니다.

4. 실험 결과

실험은 통제된 지연을 최대 $\tau=16$까지 사용하여 25 M, 1 B, 7 B 파라미터의 세 가지 규모에서 Llama 스타일 언어 모델 사전 학습을 수행했습니다.

안정성 및 수렴

• 1 B 규모: $\tau=16$에서 표준 네스테로프 레시피는 발산했습니다 (손실 $\approx 319$). 반면 CGAD 는 7.44 의 손실을 달성하여 $\tau=0$ 성능 (7.48) 과 비교 가능한 결과를 보였습니다. 네스테로프 방법은 $\tau \ge 8$에서 완전히 수렴을 잃었습니다.
• 25 M 규모: CGAD 와 다른 노후화 인식 베이스라인 (Adam-Decay, SDM, Delayed Nesterov) 은 모두 안정적으로 유지된 반면, 네스테로프는 $\tau \ge 4$에서 발산했습니다. 이 규모에서는 노후화 인식 방법들 간의 차이가 미미했습니다.

규모 보험 및 견고성

가장 중요한 발견은 메모리 제약 조건 (bf16 가중치, 8 비트 AdamW) 하의 7 B 규모에서 나타납니다:

• 시드 간 변동성: 가장 근접한 베이스라인인 Adam-Decay(코사인 컷오프가 없는 CGAD) 는 모델이 25 M 에서 7 B 로 확장됨에 따라 시드 간 표준 편차 ($\sigma$) 가 27 배 증가했습니다 ($\tau=8$에서 0.07 에서 1.89 로). 반면 CGAD 의 $\sigma$는 평평하게 유지되었습니다 (약 0.38).
• 단일 시도 위험: "평균 + $\sigma$"를 단일 시도 배포 위험의 대리 지표로 사용할 때, 7 B / $\tau=8$에서의 Adam-Decay 의 위험 (13.19) 은 우연 수준 손실 (10.83) 을 2.4 nats 초과하여, 단일 생산 실행에서 실패할 높은 확률을 시사했습니다. CGAD 의 위험 (11.28) 은 이 임계값보다 훨씬 낮게 유지되었습니다.
• 실패 메커니즘: 저자들은 Adam-Decay 의 7 B 에서의 실패를 노후화된 기울기에서 비롯된 2 차 모멘트 버퍼 ($v$) 내의 제곱 반올림 노이즈의 축적 때문이라고 귀결합니다. 이러한 기울기는 결코 완전히 0 으로 만들어지지 않기 때문입니다. CGAD 의 코사인 컷오프는 $\tau_{cut}$을 초과하는 기여도를 엄격히 0 으로 만들어 이를 방지합니다.

튜닝 민감도

• 네스테로프 vs CGAD: 네스테로프가 재튜닝되었을 때 (모멘텀을 완전히 제거), 높은 지연에서도 튜닝된 CGAD 구성의 성능을 따라잡지 못했습니다. 튜닝된 네스테로프는 다른 지연 영역마다 다른 하이퍼파라미터가 필요했던 반면, CGAD 는 전체 범위에서 견고성을 유지했습니다.
• 하이퍼파라미터 민감도: 최적의 $\alpha$는 규모에 따라 약간 이동합니다 (25 M 의 경우 0.2, $\tau=8$인 1 B 의 경우 0.4). 그러나 기본값 $\alpha=0.2$는 전체 지연 스펙트럼에 걸쳐 견고하게 유지됩니다.

5. 중요성 및 주장

이 논문은 CGAD가 대규모 제어 지연 하에서 네스테로프 모멘텀의 취약성을 해결하는 비동기 DiLoCo 학습을 위한 가장 견고한 확장 선택이라고 주장합니다.

• 규모 보험: 주요 기여는 코사인 컷오프로, 이는 "규모 보험" 역할을 합니다. 지수 감쇠만으로는 (Adam-Decay) 작은 모델이나 정밀도 설정에서는 충분할 수 있지만, 하드 컷오프는 압축된 옵티마이저 상태 (8 비트/bf16) 로 작동하는 프런티어 규모 모델 (7 B 이상) 에 필수적입니다. 이는 치명적인 시드 간 변동성과 단일 시도 실패로 이어지는 노후화된 분산의 축적을 방지합니다.
• 즉시 호환성: CGAD 는 기존 Adam 구현에 한 줄의 수정만 필요하며, 규모 간에 전이되는 기본값을 가진 두 가지 하이퍼파라미터만 추가합니다.
• 실질적 영향: 프런티어 규모 비동기 학습의 경우, 이 방법은 성공적으로 출시되는 단일 학습 실행과 실패하는 것 (우연 손실 이상으로 머무름) 사이의 차이를 제공합니다.

저자들은 이론적 범위에 대해 겸손하게, 수렴 경계가 이상화된 업데이트에 적용되며 경험적 우월성은 DiLoCo 계열과 Llama 스타일 사전 학습 컨텍스트 내에서 구체적으로 입증되었음을 인정합니다. 또한 실험에서 사용된 내부 단계 수 ($H=8$) 는 생산 DiLoCo 설정 ($H=100-500$) 보다 낮지만, 민감도 연구는 더 높은 $H$에서도 순위가 유지됨을 시사합니다.