MAI-Base-2: 작은 실험으로 거대 모델의 학습 레시피를 설계하는 방법

배경과 문제

MAI-Base-2는 Microsoft AI가 직접 사전학습한 두 번째 대형 기반 모델이다. 기반 모델은 지시 따르기나 추론 학습을 거치기 전, 다음 토큰 예측만으로 학습한 모델을 말한다. MAI-Base-2는 토큰마다 일부 전문가 네트워크만 활성화하는 전문가 혼합(Mixture of Experts, MoE) 구조로, 활성 파라미터 80B, 전체 파라미터 2.8T 규모다. 이전 세대인 MAI-Base-1은 활성 35B, 전체 1T였고, MAI는 MAI-Base-1이 크기 대비 강하지만 프런티어의 오픈 웨이트 및 비공개 모델보다 작았다고 판단했다. 따라서 다음 과제는 단순히 모델을 키우는 것이 아니라, 같은 성능을 더 적은 연산으로 내는 학습 레시피를 찾는 것이었다.

대규모 MoE 모델을 학습할 때는 여러 문제가 겹친다. 첫째, 작은 규모의 실험 결과가 큰 규모에서 그대로 유지되지 않는다. 작은 실험에서는 데이터를 거의 반복하지 않으므로 규제가 약한 아키텍처와 하이퍼파라미터가 좋아 보이다가, 큰 규모에서는 과적합으로 이어질 수 있다. 둘째, MoE 층은 전문가마다 받는 토큰 수가 달라 GPU 메모리 부족이나 메모리 단편화가 발생하기 쉽다. 셋째, Muon 옵티마이저의 직교화 계산은 모멘텀 행렬 전체를 필요로 하므로 ZeRO의 데이터 병렬 분할과 충돌한다. 넷째, 이상치 그래디언트가 배치 전체의 학습 신호를 약화시켜 학습을 불안정하게 만든다. MAI-Base-2는 이러한 문제를 스케일링 과학, 모델-인프라 공동 설계, 데이터 확장이라는 세 축으로 해결하려 했다.

기존 접근의 한계

MAI-Base-1은 밀집 FFN과 MoE 층을 번갈아 배치했다. MAI의 과거 실험에서는 모든 층에 MoE를 두는 구성이 FLOPs 기준 효율 이득(EG_FLOPs)에서 꾸준히 앞섰지만, 당시 MoE 구현의 비효율 때문에 실제 경과 시간 기준으로는 번갈아 두는 구성이 최적이었다. 구체적으로 동적 라우팅이 메모리 단편화를 늘렸고, 드물게 생기는 극단적인 전문가 부하 불균형에 대비해 메모리를 많이 잡아 두어 활성값 재계산을 줄일 기회를 놓쳤다. 또한 전문가를 더 많고 더 잘게 나누자 전문가마다 생기는 오버헤드가 커졌다.

하이퍼파라미터 탐색도 기존 방식으로는 비효율적이었다. 학습률(LR)과 가중치 감쇠(WD)는 서로 영향을 주고받는데, 모델 성능은 두 값의 곱 LR×WD에 가장 민감하다. 따라서 좋은 설정은 (log LR, log WD) 평면의 좁은 대각선 띠에 모인다. 하지만 직사각형 격자로 탐색하면 이 띠를 비효율적으로 훑게 된다. 또한 작은 모델에서 찾은 하이퍼파라미터를 큰 모델에 그대로 적용하면 데이터 반복 횟수의 차이 때문에 최적값이 어긋날 수 있다. MAI-Base-1의 학습 손실 곡선에는 전반부에 손실이 튀는 지점이 여러 번 보였는데, 이는 학습 레시피가 충분히 안정적이지 않았음을 시사한다.

연구방법을 단계별로

먼저 학습 레시피를 고정하고 크기만 바꾼 모델 묶음인 스케일링 사다리를 만들었다. 모델의 층 수에 맞춰 다른 구조를 정하고, 활성 파라미터당 학습 토큰 수를 일정하게 유지한다. 다만 작은 실험에서 전체 데이터를 조금씩 읽기만 하면 대규모 학습의 반복 학습 조건이 사라진다. 소스 스케일링은 작은 모델에 주는 고유 데이터의 양도 줄여, 규모가 달라도 각 데이터셋을 반복하는 횟수와 파라미터 대비 고유 데이터 비율을 맞춘다.

학습률 LR과 가중치 감쇠 WD는 따로 고르지 않고 함께 탐색했다. 좋은 설정이 LR과 WD의 곱에 민감하므로, LR과 LR×WD를 축으로 하는 탐색 격자를 만들었다. 여러 크기와 학습 길이에서 얻은 손실 곡면에 함수를 맞춰 최적점을 찾고, 그 최적점이 학습 스텝 수와 파라미터당 토큰 수에 따라 어떻게 바뀌는지 경험식을 구했다. 작은 실험에서 추정한 식을 중간 규모에서 확인한 뒤 큰 모델에 적용했으며, 모든 모델에 통하는 보편 공식으로 검증한 것은 아니다.

