LLM 강화학습의 학습·추론 불일치, CIS로 보정한다

들어가며

검증 가능한 보상을 활용하는 강화학습(RLVR)에서는 답변 생성과 모델 업데이트가 서로 다른 시스템에서 수행되는 경우가 많다. 추론 엔진이 롤아웃을 샘플링하고, 학습 엔진이 생성된 토큰을 바탕으로 그래디언트를 계산하는 방식이다. 두 엔진이 같은 모델 가중치를 사용하더라도 하드웨어, 수치 정밀도, 병렬화 방식, 연산 구현의 차이 때문에 동일한 토큰에 서로 다른 확률을 부여할 수 있다. 이러한 학습·추론 불일치는 정책 업데이트의 방향과 안정성에 영향을 줄 수 있다.

논문 「Rethinking Training-Inference Mismatch in LLM Reinforcement Learning」은 이 문제를 중요도 샘플링의 관점에서 분석하고, Calibrated Importance Sampling(CIS)을 제안한다. 단순히 모든 중요도 비율에 고정된 상한을 적용하는 대신, 확률 차이가 어떤 형태로 나타나는지 분석한 뒤 토큰의 신뢰도에 맞춰 보정 강도를 조절하는 것이 핵심이다.

핵심 내용

• 두 엔진의 역할 분리가 출발점이다. 롤아웃은 추론 엔진에서 나오지만 그래디언트는 학습 엔진에서 계산된다. 토큰 확률이 서로 다르면 중요도 비율이 지나치게 커지거나 작아져 정책 업데이트에 오차가 생길 수 있다.
• 로그 오즈의 변위로 불일치를 설명한다. 저자들의 실증 분석에 따르면 불일치는 로그 오즈에 더해지는 변위로 근사할 수 있다. 이 변위는 softmax 이전의 개별 logit 섭동에 의해 결정되며, 선택된 토큰의 신뢰도와 거의 무관한 분포를 보인다.
• 신뢰도에 맞는 절단 규칙을 사용한다. CIS는 변위 공간에서 큰 양의 변위를 하나의 일정한 임계값으로 제한한다. 이를 중요도 비율 공간으로 변환하면 신뢰도가 높은 토큰에는 더 엄격한 상한이 적용되고, 신뢰도가 낮은 토큰에는 더 많은 학습 신호가 남는다.
• 분산을 낮추는 대신 통제된 편향을 감수한다. 이론적으로 CIS는 정확한 중요도 샘플링에서 문제가 될 수 있는 무한대의 2차 모멘트 항을 상수로 제한되는 항으로 바꾼다. 대신 절단된 초과분에 따른 편향이 발생한다. 또한 작은 중요도 가중치를 위쪽으로 클리핑하면 홀드아웃 정확도가 낮아질 수 있다는 진단 결과도 제시된다.

실험 결과와 의미

CIS는 세 개의 혼합 전문가(MoE) 모델과 다섯 개의 수학 추론 벤치마크에서 평가됐다. 다섯 벤치마크의 평균 성능을 기준으로 CIS는 세 모델 모두에서 평가된 기준 방법 중 가장 높은 결과를 보였다. 추가 분석에서는 일반적인 절단 중요도 샘플링보다 낮은 신뢰도의 토큰에 가하는 절단 편향이 적은 것으로 나타났다.

이 연구가 주는 broader lesson은 LLM 강화학습에서 인프라의 차이가 알고리즘 문제로 이어질 수 있다는 점이다. 두 엔진 사이의 작은 확률 차이도 중요도 비율을 거치며 확대되고, 최종적으로 정책 업데이트의 크기와 방향을 바꿀 수 있다. 따라서 RLVR 시스템을 설계할 때는 학습 프레임워크와 추론 런타임뿐 아니라 수치 형식, 하드웨어, 병렬 계산 방식의 호환성도 함께 고려해야 한다.

CIS는 모든 토큰에 같은 수준의 보정을 적용하지 않고, 토큰 신뢰도에 따라 위험을 다르게 본다는 분석 틀을 제공한다. 다만 이번 결과는 논문에서 다룬 모델과 수학 과제에 기반한다. 더 다양한 작업과 엔진 조합에서 효과가 유지되는지는 추가 연구가 필요하다.

출처: Hugging Face Daily Papers