기술 요약: REAR: 보상 분해를 통한 테스트 시점 선호도 재정렬 (REAR: Test-time Preference Realignment through Reward Decomposition)

문제 정의

대규모 언어 모델(LLM)을 다양하고 개인화된 사용자 선호도에 맞게 정렬하는 것은 매우 중요한 과제입니다. 인간 피드백 기반 강화 학습(RLHF)과 같은 사후 학습(post-training) 방법은 효과적이지만, 비용이 많이 드는 데이터 큐레이션과 추가 학습이 필요하여 동적이거나 특정한 사용자 요구에 유연하게 대응하기 어렵습니다. 테스트 시점 스케일링(Test-Time Scaling, TTS)은 추론 중에 추가적인 연산을 활용하여 출력 품질을 개선함으로써, 학습이 필요 없는 대안을 제공합니다. 그러나 기존의 TTS 방법들은 주로 정답을 쉽게 검증할 수 있는 검증 가능한 도메인(예: 수학, 코딩)에 국한되어 왔습니다. 선호도 정렬에 TTS를 적용하는 것은 어렵습니다. 왜냐하면 응답의 품질은 총체적이며 단순한 검증 가능한 답으로 환원될 수 없기 때문입니다. 현재의 테스트 시점 정렬 접근 방식은 종종 고정된 계산 오버헤드를 부과하거나 가치 헤드(value head)를 학습시켜야 하는 한계가 있습니다.

방법론

본 논문은 **REAR(REAlignment Reward)**를 소개하며, 이는 테스트 시점 선호도 정렬을 재정렬(realignment) 문제로 프레임화하는 새로운 프레임워크입니다. 핵심 통찰은 베이스 모델의 학습 목적 함수가 추론 시점에 특정 사용자의 선호도를 최적으로 만족시키지 못할 수 있다는 점입니다. REAR는 근저에 있는 보상 함수를 두 가지 구성 요소로 분해하고 이들의 중요도를 동적으로 재조정함으로써 이를 해결합니다.

1. 보상 분해 (Reward Decomposition)

저자들은 최대 엔트로피 강화 학습(Maximum Entropy RL) 프레임워크 내에서 사전 학습된 LLM의 보상 구조를 공식화합니다. 저자들은 상태 $s$ (질문)와 선호도 $x_p$에 대한 보상 $r(s \oplus x_p, a)$가 다음과 같이 분해될 수 있다고 가정합니다:

• 질문 관련 보상 ($r_0$): 질문에 올바르게 답하기 위해 베이스 모델이 극대화하는 암묵적 보상.
• 선호도 관련 보상 ($r_p$): 응답이 선호도 $x_p$를 얼마나 잘 충족하는지를 포착하는 보상 구성 요소.

복합 보상은 $r(s \oplus x_p, a) = r_0(s, a) + \alpha r_p(s \oplus x_p, a)$로 모델링되며, 여기서 $\alpha$는 내재된 균형 계수입니다.

2. 재정렬 보상 (REAR)

특정 사용자를 위해 모델을 재정렬하기 위해, REAR는 추론 중에 선호도 구성 요소의 비중을 재조정하는 제어 가능한 계수 $\hat{\alpha}$ (또는 $\lambda = \hat{\alpha}/\alpha$)를 도입합니다:
$$r_{REAR}(s \oplus x_p, a) = r_0(s, a) + \hat{\alpha} r_p(s \oplus x_p, a)$$

결정적으로, 저자들은 이 보상을 위한 계산 효율적인 공식을 도출했습니다. 최적 정책과 가치 함수의 관계를 활용하여, $r_{REAR}$가 토큰 수준의 정책 로그 확률(log-probabilities)의 선형 결합으로 표현될 수 있음을 보여줍니다:
$$S_{REAR}(\tau) = \sum_{t=0}^{T} \gamma^t \left( (1-\lambda) \log \pi(a_t|s_t) + \lambda \log \pi(a_t|s_t \oplus x_p) \right)$$
이 점수 함수는 질문만 포함된 프롬프트와 선호도가 포함된 전체 프롬프트라는 두 가지 컨텍스트 하에서의 베이스 모델 로그 확률에만 의존합니다. 이는 별도의 보상 모델을 학습시키거나 점수 산정을 위해 추가적인 토큰을 생성할 필요가 없습니다.

3. 테스트 시점 스케일링과의 통합

REAR 점수는 두 가지 TTS 알고리즘에 통합됩니다:

• Best-of-N (BoN) 샘플링: $N$개의 응답을 생성하고 $S_{REAR}$ 점수가 가장 높은 것을 선택합니다.
• 다양한 검증기 트리 탐색 (Diverse Verifier Tree Search, DVTS): REAR 점수를 사용하여 브랜치 확장 및 선택을 유도함으로써, 선호도 정렬을 최적화하는 동시에 다양성을 촉진하는 트리 탐색 알고리즘입니다.

