Muon-OGD: Muon-based Spectral Orthogonal Gradient Projection for LLM Continual Learning

기술 요약: Muon-OGD

문제 정의
본 연구에서 다루는 핵심 과제는 **연속 학습(continual learning) 중 대규모 언어 모델(LLM) 의 파국적 망각(catastrophic forgetting)**입니다. LLM 이 새로운 작업에 대해 순차적으로 미세 조정될 때, 이전에 학습한 작업에 대한 성능이 급격히 저하되는 경우가 많습니다. 기존 투영 기반 연속 학습 방법들 (예: OGD, GEM, O-LoRA) 은 과거 작업과 관련된 방향에 직교하는 부분 공간으로 매개변수 업데이트를 제한함으로써 이를 완화합니다. 그러나 이러한 방법들은 일반적으로 유클리드 매개변수 기하학(euclidean parameter geometry) 하에 공식화되어 있으며, 여기서 업데이트 크기와 직교성은 **프로베니우스 노름(Frobenius norm)**에 의해 지배됩니다.

저자들은 근본적인 불일치를 지적합니다. 최근 제안된 Muon 옵티마이저와 같은 행렬 값 매개변수를 위한 현대적 옵티마이저들은 스펙트럼 노름 (연산자 노름) 기하학(spectral-norm (operator-norm) geometry) 하에서 작동합니다. Muon 은 뉴턴 - 슈르츠 (Newton–Schulz) 반복을 통해 근사된 직교화된 행렬 업데이트를 적용하며, 스펙트럼 노름 제약 하에서 가장 가파른 하강 (steepest descent) 을 수행하는 것으로 해석될 수 있습니다. 본 논문은 행렬 값 LLM 매개변수에 대한 연속 학습 제약 조건으로 프로베니우스 노름 기하학에 의존하는 것은 최적이지 않다고 주장합니다. 이는 근본적인 최적화 역학의 기하학과 일치하지 않기 때문입니다.

방법론: Muon-OGD
본 논문은 스펙트럼 노름 인식 최적화와 투영 기반 연속 학습 제약 조건을 통합하는 Muon-OGD(Muon 기반 스펙트럼 직교 기울기 투영) 프레임워크를 제안합니다.

1. 공식화: 표준 프로베니우스 노름 제약 문제를 대신하여, Muon-OGD 는 업데이트 $\Delta$를 스펙트럼 노름 제약 최적화 문제로 공식화합니다:
   $$\min_{\Delta \in \mathbb{R}^{m \times n}} \langle G, \Delta \rangle \quad \text{s.t.} \quad \|\Delta\|_2 \leq \eta, \quad \langle C_i, \Delta \rangle = 0, \forall i=1,\dots,k$$
   여기서 $G$는 기울기 (또는 모멘텀) 행렬, $\|\cdot\|_2$는 스펙트럼 노름, $\eta$는 단계 크기이며, $\{C_i\}$는 업데이트와 직교해야 하는 제약 방향 (예: 과거 작업의 기울기) 입니다.

2. 이중 최적화(Dual Optimization): 이 문제는 폐쇄형 해가 없으므로, 저자들은 이중 형식을 유도합니다. 이중 목적 함수는 이동된 행렬 $H = G + \sum \lambda_i C_i$의 핵 노름 (nuclear norm) 을 최소화하는 것을 포함합니다. 최적 업데이트는 행렬 부호 함수를 통해 복원됩니다: $\Delta^* = -\eta \cdot \text{msgn}(H)$.

3. 알고리즘: 이 방법은 이중 변수 $\lambda$에 대한 서브그래디언트 하강을 사용하여 이중 문제를 반복적으로 풉니다.

   • 내부 루프: 직교성 제약을 강제하기 위해 $\lambda$를 반복적으로 업데이트합니다.
   • 행렬 부호 근사: 행렬 부호 함수 $\text{msgn}(H)$는 명시적 SVD 의 계산 비용을 피하면서 뉴턴 - 슈르츠 반복(구체적으로 5 회 반복, $NS_5$로 표기) 을 사용하여 효율적으로 근사됩니다.
   • 저랭크 인스턴시레이션: 효율성을 위해 저자들은 과거 작업의 상위 $k$개 특이 벡터에서 유래된 제약을 사용하는 특정 인스턴시레이션을 제안합니다. 이를 통해 이중 변수를 저랭크 행렬 $\Lambda$로 조직화하여 이중 문제의 차원을 축소합니다.

