Iterative Foundation Model Fine-Tuning on Multiple Rewards

Thank you for taking the time to review our paper and sharing your valuable thoughts and suggestions. Please find our responses to your comments below.

We would like to clarify that the proposed algorithm is not a combination of MORLHF and RS. The paper provides a new perspective that MORLHF can be viewed as the case where merging is done after every batch. However, if merging happens after multiple batches, the algorithm after each merging is no longer equivalent to MORLHF.

Experiments

We performed further experiments on the text summarization task. To improve the performance of the proposed algorithm, we considered seven different sets of merging weights, including [1/3, 1/3, 1/3], permutations of [1/6, 1/6, 2/3], and permutations of [1/6, 5/12, 5/12]. During training, at each merging step, we calculated the reward of each merged model on the training data and chose the one with the highest average reward for the next iteration. To obtain the final merged model after training, we measured the average reward of all seven merged models and chose the one with the highest average reward on the validation data. We then tested the performance on the test set. We followed the same procedure for RS as well. Note that RS does not perform merging during training. The table below shows the new results. These results demonstrate the effectiveness of the flexible merging strategy in IterativeRS, as it significantly outperforms other methods on all metrics.

Diversity and GFlowNet

The focus of this paper is on reinforcement learning for LLM fine-tuning. Due to pre-training on vast amounts of unlabeled data, LLMs are known to generate diverse outputs. During fine-tuning, we monitor the KL divergence between the policy and a reference policy to ensure that the fine-tuned model does not deviate too far from the initial pre-trained model. For example, we set kl_coef = 0.05 during PPO fine-tuning.

Sequence generation with GFlowNet can be viewed as a soft Q-learning problem [15]. However, applying Q-learning to LLM fine-tuning may not be computationally feasible. Therefore, we did not consider employing or comparing against existing GFlowNet methods. Note that in all experiments, we employ LLMs for sequence generation.

Notation and Presentation:

Since Lemma 1 is taken from [19], we wanted to retain its original notation. We understand the confusion this may have caused. To address this, we will unify the notation of Lemma 1 with the rest of the paper and add its proof to the appendix. Furthermore, based on your suggestion, we will include a notation table at the beginning of the appendix.

Please let us know if you have further comments and concerns.

We would like to express our gratitude for your valuable comments and suggestions. Please find our responses to your comments below.

Theoretical Analysis

The conclusions drawn from the theoretical analysis in the paper are not guaranteed to hold for non-convex deep neural networks, which are typically used in practice. However, the goal of the theoretical analysis is not to fully explain the behavior of the proposed algorithm in practical scenarios. Instead, it demonstrates that the algorithm’s design is both sensible and grounded in a rigorous theoretical foundation. Finally, the theoretical convergence analysis of deep neural networks used in practice remains an open research problem and is beyond the scope of this paper. It is also useful to mention that assumptions A1–A3 are common in convex optimization.

Complexity

The worst-case computational and memory complexities of IterativeRS are of the same order as those of RS. Both algorithms update the model over $T$ batches. IterativeRS performs merging every $m$ rounds; however, the cost of merging is negligible compared to expert training, as it typically does not involve additional training. Furthermore, as described in line 7 of Algorithm 1, between two merging rounds, a subset of objectives can be selected for fine-tuning. This can significantly reduce the complexity of IterativeRS, especially when compared to RS. Setting $M=1$ results in a complexity comparable to that of MORLHF.

Experiments

We performed further experiments on the text summarization task. To improve the performance of the proposed algorithm, we considered seven different sets of merging weights, including [1/3, 1/3, 1/3], permutations of [1/6, 1/6, 2/3], and permutations of [1/6, 5/12, 5/12]. Previously, we only considered the uniform merging weights [1/3, 1/3, 1/3]. During training, at each merging step, we calculated the reward of each merged model on the training data and chose the one with the highest average reward for the next iteration. To obtain the final merged model after training, we measured the average reward of all seven merged models and chose the one with the highest average reward on the validation data. We then tested the performance on the test set. We followed the same procedure for RS as well. Note that RS does not perform merging during training. The table below shows the new results. These results demonstrate the effectiveness of the flexible merging strategy in IterativeRS, as it significantly outperforms other methods on all metrics.

In addition, to incorporate more standard metrics into our evaluation, we define a new metric as the average of all rewards and the ROUGE score. The results are presented in the table below. We can either replace the "Avg Reward" column in Table 3 with this new metric or update the definition of Avg Reward accordingly. As can be seen, IterativeRS achieves the highest average score.

Baselines