MoE에서는 특정 전문가에게 일이 몰리면 정해진 메모리에 토큰이 다 들어가지 않을 수 있다. 다회차 Dropless MoE는 전문가를 그룹으로 묶어 일정한 작업 공간을 주고, 넘친 토큰을 버리지 않고 추가 회차에서 처리한다. 같은 그룹의 다른 전문가가 쓰지 않은 용량도 활용한다. 평소에는 한 회차로 끝나도록 용량을 정하고, 작업을 보내기·계산하기·돌려받기를 겹쳐 통신 대기를 줄인다. 이 전문가 처리 과정에서 라우팅 불균형 때문에 필요한 메모리가 갑자기 커지는 문제를 제어하는 방식이다.

Muon은 학습 방향을 담은 모멘텀 행렬을 직교화해 업데이트를 만든다. 반면 ZeRO는 그 행렬을 여러 GPU에 나눠 저장하므로 그대로 계산하기 어렵다. MAI는 업데이트 때만 모양이 같은 행렬들의 조각을 묶어 교환하고, 각 GPU가 맡은 전체 행렬을 계산한 뒤 결과를 원래 위치로 돌려보냈다. 통신과 계산을 겹치고 조각의 위치 정보도 재사용해 비용을 줄였다. 이는 비임베딩 2D 매개변수에 적용한 구현이며, 모델의 모든 가중치를 Muon으로 바꾼 것은 아니다.

활성값 그래디언트 클리핑은 특정 내부 신호의 지나치게 큰 변화량을 제한하는 방법이다. 기존 방식은 배치 전체의 가중치 그래디언트 크기를 제한하므로, 일부 이상치 때문에 정상적인 예제의 학습 신호까지 함께 작아질 수 있다. MAI는 기존 클리핑을 유지하면서 네트워크의 정해진 위치에서 활성값의 그래디언트도 제한했다. Muon과 함께 썼을 때 그래디언트 크기가 더 좁은 범위에 머물렀고, 어떤 데이터가 이런 제한에 자주 걸리는지도 조사했다.

결과와 숫자의 의미

MAI-Base-2는 자체 보류 평가 4종(코드, QA, STEM, 수학)에서 비교한 모든 모델보다 낮은 BPB를 기록했다. BPB는 텍스트 1바이트를 예측하는 데 드는 평균 비트 수로, 토큰 단위 손실을 원문 바이트 수로 나눠 정규화하므로 토크나이저가 다른 모델끼리도 비교할 수 있다. 값이 낮을수록 모델이 그 텍스트를 잘 예측한다는 뜻이다. MAI-Base-2의 BPB는 코드 0.2203, QA 0.4162, STEM 0.3653, 수학 0.3561이었다. 2위는 모든 항목에서 사후 학습 모델인 Kimi K3였고, 수학에서는 그 차이가 0.0028로 작았다. 수학에서는 DeepSeek-V4-Pro-Base가 사후 학습 모델인 MiMo와 Qwen보다 BPB가 높았다는 점도 함께 보아야 한다.

4.2~12.3배 효율은 직접 측정한 값이 아니라 스케일링 곡선에서 추정한 값이다. MAI는 평가 4종마다 MAI 스케일링 사다리 모델들의 학습 FLOPs 대비 BPB에 스케일링 법칙을 따로 맞추고, DeepSeek-V4-Pro-Base의 BPB에 도달하는 지점의 연산량을 읽었다. 이미 결과를 아는 MAI-Base-2도 피팅에 넣었기 때문에, 도달 지점이 외삽이 아니라 피팅 범위 안에 놓인다. 이 추정치를 DeepSeek-V4-Pro-Base의 실제 학습 연산량과 비교하면 MAI 레시피가 4.2~12.3배 FLOPs 효율적이라는 결과가 나온다. 평가별 배수는 코드 12.3배, QA 4.2배, STEM 8.3배, 수학 6.6배로, 평가에 따라 3배 가까이 차이 난다. 가장 큰 값은 MAI 자체 코드베이스로 만든 코드 평가에서 나왔고, 가장 작은 값은 QA에서 나왔다. 같은 방식으로 MAI-Base-1과 비교하면 6.3~12.7배 효율이 좋아졌다.

학습 전체의 평균 MFU는 21%, 학습이 실제 진행된 시간의 비율인 굿풋은 89%였다. MFU의 분모는 GB200의 BF16 이론 최대치 2.5 PFLOPS이며 재계산과 옵티마이저 연산은 제외했다. 별도로 MoE 층 MFU가 24%에서 35%로 오른 것은 특정 층의 개선이지 전체 평균과 같은 수치가 아니다. Muon은 평균 반복 시간의 2.1%, 약 165ms를 차지했다. 이론적으로 추정한 55ms와도 차이가 남아 있어 통신·불균형·메모리 이동의 추가 개선 여지가 있다.

구체적인 예시

