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

Summary

The paper proposes rewarded soup (RS), a simple technique for combining policies trained for different rewards into a single policy performing well on a particular convex combination of those rewards. The technique consists in linearly interpolating weights of individual policies, using the fact that they share the same initialisation and remain linearly connected during finetuning on different rewards. The approach is well-motivated theoretically and thoroughly evaluated on a diverse array of tasks ranging from language generation, image captioning and image generation to robot locomotion.

Strengths

1. The paper addresses crucial problems of accounting for diverse preferences and adapting to changes in reward specification that frequently arise in the emerging and important field of aligning generative models with human preferences. 2. Rewarded soup is well-motivated theoretically as a Pareto coverage set of policies for linear combinations of individual reward functions. I found working hypotheses 1 and 2 to be very helpful in understanding how RS works. I’m also convinced by empirical evidence for these hypotheses being true. 3. The fact that linear mode connectivity also holds for RL policies trained for different rewards is an interesting finding about deep learning overall. I found it somewhat surprising that interpolated weights outperform the interpolated rewards so consistently. 4. The paper is well-written and easy to follow despite being very dense. The theory part connects with experiments very well. I also appreciate the plots (e.g. Figure 2) being easy to navigate while conveying a lot of information. 5. The experiments are very thorough and diverse and I find them compelling.

Weaknesses

I think the discussion of reward misspecification could be more nuanced. I think the claims that RS “mitigates reward misspecification” (line 75) and “If Hypothesis 2 is true, then RS can mitigate reward misspecification” (line 161) should be framed a bit more cautiously, making it clear that it’s a very particular kind of reward misspecification: when the real reward is linear in a set of proxy rewards. I don’t think this is representative of most kinds of reward misspecification that we see and that we should me worried about such as context-dependence of human preferences or biases of data workers providing feedbacks.

Questions

1. What’s the relation between scatter points and curves on the plots: are curves a smoothed based on scatter points? What smoothing technique did you use? 2. How much does the performance of RS depend on policies being close to their shared base models? How does the RS front evolve over the course of finetuning; does it degenerate after some number of gradient updates? 3. Relatedly, does using parameter-efficient finetuning (e.g. LoRA) play a role? In the paper, we claim you use LoRA only for computational efficiency reasons, but shouldn’t it also significantly help to maintain linear mode connectivity? Does RS work without LoRA?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

4 excellent

Contribution

3 good

Limitations

Limitations and societal impacts are discussed thoroughly.

Summary

This paper presents reward soups (RS) which is the idea of starting with a pre-trained network, which is finetuned to multiple proxy rewards (say, multiple different criteria), and at test time, infers a reward as a linear combination of these proxy rewards and uses this to linearly combine the corresponding weights which is then used for prediction/generation. In contrast to naive variants of multi-objective RL which trains many different policies (far greater than the number of proxy rewards) to obtain a high fidelity policy for preferences encountered at test time, the proposed approach ends up working while training a number of policies that is equal to the number of proxy rewards while showing reasonable empirical performance. ==> post rebuttal: updated score.

Strengths

The problem formation and proposed approach are topics of increasing interest and relevance to the community. The paper presents interesting results for many practically relevant and useful benchmarks.

Weaknesses

- There are not much comparisons to approaches in multi-objective RL, which makes it unclear as to how one can imagine this paper's result to be approaching notions of pareto optimal trade-offs. While the paper acknowledges this issue, it leaves open huge gaps as to what can be achieved using a suite of approaches that exist in multi-objective RL (for e.g. even starting from reference 130 cited in this paper).

Questions

Something that appears unclear is that linear mode connectivity etc. typically make sense in the context of supervised learning when training with the same set of labels. Here, the paper attempts to argue a similar perspective with different sets of labels, and in the RLHF context. Attempting to argue about both of these (particularly the latter) is non-trivial sense saying that a linear combination of parameters of a policy network optimizes the long term (i.e. generation level) linear combination of reward posits very strong assumptions on the structure of the optimal policy etc. for which I do not yet see a clear argument in this paper.

Rating

4: Borderline reject: Technically solid paper where reasons to reject, e.g., limited evaluation, outweigh reasons to accept, e.g., good evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

2 fair

Presentation

3 good

Contribution

2 fair

Limitations

yes.

Summary

In this paper, the authors propose a multi-policy strategy called "rewarded soups" to fine-tune any foundation model, embracing the heterogeneity of diverse rewards. The method combines multiple networks through linear interpolation in the weight space, despite the non-linearities in the network, which efficiently yields Pareto-optimal solutions after training. The authors demonstrate the effectiveness of the approach for text-to-text, text-image, and locomotion control tasks, showing that "rewarded soups" can mitigate reward misspecification. The proposed approach aims to enhance the alignment of deep models and how they interact with the world in all its diversity. The authors highlight the issue of aligning AI systems to specific and diverse needs while making the process more transparent and limiting the cultural hegemony of a few individuals.

