Semantic Probabilistic Control of Language Models

TL;DR

The paper tackles semantic control of language models, addressing the challenge of steering generations to satisfy non-lexical attributes defined by a sequence-level verifier. It introduces SConE, a training-free method that uses the verifier's gradient to estimate the probability a candidate continuation satisfies the target attribute and reweights the next-token distribution accordingly, all within an approximate, locally contextualized LM. By employing a pseudolikelihood-based contextual distribution, a first-order Taylor expansion of the verifier, and a circuit-based computation of expected embeddings, SConE performs tractable, inference-time constrained generation with high constraint satisfaction. Empirically, SConE yields substantial improvements in toxicity detoxification/toxification, sentiment control, and topic adherence while maintaining perplexity, outperforming baselines like random, beam search, and Best-of-N. This approach enables precise semantic control without data or fine-tuning, offering practical benefits for safer and more controllable LM deployments.

Abstract

Semantic control entails steering LM generations towards satisfying subtle non-lexical constraints, e.g., toxicity, sentiment, or politeness, attributes that can be captured by a sequence-level verifier. It can thus be viewed as sampling from the LM distribution conditioned on the target attribute, a computationally intractable problem due to the non-decomposable nature of the verifier. Existing approaches to LM control either only deal with syntactic constraints which cannot capture the aforementioned attributes, or rely on sampling to explore the conditional LM distribution, an ineffective estimator for low-probability events. In this work, we leverage a verifier's gradient information to efficiently reason over all generations that satisfy the target attribute, enabling precise steering of LM generations by reweighing the next-token distribution. Starting from an initial sample, we create a local LM distribution favoring semantically similar sentences. This approximation enables the tractable computation of an expected sentence embedding. We use this expected embedding, informed by the verifier's evaluation at the initial sample, to estimate the probability of satisfying the constraint, which directly informs the update to the next-token distribution. We evaluated the effectiveness of our approach in controlling the toxicity, sentiment, and topic-adherence of LMs yielding generations satisfying the constraint with high probability (>95%) without degrading their quality.

Figures (4)

• Figure 1: An illustration of our proposed approach. (a) Given a prefix, the LM defines a distribution over possible next-tokens. (b) For each possible next-token, we efficiently simulate future generation. (c) An LM sample induces an approximate LM distribution assigning high probability to similar samples and low probability to dissimilar samples. (d) Evaluating a verifier on a single simulated generation, we can use the first-order information to locally approximate the verifier on all possible generations, factoring in the probability of each generations w.r.t. the LM. (e) This yields a probability of the constraint, $\mathcal{A}$, the set of all generations satisfying a target attributed $\bm{a}$ being satisfied, used to reweigh the next-token distribution. (f) This results in a new distribution that discounts fluent but constraint violating generations in favor of less likely but constraint satisfying generations.
• Figure 2: A technical overview of our approach. (top left) We start by sampling an approximate generation ${\mathbf{s}}$ using Gibbs sampling $\tilde{p}$ conditioned on the prefix from the model's marginal conditionals, $p({\mathbf{y}}_{i} \mid {\mathbf{y}}_{-i}) \forall_i$. Conditioned on ${\mathbf{s}}$, the models marginal conditionals induce a distribution on all generations, assigning higher probabilities to similar sentences and lower probabilities for dissimilar sentences, which we visualize for the top-3 tokens for clarity of exposition. (bottom left) We can parameterize a circuit using the above distribution, yielding a closed-form, tractable representation of probability distribution defined in \ref{['eq:local_psl']}, where read left to right, every leaf node corresponds to a categorical distribution on ${\mathbf{y}}_i$ (right) Such a representation enables us to compute the expected embeddings w.r.t. the distribution in the neighborhood of the sample ${\mathbf{s}}$ by substituting token embedding for corresponding embeddings at leaf nodes, computing weighted sums of embeddings at sum nodes, and taking sums at product nodes. This allows us to plug the expected embedding into \ref{['eqn:tractable-taylor-approximation']} to yield the constraint probability.
• Figure : SConE
• Figure : Hogwild! Gibbs Sampling