주요 기여

1. 재정렬의 공식화: 본 논문은 테스트 시점 선호도 정렬을 재정렬 문제로 공식화하고, 분해된 목적 함수로부터 유도된 계산 효율적인 보상인 REAR를 제안합니다.
2. 학습이 필요 없는 점수 산정: REAR는 베이스 모델의 내부 확률만을 활용하여 도출되므로, 추가적인 학습, 보상 모델 학습, 또는 스펙큘레이티브 생성(speculative generation) 단계가 필요하지 않습니다.
3. 확장 가능한 알고리즘: 저자들은 다양한 샘플 예산에 따라 생성 성능을 향상시킬 수 있도록 REAR에 의해 가이드되는 두 가지 확장 가능한 테스트 시점 알고리즘(BoN 및 DVTS)을 개발했습니다.
4. 일반화 가능성: 이 방법은 작업별 선호 텍스트를 주입함으로써 수학적 추론 및 시각적 작업으로까지 일반화될 수 있음을 보여주었습니다.

실험 결과

저자들은 Qwen 및 Llama 모델을 사용하여 여러 벤치마크에서 REAR를 평가했습니다:

• 선호도 정렬 벤치마크: PrefEval (명시적/암묵적 선호도 및 선택 포함), Multifaceted Bench, Ping-Pong (역할극)에서 REAR 가이드 방식(BoN 및 DVTS 모두)은 그리디 디코딩(greedy decoding), Amulet, 선형 정렬(Linear Alignment, LA), 그리고 생성적 보상 모델을 사용하는 Best-of-N(GenRM)을 포함한 베이스라인들을 일관되적으로 능가했습니다. 예를 들어, REAR를 적용한 DVTS는 여러 지표에서 그리디 디코딩보다 평균 11.6% 향상된 성능을 보였습니다.
• 긴 컨텍스트 강건성: REAR는 컨텍스트 길이가 길어짐에 따라 성능이 크게 저하되는 GenRM에 비해 긴 컨텍스트 시나리오(최대 16k 토큰)에서 우수한 강건성을 입증했습니다.
• 수학적 추론: 적절한 선호 텍사트(예: "엄격한 수학자처럼 행동하라")와 함께 수학적 작업(MATH500, AIME, AMC)에 적용했을 때, REAR 가이드 선택 방식은 다수결 투표(majority voting)보다 뛰어난 성능을 보이며 일관된 스케일링 이득을 보여주었습니다.
• 시각적 작업: 시각적 환각 탐지를 위한 MMHal-Bench에서, REAR는 그리디 디코딩 및 GenRM에 비해 전반적인 점수를 개선하고 환각률을 낮추었습니다.
• 교차 모델 전이: 이 방법은 모델별 튜닝 없이 Llama-3.1-8B-Instruct로 일반화되었으며, 이는 보상 분해가 Qwen 제품군에 국한되지 않음을 확인시켜 줍니다.
• 효율성: REAR 가이드 방식은 외부 보상 모델이나 반복적 최적화(Amulet와 같은)를 필요로 하는 베이스라인보다 훨씬 빨랐습니다. 이는 별도의 모델을 로드할 필요 없이 단일 순전파(forward pass)만으로 점수 산정이 가능하기 때문입니다.

의의 및 주장

본 논문은 REAR가 학습이나 외부 보상 모델의 오버헤드 없이 테스트 시점에 LLM을 사용자 선호도에 맞게 정렬할 수 있는 확장 가능하고 일반적인 솔루션이라고 주장합니다. 보상을 질문과 선호도 구성 요소로 분해함으로써, REAR는 질문에 답하는 것과 선호도를 준수하는 것 사이의 트레이드오프를 동적으로 제어할 수 있게 합니다.

저자들은 REAR를 통해 전통적인 검증기들이 실패하는 주관적이고 개방적인 도메인(선호도 정렬 및 역할극 등)으로 테스트 시점 스케일링을 확장할 수 있다고 강조합니다. 또한 기본 $\lambda = 20$이 많은 작업에서 견고하게 작동하지만, 이 방법은 재정렬의 강도를 조정할 수 있게 해준다고 언급합니다. 이 연구는 테스트 시점 정렬을 더욱 접근하기 쉽고 효율적으로 만들어, 검증 가능한 도메인에서의 TTS 성공과 주관적인 선호도 정렬의 복잡성 사이의 간극을 메우는 단계로 자리매김하고 있습니다.

인정된 한계점: 저자들은 REAR가 사용자에게 튜닝 가능한 하이퍼파라미터($\lambda$)를 노출한다는 점을 언급하며, 선호도 분포가 크게 다른 작업의 경우 작은 검증 스윕(validation sweeps)이 필요할 수 있다고 밝혔습니다. 또한, TTS가 파인튜닝보다는 가볍지만, 최적의 계산 예산을 식별하는 것은 향 향후 연구 과제로 남아 있습니다.