Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

SequenceMatch: Imitation Learning for Autoregressive Sequence Modelling with Backtracking

Chris Cundy, Stefano Ermon

TL;DR

SequenceMatch reframes autoregressive sequence generation as imitation learning to address compounding OOD errors that arise under maximum likelihood training. By minimizing occupancy-divergence objectives (including a $D_{\chi^2}$-mixture divergence) and introducing a backspace action, it enables backtracking from mistakes without adversarial training or architecture changes. The method yields a fully supervised, logit-based objective over transformer policies and can be finetuned on pretrained models with a masking scheme, achieving improvements on arithmetic and text-generation tasks as measured by MAUVE and related metrics. Overall, SequenceMatch offers a practical, non-adversarial route to more robust generation with explicit handling of OOD trajectories and controlled recovery from errors.

Abstract

In many domains, autoregressive models can attain high likelihood on the task of predicting the next observation. However, this maximum-likelihood (MLE) objective does not necessarily match a downstream use-case of autoregressively generating high-quality sequences. The MLE objective weights sequences proportionally to their frequency under the data distribution, with no guidance for the model's behaviour out of distribution (OOD): leading to compounding error during autoregressive generation. In order to address this compounding error problem, we formulate sequence generation as an imitation learning (IL) problem. This allows us to minimize a variety of divergences between the distribution of sequences generated by an autoregressive model and sequences from a dataset, including divergences with weight on OOD generated sequences. The IL framework also allows us to incorporate backtracking by introducing a backspace action into the generation process. This further mitigates the compounding error problem by allowing the model to revert a sampled token if it takes the sequence OOD. Our resulting method, SequenceMatch, can be implemented without adversarial training or architectural changes. We identify the SequenceMatch-$χ^2$ divergence as a more suitable training objective for autoregressive models which are used for generation. We show that empirically, SequenceMatch training leads to improvements over MLE on text generation with language models and arithmetic.

SequenceMatch: Imitation Learning for Autoregressive Sequence Modelling with Backtracking

TL;DR

SequenceMatch reframes autoregressive sequence generation as imitation learning to address compounding OOD errors that arise under maximum likelihood training. By minimizing occupancy-divergence objectives (including a -mixture divergence) and introducing a backspace action, it enables backtracking from mistakes without adversarial training or architecture changes. The method yields a fully supervised, logit-based objective over transformer policies and can be finetuned on pretrained models with a masking scheme, achieving improvements on arithmetic and text-generation tasks as measured by MAUVE and related metrics. Overall, SequenceMatch offers a practical, non-adversarial route to more robust generation with explicit handling of OOD trajectories and controlled recovery from errors.

Abstract

In many domains, autoregressive models can attain high likelihood on the task of predicting the next observation. However, this maximum-likelihood (MLE) objective does not necessarily match a downstream use-case of autoregressively generating high-quality sequences. The MLE objective weights sequences proportionally to their frequency under the data distribution, with no guidance for the model's behaviour out of distribution (OOD): leading to compounding error during autoregressive generation. In order to address this compounding error problem, we formulate sequence generation as an imitation learning (IL) problem. This allows us to minimize a variety of divergences between the distribution of sequences generated by an autoregressive model and sequences from a dataset, including divergences with weight on OOD generated sequences. The IL framework also allows us to incorporate backtracking by introducing a backspace action into the generation process. This further mitigates the compounding error problem by allowing the model to revert a sampled token if it takes the sequence OOD. Our resulting method, SequenceMatch, can be implemented without adversarial training or architectural changes. We identify the SequenceMatch- divergence as a more suitable training objective for autoregressive models which are used for generation. We show that empirically, SequenceMatch training leads to improvements over MLE on text generation with language models and arithmetic.
Paper Structure (49 sections, 5 theorems, 19 equations, 5 figures, 8 tables)

This paper contains 49 sections, 5 theorems, 19 equations, 5 figures, 8 tables.

Table of Contents

  1. Introduction
  2. Preliminaries: Training Objectives for Autoregressive Models
  3. KL-Divergence
  4. Limitations of the KL-divergence
  5. Method
  6. Sequence Modelling as a Markov Decision Process
  7. Minimizing Occupancy Divergences
  8. Reformulating the Occupancy Divergence Minimization Problem
  9. Practical Occupancy Matching with Autoregressive Models
  10. Efficient Training
  11. Augmenting Expert Sequences with Backspace
  12. Related Work
  13. Text Degeneration in Large Language Models
  14. Matching Divergences in Imitation Learning
  15. Experiments
  16. ...and 34 more sections

Key Result

Proposition 3.1

The following equalities hold for the loss: where $\phi$ is concave, $\mathbb{E}_{\rho_{\text{data}}}$ denotes expectations over sampled states and actions $s,a$, $\mathbb{E}_{\mathcal{P}}$ denotes an expectation over successor states $s'$, and $\mathbb{E}_{\rho}$ denotes an expectation over sampled states $s$ and successor states $s'$, for a

Figures (5)

  • Figure 1: A toy model of an autoregressive generation problem, such as language modelling. Our task is to learn a set of conditional distributions that continue the sequence similarly to those sequences in the dataset (green arrows), and avoid incorrect next tokens (red arrows). Our method trains against divergences that more heavily punish out-of-distribution sequences. We additionally introduce a backspace action which can backtrack from an erroneous token (dashed purple arrows).
  • Figure 2: Transforming states and actions to single-pass inputs for parallel training.
  • Figure 3: Training an autoregressive model against a SequenceMatch objective
  • Figure 4: Accuracy on the arithmetic task against noise level (frequency with which noise tokens are added), for ground-truth noise consisting of digits and random noise consisting of random tokens. The ground-truth noise improves accuracy for the behavioral cloning and SequenceMatch models. The random noise does not improve performance for the behavioral cloning model, but does somewhat for the SequenceMatch model, likely helping the model to learn the dynamics of <bkspc>.
  • Figure 5: Algorithm A: Pseudocode for converting action sequences to masked inputs

Theorems & Definitions (10)

  • Proposition 3.1
  • proof
  • Theorem C.1
  • proof
  • Proposition C.2
  • proof
  • Proposition C.3
  • proof
  • Proposition C.4
  • proof