TRACE: 롤아웃 유도형 양자화 인식 학습으로 구현한 MoE 언어 모델의 FP4 강화학습

논문이 다루는 문제

강화학습(RL) 사후 학습은 대규모 언어 모델의 추론, 코딩, 장기 과제 능력을 높이는 표준 방법이 되었습니다. 하지만 비용이 큽니다. 매 단계마다 현재 정책으로 많은 롤아웃(샘플 생성)을 먼저 수행하고, 그다음에 모델을 갱신합니다. 저자들은 롤아웃 생성이 상당한 연산과 메모리 부담을 만든다고 밝힙니다. 그래서 저정밀 롤아웃이 주목받습니다. FP4는 현재 하드웨어가 지원하는 가장 공격적인 부동소수점 형식 중 하나로, 가중치, 활성화, KV 캐시까지 4비트로 줄일 수 있습니다.

그런데 기존 FP4 RL 방법에는 핵심 약점이 있습니다. 논문에 따르면 이 방법들은 학습 경로와 롤아웃 경로 각각의 양자화 정확도를 따로 최적화할 뿐, 두 양자화 실행 경로 사이의 차이를 직접 줄이지 않습니다. RL은 이 차이에 매우 민감합니다. 롤아웃 쪽은 FP4 수치 동작으로 궤적을 만들고, 학습 쪽은 다른 수치 동작으로 그 궤적의 기울기를 계산합니다. 이 불일치는 오프폴리시 편향처럼 작용해 학습이 불안정해지고 최종 성능이 떨어질 수 있습니다.

TRACE의 핵심 아이디어

TRACE는 Train-Rollout Quantization Alignment via Compact GuidancE의 약자이며, 전문가 혼합(MoE) 언어 모델을 대상으로 합니다. 두 가지 핵심 요소가 있습니다.

첫째는 롤아웃 유도형 양자화 인식 학습(QAT)입니다. 일반 QAT에서는 학습 쪽이 각 값을 FP4 격자에 어떻게 반올림할지 스스로 정합니다. TRACE에서는 롤아웃 쪽의 양자화 결과가 학습 쪽의 FP4 반올림 결정을 안내합니다. 쉽게 말해 학습기가 롤아웃 엔진이 실제로 수행한 반올림을 재현하려고 합니다. 이렇게 하면 학습과 롤아웃 사이의 차이가 직접 줄어듭니다. 이전 방법과의 가장 큰 차이가 여기에 있습니다. 목표가 "각 경로가 따로 정확하게 양자화한다"에서 "두 경로가 서로 일치한다"로 바뀌었습니다.

둘째는 효율적인 양자화 정보 캐싱 방식입니다. 학습 쪽을 안내하려면 롤아웃 쪽의 양자화 정보를 저장하고 전달해야 하므로 저장 공간과 통신 비용이 늘어납니다. TRACE는 더 깊은 층의 가수(mantissa)와 스케일 정보만 선택적으로 보관해 이 추가 비용을 줄입니다. 어떤 층을 보관하는지, 추가 메모리가 얼마인지는 초록에 나오지 않으므로 본문 확인이 필요합니다.

작동 원리의 공학적 이해

FP4는 보통 블록 단위 스케일링을 씁니다. 소수의 값이 하나의 스케일 계수를 공유하고, 각 값은 4비트만 사용합니다. 반올림 방향(위쪽 또는 아래쪽 격자점)은 값마다 이산적인 선택입니다. 두 경로가 같은 가중치나 활성화를 다르게 반올림하면 출력에 작지만 체계적인 오차가 생깁니다. 여러 MoE 층을 지나고 라우터의 이산적인 전문가 선택까지 더해지면 오차가 커질 수 있습니다.

TRACE는 이렇게 이해할 수 있습니다. 롤아웃 엔진이 양자화를 수행하고 결과를 기록합니다. 학습 쪽은 순전파에서 그 기록을 참고해 반올림을 정하고, 두 계산 그래프를 수치적으로 최대한 맞춥니다. 깊은 층의 오차는 최종 출력에 더 직접 영향을 주므로, 깊은 층 정보만 캐시하는 것은 정확도와 비용 사이의 절충입니다. 이는 초록에 근거한 해석이며 정확한 알고리즘은 논문 본문이 기준입니다.

