기술 요약: 그래디언트 스무딩(Gradient Smoothing): 최적화를 개선하기 위한 레이어별 업데이트 결합

문제 정의

트랜스포머(Transformer) 및 잔차 네트워크(Residual Networks)와 같이 반복적인 구조적 블록으로 구성된 현대의 심층 신경망은 학습 과정 중에 블록 간에 구조적 관계가 나타납니다. 트랜스포머 블록 결합(Transformer Block Coupling) 및 잔차 정렬(Residual Alignment)에 관한 최근 연구들은 블록 자코비안(Jacobian)과 잔차 표현(residual representations)의 특이 벡터(singular vectors)가 레이어 전반에 걸쳐 정렬된다는 것을 보여주며, 이는 네트워크 깊이 전체에 걸친 암묵적인 협응(implicit coordination)을 시사합니다. 그러나 표준 최적화 방법(예: SGD, Adam, Muon)은 일반적으로 블록 단위의 업데이트를 독립적으로 생성하고 적용하며, 이러한 층간(cross-depth) 구조적 관계를 명시적으로 활용하지 못합니다. 이는 최적화 업데이트 자체도 깊이에 따라 나타나는 구조를 활용하도록 증강될 수 있는지에 대한 의문을 제기합니다.

방법론

깊이 방향 그래디언트 증강 (Depth-wise Gradient Augmentation)

저자들은 깊이 방향 그래디언트 증강을 도입하며, 이는 각 레이어에 적용되는 업데이트를 기본 옵티마이저가 생성한 블록 단위 옵티마이저 업데이트들의 집합을 깊이 차원을 따라 변환함으로써 얻는 일반적인 최적화 패러다임입니다. 업데이트를 개별적으로 적용하는 대신, 이 프레임워크는 $\{u_1, \dots, u_L\}$ 형태의 업데이트 집합에 깊이 방향 증강 연산자를 적용합니다.

그래디언트 스무딩(Gradient Smoothing) 및 윈도우 스무딩(Window Smoothing)

이 프레임워크 내에서, 본 논문은 그래디언트 업데이트를 깊이 방향으로 매끄럽게 만드는 일련의 방법인 그래디언트 스무딩을 제안합니다. 주요 연구 대상인 윈도우 스무딩은 국소 평균 연산자(local averaging operator)입니다.

$L$개의 블록을 가진 네트워크에 대해, $u_\ell$를 레이어 $\ell$에 대해 기본 옵티마이저가 생성한 업데이트라고 합시다. 스무딩된 업데이트 $\tilde{u}_\ell$는 해당 레이어의 업데이트와 인접한 이웃 레이어들의 볼록 결합(convex combination)으로 계산됩니다:
$$\tilde{u}_\ell = \begin{cases} (1 - \frac{\alpha}{2}) u_1 + \frac{\alpha}{2} u_2 & \ell = 1 \\ (1 - \alpha) u_\ell + \frac{\alpha}{2} (u_{\ell-1} + u_{\ell+1}) & 2 \le \ell \le L-1 \\ (1 - \frac{\alpha}{2}) u_L + \frac{\alpha}{2} u_{L-1} & \ell = L \end{cases}$$
여기서 $\alpha \in [0, 1)$는 스무딩 강도를 조절합니다. 경계 조건은 연산자가 행 확률적(row-stochastic)임을 보장합니다.

이 방법은 **옵티마이저 불가지론적(optimizer-agnostic)**이며, 순수하게 기본 옵티마이저(예: Adam, AdamW, Muon)가 생성한 업데이트 벡터만을 대상으로 작동합니다. 이는 계산 오버헤드가 매우 적으며 건축적 수정이 필요하지 않습니다. 또한 저자들은 정규화 변형인 노름 보존 스무딩(Norm-Preserving Smoothing, 원래 업데이트 크기에 맞게 재조정)과 방향 스무딩(Directional Smoothing, 방향만을 스무딩한 후 크기 재조정)을 탐구합니다.

주요 기여

