LLM 강화학습을 위한 예측 분산 마스크 (Predictive Divergence Masks for LLM RL)
요약
본 논문은 LLM 강화학습(RL)에서 사용되는 신뢰 영역 마스크의 한계를 지적하고, 이를 개선한 예측 분산 마스크(Predictive Divergence Mask)를 제안합니다. 이 새로운 마스크는 정책 기울기 단계가 신뢰 영역 내에서 분산을 증가시키거나 감소시킬지 예측하여 RL 학습을 안정화하는 것을 목표로 합니다.
핵심 포인트
- PPO 스타일 접근 방식의 비율 기반 방향성 기준은 한계가 있습니다.
- 예측 분산 마스크는 분산 변화를 예측하여 RL 신뢰 영역을 개선합니다.
- 이산 소프트맥스 정책에 대한 폐쇄형 공식과 경량 top-K 추정기를 개발했습니다.
- 제안된 분산 기반 방향성은 기존 샘플링 비율보다 실현된 분산 변화와 더 잘 정렬됩니다.
대규모 언어 모델(LLMs)의 강화학습은 일반적으로 오프-폴리시 업데이트를 안정화하기 위해 신뢰 영역 마스크(trust-region masks)에 의존합니다. 지배적인 PPO 스타일 접근 방식은 두 가지 기준을 위해 샘플링된 토큰 중요도 비율(sampled-token importance ratio)을 사용합니다: 정책이 행동 정책으로부터 너무 멀리 이동했는지 묻는 근접성 기준(proximity criterion), 그리고 업데이트가 그것을 더 멀리 밀어내는지 묻는 방향성 기준(direction criterion)입니다. 최근 연구인 DPPO는 PPO의 비율 기반 테스트를 행동 정책과 학습 정책 간의 확률 분산으로 대체하여 근접성 기준을 개선했습니다. 하지만, 그 방향성 기준은 여전히 PPO에서 상속받았습니다. 토큰은 샘플링된 토큰 중요도 비율이 한쪽에서 벗어날 때만 마스크될 수 있습니다. 우리는 이 비율 기반 방향성 기준이 근접성 기준을 정의하는 분산 변화와 부호가 다를 수 있는 단일 샘플 프록시(single-sample proxy)임을 관찰했습니다. 따라서, 우리는 다음 정책 기울기 단계(policy-gradient step)가 신뢰 영역에서 사용되는 동일한 분산을 증가시키거나 감소시킬지 묻는 예측 분산 마스크(predictive divergence mask)를 제안합니다. LLM RL에 사용되는 이산 소프트맥스 정책(discrete softmax policies)에 대해, 우리는 이를 폐쇄형 공식으로 도출합니다. 프로덕션 롤아웃 엔진이 어휘의 잘린(top-K) 보기만 노출하기 때문에, 우리는 이 예측을 위한 두 가지 경량 top-$K$ 추정기(estimators)를 개발했습니다. 상세 분석에 따르면 분산 기반 방향성이 샘플링된 비율보다 실현된 분산 변화와 더 잘 정렬되며, 그 결과로 생성된 마스크는 모델 규모 및 정밀도 설정 전반의 RL 학습을 개선합니다.
AI 자동 생성 콘텐츠
본 콘텐츠는 arXiv cs.LG의 원문을 AI가 자동으로 요약·번역·분석한 것입니다. 원 저작권은 원저작자에게 있으며, 정확한 내용은 반드시 원문을 확인해 주세요.
원문 바로가기