Rewarded soups: towards Pareto-optimal alignment by interpolating weights fine-tuned on diverse rewards

Alexandre Ramé; Guillaume Couairon; Mustafa Shukor; Corentin Dancette; Jean-Baptiste Gaya; Laure Soulier; Matthieu Cord

Source-linked AI summary

Rewarded soups: towards Pareto-optimal alignment by interpolating weights fine-tuned on diverse rewards

Alexandre Ramé, Guillaume Couairon, Mustafa Shukor, Corentin Dancette, Jean-Baptiste Gaya, Laure Soulier, Matthieu Cord

arXiv:2306.04488v2Jun 7, 2023cs.LGcs.AIcs.CV

TL;DR

Proxy-reward misspecification and diverse human objectives make single-policy alignment inadequate. Rewarded soup trains one expert per proxy reward and linearly interpolates their weights, yielding a flexible multi-policy set that the paper reports validating across language, multimodal, and control tasks. Its scope remains bounded by unreliable or hackable rewards, possible overfitting, test-distribution constraints, and assumptions underlying linear mode connectivity.

Problem
Proxy rewards may not match users’ actual rewards, while real-world objectives and human preferences can be diverse and conflicting.
Method
Rewarded soup independently fine-tunes networks on diverse proxy rewards and linearly interpolates their weights for a posteriori preference selection.
Results
The paper consistently validates rewarded soup and linear mode connectivity across text-to-text, text-image, and locomotion tasks.
Takeaways & Limitations
Rewarded soup provides an efficient, flexible multi-policy strategy intended to improve transparency, fairness, and alignment across diverse preferences.
Takeaways & Limitations
The approach relies on empirical linear mode connectivity, and its theoretical guarantees are limited to linear user rewards in a simplified quadratic setting.

Abstract

from arXiv · show

Foundation models are first pre-trained on vast unsupervised datasets and then fine-tuned on labeled data. Reinforcement learning, notably from human feedback (RLHF), can further align the network with the intended usage. Yet the imperfections in the proxy reward may hinder the training and lead to suboptimal results; the diversity of objectives in real-world tasks and human opinions exacerbate the issue. This paper proposes embracing the heterogeneity of diverse rewards by following a multi-policy strategy. Rather than focusing on a single a priori reward, we aim for Pareto-optimal generalization across the entire space of preferences. To this end, we propose rewarded soup, first specializing multiple networks independently (one for each proxy reward) and then interpolating their weights linearly. This succeeds empirically because we show that the weights remain linearly connected when fine-tuned on diverse rewards from a shared pre-trained initialization. We demonstrate the effectiveness of our approach for text-to-text (summarization, Q&A, helpful assistant, review), text-image (image captioning, text-to-image generation, visual grounding, VQA), and control (locomotion) tasks. We hope to enhance the alignment of deep models, and how they interact with the world in all its diversity.

1 Introduction

The paper argues that diverse proxy rewards and human preferences make single-policy alignment inadequate. Rewarded soup addresses this by training reward-specific experts and interpolating their weights to produce preference-dependent, near-Pareto-optimal policies.

Empirical result: With two summarization reward models, λ-interpolation reveals a Pareto-optimal front and matches the more costly MORL front.MORL requires multiple trainings across linear reward weightings, whereas rewarded soup uses the two independently fine-tuned networks.
Motivation: Proxy rewards can diverge from users’ actual rewards, while real-world objectives and human opinions may conflict substantially.Examples include differing preferences about harmlessness, helpfulness, engagement, aesthetics, politics, and fairness.
Motivation: Single-policy alignment fixes trade-offs before training, potentially reducing transparency and fairness when preferences differ.The paper frames alignment as a multi-objective problem requiring a set of Pareto-optimal networks rather than one network.
Rewarded soup: Rewarded soup trains one network per proxy reward and linearly interpolates their weights according to user preferences.The approach makes preference selection a posteriori rather than committing to a single reward weighting during training.
Empirical result: The method is evaluated across language, multimodal, and locomotion tasks, consistently validating linear mode connectivity and rewarded soup’s effectiveness.Evaluations include LLaMA RLHF, image captioning, text-to-image generation, and control tasks.

2 Rewarded soups

Rewarded soup replaces costly reward scalarization with a multi-policy procedure that interpolates weights from independently reward-specialized experts. Its viability rests on linear connectivity and Pareto-optimality hypotheses, which the paper reports validating empirically and approximately justifying theoretically.

