Imitating Language via Scalable Inverse Reinforcement Learning

Summary

This paper investigates the possibility of applying inverse reinforcement learning for language imitation learning problem. Specifically, the paper reformulates IQLearn as a temporal difference regularized extension of MLE. This basically bridges inverse reinforcement learning with MLE with a coefficient $\lambda$ that can be tuned to change the weight of inverse RL loss versus MLE loss. In experiments on XSUM and GMS8K, the proposed method outperforms vanilla MLE consistently across various base models from 250M to 3B parameters. Time profiling is also reported and it is shown that the proposed offline IQLearn induces minimal computational overhead compared to vanilla MLE.

Strengths

The paper is well-written and easy to follow. The problem is well-motivated as inverse RL is already shown to work better than vanilla MLE is many robotics control problems. The idea of applying inverse RL to the language setting seems novel to me. However, I am not an expert in inverse RL so not sure about the novelty of reformulating IQLearn as a temporal difference regularized extension of MLE. The proposed method is extensively validated with various base models on two common language benchmarks: XSUM and GMS8K. The section of time profiling is particularly helpful and it is nice to see offline IQLearn enjoys a similar training efficiency as MLE.

Weaknesses

As admitted in the paper in Line 193 as well, the performance gain is not big (although consistent across different base models) and around 1% accuracy on GMS8K and 1 ROUGE-1 score. It would also be nice to see how this method works on other metrics for text generation such as BertScore[1]. Would be nice to have more intuitive understanding of why adversarial methods such as GAIL do not work very well. There are some places where the writing can be improved. Please see questions. [1]Tianyi Zhang, Varsha Kishore, Felix Wu, Kilian Q. Weinberger, & Yoav Artzi (2019). BERTScore: Evaluating Text Generation with BERT. CoRR, abs/1904.09675.

Questions

Line 31, "Including" -> "This includes" The drawing of Figure 1 can be improved, currently it is hard to tell from the figure what the differences between MLE, offline IRL, and online IRL are. It seems that MLE and offline IRL are exactly the same while online IRL has some blue boxes instead of grey boxes. It is unclear to me what that means even with the help of the caption. It seems that adversarial imitation learning method GAIL is always outperformed by non-adversarial one IQLearn? What might be a reason for that? In Figure 2 and 3, is the result of IQLearn offline or online?

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

Limitations are discussed in the Discussion section.

Summary

The paper introduces a RL-centric perspective to imitation for LLMs, with an novel imitation learning algorithm for fine tuning LLMs, that is derived from IQLearn, forming a dynamics dependent temporal difference regularized variant of MLE. The authors provide an extensive analysis of other potential inverse RL algorithms for imitation learning with LLMs, and demonstrate IQLearns ability to perform well at tasks whilst still having good diversity. Their approach is a promising offline IRL approach.

Strengths

* Clarity: The introduction is clear and well motivates the problem. Overall the paper has good clarity, bar a few typo's. * Originality: The RL-centric perspective to imitation for LLMs and the IQLearn objective and formulation appears novel and significant to the community.

Weaknesses

* The paper could benefit from more experimental tasks such as CommonGEN, ROCStories, EMNLP2017, and COCO, as done in TextGAIL. * Missing error bars in Figure 4 and Figures 2 & 3 could benefit from error bars as well, perhaps with a different version in the appendix. * Minor: The first paragraph could benefit from references to ground the reader. * Minor: L90. Reading the references cited in this paragraph, they do not mention the term $(1-\gamma)$; perhaps find a reference for it, remove this term, or explain its inclusion in a footnote or an appendix. * Minor: Equation 8 could put an enclosing bracket for the min operator to help guide the reader. * Minor: The presentation of the figures could be improved by making the font size larger and using vector graphics. Typos: * L40: “behavior cloning [7], its” -> “behavior cloning [7], and its” * L123: “further simplify” -> “further simplify it” * L168: “temporal different regularization” -> “temporal difference regularization” * L174: “we a short” -> “we use a short” * L174: “pure MLE beneficial” -> “pure MLE being beneficial” * L:202: “full trajectories” -> “full of trajectories” * L:218: “Right” -> “Left”

Questions

* Are the T5 model’s pre-trained? I presume they are, if so it could be helpful to state in the main paper in section 3.2.1 that the T5 models are pre-trained, and details of this. * As stated, IQLearn with Lambda = 0 “retrieves standard MLE.” Can we then interpret Figure 2 as IQLearn (Lambda=0.0) as MLE? If so, it is interesting how IQLearn always outperforms MLE, even for small Lambda. Could you perform ablations for small non-zero values of Lambda to show if it converges with MLE on the task results? * How do the conclusions of Figure 4 hold if Lambda is varied? What is the expectation? * Figure 6: can you quantify the uncertainty estimate used? * How would you see future work extending your approach to use human pairwise preference datasets? That is combine approaches perhaps such as Direct Preference Optimization.

