저정밀 강화학습의 수치 오차 해결, TRIAGE 방법 제안
원제 TRIAGE: Direction-Aware Mismatch Stabilization of Native NVFP4 Reinforcement Learning
추천 29댓글 2
Key Point
저정밀 연산으로 LLM 강화학습 속도를 크게 높일 수 있지만 수치 오차로 인한 불안정성이 문제인데, 이 논문이 그 원인을 정확히 분석하고 실질적인 해결책을 제시했기 때문에 LLM 최적화에 관심 있는 개발자와 연구자가 봐야 한다.
핵심 요약
- NVFP4 저정밀 실행을 통해 대형언어모델의 강화학습을 가속할 수 있지만, 학습기와 샘플러 간 실행 불일치가 정책 최적화를 불안정하게 만드는 문제가 있다.
- 연구팀은 불일치가 정책 기울기 방향과 상호작용하는 방식을 분석해, 크기만으로는 식별할 수 없는 국소 증폭 및 감소 영역을 구분했다.
- NVFP4 실행에서 초기 불균형으로 인해 음수 이점·음수 간극 업데이트가 우세하고, 이러한 토큰들이 불일치가 전역적으로 확산되기 전에 응답 세그먼트의 소수 부분에 집중된다는 패턴을 발견했다.
- 이를 바탕으로 TRIAGE라는 방향 인식형 안정화 방법을 제안했으며, 세그먼트 수준의 진단으로 정책 기울기 업데이트를 선택적으로 균형잡고 남은 심각한 불일치에 한정된 보수를 적용한다.
- TRIAGE는 최적화 목표를 수정하면서도 샘플러와 학습기 모두에서 네이티브 NVFP4 4비트 가중치·활성화(W4A3) 순전파 실행을 유지한다.
- Qwen3-4B와 Qwen3-30B-A3B 모델 실험에서 평가 기간 전체에 걸쳐 안정적인 최적화를 보였고, 5개 수학 추론 벤치마크에서 전정밀도 성능과 동등한 결과를 달성했다.
- NVFP4와 TRIAGE를 함께 사용하면 BF16 대비 최대 2.3배 높은 롤아웃 처리량을 제공한다.