행과 열을 함께 정규화하는 사전학습 옵티마이저 MeqMuon 공개

왜 지금 이 주제인가

모델 크기가 커질수록 사전학습에 드는 연산과 메모리 비용도 함께 불어납니다. 최근에는 기존 표준 옵티마이저인 AdamW 대신, 가중치 행렬을 직교화해 업데이트하는 Muon 계열 옵티마이저가 사전학습 효율을 높이는 방법으로 주목받아 왔습니다. 9월 28일 arXiv에 공개된 논문 “MeqMuon: Matrix-Equilibrating Muon for LLM Pretraining”은 이 Muon 계열의 최신 개선판을 제시하면서, 메모리 사용량과 최종 성능을 함께 검증한 수치를 내놓았습니다. 옵티마이저 설계는 눈에 잘 띄지 않지만 모든 모델 학습의 밑바닥을 이루는 요소라, 인프라 비용에 민감한 팀이라면 챙겨 볼 만합니다.

핵심 내용

Muon 옵티마이저는 파라미터 업데이트 행렬을 직교화해 학습을 안정시키는 방식으로, 최근 발표된 NorMuon 등의 변형은 행(row) 단위로 업데이트 크기를 정규화해 성능을 더 끌어올렸습니다. 논문은 이 행 단위 정규화만으로는 업데이트 행렬에 나타나는 다양한 불균형 패턴을 모두 다루지 못한다는 점을 문제로 짚습니다. 행렬에 따라 행 방향이 아니라 열(column) 방향으로 불균형이 나타나는 경우도 있는데, 기존 방식은 이를 반영하지 못한다는 것입니다.

MeqMuon이 제시하는 해법은 두 가지입니다. 첫째, 직교화된 업데이트 행렬에 대해 행 단위 정규화와 열 단위 정규화 중 어느 쪽이 맞는지를 자동으로 선택합니다. 둘째, 직교화 대상이 아닌 나머지 2차원 파라미터(임베딩, 정규화 계층 등)의 모멘텀 행렬에는 행과 열 양쪽을 동시에 정규화하는 양방향(two-sided) 정규화를 적용합니다. 이 과정에서 AdamW가 각 파라미터마다 유지하는 2차 모멘트 추정값을 아예 저장하지 않아도 되도록 설계했습니다. 두 방식을 결합해 행렬마다 다른 불균형 패턴에 맞춰 정규화 방향을 자동으로 조정하는 것이 핵심입니다.

연구팀은 Llama 60M·130M·350M, SmolLM2 135M·360M, Qwen2 0.5B까지 여섯 개 모델 규모에서 실험했습니다. 각 모델은 파라미터 수의 20배에 해당하는 토큰으로 학습했는데, 예를 들어 350M 모델은 7.4B 토큰을 사용했습니다. 검증 퍼플렉시티(PPL, 낮을수록 좋음) 결과를 보면 Llama-350M 기준 AdamW 17.02, Muon 16.04, NorMuon 15.98, MeqMuon 15.94로 MeqMuon이 가장 낮았습니다. SmolLM2-360M에서도 AdamW 18.46, Muon 17.28, NorMuon 17.20, MeqMuon 17.16, Qwen2-0.5B에서는 AdamW 19.65, Muon 18.79, NorMuon 18.76, MeqMuon 18.63으로 같은 순서가 반복됐습니다. 모델 규모가 커져도 순위가 뒤바뀌지 않았다는 점이 눈에 띕니다.

메모리 측면에서는 Qwen2-0.5B 기준으로 옵티마이저 상태(optimizer state)가 Muon 2,404.17MiB, NorMuon 2,405.33MiB였던 반면 MeqMuon은 1,884.59MiB로, 21.6% 줄었습니다. AdamW의 2차 모멘트 추정값을 저장하지 않는 설계가 그대로 메모리 절감으로 이어진 결과입니다. 실험은 PyTorch 2.6.0과 CUDA 12.4, Transformers 4.57.6 환경에서 DistributedDataParallel(DDP)을 이용해 진행됐습니다.

의미와 한계

이 결과는 옵티마이저 하나를 바꾸는 것만으로 최종 성능과 메모리 효율을 동시에 개선할 수 있다는 것을 보여줍니다. 특히 GPU 메모리가 학습 규모를 제한하는 현실에서, 옵티마이저 상태를 21.6% 줄인다는 것은 같은 하드웨어로 더 큰 배치나 더 큰 모델을 다룰 여지를 준다는 뜻입니다. PPL 개선폭 자체는 크지 않지만, 여러 모델 규모에서 일관되게 나타났다는 점은 우연이 아닐 가능성을 높입니다.

다만 실험 규모는 최대 0.5B 파라미터, 7.4B 토큰 수준에 머물러 있습니다. 최근 공개되는 상용 모델들이 수백억에서 수조 파라미터, 수조 토큰 단위로 학습되는 것과 비교하면 아직 소규모 검증 단계입니다. 이 결과가 그 규모에서도 그대로 유지되는지는 별도로 확인돼야 합니다. 또한 별도의 공개 코드 저장소가 아직 확인되지 않아, 다른 연구팀이 재현하거나 실제 파이프라인에 적용하기까지는 시간이 걸릴 것으로 보입니다.

한 줄 정리

MeqMuon은 행과 열 정규화를 함께 적용해 옵티마이저 메모리를 21.6% 줄이면서도 검증 손실을 낮췄지만, 검증 규모는 아직 0.5B 파라미터에 머물러 있습니다.