Strengths

This paper addresses the reward misspecification problem caused by single priori rewards in current RLHF frameworks for foundation models. In order to solve the problem, this paper proposes a relatively complete framework called rewarded soup (RS). RS combines multiple networks (fine-tuning on different proxy rewards) through linear interpolation in the weight space and selects relative coefficients according to the user’s preferences, yielding Pareto-optimal solutions. The content of the whole paper is complete, and the experiments are sufficient.

Weaknesses

1. The writing logic of the article could be more coherent in the reviewer's opinion. For example, in 3.3, "Moreover, RS gives a better front than MORL, validating Hypothesis 2." and in 3.5, "Moreover, the front defined by RS indicates an effective balance between risk-taking and cautiousness, providing empirical support for Hypothesis 2, although MORL with $\mu$ = 0.5 (i.e., $\alpha$ = 0.5) slightly surpasses RS's front." 2. Validations on Hypothesis 1 and 2 are based on empirical results. The authors are encouraged to give some theoretical analysis. 3. Lack of experiments on computational costs. According to the results, the proposed method has no strengths compared to MORL. The authors state that RS can reduce computational costs. However, no relative experiments showed in this paper. 4. Most experiments are assigned an N=2 reward model, which is not aligned enough with the "diverse" in the title.

Questions

1. Why does the front passing through the point obtained by MORL fine-tuning on the average of the two rewards support Hypothesis 2 in Figure 2? 2. Some formulas need clarification. For example, is $\lbrace\lambda_i\rbrace_{i}$ equals to $\lbrace\lambda_i\rbrace_{i=1}^N$? 3. $\lbrace\lambda_i\rbrace_{i=1}^N$ is selected by users according to their preferences. How do users select these coefficients? For example, the user will give a preference label over pair-wise (or k-wise) instances in the standard RLHF paradigm.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

2 fair

Contribution

2 fair

Limitations

The authors have adequately addressed the limitations.

Summary

This paper explores a model-soup strategy to efficiently adapt to diverse reward functions from various real-world users. By fine-tuning a pre-trained LLM multiple times each with a specialized reward function and interpolating their weights linearly, the proposed method is able to adapt to various reward functions without having to train a new LLM per user.

Strengths

- The proposed method is much more efficient than the baselines, which have to train a separate model when a new reward function is given. - The evaluation is through and conducted on a variety of LLM tasks. - The performance of the proposed method does not fall behind compared to the more costly baselines.

Weaknesses

- The novelty of the paper is weak. It seems the main contribution of the paper is applying the weight interpolation (model soup [1]) technique, which was well-explored in supervised learning, to RLHF. I suggest the authors clarify the paper's novelty (other than applying the model-soup technique to another domain) more clearly. - The authors point out that in RLHF, the reward function is different per each model, unlike supervised learning where the training objective is the same. I agree with the authors on this point, but the rewards used in the paper's experiments do not seem very heterogeneous to back up the authors' claim. A more realistic scenario would be to experiment with a set of rewards that contradict each other directly (e.g., different reward functions learned from users with conflicting interests). - The main strength of the proposed method is that it is more efficient in terms of training (fine-tuning) cost and inference cost. However, the paper does not provide a quantitative comparison of these costs between the proposed method and the baselines. Therefore, it is hard to assess how much efficiency gain the proposed method will provide. [1] Wortsman et al., Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time, ICML 2022.

Questions

Noted above.

Rating

4: Borderline reject: Technically solid paper where reasons to reject, e.g., limited evaluation, outweigh reasons to accept, e.g., good evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

2 fair

Contribution

2 fair

Limitations

Noted above.

Summary

The authors present a new strategy to address the heterogeneity of diverse rewards in reinforcement learning. Specifically, they propose 'rewarded soup,' which involves individually training multiple networks, each assigned to a different proxy reward, and then linearly combining these networks. Compared to the multi-objective reinforcement learning baselines, the proposed rewarded soup demonstrates its superiority on several benchmarks, including text-to-text, text-to-image, and control benchmarks.

Strengths

- The authors presented a comprehensive study on the topic and the research field, as most arguments in introduction are supported by some references. The motivation is strongly supported and the path to the proposed method is reasonable. - Presentation is clear, concise, and easy-to-understand, and the idea is simple yet effective. - The authors conducted extensive experiments to verify the effectiveness of the proposed rewarded soup, and this includes multiple text-to-text tasks (shown in Section 3.1), image-to-text tasks (Section 3.2), text-to-image tasks (Section 3.3), and control tasks (Section 3.5). For most of the experiments, the improvement against MORL is obvious.

Weaknesses

