Muon이 에이전트형 강화학습에 유용한 경우

왜 중요한가

Muon은 대규모 프리트레이닝에서 FLOP 대비 효율을 개선한 최적화기이고, 본 연구는 그 적용이 강화학습 후학습 단계에서도 이득을 줄 수 있는지를 실험적으로 검증했다. 에이전트형, 긴 시간지평의 희소 보상 환경에서 Muon을 숨겨진 행렬에만 적용하자 특정 정책 추정기와 학습률 조합에서 실질적인 성공률 향상이 관찰됐다. 이 결과는 옵티마이저 선택이 강화학습에서 보상 신호 구조와 상호작용하며 성능을 좌우할 수 있음을 시사한다.

핵심 기여

GiGPO 환경에서의 단일-시드 통제 실험으로 Muon의 실효성 확인

GiGPO를 사용한 ALFWorld 실험에서 Muon을 숨겨진 행렬에만 적용하자 최종 25-업데이트 윈도우 평균 성공률이 0.29에서 0.55로 상승하면서 상대적 개선이 88%에 달했다. 모든 비-옵티마이저 설정은 매치되어 옵티마이저만 변경한 통제 비교가 이루어졌다. 결과는 Muon이 특정 advantage estimator와 학습률 조건에서 강화학습 성능을 높일 수 있음을 보여줬다.

학습률 대조 실험으로 단순한 고학습률이 Muon 효과를 재현하지 못함을 확인

AdamW의 학습률을 1e-5와 3e-5로 높이는 대조 실험에서 두 경우 모두 포스트-업데이트 검증 성공이 0으로 소실되었다. 동일한 Nominal rate에서 Muon은 200 업데이트를 완료하면서 유의미한 성공률을 유지했다. 이 관찰은 Muon의 이득이 단순한 스텝 크기 효과가 아니라 업데이트 규칙의 스펙트럼적 성격과 연관됨을 시사한다.

스텝-레벨 크레딧 절단(ω 토글) 실험에서 Muon과 단계별 신호의 상호보완성 규명

GiGPO의 ω=0/1 매치드 2×2 절차에서 Muon은 ω=0에서도 상당한 이득(+0.220)을 보였고 ω=1에서는 이득이 더 커졌다(+0.255). 스텝-레벨 신호가 옵티마이저와 결합하여 normalized AUC 격차를 확대하는 효과가 관찰됐다. 이 결과는 advantage estimator 설계가 옵티마이저 변형과 결합해 학습 동역학을 바꿀 수 있음을 보여준다.

핵심 아이디어 이해하기

강화학습의 정책경사 업데이트는 보상 신호의 구조와 샘플링 잡음에 민감하며, 긴 시간지평의 희소 보상 문제에서는 신뢰할 만한 신호가 일부 결정적 의사결정 스텝에 국한되는 경향이 있다. Muon은 모멘텀 행렬의 스펙트럼을 평탄화하여 특이값 크기 정보를 제거하고 방향성(UV^T)만 유지하기 때문에 약한 방향을 상대적으로 증폭시킬 가능성이 있다. 이 평탄화는 약한 방향의 부호(sign)가 신뢰할 수 있을 때 유효하며, 그렇지 않으면 잡음 우세 방향의 증폭으로 이어질 위험이 있다. GiGPO는 동일 앵커 상태에서 취한 행동들을 그룹화해 스텝-레벨 대비 신호를 생성하므로 국소적 신호 대비를 증가시켜 각 토큰 수준의 score gradient에 가중치 변화를 가져온다. 이 비교적 더 정제된 신호는 Muon이 약한 방향을 증폭할 때 유효 신호를 더 잘 보존하여 업데이트 정렬을 개선할 수 있는 조건을 제공했다. 반대로 단일턴 또는 신호대잡음비가 낮은 환경에서는 Muon의 스펙트럼 평탄화가 잡음을 증폭해 학습 실패를 초래할 수 있다는 이전 보고와 본 논문의 관찰이 일치한다. 결과적으로 Muon의 효과는 옵티마이저의 행렬적 처리 방식과 advantage estimator가 만들어내는 신호 특성의 상호작용으로 이해할 수 있다. 스펙트럼 평탄화가 실제 신호의 방향성을 보존하는 상황에서 Muon은 성능을 끌어올리고, 그렇지 않으면 안정성 문제를 일으킬 수 있다. 따라서 optimizer-side 수식적 변경과 estimator-side 신호 정교화가 함께 고려되어야 한다는 직관이 핵심이다.

방법론

