기술 요약: LLM 회귀를 위한 분산 인식 보상

문제 정의

대형 언어 모델 (LLM) 은 이질적인 입력 (텍스트, 코드, 분자 문자열) 으로부터 실수 값을 예측하는 능력을 입증해 왔습니다. 그러나 LLM 회귀를 위한 표준 학습 목표인 지도 미세 조정 (SFT) 이나 포인트별 보상 (예: 음의 평균 제곱 오차) 을 사용하는 강화 학습 (RL) 은 각 디코딩된 부동 소수점 수를 독립적인 예측으로 취급합니다.

이 접근 방식은 개별 점 추정치를 최적화하지만, 샘플링된 예측 집합이 보정된 예측 분산을 형성하도록 보장하지는 않습니다. 결과적으로 모델은 개별적으로는 목표값에 가깝지만 집단적으로는 보정이 잘 되지 않거나, 지나치게 좁거나, 한쪽으로 치우친 예측을 자주 생성합니다. 이러한 "평균 추구" 행동은 예측 분산을 붕괴시켜 필요한 불확실성이나 상대적 순서를 포착하지 못하게 합니다. 이 한계는 가능한 값의 분포가 단일 최선 추정치만큼 중요한 후보 순위 매기기, 불확실성 추정, 또는 불확실성 하의 견고한 의사결정이 필요한 응용 분야를 방해합니다.

방법론: 분산 인식 보상 (DAR)

저자들은 여러 디코딩된 샘플로 형성된 전체 예측 분산의 품질을 최적화하도록 설계된 온-정책 강화 학습 목표인 **분산 인식 보상 (DAR)**을 제안합니다. 이는 개별 롤아웃을 고립적으로 최적화하는 대신, 전체 예측 분산의 품질을 최적화합니다.

핵심 메커니즘

1. 경험적 예측 분산: 주어진 입력 $x$에 대해 모델은 $K$개의 롤아웃을 샘플링하여 예측 집합 $\{p_{ik}\}_{k=1}^K$을 생성합니다. 이 집합은 경험적 예측 분산 $F_i$로 간주됩니다.
2. CRPS 를 통한 점수 매기기: 이 분산의 품질은 **연속 순위 확률 점수 (CRPS)**를 사용하여 실제 지상 진실 (ground-truth) 목표값 $y_i$와 비교하여 평가됩니다. CRPS 는 정확성 (목표값 근처에 집중됨) 과 적절한 분산 (퇴화를 피함) 을 모두 갖춘 분산에 보상을 주는 적절한 점수 규칙입니다.
   • 전체 집합에 대한 보상은 $S_i = -\text{CRPS}(F_i, y_i)$로 정의됩니다.
3. Leave-One-Out 크레딧 할당: RL 이 정책을 업데이트하기 위해 각 개별 롤아웃에 대한 스칼라 보상을 필요로 하므로, DAR 는 분산 수준의 점수를 롤아웃 수준의 피드백으로 변환합니다. 전체 집합의 CRPS($S_i$) 와 해당 롤아웃을 제거한 집합의 CRPS($S_i^{(-k)}$) 를 비교하여 각 롤아웃 $p_{ik}$의 한계 기여도를 계산합니다.
   • 롤아웃 보상은 다음과 같습니다: $R_{ik}^{\text{DAR}} = S_i - S_i^{(-k)}$.
   • 만약 특정 롤아웃을 제거하면 분산이 악화되어 (부정된 CRPS 가 낮아져) 해당 롤아웃은 개별 포인트별 오차가 높더라도 양의 보상을 받습니다. 반대로, 분산의 확산이나 중심화를 개선하지 않는 중복된 롤아웃은 더 낮거나 음의 보상을 받습니다.

이론적 차이점

이 논문은 이를 이론적으로 "평균 추구"적인 포인트별 제곱 오차 (MSE) 와 대비시킵니다. MSE 는 예측 분산이 단일 점 (조건부 평균) 으로 붕괴되도록 장려합니다. 반면, CRPS 는 "분산 추구"적이며 전체 조건부 분산을 최적화합니다. DAR 는 이를 활용하여 leave-one-out 크레딧 할당을 사용하여 잘 중심화되고 유용하게 분산된 롤아웃 집합을 장려합니다.

주요 기여

1. 새로운 RL 목표: 개별 점 오차 최소화가 아닌 경험적 예측 분산의 품질 극대화로 최적화 목표를 전환하는 DAR 의 도입.
2. Leave-One-Out 크레딧 할당: 개별 롤아웃을 집단적 분산 품질에 대한 한계 기여도에 따라 보상을 할당하는 구체적인 메커니즘으로, 개별 샘플이 덜 정확하더라도 "유용한 분산"에 보상을 부여할 수 있게 합니다.
3. 다양성 붕괴 완화: RL 미세 조정이 출력 다양성 (정책 엔트로피) 을 감소시키고 예측을 평균 쪽으로 붕괴시키는 경향을 해결하여, 샘플링된 롤아웃이 지상 진실을 충분히 커버할 수 있도록 다양성을 유지합니다.

