Title: Imitating Language via Scalable Inverse Reinforcement Learning
Venue: NeurIPS 2024
Accepted: yes
Keywords: Language Modeling

Inverse Reinforcement Learning

Imitation Learning

Supervised Fine-tuning
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

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.

1 Introduction

In recent years, the imitation of 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 applications such as robotic control [4, 29]. Recently, perspectives to language modeling have shifted towards explicit treatment as a 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. It includes a broader scope for which data contains information about rewards and preferences as well as the dynamics-aware optimization of each action based on its future impact - all while taking into account computational scalability.

MLE

Offline IRL

Online IRL

TrainingObjective

TrainingObjective

Model

Model

Model

666

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 of 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 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 a 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 1. 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 MLE. Our experiments range from 250M to 3B parameter models for the encoder-decoder T5 [51] and decoder-only PaLM2 [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 RL-centric 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 underlying language generation and enables computationally cheap offline training.

We compare MLE and IRL formulations including adversarial and non-adversarial, 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 analyze the extracted reward functions indicating the potential usefulness of IRL to obtain discriminative in addition to generative improvements from demonstrations.

2 Methods

Language generation can be modeled as a sequential decision making problem. On a token level, it is the problem of generating the next token x i given the already generated sequence of tokens ( x 0 , . . . , x i -1 ) . Wethus seek a distribution π ( x i | x 0 , . . . , x i -1 ) , which we will also refer to as a policy. For a given policy the sequence likelihood x = ( x 0 , . . . , x N ) can be computed autoregressively:

The classic maximum likelihood estimation based approach leverages this factorization in order to efficiently train the policy by maximizing the log-likelihood of the training sequences, D = { x 0 , . . . , x M } :

For fine-tuning problems, where a pre-trained model is fine-tuned to a particular set of tasks, the problem can be formulated analogously but may have additional conditioning variables which we will leave out for the sake of clarity.

Distribution matching. State-action distribution matching algorithms [25, 20], which are wellestablished in the field of imitation learning - and can be seen as solving an IRL problem see e.g. [25] - approach the problem in a different manner: They seek to minimize the divergence between the γ -discounted state-action distribution of the policy π ( a | s ) and the discounted state-action distribution of the expert policy π E ( a | s ) . We define the discounted state distribution for a Markov Decision Process (MDP) as ρ ( s ) = (1 -γ ) ∑ ∞ i =0 γ i P ( s i = s | π ) and the discounted state-action distribution is accordingly defined as µ π ( s, a ) = π ( a | s ) ρ ( s ) ; where P ( s i = s | π ) is the probability of seeing state s when acting according to the policy π . For our autoregressive generation MDP, the state corresponds to the concatenation of already generated tokens (commonly including the prefix or prompt), s i = ( x 0 , . . . , x i -1 ) and the action is the next token, a i = x i thus yielding a problem with deterministic dynamics. We omit state and action arguments in the following discussion whenever it is clear from the context.

In order to enable more straightforward algebraic manipulation, the divergence is often combined with a weighted causal entropy term H ( π ) = -E µ π log( π ( a | s )) . The goal is then to find a policy that minimizes the objective J ( π ) :

where D f = E µ E [ f ( µ π ( a,s ) µ E ( a,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 ( t ) = f RKL ( t ) = -log( t ) , we can decompose the objective into a state distribution and an MLE term:

where KL denotes the KL-divergence which corresponds to 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 it 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 -divergence can be leveraged to avoid expert densities:

where f ∗ is the convex conjugate of f and dom f ∗ is its domain. Notably, GAIL [25] can be retrieved by using the Jensen-Shannon divergence with its convex conjugate f ∗ ( g ) = -log(1 -exp( g ( s, a ))) and defining the discriminator D ( s, a ) := exp( g ( s, a )) :

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 a clearer relationship to MLE. Using a change of variable r : S × A → dom f ∗ with r ( s, a ) = -g ( s, a ) and re-arranging the terms, we start by rendering the internal RL-problem more explicit:

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 a saddle point problem [20] and the max-min can be swapped:

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]

where π r is the optimal policy for the reward r and v r and µ r are respectively the corresponding state value function and discounted state-action distribution 2 . This is applied to the right hand term of Eq. (8) after maximization to obtain (after dropping the arguments for conciseness).

Using a telescoping argument [20, 57], we can relate the value of the initial state distribution of a policy to the difference in values on the state-action distribution induced by arbitrary other policies, i.e. E ρ 0 [ v r ] = E µ π [ v r -λv ′ r ] . In particular, here we can change the second expectation's state-action distribution to one induced by the expert policy yielding:

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 a 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 ( q ∗ ) = q ∗ -E s ′ [ λv ′∗ ] [20, 42], or analogously between the reward and a state value and a policy, r ( v r , π r ) = v r + λ log π r -E s ′ [ λv ′ r ] (using q r = v r + λ log π r ). We can thus reparameterize the problem and optimize the optimal state value and policy instead of the reward.

Choosing the χ 2 -divergence with convex conjugate f ∗ ( t ) = -t 2 4 + t is particularly convenient at this point since it can be combined with rescaling the value v r λ = v r /λ to obtain our reformulated IQLearn objective to be minimized.

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

6

8

10

12

14

16

Accuracy without Calculator

75.0

77.5

80.0

82.5

85.0

87.5

90.0

92.5

95.0

Self BLEU

Base

Large

X-Large

better

10

15

20

25

30

75

80

85

90

95

Self BLEU

Base

Large

X-Large

GAIL (ent-reg=0.00)

GAIL (ent-reg=0.05)

GAIL (ent-reg=0.10)

IQLearn (

=0.10)

IQLearn (

=0.50)

IQLearn (

=1.00)

MLE (ent-reg=0.00)

MLE (ent-reg=0.05)

MLE (ent-reg=0.10)

likelihood with a 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 3 . The regularization term couples the learned policy to a value function and favors 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 λ = 0 retrieves standard MLE and where by adjusting λ the regularization strength can be flexibly increased.

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. In addition, we perform ablations on online 3 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 on the Pareto front of task performance and diversity?

For which task and dataset properties is IRL particularly relevant?

What is the impact of online data for IRL training?

How informative are rewards extracted from SFT data?

3.1 Algorithms, 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 to distinguish transitions from agent and dataset and a 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 A.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 its 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,

35.5

36.0

36.5

37.0

37.5

38.0

38.5

ROUGE 1

20

30

40

50

60

70

80

90

Self BLEU

Base

Large

better

14.5

15.0

15.5

16.0

16.5

ROUGE 2

20

30

40

50

60

70

80

90

Self BLEU

Base

Large

GAIL (ent-reg=0.0)

GAIL (ent-reg=0.01)

GAIL (ent-reg=0.1)

IQLearn (

=0.1)

IQLearn (

=0.5)

IQLearn (

=1.0)

MLE (ent-reg=0.0)

MLE (ent-reg=0.01)

MLE (ent-reg=0.05)

MLE (ent-reg=0.1)

we apply a 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 in Section 2 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-efficient fine-tuning via adapters [27] as used in prior work [15], we focus on the full 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 [51] 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 we calculate self-similarity of generated examples as measured by Self-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 add the evaluation via per-token-entropy in Appendix A.3.1.

3.2.1 T5 models

We perform experiments with the base, large, and xl T5 models [51] on the XSUM [43] and GSM8k [14] tasks. These models further serve as foundation for our later ablations. We are able to obtain small but notable gains in performance across all tasks, as shown in Figures 2 and 3, 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 a more emphasized boost in 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 A.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

10

1

10 0

10 1

Temperature

0

5

10

15

20

MLE

IQLearn(

=0.5)

10

2

10

1

10 0

10 1

Temperature

0

5

10

15

20

25

Rouge Lsum

MLE

IQLearn(

=0.5)

10

2

10

1

10 0

10 1

Temperature

0

5

10

15

20

25

30

35

BLEU

MLE

IQLearn(

=0.1)

MLE + beam search

IQLearn(

=0.1) + beam search

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-offs 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 a 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-German translation dataset WMT22 [34]. We limit these experiments to a 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 IQLearn 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 a 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 IQLearn achieves higher performance in a 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 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 an 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 4 and the comparison with the online application of IQLearn (see Appendix A.1.2). 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.

14

IQLearn

IQLearn

IQLearn, online ratio=0.1

40

IQLearn, online ratio=0.1

Accuracy with calculator

12

IQLearn, online ratio=0.5

IQLearn, online ratio=0.5

50

Self-BLEU

10

60

8

70

6

0

1000

2000

3000

4000

5000

6000

7000

0

1000

2000

3000

4000

5000

6000

7000

Size of training set

Size of training set

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-turn settings, prevents exact recovery after mistakes, which could be otherwise learned via online IRL. Therefore, extensions of the LLM action space to enable recovery is a fruitful direction for increased benefits from online data with RL or IRL methods [15].

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, IQLearn represents learned rewards implicitly via the Q-function as r t = Q ( s t , a t ) -γV ( s t +1 ) , and its online version (i.e. with a non-zero mix-in ratio α 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, we report the Spearman's rank correlation coefficient between accumulated rewards (over complete sampled trajectories for the full validation sets) for online IQLearn ( α = 0 . 1 ) and task-specific metrics. Compared to MLE rewards, which, as expected, are not strongly correlated with any metric, we see a clear increase in correlation 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 idiosyncratic metric: only the correctness of the final generated numerical expression affects the accuracy calculation, effectively ignoring most of the trajectory (and

1000

2000

3000

4000

5000

6000

7000

Size of training set

6

7

8

9

10

11

12

13

Accuracy with calculator

IQLearn (

=0.1)

IQLearn (

=0.5)

MLE

1000

2000

3000

4000

5000

6000

7000

Size of training set

3.0

3.5

4.0

4.5

5.0

5.5

6.0

6.5

Accuracy without calculator

IQLearn (

=0.1)

IQLearn (

=0.5)

MLE

3.2

1e1

num_examples=1000

num_examples=3000

num_examples =10000

3.0

ROUGE

1

2.8

IQLearn (A=0.10)

IQLearn (A=0.50)

MLE (ent-reg=0.00)

2.6

MLE (ent-reg=0.01)

MLE (ent-reg=0.05)

2.4

MLE (ent-reg=0.10)

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.00.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.00.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

Training Step

1e5

Training Step

1e5

Training Step

1e5

so its transition'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 λ .

0

10

1

5×10 1

10 0

0.1

0.0

0.1

0.2

0.3

0.4

0.5

0.6

Spearman's reward correlation

(MLE)

BLEU

ChrF

4 Related Work

General imitation learning. Imitation learning assumes a dataset of expert demonstrations, and the aim is to train a policy that matches the expert. There are two broad categories of imitation learning approaches: behavioral cloning (BC) and inverse reinforcement learning (IRL). In BC, a 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 in 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 a 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 IRL 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], IQLearn [20]), or a Lagrangian dual objective (PPIL [62]). The stable, large-scale application of inverse RL methods has been a 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 adversarial training [25, 19] or lastly saddlepoint-based 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].

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, 39], 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 a series of work using GANs for text generation [32]. These methods have been shown to work only in the temperature 1 regime [11], 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, IQLearn is 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 IQLearn. 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-BLEU 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 or 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 integrating SFT data into RL-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 a unified RL-based framework for unified LLM posttraining 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 IRL-extracted reward functions including crucial information about LLM task performance.

