Title: Learning Tractable Distributions of Language Model Continuations
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
Keywords: probabilistic circuit, tractable model, exact inference, controlled language generation

Abstract: Controlled language generation is the task of generating text conditional on some constraint: for example, valid syntax, or alignment to human values. A fundamental challenge is that these constraints may depend on future tokens, which standard autoregressive language models (LM) cannot readily account for. Prior work has tackled this problem by modeling future completions using a *tractable probabilistic circuit* as a surrogate model that aims to approximate the distribution over future completions, which is used to adjust the LM logits at decoding time. However, we surprisingly find that their distribution over future tokens is often insensitive to the past context, limiting their effectiveness in controlling LMs. By viewing tractable probabilistic circuits as latent variable models, we show via theoretical and empirical analysis that this can be largely attributed to the suboptimal context encoder used in these models. In this paper, we overcome this limitation by utilizing a separate, more powerful neural language model encoder. Building on this, we propose an efficient decoding-time controller that reuses the same language model used for generation to encode a context-rich prior over the latent states, while the tractable surrogate performs exact future lookahead from the conditioned state. Importantly, we show how this can be done without incurring significant decoding-time overhead. Experiments across diverse controlled generation tasks demonstrate significant improvements over prior approaches in terms of constraint satisfaction and perplexity of generations.