1. 프레임워크 도입: 본 논문은 옵티마이저 업데이트에 층간(cross-depth) 정보를 통합하기 위한 일반적인 패러다임으로서 깊이 방향 그래디언트 증강을 정의합니다.
2. 그래디언트 스무딩: 이 프레임워크를 그래디언트 스무딩(특히 윈도우 스무딩)으로 구체화하여, 단순하면서도 효과적이고 광범위하게 적용 가능한 방법임을 입증합니다.
3. 실증적 검증: 이 방법은 다음과 같은 다양한 환경에서 평가되었습니다:
   • LLM 사전 학습: GPT 스타일 모델(깊이 24 및 30)을 사용한 nanochat 레시피.
   • RL 사후 학습: 수학적 추론을 위해 GRPO를 사용하여 LLM(DeepSeek-R1-Distill-Qwen-1.5B) 미세 조정.
   • 이미지 분류: CIFAR-100 상의 Vision Transformer(ViT-B) 학습.
   • 확산 모델(Diffusion Modeling): 이미지 생성을 위한 U-ViT 백본 학습(CIFAR-10).
4. 이론적 분석: 저자들은 윈도우 스무딩이 업데이트 방향의 깊이 간 변동을 수축시켜, 표현 잔차(representation residuals)의 정렬에 대한 하한(lower bound)을 개선한다는 이론적 경계를 제공합니다.

실험 결과

연구는 훈련 하이퍼파라미터나 목적 함수를 수정하지 않고도 모든 테스트된 설정에서 일관된 최적화 및 일반화 성능 향상을 보고합니다:

• LLM 사전 학습: 그래디언트 스무딩은 검증 손실(validation loss)과 BPB(Bits Per Byte)의 수렴을 가속화하고 DCLM CORE 지표를 개선했습니다. 이 혜택은 더 깊은 모델(깊이 30 vs. 깊이 24)에서 더 두드러졌습니다.
• RL 추론: 수학적 벤치마크(AIME24, AIME25, AMC23, MATH-500)에서 윈도우 스무딩은 평균 pass@1 정확도를 베이스라인 55.45%에서 57.60%로 향상시켰습니다.
• 이미지 분류: CIFAR-100 상의 ViT-B는 $\alpha=0.2$일 때 75.62%의 테스트 정확도를 달성하여, 베이스라인인 74.56%를 상회했습니다.
• 확산 모델: 그래디언트 스무딩은 CIFAR-10에서 FID 점수를 10k 샘플 기준 6.58에서 5.82로, 50k 샘플 기준 4.01에서 3.74로 낮추었습니다.

표현 분석

토큰 궤적(token trajectories)에 대한 실증적 분석은 그래디언트 스무딩이 더 구조화된 표현 진화를 촉진함을 보여주었습니다:

• 정렬 증가: 연속적인 잔차 업데이트 간의 코사인 유사도가 증가하였으며, 이는 레이어 전반에 걸쳐 더 협응된 방향적 업데이트가 이루어짐을 나타냅니다.
• 선형성: 선 형상 점수(Line Shape Score, LSS)가 감소하였으며, 이는 토큰 궤적이 더 선형적(1차원 아핀 공간에 더 가까움)이 되었음을 시사합니다.
• 안정성: 이 방법은 훈련 과정 전반에 걸쳐 마이크로배치 그래디언트 분산을 감소시켜 훈련 안정성을 개선했습니다.

의의 및 주장

본 논문은 깊이 방향 그래디언트 증강이 최적화 과정에서 층간 구조를 활용하기 위한 유망한 프레임워크라고 주장합니다. 특히, 그래디언트 스무딩은 다음과 같은 특징을 가진 단순하고 광범위하게 적용 가능한 방법으로 제시됩니다:

• 건축적 변경 없이 최적화 궤적과 일반화 성능을 개선합니다.
• 구조화된 깊이 방향 프리컨디셔닝(preconditioning) 방법으로서의 해석과 일치하게, 더 구조화된 표현 진화를 촉진합니다.
• 스무딩이 업데이트 변동을 수축시켜 표현 정렬을 개선하는 방식에 대한 이론적 토대를 제공합니다.

저자들은 로컬 윈도우 평균에 집중하고 있지만, 이 프레임워크가 더 풍부한 형태의 깊이 방향 증강(예: 지수 스무딩, 적응형 스킴)을 허용한다는 점을 강조합니다. 결론적으로, 최적화 과정에서 층간 구조를 명시적으로 활용하는 것이 광범위한 현대 딥러닝 작업 전반에 걸쳐 이점을 가져다줄 수 있다고 밝히고 있습니다.