6 Conclusions

This paper presents a detailed investigation of the potential of IRL algorithms for imitation in language model tuning focusing on performance, diversity, and computational requirements. We introduce a reformulation of IQLearn which enables principled interpolation between robust, standard supervised fine-tuning and more effective IRL algorithms. Our experiments demonstrate particularly strong gains for IRL 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 MLEbased optimization. Additional correlation analysis between IRL-extracted 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.

Acknowledgments and Disclosure of Funding

We extend our gratitude to Daniele Calandriello, Doina Precup, and our 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

Rishabh Agarwal, Nino Vieillard, Yongchao 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

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

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 Meier-Hellstern, Gaurav Mishra, Erica Moreira, Mark Omernick, 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-Choo, 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, Hyeontaek 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. PaLM 2 technical report, 2023. 2, 6, 7

Brenna D Argall, Sonia Chernova, Manuela Veloso, and Brett Browning. A survey of robot learning from demonstration. Robotics and autonomous systems , 57(5):469-483, 2009. 1

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

Gregor Bachmann and Vaishnavh Nagarajan. The pitfalls of next-token prediction. ArXiv , abs/2403.06963, 2024. 2

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 arXiv:2204.05862 , 2022. 2

Michael Bain and Claude Sammut. A framework for behavioural cloning. In Machine Intelligence 15 , pages 103-129, 1995. 2, 10

