기술 요약: 데드-디렉션 컨디셔너 (Dead-Direction Conditioners, DDC)

1. 문제 정의

심층 신경망의 손실 함수는 파라미터 공간 내에서 연속적인 대칭성(게이지 대칭성, gauge symmetries)을 나타냅니다. 여기에는 로짓 시프트(logit shifts), ReLU/SwiGLU 리스케일링(rescalings), 레이어 정규화(LayerNorm) 스케일 불변성, 어텐션 헤드 회전 등이 포함됩니다. 이러한 대칭성의 궤도(orbits)를 따라 손실 지형(loss landscape)은 평탄하며(dead directions), 이는 피셔 정보 메트릭(Fisher information metric)이 0이 됨을 의미합니다.

Adam과 같은 표준 적응형 옵티마이저는 이러한 대칭성을 존중하지 못합니다. Adam의 좌표별 프리컨디셔너($1/\sqrt{\hat{v}}$)는 좌표 순열에 대해서만 등변(equivariant)합니다. 아키텍처 그룹이 좌표를 혼합할 때(예: 한 쌍의 가중치를 역으로 스케일링하거나 벡터 공간을 회전시킬 때), Adam의 업데이트는 게이지 궤도로 "누수(leakage)"됩니다. 이 누수는 두 가지 주요 문제를 야기합니다:

1. 드리프트(Drift): 옵티마이저 궤적이 고정된 섹션 $\bar{\Theta} = \Theta/G$에 머무는 대신, 평탄한 궤도를 따라 표류합니다. 이 드리프트는 옵티마이저를 대칭성이 존중되었을 때 도달했을 지점과는 다른 최소값으로 이동시킵니다.
2. 흐릿한 비율 판독성(Blurred Rate Readability): 이론적 결과(Shirodkar, 2026)에 따르면, 옵티마이저가 몫 공간(quotient) 위에서 깨끗하게 하강할 때만 단일 학습률(singular learning rate, 가장 작은 피셔 고유값의 감쇠)을 궤적으로부터 읽어낼 수 있습니다. Adam의 누수는 이러한 노이즈와 편향을 유발하여, 이 비율을 흐리게 만들고 최소값의 기하학적 구조를 판독 불가능하게 만듭니다.

2. 방법론: 데드-디렉션 컨디셔너 (DDC)

본 논문은 베이스 옵티마이저(예: Adam 또는 Muon)를 $G$-등변(G-equivariant)한 것으로 격상시키는 프리컨디셔닝 프레임워크인 DDC를 소개합니다. 핵심 메커니즘은 $G$-불변 리만 메트릭 $\rho$를 사용하여 파라미터 공간과 그래디언트를 수직(vertical)(궤도 접선) 성분과 수평(horizontal)(몫 공간) 성분으로 분해하는 것입니다.

업데이트 규칙

그래디언트 $g_t$에 대해, DDC는 이를 $g_t = \Pi_V g_t + \Pi_H g_t$로 분리합니다. 또한 별도의 2차 모멘트 추정치를 유지합니다:

• 수직 성분 ($\hat{v}_V$): 수직 그래디언트는 궤도 차원당 단일 스칼라로 압축됩니다(예: CE 시프트를 위한 행 평균 사용). 이 스칼라는 수직 업데이트를 정규화하는 데 사용되며, 효과적으로 전체 궤도 차원을 하나의 단위로 취급합니다.
• 수평 성분 ($\hat{v}_H$): 수평 그래디언트는 $G$-등변 프레임 내에서 엄격하게 좌표별로 프리컨디셔닝됩니다. 결정적으로, 좌표별 나눗셈 이후 업데이트는 궤적이 수평 상태를 유지하도록 재투영($\Pi_H$)됩니다.

업데이트 단계는 다음과 같습니다:
$$\theta_{t+1} = \theta_t - \eta \left[ \Pi_H \left( \frac{\Pi_H \hat{m}_t}{\sqrt{\hat{v}_H} + \epsilon} \right) + \text{lift}_V \left( \frac{\text{vert-to-scalar}(\Pi_V \hat{m}_t)}{\sqrt{\hat{v}_V} + \epsilon} \right) \right]$$