원문에 든 소스 스케일링 예시는 다음과 같다. 50T 토큰 전체 학습에서 어떤 자료를 8번 반복한다면, 5T 토큰 실험에서는 그 자료의 1/10만 골라 여전히 8번 반복한다. 큰 학습을 단순히 빨리 끝낸 작은 실험보다, 반복에 따른 과적합을 비슷한 조건에서 확인하려는 것이다. 그림 2는 작은 사다리에서 맞춘 곡선을 큰 모델로 외삽한 예측과 최종 측정이 가까웠음을 보여준다. 이는 앞서 설명한 DeepSeek와의 효율 비교에서 최종 모델까지 넣어 보간한 절차와 별개다.

데이터 진단 사례에서는 코드 자료의 라이선스 문장이 중간에 잘린 채 다음 내용으로 이어졌다. MAI는 활성값 그래디언트 클리핑에 걸린 이 예제를 역추적해 전처리 파이프라인의 버그를 찾고 수정했다고 보고한다. 단순히 불안정한 수치를 눌러 놓는 데 그치지 않고, 그 수치를 만든 입력도 검사한 것이다. 다만 한 사례가 데이터 전체의 오류가 제거됐음을 뜻하지 않으며, 그래디언트가 크다는 이유만으로 모든 입력을 잘못된 자료라고 단정해서도 안 된다.

한계와 오해하기 쉬운 점

현재 공개된 자료만으로 동일한 조건의 독립 재현을 하기는 어렵다. 평가는 모두 비공개 보류 세트다. 코드 평가는 MAI 자체 코드베이스이고, STEM, 수학, QA는 외부 업체에 의뢰해 만든 비공개 세트다. MAI는 모든 평가 세트를 자사 학습 데이터 기준으로 다시 오염 제거했다고 밝혔지만, 같은 세트로 다른 모델을 다시 재 볼 수는 없다. 비교 대상 중 Kimi K3, Qwen3.8-2.4T-A95B, MiMo-V2.6-Pro-RL은 사후 학습한 체크포인트만 공개되어 있어, 사후 학습이 BPB에 영향을 줄 수 있다는 단서를 달고 결과를 보고했다. 비교한 외부 공개 모델 중 사후 학습 전 기반 체크포인트는 DeepSeek-V4-Pro-Base뿐이다. 자체 이전 모델 MAI-Base-1도 별도로 비교했다.

4~12배 효율은 피팅에 의존하는 추정치이며, 추론 속도나 전체 비용, 모든 벤치마크 정답률 향상 배수가 아니다. 이 배수는 상대 모델의 학습 연산량을 같은 BPB에 도달할 MAI 레시피의 추정 연산량과 비교한 것이다. 원문은 상대 모델의 학습 연산량을 정확한 숫자로 적지 않았다. 또한 MAI-Base-2의 최종점도 피팅에 포함하므로 이 효율 비교는 보간이다. 실제 MAI-Base-2는 상대보다 더 많은 총 학습 FLOPs를 썼다. 별도의 그림 2 작은 사다리로 큰 모델 손실을 예측한 외삽과 구분해야 한다. 다운스트림 벤치마크 점수와 가중치는 공개되지 않았으며, BPB는 다음 토큰 예측 품질의 지표일 뿐 사후 학습 뒤의 실제 작업 성능은 아직 발표되지 않았다.

실제로 어떻게 활용할까

이 글에서 가져갈 것은 수치보다 방법이다. 다른 규모의 학습에 옮겨 볼 수 있는 방법으로는 소스 스케일링으로 반복 데이터를 사다리 실험에 반영하기, (log LR, log(LR×WD)) 격자로 학습률과 가중치 감쇠를 함께 탐색하기, 메모리를 고정한 다회차 Dropless MoE, ZeRO 환경에서 Muon 모멘텀을 재분할하기, 활성값 그래디언트 클리핑을 학습 데이터 디버깅에 쓰기가 있다. 특히 소스 스케일링은 작은 실험에서 데이터 반복 횟수를 맞추는 간단한 아이디어지만, 큰 모델로의 전이를 크게 개선할 수 있다. 하이퍼파라미터 탐색에서도 LR과 WD의 곱에 주목하면 탐색 공간을 크게 줄일 수 있다.

다만 적용할 때는 각 기법의 조건을 주의해야 한다. Muon은 비임베딩 2D 매개변수에 적용하며 모든 가중치를 Muon으로 바꾼 것이 아니다. 활성값 그래디언트 클리핑은 기존 전체 가중치 그래디언트 클리핑에 추가한 완화이며, 일부 이상치로 전체 배치 신호가 작아지는 문제를 줄인다. 각 개선의 독립적인 인과 기여를 전체 4~12배로 단정해서는 안 된다. 데이터는 다른 모델의 증류가 없는 organic data라는 연구팀 주장이며, 고유 학습 토큰 약 36% 증가와 비전 4.5T 고유 토큰 및 반복 포함 50T 혼합의 18%는 서로 다른 분모다. 커뮤니티의 FLOPs·LR·WD 계산값은 필자 파생 추정이며 공식 실측 하이퍼파라미터로 쓰지 말아야 한다.