Matt Barnes, Matthew Abueg, Oliver F. 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

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

Massimo Caccia, Lucas Caccia, William Fedus, Hugo Larochelle, Joelle Pineau, and Laurent Charlin. Language gans falling short. arXiv preprint arXiv:1811.02549 , 2018. 10

Stephen Casper, Xander Davies, Claudia Shi, Thomas Krendl Gilbert, J'er'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-Menell. Open problems and fundamental limitations of reinforcement learning from human feedback. ArXiv , abs/2307.15217, 2023. 1

Paul Christiano, Jan Leike, Tom B Brown, Miljan Martic, Shane Legg, and Dario Amodei. Deep reinforcement learning from human preferences. arXiv preprint arXiv:1706.03741 , 2017. 1

Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems. arXiv preprint arXiv:2110.14168 , 2021. 6, 7

Chris Cundy and Stefano Ermon. Sequencematch: Imitation learning for autoregressive sequence modelling with backtracking, 2024. 6, 8, 10, 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

Gemini Team et al. Gemini: A family of highly capable multimodal models, 2024. 1

Chelsea Finn, Sergey Levine, and P. Abbeel. Guided cost learning: Deep inverse optimal control via policy optimization. ArXiv , abs/1603.00448, 2016. 9

Justin Fu, Katie Luo, and Sergey Levine. Learning robust rewards with adversarial inverse reinforcement learning. arXiv preprint arXiv:1710.11248 , 2017. 9

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

