SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Alexandre Rame $^{1*}$ , Guillaume Couairon $^{1,2\dagger}$ , Mustafa Shukor $^{1\dagger}$ , Coretin Dancette $^{1\dagger}$ , Jean-Baptiste Gaya $^{1,2\dagger}$ , Laure Soulier $^{1}$ , Matthieu Cord $^{1,3}$ , Sorbonne Université, CNRS, ISIR, Paris, France  $^{2}$ Meta AI  $^{3}$ Valeo.ai
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), and control (locomotion) tasks. We hope to enhance the alignment of deep models, and how they interact with the world in all its diversity.
Foundation models [1] have emerged as the standard paradigm to learn neural networks' weights. They are typically first pre-trained through self-supervision [2, 3, 4, 5] and then fine-tuned [6, 7] via supervised learning [8]. Yet, collecting labels is expensive, and thus supervision may not cover all possibilities and fail to perfectly align [9, 10, 11] the trained network with the intended applications. Recent works [12, 13, 14] showed that deep reinforcement learning (DRL) helps by learning from various types of rewards. A prominent example is reinforcement learning from human feedback (RLHF) [12, 15, 16, 17], which appears as the current go-to strategy to refine large language models (LLMs) into powerful conversational agents such as ChatGPT [13, 18]. After pre-training on next token prediction [19] using Web data, the LLMs are fine-tuned to follow instructions [20, 21, 22] before reward maximization. This RL strategy enhances alignment by evaluating the entire generated sentence instead of each token independently, handling the diversity of correct answers and allowing for negative feedback [23]. Similar strategies have been useful in computer vision (CV) [14, 24], for instance to integrate human aesthetics into image generation [25, 26, 27].
Diversity of proxy rewards. RL is usually seen as more challenging than supervised training [28], notably because the real reward—ideally reflecting the users' preferences—is often not specified at training time. Proxy rewards are therefore developed to guide the learning, either as hand-engineered metrics [29, 30, 31] or more recently in RLHF as models trained to reflect human preferences
*Project lead, main contributor, correspondence to alexandre.rame@isir.upmc.fr.
Equal experimental contribution, order determined at random.
Further information and resources related to this project can be found on this website.
(a) Illustration of our proposed rewarded soup (RS).
(b) LLaMA RLHF for summarization.
Figure 1: Figure 1(a) details the different steps in rewarded soup. After unsupervised pre-training and supervised fine-tuning, we launch  $N$  independent RL fine-tunings on the proxy rewards  $\{R_i\}_{i=1}^N$ . Then we combine the trained networks by interpolation in the weight space. The final weights are adapted at test time by selecting the coefficient  $\lambda$ . Figure 1(b) shows our results (extended in Figure 2(a)) with LLaMA-7b [44] instruct fine-tuned on Alpaca [22], when RL fine-tuning for news summarization [12] with  $N = 2$  reward models assessing diverse preferences of summaries. With only two trainings ( $R_1$  and  $R_2$  rewarded on Figure 1(b)), the  $\lambda$ -interpolation  $(0 \leq \lambda \leq 1)$  reveals the green front of Pareto-optimal solutions, i.e., that cannot be improved for one reward without sacrificing the other. RS matches the costly yellow front of multi-objective (MORL) [45, 46] requiring multiple trainings on different linear weightings over the rewards  $(1 - \mu) \times R_1 + \mu \times R_2$  with  $0 \leq \mu \leq 1$ .
[15, 32, 33]. Nonetheless, designing reliable proxy rewards for evaluation is difficult. This reward misspecification [9, 34] between the proxy reward and the users' actual rewards can lead to unforeseen consequences [35]. Moreover, the diversity of objectives in real-world applications complicates the challenge. In particular, human opinions can vary significantly [36, 37, 38] on subjects such as aesthetics [39], politics or fairness [40]. Humans have also different expectations from machines: for example, while [41] stressed aligning LLMs towards harmless feedback, [42] requested helpful non-evasive responses, and others' [43] interests are to make LLMs engaging and enjoyable. Even hand-engineered metrics can be in tension: generating shorter descriptions with higher precision can increase the BLEU [29] score but decrease the ROUGE [30] score due to reduced recall.
Towards multi-policy strategies. Considering these challenges, a single model cannot be aligned with everyone's preferences [13]. Existing works align towards a consensus-based user [47, 48], relying on the "wisdom of the crowd" [49], inherently prioritizing certain principles [42, 50], resulting in unfair representations of marginalized groups [51, 52]. The trade-offs [53] are decided a priori before training, shifting the responsibility to the engineers, reducing transparency and explainability [54], and actually aligning towards the "researchers designing the study" [13, 55]. These limitations, discussed in Appendix A.1, highlight the inability of single-policy alignment strategies to handle human diversity. Yet, "human-aligned artificial intelligence is a multi-objective problem" [56]. Thus, we draw inspiration from the multi-objective reinforcement learning (MORL) literature [45, 46, 57, 58, 59, 60, 61, 62] and [54]; they argue that tackling diverse rewards requires shifting from single-policy to multi-policy approaches. As optimality depends on the relative preferences across those rewards, the goal is not to learn a single network but rather a set of Pareto-optimal networks [63].
In this paper, we propose rewarded soup (RS), an efficient and flexible multi-policy strategy to fine-tune any foundation model. As shown in Figure 1(a), we first use RL to learn one network for each proxy reward; then, we combine these expert networks according to user preferences. This a posteriori selection allows for better-informed trade-offs, improved transparency and increased fairness [54, 64]. The method to combine those networks is our main contribution: we do this through linear interpolation in the weight space, despite the non-linearities in the network. This is in line with recent findings on linear mode connectivity (LMC) [65, 66]: weights fine-tuned from a shared pre-trained initialization remain linearly connected and thus can be interpolated. This LMC inspired a plethora of weight interpolation (WI) strategies [67, 68, 69, 70, 71, 72], discussed in Section 4. Actually, the name rewarded soups follows the terminology of model soups [67], as we combine various ingredients each rewarded differently. Unlike previous works, which focused on supervised learning, we explore LMC in RL, in a challenging setup where each training run uses a different reward. Perhaps surprisingly, we show that we can trade off the capabilities of multiple weights in a
single final model, thus without any computational overhead. This enables the creation of custom weights for any preference over the diverse rewards. We summarize our contributions as follows:
- We advocate a multi-policy paradigm to align deep generative models with human preferences and reduce reward misspecification.
- We then propose a new multi-policy strategy, rewarded soup, possible when fine-tuning foundation models with diverse rewards. By weight interpolation, it defines a continuous set of (close to) Pareto-optimal solutions, approximating more costly multi-policy strategies.
In Section 3, we consistently validate the linear mode connectivity and thus the effectiveness of RS across a variety of tasks and rewards: RLHF fine-tuning of LLaMA, multimodal tasks such as image captioning or text-to-image generation with diffusion models, as well as locomotion tasks.
2 Rewarded soups
2.1 RL fine-tuning with diverse rewards
We consider a deep neural network  $f$  of a fixed non-linear architecture (e.g., with batch normalization [73], ReLU layers [74] or self-attention [75]). It defines a policy by mapping inputs  $x$  to  $f(x, \theta)$  when parametrized by  $\theta$ . For a reward  $\hat{R}$  (evaluating the correctness of the prediction according to some preferences) and a test distribution  $T$  of deployment, our goal is to maximize  $\int_{x \in T} \hat{R}(f(x, \theta))$ .
For example, with  $f$  a LLM,  $x$  would be textual prompts,  $\hat{R}$  would evaluate if the generated text is harmless [76], and  $T$  would be the distribution of users' prompts. Learning the weights  $\theta$  is now commonly a three-step process: unsupervised pre-training, supervised fine-tuning, and reward optimization. Yet  $\hat{R}$  is usually not specified before test time, meaning we can only optimize a proxy reward  $R$  during training. This reward misspecification between  $R$  and  $\hat{R}$  may hinder the alignment of the network with  $\hat{R}$ . Moreover, the diversity of human preferences complicates the design of  $R$ .
Rather than optimizing one single proxy reward, our paper's first key idea is to consider a family of  $N$  diverse proxy rewards  $\{R_i\}_{i=1}^N$ . Each of these rewards evaluates the prediction according to different (potentially conflicting) criteria. The goal then becomes obtaining a coverage set of policies that trade-off between these rewards. To this end, we first introduce the costly MORL baseline. Its inefficiency motivates our rewarded soups, which leverages our second key idea: weight interpolation.
MORL baseline. The standard MORL scalarization strategy [45, 46] (recently used in [62] to align LLMs) linearizes the problem by interpolating the proxy rewards using  $M$  different weightings. Specifically, during the training phase,  $M$  trainings are launched, with the  $j$ -th optimizing the reward  $\sum_{i=1}^{N} \mu_i^j R_i$ , where  $\forall j \in \{1, \dots, M\}$ ,  $\{\mu_i^j\}_{i=1}^N \in \Delta_N$  the  $N$ -simplex s.t.  $\sum_{i=1}^{N} \mu_i^j = 1$  and  $0 \leq \mu_i^j \leq 1$ . Then, during the selection phase, the user's reward  $\hat{R}$  becomes known and the  $j$ -th policy that maximizes  $\hat{R}$  on some validation dataset is selected. We typically expect to select  $j$  such that  $\sum_{i=1}^{N} \mu_i^j R_i \approx \hat{R}$  linearly approximates the user's reward. Finally, this  $j$ -th weight is used during the inference phase on test samples. Yet, a critical issue is that "minor [preference] variations may result in significant changes in the solution" [77]. Thus, a high level of granularity in the mesh of  $\Delta_N$  is necessary. This requires explicitly maintaining a large set of  $M \gg N$  networks, practically one for each possible preference. Ultimately, this MORL strategy is unscalable in deep learning due to the computational, memory, and engineering costs involved (see further discussion in Appendix A.2).
Rewarded soup (RS). In this paper, we draw inspiration from the weight interpolation literature. The idea is to learn expert weights and interpolate them linearly to combine their abilities. Specifically, we propose RS, illustrated in Figure 1(a) and whose recipe is described below. RS alleviates MORL's scaling issue as it requires only  $M = N$  trainings while being flexible and transparent.
1. During the training phase, we optimize a set of  $N$  expert weights  $\{\theta_i\}_{i=1}^N$ , each corresponding to one of the  $N$  proxy rewards  $\{R_i\}_{i=1}^N$ , and all from a shared pre-trained initialization.
2. For the selection phase, we linearly interpolate those weights to define a continuous set of rewarded soups policies:  $\{\sum_{i=1}^{N} \lambda_i \cdot \theta_i\}_{\{\lambda_i\}_{i=1}^N \in \Delta_N}$ . Practically, we uniformly sample  $M$
interpolating coefficients  $\{\{\lambda_i^j\}_{i = 1}^N\}_{j = 1}^M$  from the  $N$ -simplex  $\Delta_N$  and select the  $j$ -th that maximizes the user's reward  $\hat{R}$  on validation samples, i.e.,  $\operatorname*{argmax}_{j = 1}^{M}\hat{R}\Big(\sum_{i = 1}^{N}\lambda_{i}^{j}\theta_{i}\Big)$ .
3. For the inference phase, we predict using the network  $f$  parameterized by  $\sum_{i=1}^{N} \lambda_i^j \theta_i$ .
While MORL interpolates the rewards, RS interpolates the weights. This is a considerable advantage as the appropriate weighting  $\lambda$ , which depends on the desired trade-off, can be selected  $a$  posteriori; the selection is achieved without additional training, only via inference on some samples. In the next Section 2.2 we explicitly state the Hypotheses 1 and 2 underlying in RS. These are considered Working Hypotheses as they enabled the development of our RS strategy. Their empirical verification will be the main motivation for our experiments on various tasks in Section 3.
2.2 Exploring the properties of the rewarded soups set of solutions
2.2.1 Linear mode connectivity of weights fine-tuned on diverse rewards
We consider  $\{\theta_i\}_{i=1}^N$  (or  $\{\theta_i\}_i$  for brevity) fine-tuned on  $\{R_i\}_i$  from a shared pre-trained initialization. Previous works [65, 66, 67, 72] defined linear mode connectivity (LMC) w.r.t. a single performance measure (e.g., accuracy or loss) in supervised learning. We extend this notion in RL with  $N$  rewards, and define that the LMC holds if all rewards for the interpolated weights exceed the interpolated rewards. It follows that the LMC condition which underpins RS's viability is the Hypothesis 1 below.
Working Hypothesis 1 (LMC).  $\forall \{\lambda_i\}_i \in \Delta_N$  and  $k \in \{1, \dots, N\}$ ,  $R_k(\sum_i \lambda_i \cdot \theta_i) \geq \sum_i \lambda_i R_k(\theta_i)$ .
2.2.2 Pareto optimality of rewarded soups
The Pareto front (PF) is the set of undominated weights, for which no other weights can improve a reward without sacrificing another, i.e.,  $\{\theta \mid \not\in\theta' \in \Theta$  s.t.  $\{R_i(\theta')\}_i > _N\{R_i(\theta)\}_i\}$  where  $>_N$  is the dominance relation in  $\mathcal{R}^N$ . In practice, we only need to retain one policy for each possible value vector, i.e., a Pareto coverage set (PCS). We now introduce the key Hypothesis 2, that state the Pareto-optimality of the solutions uncovered by weight interpolation in RS.
Working Hypothesis 2 (Pareto optimality). The set  $\{\sum_{i}\lambda_{i}\cdot \theta_{i}|\{\lambda_{i}\}_{i}\in \Delta_{N}\}$  is a  $PCS$  of  $\{R_i\}_i$
Empirically, in Section 3, we consistently validate Hypotheses 1 and 2. Theoretically, in Appendix C.2, we prove they approximately hold, in a simplified setup (quadratic rewards with co-diagonalizable Hessians) justifiable when weights remain close.
Remark 1. Hypotheses 1 and 2 rely on a good pre-trained initialization, making RS particularly well-suited to fine-tune foundation models. This is because pre-training prevents the weights from diverging during training [66]. When the weights remain close, we can theoretically justify Hypotheses 1 and 2 (see Appendix C.2) and, more broadly, demonstrate that WI approximates assembling [78, 79] (see Lemma 4). In contrast, the LMC does not hold when training from scratch [66]. Neuron permutations strategies [80, 81] tried to enforce connectivity by aligning the weights, though (so far) with moderate empirical results: their complementarity with RS is a promising research avenue.
Remark 2. Pareto-optimality in Hypothesis 2 is defined w.r.t. a set of possible weights  $\Theta$ . Yet, in full generality, improvements in initialization, RL algorithms, data, or specific hyperparameters could enhance performances. In other words, for real-world applications, the true  $PF$  is unknown and needs to be defined w.r.t. a training procedure. In this case,  $\Theta$  represents the set of weights attainable by fine-tuning within a shared procedure. As such, in Section 3 we analyze Hypothesis 2 by comparing the fronts obtained by RS and scalarized MORL while keeping everything else constant.
2.2.3 Consequences of Pareto optimality if the user's reward is linear in the proxy rewards
Lemma 1 (Reduced reward misspecification in the linear case). If Hypothesis 2 holds, and for linear reward  $\hat{R} = \sum_{i}\hat{\mu}_{i}R_{i}$  with  $\{\hat{\mu}_i\}_i\in \Delta_N$ , then  $\exists \{\lambda_i\}_i\in \Delta_N$  such that  $\sum_{i}\lambda_{i}\cdot \theta_{i}$  is optimal for  $\hat{R}$ .
The proof outlined in Appendix C.1 directly follows the definition of Pareto optimality. In simpler terms, Lemma 1 implies that if Hypothesis 2 holds, RS mitigates reward misspecification for linear rewards: for any preference  $\hat{\mu}$ , there exists a  $\lambda$  such that the  $\lambda$ -interpolation over weights maximizes the  $\hat{\mu}$ -interpolation over rewards. In practice, as we see in Figure 5(a), we can set  $\lambda = \hat{\mu}$ , or cross-validate  $\lambda$  on other samples.
In this section we implement RS across a variety of standard learning tasks: text-to-text generation, image captioning, image generation, visual grounding, visual question answering, and locomotion. We use either model or statistical rewards. We follow a systematic procedure. First, we independently optimize diverse rewards on training samples. For all tasks, we employ the default architecture, hyperparameters and RL algorithm; the only variation being the reward used across runs. Second, we evaluate the rewards on the test samples: the results are visually represented in series of plots. Third, we verify Hypothesis 1 by examining whether RS's rewards exceed the interpolated rewards. Lastly, as the true Pareto front is unknown in real-world applications, we present empirical support for Hypothesis 2 by comparing the front defined by RS (sliding  $\lambda$  between 0 and 1) to the MORL's solutions optimizing the  $\mu$ -weighted rewards (sometimes only  $\mu = 0.5$  for computational reasons). Implementations are released on github, and this website provides additional qualitative results.
3.1 Text-to-text: LLaMA with diverse RLHFs
Given the importance of RLHF to train LLMs, we begin our experiments with text-to-text generation. Our pre-trained network is LLaMA-7b [44], instruction fine-tuned [20, 83] on Alpaca [22]. For RL training with PPO [84], we employ theTRL package [85] and the setup from [86] with low-rank adapters (LoRA) [87] for efficiency. We first consider summarization [12, 17] tasks on two datasets: Reuter news [88] in Figures 1(b) and 2(a) and Reddit TL;DR [89] in Figure 2(b). We also consider answering Stack Exchange questions [90] in Figure 2(c), movie review generation in Figure 2(d), and helpfulness as a conversational assistant [49] in Figures 2(e) and 2(f). To evaluate the generation in
(a) News summary.
(b) Reddit TL;DR summary.
(c) Stack Exchange Q&A.
(d) Useful movie review.
(e) Helpful assistant.
(f) Helpful assistant: spider map.
Figure 2: RLHF results in NLP with LLaMA-7b [44] and reward models  $R_{i}$  from HuggingFace [82]. The blue line reports checkpoints' results along the training trajectory of  $\theta_{1}$  rewarding  $R_{1}$ , the red line  $\theta_{2}$  rewarding  $R_{2}$ , and the purple line the MORL rewarding  $\frac{R_1 + R_2}{2}$ . Our rewarded soup (RS) linearly interpolates between the weights  $\theta_{1}$  and  $\theta_{2}$ ; sliding the interpolation coefficient  $\lambda$  from 0 to 1 reveals the green solid front of rewarded soups solutions. In Figures 2(a) and 2(b), we additionally show the multiple MORL runs rewarding  $(1 - \mu) \times R_{1} + \mu \times R_{2}$  with preferences  $0 \leq \mu \leq 1$ . It reveals a similar yellow front, yet more costly. In Figure 2(f), we uniformly  $(\lambda_i = \frac{1}{4})$  average the weights fine-tuned for the assistant task on  $N = 4$  reward models.
the absence of supervision, we utilized  $N = 2$  different reward models (RMs) for each task, except in Figure 2(f) where  $N = 4$ . These RMs were trained on human preferences datasets [15] and all open-sourced on HuggingFace [82]. For example in summarization,  $R_{1}$  follows the "Summarize from Human Feedback" paper [12] and focuses on completeness, while  $R_{2}$  leverages "contrast candidate generation" [91] to evaluate factuality. For other tasks, we rely on diverse RMs from OpenAssistant [92]; though they all assess if the answer is adequate, they differ by their architectures and procedures. Table 1 details the experiments.
The results are reported in Figure 2. The green front, defined by RS between the two weights specialized on  $R_{1}$  and  $R_{2}$ , is above the straight line connecting those two points, validating Hypothesis 1. Second, the front passes through the point obtained by MORL fine-tuning on the average of the two rewards, supporting Hypothesis 2. Moreover, when comparing both full fronts, they have qualitatively the same shape; quantitatively in hypervolume [93] (lower is better, the area over the curve w.r.t. an optimal point), RS's hypervolume is 0.367 vs. 0.340 for MORL in Figure 2(a), while it is 1.176 vs. 1.186 in Figure 2(b). Finally, in Figure 2(f), we use  $N = 4$  RMs for the assistant task and uniformly average the  $N = 4$  weights, confirming that RS can scale and trade-off between more rewards.
3.2 Image-to-text: captioning with diverse statistical rewards
RL is also effective for multimodal tasks [14] such as in image captioning [24], to generate textual descriptions of images. Precisely evaluating the quality of a prediction w.r.t. a set of human-written
(a) BLEU1 and ROUGE.
(b) BLEU1 and BLEU4.
(c) Captioning: spider map.
Figure 3: Results in image captioning on COCO [94]. As rewards  $R_{1}$  (blue stars every epoch) and  $R_{2}$  (red stars), we consider standard statistical metrics: BLEU1 (1-gram overlap), BLEU4 (4-grams overlap), ROUGE, METEOR and CIDEr. Figure 3(a) include the MORL training trajectories optimizing  $(1 - \mu) \times BLEU1 + \mu \times ROUGE$ , uncovering a yellow front similar to RS's green front. In Figure 3(c), RS uniformly averages the 5 weights (one for each reward), resulting in the largest area and the best trade-off between the 5 rewards.
(a) BLEU1 first.
(b) BLEU4 first.
(c) ROUGE first.
Figure 4: Those spider maps uniformly average  $1 \leq M \leq 5$  weights for captioning, where  $\theta_{1}$  is fine-tuned on BLEU1 (B1),  $\theta_{2}$  on BLEU4 (B4),  $\theta_{3}$  on ROUGE (R),  $\theta_{4}$  on METEOR (M) and  $\theta_{5}$  on CIDEr (C). To show different combinations among the  $\binom{5}{M}$  possible, we iterate in a clockwise direction starting in Figure 4(a) from  $i = 1$  (always including  $\theta_{1}$  optimized on BLEU1), in Figure 4(b) from  $i = 2$  (always including  $\theta_{2}$  optimized on BLEU4), and in Figure 4(c) from  $i = 3$  (always including  $\theta_{3}$  optimized on ROUGE).
(a) Analysis of Lemma 1.
(b) Performances w.r.t.  $\lambda$
(c) RS vs. prediction ensembling.
Figure 5: Results in captioning for  $R_{1} = BLEU1$  and  $R_{2} =ROUGE$ . When normalized, rewards are set to 1 for the init and 0 for the worst model. Figure 5(a) validates Lemma 1 by reporting results of RS (for varying  $\lambda$ ) and of MORL (for varying  $\mu$ ) for varying user's preference  $\hat{\mu}$ . Figure 5(b) evaluates different rewards as a function of the interpolating coefficient. Figure 5(c) reports ensembling scores when interpolating predictions.