Rewards-in-Context: Multi-objective Alignment of Foundation Models with Dynamic Preference Adjustment

We consider the problem of multi-objective alignment of foundation models with human preferences, which is a critical step towards helpful and harmless AI systems. However, it is generally costly and unstable to fine-tune large foundation models using reinforcement learning (RL), and the multi-dimensionality, heterogeneity, and conflicting nature of human preferences further complicate the alignment process. In this paper, we introduce Rewards-in-Context (RiC), which conditions the response of a foundation model on multiple rewards in its prompt context and applies supervised fine-tuning for alignment. The salient features of RiC are simplicity and adaptivity, as it only requires supervised fine-tuning of a single foundation model and supports dynamic adjustment for user preferences during inference time. Inspired by the analytical solution of an abstracted convex optimization problem, our dynamic inference-time adjustment method approaches the Pareto-optimal solution for multiple objectives. Empirical evidence demonstrates the efficacy of our method in aligning both Large Language Models (LLMs) and diffusion models to accommodate diverse rewards with only around 10% GPU hours compared with multi-objective RL baseline.

한 줄 요약

RiC는 복수의 사용자 선호를 고려한 기초 모델 정렬 문제를 해결하기 위해, 단일 모델의 지도 학습과 추론 시 동적 조정을 결합한 효율적인 알고리즘을 제안한다.

핵심 기여도

**RiC 알고리즘**을 제안하여, 기존 MORLHF 대비 약 10%의 GPU 시간으로도 우수한 정렬 성능을 달성.
**지도 학습 기반 정렬**과 **추론 시 선호도 조정**을 결합하여, 복수의 보상 함수를 동시에 고려하는 Pareto 최적 접근.
**Rewarded Soups**와 비교해, 기저 모델의 능력을 유지하면서 보상 간 균형을 개선.
**3개 이상의 목표**를 처리하는 확장성을 실험적으로 검증.

핵심 아이디어

RiC는 기존의 복수 목표 정렬 문제에서 사용되는 강화 학습 기반 접근법(MORLHF)의 높은 계산 비용과 불안정성을 해결하기 위해, **지도 학습(SFT)**과 **추론 시 조정**을 결합한 새로운 프레임워크를 제안한다. 기존 MORLHF는 사용자 선호도에 따라 복수의 모델을 학습하거나 가중치를 조합하는 방식을 사용하지만, 이는 계산 비용이 높고, 모델의 기존 능력을 유지하기 어렵다. RiC는 대신, **prompt에 여러 보상 함수를 조건으로 삼아 단일 모델을 학습**하고, 추론 시에는 **선호도 벡터를 기반으로 보상 가중치를 동적으로 조정**한다. 이는 볼록 최적화 문제의 해석적 해법에서 영감을 받아, Pareto 최적에 가까운 결과를 도출한다.

기술적 접근법

RiC는 다음과 같은 세 단계로 구성된다:
1. **오프라인 학습 단계**: 여러 보상 함수를 조건으로 삼은 **multi-reward conditional SFT**를 수행.
2. **온라인 학습 단계**: 경험적 Pareto front 근처의 데이터를 보강하여 추가 학습.
3. **추론 단계**: 사용자 선호도 벡터 **𝒘 ∈ Ω**를 기반으로, **선호도-보상 매핑**을 통해 동적 조정.

RiC는 별도의 손실 함수 수정 없이 기존 SFT를 활용하며, **Rewarded Soups**와 달리 가중치 조합이 아닌 **prompt 내 조건 조정**을 통해 보상 간 균형을 유지한다. 실험에서는 **LLM**과 **diffusion 모델** 모두에서 적용 가능함을 보여준다.

주요 결과

**Helpful Assistant** 태스크에서, RiC는 MORLHF 대비 약 **10.4%**, Rewarded Soups 대비 **24.6%**의 GPU 시간만으로 동일한 수준의 정렬 성능을 달성.
**Reddit Summary** 태스크에서, 기저 모델의 'faithful' 보상 점수를 유지하면서 'pref1' 및 'pref2' 보상도 향상.
**3개 이상의 목표**를 처리하는 실험에서, RiC는 **Pareto 최적에 가까운 성능**을 보이며, 기존 방법들보다 균형 잡힌 결과를 제공.
**RiC w/ online iter2**는 총 **153.2 GPU 시간**만으로, MORLHF(1,477.1 GPU 시간) 대비 약 10% 수준의 성능을 달성.

의의 및 한계

RiC는 복수의 사용자 선호를 고려한 기초 모델 정렬 문제에서 **계산 효율성**과 **유연성**을 동시에 달성한 점에서 학술적·실용적 가치가 있다. 특히, **단일 모델 학습**과 **추론 시 동적 조정**을 결합함으로써, 기존 MORLHF의 높은 계산 비용 문제를 해결하고, Rewarded Soups의 가중치 조합 방식의 한계를 극복한다. 그러나, RiC는 **선호도 벡터의 정확한 입력**을 전제로 하므로, 사용자 입력의 불확실성이 있을 경우 성능이 저하될 수 있다. 또한, **복수 보상 간의 충돌이 극심한 경우**에도 Pareto 최적에 도달하기 어려울 수 있다.

실용적 활용

RiC는 **대규모 언어 모델**과 **이미지 생성 모델** 모두에서 사용 가능하며, **개인화된 AI 서비스**(예: 챗봇, 콘텐츠 생성)에서 사용자 선호도에 따라 유연하게 반응을 조정할 수 있다. 또한, **다중 목적 최적화**가 필요한 산업 분야(예: 자율 주행, 의료 AI)에서도 활용 가능하다.