Shivam Garg, Dimitris Tsipras, Percy S 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-30598, 2022. 16

Matthieu Geist, Bruno Scherrer, and Olivier Pietquin. A theory of regularized markov decision processes. CoRR , abs/1901.11275, 2019. 4

Seyed Kamyar Seyed Ghasemipour, Richard S. Zemel, and Shixiang Gu. A divergence minimization perspective on imitation learning methods. CoRR , abs/1911.02256, 2019. 2, 3

Joey Hejna and Dorsa Sadigh. Inverse preference learning: Preference-based rl without a reward function. ArXiv , abs/2305.15363, 2023. 10

Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. Advances in neural information processing systems , 29:4565-4573, 2016. 2, 3, 4, 5, 9

Luca Hormann and Artem Sokolov. Fixing exposure bias with imitation learning needs powerful oracles. arXiv preprint arXiv:2109.04114 , 2021. 10

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. 6

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

Ahmed Hussein, Mohamed Medhat Gaber, Eyad Elyan, and Chrisina Jayne. Imitation learning: A survey of learning methods. ACM Comput. Surv. , 50(2), apr 2017. 1

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

Timo Kaufmann, Paul Weng, Viktor Bengs, and Eyke Hüllermeier. A Survey of Reinforcement Learning from Human Feedback, 2023. 9

Pei Ke, Fei Huang, Minlie Huang, and Xiaoyan Zhu. Araml: A stable adversarial training framework for text generation. arXiv preprint arXiv:1908.07195 , 2019. 10

Geon-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 on Learning Representations , 2022. 10, 16

Tom Kocmi, Rachel Bawden, Ondˇ rej 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´ c. Findings of the 2022 Conference on Machine Translation (WMT22). In Proceedings of the Seventh Conference on Machine Translation (WMT) , 2022. 6, 7

Ilya Kostrikov, Kumar Krishna Agrawal, Debidatta Dwibedi, Sergey Levine, and Jonathan Tompson. Discriminator-actor-critic: Addressing sample inefficiency and reward bias in adversarial imitation learning. In International Conference on Learning Representations , 2019. 9

