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

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

강화학습의 정책경사 업데이트는 보상 신호의 구조와 샘플링 잡음에 민감하며, 긴 시간지평의 희소 보상 문제에서는 신뢰할 만한 신호가 일부 결정적 의사결정 스텝에 국한되는 경향이 있다. 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 가설은 상관관계 기반의 가설로 남아 있다.