AutoIndexer: Flexible-Order Decoding by Training Causal Models on Chains of Edits

Autoregressive decoding faces challenges for data that are not causally ordered, e.g., containing cross-references or future dependencies. Furthermore, once tokens are generated, they are fixed and cannot be modified or improved, unlike how humans can iteratively refine their thoughts. While diffusion LLMs partially address this issue, they do not naturally lend themselves to adding or removing arbitrary numbers of tokens to/from a pre-generated context. We design AutoIndexer as a generative model with a novel training and decoding mechanism that simulates a Chain of Edits, allowing the model to flexibly modify its prior outputs in hindsight via token insertions, substitutions, and deletions. The model learns to predict marker tokens for opening and closing edits, simultaneously moving a cursor to positions in the existing token sequence, analogous to a text editor. Training sequences are obtained by perturbing ground-truth labels with inverse chained edit operations, and the model learns to edit by undoing those inverses. The perturbation training allows AutoIndexer to be robust to erroneous context and adaptable to editing tasks such as code correction. We demonstrate that our method exhibits post-generation context-editing capabilities that most autoregressive and diffusion models lack without compromising performance on standard LLM benchmarks.