Ilya Kostrikov, Ofir Nachum, and Jonathan Tompson. Imitation learning via off-policy distribution matching. arXiv preprint arXiv:1912.05032 , 2019. 9

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

Jiaxiang Li, Siliang Zeng, Hoi-To Wai, Chenliang Li, Alfredo García, and Mingyi Hong. Getting more juice out of the sft data: Reward learning from human demonstration improves sft for llm alignment. ArXiv , abs/2405.17888, 2024. 10

Ziniu Li, Congliang Chen, Tian Xu, Zeyu Qin, Jiancong Xiao, Ruoyu Sun, and Zhimin Luo. Entropic distribution matching in supervised fine-tuning of llms: Less overfitting and better diversity, 2024. 10

Alexander Lin, Jeremy Wohlwend, Howard Chen, and Tao Lei. Autoregressive knowledge distillation through imitation learning. arXiv preprint arXiv:2009.07253 , 2020. 10

Yecheng Ma, Andrew Shen, Dinesh Jayaraman, and Osbert Bastani. Versatile offline imitation from observations and examples via regularized state-occupancy matching. In International Conference on Machine Learning . PMLR, 2022. 10, 16

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

Shashi Narayan, Shay B. Cohen, and Mirella Lapata. Don't give me the details, just the summary! topic-aware convolutional neural networks for extreme summarization. ArXiv , abs/1808.08745, 2018. 6

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

OpenAI. Gpt-4 technical report, 2023. 1

Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to follow instructions with human feedback. Advances in neural information processing systems , 35:27730-27744, 2022. 1, 5, 10

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

Bilal Piot, Matthieu Geist, and Olivier Pietquin. Boosted and reward-regularized classification 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, 9

Dean A Pomerleau. Alvinn: An autonomous land vehicle in a neural network. Advances in neural information processing systems , 1, 1988. 9

Rafael Rafailov, Joey Hejna, Ryan Park, and Chelsea Finn. From r to q*: Your language model is secretly a q-function, 2024. 10

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

Marc'Aurelio Ranzato, Sumit Chopra, Michael Auli, and Wojciech Zaremba. Sequence level training with recurrent neural networks, 2016. 2

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

Siddharth Reddy, Anca D. Dragan, and Sergey Levine. Sqil: Imitation learning via reinforcement learning with sparse rewards, 2019. 10

Juntao Ren, Gokul Swamy, Zhiwei Steven Wu, J. Andrew Bagnell, and Sanjiban Choudhury. Hybrid inverse reinforcement learning. ArXiv , abs/2402.08848, 2024. 19

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 international conference on artificial intelligence and statistics , pages 627-635. JMLR Workshop and Conference Proceedings, 2011. 2, 8, 9

Harshit Sikchi, 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

Nisan Stiennon, Long Ouyang, Jeffrey Wu, Daniel Ziegler, Ryan Lowe, Chelsea Voss, Alec Radford, Dario Amodei, and Paul F Christiano. Learning to summarize with human feedback. Advances in Neural Information Processing Systems , 2020. 6, 7

Hao Sun. Supervised fine-tuning as inverse reinforcement learning. arXiv preprint arXiv:2403.12017 , 2024. 10

Umar Syed and Robert E. Schapire. A game-theoretic approach to apprenticeship learning. In Advances in Neural Information Processing Systems , 2007. 9

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 arXiv:2307.09288 , 2023. 2

Luca Viano, Angeliki Kamoutsi, Gergely Neu, Igor Krawczuk, and Volkan Cevher. Proximal point imitation learning. In Advances in Neural Information Processing Systems , 2022. 9

Joe Watson, Sandy H. Huang, and Nicolas Heess. Coherent soft imitation learning. In Advances in Neural Information Processing Systems (NeurIPS) , 2023. 10

Ronald J. Williams and David Zipser. A learning algorithm for continually running fully recurrent neural networks. Neural Computation , 1:270-280, 1989. 2

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

Markus Wulfmeier, Arunkumar Byravan, Sarah Bechtle, Karol Hausman, and Nicolas Manfred Otto Heess. Foundations for transfer in reinforcement learning: A taxonomy of knowledge modalities. ArXiv , abs/2312.01939, 2023. 1