To the best of our knowledge, we have compared the proposed algorithm with all relevant state-of-the-art multi-objective reinforcement learning methods for LLM fine-tuning. Sequence generation with GFlowNet can be viewed as a soft Q-learning problem [15]. However, applying Q-learning to LLM fine-tuning may not be computationally feasible. Therefore, we did not consider employing or comparing against existing GFlowNet methods. Note that in all experiments, we employ LLMs for sequence generation.

IterativeRS and MolMoE have distinct and non-overlapping functionalities and applications. MolMoE is an expert merging method designed for molecular applications, whereas IterativeRS focuses on expert training. Therefore, these two methods can be used in conjunction. We applied MolMoE to both IterativeRS and RS, and the results are presented below. One of the main advantages of using MolMoE is its ability to handle scenarios where preferences over objectives change dynamically over time, which is outside the scope of this paper's experimental study.

Please let us know if you have further comments and concerns.

We would like to express our gratitude for reviewing our paper and letting us know your valuable comments. Please find below our responses to your comments.

Fine-Tuning Frameworks

The experiments are conducted using PPO, a commonly used reinforcement learning (RL) method for LLM fine-tuning. Based on your comment, we performed additional experiments on the DNA sequence generation task to evaluate the performance of RL-based fine-tuning methods MORLHF, RS, and IterativeRS, using RLOO, which does not rely on a value model. The experimental setup for RLOO is the same as that used with PPO. The results are presented in the table below. They indicate that using PPO leads to better performance for IterativeRS. It is possible that improved tuning of the hyperparameter $m$ could enhance the performance of IterativeRS when using RLOO. Furthermore, the results show that IterativeRS achieves a higher average reward than both MORLHF and RS when using RLOO.

Reinforcement Learning from Scratch

The focus of this paper is on RL for LLM fine-tuning. However, the proposed method can also be applied to RL from scratch. It is worth noting that even in relatively small-scale applications where training from scratch is feasible, fine-tuning LLMs can offer significant advantages. First, LLMs can be pre-trained on both unsupervised and supervised data, while RL from scratch utilizes labeled data. This pre-training provides LLMs with prior knowledge that can be beneficial for downstream tasks. Moreover, fine-tuning an LLM might be more computationally efficient than training an RL model from scratch. For these reasons, we chose to focus on LLM fine-tuning in this work.

Applications

There are many practical applications for multi-objective RL in LLM fine-tuning. For example, LLMs are currently used for small molecule generation in drug discovery. These models learn to generate valid SMILES representations of molecules through pre-training on large molecular datasets. During generation, the goal is to produce molecules that satisfy multiple desirable properties. Each property typically has its own labeled dataset. An effective fine-tuning approach should enable the LLM to optimize all target properties simultaneously. Similarly, in text summarization, multiple criteria such as helpfulness, conciseness, and faithfulness are important for evaluating the quality of a summary. Relying on a single reward model to capture all these criteria may be suboptimal. Therefore, multi-objective RL offers a necessary solution for a range of important tasks in LLM fine-tuning.

Thank you so much for reviewing our paper and letting us know your valuable comments. Please find below our responses to your comments.

Theory

The theoretical analysis in this paper focuses on convex settings. The statement regarding generalization in practice refers to extending the results to non-convex cases, which is the relevant scenario for LLM fine-tuning. The theoretical convergence analysis of deep neural networks remains an open research problem and is beyond the scope of this paper. The purpose of our theoretical analysis is not to fully explain the algorithm’s behavior in practical settings, but rather to show that the algorithm’s design is well-motivated and grounded in a rigorous theoretical foundation.

Furthermore, the goal of this paper is not to propose an algorithm whose superiority over RS and MORLHF is theoretically proven. Instead, the paper suggests that better empirical results than both RS and MORLHF can be achieved by appropriately tuning the hyperparameter $m$ of the proposed IterativeRS.

Metrics and Interpretation of Results

Sharp-like score: For a given sample, the Sharpe-like score is defined as the average reward across all objectives divided by the standard deviation of those rewards. A higher Sharpe-like score indicates lower variability and more balanced performance across objectives. Multiplying the reward vector by a scalar $\lambda>1$ increases the standard deviation, but the Sharpe-like score remains unchanged. In this case, the relative performance across objectives does not change, even though all objective values improve. A higher Sharpe-like score reflects more consistent performance across objectives, which is generally desirable. However, in practice, higher average performance (i.e., average reward) may be preferred over higher consistency. As an example, if two objectives remain unchanged while one objective improves, this leads to a reduction in the Sharpe-like score, while the overall performance has improved. The definition of the Sharpe-like score is provided in the first paragraph of Section 5. We can highlight the definition by providing math formulation.