Although there are many strengths in the paper, there is a weakness that can be further enhance the overall quality. - Ablation studies could be added: While the authors presented many reinforcement learning applications and also showed the improvement, it would be better to include more fine-grained ablation studies, such as how the difference of rewards affects the effectiveness of MORL baseline and the rewarded soup or how the number of networks affects the results.

Questions

- This is aligned with the weakness part. I am curious about the effectiveness of the rewarded soup under different scenarios: 1. How does the difference/gap of rewards affect the effectiveness of the MORL baseline and the rewarded soup? 2. How does the number of networks affect the results? The reason why I am particularly interested in these two questions is that a prior work [1] indicates when the models are quite different (based on their objectives), the linear combination is likely to produce less favorable results. I am wondering about "the limit" of the rewarded soup and in what situations the proposed method might fail. [1] "Robust fine-tuning of zero-shot models" Wortsman et al., CVPR 2022

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The authors have provided sufficient information of limitations and societal impact in the paper.

Summary

This manuscript studies a way to interpolate trained networks' parameters for diverse rewards in a reinforcement learning manner. To be specific, the proposed method introduces a way to achieve Pareto-optimal solutions through linearly weighted parameters after training. Extensive experiments showed the effectiveness of the proposed method on various domains such as text2text, image2text, and text2image.

Strengths

- The manuscript is well-written and organized overall. - The proposed idea is effective yet efficient as it does not require additional training. - Extensive experimental results demonstrate the effectiveness of the proposed method on various domains; text generation, image captioning, and diffusion model.

Weaknesses

Interpolating weights for better performance is not a new concept; model soups, which is mentioned in the manuscript. However, the authors did not provide any comparison with it. When we have N fine-tuned models, rewarded soup can perform better than model soup? For example, in the case of an image captioning task, the experimental setup assumes only two rewarded models, differently fined-tuned models on AVA and cafe datasets, respectively. I think there is no reason to hesitate to apply a way of model soup.

Questions

I already mentioned above in the weakness section. I would like to ask the authors how the proposed method is significantly better than other weight interpolation methods like model soup, not only the reward interpolation method, MoRL.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

I think there is no potential negative societal impact.

Summary

This paper proposes a method of using linear interpolated weight finetuned on different rewards instead of using linear combination of rewards to finetune weight, which is a solution to applying model under different and multiple preference scenarios. The idea is intuitive but works well, it (the RS) can achieve similar performance compared with MORL, while RS reduces the computational cost significantly. The author makes some mathematical hypotheses which empirically hold when all the weight to be interpolated is finetuned from a pretrained model. Also, the author has done a lot of experiments to show the feasibility of the proposed method.

Strengths

The paper is clearly written. The experimental work is sufficiently done. The author did a lot of experiments on different task which shows that this interpolating strategy is universal under different application scenarios, while with good performance. The proposed method reduces the heavy computational requirements for pretuning compared with previous work, which makes it much more applicable and flexible to complex application scenarios.

Weaknesses

Novelty: The most heavy workload in this paper is applying the strategy to different tasks and testing their performance, while less novel concepts or theories is presented. Condition for Hypothesis: The hypothesis used in the method, such as the linear mode connectivity which states that the combined reward is concave to the weights requires research on the model structure and activation function. There should be some limitations for the design of networks which ensures that the hypothesis holds.

Questions

Does the method harm the absolute quality of the produced samples? Evaluation other than reward functions should be provided. May be you can show some extrapolated examples generated through this method.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

All right in total.

Response to author's rebuttal

Thank you for your response. I understand the relative merits offered by a simple strategy such as rewarded soups; while there are advantages to RS from the perspective of issues such as alignment, re-training, forgetting etc. that appear to be interesting, I was looking to understand what this method offered compared to an RL solution, and that hasn't been adequately addressed by the rebuttal. 1. I find notions of approximate optimality to not be well qualified without understanding what is achievable. This can be obtained through benchmarking some multi-objective RL methods which are directly related to the contributions of this paper. Since the method described in the paper is obviously simple, which is a positive, it will obviously serve the interest of the broader community to attempt to benchmark MORL methods from the literature. If this still is prohibitively infeasible as the authors appear to indicate, consider two possibilities: (a) I think running an RL method on a few grid points as a reward function and understanding what fraction of headroom is left out using this strategy of reward soups is necessary to round up on how good the proposed approach is. With regards to results that are presented in the paper in this front, why does running RL on the scalarized reward sometimes appear to be inferior to RS when RL can pretty much realize the RS solution? (b) It will still be interesting to understand rewarded soups a bit more on locomotion tasks, where there is broad precedent of running multi-objective RL methods Please note that RL is one of the central approaches that has helped improve alignment, and this disconnect with RL literature is a gap that I believe falls well within this paper's scope to be addressed. If the authors can address this comment (for instance, as they have presented results with (a)), I am open to improving my score. 2. Thanks for pointing to this. Again, can the authors clarify how sub-optimal the RS policy will be as a function of generation length and how this is captured by the bound?