실험 결과

저자들은 세 가지 다른 작업인 합성 가우스 혼합 회귀 작업, 코드 성능 예측, 분자 특성 예측에서 DAR 를 평가했습니다.

1. 합성 분산 회귀

• 작업: 이분산성 두 성분 가우스 혼합에서 목표값 예측.
• 결과: DAR 는 LLM 기반 방법들 중에서 가장 높은 스피어만 상관관계를 달성하여 SFT 와 포인트별 MSE 보상을 크게 능가했습니다. 이는 목표 함수의 전역적 순서를 더 잘 보존하고 다른 신경망 베이스라인이 어려움을 겪는 외삽 영역에 잘 적응함을 보여주었습니다.

2. 코드 성능 예측

• 작업: 소스 코드에서 커널 지연 시간 (KBSS) 및 피크 메모리 사용량 (APPS) 예측.
• 결과: DAR 는 SFT 와 포인트별 RL 베이스라인에 비해 가장 강력한 순위 상관관계 (스피어만 $\rho$) 를 달성했으며, 경쟁력 있거나 더 낮은 평균 절대 오차 (MAE) 를 보였습니다.
  • KBSS 에서 DAR 는 가장 강력한 베이스라인 대비 스피어만 상관관계를 6 포인트 개선했습니다.
  • APPS 에서는 스피어만 상관관계를 4 포인트 개선했습니다.
  • 이는 입력 의존적 변동에 대한 우수한 민감도와 더 신뢰할 수 있는 후보 순위 매기기를 나타냅니다.

3. 분자 특성 예측 (MoleculeNet)

• 작업: SMILES 문자열에서 물리화학적 특성 (FreeSolv, ESOL, 친유성) 예측.
• 결과: 강력한 GNN 베이스라인이 사용하는 명시적 그래프 또는 3 차원 유도 편향을 포함하지 않는 SMILES 문자열만 사용했음에도 불구하고 DAR 는 경쟁력을 유지했습니다.
  • 다른 LLM 학습 전략 (SFT, MSE, Value Head) 에 비해 FreeSolv 와 친유성에서 RMSE 에 통계적으로 유의미한 개선을 달성했습니다.
  • 이는 특수한 분자 표현 없이도 분산 수준 최적화가 견고한 회귀 행동을 이끌어낼 수 있음을 보여주었습니다.

불확실성 진단

• 보정: 분석 결과, DAR 의 롤아웃 표준 편차와 예측 오차 사이에 강한 로그-피어슨 상관관계 (0.91) 가 있음을 보여주어 모델의 불확실성 추정이 실제 오차와 잘 보정되었음을 나타냈습니다.
• 다양성: 훈련 기간 동안 DAR 는 MSE 기반 RL 에 비해 더 높은 롤아웃 엔트로피와 더 높은 "브래킷 비율" (지상 진실이 최소 및 최대 롤아웃 예측 사이에 있는 예제의 비율) 을 유지하여 다양성 붕괴가 완화되었음을 확인했습니다.

중요성과 주장

이 논문은 예측 분산을 직접 최적화함으로써 LLM 기반 회귀를 더 견고하고 잘 보정되게 만든다고 주장합니다. 여러 디코딩된 샘플을 분산으로 간주하고 그들의 집단적 품질에 보상을 부여함으로써, DAR 는 포인트별 정확성을 희생하지 않으면서 순위 기반 성능 (후보 선택에 중요) 과 불확실성 추정을 개선합니다.

저자들은 이 작업을 텍스트 생성기나 점 추정기가 아닌 과학 분야를 위한 유연한 확률적 회귀기로 LLM 을 사용하기 위한 한 걸음으로 위치시킵니다. 그들은 LLM 회귀가 이질적인 입력 (텍스트, 코드, 문자열) 처리에 유연성을 제공하지만, 불확실성 정량화와 신뢰할 수 있는 순서가 필요한 작업에서 전통적인 포인트별 학습이 실패하는 영역에서 이러한 모델의 잠재력을 완전히 발휘하려면 제안된 분산 인식 접근 방식이 필요하다고 강조합니다. 이 작업은 모든 분야에서 GP 나 GNN 과 같은 전문 회귀 모델을 대체한다고 주장하지는 않지만, 분산 인식 RL 이 범용 LLM 을 과학적 회귀 설정에서 경쟁력 있게 만들 수 있음을 입증합니다.