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

arXiv: https://arxiv.org/abs/2504.18535 · HTML: https://mlsanity.com/paper/2504.18535/

## Key facts

- Cited by 8 papers in the map (data through 14 Oct 2026).
- Placed in Test-time compute and efficient thinking, one of 17 directions in LLM reasoning.
- Builds on 31 of its 41 references that are in the map.
- First posted on arXiv 25 Apr 2025.

## Placed in

- LLM reasoning / [Test-time compute and efficient thinking](https://mlsanity.com/lr/test-time-compute-and-efficient-thinking/)

## Builds on (most cited first, 30 of 31)

- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Direct Preference Optimization: Your Language Model is Secretly a Reward Model](https://mlsanity.com/paper/2305.18290/) (2023)
- [The Curious Case of Neural Text Degeneration](https://mlsanity.com/paper/1904.09751/) (2019)
- [Don't Stop Pretraining: Adapt Language Models to Domains and Tasks](https://mlsanity.com/paper/2004.10964/) (2020)
- [RealToxicityPrompts: Evaluating Neural Toxic Degeneration in Language Models](https://mlsanity.com/paper/2009.11462/) (2020)
- [Plug and Play Language Models: A Simple Approach to Controlled Text Generation](https://mlsanity.com/paper/1912.02164/) (2019)
- [Compositional Visual Generation with Composable Diffusion Models](https://mlsanity.com/paper/2206.01714/) (2022)
- [Towards a Human-like Open-Domain Chatbot](https://mlsanity.com/paper/2001.09977/) (2020)
- [DExperts: Decoding-Time Controlled Text Generation with Experts and Anti-Experts](https://mlsanity.com/paper/2105.03023/) (2021)
- [FUDGE: Controlled Text Generation With Future Discriminators](https://mlsanity.com/paper/2104.05218/) (2021)
- [Quark: Controllable Text Generation with Reinforced Unlearning](https://mlsanity.com/paper/2205.13636/) (2022)
- [A Mechanistic Understanding of Alignment Algorithms: A Case Study on DPO and Toxicity](https://mlsanity.com/paper/2401.01967/) (2024)
- [COLD Decoding: Energy-based Constrained Text Generation with Langevin Dynamics](https://mlsanity.com/paper/2202.11705/) (2022)
- [Controlled Decoding from Language Models](https://mlsanity.com/paper/2310.17022/) (2023)
- [Detoxifying Language Models Risks Marginalizing Minority Voices](https://mlsanity.com/paper/2104.06390/) (2021)
- [Einsum Networks: Fast and Scalable Learning of Tractable Probabilistic Circuits](https://mlsanity.com/paper/2004.06231/) (2020)
- [Mix and Match: Learning-free Controllable Text Generation using Energy Language Models](https://mlsanity.com/paper/2203.13299/) (2022)
- [Tractable Control for Autoregressive Language Generation](https://mlsanity.com/paper/2304.07438/) (2023)
- [Gradient-Based Constrained Sampling from Language Models](https://mlsanity.com/paper/2205.12558/) (2022)
- [Transfer Q Star: Principled Decoding for LLM Alignment](https://mlsanity.com/paper/2405.20495/) (2024)
- [News Category Dataset](https://mlsanity.com/paper/2209.11429/) (2022)
- [Controllable Text Generation with Neurally-Decomposed Oracle](https://mlsanity.com/paper/2205.14219/) (2022)
- [Scaling Up Probabilistic Circuits by Latent Variable Distillation](https://mlsanity.com/paper/2210.04398/) (2022)
- [Adaptable Logical Control for Large Language Models](https://mlsanity.com/paper/2406.13892/) (2024)
- [A Pseudo-Semantic Loss for Autoregressive Models with Logical Constraints](https://mlsanity.com/paper/2312.03905/) (2023)
- [Building Efficient Universal Classifiers with Natural Language Inference](https://mlsanity.com/paper/2312.17543/) (2023)
- [Image Inpainting via Tractable Steering of Diffusion Models](https://mlsanity.com/paper/2401.03349/) (2023)
- [LiSeCo: Linear Semantic Control for Language Generation](https://mlsanity.com/paper/2405.15454/) (2024)
- [On Tractable Computation of Expected Predictions](https://mlsanity.com/paper/1910.02182/) (2019)
- [Unlocking Anticipatory Text Generation: A Constrained Approach for Large Language Models Decoding](https://mlsanity.com/paper/2312.06149/) (2023)

## Built on by (most cited first, 8 of 8)

- [Dream-VL & Dream-VLA: Open Vision-Language and Vision-Language-Action Models with Diffusion Language Model Backbone](https://mlsanity.com/paper/2512.22615/) (2025)
- [Entropy-Aligned Decoding of LMs for Better Writing and Reasoning](https://mlsanity.com/paper/2601.01714/) (2026)
- [ExplainFuzz: Explainable and Constraint-Conditioned Test Generation with Probabilistic Circuits](https://mlsanity.com/paper/2604.06559/) (2026)
- [Learning Tractable Distributions Of Language Model Continuations](https://mlsanity.com/paper/2511.16054/) (2025)
- [The Expressivity Boundary of Probabilistic Circuits: A Comparison with Large Language Models](https://mlsanity.com/paper/2605.12940/) (2026)
- [Conditional Attribute Estimation with Autoregressive Sequence Models](https://mlsanity.com/paper/2605.14004/) (2026)
- [Grab a Coffee: Future-Aware Guidance for Discrete Diffusion with Compiled Objectives](https://mlsanity.com/paper/2609.35924/) (2026)
- [Provably Tractable NFA-Constrained Language Generation via HMMs](https://mlsanity.com/paper/2609.40185/) (2026)

## Abstract (arXiv)

As large language models (LMs) advance, there is an increasing need to control their outputs to align with human values (e.g., detoxification) or desired attributes (e.g., personalization, topic). However, autoregressive models focus on next-token predictions and struggle with global properties that require looking ahead. Existing solutions either post-train LMs for each new attribute--expensive and inflexible--or approximate the Expected Attribute Probability (EAP) of future sequences by sampling or training, which is slow and unreliable for rare attributes. We introduce TRACE (Tractable Probabilistic Reasoning for Adaptable Controllable gEneration), a novel framework that efficiently computes EAP and adapts to new attributes through tractable probabilistic reasoning and lightweight control. TRACE distills a Hidden Markov Model (HMM) from an LM and pairs it with a small classifier to estimate attribute probabilities, enabling exact EAP computation over the HMM's predicted futures. This EAP is then used to reweigh the LM's next-token probabilities for globally compliant continuations. Empirically, TRACE achieves state-of-the-art detoxification results with only 20% decoding overhead, yields 76 low-resource personalized LMs within seconds, and seamlessly extends to composite attributes. Our code is available at: https://github.com/yidouweng/trace.