Rating

8

Confidence

3

Soundness

4

Presentation

3

Contribution

4

Limitations

These were adequately discussed in Section 5.

Summary

This paper investigates using inverse reinforcement learning (IRL) to directly optimize sequences for fine-tuning large language models. Moreover, this work reformulates inverse soft-Q-learning as a temporal difference regularized extension of maximum likelihood estimation (MLE), which bridges IRL and MLE in supervised fine-tuning tasks. The experiments demonstrate the clear benefits of IRL-based imitation for retaining diverse responses while maximizing task performance.

Strengths

1. This work investigates the potential of inverse reinforcement learning for tuning LLMs. The strengths of IRL are evaluated in terms of performance, diversity, and computational requirements. The experiments demonstrate the computationally cheaper offline IRL can obtain crucial performance gain over MLE-based LLM tunning. 2. IQLearn is a method that can work online and offline. The re-formulation of IQLearn to LLM tunning enables large performance gains compared to MLE methods. Moreover, IQLearn can be regarded as a regularized extension of MLE, which bridges IRL and MLE for LLM tunning.

Weaknesses

1. IRL methods, even for offline IRL methods, cost more time than MLE for LLM tunning. 2. More experimental details are required to reproduce the results.

Questions

1. The original GAIL is an online IRL method, does GAIL execute offline or online in the experiments?

Rating

5

Confidence

3

Soundness

3

Presentation

2

Contribution

3

Limitations

The discussion section discussed the potential limitations of this work.

Summary

In this paper the authors aim to cast fine-tuning LM as inverse RL from the perspective of distribution matching. In order to avoid online generation in existing IRL algorithms such as GAIL, the authors leverage an offline IRL algorithm, IQL, and reformulate it as a temporal difference regularized extension of MLE. This derivation enables the use of expert samples solely in the form of distribution matching. Experiments in XSUM and GSM8K show that their method outperforms both MLE and online IRL algorithms in terms of task performance and diversity.

Strengths

The reformulation of IQL into distribution matching is novel and it leads to a simple and efficient version of IQL. The simplicity of the new derivation of IQL may potentially make it an impactful learning objective like MLE.

Weaknesses

- The improvement is somewhat marginal especially in the task performance on the summarization benchmark. - It is not clear the proposed method is better than MLE in terms of quality-diversity trade-off. - The diversity evaluation for lower open-endedness tasks like GSM8K does not make much sense.

Questions

- Could the authors provide an analysis for quality-diversity trade-off [1] of the algorithms? - Why does the temporal regularization term encourages higher diversity? - It would be helpful if the paper contained pseudo code of the proposed algorithm. [1] LANGUAGE GANS FALLING SHORT

Rating

4

Confidence

3

Soundness

2

Presentation

2

Contribution

2

Limitations

see weaknesses

Summary

The paper looks at the existing IQLearn algorithm and makes some connections with standard maximum likelihood learning in the context of sequence based language models. On the empirical front, results are given showing different model performance along diversity (measured via self-Bleu) and some accuracy measure (such as rouge). These are given on a few public datasets. The empirical part of the paper lacks focus on what is trying to be optimised, and the actual metrics that are reported leave open some potential that the performance is not necessarily better. For example self-bleu may increase, but it isn't clear this is a better model. The diversity may be erroneous, but still more diverse and still scoring well on the token overlap measures of rouge etc. The other graphs and figures fail to leave the reader with a clear story of why using imitation learning at the sequence level, rather than MLE or RLHF (imitation learning at the token level) is necessarily better.

Strengths

* A solid theoretical analysis is given which draws links between MLE and forms of IRL.

Weaknesses

* The story of the paper is not that clear, and the empirical results do not tell a clear story. The paper shows that the IRL methods can be used to obtain better models on the Pareto type plots of quality and diversity, however it's not clear that these are actually better models. * The data used in the empirical section only appears a long way into section 3. It would be much clearer if the empirical section had a clearer introduction of what types of problems are being looked at, and what data is being used for such.

Questions

What was the conclusion for which tasks and dataset properties is IRL most relevant for? (This is a question at line 160, however it's not clear that it is answered).

Rating

4

Confidence

3

Soundness

2

Presentation

2

Contribution

2

Limitations

No concerns.

Follow up to reviewer tAhh

Thank you for your original review. Our team has worked hard to answer your questions and address your comments. Since the original rebuttal, we further expanded the reward analysis to GSM8k and TLDR. With the discussion ending soon please ask any remaining questions or comments soon and we will answer ASAP; or share if previous answers and additional experiments have addressed your concerns.

Follow up for reviewer aRfw

Thank you for your original review. Our team has worked hard to answer your questions and address your comments. Since the original rebuttal, we further expanded the reward analysis to GSM8k and TLDR. With the discussion ending soon please ask any remaining questions or comments soon and we will answer ASAP; or share if previous answers and additional experiments have addressed your concerns.