처리되는 아키텍처 게이지

DDC는 네 가지 아키텍처 게이지를 위한 특정 투영기와 축소 연산자를 구축합니다:

1. Cross-Entropy Shift: 로짓(bias 또는 unembed 가중치)의 가산적 시프트.
2. ReLU/SwiGLU Rescaling: 인접한 가중치 쌍($W_\ell, W_{\ell+1} \to cW_\ell, c^{-1}W_{\ell+1}$)의 곱셈적 스케일링.
3. LayerNorm/RMSNorm Scale: 정규화 파라미터 및 다운스트림 가중치의 채널별 스케일링.
4. Attention Head Rotation: QK 및 VO 가중치의 비가환(non-abelian) $O(d_{head})$ 회전.

비가환 회전 게이지의 경우, 표준적인 좌표별 $\hat{v}$는 작동하지 않습니다. DDC는 네 가지 $O(d_{head})$-인식 конструк션(construction)을 제안하며, 배포된 기본값은 **body_frame_topk**입니다. 이는 Gram 행렬 $M = W_Q^T W_Q + W_K^T W_K$의 고유 기저(eigenbasis) 내에서 좌표별 적응성을 유지하되, "활성(live)" 고유 방향(고유값이 임계값 이상인 방향)에 대해서만 적용하고, 데드 디렉션은 모멘텀으로만 처리합니다.

베이스 옵티마이저 구성

• DDCAdam: Adam에 전체 DDC 구성을 적용합니다. 이는 네 가지 게이지를 모두 정확하게(또는 body-frame 회전에 대해 1차 근사적으로) 처리합니다.
• DDCMuon: Muon 옵티마이저(업데이트를 직교화함)에 DDC를 적용합니다. Muon의 직교화는 회전 게이지와 정확히 교환(commute)되므로, DDCMuon은 회전에 대해 정확합니다. 그러나 Muon의 직교화는 per-channel 스케일 자유도를 제거하므로, DDCMu건은 리스케일링 게이지(ReLU/LN scale)를 조건화할 수 없으며 해당 성분에 대해서는 Adam을 기본값으로 사용합니다.

3. 주요 기여

1. 등변 프리컨디셔닝(Equivariant Preconditioning): 적응형 옵티마이저를 $G$-등변한 것으로 격상시켜, 궤적이 대칭 몫 공간 $\bar{\Theta}$ 위의 프리컨디션된 그래디언트 흐름을 유지하도록 하는 일반적인 구성을 제공합니다.
2. 비율 판독성(Rate Readability): DDC가 "궤적-비율 판독성(trajectory-rate readability)" 속성을 복구한다는 이론적 증명(Theorem 1)을 제시합니다. 몫 공간 위에서 가장 작은 피셔 고유값은 $\bar{t}^{2(k-1)}$로 감쇠하며, 이를 통해 Adam에서는 가려졌던 단일 차수(singular order) $k$를 궤적으로부터 직접 읽을 수 있습니다.
3. 게이지 억제(Gauge Containment): DDC가 Adam에서 발생하는 게이지 모드의 드리프트(예: 로짓 시프트, 가중치 리스케일링)를 방지하여, 이들을 초기값 또는 고정점으로 유지함을 입증합니다.
4. 개선된 최적화 결과:
   • 과적합 저항성(Over-training Resistance): 모델이 적합(fit) 지점을 지나 훈련될 때, DDCAdam은 검증 손실이 훈련 손실에서 멀어지는 "과적합 붕괴(over-training collapse)" 현상을 방지합니다.
   • 그로킹 신뢰성(Grokking Reliability): 심층 트랜스포머(depth 24)에서 DDCMuon은 일반적인 Muon이 완전히 실패하는 것(0/11)과 달리, 11개 시드 중 10개에서 "그로킹"(갑작스러운 일반화)을 달성합니다.
   • 압축(Compression): Vision Transformer에서 DDCAdam은 동일한 weight decay 하에서 AdamW보다 더 공격적으로 여유 FFN 용량을 압축하여 더 낮은 검증 손실에 도달합니다.