주요 기여

1. 연속 학습을 위한 스펙트럼 노름 기하학: 본 논문은 현대적 옵티마이저 (Muon 등) 의 연산자 노름 기하학과 연속 학습의 안정성 제약을 정렬하는 스펙트럼 노름 기하학 하에서 투영 기반 연속 학습을 공식화합니다.
2. 실용적 업데이트 규칙: 과거 작업 직교성 제약과 Muon 스타일 행렬 부호 업데이트를 결합한 실용적 알고리즘 (Muon-OGD) 을 유도하며, 이중 형식과 뉴턴 - 슈르츠 근사를 통해 제약 문제를 효율적으로 해결합니다.
3. 실증적 검증: 이 프레임워크는 LLM 연속 미세 조정에 인스턴시레이션되어 다양한 벤치마크에서 평가되었으며, 프로베니우스 노름 기반 대응 방법들보다 일관된 개선을 입증했습니다.

실험 결과
저자들은 Muon-OGD 를 다음과 같이 평가했습니다:

• 표준 벤치마크: T5-Large 를 사용한 5 작업 및 15 작업 연속 학습 벤치마크.
• TRACE 벤치마크: 코딩, 수학, 다국어 이해 등 다양한 기술을 다루는 지시 미세 조정 LLM 을 위한 벤치마크 (LLaMA-2-7B-Chat 사용).
• 도메인별 커리큘럼: Qwen2.5 및 LLaMA3.2 모델을 사용한 코딩 (BigCodeBench) $\to$ 수학 (GSM8K) $\to$ 의료 (HuatuoGPT-o1) 에 대한 순차 학습.

결과:

• 성능: Muon-OGD 는 평균 정확도 (AA) 및 역전달 (Backward Transfer, BT) 측면에서 순차 미세 조정 (SeqFT) 및 경쟁적인 직교 기울기 기준 방법들 (예: O-LoRA, Sculpting Subspaces, OSFT) 보다 일관되게 우수한 성능을 발휘했습니다.
  • T5-Large 15 작업 벤치마크에서 Muon-OGD 는 **72.0%**의 평균 정확도를 달성하여 O-LoRA(69.6%) 와 OSFT(71.3%) 를 능가했습니다.
  • TRACE 벤치마크에서 Muon-OGD 는 비교된 방법들 중 가장 높은 AA(49.4%) 와 가장 낮은 망각률 (-4.7%) 을 달성했습니다.
• 확장성: 이 방법은 뉴턴 - 슈르츠 반복의 효율성과 저랭크 제약 표현을 활용하여 계산적으로 확장 가능한 상태를 유지합니다.
• 아키텍처 무관성: 이 접근법은 인코더 - 디코더 (T5) 와 디코더 전용 (LLaMA, Qwen) 아키텍처 모두에서 효과적입니다.

의의 및 주장
본 논문은 스펙트럼 노름 인식 업데이트 기하학이 LLM 의 연속 학습을 위한 프로베니우스 노름 투영에 대한 실용적이고 효과적인 대안이라고 주장합니다. 현대적 옵티마이저 (Muon 등) 의 기하학과 연속 학습의 안정성 요구 사항 사이의 간극을 메움으로써, Muon-OGD 는 **안정성 - 가소성 트레이드오프(stability-plasticity trade-off)**를 개선합니다. 저자들은 기존 투영 기반 방법들을 체계적으로 업그레이드하여 스펙트럼 노름 기하학 하에서 작동하게 함으로써 파국적 망각을 방지하면서도 Muon 에서 관찰된 조건화 및 안정성 이점을 더 잘 보존할 수 있다고 제안합니다. 본 연구는 모든 연속 학습 과제를 해결한다고 주장하지는 않지만, 다양한 도메인과 모델 규모에서 유지 및 적응 측면에서 실증적 이득을 제공하는 원칙적인 기하학적 대안을 제시합니다.