Research map / Large language models / Language, safety and interpretability

TRACE Back from the Future: A Probabilistic Reasoning Approach to Controllable Language Generation

ICML 2025 · Large language models › Language, safety and interpretability · approach: Controlled text generation · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes TRACE, which distills an HMM from a base LM and uses a factorized log-linear attribute classifier to compute expected attribute probability exactly, reweighting LM next-token probabilities for controllable generation.

Problem

Controllable language generation requires global, look-ahead attribute satisfaction (e.g., non-toxicity, persona, topic), but autoregressive decoding is myopic and existing EAP approximations via training or sampling are costly, high-variance, and slow to adapt to rare attributes.

Approach

Distill a tractable HMM once from the base LM, restrict attribute classifiers to factorized (log-linear) token-weight form so that expected attribute probability over HMM futures can be computed exactly via forward-backward recursion, and fit new attributes with a log-MSE loss on oracle scores; composition is done by multiplying token weights.

Key results

• On GPT-2 Large detoxification, TRACE reaches avg. max toxicity 0.163 and toxicity probability 0.016, versus 0.180 / 0.026 for DPO and 0.385 / 0.254 for the base GPT2, with dist-2 0.85.
• On Gemma-2B, TRACE lowers avg. max toxicity from 0.359 to 0.189 (DPO: 0.222) while keeping dist-2 at 0.86.
• TRACE adapts to 76 RoleBench characters with about 10 seconds of training each and reports about 1.2x inference time relative to the baseline LM.

Measured effects reported in the paper

Surprising findings

• A simple factorized (log-linear) classifier, despite toxicity being non-factorizable (neural CE 0.007 vs factorized 0.386), outperforms neural-classifier decoding methods like GeDi and FUDGE on detoxification.

Limitations

• Attributes like toxicity are not fully factorizable, so the factorized classifier is an approximation.
• Control quality depends on the quality of the distilled HMM; TRACE trades toxicity for fluency via the decoding transformation scale.
• Composition assumes attribute independence given the text.

Experimental setup

Datasets: RealToxicityPrompts RoleBench News Category dataset

Benchmarks: RealToxicityPrompts

Metrics: average maximum toxicity toxicity probability perplexity distinct-2 distinct-3 role quality political relevance inference time ratio

Baselines: DAPT PPO Quark DPO GeDi FUDGE DExperts PPLM MuCoLa Mix and Match Prompting

Builds on: hidden Markov models probabilistic circuits GeDi FUDGE DExperts

Closest papers by meaning

• Adaptable Logical Control for Large Language Models (NeurIPS 2024 · similarity 0.88)
• Syntactic and Semantic Control of Large Language Models via Sequential Monte Carlo (ICLR 2025 · similarity 0.88)
• Latent Principle Discovery for Language Model Self-Improvement (NeurIPS 2025 · similarity 0.87)
• Markovian Transformers for Informative Language Modeling (ICLR 2026 · similarity 0.87)
• Revealing Behavioral Plasticity in Large Language Models: A Token-Conditional Perspective (ICML 2026 · similarity 0.87)
• XTC: Head-Aware Sampling by Excluding Top Choices (NeurIPS 2026 · similarity 0.87)