Problem setup: The alignment objective uses a user reward and deployment distribution, but training optimizes proxy rewards because the user reward is usually unknown beforehand.This creates the paper’s reward-misspecification setting.
MORL baseline: MORL interpolates proxy rewards across M trainings, requiring a finely spaced set of networks because small preference changes can alter solutions significantly.The paper identifies this requirement as unscalable for deep learning.
Rewarded soup: Rewarded soup requires only M = N trainings by independently fine-tuning one expert weight set for each proxy reward, then interpolating those weights.All experts are fine-tuned from a shared pre-trained initialization.
Selection and inference: The interpolation coefficient can be selected after training using validation samples, without additional training, and the resulting weights are used for inference.This supports custom weights for different preference trade-offs.
Working hypotheses: The LMC hypothesis requires interpolated policies’ rewards to exceed the corresponding interpolated rewards, while the Pareto hypothesis treats the interpolated set as Pareto-optimal.These hypotheses underpin rewarded soup’s viability.
Evidence and assumptions: The hypotheses are empirically validated across tasks and approximately proved only in a simplified quadratic-reward setting with co-diagonalizable Hessians.The theory relies on a good pre-trained initialization and weights remaining close during fine-tuning.

3 Experiments

The experiments evaluate rewarded soup across text, multimodal, and locomotion tasks, finding that interpolated expert weights produce flexible Pareto trade-offs across diverse rewards. RS generally matches or exceeds multi-objective reinforcement learning fronts while requiring fewer trainings.

Experimental setup: Across text-to-text, image captioning, image generation, visual grounding, visual question answering, and locomotion, the experiments independently optimize diverse rewards before evaluating weight interpolation.The procedure keeps architecture, hyperparameters, and RL algorithm fixed while varying the reward across runs.
Image captioning: 0.140 hypervolume: interpolating BLEU1- and ROUGE-specialized captioning weights qualitatively matches MORL while producing a convex set of trade-off solutions.Uniformly averaging five independently fine-tuned weights yields the largest area and best trade-off across BLEU1, BLEU4, ROUGE, METEOR, and CIDEr.
Image captioning: Adding each additional reward-specialized network extends the set of rewards that rewarded soup can address Pareto-optimally in five-metric captioning.The experiment averages 1 ≤ M ≤ N networks, sequentially adding specialists trained on BLEU1, BLEU4, ROUGE, METEOR, and CIDEr.
Captioning analysis: RS improves reward endpoints and can mitigate reward misspecification when evaluation rewards differ from the training rewards.The optimal λ depends on similarity between evaluation and training rewards, while weight interpolation can approximate prediction ensembling when weights remain close.
Other tasks: In text-to-image generation, RS produces a convex front above MORL, while visual grounding uses λ to trade accuracy across small, medium, and large objects.The locomotion and VQA experiments likewise report Pareto fronts, although MORL with µ = 0.5 slightly surpasses RS in locomotion.
Efficiency: MORL needs M ≫ 2 trainings to match RS, which uses two trainings for two rewards.This efficiency advantage follows from selecting λ after training rather than launching separate trainings for many preference weightings.

4 Related work

The related work positions rewarded soup alongside proxy-reward reinforcement learning and linear mode connectivity. Its distinction is applying weight interpolation to RL-trained policies across diverse rewards.

Proxy rewards: Rewarded soup builds on proxy rewards and existing reinforcement-learning algorithms, including REINFORCE, to optimize objectives that are difficult to specify directly.The paper contrasts hand-engineered statistical metrics with reward models learned from human comparisons.
Linear mode connectivity: The paper investigates linear mode connectivity in reinforcement learning, extending prior connectivity work across tasks, modalities, and losses.Related control work averages decision transformers or explicitly enforces connectivity for policies trained from scratch on a single reward.

5 Discussion: limitations and societal impacts

The paper identifies unresolved limitations involving proxy-reward reliability, theoretical guarantees, preference scope, and personalization risks. These boundaries constrain how broadly rewarded soup’s alignment claims should be interpreted.

Proxy rewards may be non-robust or adversarially hacked, while overfitting can cause poor generalization and goal misgeneralization.The paper also states that complete alignment may be impossible without constraints on the test distribution, such as arbitrary prompt lengths.
Rewarded soup relies on empirical linear mode connectivity, whose full theoretical guarantees are currently unavailable.The paper provides near-optimality guarantees only for quadratic rewards with co-diagonalizable Hessians and weights remaining close.
The theoretical analysis fixes the user reward as linear over proxy rewards, which cannot represent all human preferences.The authors report improvements even when the reward is not linear, but the formal limitation remains.
Personalization may improve usefulness for under-represented groups but risks reinforcing individual biases and narrowing information exposure.The distributed implementation may also support privacy-preserving federated learning, but personalization introduces these societal risks.

6 Conclusion

The paper presents rewarded soup as a weight-interpolation strategy for efficiently producing Pareto-optimal solutions across diverse needs. Experiments across significant large-scale learning tasks support its working hypotheses and indicate mitigation of reward misspecification.