Markus Wulfmeier, Ingmar Posner, and P. Abbeel. Mutual alignment transfer learning. ArXiv , abs/1707.07907, 2017. 5

Markus Wulfmeier, Dominic Zeng Wang, and Ingmar Posner. Watch this: Scalable cost-function learning for path planning in urban environments. 2016 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS) , pages 2089-2095, 2016. 9

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 J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems , volume 34, pages 26523-26535. Curran Associates, Inc., 2021. 5

Tian Xu, Ziniu Li, and Yang Yu. Error bounds of imitating policies and environments. Advances in Neural Information Processing Systems , 2020. 9

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

Tianyi Zhang, Varsha Kishore, Felix Wu, Kilian Q. Weinberger, and Yoav Artzi. Bertscore: Evaluating text generation with bert. ArXiv , abs/1904.09675, 2019. 17

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 41st international ACM SIGIR conference on research & development in information retrieval , pages 1097-1100, 2018. 6

Brian D Ziebart. Modeling Purposeful Adaptive Behavior with the Principle of Maximum Causal Entropy . PhD thesis, Carnegie Mellon University, 2010. 4

Brian D 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, IL, USA, 2008. 9

A Appendix

A.1 Implementation Details

A.1.1 IQLearn

For IQLearn, fine-tuning starts with a policy which is subsequently finetuned as a 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 Qfunction Q ∗ via π ∗ ( a | s ) = 1 / Z s exp Q ∗ ( s, a ) with normalization factor Z s = ∑ a ′ ∈ A exp Q ∗ ( s, a ′ ) . We further obtain V via V ( s ) = log ∑ a ∈ A exp Q ( 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 statebased value V can be arbitrarily inaccurate. In practice, we can add further KL regularization, or 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].

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.

A.1.2 Online IQLearn

We also implemented an online version of IQLearn, which makes use of additional (non-expert) samples. In contrast to IQLearn [21] which makes use of these examples to estimate the initial value, E ρ 0 v ( s ) (which is equivalent between online and offline data when the prompt is the same as in our case), we integrate the additional samples to relax the distribution matching loss:

where µ B is the discounted state-action distribution of the additional samples and where α ∈ [0 , 1) is the strength of the mix-in. The problem is thus relaxed and allows the policy to match the mix-in distribution µ B in case the expert is too difficult to match.

From the adapted distribution matching loss the same steps can then be taken as in Section 2. This results in the generalised loss:

with

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 to 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].

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 a final cost weight displayed below. The value network is also initialized from the initial LLM. The discriminator is trained with a crossentropy objective. The reward is re-shaped from the the discriminator output to be a positive reward: r t = log(1 + exp( D ( s t ))) .

We explore the heuristic combination of GAIL with standard MLE training by using a weighted combination of GAIL and MLE losses for the policy. While this is 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.

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

where we compute the entropy at each token of the sequences form µ E .

A.2.1 Computational Requirements

Our experiments with T5 models use TPU v3 infrastructure and are running between approximately 3 days and 2 weeks. Our experiments with PaLM2 models use TPU v4 infrastructure and are running under 1 week.

A.3 Additional Experiments

We include a set of further experiments to complement the results in the main paper.

A.3.1 Quality-Diversity Evaluation

In addition to the Self-BLEU metric, we further add plots for model entropy and add the ROUGELSUM 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 can consistently outperform MLE in accuracy, even when sampling with a lower temperature than the training one.

A.3.3 Toy Experiments for Offline 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 autoregressive language generation. The agent starts in the left

28.5

29.0

29.5

30.0

30.5

31.0

ROUGE-LSUM

30

40

50

60

70

80

90

Self BLEU

Base

Large

GAIL (ent-reg=0.0)

GAIL (ent-reg=0.01)

GAIL (ent-reg=0.1)

IQLearn (

=0.1)

IQLearn (

=0.5)

IQLearn (

=1.0)

MLE (ent-reg=0.0)

MLE (ent-reg=0.01)

MLE (ent-reg=0.05)

MLE (ent-reg=0.1)

28.5

29.0

29.5

30.0

30.5

31.0

ROUGE-LSUM

0.0

0.2