실험 결과와 성능

저자들은 추론, 코딩, 장기 RL 과제에서 네 개의 대규모 MoE 언어 모델로 평가했습니다. 초록이 전하는 주요 결과는 세 가지입니다. 첫째, TRACE는 FP4 가중치와 활성화, 그리고 FP4 KV 캐시를 함께 쓰는 롤아웃을 가능하게 하면서 RL 성능은 BF16 롤아웃과 비슷했습니다. 저정밀로 인해 최종 품질을 잃지 않았다는 뜻입니다. 둘째, 롤아웃 속도는 최대 5.4배 빨라졌습니다. "최대"라는 말에 주의하세요. 최상의 경우 수치이며 실제 효과는 모델, 시퀀스 길이, 배치 크기, 하드웨어에 따라 달라집니다. 평균값으로 받아들이면 안 됩니다. 셋째, BF16으로 학습한 정책에 사후 FP4 양자화를 적용하는 방식과 비교해 TRACE의 최종 FP4 성능이 더 강했습니다. 이 비교는 중요합니다. 학습 중에 FP4에 적응시키는 편이 학습 후에 압축하는 것보다 낫다는 점을 보여 줍니다.

초록에는 구체적인 벤치마크 점수, 메모리 절감 비율, 모델 이름이 없습니다. 여기서는 추측하지 않습니다.

개발자와 기업에 주는 영향

RL 사후 학습을 하는 팀에서는 특히 긴 사고 사슬과 장기 과제에서 롤아웃이 실제 소요 시간의 큰 부분을 차지합니다. 품질 저하 없이 FP4 롤아웃으로 몇 배의 속도를 얻는다면, 같은 GPU 예산으로 더 많은 반복을 돌리거나 실험 기간을 줄일 수 있습니다. MoE 모델은 가중치가 크고 KV 캐시 부담도 높아 FP4의 이점이 더 클 것입니다.

생태계 측면에서는 FP4를 지원하는 하드웨어와 추론 커널이 전제입니다. 도입 장벽도 있습니다. 추론 엔진이 양자화 결과를 내보내야 하고, 학습 프레임워크가 이를 받아 써야 하며, 둘 사이의 데이터 경로를 구현해야 합니다. 또한 이 연구는 RL 시스템에서 학습과 추론의 수치 일관성을 부수 효과가 아니라 일급 목표로 다뤄야 한다는 일반적인 설계 원칙을 뒷받침합니다.

한계와 향후 방향

주의할 점이 몇 가지 있습니다. 5.4배는 상한 수치입니다. 실험은 네 개의 MoE 모델이 대상이므로 밀집 모델이나 더 작은 규모로 확장되는지는 검증되지 않았습니다. 롤아웃 유도는 두 쪽의 결합을 강화해 학습기가 롤아웃 출력에 의존하게 되고, 캐시를 압축하더라도 동기화, 저장, 통신 비용이 남습니다. 또한 이는 막 공개된 프리프린트로 동료 심사를 거치지 않았으므로 독립적인 재현을 기다려야 합니다.

앞으로는 정렬 아이디어를 다른 수치 형식으로 넓히는 것, 비동기 RL 시스템과의 결합, 더 긴 시퀀스에서 캐시의 확장성 검증, 안전성과 강건성이 그대로 유지되는지 확인하는 것이 과제입니다.

요약

TRACE는 FP4 RL의 진짜 병목을 겨냥합니다. 문제는 한 경로의 양자화가 거칠다는 것이 아니라 학습 경로와 롤아웃 경로가 맞지 않는다는 점입니다. 롤아웃 유도형 QAT와 간결한 양자화 정보 캐시로, 초록은 BF16 롤아웃에 견줄 RL 성능과 최대 5.4배 롤아웃 속도 향상을 보고합니다. 도입 여부는 하드웨어, 프레임워크, 모델 규모에 달려 있으므로 소규모 과제에서 먼저 재현해 보는 것이 현실적입니다.