Rewarded soup efficiently yields Pareto-optimal solutions through weight interpolation after training.The method is presented as a response to the need for transparent alignment with diverse user needs.
Experiments consistently validate rewarded soup’s working hypotheses across various significant large-scale learning tasks.
The experiments demonstrate that rewarded soup can mitigate reward misspecification.

Supplementary material

The supplementary discussion contrasts rewarded soup with single-policy, multitask, and model-soup strategies. It presents flexibility, personalization, broad empirical connectivity, and computational advantages alongside the shared-initialization condition.

Single-policy strategies: Single-policy methods optimize one preference set, whereas rewarded soup builds a coverage set of Pareto-optimal policies.
Single-policy strategies: Rewarded soup postpones trade-off selection until after training, making personalization explicit and supporting changing human preferences.
Multi-policy strategies: Unlike multitask approaches that require dense policy sets, rewarded soup independently trains proxy rewards and selects the interpolation coefficient a posteriori.The authors emphasize this advantage when the number of rewards and possible combinations grows.
Model soups: Rewarded soup and model soups both average weights from a shared initialization but address different objectives and use different interpolation choices.Rewarded soup varies training rewards and uses non-uniform coefficients selected a posteriori; model soups usually combine models trained for one objective with uniform coefficients.
Model soups: Figure 14(b) indicates that rewarded soup and model soups are complementary, combining reward-misspecification reduction with variance reduction.
Linear mode connectivity: Linear mode connectivity is reported as architecture-, procedure-, task-, and modality-agnostic across transformers, CNNs, MLPs, adapters, partial fine-tuning, and end-to-end fine-tuning.
Linear mode connectivity: The main required condition is a shared pre-trained initialization, which keeps fine-tuned weights close; sufficient trainable parameters may also help.

B.2.2 Limitations for the number of training steps for the LMC?

The discussion reports that rewarded soup remains effective during longer fine-tuning and across multiple rewards, while its Pareto-optimality evaluation relies on an unavailable oracle and linearized MORL reference. Performance depends on reward similarity and can weaken for antagonistic rewards.

Training-step limitations: Doubling training steps still produced good rewarded-soup performance throughout fine-tuning for news summarization and image captioning.The corresponding fronts are shown in Figure 9, including the previously used two and six training epochs.
Reward dimensionality: Rewarded soup scales empirically to up to N = 4 text-generation, N = 5 captioning, and N = 3 visual-grounding rewards.
Pareto-optimality evaluation: The true Pareto front and oracle policy are unknown in real-world applications, so linearized MORL solutions serve as the practical reference.
Pareto-optimality evaluation: Rewarded soup and MORL usually perform similarly, supporting rewarded soup as an empirical solution toward Pareto-optimality.
Reward diversity: RS gains over MORL are higher when rewards are more similar, while including an image-quality-inversely-correlated nsfw reward worsens image-generation performance.
Reward diversity: Fully antagonistic rewards may yield less favorable results because fine-tuned weights can become more distant and weaken linear mode connectivity.The theoretical bound also becomes looser as reward-Hessian eigenvalue ratios grow with reward diversity.
Reward diversity: MORL and rewarded soup may be complementary for antagonistic rewards, providing a proposed direction for addressing this limitation.
Comparison with MORL: Rewarded soup sometimes exceeds linearized MORL, which the authors speculate may reflect implicit regularization and variance reduction from weight averaging.

C.2.1 Simple case with Hessians proportional to the Identity matrix

Under simplified quadratic-reward assumptions, the paper shows that linear weight interpolation can optimize weighted user rewards. With diagonal Hessians, it derives a bound on the interpolation gap and identifies when rewarded soups are optimal.

Simple case with Hessians proportional to the Identity matrix: Under Assumption 1, each reward is quadratic with a Hessian proportional to the identity, and θ_i is its global maximum.
Simple case with Hessians proportional to the Identity matrix: The weighted reward R̂_µ = Σ_i ˆµ_i × R_i is maximized on the convex hull of the independently optimized weights {θ_i}.
Simple case with Hessians proportional to the Identity matrix: The optimal rewarded-soup coefficients satisfy λ_i ∝ ˆµ_iη_i, making the interpolated weights optimal for the weighted reward.
Advanced case with diagonal Hessians: Under diagonal Hessians, Lemma 3 bounds the gap between the global optimum and the best single-coefficient interpolation.
Advanced case with diagonal Hessians: When the Hessians are equal, M = 1 and the interpolation gap is zero, so rewarded soups are optimal.
Advanced case with diagonal Hessians: The proof maximizes the weighted reward over all weights, maximizes it along the interpolation path, and applies the Bhatia-Davis inequality to bound their difference.
Advanced case with diagonal Hessians: Empirically, the bound is tight for M = 1 and the rewarded-soup value remains close to the global optimum for M = 2, but is less tight for M = 10.
Advanced case with diagonal Hessians: The interpolation approximation between weights and prediction ensembling holds when the optimized weights remain close, with the discrepancy represented by a higher-order term.

