Imitating Language via Scalable Inverse Reinforcement Learning
Authors: Markus Wulfmeier
Michael Bloesch
Nino Vieillard
Arun Ahuja
Jorg Bornschein
Sandy Huang
Artem Sokolov
Matt Barnes
Guillaume Desjardins
Alex Bewley
Sarah Maria Elisabeth Bechtle
Jost Tobias Springenberg
Nikola Momchev
Olivier Bachem
Matthieu Geist
Martin Riedmiller
Venue: NeurIPS 2024 poster (2024)
Keywords: Language Modeling
Inverse Reinforcement Learning
Imitation Learning
Supervised Fine-tuning
Primary area: natural_language_processing
TL;DR: Investigating an effective, computationally scalable reinforcement learning perspective to imitation for language modeling.

Abstract: The majority of language model training builds on imitation learning. It covers pretraining, supervised fine-tuning, and affects the starting conditions for reinforcement learning from human feedback (RLHF). The simplicity and scalability of maximum likelihood estimation (MLE) for next token prediction led to its role as predominant paradigm. However, the broader field of imitation learning can more effectively utilize the sequential structure underlying autoregressive generation. We focus on investigating the inverse reinforcement learning (IRL) perspective to imitation, extracting rewards and directly optimizing sequences instead of individual token likelihoods and evaluate its benefits for fine-tuning large language models.  We provide a new angle, reformulating inverse soft-Q-learning as a temporal difference regularized extension of MLE. This creates a principled connection between MLE and IRL and allows trading off added complexity with increased performance and diversity of generations in the supervised fine-tuning (SFT) setting. We find clear advantages for IRL-based imitation, in particular for retaining diversity while maximizing task performance, rendering IRL a strong alternative on fixed SFT datasets even without online data generation. Our analysis of IRL-extracted reward functions further indicates benefits for more robust reward functions via tighter integration of supervised and preference-based LLM post-training.

Imitating Language via Scalable Inverse Reinforcement Learning
Markus Wulfmeier
Michael Bloesch
Nino Vieillard
Arun Ahuja
Jörg Bornschein
Sandy Huang
Artem Sokolov
Matt Barnes
Guillaume Desjardins
Alex Bewley
Sarah Maria Elisabeth Bechtle
Jost Tobias Springenberg
Nikola Momchey
Olivier Bachem
Matthieu Geist
Martin Riedmiller
Google DeepMind
Abstract
The majority of language model training builds on imitation learning. It covers pretraining, supervised fine-tuning, and affects the starting conditions for reinforce- ment learning from human feedback (RLHF). The simplicity and scalability maximum likelihood estimation (MLE) for next token prediction led to its role as predominant paradigm. However, the broader field of imitation learning can effectively utilize the sequential structure underlying autoregressive generation. We focus on investigating the inverse reinforcement learning (IRL) perspective to imi- tation, extracting rewards and directly optimizing sequences instead of individual token likelihoods and evaluate its benefits for fine- totettuning large language models. We provide new angle, reformulating inverse soft-Q- soft-Q-learning as a temporal differ- ence regularized extension MLE. This creates a principled connection between MLE and IRI. and allows trading off added complexity with increased performance and diversity of generations in the supervised fine-tuning (SFT) setting. We find clear advantages for IRL-based imitation. particular for retaining diversity while maximizing task performance, rendering IRL. a strong alternative on fixed SFT datasets even without online data generation. Our analysis of IRL-extracted reward functions further indicates benefits for more robust reward functions via tighter integration of supervised and preférence-based LLM post-training.
1 Introduction
In recent years, the imitation existing human knowledge via large datasets has become a key mechanism underlying increasingly capable and general artificial intelligence systems [17, 45, 9]. Pretraining and supervised fine-tuning phases for large language models (LLMs) predominantly rely on imitation learning, in particular next token prediction via maximum likelihood estimation (MLE). In addition, preference-based fine-tuning is affected by imitation via initial online data generation and optimization objectives such as regularization towards the previously fine-tuned LLM [46, 12].
The field of imitation learning for sequential decision making has a long-standing history for appli- cations such as robotic control [4, 29]. Recently, perspectives to language modeling have shifted sequential decision making problem in particular for later stages of model adaptation via reinforcement learning from human feedback (RLHF) [46, 13, 17, 66). This vantage point opens up new opportunities for the effective use of different data sources and obtaining aligned models that better represent human intent. includes a broader scope for which data contains information about rewards and preferences as well as the dynamics Syyamiics-ware aware optimization of each action based on its future impact all while taking into account computational scalability.
towards explicit treatment
Now at Cohere.
38th Conference Neural Information Processing Systems (NeurIPS 2024).

---