0.4

0.6

0.8

1.0

1.2

1.4

Avg. per token entropy

Base

Large

GAIL (ent-reg=0.0)

GAIL (ent-reg=0.01)

GAIL (ent-reg=0.1)

IQLearn (

=0.1)

IQLearn (

=0.5)

IQLearn (

=1.0)

MLE (ent-reg=0.0)

MLE (ent-reg=0.01)

MLE (ent-reg=0.05)

MLE (ent-reg=0.1)

10

1

10 0

10 1

Temperature

0

5

10

15

20

25

Accuracy with calculator

MLE

IQLearn(

=0.5)

10

2

10

1

10 0

10 1

Temperature

0

5

10

15

20

25

30

35

Rouge 1

MLE

IQLearn(

=0.5)

10

2

10

1

10 0

10 1

Temperature

0

2

4

6

8

10

12

Rouge 2

MLE

IQLearn(

=0.5)

black state and has to reach the green on the right. In each of the bottom states the agent has 2 actions, move left or move up. The transition dynamics are noisy so 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. In a way, it 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 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 to 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

I

I

I

1

I

1

I

1

1

1

the discriminator provides negative rewards. If we provide positive rewards, tuning those rewards becomes a complex task as seen in Figure 12, where without additional MLE objective the GAIL agent often uses the maximum sample length with correlated loss of 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].

50000

100000

150000

200000

step

0

5

10

15

20

25

30

35

40

ROUGE 1

mle_weight = 0.0

50000

100000

150000

200000

step

mle_weight = 1.0

GAIL (ent-reg=0.05)

GAIL (ent-reg=0.10)

GAIL (ent-reg=0.00)

0

50000

100000

150000

200000

step

30

35

40

45

50

55

Avg. number of tokens

mle_weight = 0.0

0

50000

100000

150000

200000

step

mle_weight = 1.0

GAIL (ent-reg=0.05)

GAIL (ent-reg=0.10)

GAIL (ent-reg=0.00)

A.4 Expanded Experiments on WMT22

For WMT22, we additionally evaluate IQLearn performance for a wide range of λ values. Here, IQLearn consistently gains (up to + 2.16 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 4) is used with a length penalty 0.6.

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 set 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 experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: Experimental recreation of the results might be limited by required computational resources. Other than that, algorithms are accurately described including hyperparameters 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 instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [No]

Justification: We have been unable to provide public access to our code.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: All experimental detailed information is either given in the main paper or more likely in the appendix.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: Where computationally possible, the experiments include error bars and related information.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: Approximate details are given in the appendix

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

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 describe 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 or 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 as 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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nAs 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.\n\nQUESTIONS:\nLine 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?\n\nLIMITATIONS:\nLimitations are discussed in the Discussion section.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-06-19"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n* 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.\n\nWEAKNESSES:\n* 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”\n\nQUESTIONS:\n* 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.\n\nLIMITATIONS:\nThese were adequately discussed in Section 5.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-11"}

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

STRENGTHS:
1.	This work investigates the potential of inverse reinforcement learning for tuning LLMs. The strengths of IRL are evaluated in terms of performance, diversity, and computational requirements. The experiments demonstrate the computationally cheaper offline IRL can obtain crucial performance gain over MLE-based LLM tunning. 

2.	IQLearn is a method that can work online and offline. The re-formulation of IQLearn to LLM tunning enables large performance gains compared to MLE methods. Moreover, IQLearn can be regarded as a regularized extension of MLE, which bridges IRL and MLE for LLM tunning.

WEAKNESSES:
1.	IRL methods, even for offline IRL methods, cost more time than MLE for LLM tunning. 

2.	More experimental details are required to reproduce the results.

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

LIMITATIONS:
The discussion section discussed the potential limitations of this work.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-11

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

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

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

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

[1] LANGUAGE GANS FALLING SHORT

LIMITATIONS:
see weaknesses
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2024-07-11

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n* A solid theoretical analysis is given which draws links between MLE and forms of IRL.\n\nWEAKNESSES:\n* 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.\n\nQUESTIONS:\nWhat 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).\n\nLIMITATIONS:\nNo concerns.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-15"}