D.1 Experimental details

The text-to-text experiments fine-tune an instruction-following LLaMA model with PPO under multiple rewards, then inspect summaries and quantitative quality metrics produced by interpolated models.

Experimental setup: The experiments use LLaMA-7b with low-rank adapters fine-tuned on Alpaca before PPO fine-tuning on each considered task.
Quantitative evaluation: Interpolating the weights also interpolates the evaluated metrics, and the authors report no detected reduction in generated-summary quality.
Prediction inspection: The supplementary material provides generated news-summary examples interpolated between models trained on completeness and faithfulness rewards.
Quantitative evaluation: Figure 11 evaluates interpolated summaries using MLMS or GPT2 perplexity and a news-quality classifier.

E.2 Additional results

Additional captioning experiments test rewarded soups across alternative reward pairs and interpolation settings, revealing both metric-dependent fronts and limitations when a reward duplicates pre-training.

Additional captioning results: Figure 12 evaluates BLEU4 with ROUGE, and BLEU1 with METEOR or CIDEr, extending the captioning reward-pair analysis.
Additional captioning results: Optimizing CIDEr a second time fails to improve CIDEr or the other rewards because the pre-trained initialization was already optimized for CIDEr.
Additional captioning results: The optimal interpolation coefficient depends on the similarity between the evaluation metric and the two training rewards.
Additional captioning results: Interpolating fine-tuned networks with the pre-trained initialization exposes only a small portion of the front, whereas adding a MORL solution as an intermediate weight improves the rewarded-soup front.

F.1 Experimental details

The image-generation experiments combine aesthetic and safety-related rewards using an offline reward-weighted fine-tuning procedure. Interpolation preserves coherent generation and enables aesthetic trade-offs, but performs worse than MORL when safety conflicts with image quality.

Task description: The experiments use ava and cafe aesthetic models, plus an nsfw detector, to represent diverse image preferences and safety evaluation.
Implementation details: A 2.2B-parameter diffusion model is fine-tuned with reward-weighted negative log-likelihood using image rewards and text prompts.
Implementation details: Only 10% of the diffusion model’s weights are fine-tuned, specifically cross-attention layers and bias or scaling parameters.
Additional results: Rewarded soups trade off the ava and cafe aesthetic rewards, allowing user-preference adaptation at test time.
Additional results: With ava, cafe, and nsfw rewards together, MORL scores higher than rewarded soups, plausibly because nsfw differs from aesthetic preferences and is inversely correlated with image quality.
Quantitative evaluation: FID and CLIPScore evaluations confirm that images from interpolated models remain coherent.
Qualitative interpolation: Interpolating models fine-tuned on ava and cafe produces images of similar quality to the fine-tuned models, demonstrating linear mode connectivity.

G.1 Experimental details

The experiments apply reward-specialized fine-tuning to visual grounding, VQA, and humanoid locomotion, using task-specific models, rewards, and reinforcement-learning procedures. Locomotion additionally tests robustness across hyperparameter settings, with results remaining close and validating the working hypotheses.

Visual grounding: Visual grounding uses UnIVAL fine-tuned on RefCOCO+ with cross-entropy plus REINFORCE rewards specialized for small, medium, or large objects.The resulting RS model averages three weights, while predictions are evaluated by whether their IoU exceeds 0.5 against ground-truth boxes.
Visual grounding: Visual grounding predicts a bounding box from one text description and evaluates object-size-specific predictions against corresponding validation-set ground truths.Images are discretized into 1000 × 1000 bins before calculating box areas.
Visual question answering: VQA fine-tunes OFA on VQA v2 with cross-entropy, then applies REINFORCE on different rewards using a held-out set for reinforcement-learning fine-tuning.The supplied implementation description identifies BLEU1 among the rewards used.
Locomotion: The locomotion experiment controls a running humanoid in Brax and fine-tunes a pre-trained policy on risky and cautious reward functions.Pre-training uses Brax’s default dense reward with PPO, collects observation statistics for normalization, saves the policy, and discards the value function; fine-tuning freezes those statistics.
Locomotion: Results remain close across tried locomotion hyperparameters and consistently validate the working hypotheses.The hyperparameter analysis interpolates each risky policy toward each cautious policy, with Figure 20 reporting the resulting variance analysis.