MLE
Offline IRL
Online IRL
of Training Objective
Training Objective
Training Objective
Model
Model
Model
Data usage and optimization flow in MLE, offline and online IRL. Independent of the method, current models use the history of past tokens to predict the next. However, MLE purely optimizes the current output for exact matching the corresponding datapoint while IRL-based methods take into account the impact on future tokens. Online optimization additionally conditions past model generations rather than the original dataset. Grey and blue objects respectively represent training data and model generations. The impact of future datapoints is often indirect and mediated via learned functions (e.g. the discriminator in GAIL [25] and the Q-function in IQLearn [20]).
Figure 1:
Due to the importance of imitation for language modelling, we believe detailed analysis of the underlying problem for sequential decision making is warranted. Despite the widespread perspective of RL for aligning and fine-tuning language models (via RLHF), supervised learning via maximum likelihood estimation for next token prediction remains the dominant component of our imitation learning pipelines due to its simplicity and scalability. However, pure MLE for next token prediction (including "teacher forcing" [64]) can create challenges in autoregressive models, many of which being related to classic challenges with behavior cloning [8], its equivalent in the context sequential decision making. Compounding errors can occur due to iterative model application creating data sequences which further shift from the training distribution [56, 30), growing increasingly likely for longer sequences [16]. In particular, a model's 's own samples can cause such distribution shifts and exposure bias [52. 6, 5]. Taking the RL perspective to imitation aims to mitigate these issues via dynamics-aware optimization, where each action is optimized for the impact on the whole future trajectory. It further enables a shift from passive to active learning, where the system actively generates data. Furthermore, best performance of the fine-tuned model is only one metric of importance. Indeed, continued alignment of language models to human preferences requires sampling for given prompt. collecting human preferences over these completions, and finally aligning the model via preference fine-tuning. Improving the diversity of sampled completions via temperature sampling [7], or using mixtures of past models [61] has been linked to performance improvements. Studying inverse RL. potential divergences, and regularizations provides another angle to increasing diversity [23].
In this paper, we investigate RL-based optimization, in particular the distribution matching perspective to inverse reinforcement learning (IRL., for fine-tuning language models; which can be contrasted with standard MLE as depicted in Figure Our goal is improved understanding of when, and how, IRL can be used as an effective alternative for supervised MLE in the fine-tuning pipeline. The evaluation covers both adversarial and non-adversarial, offline and online methods. We further extend inverse soft Q-learning [20] to explicitly provide a principled connection to classical behavior cloning or MIF.. Our experiments range from 250M to 3B parameter models for the encoder- T5 [51] and decoder-only MM [3] models. Throughout evaluation, we investigate task performance and diversity of model generations illustrating clear benefits of inverse RL over behavior cloning for imitation. A further, principal benefit of RI- imitation close to RLAIF,, RL from Artificial Intelligence Feedback- - is the natural connection to later RLHF stages via rewards obtained from demonstration data and we take first steps to analyse the value of IRL-obtained reward functions.
Our key contributions are:
We investigate the RL-centric perspective to imitation for LLMs, extracting rewards and directly optimizing actions for sequence generation instead of individual token likelihood.
We reformulate inverse soft Q-learning as temporal difference regularized extension of MLE. This explicitly bridges between MLE and algorithms exploiting the sequential nature underly- ing language generation and enables computationally cheap offline training.
We compare MIF and IRI. formulations including adversarial and non-avversarial, offline and online methods to improve our understanding of imitation in LLMs. Our main results show better or on par task performance, with increased diversity of model generations, in particular demonstrating that key improvements can be obtained via (better scalable) offline IRL.

---

Finally, we amalyze the extracted reward functions indicating the potential usefulness of IRL to obtain discriminative in addition generative improvements from demonstrations.
2 Methods
Language generation can be modeled as a sequential decision making problem. On a token level, is the problem of generating the next token I; given the already generated sequence tokens (zo..... (20....;; (.0....;.ii11) We thus seek distribution (20 ................................ Ii-1), which we will also refer to policy, For given policy the sequence likelihood æ (To. : IN can be computed autoregressively;
p II (2) 2111)
(1)
0
The classic maximum likelihood estimation based approach leverages this factorization in order to efficiently train the policy by maximizing the log-likelihood the training sequences, D
(a x M3
N log n (2: 3o. -............................... 2t-1)-
arg
logp(x) arg max
For fine-tuning problems, where a pre-trained model is fine-tuned to particular set of tasks, the problem can be formulated analogously but may have additional conditioning variables which will leave out for the sake of clarity.
Distribution matching. State-action distribution matching algorithms [25, 20], which are well- established in the field of imitation learning and can be seen as solving IRL problem see e.g. [25] approach the problem in a different manner: They seek minimize the divergence between the dddiscuuteed state-action distribution of the policy mTass) and the discounted state-action distribution of the expert policy ( We define the discounted state distribution for Markov Decision Process (MDP) as p(s) = (1-+) -i P'p s|π) and the discounted state-action distribution accordingly defined as μ.. (s. a) = {{sss)((s); where P(s, = s/T) is the probability of seeing state when acting according to the policy æ. For our autoregressive generation MDP, the state corresponds the concatenation of already generated tokens (commonly including the prefix or prompt), & = (Ip..... Z;- and the action is the next token, a, = I; thus yielding a problem with deterministic dynamics. We omit state and action arguments in the following discussion whenever it clear from the context.
In order to enable more straightforward algebraic manipulation, the divergence is often combined with a weighted causal entropy term H (TT) = EE,, log(T(als)). The goal is then to find a policy that minimizes the objective J(x):
J J(T) =
-M
where D, = Eps [f [f(aala.))
(gla,s) )] is an f-divergence. Different f-divergences have been used in the literature [25, 23]. When taking the example of the reverse KL divergence, with f(() SRKI. (1)
lt), we can decompose the objective into a state distribution and an MLE term:
min FExI (M~l)uR) (T(pello) + Eps ([T[]|||]))
(4)
where KL denotes the KL diveiveree which corresponds maximum likelihood (on the discounted state distribution) for the second part after dropping terms independent of In comparison to MLE, IRL algorithms thus also try to match expert actions but additionally attempt to match the state visitations of the expert. The use of different f-divergences can influence mode seeking and mode covering properties of the optimal policy [23].
Adversarial imitation. The state-action divergence can be hard to evaluate for two reasons: first requires many samples from the current policy and second it requires access to the unknown expert densities. In a first step, a variational representation of the f f-divergence can be leveraged to avoid expert densities:
min ((T) = min max
- EE (g)) + E,,, [g] + AE [loga],
(5)
g -domy.
3

---

where f is the convex conjugate of f and domy. is its domain. Notably, GAIL [25] can be retrieved by using the Gensen-Shannon divergence with its convex conjugate (*) log(1 exp(g(>, a)))
and defining the discriminator D a) == exp(g(s.a)) a)):
min
max D-SA--
p [oog(1 - D)] + Epm [lo D) + X log T].
(6)
Non-adversarial imitation. From here, we can re-derive IQLearn [20], but instead consider state value rather than state-action value functions. This representation further allows us to establish clearer relationship to MLE. Using a change of variable r 8 x A + -dom with r(s.
g(s, (8, a) and re-arranging the terms, we start by rendering the internal RI-problem more explicit:
min (T) = max min Eps (-r)) + Ey. Ir X 1)
where r(s, a) can be interpreted as an 'implicit' reward function and we omit its arguments in the following for brevity. Due to the convexity of f and the concavity of the causal entropy (other terms are linear) this is saddle point problem [20] and the max-min can be swapped:
min ((T) = - min (-r)) + max Ep, Ir- A Alog m].
(8)
where the second term now corresponds to the discounted cumulative reward of a soft, or entropy regularized, RL problem [74]. The soft-RL problem is well understood [22] and we can use analytical identities regarding its solution for further simplification. In particular, we have that [42]
r(s, a) Aogg x" (als) = v (s) - Ea.n((laas)) (Au" (s') As E S,
(9)
where x² the optimal policy for the reward r and v" and p" are respectively the corresponding state value function and discounted state-action distribution'. This is applied to the right hand term of Eq. (8) after maximization to obtain (after dropping the arguments for conciseness).
min min (TT) = min Eps [fr)) + Ey-le" E.
(10)
Using a telescoping argument [20, 57J, policy to the difference in values on the state-action distribution induced by arbitrary other policies, i.e. Ep [u*] Fy,,, [o" 1o"-~xe >xe"]] In particular, here we can change the second expectation's 's state-action distribution to one induced by the expert policy yielding:
relate the value of the initial state distribution of
min [(r) = min F (-r) + Fsg'
(11)
r
= min E,g (-r) +r -Alog
(12)
where we use Equation (9) again in the second line. Here we obtain an explicit MLE term (the last term in the above equation), albeit via the intermediate of the optimal policy of the reward. We currently also have minimization problem in the reward r and the optimal policy being a function thereof. This can be changed by leveraging the bijective relationship between the reward and the optimal Q-value, r(g") 9² E. (Xe 120, 42], or analogously between the reward and state value and a policy, r(e", TT") = o" + A log π- Eg [Xe"]] (using qº = =v" + Alog T"). We can thus reparameterize the problem and optimize the optimal state value and policy instead of the reward.
min ((T) min E If" (-r(v", ,,))
x²)) rro", T²) Alog(=")]
(13)
Choosing the x diveergeee with convex conjugate (1()= - + is particularly convenient at this point since it can be combined with rescaling the value vx, = v" /> to obtain our reformulated IQLeam objective to be minimized.
Jigi cam (1x, -²)= = Eun [) (v + log n- Es Ex log(TT))
(14)
regularization
MLE
With this derivation we show the clear relation between MLE and the inverse RL perspectives (expanding on [48]): we can see distribution matching (and thus inverse RL) as performing maximum
Note that in practice, parameterize v based on the logits of using the identities x² <a(s)) exp.q (8,
and v" (s) = log E. exp (s, a) explained in the appendix.

---

73.0
77.5
better
GAL (ent-regin 0.00)
GAL (cnk ronkroog0000
89.0
GAL (ert-rag-0.11)) 0.130
10Learn (1 (k-0.50)
85.0
83
loccamn (1-1
87.5
(entceg--
90.0
B8
Large
X-Large
50
Base
Large
X-Large
NLE (ont rog-0.00-
(ent-reg-0 05)
14
20
Accuracy without Calculator
Accuracy with Calculator
Figure 2: GSM8K results for fine-tuning with MLE, IQLearn, and GAIL across different regularization strengths. particular MLE shows strong performance reduction with higher entropy cost. Larger models demonstrate higher performance but also stronger self similarity across generations, rendering effective trading of between task performance and diversity highly relevant. Error bars indicate the standard error of the mean after repeating the experiment with 3 different seeds.
likelihood with dynamics dependent temporal difference regularization term. In contrast to the adversarial setting, this objective does not require samples from the current policy but uses expert samples only². The regularization term couples the learned policy policies where the log probability of actions matches the difference in state values. An advantage of the above formulation is that it allows annealing of the regularization term where setting À = O retrieves standard MLE and where by adjusting À the regularization strength can be flexibly increased.
value function and favors
3 Experiments
In this section, we evaluate the benefits of different inverse RL based methods in comparison to MLE for training large language models. We assess their impact on task performance, diversity of model generations, and computational requirements. Concretely, we compare MLE-based next token prediction and different IRL methods for fine-tuning LLMs on common benchmarks. addition, we perform ablations on online? vs offline versions of IQLearn, showing results across dataset and model sizes. We finally add analysis of the implicitly learned rewards extracted from SFT data via IRL. methods (which bring the potential downstream use to aid RLHF/RLAIF [46, 37] training stages).
These experiments mainly aim to answer the following questions:
Do IRL methods provide a scalable, effective alternative to MLE for fine-tuning?
How are different algorithms placed the Pareto front of task performance and diversity?
For which task and dataset properties is IRL particularly relevant?
What is the impact online data for IRL training?
How informative are rewards extracted from SFT data?
3.1 Algorithmss, baselines, and datasets
In addition to naive maximum likelihood estimation for next token prediction, we evaluate the following IRL methods. Generative adversarial imitation learning (GAIL) [25] represents a common adversarial framework training a discriminator distinguish transitions from agent and dataset and separate policy. We heuristically adapt the algorithm to mitigate instability for adversarial training [69, 67] including the start for MLE-trained checkpoints and additional loss terms (including MLE) with further details in Appendix 1.1.3. IQLearn [20] departs from adversarial learning and our reformulation from Eq. 14 enables us principled control of the temporal difference regularization component to retain stable training. We will use the reformulated offline variant of the algorithm in all experiments and further add an ablation to online version in Section 3.3.1. Since we compare all methods with respect to task performance and diversity of generations, we additionally evaluate further entropy regularization terms for the MLE baseline; clearly denoted with 'ent-reg' in all plots, corresponding to the respective regularization parameter in GAIL and MLE (see Appendix A.2 for a more detailed description). In line with previous work on inverse RL for language modelling,
An online version of IQLeam is derived in Appendix A.1.2.
5

---

23
20
GAL (erk-rag-
better
GAL (ent-reguu000
(ork-rg--011
Self BUEU
53
10Uearn (4m0.5)
60
10Ccarn (1
NLE (enl-veg--00)
NLE (ont reg-0.011
(ent-reg-)) (.25)
89 Bage
80
NLE (ant-eg-0.3)
Large
Base
Large
15.5
ROUGE
ROUGE
Figure 3: XSUM results for models trained with MLE, IQLearn, and GAIL across different regu- larization strengths. ROUGE and ROUGE 2 are used as performance metrics on the x-axes with Self-BLEU as diversity measure on the y-axis. Entropy regularizing large MLE and GAIL trained models with 0.1 leads to catastrophic results outside the limits of the plot. Figure 9 in the appendix shows the corresponding plots for ROUGE-LSUM.
apply short warm-up phase with pure MLE [15, 65]. Rather than separate experiments or heuristic combinations, the explicit MLE term emerging out of our IRL. objective Section enables principled integration of this mechanism.
We use the following datasets and subsets for ablation in the following sections: XSUM [43],
GSM8k [14]. TLDR [58]. and WMT22 [34]. Unlike parameter-effiiient fine-tuning via adapters [27] as used in prior work [15], we focus on the full fine-tune fine-tuning setting to decouple our analysis from the specifics of adapter-based optimization dynamics [10].
3.2 Quality-diversity evaluations
We evaluate both encoder-decoder and decoder-only model classes, respectively using the T5 [5]] and PALM2 [3] models. Our main visualizations focus on task performance and diversity of model generations. For task performance, we use the standard metrics for the respective benchmarks (e.g. ROUGE-1 and accuracy percentage). To measure diversity of model generations self-similarity of generated examples as measured by Belf-BLEU [73]. A high score denotes low diversity and vice versa. Using these different axes of evaluation allows us to visualize Pareto fronts between performance and diversity which we will use to assess algorithmic differences. We further entropy Appendix A.3.1.
calculate
add the evaluation via
3.2.1 T5 models
We perform experiments with the base, large, and T5 models [51] on the XSUM [43] and GSM8k [14] tasks. These models further serve as foundation for our later ablations. We are able obtain small but notable gains in performance across all tasks, shown Figures 2 and in particular math and reasoning tasks show clear improvements in accuracy when fine-tuning with IRL compared to MLE. We hypothesize that specific and shared structure of responses is better exploited via IRL. methods. There is more emphasized boost the diversity of model generations for IQLearn over MLE. In comparison to prior work [65], we have been able to stabilize GAIL training, with minor changes described in Appendix 1.1.3, but are required to start from a checkpoint previously trained via MLE to ensure strong similarity between dataset and model generations starting GAIL training. This only applies to the T5 model class and for PaLM2 models GAIL is highly
NLE IQsami
VLE
qieamd 0,5)
109 Shrombbr
Figure 4: PaLM2 results for various sampling temperatures with MLE and IQLearn. Left: GSM8K. Mid: TLDR, Right: WMT22, including beam search. By propagating sequence information during training, IQLearn reduces inference time dependency beam search for improving performance.

---

challenging to stabilize with details in Appendix A.1.3, While entropy terms can further be added to MLE optimization, we hypothesize that better trade-off's between diversity and performance can be obtained via methods able to aggregate information across trajectories to optimize entropy over a different space, complete trajectories rather than just per step policy entropy.
3.2.2 PaLM2 models
We also perform experiments with PALM2 models [3], specifically fine-tuning from pre-trained PALM2 'Gecko' model. We evaluate offline IQLearn on the summarization task TLDR [58], the mathematical reasoning task GSM8k [14]. and the large (285M examples) English-to- Germann translation dataset WMT22 [34]. We limit these experiments to single choice of the regularization parameter per task to save computational costs and instead include an analysis of the effect of the sampling temperature parameter during sampling (noting that we similarly observed e outperforming MLE at varying temperatures for the T5 models). We selected the checkpoints with early stopping for WMT22 as BC was overfitting on the task, unlike TLDR and GSM8K for which we evaluated their latest checkpoints.
Similar to the previous section, we perceive improvements over MLE on all three benchmarks, though for lower accuracy values MLE covers part of the front. Figure 4 summarizes these improvements, showing the performance of the trained models depending on the temperature used during sampling responses for the test set prompts. These results show a similar behavior between all three tasks, where IQLear achieves higher performance in low temperature regime.
3.3 Analysis and ablations
We perform additional experiments and post-training analyses to better understand the impact of dataset size, initial checkpoints and computational demands of offline and online al algorithms.
3.3.1 Computational efficiency and accuracy for online & offline inverse RL
One of the key benefits of (offline) MLE compared with (online) IRL-based methods is its lower computational costs. These are principally related to online rollouts - slow, autoregressive sampling from the model. Our re-formulation of IQLearn results in algorithm that can be applied offline on a fixed dataset, which underlies all IQLearn results presented, mitigating this limitation. In this section, we additionally present update and sampling times across algorithms* and the comparison with the online application
IQLearn (see Appendix 1.1.)). Figure 5 demonstrates minimal task performance gains (for the T5 base mode on GSM8k), though considerably improved diversity for model generations. At the same time, Table 1 visualizes the relative cost of sampling in comparison to different algorithm updates, excluding sampling. Note that MLE batch sizes are smaller than IRL ones as we add additional online samples to the batch. Experiment times for online IRL can be reduced, but at the cost of additional hardware, via distributed training with separate sampling infrastructure. The application choice of online or offline IRL finally lies with the practitioner trading of additional computational cost with diversity benefits.
*Computational resources and time in practice highly dependent hardware and implementation but always include crucial, additional sampling costs.
Elsert Kleam colino 100p-03 Cleam celino 1alln0
1OLEITL oeline FIDOH.L
32
2
1300 2000
3000
2000
3200
K000
1000
1000
2000
3052
4200
5000
5000
7030
Size training
Size
training 541
Figure 5: Left: performance of offline and online inverse RL performance with online ratio describing the ratio of offline data used. Right: diversity of model generations. While only showing limited gains in performance, diversity clearly improves.
7

---

We provide intuition for low differences between offline and online Inverse RL via toy experiments in Appendix A.3.3. At its core, the dynamics of autoregressive language generation, in particular the concatenation underlying single-tum settings, prevents exact recovery after mistakes, which could be otherwise learned via online IRL. Therefore, extensions of the LLM action space to enable recovery a fruitful direction for increased benefits from online data with RL or IRL methods [15].
Table I:
Algorithm profiling with computation in milliseconds. 'Sampling' refers to generating number of sequences equivalent to batch size and often uses equal or more time than updates. These times depend on hardware, implementation and code optimization.
15-base
T5-large
T5-XL
MLE Update
189 + 11
422 + 17
1031 + 22
GAIL Update
451 13
1064 + 25
1410 + 17
IQLearn Update
196 - 13
606 = 13
1355 + 46
+ Sampling (for online IRI.)
443 + 45
1345 + 211
1823 + 327
3.3.2 Dataset size ablations
Evaluating training on smaller subsets of GSM8K and XSUM with T5 base is respectively pictured in Figures 6 and 7. Performance improvements are consistent across dataset sizes. The analysis of subsets of XSUM further demonstrates increased robustness against overfitting with IQLearn not showing any of the performance loss over time that plagues MLE in particular with the smallest subsets which cannot be overcome with simple entropy regularization. In line with arguments around the compounding of errors in imitation learning [56], we find that both datasets with longer targets and smaller datasets show stronger task performance gains for IRL.
3.4 Reward analysis
In comparison to related work on applying IRL. to classical control domains, there is no access to ground truth reward functions underlying the process of data generation. Instead, we measure the correlation between IRL. extracted rewards and other task-specific performance metrics. High correlation here tells us how informative a reward function is w.r.t. task performance.
In particular, IQLeam represents learned rewards implicitly via the Q-function as rt Q(se, az) yV (s(+1)) and its online version (i.e. with a non-zero mix-in ratio a of on-policy examples, see Appendix A. 1.2)) additionally exposes the algorithm to (initially, low reward) policy rollouts to help discriminate between them and (high reward) SFT data. In Table 2, report the Spearman's 's rank correlation coefficient between accumulated rewards (over complete sampled trajectories for the full validation sets) for online IQLean (a = 0.1) and task-specift metrics. Compared to MLE rewards, which, as expected, are not strongly correlated with any metric, pointing to IQLearn incorporating the task-relevant quality into the extracted rewards. We find that using online data is important for consistent correlations across all tasks, in particular as we evaluate over agent generated rollouts. We hypothesize that the comparably lower correlations for GSM8K are likely to be explained by the task's 's idiosyncratic metric: only the correctness of the final generated numerical expression affects the accuracy calculation, effectively ignoring most of the trajectory (and
clear increase in correlation
Acormcc cabou Abor
02
with
3.0
1302
2200
2000
40000
5230
2200
4322
2035
4020
7006
Figure 6: Different subsets of GSMSK Performance gains persist across dataset scales with larger datasets demonstrating minimal preference for larger regularization coefficients. Each experiment was run with 3 random seeds to obtain uncertainty estimates.
8

---

10.111 Mccrrppes 1000
ruin wetetiiies 3000
rum wtamples 10000
3.0
2.0
10Leam ul-0 50)
MLE lent
2.6
MLE lemt mw..
MLE lemt mg-00.01
2.4
MIE (emt-nng-0000
0.0
0.5
1.0
3.5
2.0
2.5
300
3.5
4.0
0.0
0.5
10
1.3
2.0
2.5
3.0
3.5
4.0
0.0
0.5
3.0
1.5
2.0
25
30
35
4.9
Training Stop
165
Taaiinng Stop
165
Tranning
165
Figure 7: Learning curves for subsets of the XSUM training data. The smallest subsets demonstrate strong overfitting for pure MLE which the TD regularization in IQLearn mitigates. Pure entropy regularization is unable to obtain similar robustness and directly conflicts with task performance.
so its transition's 's rewards) up to a specific answer segment. This highly targeted reward function becomes harder to learn. Finally, Figure 8 displays how the reward alignment with the task metrics increases with larger TD regularization A.
Task
Metric
TQLearn
ML.
CheF
0.5
TLDR
ROUGE-1
0.64
-0.05
0.4
ROUGE-2
0.41
-0.04
ROUGE-Lsum
0.65
-0.05
0.3
WMT22
BLEU
0.43
-0.05
0.2
ChrF
0.43
-0.01
0.1
GSM8k
Acc. w/ calculator
0.17
-0.02
0.0
Acc. w/o calculator
0.17
0.04
0.1
10-1
10a
Table 2: The Spearman's 's rank correlation for online IQLearn (a 0.1) with À 0.1 (for GSM8k, WMT22) and = 0.5 (for TLDR), compared to MLE (i.c. À 0.0).
Figure 8: Reward correlation on WMT22 as a func-
tion of X for a fixed mix-in 0.1 for online data.
4 Related Work
General imitation learning. Imitation learning assumes a dataset of expert demonstrations, and the aim is to train policy that matches the expert. There are two broad categories of imitation learning approaches: behavioral cloning (BC) and inverse reinforcement learning (IRL). In BC. policy is trained using regression to directly mimic the expert demonstrations [49]. This is analogous to supervised fine-tuning of language models via MLE. BC requires sufficient data coverage to perform well, and suffers from compounding errors at evaluation time, as the policy deviates from the state distribution covered by the demonstrations. This can be alleviated by additionally querying the expert for the correct action in the states visited by the agent [56].
Inverse reinforcement learning. In contrast, IRL, jointly infers the policy and reward function, such that the provided expert demonstrations are optimal under the reward function, and the learned policy maximizes this reward function [44]. By using additional environment interactions beyond the demonstrations, IRL. can theory overcome the compounding errors observed with BC [70]. Note that IRL is related to RL from human feedback (RLHF) but differs in key aspects: IRL also learns both reward and policy, but extracts information from demonstration and agent data rather than paired preference data [31]. The game-theoretic approach to IRL treats the optimization problem as a zero-sum two-player game [60]. A subset of recent IRL methods can be seen as combining game-theoretic IRI. with entropy regularization of the policy, where the doubly-nested optimization is implemented with a classifier (GAIL [25]. DAC [35], AIRL [19]), implicit reward functions (ValueDICE [36], [Qeearn [20], or Lagrangian dual objective (PPII. [62]). The stable, large-scale application of inverse RL methods has been persistent goal throughout these developments. The classical requirement for complete RL optimization before updating the reward function has presented a limitation [75] which can be overcome via abstracted [68, 9] or linearized models [18], iterative value function based formulations [20]. We expand on the insights of the latter to evaluate the competitive performance of computationally cheap offline IRL and emphasise the connection between MLE and IRL [48].
adversarial training [25, 19] lastly saddlepoint-bar

---

Imitation learning for language modeling. Understanding language modeling as an imitation problem has been previously explored. Indeed, the link can already be made from MLE, commonly referred to as Behavioral Cloning (BC) [8] from an imitation perspective. Although the link to imitation has been made explicit recently, either theoretically [59], or in the case of distillation and divergence minimisation [1, 28, 40, 26, 391, some works had already tackled the issues of MLE with imitation- like techniques. For example, adversarial training of text generation an alternative to MLE was first proposed in SeqGAN [71], and followed-up by series of work using GANs for text generation [32]. These methods have been shown to work only in the temperature regime 1], a possible shortcoming that we address in Section 3. Then, leveraging the literature of imitation learning, GAIL was successfully adapted to language [65], showing an improvement over MLE. Closer to our contributions, IQLear utilized for language in SequenceMatch [15]. Key differences to our work include the reformulation as temporal difference regularized MLE, large-scale analysis including other inverse RL methods and focus on computational costs via the application of offline IQLeam. Indeed, SequenceMatch requires the use of online data, via the introduction of the "backward" token, that allows the model to change a previously chosen token during sampling.
5 Discussions
Our investigation focuses on diversity measures such as Self- or model entropy which are easily calculable but limited with respect to their ability to describe the impact on later training stages. Future evaluation and practical application will demonstrate if the increased diversity is relevant to RLHF such as for human raters in preference data evaluation or improved exploration during subsequent RL optimization [53].
The field of imitation learning has led to a gamut of algorithms, many of which are intuitively simple to implement with existing RL or RLHF infrastructure (e.g. [54, 63, 24, 50]). Ease of adaptation and hyperparameter tuning have principal impact on our practical algorithm choices and the methods and extensions discussed in this work enabled quick first results and iteration. Looking forward, the evaluation and adaptation of further imitation learning methods to LLMs is likely to lead to fruitful results in the coming years. While our analysis focuses on specific algorithms, our key arguments should be seen in the light of the benefits of underlying mechanisms rather than the specific methods.
The sampling- -free application of RL mechanism can eventually extend to even larger datasets such as pretraining data, domains with high requirements for computational efficiency. While these datasets are, in a certain sense, less optimal and representative of preferred model behavior, different aspects of our analysis have the potential to generalize, such as increased efficiency for modelling sequential decision making increased diversity of model generations. Pretraining and other sub-optimal data could further be used as additional data source for IQLearn and related algorithms [41, 33, 57] instead of autoregressive sampling data.
Finally, RLHF's key role lies in the alignment of models with respect to user preferences. By inte- grating SFT data into RI.-based optimization, we hope to expand data used for preference description from paired comparisons to individual demonstrations. Integrating generative and discriminative information from different data sources into unified RL-based framework for unified LLM post- training has further practical potential as previously demonstrated for the generative side [46] with first steps towards reward transfer [38]. Our reward analysis in Section 3 provides an initial signal about - reward functions including crucial information about L.LM task performance.
6
Conclusions
This paper presents a detailed investigation of the potential IRL algorithms for imitation in language model tuning focusing on performance, diversity, and computational requirements. We introduce a reformulation of IQI Lean which enables principled interpolation between robust, standard supervised fine-tuning and more effective IRL algorithms. Our experiments demonstrate particularly strong gains for IRI. on the Pareto front of task performance and diversity of model generations. While prior work primarily focused on online IRL, we demonstrate that computationally cheaper offline IRL. without the requirement of online sampling. already obtains crucial performance gains over MLE- based optimization. Additional correlation analysis between IRL-extrected rewards and performance metrics further emphasises the potential to obtain more accurate and robust reward function for language modelling. We hope this work will help to pave the way for better compromises between data and compute efficiency via RL-based algorithms across the complete LLM training pipeline.
10

---

Acknowledgments and Disclosure of Funding
We extend our gratitude to Daniele Calandriello, Doina Precup, and anonymous reviewers for their invaluable feedback on the final manuscript drafts. We also would like to acknowledge the contributions of our internal engineering teams, whose robust and flexible RLHF infrastructure significantly accelerated algorithmic development and evaluations. Special thanks are due to Piotr Stanczyk and Leonard Hussenot for their swift and effective solutions to numerous infrastructure challenges. Finally, we deeply appreciate the strategic feedback and guidance provided by Satinder Singh, Dale Schuurmans, and David Silver.
References
[1] Rishabh Agarwal, Nino Vieillard, Yongehao Zhou, Piotr Stanczyk, Sabela Ramos Garea, Matthieu Geist, and Olivier Bachem. On-policy distillation of language models: Learning from self-generated mistakes. In International Conference on Learning Representations, 2024. 10
[2] Firas Al-Hafez, Davide Tateo, Oleg Arenz, Guoping Zhao, and Jan Peters. Ls-iq: Implicit reward regularization for inverse reinforcement learning. arXiv preprint arXiv:2303.00599 2023. 16
[3] Rohan Anil, Andrew M. Dai, Orhan Firat, Melvin Johnson, Dmitry Lepikhin, Alexandre Passos, Siamak Shakeri, Emanuel Taropa, Paige Bailey, Zhifeng Chen, Eric Chu, Jonathan H. Clark, Laurent El Shafey, Yanping Huang, Kathy Mcier-Houe Hellstern, Gaurav Mishra, Erica Moreira, Mark Omemnic,, Kevin Robinson, Sebastian Ruder, Yi Tay, Kefan Xiao, Yuanzhong Xu, Yujing Zhang, Gustavo Hernandez Abrego, Junwhan Ahn, Jacob Austin, Paul Barham, Jan Botha, James Bradbury, Siddhartha Brahma, Kevin Brooks, Michele Catasta, Yong Cheng, Colin Cherry, Christopher A. Choquette-Chop Aakanksha Chowdhery, Clément Crepy. Shachi Dave, Mostafa Dehghani. Sunipa Dev, Jacob Devlin, Mark Díaz, Nan Du, Ethan Dyer, Vlad Feinberg. Fangxiaoyu Feng, Vlad Fienber, Markus Freitag, Xavier Garcia, Sebastian Gehrmann. Lucas Gonzalez, Guy Gur-Ari, Steven Hand, Hadi Hashemi, Le Hou, Joshua Howland, Andrea Hu, Jeffrey Hui. Jeremy Hurwitz, Michael Isard, Abe Ittycheriah, Matthew Jagielski, Wenhao Jia, Kathleen Kenealy, Maxim Krikun, Sneha Kudugunta, Chang Lan, Katherine Lee, Benjamin Lee, Eric Li, Music Li, Wei Li, YaGuang Li, Jian Li, Hyeontack Lim, Hanzhao Lin, Zhongtao Liu. Frederick Liu. Marcello Maggioni, Aroma Mahendru, Joshua Maynez, Vedant Misra, Maysam Moussalem, Zachary Nado, John Nham, Eric Ni, Andrew Nystrom, Alicia Parrish, Marie Pellat. Martin Polacek, Alex Polozov, Reiner Pope, Siyuan Qiao. Emily Reif, Bryan Richter, Parker Riley, Alex Castro Ros, Aurko Roy, Brennan Saeta, Rajkumar Samuel, Renee Shelby, Ambrose Slone, Daniel Smilkov, David R. So, Daniel Sohn, Simon Tokumine, Dasha Valter, Vijay Vasudevan, Kiran Vodrahalli, Xuezhi Wang. Pidong Wang. Zirui Wang, Tao Wang, John Wieting, Yuhuai Wu, Kelvin Xu, Yunhan Xu, Linting Xue, Pengcheng Yin, Jiahui Yu, Qiao Zhang, Steven Zheng, Ce Zheng, Weikang Zhou, Denny Zhou, Slav Petrov, and Yonghui Wu. Pal. 2 technical report, 2023. 6,
[4] Brenna D Argall, Sonia Chernova, Manuela Veloso, and Brett Browning. A survey of robot learning from demonstration. Robotics and autonomous systems, 57(5) 483, 2009.
[5] Kushal Arora, Layla El Asri, Hareesh Bahuleyan, and Jackie Chi Kit Cheung. Why exposure bias matters: An imitation learning perspective of error accumulation in language generation. In Findings, 2022. 2
[6] Gregor Bachmann and Vaishnavh Nagarajan. The pitfalls of next-token prediction. ArXiv, abs/2403.06966 06666, 2024. 2
[7] Yuntao Bai, Andy Jones, Kamal Ndousse, Amanda Askell, Anna Chen, Nova DasSarma, Dawn Drain, Stanislav Fort, Deep Ganguli, Tom Henighan, et al. Training a helpful and harmless assistant with reinforcement learning from human feedback. arXiv preprint arXhv:200.05886. 2022. 2
[8] Michael Bain and Claude Sammut. A framework for behavioural cloning. In Machine Intelli- gence 15, pages 103-129, 1995. 2, 10
[9] Matt Barnes, Matthew Abueg, Oliver Lange. Matt Deeds, Jason Trader, Denali Molitor, Markus Wulfmeier, and Shawn O' Banion. Massively scalable inverse reinforcement learning in google maps. In The Twelfth International Conference on Learning Representations, 2024. 1, 9
11

---

[10] Dan Biderman, Jose Gonzalez Ortiz, Jacob Portes, Mansheej Paul, Philip Greengard, Connor Jennings, Daniel King, Sam Havens, Vitaliy Chiley, Jonathan Frankle, Cody Blakeney, and John P. Cunningham. Lora learns less and forgets less, 2024. 6
[11]
Massimo Caccia, Lucas Caccia, William Fedus, Hugo Larochelle, Joelle Pineau, and Laurent Charlin. Language gans falling short. arXiv preprint arXiv:1811.02545 2018. 10
[12] Stephen Casper, Xander Davies, Claudia Shi, Thomas Krendl Gilbert, J'yr'emy Scheurer, Javier Rando, Rachel Freedman, Tomasz Korbak, David Lindner, Pedro Freire. Tony Wang. Samuel Marks. Charbel-Raphaél Ségerie, Micah Carroll, Andi Peng, Phillip J. K. Christoffersen, Mehul Damani, Stewart Slocum, Usman Anwar, Anand Siththaranjan, Max Nadeau, Eric J. Michaud, Jacob Pfau, Dmitrii Krasheninnikov, Xin Chen, Lauro Langosco di Langosco, Peter Hase, Erdem Biyik, Anca D. Dragan, David Krueger, Dorsa Sadigh, and Dylan Hadfield-Ment Open problems and fundamental limitations of reinforcement learning from human feedback. ArXiv, abs/2307.1522 2023.
[13] Paul Christiano, Jan Leike, Tom B Brown, Miljan Martic, Shane Legg, and Dario Amodei. Deep reinforcement learning from human preferences. arXiv preprint
2017.
[14] Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert. Jenry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems. arXiv preprini 10.v:2110.141660 2021. 6,7
[15] Chris Cundy and Stefano Ermon. Sequencematch: Imitation learning for autoregressive sequence modelling with backtracking, 2024. 6, 8, 10, 16
[16] Anna Dawid and Yann LeCun. Introduction to latent variable energy-based models: A path towards autonomous machine intelligence. ArXiv, abs/2306.02572, 2023. 2
[17] Gemini Team et al. Gemini: A family of highly capable multimodal models, 2024. 1
[18] Chelsea Finn, Sergey Levine, and P. Abbeel. Guided cost learning: Deep inverse optimal control 2016.
via policy optimization. ArXiv,
[19] Justin Fu, Katie Luo, and Sergey Levine. Learning robust rewards with adversarial inverse reinforcement learning. arXiv preprint arXiv:1710.11244 2017. 9
[20] Divyansh Garg, Shuvam Chakraborty, Chris Cundy, Jiaming Song, Matthieu Geist, and Stefano Ermon. Iq-learn: Inverse soft-q learning for imitation, 2022. 2, 3, 4, 5, 9, 16
[21] Shivam Garg, Dimitris Tsipras, Percy Liang, and Gregory Valiant. What can transformers learn in-context? a case study of simple function classes. Advances in Neural Information Processing Systems, 35:30583-30599 2022. 16
[22] Matthieu Geist, Bruno Schemer, and Olivier Pietquin. A theory of regularized markov decision processes. CoRR, abs/1901.112755 2019. 4
[23] Seyed Kamyar Seyed Ghasemipour, Richard S. Zemel, and Shixiang Gu. A divergence mini- mization perspective on imitation learning methods. CoRR, 0tx/1911.02255. 2019. 2, 3
[24] Joey Hejna and Dorsa Sadigh. Inverse preference learning: Preference-based without reward function. ArXiv, abs/2305. 153633, 2023. 10
[25] Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. Advances in neural information processing systems, 29:4565-4573 2016. 2, 3, 5,
[26] Luca Hormann and Artem Sokolov. Fixing exposure bias with imitation learning needs powerful oracles. arXiv preprint arXiv:2 109, 099.0441. 14, 2021. 10
[27] Edward J. Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models, 2021.
[28] Rebekka Hubert, Artem Sokolov, and Stefan Riezler. Improving end-to-end speech translation by imitation-based knowledge distillation with synthetic transcripts. In Int, Workshop on Spoken Language Translation, 2023. 10
12

---

[29] Ahmed Hussein, Mohamed Medhat Gaber, Eyad Elyan, and Chrisina Jayne. Imitation learningg survey of learning methods. ACM Comput. Surv., 50(2), apr 2017.
[30] Ziwei Ji, Nayeon Lee, Rita Frieske, Tiezheng Yu, Dan Su. Yan Xu, Etsuko Ishii, Ye Jin Bang. Andrea Madotto, and Pascale Fung. Survey of hallucination in natural language generation. ACM Comput. Surv. 55(12), mar 2023. 2
[31] Timo Kaufmann, Paul Weng, Viktor Bengs, and Eyke Hüllermeier. Survey Reinforcement Learning from Human Feedback, 2023. 9
[32] Pei Ke, Fei Huang, Minlie Huang, and Xiaoyan Zhu. Araml: A stable adversarial training framework for text generation. arXiv preprini urXiv:/908.000 07908, 2019. 10
[33] Beon-Hyeong Kim, Seokin Seo, Jongmin Lee, Wonseok Jeon, HyeongJoo Hwang, Hongseok Yang, and Kee-Eung Kim. Demodice: Offline imitation learning with supplementary imperfect demonstrations. In International Conference Learning Representations, 2022. 10, 16
[34] Tom Kocmi, Rachel Bawden, Ondřej Bojar, Anton Dvorkovich, Christian Federmann, Mark Fishel, Thamme Gowda, Yvette Graham, Roman Grundkiewicz, Barry Haddow, Rebecca Knowles, Philipp Koehn, Christof Monz, Makoto Morishita, Masaaki Nagata, Toshiaki Nakazawa, Michal Novák, Martin Popel, and Maja Popović. Findings of the 2022 Conference on Machine Translation (WMT22). In Proceedings of the Seventh Conference on Machine Translation (WMT), 2022. 6.
[35] Iya Kostrikov, Kumar Krishna Agrawal, Debidatta Dwibedi, Sergey Levine, and Jonathan Tompson. Discriminator-actor-noo Addressing sample inefficiency and reward bias in adversarial imitation learning. In International Conference on Learning Representations, 2019.
9
[36] Ilya Kostrikov, Ofrr Nachum, and Jonathan Tompson. Imitation learning via off- toflcol distribu- tion matching. arXiv preprint arXiv: 191200033. 2019. 9
[37] Harrison Lee, Samrat Phatale, Hassan Mansoor, Thomas Mesnard, Johan Ferret, Kellie Lu, Colton Bishop, Ethan Hall, Victor Carbune, Abhinav Rastogi, and Sushant Prakash. Rlaif: Scaling reinforcement learning from human feedback with ai feedback. arXiv, abs/2309.00267, 2023. 5
[38] Jiaxiang Li, Siliang Zeng, Hoi-To Wai, Chenliang Li, Alfredo Gareía, and Mingyi Hong. Getting more juice out of the sft data: Reward learning from human demonstration improves sit for 11m alignment. ArXiv, abs/2405.1788 17888, 2024. 10
[39] Ziniu Li, Congliang Chen, Tian Xu, Zeyu Qin, Jiancong Xiao, Ruoyu Sun, and Zhimin Luo. Entropic distribution matching in supervised fine-tuning of 11ms: Less overfitting and better diversity, 2024. 10
[40] Alexander Lin, Jeremy Wohlwend, Howard Chen, and Tao Lei. Autoregressive knowledge distillation through imitation learning. arXiv preprint arXiv:2009.0722 2020. 10
[41] Yecheng Ma, Andrew Shen, Dinesh Jayaraman, and Osbert Bastani. Versatile offline imitation from observations and examples via regularized state-occupancy matching. In International Conference Machine Learning. PMLR, 2022. 10, 16
[42] Ofir Nachum, Mohammad Norouzi, Kelvin Xu, and Dale Schuurmans. Bridging the gap between value and policy based reinforcement learning. In Advances in Neural Information Processing Systems, 2017. 4
[43] Shashi Narayan, Shay B. Cohen, and Mirella Lapata. Don't give the details, just the summary! topic-aware convolutional neural networks for extreme summarization. ArXiv, abs/ Abs/1808.08845, 2018. 6
[44] Andrew Y. Ng and Stuart J. Russell. Algorithms for inverse reinforcement learning. In Proceedings of the Seventeenth International Conference on Machine Learning, ICML '00, page 663-670, San Francisco, CA, USA, 2000, Morgan Kaufmann Publishers Inc. 9
[45] OpenAl. Gpt-4 technical report, 2023, 1
13

---

[46] Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida. Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et. Training language models to follow instructions with human feedback. Advances in neural information processing systems, 35:27730-27744, 2022. 1,
[47] Krishna Pillutla, Lang Liu, John Thickstun, Sean Welleck, Swabha Swayamdipta, Rowan Zellers, Sewoong Oh, Yejin Choi, and Zaïd Harchaoui. Mauve scores for generative models: Theory and practice. ArXiv, abs/2212.14578, 2022. 17
[48] Bilal Piot, Matthicu Geist, and Olivier Pietquin. Boosted and reward-regularized classifica- tion for apprenticeship learning. In Proceedings of the 2014 International Conference on Autonomous Agents and Multi-Agent Systems, AAMAS 14, page 1249-1256, Richland, SC, 2014. International Foundation for Autonomous Agents and Multiagent Systems. 4,
[49] Dean Pomerleau. Alvinn: An autonomous land vehicle in a neural network. Advances in neural information processing systems, 1, 1988. 9
[50] Rafael Rafailov, Joey Hejna, Ryan Park, and Chelsea Finn. From to q*: Your language model secretly a q-function, 2024. 10
[51] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer, 2023. 2, 6
[52] Mare' Aurelio Ranzato, Sumit Chopra, Michael Auli, and Wojeiech Zaremba. Sequence level training with recurrent neural networks, 2016. 2
[53] Noam Razin, Hattie Zhou, Omid Saremi, Vimal Thilak, Arwen Bradley, Preetum Nakkiran, Joshua Susskind, and Etai Littwin Vanishing gradients in reinforcement finetuning of language models. arXiv preprint arXiv:2310.20703 2023. 10
[54] Siddharth Reddy, Anca D. Dragan, and Sergey Levine. Sqil: Imitation learning via reinforcement learning with sparse rewards, 2019. 10
[55] Juntao Ren, Gokul Swamy, Zhiwei Steven Wu, J. Andrew Bagnell, and Sanjiban Choudhury. Hybrid inverse reinforcement learning. ArXiv, 000b/2402.08888, 2024. 19
(56) Stéphane Ross, Geoffrey Gordon, and Drew Bagnell. A reduction of imitation learning and structured prediction to no-regret online learning. In Proceedings of the fourteenth interna- tional conference on artificial intelligence and statistics, pages 627-635. JMLR Workshop and Conference Proceedings, 2011. 1. 2, 8,9
[57] Harshit Sikehi, Qinqing Zheng, Amy Zhang, and Scott Niekum. Dual rl: Unification and new methods for reinforcement and imitation learning. In International Conference on Learning Representations, 2023. 4, 10, 16
[58] Nisan Stiennon, Long Ouyang, Jeffrey Wu, Daniel Ziegler, Ryan Lowe, Chelsea Voss, Alec Radford, Dario Amodei, and Paul Christiano. Learning to summarize with human feedback. Advances in Neural Information Processing Systems, 2020. 6,
[59] Hao Sun. Supervised fine-tuning inverse reinforcement learning. arXiv preprint 24Xiv:2403.1 /2017, 2024. 10
[60] Umar Syed and Robert E. Schapire. A game-theoretic approach to apprenticeship learning. In Advances in Neural Information Processing Systems, 2007. 9
[61] Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint wwXi::2307.09288 2023. 2
[62] Luca Viano, Angeliki Kamoutsi, Gergely Neu, Igor Krawczuk, and Volkan Cevher. Proximal point imitation learning. In Advances in Neural Information Processing Systems, 2022.
[63] Joe Watson, Sandy H. Huang, and Nicolas Heess. Coherent soft imitation learning. In Advances in Neural Information Processing Systems (NeurIPS), 2023. 10
14

---

[64] Ronald J. Williams and David Zipser. A learning algorithm for continually running fully recurrent neural networks. Neural Compulation, 1:270-280, 1989. 2
[65] Qingyang Wu, Lei Li, and Zhou Yu. Textgail: Generative adversarial imitation learning for text generation. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, 2021.
6, 10, 19
[66] Markus Wulfmeier, Arunkumar Byravan, Sarah Bechtle, Karol Hausman, and Nicolas Man- fred Otto Heess. Foundations for transfer in reinforcement learning: A taxonomy of knowledge modalities. ArXiv, abs/2312.01939, 2023.
[67] Markus Wulfmeier, Ingmar Posner, and Abbeel. Mutual alignment transfer learning. ArXiv, (obs/170..00 079077 2017. 5
[68] Markus Wulfmeier, Dominic Zeng Wang, and Ingmar Posner. Watch this: Scalable cost-function learning for path planning in urban environments. 2016 IEEF/RSJ International Conference on Intelligent Robots and Systems (IROS), pages 2089-2095, 2016.
[69] Yue Xing, Qifan Song, and Guang Cheng. On the algorithmic stability of adversarial training. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and Wortman Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 26523-26535 Curran Associates, Inc., 2021.
[70] Tian Xu, Ziniu Li, and Yang Yu. Error bounds of imitating policies and environments. Advances in Neural Information Processing Systems, 2020.
[71] Lantao Yu, Weinan Zhang, Jun Wang, and Yong Yu. Seqgan: Sequence generative adversarial nets with policy gradient. In Proceedings of the AAAI conference on artificial intelligence, volume 31, 2017. 10
[72] Tianyi Zhang, Varsha Kishore, Felix Wu, Kilian Q Weinberger, and Yoav Artzi. Bertscore: Evaluating text generation with bert. ArXiv, abs/1 htss/994.09675. 2019. 17
[73] Yaoming Zhu, Sidi Lu, Lei Zheng, Jiaxian Guo, Weinan Zhang, Jun Wang, and Yong Yu. Texygen: A benchmarking platform for text generation models. In The 4 international ACM SIGIR conference research & development in information retrieval, pages 1097-1100, 2018.
[74] Brian D Ziebart. Modeling Purposeful Adaptive Behavior with the Principle of Maximum Causal Entropy. PhD thesis, Carnegie Mellon University, 2010.
[75] Brian Ziebart, Andrew L Maas, J Andrew Bagnell, Anind K Dey, et al. Maximum entropy inverse reinforcement learning. In Aaai, volume 8, pages 1433-1438. Chicago, II, USA, 2008
15

---

A Appendix
A.1
Implementation Details
A.1.1
IQLearn
For IQLearn, fine-tuning starts with policy which is subsequently finetuned Q-function. In order to do so, we take the logits underlying the LLM softmax layer and continue training as Q-values. In entropy regularized RL, we can build on the equality between optimal policy л and optimal Q- function Q" via x (a|s) 1/z, exp Q" (s, a) with normalization factor Z, Eo'eA exp Q (s. a'). We further obtain V via V(s) log EaE exp ((s,a)
Due to the translation invariance of the softmax function, the there can be a considerable state-based offset between initial logits and final values. In other words, with Q(s, a) V(s) + A(s, a) only the advantage function A has to be accurately reflected by the initial policy logits, while the state- based value V can be arbitrarily inaccurate. separate the representation of state and advantage function, to stabilize training during the additional identification of correct offsets but found it to be unnecessary for our experiments in comparison to recent work [15].
practice, we can add further KL regularization, or
We handle terminal states by setting their values to zero in line with the original IQLearn paper [20]. Improvements such as learned values for the terminal states [2, 15] did not contribute to improved performance in this setting.
Table 3: IQLearn and MLE hyperparameters
Hyperparameter
Value
Learning rate T5
lc-4
Learning rate PaI.M2
le-4
Warmup steps
2000
Batch size T5 (base/large/xl)
32/32/16
Batch size PaLM2 Random seeds experiment
16
A.1.2 Online IQLearn
We also implemented an online version of IQLearn, which makes use of additional (non-expert) samples. In contrast to IQLear [21] which makes use of these examples to estimate the initial value, Epo v(s) (which is equivalent between online and offline data when the prompt the same as in our case), we integrate the additional samples to relax the distribution matching loss:
min Dy((1 - a) + op ())L ouB) (H()..
(15)
where un is the discounted state-action distribution of the additional samples and where a E [0, 1) is the strength of the mix-in. The problem is thus relaxed and allows the policy to match the mix-in distribution P& in case the expert is too difficult to match.
From the adapted distribution matching loss the same steps can then taken as Section 2. This results in the generalised loss:
min F(l-a)uetapall ( x²)) ((",*")) E log(("))
(16)
pr
with
r("",=)) (T) =
1 log n. E
(17)
This formulation strongly relates to the offline version of the algorithm but additionally applies the temporal difference based regularization term on the additional non-expert transitions. A similar algorithm is derived in Appendix D.1 [57] with strong results on classical, continuous control domains of its state-action value function based version. It further connects prior methods enabling the integration of sub-optimal data into imitation learning [41, 33] with a useful overview on the broader field of related methods in [57].
16

---

A.1.3 GAIL
For GAIL,, both the policy and discriminator are represented via separate networks initialized from the initial, pre-trained, LLM. Policy optimization is performed similar to PPO with an A2C update, with re-scaled advantage and KL constraint to the initial policy. The KL constraint has a weight hyperparameter that is annealed over 10,000 steps to final cost weight displayed below. The value network is also initialized from the initial LLM. The discriminator is trained with cross- entropy objective. The reward is re-shaped from the the discriminator output to be a positive reward: = logg1 + exp(D(s:))).
We explore the heuristic combination of GAIL with standard MLE training by using weighted combination of GAIL and MLE losses for the policy, While this not required for GSM8k, it leads to considerable improvements for XSUM.
Policy, value function and discriminator are all updated with the Adam optimizer, with a constant learning rate and linear warm-up. The discriminator is updated after every step of policy optimization.
Table 4: GAIL hyperparameters
Method
T5- T5-base
T5-large
T5-xl
Batch size
32
32
16
Learning rate
le-4
le-4
le-4
Warmup steps
2000
2000
2000
KL strength
1e-3
1e-3
1e-3
Random seeds / experiment
3
3
A.2 MLE & Entropy Regularization
As mentioned in the experiments section, we compare our methods to an entropy regularized version of MLE, to disentangle between the imitation contibution and the regularization. This algorithm simply follows the objective
min E log(T) XH(()),
(18)
where we compute the entropy at each token of the sequences form Mgg
A.2.1 Computational Requirements
Our experiments with T5 models use TPU v3 infrastructure and are running between approximately days and weeks. Our experiments with PaLM2 models use TPU v4 infrastructure and are running under week.
A.3 Additional Experiments
We include a set of further experiments to complement the results in the main paper.
A.3.1 Quality-Diversio Evaluation
In addition to the Self-BLEU metric, we further add plots for model entropy and add the ROUGE- LSUM results in Figure 9. Further performance metrics like MAUVE [47] and BertScore [72] can be of use in the future to further represent human judgement and preferences.
A.3.2 PALM2 Additional Results
We complement results on the temperature sweep over PALM2 models in Figure 10, with additional metrics for GSM8k and TLDR (accuracy with calculator and rouge scores). This confirms the results of the main paper experiment, showing that IQLearn consistently outperform MLE in accuracy, even when sampling with a lower temperature than the training one.
A.3.3 Toy Experiments for Offfine and Online IRL with Autoregressive Generation
The toy scenario displayed by the MDP in Figure 11 is used to represent a key difference between many classical control settings and autorcgressive language generation. The agent starts in the left
17

---

GAIL rent-reg-0.00
40
(ent-reg=0001
GAIL (ent-reg=0.1)
Seff BLEU
50
IQLearn (-0.11)
CQLearn (A-055)
60
(A-10)
MLE (ent-reg=0.0
70
MLE (ent-reg=0.01)
(ent-reg-0.03)
80
MLE (ent-reg=0.1)
Base
Large
90
28.5
29.0
29.5
30.0
30.5
31.0
ROUGE-LSUM
1.4
GAIL (ent-reg-0.0)
Avg. per token entropy
1.2
GAIL (ent-reg=0.01
1.0
GALL (ent-reg=0.1)
IQLearn (1-
0.8
(A-0.5)
0.6
IQLearn (=110)
MLE (ent-reg=0.0)
0.4
MLE (ent-reg=0 01)
MLE (ent-reg-0.05
MLE (ent-reg=0.1)
0.0
Base
Large
28.5
30.0
30.5
31.0
ROUGE-LSUM
Figure 9: XSUM results for models trained with MLE, IQLearn, and GAIL across different regular- ization strengths. Top: we show ROUGE-LSUM performance metric on the x-axes and Self-BLEU diversity measure on the y-axis. Bottom: with the per-token entropy as diversity metric. Error bars indicate the standard error of the mean after repeating the experiment with 3 different seeds. Compare to Figure for more results.
KLsamdi
NLC
Figure 10: Additional scores against temperature for PALM2 models. From left to right: GSM8k,
TDLR (ROUGE-1) and TLDR (ROUGE-2).
black state and has to reach the green on the right. each the bottom states the agent has actions, move left move up. The transition dynamics noisy that the agent executes the unwanted action 10% of the time. There are two variants of the MDP, one where the agent can return from the top states by learning to execute the right action and one where it cannot. The latter represents one aspect of concatenation dynamics, the agent cannot return to the exact same state after sampling the wrong action.
way, cannot correct its behavior exactly.
When training offline and online variants IQLearn in this setting with demonstrations without mistakes (and their correction), we clearly see that the online version of IQLearn outperforms offline learning by over 1 11% in success rates, while in the setting without recovery, the difference is considerably smaller and results are within each others confidence bounds with more variance for the offline agent.
A.3.4 Analyzing GAIL Stabilization
We empirically find GAIL overall more complex to tune and stabilize and aim provide further insights here. Intuitively, control for language modelling differs from many classical applications via large discrete action space and terminations not being state but action conditioned. In other words, the agent can directly chose to terminate. Therefore the agent can learn to directly terminate if
18

---

Figure 11: Simplified MDP to represent characteristics of concatenation-base autoregressive genera- tion comparison to many classical control domains. Dashed lines visualize the potential to return or correct mistakes, missing from autoregressive generation. The agent starts in the black state and has to complete the sequence to reach green to receive rewards.
the discriminator provides negative rewards. If we provide positive rewards, tuning those rewards becomes complex task seen in Figure 12, where without additional MLE objective the GAIL agent often uses the maximum sample length with correlated loss performance (visualized by the drop on ROUGE-1 values). Additional MLE training can help to both shape policy behavior but further also provide more relevant data for the discriminator [55]. Especially in high dimensional action spaces it otherwise becomes challenging to obtain a useful reward signal from the discriminator as seen in prior work [65].
mil__eeight = 0.0
mie weight 1.0
36
29
ROUSE1
20
GALLL (wttlesg-0.3) 00)0
GALL (ontrrog-0.05
(ent-rog-0.00
10
50000
100000
150000
200000
50000
100000
150000
200030
mie weight - 0.0
mie weight 1.0
Peg. number tekens
50
GAL (entrreg 0.05)
GAL (ont (ontirog-.100
40
eentirag-0.00
50000
100000
150000
200000
50000
100000
150000
200000
Figure 12: Effect of adding a standard MLE loss (mle weight=1) on the training data combined with GAIL. XSUM. We show the ROUGE metric and the average length of the generated summaries when training T5-Large model with GAIL.
A.4 Expanded Experiments on WMT22
For WMT22, we additionally evaluate IQLearn performance for a wide range of À values. Here, eQLearn consistently gains (up to +2. BLEU) over BC until very high values of À (Table 5). Note, that unlike the sampling performance curves in Figure 4, here beam search decoding (size used with length penalty 0.6.
19

---

Table 5: WMT22 results for offline eeaa initialised with PaLM2 checkpoint. Italic - best dev BLEU in group (i.e. same mix value), bold best overall.
A
mixin
dev-BLEU
test-BLEU
0.0
0.0
26.92
32.34
0.05
0.0
29. 14
34.84
0.1
0.0
28.93
34.51
0.3
0.0
29.07
34.79
0.5
0.0
29.20
34.63
0.7
0.0
28.89
35.68
0.9
0.0
28.89
34.35
1.0
0.0
28.92
33.69
0.0
0.1
27.43
32.77
0.05
0.1
29.43
34.50
0.1
0.1
29.23
34.56
0.3
28.89
34.84
0.5
0.1
28.95
34.39
0.7
0.1
28.78
34.17
0.9
0.1
28.83
34.41
1.0
0.1
28.73
34.29
0.0
0.2
27.10
32.31
0.05
0.2
29.08
34,40
0.1
0.2
29.46
34.54
0.9
0.2
28.85
33.99
20

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
The abstract provides key information about methods and results.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
Parts of the discussion section explicitly discuss limitations of the performed
work.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification: All methods are properly derived.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Answer: Yes]
Justification: Experimental recreation of the results might be limited by required computa- tional resources. Other than that, algorithms are accurately described including hyperparam- eters and the T5 model classes should enable the recreation of results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Answer: [No]
Justification:
We have been unable to provide public access our code.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: All experimental detailed information either given in the main paper or more likely in the appendix.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Ye]]
Where computationally possible, the experiments include error bars and related information.
Justification:
8. Experiments Compute Resources
Question:
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Answer: [Yes]
Justification: Approximate details given in the appendix
21

---

9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics stttps//neurr/
Answer: [Yes]
Justification: The NeurIPS Code of Ethics has been read by the authors and all applicable areas have been verified.
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: This paper includes fundamental research on imitation learning with respect to the application to language modelling and does not itself introduce new risks and concerns in addition to those related to existing optimisation methods for language models.
11. Safeguards
Question: Does the paper deseribe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification: No data model are publicly released.
12. Licenses for existing assets
Question:
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: References are included for all datasets and models.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: We release no new assets.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: The paper does not involve crowdsourcing or research with human subjects. 15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not include research with human subjects
22


Reviews:
[{"review_id": "dOJwlonjHI", "reviewer": "Reviewer_tAhh", "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. \n\nOn 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. \nThe 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.", "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).", "limitations": "No concerns.", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "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. \n\n* 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the constructive feedback. By addressing these helpful comments and questions, the paper has improved and gained in clarity. Please mention further questions or clarifications and we will work to address them ASAP.\n\nThe included experiments show that IRL methods (GAIL/IQLearn) improve both performance and, crucially, diversity metrics, compared to common next token likelihood. The performance improvements are small for XSUM but consistent across domains. If the reviewer could specify what additional information or results would be needed to fulfill specific additional criteria for 'actually better models', we would be glad to provide further clarification. In addition, in the paper we have clarified the exact metrics used, and how to interpret the results. In particular for Self-BLEU, we will expand the captions to emphasize that smaller numbers indicate better diversity (lower self-similarity) and that the y-axis is flipped.\n\nFurthermore, if the point above is meant in relation to the comment “The empirical part of the paper lacks focus on what is trying to be optimised”, we can clarify. Section 2 describes the optimization (aspects such as the relation of reverse KL to MLE, regularized chi2 for Iq-learn). However, evaluating these divergences once trained is hard, because we do not have access to the distribution underlying the dataset. In classic imitation learning papers, the underlying ground truth reward function can be used for evaluating the quality of the imitating agent (typically, the expert is an RL agent trained on some known reward). Yet, in our case, we do not have such a reward function, and the expert is a human. We therefore have to rely on these proxy metrics, and we consider classic ones in the context of language models.\n\nIn parallel to clarifying this point, we have expanded our experiments to include results on WMT22 (with results in the general rebuttal). Furthermore, IRL-based models have the added benefit of extracting reward information from demonstration data. We have added further analysis on the correlation between extracted rewards and performance metrics to visualize this advantage, with exact numbers added in the general rebuttal.\n\nRegarding the question of most relevant tasks for IRL, we have emphasized the theoretic connection to target/trajectory length known through the imitation learning literature (e.g. DAgger[1]). Practically, the performance gains are stronger with longer target lengths (GSM8k) and partially with smaller datasets (ablations on smaller subsets of XSUM), while improved diversity of generations is persistent across tasks. We have expanded our discussion of these aspects in the paper. \n\nMinor points:\n- We will introduce the used datasets and tasks earlier in Section 3.\n- RLHF: We do not provide an alternative to RLHF but rather emphasize the RL perspective during SFT, which could in the future enable better connections between SFT and RLHF data sources. As both method types use different data sources, a direct, fair comparison between IRL-based approaches and RLHF is not possible. Instead, these are complementary.\n\n[1] A Reduction of Imitation Learning and Structured Prediction to No-Regret Online Learning, AISTATS 2011, Stéphane Ross, Geoffrey J. Gordon, J. Andrew Bagnell", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "alOe4Ww3Dr", "reviewer": "Reviewer_aRfw", "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.", "questions": "- Could the authors provide an analysis for quality-diversity trade-off [1] of the algorithms?\n- Why does the temporal regularization term encourages higher diversity?\n- It would be helpful if the paper contained pseudo code of the proposed algorithm.\n\n[1] LANGUAGE GANS FALLING SHORT", "limitations": "see weaknesses", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "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.\n- It is not clear the proposed method is better than MLE in terms of quality-diversity trade-off.\n- The diversity evaluation for lower open-endedness tasks like GSM8K does not make much sense.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the constructive feedback. By addressing these helpful comments and questions, the paper has improved and gained in clarity. Please mention further questions or clarifications and we will work to address them ASAP.\n\nOne of the review’s key points regards limited performance improvements on summarization. Here, we would like to emphasize that the goal of our paper is to establish in which tasks IRL-like learning is most impactful. Overall, the performance gains are stronger across the other domains, and in summarization we additionally see considerably improved diversity of generations. We have further emphasized the limitations of MLE/BC with respect to long target/trajectory lengths known from the imitation learning literature (e.g. DAgger [A]).\n\nTo further improve the empirical evaluation, we have added results on WMT22 where offline IQ-Learn shows increased performance in comparison to MLE-based training (with results in the general rebuttal). Finally, a key benefit for IRL-based models lies in additionally extracting reward information. To explore this, we have added analysis for the correlation between extracted rewards and performance metrics, with results summarized in the general rebuttal.\n\nRegarding the connection between IRL and higher diversity, equation 3 in the paper describes this best. The original state-distribution-matching formulation builds on entropy-regularized RL and in particular includes a term for causal entropy. We have further emphasized this connection in the paper. \n\nThank you in particular for enabling us to strengthen the connection to [1]. Our temperature sweeps as well as the sweep over different entropy regularisations enable the generation of the kind of quality diversity evaluations that [1] is arguing for. While we move away from purely adversarial inverse RL and expand to saddle-point-based methods like IQLearn, their insights are highly relevant and we have clarified additional heuristics required to enable better performance for our implementation of GAIL (which are already included in the appendix but were previously skipped in the main paper).\n\nMinor:\n- We add pseudo code to the appendix to minimize required space in the main paper.\n- We further expanded our discussion of related work to include the suggested paper.\n\n[A] A Reduction of Imitation Learning and Structured Prediction to No-Regret Online Learning, AISTATS 2011, Stéphane Ross, Geoffrey J. Gordon, J. Andrew Bagnell", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "iCxZvdlwuf", "reviewer": "Reviewer_6J5H", "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.", "questions": "1.\tThe original GAIL is an online IRL method, does GAIL execute offline or online in the experiments?", "limitations": "The discussion section discussed the potential limitations of this work.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1.\tThis 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. \n\n2.\tIQLearn 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.\tIRL methods, even for offline IRL methods, cost more time than MLE for LLM tunning. \n\n2.\tMore experimental details are required to reproduce the results.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the constructive feedback. By addressing these helpful comments and questions, the paper has improved and gained in clarity. Please mention further questions or clarifications and we will work to address them ASAP.\n\nThank you for pointing out computational requirements. We crucially reduce computational costs by training offline and thus removing the additional sampling cost for IQLearn in all sections but the online to offline comparison in section 3.3.1. The remaining difference in computational costs are partially due to a non-optimized implementation in our codebase, which shares many components between offline and online IRL algorithms. We have further emphasized the implementation dependence in the paper. We have started to recreate a pure and clean offline implementation but will be unable to verify results with the new implementation before the rebuttal ends.\n\nWe have added more general details on the experiments. We would be glad to provide further clarifications and experiment details if the reviewer could specify what additional information or results would be needed from their perspective.\n\nWe have further expanded the results, adding WMT22 where offline IQ-Learn shows increased performance in comparison to MLE-based training (with results in the general rebuttal). Finally, a key benefit for IRL-based models lies in additionally extracting reward information. To explore this, we have added analysis for the correlation between extracted rewards and performance metrics, with results summarized in the general rebuttal.\n\nMinor points:\n- GAIL is an online method that requires training the discriminator between the online agent and offline demonstration dataset. We have emphasized this in the paper.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Elrtyktwel", "reviewer": "Reviewer_UqkB", "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.", "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.\n* 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?\n* How do the conclusions of Figure 4 hold if Lambda is varied? What is the expectation?\n* Figure 6: can you quantify the uncertainty estimate used?\n* 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.", "limitations": "These were adequately discussed in Section 5.", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "* Clarity: The introduction is clear and well motivates the problem. Overall the paper has good clarity, bar a few typo's.\n* 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.\n* 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.\n* Minor: The first paragraph could benefit from references to ground the reader.\n* 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.\n* Minor: Equation 8 could put an enclosing bracket for the min operator to help guide the reader.\n* Minor: The presentation of the figures could be improved by making the font size larger and using vector graphics.\n\nTypos:\n\n* L40: “behavior cloning [7], its” -> “behavior cloning [7], and its”\n* L123: “further simplify” -> “further simplify it”\n* L168: “temporal different regularization” -> “temporal difference regularization”\n* L174: “we a short” -> “we use a short”\n* L174: “pure MLE beneficial” -> “pure MLE being beneficial”\n* L:202: “full trajectories” -> “full of trajectories”\n* L:218: “Right” -> “Left”", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the constructive feedback. By addressing these helpful comments and questions, the paper has improved and gained in clarity. Please mention further questions or clarifications and we will work to address them ASAP.\n\nUncertainty estimates and error bars: We always plot the standard error (standard-deviation / sqrt(number-of-seeds) in all plots. We added a comment to the paper explaining our uncertainty estimation. Figures 2 & 3 already include error bars; these are very small due to only minimal variability from the random seeds. We updated the color of these error bars to make them more legible. Due to computational requirements, we cannot add further experiments to add error bars for Figure 4 before the rebuttal deadline, but have added it to our prioritization list to follow-up when other experiments are completed.\n\nWe strongly acknowledge the benefits of extending analysis and a broader task set. While experiments, in particular with hyperparameter sweeps and multiple seeds are expensive, we have added results on WMT22 training (with results in the general rebuttal). We continue looking into further experiments that can be performed with limited costs and have, in parallel, extended our analysis to obtain further insights without additional training runs. For this, we have investigated the rewards extracted by IRL as described in the general rebuttal comment.\n\nRegarding the extension to preference learning, this is a very good point and we quickly discuss opportunities in the discussions section. A key benefit is the ability to extract reward information from both demonstration and preference datasets, which could bring benefits regarding the robustness of RLHF. While these are untested for language modeling, there exist early examples for classic robotics tasks [1].\n\nAblations on \\lambda to interpolate between MLE and IQLearn are of crucial interest and included (with limited settings), for example in Figure 2. We know from early tests that much smaller values critically reduce performance and \\lambda = 0 algorithmically reduces to pure MLE. However, since these experiments are computationally expensive we are unable at this point to add a more detailed study.\n\nSimilarly, we made the decision to sweep over the quality-diversity tradeoff for PaLM2 models via inference time settings, rather than separate training experiments, to minimize computational costs and provide better coverage of different mechanisms for the tradeoff. We agree this would be interesting to investigate, but have to rely on the related sweeps for Figures 2 and 3 for the time being.\n\nMinor points:\n- The T5 models are in fact pretrained. Thank you for pointing this out; we have expanded the model discussion in the paper and clarified this aspect.\n- We have fixed the remaining spelling and other mistakes, as well as added further clarifications as suggested. Thank you for pointing these out.\n\n[1] Learning Reward Functions from Diverse Sources of Human Feedback: Optimally Integrating Demonstrations and Preferences, 2020, Erdem Bıyık, Dylan P. Losey, Malayandi Palan, Nicholas C. Landolfi, Gleb Shevchuk, Dorsa Sadigh", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "WEAqKGpt4v", "reviewer": "Reviewer_pjSm", "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.", "questions": "Line 31, \"Including\" -> \"This includes\"\n\nThe 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.\n\nIt seems that adversarial imitation learning method GAIL is always outperformed by non-adversarial one IQLearn? What might be a reason for that?\n\nIn Figure 2 and 3, is the result of IQLearn offline or online?", "limitations": "Limitations are discussed in the Discussion section.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "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.\n\nThe 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.\n\nThe 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].\n\nWould be nice to have more intuitive understanding of why adversarial methods such as GAIL do not work very well.\n\nThere are some places where the writing can be improved. Please see questions.\n\n[1]Tianyi Zhang, Varsha Kishore, Felix Wu, Kilian Q. Weinberger, & Yoav Artzi (2019). BERTScore: Evaluating Text Generation with BERT. CoRR, abs/1904.09675.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the constructive feedback. By addressing these helpful comments and questions, the paper has improved and gained in clarity. Please mention further questions or clarifications and we will work to address them ASAP.\n\nAs the review mentions, our small performance gains are consistent and considerably larger in some domains, in particular for GSM8k. They further go hand-in-hand with crucial diversity improvements across tasks. We took this feedback seriously and additionally looked into further analysis as well as additional tasks. In addition to XSUM, GSM8k, and TLDR, we added WMT22 results (with results in the general rebuttal). We further analyzed the extracted rewards as described in the general rebuttal and are able to demonstrate correlation between learned rewards and task performance metrics. The particular suggestion of BertScore as an additional metric is highly appreciated. We are looking into experiments with BertScore but will unlikely be able to generate results before the rebuttal session ends due to the computational cost of recreating experiments.\n\nAdversarial imitation is generally known to be less stable (e.g. [1,2]) and often additional heuristics are required for stabilization [3] while hyperparameter tuning becomes increasingly expensive. Due to these aspects, we have only been able to obtain results for the T5 models. This is mentioned in parts of the paper already (see e.g. Sections 3.1 and A.3.3) and we have further expanded the corresponding sections in the submission.\n\nMinor points:\n- We have further expanded our discussion of related work to include the suggested papers.\n- We have improved figure 1 and the corresponding caption.\n- We have fixed the remaining spelling and other mistakes. Thanks for pointing these out.\n- The main IQLearn experiments are offline as we mention in sections 2 and 3. We have emphasized and clarified this throughout all sections.\n\n[1] On the Algorithmic Stability of Adversarial Training, NeurIPS 2021, Yue Xing, Qifan Song, Guang Cheng\n\n[2] Language GANs falling short, ICLR 2020, Massimo Caccia, Lucas Caccia, William Fedus, Hugo Larochelle, Joelle Pineau, Laurent Charlin\n\n[3] Creating Multimodal Interactive Agents with Imitation and Self-Supervised Learning, 2022, DeepMind Interactive Agents Team: Josh Abramson, Arun Ahuja, Arthur Brussee, Federico Carnevale, Mary Cassin, Felix Fischer, Petko Georgiev, Alex Goldin, Mansi Gupta, Tim Harley, Felix Hill, Peter C Humphreys, Alden Hung, Jessica Landon, Timothy Lillicrap, Hamza Merzic, Alistair Muldal, Adam Santoro, Guy Scully, Tamara von Glehn, Greg Wayne, Nathaniel Wong, Chen Yan, Rui Zhu", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We would like to thank our reviewers for their valuable feedback. The paper has already considerably improved and we will work hard to address further comments and questions to conclude the rebuttal process to everyone’s satisfaction.

The reviews generally appreciate the discussion of inverse RL in the context of language modeling as well as the derivations and reformulation of IQLearn as temporal difference regularized MLE. Multiple reviews ask for the addition of further tasks and, while the computational cost is non-negligible, we have added results on the large WMT22 English-German task (285M examples). IQLearn gains are between 2.2-2.4 over MLE with beam search decoding. Early results with computationally cheaper decoding via temperature sampling indicate larger gains (evaluation is ongoing and we will add the complete results via comments ASAP). We are investigating further options for tasks as well. 

We aim to address questions around the key benefits of IRL. Here, in addition to training a policy (i.e. the generative LLM), IRL enables the extraction of reward functions from demonstration data. We discuss multiple benefits for this additional source of reward information in the discussion section and this includes the potential to mitigate reward model overoptimization during RLHF [1], and we have added more analysis on this (see below). In IQLearn, in particular, we can recover rewards from the Q-function, as we describe in line 136 (via $r_t=Q(s_t, a_t) - \gamma V(s_{t+1})$). We have added correlation analysis between the extracted rewards and task performance metrics. As an example, current results show both Pearson Correlation Coefficient and Spearman's Rank Correlation Coefficient between 0.21 and 0.44 for IRL-extracted rewards and BLEU and between 0.27 and 0.48 for ChrF on WMT. For comparison, the corresponding correlations for MLE are respectively around 0.04-0.05 and 0.004-0.005. The complete results have been added to the paper. PDF updates are unfortunately not possible during the rebuttal process, so please reach out if there are further questions about these results.


Further specific details for each review are included in our individual answers. Please do not hesitate to point out further questions or comments.

[1] Scaling Laws for Reward Model Overoptimization, 2022, Leo Gao, John Schulman, Jacob Hilton

Meta-review: This paper proposes a technique based on Inverse Reinforcement Learning for LLM training, as alternative to the standard MLE estimator for next token prediction. The proposed approach, that builds on existing work on inverse soft-Q learning, is shown to help increase response diversity while maintaining or improving quality (depending on the benchmark).

Reviews were somewhat mixed (4,4,5,6,8), with reviewers being generally appreciative of the novelty of the approach (applying IRL to LLM training), but with an overall concern towards empirical results that do not show strong improvements over standard MLE across all benchmarks. In their response, the authors added more experiments showing benefits on the WMT22 task.

Among the two reviewers who rated 4, one didn't follow-up on the author response, while the other did (privately) and decided to keep their score due to quantiative improvements not being strong enough.

Although I agree it would be nice to see stronger improvements across the board, I consider the gains in diversity to be significant and potentially quite relevant in some applications. In addition, being able to improve on some tasks (while not degrading on others) remains meaningful, and I believe this work also has potential to encourage further research on IRL for LLM training. As a result, I am recommending acceptance.
Decision: Accept (poster)