Learning Tractable Distributions Of Language Model Continuations

TL;DR

The paper introduces LTLA, a hybrid approach combining neural and HMM models to improve tractable language model continuations, enabling better control, syntactic accuracy, and safety with minimal decoding overhead.

Contribution

LTLA is a novel hybrid method that conditions a shared HMM on neural embeddings to improve sequence continuation control and efficiency.

Findings

LTLA outperforms standard HMM surrogates in likelihood.

Enables lookahead control in vision-language models.

Achieves 100% syntactic constraint satisfaction with minimal overhead.

Abstract

Controlled generation imposes sequence-level constraints (syntax, style, safety) that depend on future tokens, making exact conditioning of an autoregressive LM intractable. Tractable surrogates such as HMMs can approximate continuation distributions and steer decoding, but standard surrogates are often weakly context-aware. We propose Learning to Look Ahead (LTLA), a hybrid method that uses base-LM embeddings to condition a globally learned tractable surrogate: a neural head predicts only a prefix-dependent latent prior, while a shared HMM answers continuation queries exactly. LTLA is designed to avoid two common efficiency traps when adding neural context. First, it avoids vocabulary-sized prefix rescoring (V extra LM evaluations) by scoring all next-token candidates via a single batched HMM forward update. Second, it avoids predicting a new HMM per prefix by learning one shared HMM and conditioning only the latent prior, which enables reuse of cached future-likelihood (backward) messages across decoding steps. Empirically, LTLA improves continuation likelihood over standard HMM surrogates, enables lookahead control for vision--language models by incorporating continuous context, achieves 100% syntactic constraint satisfaction, and improves detoxification while adding only a 14% decoding-time overhead.