Gradient-Based Constrained Sampling from Language Models

TL;DR

MuCoLa introduces a gradient-based, non-autoregressive constrained sampling framework for language models by performing Langevin Dynamics in embedding space. It seamlessly combines LM likelihood with differentiable constraints through a Lagrangian formulation, enabling soft and hard constraints without compromising fluency or task performance. Empirical results across toxicity avoidance, sentiment control, and keyword-guided generation demonstrate competitive constraint satisfaction, strong fluency, and preserved diversity, with scalable memory advantages over vocabulary-space approaches. The approach offers a flexible, plug-in method for controllable generation without fine-tuning the base LM.

Abstract

Large pretrained language models generate fluent text but are notoriously hard to controllably sample from. In this work, we study constrained sampling from such language models: generating text that satisfies user-defined constraints, while maintaining fluency and the model's performance in a downstream task. We propose MuCoLa -- a sampling procedure that combines the log-likelihood of the language model with arbitrary (differentiable) constraints in a single energy function, and then generates samples in a non-autoregressive manner. Specifically, it initializes the entire output sequence with noise and follows a Markov chain defined by Langevin Dynamics using the gradients of the energy function. We evaluate MuCoLa on text generation with soft and hard constraints as well as their combinations obtaining significant improvements over competitive baselines for toxicity avoidance, sentiment control, and keyword-guided generation.

Figures (2)

• Figure 1: MuCoLa, our proposed method, stylized as $\mu$CoLa. Given a language model, a prompt/input $\mathbf{x}$, and desired constraints defined as thresholds on differentiable functions, we perform Langevin Dynamics updates to generate the entire output sequence $\mathbf{y}$ non-autogressively. We show experiments highlighting both hard and soft constraints (§\ref{['sec:experiments']}).
• Figure 2: Different kinds of functions can be incorporated into MuCoLa defined on a shared embedding table $\mathbf{E}$. (Left) Language Modeling objective defines a per-token loss directly on the sequence of embeddings. For every token this loss provides gradients to update $\tilde{e}_i$ via backpropagation through the transformer layers and directly to $\tilde{e}_{i+1}$ through the negative loss likelihood loss as computed in §\ref{['subsec:lagrangian']}. This is used as a primary objective for the underlying LM and can also be used for classification as discussed in §\ref{['subsec:sentiment-controlled-generation']} (Center) Classification objective defined on probability of the desired label. The classifier gets the token embeddings $\tilde{\mathbf{e}}$ directly as input and updates the embedding using gradients obtained via backpropagation from the transformer layers (Right) Lexical loss defined on the embeddings directly (without the use of additional models) to include desired keywords or phrases in the output sequence (§\ref{['sec:hard-constraints']}). In practice any combination of these constraints can be used.