실험은 ALFWorld의 여섯 작업 범주를 대상으로 Qwen2.5-0.5B-Instruct 에이전트를 훈련시키는 설정으로 이루어졌고, 모든 비교는 매치드 설정에서 옵티마이저만 변경하는 방식으로 설계되었다. 정책의 2차원(hidden) 행렬 파라미터에 대해서만 Muon을 적용하고 임베딩·노름·헤드 등 비행렬 파라미터는 AdamW로 유지하여 Muon 영향이 행렬 업데이트에 국한되도록 했다. 각 업데이트는 그룹화된 trajectory 배치를 사용했고 검증은 5업데이트마다 128 에피소드를 평가하여 원시 체크포인트 곡선을 기록했다. 옵티마이저별 학습률 범위를 따로 둔 이유는 Muon과 AdamW의 학습률 해석이 직접적으로 비교 가능하지 않기 때문이다. GiGPO와 GraphGPO에서는 Muon 학습률로 1e-5와 3e-5를, GRPO에서는 3e-5를 주로 사용했고 모든 실험에서 baseline AdamW는 1e-6를 사용했다. Muon은 각 2D 모멘텀 행렬에 대해 Newton–Schulz 반복(논문에서 NS5, momentum 0.95, Nesterov)을 적용해 UV^T 형태의 업데이트를 계산했고 비-행렬 파라미터는 기존 AdamW 규칙으로 갱신되었다. 성능 측정은 최종 윈도우(steps 175–200)의 평균 검증 성공률과 0–200 단계의 normalized validation AUC로 나누어 이루어졌다. 학습률 대조 실험은 AdamW의 nominal rate만 상향 조정했을 때 Muon과 유사한 이득이 재현되는지를 확인하기 위해 수행되었고, 추가적으로 ω 토글을 통해 스텝-레벨 신호의 유무가 결과에 미치는 영향을 동일한 코드 경로에서 비교했다. 구현상 Muon은 전체 2D 행렬을 NO_SHARD로 노출해 Newton–Schulz를 적용했고 이로 인한 메모리 요구량 증가는 별도로 보고되었다.

주요 결과

GiGPO에서 Muon(3×10^-5)을 숨겨진 행렬에 적용하자 최종 윈도우 평균 검증 성공률이 0.29에서 0.55로 상승하며 상대적 개선이 88%로 나타났다. 동일 설정에서 AdamW를 높은 학습률(1e-5, 3e-5)로 올린 대조군은 모든 포스트-업데이트 체크포인트에서 검증 성공이 0으로 소실되어 Muon의 이득이 단순한 스텝 크기 효과가 아님이 확인됐다. GraphGPO에서는 Muon이 normalized AUC를 크게 높였으나 늦은 윈도우에서는 포화 현상으로 차이가 감소했고 GRPO에서는 늦은 구간에서 소폭 향상이 관찰되었다. 스텝-레벨 크레딧(ω) 토글 실험에서는 ω=0일 때도 Muon이 큰 이득을 보였고 ω=1일 때 두 옵티마이저 모두 성능이 개선되어 step-level 신호와 Muon이 상호보완적으로 작용함이 관찰됐다. 학습 동역학 진단에서 Muon은 GiGPO 환경에서 더 큰 정책 이동(KL)을 유도했으나 GraphGPO에서는 KL과 gradient-norm 스케일이 대체로 유사한 범위를 유지했다. 결론적으로 Muon의 효과는 advantage estimator 구조와 학습률에 민감하게 의존하며, 일부 조합에서는 빠른 학습과 높은 최종 성능을 동시에 달성할 수 있음이 확인됐다.

기술 상세

Muon은 각 2D 숨겨진 가중치 행렬의 모멘텀 누적 행렬 M에 대해 Newton–Schulz 반복을 적용하여 근사적으로 극좌표 UV^T를 계산하고, 업데이트를 이 극좌표에 비례하도록 적용한다. 논문은 NS 반복 횟수 5회, momentum 0.95, Nesterov를 사용한 설정을 상세히 명시했고 비행렬 파라미터는 AdamW에 위임하는 하이브리드 구성을 사용했다. 이 방식은 singular-value 크기 정보를 제거함으로써 업데이트 스펙트럼을 평탄화하는 역할을 한다. GiGPO는 그룹 내 에피소드 정규화된 전체 반환 A^E와 앵커 상태별 스텝-레벨 A^S를 결합한 이점을 사용하며 전체 advantage는 A = A^E + ω A^S로 표현됐다. ω를 0으로 놓으면 GiGPO는 GRPO와 동일한 episode-level advantage를 재현하고 ω=1이면 스텝-레벨 신호가 활성화되어 동일 코드 경로 안에서 두 행동 신호 구조를 비교할 수 있도록 설계됐다. GraphGPO는 롤아웃을 상태-전이 그래프로 통합하여 successor-state distance에 기반한 전이 수준의 advantage를 산출함으로써 더 강한 베이스라인을 제공했다. 실험적 설정은 ALFWorld의 최대 호라이즌 50, 각 업데이트는 8-trajectory 그룹, 검증은 128 에피소드로 5업데이트마다 수행하는 고정 스케줄을 따랐다. normalized validation AUC는 0–200 단계의 원시 체크포인트 곡선에 대해 사다리꼴 적분을 취하고 200으로 나누어 계산했다. 구현상 Muon은 전체 2D 행렬을 NO_SHARD로 처리하여 NS 반복을 수행했으며, 이로 인한 per-device 메모리 증가는 분산 Muon 변형이 필요함을 제약으로 남겼다.

