Source-linked AI summary
Rewards-in-Context: Multi-objective Alignment of Foundation Models with Dynamic Preference Adjustment
Rui Yang, Xiaoman Pan, Feng Luo, Shuang Qiu, Han Zhong, Dong Yu, Jianshu Chen
TL;DR
Multi-objective alignment must accommodate heterogeneous and conflicting human preferences without the high cost of reinforcement-learning fine-tuning. RiC conditions a single foundation model on multiple rewards, using supervised fine-tuning, online augmentation, and inference-time preference mapping. Across reported alignment experiments, it achieves strong empirical fronts with substantially lower computational cost, while prompt-controlled conditioning introduces a safety risk requiring additional checks and filters.
Problem
Heterogeneous, multidimensional, and conflicting human preferences make single-reward alignment inadequate, while multi-objective reinforcement-learning alignment is costly.
Method
RiC uses multi-reward conditional supervised fine-tuning, online Pareto-front data augmentation, and dynamic inference-time preference-to-reward mappings for one model.
Results
RiC achieves stronger empirical alignment fronts than baselines while RiC with two online iterations uses 10.4% of MORLHF’s GPU hours.
Takeaways & Limitations
RiC provides a scalable and computationally efficient approach for adapting foundation-model behavior across diverse reward preferences.
Takeaways & Limitations
Malicious users may set harmful desired rewards through the prompt, so additional checks and filters are needed.
Abstract
from arXiv · showhide
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.
1. Introduction
Foundation models can be aligned with human feedback, but heterogeneous, multidimensional, and conflicting preferences make single-reward alignment insufficient and multi-objective RL expensive. RiC addresses this with conditional supervised fine-tuning, online data augmentation, and inference-time preference adaptation.
- Human preferences are heterogeneous, multidimensional, and potentially conflicting, so a single reward model may not adequately align language models.
- Linear-scalarization MORLHF requires substantial computation because the preference space can be large, even when preferences are quantized.
- RiC restructures multi-objective alignment into offline multi-reward conditional SFT, online Pareto-front data augmentation, and inference-time preference adaptation.
- RiC avoids modified loss functions and structured per-objective preference data while extending to more rewards with minimal computational-cost growth.
- 10.4% of MORLHF’s GPU hours is required by RiC with two online iterations in the reported two-objective Helpful Assistant experiment.
2. Background
The background contrasts supervised fine-tuning and RLHF with multi-objective alignment methods. MORLHF scalarizes multiple rewards using preference weights, while Rewarded Soups interpolates policies to reduce training burden.
- Supervised fine-tuning trains language models from labeled prompt-response demonstrations using a token-level likelihood objective.
- RLHF trains a reward model from preferred and dispreferred responses, then applies reinforcement learning with a KL penalty relative to the SFT policy.
- MORLHF represents multiple rewards as a vector and uses preference weights on the simplex to form a scalarized optimization objective.
- Rewarded Soups linearly combines the weights of reward-specialized policies according to the user preference vector.
- Rewarded Soups reduces the burden associated with preference enumeration because the number of discretized preferences is generally larger than the number of reward models.
3. RiC Algorithm
RiC trains a single conditional policy on multiple reward annotations, augments scarce Pareto-front data online, and maps user preferences to reward prompts at inference time. Its optimization-based mappings and conditional SFT support flexible multi-objective alignment with limited added cost.
- RiC trains one model through offline conditional SFT, online data augmentation, and inference-time adaptation instead of separately training for each preference.
- Offline Training: Reward-conditional training teaches the policy to ground responses in desired rewards, with obtained test rewards positively correlated with desired rewards.
- Offline Training: RiC relabels each prompt with reward-model outputs, normalizes reward dimensions, and performs multi-reward conditional SFT.
- Offline Training: Multi-reward conditional SFT uses positive and negative responses without explicit filtering and scales to additional rewards with minimal prompt-cost increase.
- Online Training: Scarcity of original samples near the empirical reward frontier motivates generating responses near that frontier for online data augmentation.
- Inference Stage: Linear preference-to-reward mapping can miss Pareto optimality, so RiC uses optimization-derived mappings to dynamically set reward conditions from user preferences.
- Determining the Preference-to-Reward Mapping: The mapping optimization imposes a regularized trade-off across normalized reward outcomes and gives higher outcomes to dimensions with larger preferences.
- Determining the Preference-to-Reward Mapping: For the practical regularization set, the constrained problem has a closed-form solution, and experiments report resilience to the choice of p.
4. Experiments
Experiments evaluate RiC on two text-generation tasks and one text-to-image task, testing multi-objective alignment, scalability, efficiency, and component ablations. RiC balances competing rewards, improves frontier quality with online training, and supports preference-controlled trade-offs across modalities.
- Text generation: RiC effectively aligns Helpful Assistant models across harmless-helpful and humor-helpful preferences, achieving a superior frontier to MORLHF and Rewarded Soups.
- Text generation: On Reddit Summary, RiC significantly outperforms baselines by restoring base-model faithfulness while improving preference rewards.
- Three-objective alignment: RiC provides the most balanced performance across harmless, helpful, and humor rewards, with a frontier closer to Pareto efficiency than Rewarded Soups.
- Computational cost: RiC w/ online iter2 uses 10.4% of MORLHF’s GPU hours and 24.6% of Rewarded Soups’ GPU hours in the two-objective Helpful Assistant experiment.
- Ablations: Offline training alone achieves a considerable empirical front, while additional online training improves that front.
- Ablations: Mappings with p = 2 and p = ∞ outperform linear mappings, while showing similar performance and robustness to p > 1.
- Text-to-image generation: In text-to-image generation, increasing w1 produces more beautiful but less compressible images, whereas decreasing w1 increases compressibility.
5. Related Work
Related work spans RLHF, multi-objective RLHF, reward-conditioned supervised fine-tuning, and multi-attribute conditioning. RiC differs by targeting multi-objective alignment and optimizing an empirical Pareto front with adaptive reward conditioning.
- RLHF: RLHF uses human feedback to guide reinforcement learning but faces instability and inefficiency challenges.
- Reward-conditioned training: Reward-conditioned training augments prompts with rewards and applies supervised fine-tuning, whereas RiC extends this setting to multiple objectives.
- MORL and MORLHF: MORLHF addresses heterogeneous preferences through multi-objective reinforcement learning, including approaches such as Rewarded Soups and MODPO.
- Multi-attribute conditioned SFT: Unlike SteerLM, RiC considers reward trade-offs, targets an optimal empirical Pareto front, and assigns desired rewards closer to that frontier during online generation.
6. Conclusion
RiC addresses multi-objective alignment with supervised fine-tuning and simple inference-time adaptation, reducing computational cost while supporting diverse objectives and user preferences.
- RiC uses only supervised fine-tuning and simple inference-time adaptation for multi-objective alignment.The approach is designed as a scalable solution for heterogeneous objectives.
- RiC reduces computational cost while maintaining strong alignment performance across a range of objectives.
- The authors identify context-dependent inference-time adaptation as a direction for further improving alignment performance.
Impact Statement
The paper identifies potential misuse of reward-conditioning prompts as a safety concern and recommends proactive safeguards for trustworthy deployment.
- Malicious users may set harmful desired rewards through prompts, potentially leading to negative outcomes.
- The paper states that this safety issue is encountered by all conditional training methods.
- Additional checks and filters are recommended as proactive measures for developing trustworthy AI systems.
A.2. Proof of Theorem 3.1
The proof reduces the relevant optimization to a convex problem, invokes strong duality and KKT conditions, and verifies solution feasibility across cases including 1 < p < ∞ and p = ∞.
- The regularization set ensures zi ≤ 1, allowing that constraint to be omitted from the optimization problem.
- The case 1 < p < ∞ is solved through a Lagrangian formulation with dual variables.
- The optimization is treated as a convex problem satisfying Slater’s condition, so strong duality applies.
- The proof analyzes the primal solution using KKT conditions and its relation to the max-min formulation.
- The proof establishes u* > 0 under the nonnegative preference-vector assumptions by contradiction.
- For 1 < p < ∞, the derived solution satisfies the ordered feasibility constraints and therefore solves the optimization problem.
- When p = ∞, the proof converts the norm constraint into coordinatewise bounds and selects z*i = 1/λζi.
- The resulting p = ∞ solution also satisfies the required ordering constraints, completing the case analysis.
B. Implementation Details
The text-generation experiments use Llama 2-7B and distinguish RiC’s direct supervised fine-tuning from PPO-based baselines, with reward markers and normalization supporting conditional generation.
- RiC directly fine-tunes the Llama 2-7B base model, whereas MORLHF and Rewarded Soups fine-tune the SFT model with PPO.The implementation is based primarily on trl and uses the Llama 2-7B base model.
- Reward values are represented in prompts using distinct position marks such as <harmless score> and <helpful score>.
- RiC normalizes rewards using the offline dataset’s mean and standard deviation before inserting them into prompts.
- Table 3 summarizes the key implementation details of the text-generation experiments.
- During online generation and evaluation, sampled normal-distribution extrema replace dataset extrema to reduce the influence of extreme reward values.
C.1. Text-to-image Generation: diffusion models with RiC
RiC applies multi-reward conditioning and supervised fine-tuning to text-to-image generation, enabling adaptation to different reward preferences. Experiments show alignment across preferences and model sizes, with RiC outperforming MODPO in the reported language-task comparisons.
- Task Setup: RiC aligns text-to-image generation using the LAION aesthetic predictor and compressible score as two reward objectives.Samples use a fixed 512×512 resolution, with compressibility defined through JPEG file-size minimization.
- Experimental Results: 1,000 generated images evaluated on COCO captions show RiC adapts to different preferences, unlike the Stable Diffusion v1.5 base model.Each plotted point represents average rewards for a specific user preference.
- Experimental Results: RiC generates images with superior performance on both the LAION aesthetic and compressible score reward models.
- Model Size: RiC achieves successful multi-objective alignment across 1B, 3B, and 7B models, with larger models typically yielding better alignment results.
- Comparison: RiC outperforms MODPO in both the Helpful Assistant and Reddit Summary settings.MODPO aligns two objectives in Helpful Assistant but fails to align both objectives in Reddit Summary.
C.4. Additional Ablation Study
The ablations examine online fine-tuning, MORS, and offline-data regularization, while an additional experiment tests positively correlated rewards. Online fine-tuning is most influential, whereas positively correlated rewards remain a clear limitation.
- MORS Ablation: Online fine-tuning without MORS produces the largest performance improvement, while MORS adds a slight further enhancement.The comparison includes pure offline training, online fine-tuning without sample rejection, and default RiC with 0.7-quantile MORS thresholds.
- Offline Data Regularization: Offline-data regularization does not directly produce significant performance gains and is therefore optional.RiC mixes offline data equal to 1/2 the online-sample size into the online buffer.
- Offline Data Regularization: The regularization can help prevent online fine-tuning from deviating too far from the offline policy, supporting learning stability.
- Positively Correlated Rewards: RiC performs poorly with positively correlated rewards and does not produce a clear Pareto front.The authors attribute this to SFT capturing reward correlation and focusing primarily on one reward.
C.6. Examples in the Text Generation Tasks
Text-generation examples show that RiC adjusts Helpful Assistant outputs according to the relative weighting of helpfulness and harmlessness. Lower helpfulness weight favors helpful responses that may be harmful, while higher weight favors harmless but less helpful responses.
- Helpful Assistant: RiC adjusts Helpful Assistant outputs by changing the relative weighting of harmless and helpful rewards.
- Helpful Assistant: Smaller w1 values tend to produce helpful responses that can be potentially harmful for malicious prompts.
- Helpful Assistant: Larger w1 values tend to produce harmless sentences that are not necessarily helpful.The examples include responses such as “I don’t unders...”.
- Related Work: RiC is related to goal-reaching in multi-reward settings by setting desired rewards on the Pareto front.The approach uses generalization and extrapolation of the learned LLM policy to achieve a superior empirical front.
- Related Work: During online training, RiC relabels desired rewards with obtained rewards, further enhancing performance.