Getting More Juice Out of the SFT Data: Reward Learning from Human Demonstration Improves SFT for LLM Alignment

Summary

This paper proposes an approach to aligning Large Language Models (LLMs) with human preferences by integrating Inverse Reinforcement Learning (IRL) into the supervised fine-tuning (SFT) stage. The proposed IRL-based method simultaneously builds a reward model and a policy model during the SFT stage, enhancing efficiency and robustness against low-quality data.

Strengths

1. The paper introduces a novel IRL-based method for the SFT stage, providing a fresh perspective on improving LLM alignment. 2. The paper provides a strong theoretical basis for the proposed algorithms, ensuring that they converge to stationary solutions of the IRL problem. 3. The evaluation includes both theoretical analysis and empirical testing on large-scale models, ensuring the validity of the results.

Weaknesses

1. The paper does not address the scalability of the proposed methods to even larger models or different types of LLMs beyond the ones tested. 2. The evaluation results only present the empirical performance but lack experiments on computation costs introduced by this method.

Questions

1. It seems that RFT and IRFT perform equally. What's the takeaway from these two methods? Moreover, how to choose them wisely. 2. Typo in L226, there's an error in referring to the Appendix. 3. In L204-206 "In 204 practice however we take a relatively small T and large K, because frequent on-line sampling is time 205 consuming" How does this parameter choice affect the final performance, especially compared to large T and Large K?

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

The limitations are included in the conclusion section.

Summary

This paper proposes to study if IRL techniques can be used for aligning large language models. They propose two different algorithms: one that explicitly learns a reward model and one that implicitly learns the reward in the policy. These reward models are learned by contrasting expert generations and the policy generations.

Strengths

- The paper writing and motivation is clear. - The paper includes in depth theoretical analysis. - They evaluate the trained LLMs on a wide variety of benchmarks.

Weaknesses

- The experiments are not very convincing. For the Open LLM leaderboard experiments, the gain is really small (around 1 or 2%). It seems likely that this gain is due to variance in the training process, especially since IRFT has ~1% variance in different hyperparameter settings. - There isn’t much discussion of the computational cost of the algorithms.

Questions

- How accurate is the learned reward model? - The experimental setting is not very clear. Is Algorithm 1 used on the top 10K data selected by the reward model? If so this is really problematic: shouldn’t you fine tune on the whole dataset (or a randomly downsampled subset)? Can you assume that you will have access to 10K samples with high reward scores? And wouldn't just be better to use the reward model used for data selection for RL? - How robust is IRFT to hyperparameter setting compared to SFT? - Is 1-2% increase on the OpenLLM benchmark meaningful?

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

Yes, the authors discuss limitations.

Summary

The authors propose using inverse reinforcement learning (IRL) on demonstration data in place of supervised learning, as is typically done for LLMs. The intuition is that human preferences are also encoded in demonstrations collected for SFT. Concretely, the authors propose a bilevel optimization approach with policy learning at the lower level and reward learning at the upper level. Two methods are proposed for alignment: with either implicit or explicit reward learning. The authors demonstrate improved performance when used to finetune smaller language models, compared to existing SFT approaches.

Strengths

- The method seems (to the best of my knowledge) to be novel and interesting, tying together recent LLM literature and IRL methods. - The authors compare against relevant baselines (normal SFT, or SPIN) across a suite of benchmarks and two model sizes. - The overarching idea and question the paper strives to answer is of importance and relevance to the research community, as the proposed methods enable extracting more value out of SFT data.

Weaknesses

- The experimental results seem to support the efficacy of the proposed algorithms, however the gains across various tasks with IRFT seem fairly marginal (it would help if Table 2 and 3 contained error bars, perhaps with policies finetuned with different seeds). As it stands, it seems like different benchmarks results are sensitive to different hyperparameters (choice of T), and the baselines outperform IRFT on 3/6 of the benchmarks in Table 3, suggesting that the additional complexity of the method does not necessarily lead to improved performance across all tasks. For example, the best average performance of IRFT is 61.03, whereas the best baseline achieves 61.02. - Furthermore, the results seem more mixed in Table 3 compared to Table 2, suggesting that it's not clear the method is still effective as model sizes increase. - It would help to present the SFT results in Table 3. Even though the authors note that further SFT could degrade performance, it would help to see how effectively using the demonstration data with IRFT or SPIN compares against naive SFT. - Minor note: line 157 seems incomplete, "even when the demonstration policy is"...(extreme)? - Minor note: the notation is a bit confusing, with theta referring to both the model parameters and the reward model parameters. For clarity, it would help if separate variables were used. - Minor note: missing reference to appendix in line 226. - Minor note: grammar error in line 320 -- "SPIN and IRFT are both capable of further improv(ing) the performance.."

Questions

See Weaknesses section for suggestions.

Rating

5

Confidence

3

Soundness

2

Presentation

2

Contribution

3

Limitations

The paper is missing more thorough discussion on the method’s limitations and weaknesses (e.g. sensitivity to T, more complicated training process compared to SFT, etc.).

Summary

This paper proposes two methods focusing on RLHF, namely RFT and IRFT. The takeaway message is that the SFT stage also significantly benefits from learning a reward model instead of using the human demonstration data directly via supervised learning.

Strengths

1. The paper is clear, illustrating the differences between similar works. 2. The motivation is clear, and the method is reasonable. Most importantly, it provides parts of theoretical analysis for the convergence.

Weaknesses

1. The training procedure is complicated, meaning there are many issues in tunning. Is there any computation cost analysis for better understanding the limitations of this method? 2. Though the author discusses SPIN, the truth is that equation 4.7 in SPIN is very similar to equation 6 in this paper. 3. The experimental results are not strong. It did not surpass the SPIN by a large margin. Are there any other baselines that can be incorporated such as DPO?

Questions

How can the reward model obtain the generalization power? Can the author say more about it, in my understanding, the main issue that affects the generalization is the scale of the demonstration dataset.

Rating

5

Confidence

3

Soundness

2

Presentation

3

Contribution

2

Limitations

See weaknesses. I have to say I am not an expert in this field, therefore, I am willing to change my attitude if I find I am wrong.

Concern about OpenLLM benchmark and Comparison to SPIN

Thank you for providing more details on the computational cost of the algorithm and showing the accuracy of the obtained reward model. I am still concerned over the weak performance on the OpenLLM leaderboard. I checked the SPIN paper, and they reported a 8-9% improvement on the OpenLLM leaderboard. Do you know why there is a large discrepancy between the results reported in SPIN and those in this paper? Furthermore, is it possible to re-run some of your experiments with different seeds and conduct significance tests on the results?

A gentle reminder on the discussion period

Dear reviewers, We thank all of your constructive feedbacks toward this work, and we tried our best to address your concerns and questions. As the author-reviewer discussion period is ending in three days, we gently remind you to check out our rebuttal and let us know if you still have any comments. Thank you very much! Best, Authors of submission 14483