Imitating Language via Scalable Inverse Reinforcement Learning

TL;DR

This paper explores inverse reinforcement learning (IRL) for language model fine-tuning, offering a new perspective that improves diversity and robustness over traditional maximum likelihood methods.

Contribution

It reformulates IRL as a temporal difference extension of MLE, establishing a principled link and demonstrating IRL's advantages in diversity and performance for language model fine-tuning.

Findings

IRL enhances diversity in generated sequences.

IRL achieves comparable or better task performance.

IRL-based rewards improve robustness of language models.

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.