4. 실험 결과

• 언어 모델 (Depth-12, 과적합 상황):
  • 비율 판독성: DDCAdam은 65개의 레이어 관측 가능 셀 중 32개에서 단일 비율을 읽어내는 반면, AdamW는 단 7개에서만 읽어냅니다. AdamW에서는 피셔 고유값이 최소값으로부터 상승하지만, DDCAdam에서는 바닥(floor)을 향해 붕괴합니다.
  • 과적합: DDCAdam은 AdamW의 5.88에 비해 0.67의 검증-훈련 손실 격차를 유지합니다.
• 그로킹 트랜스포머:
  • Depth-24: DDCMuion은 정확도 >0.9에서 11개 중 10개의 그로킹 시드를 달성하는 반면, 일반 Muon은 0/11을 기록합니다.
  • Depth-8 (RoPE): 회전 위치 임베딩(Rotary Position Embeddings, $SO(2)^{d/2}$)의 특정 토러스 대칭에 맞춰 게이지를 조정하면, 일반적인 $O(d_{head})$ 게이지(0.955)보다 높은 정확도(0.977)와 낮은 분산을 얻습니다.
• Vision Transformer (ViT):
  • ImageNet-100에서 처음부터 훈련할 때, DDCAdam은 AdamW가 용량을 그대로 남겨두는 동일한 weight decay 조건에서도 더 낮은 검증 손실(1.71 vs 2.12 nats)을 달성하고 네트워크를 더 효과적으로 압축합니다(활성화 단일값 4배 작음).
• 게이지 억제:
  • 합성 태스크에서 DDCAdam은 AdamW 대비 게이지 모드의 교차 시드 드리프트(cross-seed drift)를 ReLU 리스케일의 경우 14배, LayerNorm 스케일의 경우 65배 감소시키면서도 최종 태스크 손실은 동일하게 유지합니다.

5. 의의 및 주장

본 논문은 대칭성을 옵티마이저에 구축하는 것이 발견되는 솔루션의 기하학적 구조를 근본적으로 변화시킨다고 주장합니다. 게이지 대칭성을 존중함으로써:

• 최소값은 더 날카롭습니다(Sharper): 옵티마이저는 덜 퇴화된 최소값(단일 바닥 위에서 더 높은 피셔 고유값)에 안착합니다.
• 기하학은 측정 가능합니다: 궤적은 밑바닥의 단일 계수(singular coefficient)를 나타내는 읽을 수 있는 신호가 되며, 이를 통해 훈련 역학으로부터 모델 복잡도를 직접 진단할 수 있습니다.
• 신뢰성이 높아집니다: 옵티마이저는 깊은 모델이나 과적합된 환경에서 발생하는 "드리프트"로부터 견고해지며, 특히 그로킹(grokking)이 필요한 작업에서 유리합니다.

저자들은 이 구성이 프리컨디셔너에 범용적으로 적용 가능하며(Adam, Muon, Shampoo를 격상), 이 이점이 단순히 weight decay 때문이 아니라 몫 공간(quotient)에 머물도록 하는 기하학적 제약 때문임을 강조합니다. 결론적으로, 네트워크의 게이지 대칭성을 옵티마이저가 존중할 때, 이는 "찾아낸 최소값을 더 날카롭게 만들고, 그 최소값의 기하학적 구조를 궤적이 측정할 수 있는 대상으로 바꿉니다."

참고된 한계점:

• Grokking에서의 Muon 기반 속도 향상은 등변 업데이트에 내재된 번들된 weight decay의 영향도 일부 포함되어 있으나, 투영(projection) 자체는 신뢰성과 더 깨끗한 구조를 제공합니다.
• 고유 벡터 부호 모호성으로 인해 body-frame 회전 구성에 대한 정확한 등변성은 근사적(1차 근사)이지만, 이러한 드리프트는 매우 작습니다($20$스텝 동안 약$10^{-3}$).
• 정렬(alignment)을 통해 전체 단일 구조(차수 및 다중도)를 읽어내는 데는 본 연구에서 아직 완전히 개발되지 않은 추가적인 메커니즘이 필요합니다.