한계점

이 실험은 단일 시드(single-seed) 탐색으로 수행되어 개별 학습 궤적을 보고하며 모집단 수준의 불확실성은 반영되지 않았다. 모델 크기는 0.5B급 하나에 한정되었고 환경은 ALFWorld 하나뿐이어서 결과의 일반성은 추가 시드·모델·환경 검증이 필요하다. 또한 실험은 gradient SNR이나 업데이트 스펙트럼을 직접 측정하지 않았으므로 제안된 SNR 가설은 상관관계 기반의 가설로 남아 있다.

키워드

Muon(Muon 최적화)GiGPO(GiGPO 그룹 기반 정책 최적화)GraphGPO(GraphGPO 전이 그래프 기반 정책 최적화)Qwen2.5-0.5B-Instruct(Qwen2.5-0.5B-Instruct 모델)ALFWorld(ALFWorld 시뮬레이션 환경)

용어 해설

어드밴티지 추정기(Advantage estimator): — 정책경사 강화학습에서 각 행동의 가치 기여도를 계산하는 함수로, 에피소드 전체 보상이나 상태·행동별 보상 등을 이용해 신호를 생성한다. 이 논문 맥락에서는 episode-level, step-level, transition-level로 구분되어 각 방식이 업데이트 스펙트럼과 SNR에 다른 영향을 미친다. 어드밴티지 설계는 Muon 적용 시 학습 안정성과 이득 발생 여부를 결정하는 핵심 요소로 작용했다.
스텝 레벨 크레딧(Step-level credit): — 동일한 앵커 상태에서 취한 행동들을 그룹화하여 각 행동에 할당하는 시점별 보상 신호로, 에피소드 단위 보상보다 더 국지적인 신호를 제공한다. GiGPO의 핵심 구성요소로 작동하여 반복 관찰된 상태 간 대비를 만들어 gradient 스펙트럼의 약한 방향에서 신호를 개선할 가능성이 제기됐다. 본 논문 실험에서는 step-level credit의 유무가 Muon의 효과 크기에 영향을 미치는 것으로 확인됐다.
뉴턴-슐츠 반복(Newton–Schulz): — 행렬의 정규화 또는 역근사에 쓰이는 반복 알고리즘으로, Muon은 이 반복을 이용해 모멘텀 행렬의 극분해 근사에서 극좌표(P = UV^T)만 남긴다. Muon은 다수의 NS 반복(논문에서 5회)을 적용하여 singular-value 정보를 제거하고 방향성만 보존하는 업데이트를 생성했다. 이 과정은 업데이트의 스펙트럼을 평탄화하여 약한 방향을 상대적으로 증폭시키는 역할을 했다.
스펙트럼 평탄화(Spectral flattening): — 행렬 업데이트에서 특이값의 크기 차이를 줄여 주된 방향과 약한 방향 간 격차를 좁히는 연산을 의미한다. Muon은 singular-value를 제거하고 UV^T 형태로 업데이트를 적용하여 스펙트럼 평탄화를 수행했고, 이로 인해 잡음 우세 방향의 상대적 증폭이 발생할 수 있음이 논의되었다. 평탄화가 유효하려면 약한 방향의 부호가 신뢰할 만해야 한다는 수리적 조건이 제기됐다.
그래디언트 신호대잡음비(Gradient SNR): — 샘플링 잡음과 실제 정책 신호 간의 비율로, 값이 클수록 유효한 업데이트 방향이 잡음에 비해 뚜렷하다. 논문은 Muon의 효과가 이 비율에 의존할 수 있다고 가정했고, 단일턴 RLVR에서 실패 보고와 장기간 에이전트 환경에서의 성공 사례를 이 가설로 연결했다. 실험은 SNR을 직접 측정하지는 않았으므로 가설로 남아 있다.

Muon이 에이전트형 강화학습에 유용한 경우

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

키워드

용어 해설

관련 토론

댓글

관련 기사

Adam의 좌표별 적응성이 저랭크 복원을 깨뜨리는 이유

숙련된 것은 내려놓고 남은 것을 학습하기

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

키워드

관련 토론

댓글

관련 기사

Adam의 좌표별 적응성이 저랭크 복원을 깨뜨리는 이유

숙련된 것은 내려놓고 남은 것을 학습하기