분포 인식 보상: LLM 회귀를 위한 예측 분포 기반 강화학습 (Distribution-Aware Reward: Reinforcement

대규모 언어 모델 (Large language models)은 텍스트, 코드, 분자 문자열과 같은 이질적인 입력으로부터 실수 값을 예측할 수 있지만, 대부분의 학습 목표 (training objectives)는 각각 디코딩된 부동 소수점 숫자를 독립적으로 평가합니다. 이는 점 추정치 (point estimates)를 개선할 수는 있으나, 보정된 예측 분포 (calibrated predictive distributions)를 보장하지는 못합니다. 이는 후보 순위 지정 (candidate ranking)이나 불확실성 추정 (uncertainty estimation)이 필요한 응용 분야를 제한합니다. 본 논문에서는 분포 인식 보상 (Distribution-Aware Reward)을 소개합니다. 이는 온폴리시 강화학습 (on-policy reinforcement learning) 목표로, 주요 기여점은 언어 모델이 개별 디코딩 출력만을 스칼라 타겟 (scalar targets)에 맞춰 최적화하는 대신, 회귀 (regression) 작업을 위해 더 나은 예측 분포를 생성하도록 학습시키는 것입니다. 우리의 방법은 여러 개의 디코딩된 샘플을 경험적 예측 분포 (empirical predictive distribution)로 취급하고, 이를 연속 순위 확률 점수 (Continuous Ranked Probability Score, CRPS)로 평가하며, 각 롤아웃 (rollout)이 분포 품질에 기여하는 한계 기여도 (marginal contribution)를 기반으로 leave-one-out 신용 (credit)을 할당하여, 정확하면서도 적절하게 분산된 예측에 보상을 부여합니다. 우리는 통제된 가우시안 혼합 (Gaussian-mixture) 작업, 코드 성능 예측, 그리고 SMILES 문자열로부터의 분자 특성 예측에서 우리의 방법을 평가합니다. 여러 작업에 걸쳐, 우리의 방법은 지도 미세 조정 (supervised fine-tuning) 및 포인트와이즈 강화학습 (pointwise reinforcement learning) 베이스라인보다 성능을 개선하였으며, KBSS에서 6포인트의 Spearman 개선을 포함하여 강력한 순위 상관관계 (rank-correlation) 이득을 보였습니다. MoleculeNet에서 우리의 방법은 SMILES 문자열만을 사용함에도 불구하고 강력한 그래프 기반 및 3D 분자 모델들과 경쟁할 만한 성능을 유지합니다. 추가 분석을 통해 우리의 방법이 롤아웃 다양성 붕괴 (rollout diversity collapse)를 완화하고 불확실성 진단 (uncertainty diagnostics)을 개선함을 보여주며, 이는 예측 분포를 직접 최적화하는 것이 언어 모델의 회귀를 더 견고하고 더 잘 보정되게 만든다는 것을 시사합니다.