Pareto-front: Let us first explain the difference between the goals in text summarization and biological applications. In text summarization, each input consists of a prompt, and the goal is to generate a high-quality summary for that prompt. In contrast, for small molecule and DNA sequence generation, there is no prompt; the objective is to discover a small number of high-performing sequences using the LLM. In our experiments, we used the fine-tuned model to generate 10,000 molecules or DNA sequences, which were then evaluated using oracle models. Note that in our setup, oracles are distinct from reward models. Reward models are trained on labeled training data, where the labels are provided by oracle models. In biological applications, if one generated molecule outperforms another (both produced by the same model) across all objectives, the former dominates the latter and is included in the Pareto front, while the dominated sample is considered suboptimal. Therefore, for molecule and DNA sequence generation, we report the performance based on the Pareto front. In contrast, for text summarization, we generate a response for each prompt in the test set and report the scores assigned by the oracle reward models. As with the biological tasks, the LLMs did not interact with the oracle models during fine-tuning for text summarization.

Figures: The comparison between RiC and IterativeRS is emphasized in the figures because RS and MORLHF can be viewed as special cases of IterativeRS. In the results shown in Figure 1, it is desirable to have samples in the top-right corner. Since the task is molecule generation, generating a few dominant samples is considered optimal. As shown in Figure 1, IterativeRS achieves the highest density of samples in the top-right corner.

Experiments

We performed further experiments on the text summarization task. To avoid relying on predefined merging weights, we considered seven different sets of merging weights, including [1/3, 1/3, 1/3], permutations of [1/6, 1/6, 2/3], and permutations of [1/6, 5/12, 5/12]. Previously, we only considered the uniform merging weights [1/3, 1/3, 1/3]. During training, at each merging step, we calculated the reward of each merged model on the training data and chose the one with the highest average reward for the next iteration. To obtain the final merged model after training, we measured the average reward of all seven merged models and chose the one with the highest average reward on the validation data. We then tested the performance on the test set. We followed the same procedure for RS as well. Note that RS does not perform merging during training. The table below shows the new results. These results demonstrate the effectiveness of the flexible merging strategy in IterativeRS, as it significantly outperforms other methods on all metrics.

Hyperparameter Tuning: To tune the hyperparameter $m$, we perform a few experiments with different values of $m$ and evaluate performance on the validation set. Based on these results, we select the value of $m$.

Please let us know if you have further comments and concerns.

Thank you for reviewing our paper and letting us know your valuable comments. Please find below our responses.

Experiments

We performed further experiments on the text summarization task. To improve the performance of the proposed algorithm, we considered seven different sets of merging weights, including [1/3, 1/3, 1/3], permutations of [1/6, 1/6, 2/3], and permutations of [1/6, 5/12, 5/12]. Previously, we only considered the uniform merging weights [1/3, 1/3, 1/3]. During training, at each merging step, we calculated the reward of each merged model on the training data and chose the one with the highest average reward for the next iteration. To obtain the final merged model after training, we measured the average reward of all seven merged models and chose the one with the highest average reward on the validation data. We then tested the performance on the test set. We followed the same procedure for RS as well. Note that RS does not perform merging during training. The table below shows the new results. These results demonstrate the effectiveness of the flexible merging strategy in IterativeRS, as it significantly outperforms other methods on all metrics.

Theory

Please note that Assumption A.3, which is the boundedness of gradients, is a standard and common assumption in convex optimization (see for example [19]). The conclusions drawn from the theoretical analysis in the paper are not guaranteed to hold for non-convex deep neural networks, which are typically used in practice. However, the goal of the theoretical analysis is not to fully explain the behavior of the proposed algorithm in practical scenarios. Instead, it demonstrates that the algorithm’s design is both sensible and grounded in a rigorous theoretical foundation. Finally, the theoretical convergence analysis of deep neural networks used in practice remains an open research problem and is beyond the scope of this paper.

Thanks for the response. Regarding the theory, I don't fully agree that the assumption A.3 is standard. In the reference [19] you mentioned, the assumption is applied to $\mathbb{E}[\|G\|^2]$ instead of $\|G\|^2$. For example, a Gaussian random variable $X$ can have a finite $\mathbb{E}[X^2]$ while $X^2$ is unbounded. This gap is not closed in this paper. However, I acknowledge that the paper's contribution is not theoretical. Therefore, I maintain my original score of 4 - Weak Accept.