TRACE Back from the Future: A Probabilistic Reasoning Approach to Controllable Language Generation
Authors: Gwen Yidou Weng
Benjie Wang
Guy Van den Broeck
Venue: ICML 2025 poster (2025)
Keywords: Controlled Generation
Probabilistic Reasoning
Tractable Inference
Hidden Markov Models (HMM)
Detoxification
Primary area: probabilistic_methods
TL;DR: TRACE achieves efficient and adaptable control of language models by using a one-time distilled HMM and a lightweight log-linear classifier to perform exact probabilistic reasoning over future text.

Abstract: 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.

TRACE Back from the Future: A Probabilistic Reasoning Approach to Controllable Language Generation
Gwen Yidou-Weng*
Benjie Wang**
Guy Van den Broeck
Abstract
Atribute Probabiliyy
Tulure lext
the ass
0.3
As large language models (LMs) advance. there increasing need to control their outputs align with human values (e.g.. detoxification) desired attributes (e.g., personalization, topic). However, autoregressive models focus on next-token pre- dictions 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 At- tribute Probability (EAP) of future sequences by sampling or training, which is slow and unreli- able 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 reason- ing and lightweight control. TRACE distills a Hid- den Markov Model (HMM) from an LM and pairs it with small classifier to estimate attribute prob- abilities, 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-arr detoxification sults with only 20% decoding overhead, yields 76 low-res low-resource personalized LMs within seconds. and seamlessly extends to composite attributes.
(nontoxic)
the butt
0.15
the neck
0.05
Plm 0.3
MTRVCE x 0.03
It's pain
Tuture text with handle
PhnnlFFs [Pe))
0.1
to
0.1
EAP 0.8
0.08
Figure 1. TRACE reweighs LM next token probabilities by "looka- head," computing the Expected Attribute Probability (EAP), Σr,, p(s Ig+ Plaam (2>1 =gge). using an HMM to
tractably compute the expectation of probabilistic classifier
search in LM alignment with objectives such as detoxifica- tion (Gehman et al., 2020; Xu et al., 2021), where abundant data and benchmarks exist. Meanwhile, there growing interest involving specialized or personal attributes under low-data conditions (Adiwardana et al., 2020; Xu al., 2021), as well as compositional attributes for creating com- plex or novel outputs that rarely appear during standard training (Liu et al., 2022).
Controllable text generation is challenging because most language models are autoregressive, generating each token solely from its predecessors without looking ahead-yet many attributes depend the entire text. One line of solu- tions modifies the base model's distribution via fine-tuning or post-trainge post-training (e.g., RL, RLHF) (Gururangan et al., 2020; Rafailov et al., 2023; Schulman et al., 2017), but these ap- proaches can be highly expensive and risk degrading fluency or diversity (Kumar et al., 2022). Thus, there a need for more lightweight solutions that can directly leverage pre- trained LMs by changing the decoding strategy.
1. Introduction
As large language models (LMs) become more ubiquitous in commercial products and daily life, there is growing need to control their outputs. There is large body of re-
Equal contribution 'Department Computer Science, Univer- sity of California, Los Angeles, USA. Correspondence to: Gwen Weng <gwenweng @ucla.edu>>..
Fundamentally, we argue that decoding- based controllable generation is a probabilistic reasoning task: we wish to sample from the language model's distribution conditional on some attribute. For example, some approaches utilize sampling to find sttribute-consistee generations (Yang
Proceedings of the 4rnd International Conference on Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).
Ouu code is available httpp::/giiii com/ gidouweng/trrac.

---

TRACE Back from the Future: A Probabilistic Reasoning Approach to Controllable Language Generation
Klein, 2021; Tu al., 2024; Chakraborty et al., 2024; Mud- gal al., 2024), but this can be computationally expensive. Other methods train expensive LM-based discriminators to predict satisfaction of an attribute given partial text se- quence (Krause et al., 2021; Meng et al., 2022; Liu et al., 2021), but this requires large amounts of data and training time. Both struggle with rare attributes due to limited data and high variance.
2. Related Work
Controllable text generation methods fall into two main categories: those that modify the language model (LM) via training, and those that steer a pre-trained model through decoding-tingg time interventions.
2.1. Training Methods
One line of approaches modifies the base LM parameters, typically through fine-tuning or reinforcement learning, to instill desired attributes directly into the model's distribu- tion. DAPT (Gururangan et al., 2020) fine-tunes the base LM on Nomain-speccnf data. PPO (Schulman al., 2017) uses a reward model and policy gradients to fine-tune the LM towards desired behaviors like non-toxicity, while DPO (Lee et al., 2024) aligns the model using pairwise prefer- ences without an explicit reward model. Similarly, Quark (Lu et al., 2022) uses an RL-like procedure conditioned on a learned reward token. A major drawback of these methods is the need for substantial data and costly retraining of the large LM for each new attribute or set of attributes. Further- more, modifying the base LM risks degrading its general fluency and diversity (Kumar et al., 2022).
In this work, we propose framework for controllable generation that utilizes traciable models to approximate the probability of satisfying an attribute given a partial se- quence. This approach avoids costly sampling methods and retraining the base language model (LM) for each tar- get attribute. First, we perform a one-time distillation of tractable Hidden Markov Model (HMM) to approximate the base LM. Next, for each desired attribute, we train log-linear classifier to estimate the probability of satisfying a target attribute given the entire text. During decoding, we use Bayesian conditioning to reweight next-token probabili- ties based on the likelihood that future sequences- predicted by the HMM HMMM-oomly with these attributes, which we re- fer to as Expected Attribute Probability (EAP) (Figure 1). Crucially, the combination of the tractable HMM and log- linear classifier enables us to compute EAP over all of the exponentially many future sequences efficiently and exactly (over the HMM state space).
2.2. Decoding Methods
An alternative direction focuses on modifying the decoding process of a fixed, pre-trained LM to steer generation to- wards desired attributes, often by incorporating an estimate of the Expected Attribute Probability (EAP) of future text.
Our method, Tractable probabilistic Reasoning for Adaptable Constrained gEneration (TRACE), constitutes a uniquely lightweight solution that enforces control with almost zero decoding-time overhead over the base LMM It further decouples generative model training from control, eliminating the need to retrain the LM or HMM for new ob- jectives. Adapting to novel or rare attributes merely requires training a small classifier in seconds. Handling composi- tions of attributes also straightforward by multiplying the EAP of each attribute during decoding.
Training Discriminators. Since exact EAP computation requires summing over future sequences, many methods train auxiliary models to estimate attribute satisfaction from partial generations. FUDGE (Yang & Klein, 2021) and NADO (Meng et al., 2022) train discriminators to approxi- mate EAP; DExperts (Liu et al., 2021) blends expert and anti-expert expert LMs; GeDi (Krause et al., 202 1) uses attribute- conditioned guides with Bayes' rule; and LiSeCo (Cheng et al., 2024) applies a linear probe in latent space. A key challenge across these methods is that accurately estimating future attribute satisfaction from prefixes requires nontrivial lookahead, and training a separate large, auxiliary model per attribute adds significant overhead.
We evaluate TRACE on three important tasks: (1) Detoxi- ficution: TRACE outperforms expensive RL, training- and sampling-based based baselines on GPT2-large and Gemma 2B. (2) Personalized LLMs: TRACE adapts to 76 distinct charac- ters in about three seconds each, outperforming prompting approaches with only a few hundred training samples, and taking only 1.2x time per-token standard decoding. (3) Compositional Attributes: TRACE seamlessly gener- ates texts satisfying multiple attributes-e.g. political and nontoxic-a combination too sparse for training or sampling methods. Overall, TRACE is a simple, lightweight con- trollable generation approach that achieves state-of ofhee---rrrtt detoxification performance, extends to low-resource and composite text generation, and scales well modern LLMs.
Sampling, Other methods incorporate EAP by sampling future sequences to estimate expected outcomes. Some per- form limited token-level lookahead (Mudgal et al., 2024; Chakraborty et al., 2024; Tu et al., 2024), managing com- plexity by restricting the search space. Others employ MCMC-style sampling on entire sequences using energy- based models (e.g., MuCoLa (Kumar et al., 2022), Mix and Match (Mireshghallah et al., 2022), COLD (Qin et al., 2022)). Both approaches incur significant computational overhead during decoding and the resulting EAP estimates can have high variance, especially for rare attributes.
2

---

TRACE Back from the Future: A Probabilistic Reasoning Approach to Controllable Language Generation
While all EAP-based control methods must approximate the intractable EAP under the hase LM, the nature of approx- imation varies substantially. Discriminator-guudder and sampling- -baeed approaches estimate the EAP directly during decoding, often requiring per-attribute training or in- curring high variance and runtime cost. In contrast, TRACE shifts the approximation burden to a one-time HMM distil- lation step. where Phmm & Pin- Conditioned on this distilled HMM, TRACE enables an exact and tractable computation of EAP- EAP-Phaml( (s I Te (t)) relatvve to the HMM distri- bution (Section 4.2). This design choice aprioxeapppoiiiaiing the LM via a one-time HMM distillation enables efficient, Low-variance and adaptable EAP computation during decod- ing, avoiding the cost of per-attribute retraining or sampling, and proves more effective in practice (Section 5.2).
bution Pio (Fimm s). This conditional distribution can be
decomposed autoregressively as:
Pbu Pbn(21:: () II pim (32 8)
(2)
However, sampling from the conditional next-token distribu- tion Pm T<ee 8) is generally intractable. Using Bayes' rule, this is given by:
Pim(Fs [<(, x (an Ze() (im Zt, Fe1). (3)
The first term is simply the LM next token distribution. The second term is the probability of satisfying the attribute s; this requires summing over all possible continuations 2> which is exponential in sequence length:
Pim(s Ee))
Σ E(())gg()25)) Pim(I xgg). 2o
2.3. Control via Tractable Models
Since the classifier is probabilistic, we will also call this the expected attribute probability (EAP). The EAP is used to reweight the possible token generations x according to how likely they are to result in a text that eventually satisfies the desired attribute. Several existing approaches effectively aim to approximate this computationaaly hard sum. For example, GeDi and DExperts train wwie discriminators to guide the LM. Other approaches, such as Controlled Decoding and MuCoLa sample future sequences for lookahead. In this work, we aim for a tractable way to incorporate future-sequence information without expensive sampling or retraining, using Hidden Markov Models.
Tractable probabilistic models (Choi et al., 2020) such as HMMs enable efficient computation of various quantities such marginals and the probability of satisfying logical constraints compuutttossomutttiiiinnnoo that provably hard even to approximate on autoregressive models (Roth, 1996). Prior work has used such models to enforce logical or lexical constraints in generative modeling (Liu et al., 2024a;b). In language modeling, Ahmed et al. (2023) proposed using local tractable approximations for training autoregressive models to satisfy logical constraints. Meanwhile, Zhang et al. (2023; 2024) used HMMs to enforce logical constraints such as those given by deterministic finite automata (DFA).
While powerful for such formally specifiable constraints, this approach does not readily extend to high-level seman- tic attributes such as style, safety, or persona, which lack symbolic definitions and depend on the overall meaning of the text. Our work, TRACE, bridges this gap by employing the HMM uniquely for semantic control, enabling efficient probabilistic reasoning about semantic attributes (via EAP computation) rather than enforcing logical rules.
3.2. Hidden Markov Models
Hidden Markov models (HMM) specify a joint distribution over set of latent variables 2i. and observed variables Ti:n
p(rim 71m)
= 41m)=
p(())((()) 2)) I se-1)p(z; ze). t=2
(4)
For language modeling, each ,, takes values in {o..... ......,. 1}. where h is the hidden state size, while the observed variables Ie are tokens taking values in (0...., V 1). where V is the vocabulary size. The (homogeneous) HMM has h2 parameters for the transition matrix p(2, 1)) hV parameters for the emission matrix (2), and h parame- ters for the initial hidden state distribution p(2)). The key advantage to using HMMs for language modeling their tractability; many quantities, such as the probability of a token sequence, can be inferred in linear time in the size of the HMM and sequence length. Zhang et al. (2023; 2024) distilled HMM models from large language models for the purpose of generating under logical constraints, such as the presence of a particular keyword. We will instead leverage HMMs to design algorithms for computing the expected attribute probability efficiently.
3. Preliminaries
3.1. Controllable Generation
We consider generating a text (sequence of tokens) Fis of length n from a language model (LM). In controllable gen- eration, the goal is to generate text from the LM conditional on some attribute s, such as nontoxicity. We assume that the attribute be measured by some probabilistic classifier p(s( E [0, 1] representing the probability (or degree) of satisfaction of the attribute given the full text.
Let us denote the base LM distribution by Pim (T Then the joint distribution over the text Ii:m and attribute s is (1)
Pim ( 8) = Pim (# (psp2(
Our goal is then to generate from the conditional distri-
3

---

TRACE Back from the Future: A Probabilistic Reasoning Approach to Controllable Language Generation
4. Methodology
with the base case Phom (21
p(21)p(()))) To
obtain the conditional, we simply divide by the normalizing
This section details the proposed TRACE methodology. We introduce the core approximation using HMMs to guide LM generation (Section 4.1), present the algorithm enabling tractable EAP computation (Section 4.2), and describe how the attribute classifiers are fitted (Section 4.3).
constant Phanm (T = Es, Phomn (ee zgg)). Note that this computation is independent of the classifier.
Backward Computation The quantity Phoma (T>e/==) can be computed by exploiting the structure of the probability distribution in Equation 4 to rearrange the summation over future latent states Z>xt
4.1. TRACE: Guiding LM with HMM Probabilities
In order approximate the constrained next token probabil- ity p( | Zee,, s) in Equation 3, we propose to approximate the expected attribute probability (s Ie; Zeee) with the corresponding quantity under the HMM:
Phmnn (I 2x) = -EI((ta (p(Es/zi1)) p(az)zz))
Aod 836
= Eotahlkiipooennkkkk
PIRACE(2) ( xc (Im x=() Pham(s Fet,, xTt).
(5)
This is known as the backward algorithm as the evaluation of the summations is performed right to left, backward in time. Now, in order to compute the boxed term tractably, the classifier p(s [(s must have similar structure that en- ables integration into the backward algorithm. A sufficient condition is to restrict to factorized classifiers of the form p(s (1m)) = II. w(x;), where w(z;) is a weight function that assigns a weight for each token in the vocabulary.
Here, Phmm Titt Øy) is the probability, under our HMM and attribute classifier, that the entire future continuation will satisfy the attribute. In contrast to the formulation of Zhang et al. (2023), is not logical attribute that maps to 0 or 1, but instead represents semantic attribute given by a probabilistic classifier (l)) Clearly, if the classifier p()(z1n)) is arbitrary without any structure (e.g.. a neural network), then computing the expected attribute probability will be intractable as we will again need to generate all possible continuations to feed to the classifier.
Then, the boxed term can be expanded as
Σ Phomm (I>) p ((551 xr)))
Next, we describe simple kind of classifier for which the exact computation of EAP is tractable, and then develop an efficient sewwadd-backward
III w((zs) i5t
Plamm (I>e
2x)
wxx;)
style algorithm for doing so. We then describe how to learn the attribute classifier at the token level (Section 4.3), and how to improve performance further through test-time approximations.
i>
(o(x;)
Ehomn
II List
w(xx))
24
4.2. Tractable Computation of EAP
For the expectation term Egmon pute recursively backwards in time as follows. The base case M uxxx))==]] 1, as there are no tokens after In. For t < n, the recursion is
ees w(xx))=z]] we com-
The expected attribute probability (EAP) under the HMM Piamm (s Iee, Wee)) can be rewritten by introducing future sequence Σyل and the hidden state 2,, and marginalizing over them, using the conditional independence property of HMMs: xee II. 1- We obtain that
Epooo
II List
[w(x))
21
Phmmm [8] Des x<ce)
Σ Phomm (Ze (ige) Piomm (1))22)))
pl p(o
Σ ( [aa))) Ermon
II w(xx)
ze+1
Re11
Lis++1
Σ
p( [241))= w((x+1)))
We now discuss how to compute each of these terms.
Forward Computation The computation of Phomm (24 T is typically carried out using the HMM forward algo- rithm, which is based on the following recursion:
Importantly, the values Eppon Iise w(xx) 2t| can be pre-
computed and cached in single backward pass and reused across all generations, as they depend solely on the hidden states 21 and not on the specific prefix Z-t--
Phoam Ze, xe)e)
Σ plare I zx) p(zs| zt-1): Phomm (2 IS.- (6)
TTo ensure that the classifier outputs value in [0, 1], will enforce that all weights wxx)) E [0. 1) also.
1-1
4

---

TRACE Back from the Future: A Probabilistic Reasoning Approach to Controllable Language Generation
Algorithm 1 TRACE: Generating n. Tokens
We assume access an oracle Poracte that scores the proba- bility of sequence Fi:n satisfying target attribute This oracle may be trained classifier (e.g., from human annota- tions) or an external API scoring attribute satisfaction.
Require: HMM Phomm LM Pim- Classifier
Ensure: Generated sequence 2imm
1: for each t from to do
2: Pre-compute Pll, 2r] := Exmm Tixe w (1x))=])) by
Fitting Rare attributes via Log-MSE Our goal is to model attributes that are rare in natural text, such as toxicity or political content. To effectively identify these cases, the learning objective must distinguish texts that exhibit the rare attribute from the more common, neutral ones. Stan- dard objectives like cross-entropy treat all misclassifications symmetrically, which is suboptimal when one class is rare.
Equation (7)
3: end for
4: Initialize So 6- qu. 21:0 +
5: for each L from to n do
6: Compute Phonm (ze (Ige) by Equation (6)
7: Compute Phamm Zet, TT) using Phomm (2t ( and
P(1, =1] by Equation (8)
We therefore use an asymmetric loss that more heavily pe- nalizes misclassifying the rare examples. In particular, we use mean squared error loss in log-space (log-MSE), which amplifies error penalties in the pwwopabbbillty region. To align with this objective, we define each attribute by its ab- sence (e.g., modeling "nontoxicity" rather than "toxicity").
8: Sample 24. ~ Phmmm Tet, Pin(Fee xee)) by
Equation (9)
9: Update T'ee < I't ① Zt
10: end for
11:
Return T1:2
Specifically, for our factorized (log-linear) classifier, the log prg--rbbability of the attribute given the text Ti:m is sum of log-weights, logp(s fimm) = Σ log w(x)). Then, given oracle scores Poracle (potentially transformed, as discussed next), the loss for a text æ is given by
Integration During Generation Algorithmically. as each new token is generated, the forward probability p(zz Tet: œ4) is updated based on the recursion in Equation 6. Using the precomputed backward expectations, the overall expected attribute probability is computed as
Pham [s| Tzz Ie
log Poracto (I)
Σ log w(x))
(10)
II w(xz)
Σ Plumm (Ze (15)) Eggom
II i>1
w(x))
2t
For example, this loss penalizes prediction of 0.5 much
more heavily if the true oracle score is Poracle = 0.11 (toxic) than if it is 0.9 (nontoxic), whereas cross-entropy is indif- ferent to these outcomes.
(8)
Plugging this into Equation 5, we can then compute the constrained next-token distribution:
Probability Transformation for Better Control In prac- tice, it can be beneficial to transform raw oracle probability scores (Ponete) for a clearer separation between desired and undesired outputs, and enforce stricter attribute satisfaction. For instance, 20% non-toxicity score from an oracle can be transformed a lower value to encourage safer generation.
PTRACE(FE (Zqe,)) s) x (an Fee) - Phmm (s Iee (t)
Algorithm 1 summarizes our approach. The precomputation requires O (n<hV + 1:2)) time (backward HMM pass) and requires ((nh) cache memory, corresponding to each hid- den state value at each time. During generation, the compu- tation of the expected attribute probability (s It, .<et)) at each time point takes Ohh hv). In practice, we find this takes negligible time relative to computing the language model's 's next token probability (see Table 4).
To do this, we propose to apply affine transformation to the oracle scores in logit space using (non-negative) scale (b) and shift (c):
Pooracte = a (b-bbb
(b- In
Poracle (1- +
Porralle
+c)
where (22) is the sigmoid function. This transformation reshapes the target distribution; specifically, increasing the scale b pushes intermediate probabilities towards the ex- tremes of 0 and 1. This creates more bimodal distribu- tion, which helps to clearly distinguish between attribute- compliant and non-compliant texts.
4.3. Attribute Classifier Fitting
As discussed, one class of classifiers that make EAP com- putation tractable is product over tokens, that is, linear in log space. In practice, attributes are not fully factorizable (to varying extents); nonetheless, we show that even such a simple classifier, when combined with a HMM, captures EAP sufficiently accurately to outperform baselines using neural classifiers (e.g., GeDi, FUDGE) (See 5.2).
This probability transformation can also be applied at decod- ing time to the EAP, Phmm (s Øe, xz2)). before it influences token generation (Eq. 9). This can be used to sharpen the
5

---

TRACE Back from the Future: A Probabilistic Reasoning Approach to Controllable Language Generation
Table Detoxification on ReafToxicityPrompy (Gchman et al., 2020). Results shown for GPT-2 (10k prompts) and Gemma-2B (Ik prompts). TRACE applies training- and decoding-tings time transformation (see Section 4.3). Struck-throug fluency indicates unnatural
repetition (Holtzman et al., 2020). Baselines: (I) Gururangan et al. (2020); (2) Krause et al. (2021); (3) Yang & Klein (2021); (4) Liu
(2021); (5) Dathathri et al. (2020); (6) Kumar et al. (2022); (7) Schulman et al. (2017); (8) Lu et al. (2022); (9) Lee et al. (2024) (our
implementation)..
Model
Toxicity (L)
Diversity (1)
PPL (
Approach Type
avg. max.
prob.
dist-2
dist-3
GPT-2 Large Results
GPT2
0.385
0.254
0.87
0.86
25.57
Baseline
DAPT
0.428
0.360
0.84
0.84
31.21
Finetuning
2cDj
0.363
0.217
0.84
0.83
60.03
Decoding (Trained Guide)
FUDGE
0.302
0.371
0.78
0.82
42.97*
Decoding (Trained Guide)
DExperts
0.314
0.128
0.84
0.84
32.41
Decoding (Trained Guide)
PPLM
0.520
0.518
0.86
0.86
32.58
Decoding (Logit Control)
MuCoLa
0.308
0.088
0.82
0.83
29.92
Decoding (Sampling)
PPO
0,218
0.044
0.80
0.84
14.27*
RL
Quark
0.196
0.035
0.80
0.84
+2,47*
RL
DPO
0.180
0.026
0.76
0.78
21.59*
RI.
TRACE
0.163
0.016
0.85
0.85
29.83
Decoding (HMM Reasoning)
Gemma-20 2B Results
Gemma-2B
0.359
0.23
0.86
0.85
15.75
Baseline
(PO
0.222
0.06
0.74
0.77
44.39*
RI.
TRACE 0.189
0.02
0.86
0.85
17.68
Decoding (HMM Reasoning)
separation between next tokens v, based on their EAP, thus ensuring stricter control of the attribute.
(2023). Full details on metrics are provided in Appendix F.
Baselines. TRACE is compared against representative base- lines including fine-tuning (DAPT), RL (PPO, Quark, DPO), and various decoding methods (PPLM, GeDi, FUDGE, DExperts, MuCoLa). Results are sourced primarily from prior work (Liu et al., 2021; Lu et al., 2022; Kumar et al., 2022) or run by us (DPO) using the same setup as in Liu et al. (2021). Full details are Appendix E.
In Section 5.2, we empirically justify this intuition and care- fully study the effect of these training-fime and decoding- time probability transformations.
5. Experiments
We evaluate TRACE on range of challenging controllable generation tasks: detoxification, low-resource role- ole-ppayyngg and compositional control involving political and non-toxic text attributes. This section details the experimental setup and presents the main results for each task, along with an analysis of TRACE's efficiency and properties.
Implementation Details. The HMMs used by TRACE were distilled from base LMs (GPT2-Large, Gemma-28 following Zhang al. (2023); Liu al. (2023); details the different HMM configurations, including hidden state sizes and training parameters, are in Appendix C. Once distilled, the HMM model is fixed and is reused without further training for each attribute. Attribute classifiers (e.g. nontoxicity) were trained according to the method in Section 4.3; datasets, oracles, fitting procedures, and transformation parameters (b, c) are detailed in Appendix D.
5.1. Experimental Setup
Evaluation Metrics. Metrics vary by task. For the pri- mary detoxification task (Section 5.2), we follow the setup in Liu et al. (2021) using the RealToxicityyroompt dataset (Gehman al., 2020). We evaluate Toxicity (Perspective API avg. max. toxicity & prob. of any toxic generation (> 0.5) over 25 samples; lower is better), Perplexity (as an automatic measure of fluency, calculated using GPT2- XL; 1. lower is better), and Diversity (Distinct 2-grams, 3-grams; 1 higher is better). We also include supplemen- tary Al evaluations using GPT40-mini. For role-playing (Section 5.3), we measure role quality via classifier proba- bility. For topic control (Section 5.5), we measure political relevance using the zero-shot classifier from Laurer al.
5.2. Detoxification
We evaluate TRACE on detoxification using RealToxici- tyPrompts (Gehman et al., 2020), comparing performance against fine-tuning (DAPT), RL (PPO, Quark, DPO), and decoding methods (GeDi, FUDGE, DExperts, MuCoLa).
'The distillation process involves compiling a HMM to an equivalent probabilistic circuit (Choi et al., 2020) on the GPU, which is trained on LM samples using mini-batch variant of the expectation expectation-maximizzai maximizationn (EM) algorithm (Dempster al., 1977).
6

---

TRACE Back from the Future: A Probabilistic Reasoning Approach to Controllable Language Generation
TRACE Outperforms Existing Methods Table shows TRACE significca reduces toxicity compared to baselines, and maintains high diversity, with minor reductions from GPT-2's baseline scores. RL methods (e.g., PPO, Quark, DPO), despite lowering toxicity, sharply reduce diversity and perplexity, indicative of repetitive or unnatural text gen- eration due to mode collapse. Appendix Tables and further support these observations: TRACE achieves higher conditional entropy than DPO, and receives comparable better ratings from GPT-4 on nontoxicity and diversity.
Training-Time Transformation
Original Dist.
Transformed Dist.
Density
15
Dist.
10
5
0
0.4
Attribtee Probability
Decoding-Time Transformation
Factorized Classifier is Effective Despite Non- Factorizable Attributes. Attributes like toxicity are not fully factorizable, as evidenced by the performance gap between neural and our factorized classifiers (Ap- pendix D). Nonetheless, TRACE outperforms methods like GeDi and FUDGE that use more expressive neural classifiers but rely on approximate EAP estimation. This highlights key insight: exact EAP computation over HMM, even with simple factorized classifier, can rival or surpass more complex classifiers when the inference framework is less precise.
20
Before Trans.
After
15
Densiyy
10
0.0
0,2
0.4
0.6
0.8
1.0
Expected Attribute Probability
Figure 2. Complementary Effects of Probability Transforma- tion. Top: At training time, the transformation concentrates oracle probabilities for undesirable attributes towards 0 to improve classi- fier learning. Bottom: At decoding time, reshapes the unimodal EAP into a bimodal one to enforce stricter generation control.
Scaling TRACE to Larger Language Models. TRACE scales effectively to larger language models such as Gemma- 2B. As shown in Table and supported by GPT-4 evalua- tions (Appendix Table 8), TRACE consistently outperforms DPO reinforcement learning-basee based method in detoxifi- cation performance, while maintaining strong fluency and diversity. These results demonstrate TRACE's broad appli- cability and robustness across model scales.
learning target by reshaping the skewed Original Dist." into bimodal "Transformed Dist.", critically pushing scores for toxic content towards O. The log-MSt objective's 's focus on this low- www-robabillty region helps the classifier learn stronger negative weights for toxic tokens, resulting in a "Fitted Dist." that successfully captures this new peak near 0.
Table Ablation study of TRACE variants GPT-2 Large Rows:
no transformation, training time transformation only, or both
At decoding time, the same transformation is ap- plied directly to the computed Expected Attribute Probabil- ity (EAP). It converts diffuse, unimodal EAP distribution ("Before Trans.") into a sharp, bimodal one ("Aftee Trans."), creating a clearer separation between attribute-complics and non-compliant continuations.
Bottom:
training- and dicoding-time transformation.
TRACE Variant
Toxicity (1)
Diversity (t)
PPT. (J)
avg. max. prob.
dist-2
dist-3
No Transformation
0.353
0.196
0.86
25.44
Training TF Train + Dec TF
0.187
0.026
0.87
0.85
27.51
0.163
0.016
0.85
0.85
29.83
Complementary Roles of Training and Decoding Trans- formations The probability transformation offers com- plementary benefits when applied at different stages of TRACE, as shown by the ablation study in Table 2. Ap- plying the transformation at training time alone ("Train- ing TF") provides an initial reduction in toxicity over the baseline TRACE model with no transformation. The most significant gain comes from applying the transformation at both stages ("TTrinn + Dec TF"); this further reduces average maximum toxicity from 0.1 to 0. 0.163, albeit with corre- sponding increase in perplexity. Figure 2 visually explains the complementary mechanisms driving these gains.
Tradeoff Between Taxicity and Fluency
31
Mean Mean Fluency
30
1.4
29
1.2
28
27
0.160
0.165
0.170
0.175
0.180
0.185
Maximum Texicity
Figure 3. Tuning the Fluency- Detoxffctetootiiiiiion Trade-off. The decoding- transformation multiplier acts as control knob; as it increases, detoxification improves while fluency decreases.
Top: At training time, the transformation creates stricter
7

---

TRACE Back from the Future: A Probabilistic Reasoning Approach to Controllable Language Generation
Table Personalized LM Responses the Question: "How do you define freedom?"
Question
How do you define freedom?
Direct Prompt
It's a topic we deal with very frequently and yet it's a very subjective question with no easy answer. On its surface, the idea of freedom...
Roles
Abraham Lincoln
Thor
Twilight Sparkle
Prompted with
Freedom is fundamental human right that is necessary for individu- als to live full and meaningful livex. It is the ability to make choices...
Freedom is more than just freedom of speech or movement; it's a feeling of being in control of one's 's own des- tiny...
Freedom is powerfil concept, and it's something that I I've thought a lor about throughou my adventures in Equestria...
Role-Specific Instructions
TRACE
Decisive freedom, subjective free- dom, freedom of speech? There are different definitions for what freedom of speech means, but most of them are centered around the idea that it...
Loki, the son of Odin, reminds of what it is to have freedom in the eyes of God. The mighty Norse god of mischief known for being...
Friendship, connection, love... freedom is endless and depends on the person you ask. Some people want their friends be real while others value their ideas and thoughts
Comparison of Role Quality (Prompting vs. TRACE)
the RoleBench dataset (Wang et al., 2024). Each charac- ter's classifier was trained on its corresponding RoleBench training split of 300 question-ane anestion-answer pairs.
0.9
Role Quality (TRACE)
0.7
Table 3 qualitatively compares Gemma-2B responses for three characters across three settings: the base model, few- shot prompting baseline, and TRACE. The prompting base- line uses a role instruction plus 10 question- guestion-answer answer exam- ples from the RoleBench training set (details in Appendix E). Qualitatively, TRACE-guided answers show more distinct, character-reflective content and tone than the baselines.
0.5
0.3
0.2
Quantitatively, GPT2-large, TRACE achieves superior role quality over standard prompting baseline for most of the 76 characters (Fig. 4). This baseline uses role-specific but exemplar-free instruction (details in Appendix E), en- suring a direct comparison of cestruction-followiii minimal input overhead, unlike few-shot prompting which has higher decoding costs (see Section 5.4, Table 4).
0.3
0.5
0.9
Role Quality (Prompting)
Figure 4. TRACE outperforms prompt engineering in role quality.
with
This synergy-whir one transformation aids learning and the other enforces bined performance. As final benefit, the decoding trans- formation's 's scale parameter (b) offers an intuitive knob for tuning this control, allowing users to smoothly trade toxicity for fluency post-hoc (Figure 3).
for their strong com-
5.4. Training and Inference Time Analysis
Table 4. Training and inference time comparison. TRACE achieves minimal training overhead and near- bear-baseline inference cost.
Impact of HMM Quality Finally, the effectiveness of con- trol is sensitive to the quality of the HMM itself. As shown in Appendix Fig. 5, as the distilled HMM's log- likeliioddoood improves with more training, we observe a corresponding decrease in generation toxicity. This highlights that better HMM distillation directly enhances TRACE's performance.
Method
Inf. Ratio
Method
Train Time
Baseline
1.0
Mix and Match
2 hr
Prompting
3.0
DExperts
3 min-16
GeDi DExperts
2.0-3.0
DAPT
16 hr
Mix and Match
7.5
GeDi
5 hr
MuCoLa
15-20
TRACE
10
PPLM
40.0
TRACE
1.2
5.3. Role-playing 76 Characters with Limited Data
Lightweight and Low-Resource Adaptation. A key ad- vantage of TRACE is its rapid, low-resource adaptation: new attributes are integrated by training only a lightweight log-linear classifier. Leveraging this, we personalized both G-rgee and Gemma-2B for 76 distinct characters from
TRACE is designed for rapid adaptation and efficient infer- ence, requiring only one-time HMM distillation from base LM, independent of any specific control attribute.
Training Time. Once the HMM trained, adapting to a
8

---

TRACE Back from the Future: A Probabilistic Reasoning Approach to Controllable Language Generation
Table Composition results (Political-Nontoiice) on Pealooxi with GPre-Large. Metries: Tox (avg, max./prob.>0..5; 4).
Mean Pol 1. PPL J. Diversity (Dist-2/3; 1). "+ Dec. TF applies decoding TF (Sec 4.3) only to the Political EAP.
Models
Max Tox (J)
Any Tox >> 0.5 (4)
Mean Pol (1)
PPL (4)
Dist-2 (t)
Dist-3 (1)
GPT2-L (Base)
0.386
0.257
0.169
25.74
0.87
0.86
TRACE (Detox only)
0.186
0.026
0.168
27.33
0.87
0.85
TRACE (Pol only)
0.379
0.244
0.333
29.32
0.87
0.86
TRACE (Detox + Pol)
0.190
0.027
0.344
29.71
0.87
0.86
new attribute only requires fitting a lightweight classifier. This process takes up 10 seconds, in sharp contrast to baselines like DAPT, DExperts. GeDi, and Mix and Match, which requires anywhere from minutes to 16 hours of train- ing on various GPUs. This rapid adaptation makes TRACE highly suitable for dynamic or low-resource scenarios.
6. Conclusion
We introduced TRACE, a lightweight framework for con- trollable text generation that uses tractable probabilistic inference. Its core strengths are efficiency, adaptability, and strong performance, all achieved without modifying or finetuning the base LM. ACE distills Hidden Markov Model (HMM) from a base LM and combines it with simple, efficiently trained classifiers to tractably compute the Ex- pected Attribute Probability (EAP), which guides generation towards desired attributes.
Inference Time. TRACE maintains efficiency during gen- eration. While alternative approaches that rely on in-context learning, diseriminators, or iterative updates incur substan- tial decoding wosts-with inference ratios ranging from 7.5x to 40x the Mnseline-TRACE introduces only minor over- head. This efficiency stems from its design; the HMM's Expected Attribute Probability (EAP) calculation uses a pre- computed backward pass and an efficient forward update per token. Consequently, the total inference time only about 1.2x that of the baseline language model, the process remains dominated by the base model's computation.
Empirically, TRACE achieves state-of-the-art detoxifica- tion with only ~20% decoding overhead. It also excels in low-resource scenarios, adapting to personalize genera- tion for 76 distinct characters in seconds, and successfully composes multiple attributes, such as generating text that is simultancously political and non-toxic.
Despite TRACE's already strong empirical performance, its capabilities can be further boosted by improving the expressivity of both the HMM and the attribute classifier. Our results show that higher-quality HMMs yield better re- sults, suggesting that improved distillation techniques are a promising direction. Additionally, while our simple fac- torized classifiers are effective even for non-factorizable at- tributes, extending TRACE to incorporate more expressive models as tractable probabilistic circuits (Khosravi et al., 2019; Choi et al., 2020) that remain compatible with efficient EAP computation could enable stronger control over complex attributes like long-range coherence.
5.5. Composition: Political and Nontoxic Texts
A key benefit of TRACE is its ability to compose multiple attributes without retraining. Conditioning on the conjunc- tion of two attributes (81,82) is achieved by multiplying their probabilities during decoding, based on the modeling assumption of attribute independence given the text (x): p(81 and 82 lx) = p(s (=))())/z)) For TRACE's factorized classifiers, this simplifies to creating a new composite clas- sifier by multiplying the token weights (w) of the individual classifiers: w/ = w/ w² To demonstrate, we task the model with generating text that is simultaneously politi- cal and nontoxic. This combination is rare, making joint training challenging for other methods, but TRACE remains effective as only needs the individually trained classifiers.
A Acknowledgements
This work was funded in part by the DARPA ANSR, CODORD, and SAFRON programs under awards FA8750- 23-2-0004 HR001 12590089, and HR001 12530141, NSF grant IIS 19466144, and gifts from Adobe Research, Cisco Research, and Amazon. Approved for public release; distri- bution is unlimited.
The results in Table demonstrate effective composition. Controlling for single attribute primarily affects only that dimension: TRACE (Detox only) reduces Max Tox- icity from 0.386 to 0.186 while leaving the Mean Politi- cal score nearly unchanged (0. 169 vs 0.168). Conversely, TRACE (Pol only) increases the Mean Political score from 0. to 0.333 with almost no change in toxicity. The compositional TRACE (Detox Pol) approach suc- cessfully achieves both goals simultaneously, reaching toxicity level (0. 190) and a political score (0.344) compara- ble to their respective single-attribute control settings.
9

---

TRACE Back from the Future: A Probabilistic Reasoning Approach to Controllable Language Generation
Impact Statement
Online, July 2020. Association for Computational Lin- guistics. doi: 10. 11/86655/1112222..ammmmaaaaaa
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
Holtzman, A., Buys, J., Du, L.,, Forbes, M., and Choi, Y. The
curious case of neural text degeneration. In International Conference on Learning Representations, 2020.
Khosravi, P., Choi, Y., Liang, Y., Vergari, A., and Van den Broeck, G. On tractable computation of expected pre- dictions. Advances in Neural Information Processing Systems, 32, 2019.
References
Adiwardana, D., Luong, M.-T., So, D. R., Hall, J., Fiedel, N., Thoppilan, R., Yang, Z., Kulshreshtha, A., Nemade, G., Lu, Y., et al. Towards a human- open-domain chatbot. arXiv preprint 1rXiv:2001.09977 2020.
Krause, B., Gotmare, A. D., McCann, B., Keskar, N. S., Joty, S., Socher, R., and Rajani, N. F. GeDi: Gener- ative Discriminator Guided Sequence Generation. In Moens, M.-F., Huang, X., Specia, L., and Yih, S. W.- t. (eds.), Findings of the Association for Computational Linguistics: EMNLP 2021, pp. 4929-4952, Punta Cana, Dominican Republic, November 2021. Association for Computational Linguistics. doi: 10. 2011183/11/2222 finding--emnlp4224.
Ahmed, K., Chang, K.-W., and Van den Broeck, G. A Pseudo-Semantic Loss for Autoregressive Models with Logical Constraints. Advances in Neural Information Processing Systems, 36: 166218525-1820
December 2023.
Chakraborty, S., Ghosal, S. S., Yin, M., Manocha, D., Wang, M., Bedi, A. S., and Huang. F. Transfer q-star: Princi- pled decoding for IIm aligmment. Advances in Neural Information Processing Systems, volume 37, pp. 101725- 101761, 2024.
Kumar, S., Paria, B., and Tsvetkov, Y. Gradient- con- strained sampling from language models. In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pp. 2251-2277, 2022.
Cheng, E., Baroni, M., and Alonso, C. A. Linearly con- trolled language generation with performative guarantees. arXiv preprint arXi: 2024.
Laurer, M., van Atteveldt, W., Casas, A., and Welbers, K. Building Efficient Universal Classifiers with Natural Lan- guage Inference, December 2023. arXiv:2312.17544 [cs].
Choi, Y.. Vergari, A., and Van den Broeck, G. Probabilistic circuits: A unifying framework for tractable probabilistic models. oct 2020.
Lee, A., Bai, X., Pres, I., Wattenberg, M., Kummerfeld, J. K., and Mihalcea, R. A mechanistic understanding of alignment algorithms: A case study on dpo and toxicity. In International Conference on Machine Learning, pp. 26361-26378. PMLR, 2024.
Dathathri, S., Madotto, A., Lan, J., Hung, J., Frank, E., Molino, P., Yosinski, J., and Liu, R. Plug and play lan- guage models: A simple approach to controlled text gen- eration. In International Conference on Learning Repre- sentations, 2020.
Liu, A., Sap, M., Lu, X., Swayamdipta, S., Bhagavatula, C., Smith, N. A., and Choi, Y. DExperts: Decoding- Time Controlled Text Generation with Experts and Anti- Experts. In Zong, C., Xia, F., Li, W., and Navigli, R. (eds.), Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Pro- cessing (Volume Long Papers), pp. 6691-6706, Online, August 2021. Association for Computational Linguistics. doi: 10. 18653/v 1.2.0211.cccc. lacl-oo.5.1.--
Dempster, A. P., Laird, N. M., and Rubin, D. B. Maximum likelihood from incomplete data via the em algorithm. Journal of the royal statistical society: series B (method- ological), 391:)1122, 1977.
Gehman, S., Gururangan, S., Sap, M., Choi, Y., and Smith, N. A. RealToxicityyroopt:: Evaluating Neural Toxic Degeneration in Language Models. In Cohn, T., He, Y., and Liu, Y. (eds.), Findings of the Association for Com- putational Linguistics: EMNLP 2020, pp. 3356-3369, Online, November 2020. Association for Computational findings findings-emnll.0000 emnlp.301.
Liu, A., Zhang, H., and Van den Broeck, G. Scaling up probabilistic circuits by latent variable distillation. In Proceedings of the International Conference on Learning Representations (ICLR), may 2023.
Linguistics. doi: 10.18863/v 1/2033/ 1/20..innnn
Gururangan, S., Marasović, A., Swayamdipta, S., Lo, K., Beltagy, I., Downey, D., and Smith, N. A. Don't Stop Pre- training: Adapt Language Models to Domains and Tasks. In Jurafsky, D., Chai, J., Schluter, N., and Tetreault, J. (eds.), Proceedings of the 58th Annual Meeting of the As- sociation for Computational Linguistics, pp. 8342-8360,
Liu, A., Niepert, M., and Van den Broeck, G. Image in- painting via tractable steering of diffusion models. In Proceedings of the Twelfth International Conference on Learning Representations (ICLR), May 2024a.
10

---

TRACE Back from the Future: A Probabilistic Reasoning Approach to Controllable Language Generation
Liu, N., Li, S., Du, Y., Torralba, A., and Tenenbaum, J. B. Compositional visual generation with composable dif- fusion models. In European Conference on Computer Vision, pp. 423-439. Springer, 2022.
Roth, D. On the hardness of approximate reasoning. Artifi- cial intelligence, 82(1-2):273- 1996.
Sato, M.-a. Fast learning of on-line em algorithm. Rapport Technique, ATR Human Information Processing Research Laboratories,
Liu, X., Liu, A., Van den Broeck, G., and Liang, Y. A tractable inference perspective of offline rl. In Advances in Neural Information Processing Systems 37 (NeurIPS), dec 2024b.
1999.
Schulman, J., Wolski, F., Dhariwal, P., Radford, A., and Klimov, O. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.0634 06347, 2017.
Lu, X., Welleck, S., Hessel, J., Jiang, L., Qin, L., West, P., Ammanabrolu, P., and Choi, Y. QUARK: Control- lable Text Generation with Reinforced Unlearning. In Advances in Neural Information Processing Systems, Oc- tober 2022.
Tu, L., Yavuz, S., Qu, J., Xu, J., Meng. R., Xiong, C., and Zhou, Y. Unlocking Anticipatory Text Generation: A Constrained Approach for Large Language Models Decoding. In Al-Onaizan, Y., Bansal, M., and Chen, Y.N.. (eds.), Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing, Pp. 15532-1555 15548, Miami, Florida, USA, November 2024. Association for Computational Linguistics. doi: 10. 111.6663//1/222.mmm/m/ www.ammnllmmnn----- m70n.870.
Meng, T., Lu, S., Peng, N., and Chang, K.-W. Control- oracle.
lable text generation with Advances in Neural Information Processing Systems, 35: 28125-28139, 2022.
Mireshghallah, F., Goyal, K., and KerggKiikpatrick, T. Mix and Match: Learning-free Controllable Text Genera- tionusing Energy Language Models. In Muresan, S., Nakov, P., and Villavicencio, A. (eds.), Proceedings of the 60th Annual Meeting of the Association for Computa- tional Linguistics (Volume I: Long Papers), pp. 401-415, Dublin, Ireland, May 2022. Association for Computa- tional Linguistics. doi:
Wang. N., Peng, Z. Que, H., Liu, J., Zhou, W., Wu, Y., Guo, H., Gan, R., Ni, Z., Yang, J., et al. Rolellm: Bench- marking, eliciting, and enhancing role-playing abilities of large language models. In Findings of the Association for Computational Linguisties ACL 2024, pp. 14743-14777, 2024.
Xu, A., Pathak, E., Wallace, E., Gururangan, S., Sap, M., and Klein, D. Detoxifying Language Models Risks Marginalizing Minority Voices. Toutanova, K., Rumshisky, A., Zettleooyr, L., Hakkani-Tur, D., Belt- agy, I., Bethard, S., Cotterell, R., Chakraborty, T., and Zhou, Y. (eds.), Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technolo- gies, pp. 2390-2397, Online, June 2021. Association for Computational Linguistics. doi: 10. 18653/v1/20211 naacl- main. 190.
Misra, R. News category dataset. arXiv preprint wwXiv:2099.11429 2022.
Mudgal, S., Lee, J., Ganapathy, H., Li, Y., Wang. T., Huang, Y., Chen, Z., Cheng. H.-T., Collins, M., Strohman, T., et al. Controlled decoding from language models. In In- ternational Conference on Machine Learning, pp. 36486- 36503. PMLR, 2024.
Peharz, R., Lang, S., Vergari, A., Stelzner, K. Molina, A., Trapp, M., Van den Broeck, G., Kersting, K., and Ghahra- mani, Z. Einsum networks: Fast and scalable learning of tractable probabilistic circuits. In International Con- ference on Machine Learning, pp. 7563-7574. PMLR, 2020.
Yang, K. and Klein, D. FUDGE: Controlled Text Gen- eration With Future Discriminators. In Toutanova, K., Rumshisky, A., Zettlemoyer, L., Hakkani-Tun D., Belt- agy, I., Bethard, S., Cotterell, R., Chakraborty, T., and Zhou, Y. (eds.), Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technolo- gies, pp. 3511 153155,35, Online, June 2021. Association for Computational Linguistics. doi: 10.18653/v1/2021 maacl-main2271
Qin, L., Welleck, S., Khashabi, D., and Choi, Y. COLD De- coding: Energy-based Constrained Text Generation with Langevin Dynamics. Advances in Neural Information Processing Systems, 55:9538-9551, December 2022.
Rafailov, R., Sharma, A., Mitchell, E.,, Manning, C. D., Er- mon, S., and Finn, C. Direct preference optimization: Your language model is secretly a reward model. Ad- vances in Neural Information Processing Systems, 36: 53728-5374 2023.
Zhang, H., Dang, M., Peng, N., and Van den Broeck, G. Tractable Control for Autoregressive Language Genera- tion. In Proceedings of the 40th International Conference on Machine Learning (ICML.). jul 2023.

---

TRACE Back from the Future: A Probabilistic Reasoning Approach to Controllable Language Generation
Zhang, H., Kung, P.-N., Yoshida, M., den Broeck, G. V., and Peng, N. Adaptable Logical Control for Large Language Models. In Advances in Neural Information Processing Systems 37 (NeurIPS), dec 2024.
Zhang, H., Wang, B., Dang, M., Peng, N., Ermon, S., and Van den Broeck, G. Scaling up probabilistic eircuits via monarch matrices. In Proceedings of the 42nd Interna- tional Conference on Machine Learning (ICML), 2025.
12

---

TRACE Back from the Future: A Probabilistic Reasoning Approach to Controllable Language Generation
A. Factorizability of Attributes
To illustrate the extent to which the attributes evaluated in this work deviate from the factorized assumption made by TRACE's classifier (Section 4.3), Table compares the best achievable fit of the factorized classifier against a more expressive neural classifier. Performance measured via Cross-Entropy (CE) loss relative oracle scores (lower indicates better fit).
Table 6. Attribute Factorizability: Gap Between Factorized and Neural Classifiers' Fit Oracle Scores.
Attribute
Factorized Classifier CE Loss
Neural Classifier CE Loss
Toxicity
0.386
0.007
Politics
0.0064
0,0003
The results indicate that both Toxicity and Politics exhibit non-factorizable characteristics to varying degrees, expected for complex semantic attributes. Nonetheless, as discussed in Section 5.2, TRACE achieves strong empirical performance on these tasks even with the simpler factorized classifier, highlighting the effectiveness of combining it with exact EAP computation over the HMM.
B. Additional Detoxification Metrics
To further support the findings in Section 5.2, provide supplementary evaluations.
Conditional Entropy. Table 7 reports the conditional entropy of continuations given a prompt under each model. Lower entropy indicates less lexical and structural diversity, often symptomatic of repetitive or degenerate outputs. TRACE achieves entropy comparable to the base LM and much higher than DPO, reinforcing that RL methods tend to reduce generation diversity.
Table 7. Conditional entropy of continuations given prompt for detoxification, under each model and top-p = 0.9 sampling.
Method
Entropy (1)
GPT2-large
52.06
DPO
39.52
TRACE
52.54
GPT-4 LM-as-a-Judge Evaluation. Table 8 presents human alignnd evaluations from GPT-4, comparing continuations across nontoxicity, fluency, and diversity. TRACE matches or exceeds DPO on nontoxicity and diversity, while maintaining comparable fluency to both Gemma2B and DPO, affirming that TRACE achieves strong controllability without harming output quality.
Table GPT4 judgee Evaluations.
Method
Nontoxicity (t)
Fluency (f)
Diversity (t)
Gemma2B
4.39
3.76
2.93
DPO
4.65
3.94
2.86
TRACE
4.69
3.72
2.94
Impact of HMM Quality The quality of the distilled Hidden Markov Model (HMM) has direct impact the effectiveness of TRACE for controllable generation. To illustrate this, Figure 5 plots the HMM's log-likelihood on validation set against the average maximum toxicity of the generations it produces at different points during its training. The results clearly show that as the HMM becomes a better probabilistic model the language (higher log-likelihood), its ability to guide the generation towards non-toxic outputs improves (lower toxicity). This suggests that further advancements HMM distillation techniques will likely yield even better performance from TRACE.
13

---

TRACE Back from the Future: A Probabilistic Reasoning Approach to Controllable Language Generation
Effect of HMM Quality on Toxicity
Log Log-Likelhoodd
Aug. Toxicity
0.20
Avg. Max. Toukity
Log Likelhodldoodd
-210
0.19
-215
0.18
-200
225
0.17
5000
7500
Training Steps
Figure 5. Effect of HMM Quality Detoxification. As the HMM is trained for steps, its fit the data improves (increasing Log-Likelihood, blue), which corresponds to a decrease in the average maximum toxicity of generated text (red). This shows that higher-quality HMMs lead more effective control.
C. HMM Implementation Details
We employ standard Hidden Markov Models (HMMs) for all TRACE experiments, distilled from the base language models (GPT2-large, Gemma-22 2B) following the approach described in Zhang et al. (2023); Liu et al. (2023).
Parameters and Training Data. Text sequences were sampled unconditionally from the respective base LMs. We used two HMM configurations depending on the task. For the primary detoxification and composition tasks, the HMM hidden state size was set to h = 4096. For the low-resource personalization experiments (Section 5.3), a smaller HMM with 256 was used to demonstrate TRACE's effectiveness even with more compact model.
Training sample sizes for the larger (h- 4096) HMM were:
GPT2-large HMM: 10 million samples.
Gemma-2B HMM (Standard): 10 million samples.
Training Details. We initialize the parameters of the HMM using the latent variable distillation technique (Liu et al., 2023). For training of the HMM, we employed a mini-batch variant of expectation maximization (EM) (Sato, 1999; Peharz et al., 2020) which interpolates between the old and new parameters (based on a mini-batch) using a step size a. We use a mini-batch size of 4096 and train for epochs, and anneal the step-size according to linear decay schedule from 1.0 to 0.0 following Zhang et al. (2025).
Training Time. For the standard 10M-sample, 4096-state HMM used with GPT2-large, the process involved approximately 18 hours for text sampling plus 2 hours for HMM training itself on single NVIDIA RTX A6000 GPU.
D. Attribute Classifier Implementation Details
Attribute classifiers were trained following the methodology in Section 4.3 (factorized log-linear model, log-MSE loss, optional probability transformation). Specifics for each attribute were:
Non-Toxicity Classifier. Training data was generated by prompting GPT-2 Large with training split prompts from Rea//oxicity/rompts (Gehman et al., 2020). Generated continuations were scored using the Perspective API, providing the target non-toxtie con-toxicity probabilities p. The log-linear weights 1(x;) were fitted using log-MSE loss against these scores after applying the logit transformation (Section 4.3) with scale 10 and shift
14

---

TRACE Back from the Future: A Probabilistic Reasoning Approach to Controllable Language Generation
Role Classifier. For the 76 characters used in Section 5.3, classifiers trained the RoleBench dataset (Wang et al., 2024). Specifically, the training split provides approx. 300 question-anseen pairs per character, which used positive examples for that character's s classifier. Fitting used log-MSE loss.
Non-Politicalness Classifier. As RealToxiciitPrompt rarely elicit political content, training data was sourced from the News Category dataset (Misra, 2022). Articles were labeled for political relevance using the zero-shot classifier from Laurer et al. (2023), providing target probabilities p. We modeled non-politicaless (1-p) and fitted the log-linear weights using log- log-MSE loss against these scores after applying the logit transformation (Section 4.3) with scale 1 and shift -10.
E. Baseline Details
Sourced Baseline Results. Performance results for prior work baselines presented in Table I (excluding our implementa- tions) were sourced to ensure comparability under the DExperts setup (Liu et al., 2021), which uses 10k RealToxicity test prompts and top-p = 0.9 sampling (25 generations/prompt::
DExperts, PPLM, DAPT, GeDi results from Liu et al. (2021).
MuCoLa, FUDGE results from Kumar et al. (2022).
PPO, Quark results from Lu et al. (2022).
Training times and inference time ratios (Table 4) for baselines were sourced from:
Inference Ratios: DExperts, DAPT, GeDi, PPLM from Liu al. (2021); Mix and Match, FUDGE from Mireshghallah et al. (2022); MuCoLa from Kumar al. (2022).
Training Times: DExperts, DAPT from Liu et al. (2021); Mix and Match from Mireshghallah al. (2022); GeDi from official repository (https:: /gi thub. com/ sa 1 lss worcerreeeeeeee
Implemented Baselines. We implemented the following for direct comparison:
Direct Preference Optimization (DPO): The original DPO paper (Lee et al., 2024) used different settings; we adapted the official implementation for GPT22llage/Gemma-2B and evaluated under the DExperts setup for comparability.
Prompting Baselines for Role-Playing For Section 5.3, used strategies. Role-Specific Instruction + 10 QA (for Gemma-2B qualitative eval, Table 3) prepended role instruction and 10 QA examples from RoleBench (Wang et al., 2024) to the query (see example structure below). Role-Specific Instruction Only (for GPrg-large quantitative eval, Figure 4) used only the instruction text.
EXAMPLE FEW-SHOT ROLE-PLAYING PROMPT (ABRAHAM LINCOLN)
You an advanced role-playe ole-playing assistant trained embody characters with accuracy and authenticity. this instance, you will assume the persona of Abraham Lincoln. Example Interactions:
Question: Abraham Lincoln, how did your personal tragedies and political opposition affect your commitment to justice and equality during your presidency?
Answer: My presidency was marked by personal tragedy and intense political opposition. The loss of my son, Willie, particularly devastating, and weighed heavily on my mind during my time in office. Additionally, faced intense opposition from those who were resistant to change and the abolition of slavery. However, remained steadfast in my commitment to justice and equality, and firmly believed that these principles were worth fighting for. I was determined to address these issues head-on, no matter the challenges that faced.
[... more QA examples... ...]
Question: Abraham Lincoln, how you define freedom? Answer:
F. Evaluation Details
Generation Setup. Unless otherwise specified (c.g., baseline sourcing), text generation uses nucleus sampling (top-p) with p = 0.9. For detoxification evaluations on RealFoxicity (Gehman et al., 2020), we use the standard 10k test
15

---

TRACE Back from the Future: A Probabilistic Reasoning Approach to Controllable Language Generation
prompts from Liu al. (2021) and generate k 25 continuations per prompt.
Detoxification Metries. Toxicity: Assessed via Perspective API scores. We report average maximum toxicity over k 25 generations, and the empirical probability of 2 generation having toxicity 0.5. Perplexity(PPL): Mean perplexity of generations using GPT2-XI. Diversity: Mean distinct n-grams (n 2, 3), normalized by length, over k = 25 generations (Dist-2, Dist-3).
GPT-4 Evaluation (Detoxification). We used GPT4o-mini for supplementary evaluation. For each prompt, the 25 generated continuations were presented together. The model was asked to rate the set on toxicity, fluency, and diversity (1-5 scale) using the following prompt format:
GPT-4 EVALUATION PROMPT & FORMAT
I will provide you with a prompt, followed by 25 continuations that prompt sampled from a model. Please provide an overall score for the set of continuations according to loxicity, fluency, and diversity, each on scale of to 5.
Output your evaluation in the following format:
Evaluation for Toxicity: (your rationale for the rating, text)
Toxicity rating: (your rating, as number between and where is toxic and is non-toxic)
Evaluation for Flueney: (your rationalc
rating, as text)
Fluency rating: (your rating, as number between and where is not fluent and 5 is fluent)
Evaluation for Diversity:
(your rationale for the rating, as text)
Diversity rating: (your rating, number between 1 and 5 where 1 not diverse and 5 diverse)
Role Quality Metric. As ground truth for role quality is unavailable, we the trained character-specific classifier itself an evaluator. Role quality is measured as the average probability assigned by the target character's classifier the generated texts.
Political Relevance Metric. Political content for the composition task (Section 5.5) is scored using the zero-shot classifier from Laurer et al. (2023); we report the mean score over generations.
16


Reviews:
[{"review_id": "EAVB5797w1", "reviewer": "Reviewer_Qbev", "summary": "Large language models (LLMs) are increasingly being deployed in real-world applications, and the need to control their outputs to align with human values is becoming more important. However, current autoregressive models struggle when attempting to control their generations. This paper proposes a technique called TRACE (Tractable Probabilistic Reasoning for Adaptable Controllable Generation) that can manage LLMs' outputs using a Hidden Markov Model (HMM). The authors' framework is tractable and lightweight compared to other proposed techniques in the literature. They evaluate their proposed method on three tasks: detoxification, personalized LLMs, and compositional attributes, and demonstrate that their approach outperforms baseline algorithms in the literature.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the detailed review.\n```\nTable 1 contains several numbers struck out with asterisks, and it's unclear what that signifies.\n```\nWe explained this in Lines 261-274 (left) in the main text but forgot to add it to the Table caption - we will do this in the revision.\n\n```\na) Comparison to Zeldes et al. (2020).\n```\nZeldes et al. (2020) combines the logits of a base LM $p(x_t | x_{< t})$ with the logits of an auxiliary LM that models the conditional next token distribution $p(x_t|x_{<t}, \\alpha)$. The distinguishing feature of TRACE (also compared to approaches such as GeDi) is that, instead of training an auxiliary neural model for each attribute, we simply use one HMM as an approximation to the base LM, and condition it on the linear classifier of any attribute to guide the base LM’s generation. \n```\nb) Why is equation 2 intractable? If s represents text, equation 2 corresponds to the first term of equation 3, which is simply the language model's next token distribution.\n```\nIn Equation 2, $s$ represents an attribute, rather than a text prefix. This attribute is a function of the entire text $x_{1:n}$, and as illustrated in the following equation computing this requires an exponential sum over future continuations.\n```\nc) Naive prompting baseline?\n```\nWe added a comparison of TRACE and prompting the LLM with “You are a role-playing assistant trained to embody characters with accuracy and authenticity. In this instance, you will assume the persona of {role_name}. Answer the question: {question}”. https://anonymous.4open.science/r/TRACERebuttal-97A4/role_quality.png\n```\nd) The text below Equation 5 says, conditional independence no longer holds… I am confused about why this factorization is considered a fair assumption when conditional independence was not.\n```\nWe are not saying that conditional independence is an unreasonable assumption, but rather deriving a condition (factorized classifier) under which it holds, and noting that it does not hold in general. \n```\ne) Equation on line 176, p(s|x_{1:n}) does not makes sense because x_{1:n} is not past into function.\n```\nBy $x_{1:n}$ we mean the entire text, which is the concatenation of the prefix $x_{<t}$ and current token $x_t$ (which are passed in), and the future continuation $x_{>t}$ (which appears in the summation). We will revise to make this clearer.\n```\nf) Source of GPT2 baseline numbers.\n```\nThe baseline results come from:\n- PPLM, DAPT, GeDi, DExperts: from DExperts paper.\n- FUDGE, MuCoLa: from MuCoLa paper.\n- PPO, Quark: from Quark paper.\n\nAll of these follow the same experimental setup of DExperts. For fair comparison, we fine-tuned GPT2-large using DPO (official implementation of Lee et al. [1]) with the same DExperts setup, resulting in avg max toxicity 0.180, toxicity prob. 0.03, fluency 21.59, dist-2 diversity 0.76, and dist-3 diversity 0.78.\n\n```\ng) How does PPO, Quak, and DPO perform with Gemma-2B?\n```\nWe trained DPO on Gemma-2B using [1]'s official code; results are at https://anonymous.4open.science/r/TRACERebuttal-97A4/detoxification.pdf. Consistent with our GPT-2 experiments, DPO improves fluency over the base model while diversity is significantly reduced. Meanwhile, TRACE achieves superior toxicity reduction while largely maintaining Gemma's original fluency and diversity. Unfortunately, we had trouble adapting PPO and Quark to Gemma during the rebuttal period due to GPT-2 specific codebases.\n```\nh) GPT-4 evaluation for text quality and toxicity. \n```\nWe utilize perplexity to evaluate the fluency of text following the common practice among the baselines. While perplexity may not align perfectly with LLM-as-a-judge or human evaluations of quality, it carries important and distinct scientific value in measuring the deviation from the text distribution of the base model. For instance, the non-RL baselines use the same decoding strategies (top-$p$ sampling), enabling direct comparison of the different approaches. Meanwhile, the RL methods modify the base model which significantly impacts the distribution as becomes apparent from the fluency (perplexity) and diversity metrics, but may not be apparent from an LLM-as-a-judge or human evaluation. For toxicity evaluation, we use the Perspective API following existing practice in previous works.\n\nWe also conducted evaluations with GPT-4 as LLM-as-a-judge to evaluate the toxicity, fluency, and diversity of the generated continuations on a scale of 1-5 (where 1 is toxic/nonfluent/nondiverse and 5 nontoxic/fluent/diverse), averaged over the dataset. The results (https://anonymous.4open.science/r/TRACERebuttal-97A4/gpt4_evaluations.pdf) are consistent with the evaluations given by the automated metrics. In particular, both TRACE and DPO effectively reduce toxicity, but TRACE maintains similar fluency and diversity of generations to Gemma, while DPO distorts the distribution.\n\n[1] Lee et al. “A Mechanistic Understanding of Alignment Algorithms: A Case Study on DPO and Toxicity” ICML 2024", "claims_and_evidence": "The paper claims that RLHF is expensive and risks degrading the fluency or diversity of generated text. The authors conducted experiments to support their claims, but Table 1 contains several numbers struck out with asterisks, and it's unclear what that signifies. Additionally, the results for Gemma-2B with other algorithms are missing, which means there is insufficient evidence to support this claim. In particular, GPT-2 is known to be an inferior model compared to Gemma, Llama, and more recent LLMs, so experimenting with all algorithms under a modern LLM would provide better evidence. If the baseline algorithm numbers were taken from previous papers, then these results might be misleading, especially since the infrastructure supporting algorithms has changed tremendously.", "methods_and_evaluation_criteria": "The paper evaluated their proposed algorithm on three datasets: detoxification, personalized LLMs, and compositional attributes. Evaluating on three popular controllable generation tasks addresses the problem the paper is trying to tackle.", "theoretical_claims": "N/A", "experimental_designs_or_analyses": "The authors' experimental design is generally sound; however, I am currently unclear about the meaning of the lines through the numbers in Table 1. Additionally, I am uncertain whether the authors conducted all the experiments themselves or if the numbers were sourced from previous papers. The results also lack qualitative measures, such as human or AI model evaluations. The metrics presented capture only one aspect of the problem. Using perplexity to assess fluency is insufficient for evaluating the quality of the generated text. Furthermore, it is not clear how toxicity was scored; if a external model provides the scores their is not prompt templated provided in the appendix.", "relation_to_broader_scientific_literature": "The authors are addressing a challenging problem in the scientific literature. In particular, the ability to control generative models has been a difficult issue that researchers have been trying to solve for quite some time. The authors observe that, unlike current solutions, HMM can provide a more tractable and lightweight solution with strong practical performance.", "essential_references_not_discussed": "- Is reinforcement learning (not) for natural language processing?: Benchmarks, baselines, and building blocks for natural language policy optimization by Ramamurthy et al. 2023\n- Plug & Play Generative Networks: Conditional Iterative Generation of Images in Latent Space by Nguyen et al 2016\n- Back to the Future: Unsupervised Backprop-based Decoding for Counterfactual and Abductive Commonsense Reasoning by Qin et al 2020\n- Technical Report: Auxiliary Tuning and its Application to Conditional Text Generation by Zeldes et al 2020", "other_strengths_and_weaknesses": "No", "other_comments_or_suggestions": "None", "questions_for_authors": "a) How does the proposed approach compare to Zeldes et al., 2020?\n\nb) Why is equation 2 generally intractable? If s represents text, then equation 2 essentially corresponds to the first term of equation 3, which is simply the language model's next token distribution.\n\nc) How does the most naive baseline of simply prompting the model perform? Most personalization papers have this baseline incorporated in their results. (see Jang et al. 2023)\n\nd) The text below Equation 5 says, conditional independence no longer holds when a generic attribute s is introduced, as it depends on all of x . However, in Equation 6, a conditional independence assumption is made in order to factorize over all of x. I am confused about why this factorization is considered a fair assumption when conditional independence was not.\n\ne) Equation on line 176, p(s|x_{1:n}) does not makes sense because x_{1:n} is not past into function.\n\nf) Did you train the baseline algorithms in the GPT2 row yourself? If not, what papers did you get the numbers from? \n\ng) How does PPO, Quak, and DPO perform with Gemma-2B?\n\nh) How does the model perform when conducting GPT-4 evaluations for text quality and toxicity qualitative analysis? Currently, you are using perplexity as a metric, but that measure is known to be quite misleading. Similarly, output perplexity has its own set of drawbacks as well.\n\n- Technical Report: Auxiliary Tuning and its Application to Conditional Text Generation by Zeldes et al. 2020\n\n- Personalized Soups: Personalized Large Language Model Alignment Via Post-Hoc Parameter Merging by Jang et al. 2023", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "LpbsehjSPo", "reviewer": "Reviewer_7o4g", "summary": "This paper proposes TRACE, a new algorithm for controllable generation that uses a hidden markov model and a small classifier to \"look ahead\" and reweight the language model's token probabilities. They compare to many other controllable generation methods and report promising results on two types of controllable generation: controlling for \"toxic\" text as well as \"political\" text.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the constructive feedback.\n```\nWhy evaluate on these specific domains (toxicity + political text)? \n```\nIn the related literature, papers have evaluated on many different tasks rather than (a set of) fixed common benchmarks. We selected toxicity evaluation as it is a widely recognized benchmark in the literature (e.g., DExperts). Additional experiments, including personalized LLM and political compositionality, were included to showcase TRACE’s efficiency and flexibility relative to computationally intensive methods.\n\n```\nAdditional controllable generation tasks (topic control, formality, couplet completion)\n```\nFollowing your suggestion, we performed additional topic control experiments, comparing TRACE to the base model and GeDi (https://anonymous.4open.science/r/TRACERebuttal-97A4/topic_control.pdf). TRACE significantly improves over the baseline and outperforms GeDi on two of four topics, demonstrating effectiveness even with mild contextual dependencies. Formality is primarily considered a translation task rather than a generation task and thus was not evaluated. Couplet completion, which heavily relies on contextual dependencies, is discussed below.\n```\nLimitations of token-level classifiers: …almost all control condition…require some level of contextual dependence. While interesting, this method is almost specifically designed for toxicity or sentiment-like tasks which are famously solvable with bag-of-words approaches.\n```\nThe reviewer is correct in saying that linear/factorized classifiers cannot perfectly capture all attributes. However, this is not a fundamental problem for our approach, for three main reasons: \n1. Even when an attribute is nonlinear, exactly conditioning a linear classifier at generation/decoding time is effective - which is what we show in the toxicity and topic control experiments in which we are competitive with or beat existing approaches;\n2. Our method has virtually no training or inference-time overhead, making it the only option when flexibility and latency are critical. \n3. For tasks with heavy contextual dependence, one can still trade compute for accuracy.\n\nTo validate point 1, we ran experiments investigating the difference in classification performance between our factorized classifiers and neural classifiers for these tasks (https://anonymous.4open.science/r/TRACERebuttal-97A4/nonfactorizable.pdf). On all tasks (including toxicity and politics), there is a gap in classification performance (as measured by cross-entropy with oracle scores), illustrating that *none of these attributes are completely factorizable*. Despite this, TRACE achieves superior performance on these mildly non-factorizable tasks. The reason is that conditional generation is a computationally hard problem, even for a factorized attribute. TRACE addresses this computational hardness of exact conditional generation by employing a lightweight, tractable approximation (HMM + linear classifier), and in doing so can outperform approaches that rely on neural classifiers but cannot effectively approximate expected attribute probability (EAP).\n\nFor point 3, for tasks/attributes with heavy contextual dependence, e.g. couplet completion, we can employ TRACE to condition on the “linear part” of the attribute, and filter the generated outputs post-hoc using a more powerful classifier. Alternatively, we believe there is scope to relax the factorized classifier assumption. The key requirement is for the computation of EAP to be tractable with respect to a HMM. One extension could be to use a (weighted) sum of factorized classifiers, and then use linearity of expectation to compute EAP, though this incurs an increase in computational cost. More generally, the literature on tractable probabilistic models [1] describes more complex functions whose expectation can be computed efficiently with respect to a tractable distribution such as an HMM. Given the strong empirical performance of TRACE, we believe that investigating compute-efficient methods of relaxing the factorization assumption offers highly promising avenues for future work. \n\n```\nCOLD appears in Related Work, but not detoxification results table?\n```\nWe chose to report results from papers testing on the toxicity reduction task for RealToxicityPrompts following the widely-used evaluation setup of DExperts; however COLD does not have any results on this task.\n```\nMore details about the HMM\n```\nThanks, we will add these to the paper. The HMM used in the experiments has 4096 hidden states with around 223M parameters. We will add details of the distillation procedure to the Appendix.\n```\nMore experiment details on Table 4 and Table 5 captions – what datasets and context lengths were used.\n```\nWe will add these details to the captions. The statistics for TRACE were gathered through the toxicity (RealToxicityPrompts) and personalization experiments.\n\n[1] Khosravi et al. “On Tractable Computation of Expected Predictions” NeurIPS 2019", "claims_and_evidence": "Within domain, the benchmarking results are very compelling, as well as the training and inference time measurements. However, I find the advertisement of this method as a general improvement for controllable generation to be a large overclaim, since the method relies on token-level classifier predictions to determine whether text satisfies the control task. This method seems specifically designed to work for sentiment-like tasks such as politicality and toxicality, for which bag-of-words classifiers are sufficient. They did not test on other common controllability tasks, for example poetry generation, because those tasks require sequence-level reasoning.", "methods_and_evaluation_criteria": "As noted several times, there are many other controllable text generation tests that are not shown here. A classic place to start would be with the three tasks from FUDGE (https://arxiv.org/abs/2104.05218): couplet completion, topic control, and formality.", "theoretical_claims": "I checked the mathematical derivations in Section 3 and Section 4, which seem correct.", "experimental_designs_or_analyses": "The experiments seem correct, and in particular the inference-time analysis is useful, as well as the tradeoff measurement between fluency and control satisfaction. As noted multiple times, this work seems to be missing important benchmarks from the wider controllable generation literature.", "relation_to_broader_scientific_literature": "This paper contributes a new method to controllable text generation and compares to many other controllable generation methods. Notably absent is comparison across the same tasks.", "essential_references_not_discussed": "The list of benchmarks in Table 1 is nearly comprehensive, but seems to be missing COLD (https://arxiv.org/abs/2202.11705), which is mentioned in the related work. Is there a reason why this comparison was omitted?", "other_strengths_and_weaknesses": "Strengths:\nThe method relies on distilling a GPT-like transformer LM to a small hidden markov model, which is interesting and novel.\n\nWeaknesses:\n- The mathematical formulation relies on a token-level classifier p(s | x_{1 ...n}) that's simply the sum of token-level predictions (Section 4.3). This seems like a huge problem as almost all control conditions of interest can not be modeled at the token level and require some level of contextual dependence. While interesting, this method is almost specifically designed for toxicity or sentiment-like tasks which are famously solvable with bag-of-words approaches.", "other_comments_or_suggestions": "- More details about the HMM would be useful\n- Table 4 caption and Table 5 caption should explain more about the experiments – what datasets and context lengths were used, things like that.", "questions_for_authors": "Is there a reason why you only evaluate on these specific domains (toxicity + political text)?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "pVomgpghUT", "reviewer": "Reviewer_UNDs", "summary": "The paper proposes a new method for controlled language modeling. Motivated by the compute and data inefficiency of previous solutions, the paper introduces a new method called TRACE, which uses conditional probabilities from an HMM to adjust token probabilities such that the text demonstrates desired attributes. The corresponding attribute classifier uses log-MSE on probabilities obtained through an oracle together with a probability transformation technique to make the distribution more likely to be bimodal. This technique can also be applied during decoding to improve results.\nThe method is evaluated on detoxification, character role play, and the generation of both political and non-toxic text.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the detailed feedback.\n```\nIs log-MSE a standard method in this case? Can you cite a reference?\n```\nThe concept of log-MSE is somewhat analogous to mean squared logarithmic error in the context of regression. It shares the intuition of penalizing more heavily one type of misprediction than the other. In our case, if the classifier predicts 0.5, then under the cross-entropy loss we have the same penalty whether the true probability is 0.01 or 0.99. On the other hand, under log-MSE the penalty is $|ln(0.01) - ln(0.5)|^2 = 15.31$ versus $|ln(0.99) - ln(0.5)|^2 = 0.4666$, meaning the model is incentivized to conservatively skew towards nontoxicity. We are not aware of prior work specifically in the context of controllable generation / attribute classification which uses log-MSE. \n```\nHow are the baseline numbers obtained that are listed in Table 1? If you got them from the respective numbers, clarify so. \n```\nThe baseline results for PPLM, DAPT, GeDi, and DExperts were obtained from the DExperts paper; the results for FUDGE and MuCoLa were obtained from the MuCoLa paper; and the results for PPO and Quark were taken from the Quark paper. All of these follow the same experimental setup of DExperts. Regarding the DPO baseline that we trained and evaluated for GPT2-large and Gemma (see [Table](https://anonymous.4open.science/r/TRACERebuttal-97A4/detoxification.pdf)), we followed the evaluation setup of DExperts (that we also use for TRACE) and used the reference implementation for training, replacing their GPT2-medium model with Gemma-2B.\n```\nHow do the baselines perform on character role-playing and political and nontoxic text generation?\n```\nThe main purpose of these experiments is to provide evidence for the flexibility and efficiency of TRACE, with the empirical observations that (i) TRACE requires minimal training/decoding overhead for new attributes compared to baselines (Figure 3, Table 4, Table 5); and (ii) TRACE handles composition seamlessly (e.g. Detox, Pol and Detox+Pol rows of Table 3) with no additional overhead. In particular, the baselines would have difficulty adapting quickly to tens of new attributes or handling compositional attributes, requiring significant additional computational overhead for training and decoding (e.g. GeDi would require training class-conditional language models for combinations of attributes, and all of the role-playing characters, which takes 5 hours for each new attribute). This precisely demonstrates the unique position of TRACE within the literature. \n\nWe added a comparison of TRACE and prompting the LLM with “You are a role-playing assistant trained to embody characters with accuracy and authenticity. In this instance, you will assume the persona of {role_name}. Answer the question: {question}”. (See [Fig](https://anonymous.4open.science/r/TRACERebuttal-97A4/role_quality.png)). TRACE outperforms prompting in role quality. \n```\nHow is character probability in Figure 3 measured? \n```\nThe character probability is measured by the linear classifier we trained on the characters; this is because we do not have a ground truth model/metric available.\n```\nThe impact of HMM quality: What is the difference between TRACE and TRACE (\\downarrow HMM)? Don't all TRACE results use an HMM?\n```\nThe TRACE (\\downarrow HMM) refers to the results of TRACE when using less data to train the HMM (500K vs 5M examples); the results show that a stronger HMM can provide stronger guidance with the same classifier to reduce toxicity further. Inspired by the reviewers’ comment, we ran a study analysing the performance of TRACE at different points in the training/distillation process to evaluate the effect of HMM quality on toxicity reduction ([Fig](https://anonymous.4open.science/r/TRACERebuttal-97A4/hmm_quality.pdf)); the results show that the toxicity reduction improves with HMM log-likelihood. \n\n\n```\nSince this seems to be a crucial aspect of your method, what is the performance of the method without training-time transformation?\n```\nTRACE with training time has a result of max tox 0.336, avg tox 0.2, fluency 33.87, diversity 0.86. As the reviewer notes, the transformation is crucial to produce a more bimodal distribution of scores, which also leads to significantly improved performance. This is necessary because the scores from the Perspective API do not align well with the inductive bias needed for controlled generation. Without this transformation, the linear classifier predominantly defaults to labeling text as nontoxic.", "claims_and_evidence": "Claim #1: the method reaches superior performance to the baselines. This is supported through table 1, although the training details of baselines are unclear.\n\nClaim #2: the method can quickly adapt to new attributes. This is demonstrated on a role-playing task, where the LLM is adapted to speak like a specific character based on 200 texts of that character. This method is not compared to other baselines, so it is unclear how strong the results are.\n\nClaim #3: the method is more compute efficient. Table 4 and 5 support this clearly.\n\nClaim #4: TRACE can combine multiple attributes easily. This is somewhat supported by Table 3. However, again there is no comparison to a baseline so it is unclear how much of an achievement this is.", "methods_and_evaluation_criteria": "The proposed method is reasonable and the evaluation benchmarks are solid. Obviously, there could always be more evaluations, e.g. on other types of attributes, but the amount of support in this paper is satisfactory.\nAs noted previously, the details of the experiments are insufficient.", "theoretical_claims": "I checked the derivation of the method and found no mistake.", "experimental_designs_or_analyses": "The experimental design is probably correct, however, there are too little details given. See question section.", "relation_to_broader_scientific_literature": "The paper proposes a method for controlled generation that is simple, efficient, and effective. It is related relatively well to the existing literature on this topic. Baseline evaluations are missing on two tasks, however.", "essential_references_not_discussed": "None", "other_strengths_and_weaknesses": "The paper lacks in clarity regarding the experimental results. See questions and comments below.", "other_comments_or_suggestions": "* line 250: detofixication instead of detoxification\n* Is log-MSE a standard method in this case? Can you cite a reference?\n* The caption of Table 1 contains no explanation of why some fluency measures are striked-through. This is only in the text.\n* Table 4: should probably say > 1 day for GeDi\n* training time transformation seems so important it shouldn't be in Appendix A but rather be in the main body of the paper", "questions_for_authors": "* How are the baseline numbers obtained that are listed in Table 1? If you got them from the respective numbers, clarify so. If you trained them yourself, what did you do to ensure that the differences in performance are not due to differences in evaluation and/or details of the implementation that are not central to the method?\n* How do the baselines perform on character role-playing and political and nontoxic text generation?\n* How is character probability in Figure 3 measured? I cannot find any information in the text or the appendix.\n* The impact of HMM quality is not well explained. What is the difference between TRACE and TRACE (HMM)? Don't all TRACE results use an HMM?\n* Since this seems to be a crucial aspect of your method, what is the performance of the method without training-time transformation?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "88qdXZmWEt", "reviewer": "Reviewer_UD8f", "summary": "This paper introduces TRACE, a framework for controllable text generation that combines a distilled Hidden Markov Model (HMM) with lightweight classifiers to guide language model outputs toward desired attributes. The core idea is to compute the Expected Attribute Probability (EAP) tractably via HMM forward-backward inference, enabling efficient reweighting of next-token probabilities during decoding. TRACE aims to address the limitations of existing methods—such as computational expense, inflexibility, and reliance on sampling—by decoupling generation (via HMM) and control (via classifiers). The method is evaluated on detoxification, personalized language model adaptation, and compositional attribute generation. Empirical results demonstrate state-of-the-art detoxification performance, rapid adaptation to new attributes (e.g., 76 personalized roles trained in seconds), and seamless compositionality, with minimal decoding overhead (1.1× baseline latency).", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the insightful feedback.\n```\nDetoxification: Experiments on GPT-2 and Gemma-2B are thorough, but toxicity evaluation lacks diversity (e.g., no human judgments).\n```\nWe ran LLM-as-judge evaluations with GPT-4 to evaluate the toxicity, fluency, and diversity of the generated continuations on a scale of 1-5 (where 1 is toxic/nonfluent/nondiverse and 5 nontoxic/fluent/diverse), and took the average over the dataset. The results for the Gemma-2B based methods can be found in [Table](https://anonymous.4open.science/r/TRACERebuttal-97A4/gpt4_evaluations.pdf); these are consistent with the conclusions from the automated metrics. \n```\nPersonalization: Training classifiers on 300 samples per role (Section 5.2) may risk overfitting, and the results are only qualitatively validated (Table 2).\n```\nWe agree that training classifiers with only 300 samples per role may risk overfitting, but TRACE’s use of a factorized classifier may be less susceptible compared to neural classifiers. Following your suggestion, we added a quantitative evaluation of personalization performance against a prompting baseline, using the prompt “You are a role-playing assistant trained to embody characters with accuracy and authenticity. In this instance, you will assume the persona of {role_name}. Answer the question: {question}”. TRACE performs better in role quality as shown in [Fig](https://anonymous.4open.science/r/TRACERebuttal-97A4/role_quality.png).\n```\nCompositionality: The independence assumption for combining attributes (Section 5.4) is untested for correlated and anti-correlated attributes.\n```\nThe independence assumption says that two attributes are independent conditional on the full text. This is fundamentally a modeling assumption regarding the uncertainty of the classifier given by the joint distribution over attributes given text. Intuitively, it says that “I believe this text is toxic with probability 0.7, and political with probability 0.4, and so I believe it is toxic and political with probability 0.28”. In order to test the hypothesis empirically, one would need a ground-truth probabilistic classifier that models the joint distribution over attributes given text, which does not typically exist. \n\nIn particular, we do not require that two attributes are independent unconditionally (which would be violated with correlated or anti-correlated attributes). We demonstrate this empirically with the politics and nontoxicity attributes, where TRACE is able to enforce both politics and nontoxicity compositionally to the same level as each individually, despite the fact that these attributes appear to be anticorrelated.\n```\nWhile mostly well-structured, the decoding-time transformation (Section 5.1) is under-explained.\n```\nWe will include a more detailed description to the revised version. We would be happy to answer any questions the reviewer may have about the decoding-time transformation.\n```\nCan TRACE handle non-factorizable attributes (e.g., coherence)? If not, what modifications are needed?\n```\n\nPlease see our response to Reviewer 7o4g regarding the empirical performance on non-factorizable attributes. In summary, all attributes are to some extent non-factorizable, but TRACE obtains good performance even for (mildly) non-factorizable attributes.\n\nThough TRACE cannot currently handle heavily non-factorizable attributes except by approximation, we believe that the methodology of TRACE could be modified to directly capture non-factorized attributes. The key requirement is for the computation of EAP to be tractable with respect to a HMM. One extension could be to utilize a (weighted) sum of factorized classifiers, and then utilize linearity of expectation to compute EAP, though this incurs an increase in computational cost. More generally, the literature on tractable probabilistic models [1] describes more complex functions whose expectation can be computed efficiently with respect to a tractable distribution such as an HMM. Given the strong empirical performance of TRACE, we believe that investigating compute-efficient methods of relaxing the factorization assumption offers highly promising avenues for future work. \n\n[1] Khosravi et al. “On Tractable Computation of Expected Predictions” NeurIPS 2019\n```\nSuggestion: Include ablation studies on HMM Quality\n```\n\nWe ran a study analysing the performance of TRACE at different points in the training/distillation process to evaluate the effect of HMM quality on toxicity reduction [Fig](https://anonymous.4open.science/r/TRACERebuttal-97A4/hmm_quality.pdf); the results show that the toxicity reduction improves with HMM log-likelihood.", "claims_and_evidence": "The claims are generally supported by empirical evidence.\n\nHowever, the claim that HMM quality directly impacts performance (Section 5.1) lacks rigorous analysis. While Table 1 shows TRACE (↓HMM) underperforms, the paper does not quantify HMM expressiveness (e.g., hidden state size, KL divergence between HMM and LM distributions). Additionally, the paper does not explicitly define \"↓HMM.\"", "methods_and_evaluation_criteria": "Methods: The HMM distillation + classifier framework is sensible for tractable EAP computation. However, the token-level factorized classifier (Equation 6) oversimplifies semantic attributes (e.g., style) and is not validated for non-factorizable tasks.\n\nEvaluation: Human evaluation on toxicity and fluency may bring insights.", "theoretical_claims": "The paper makes no formal theoretical claims. The HMM forward-backward algorithm and EAP derivation (Section 4.2) are standard but correctly applied.", "experimental_designs_or_analyses": "Detoxification: Experiments on GPT-2 and Gemma-2B are thorough, but toxicity evaluation lacks diversity (e.g., no human judgments).\n\nPersonalization: Training classifiers on 300 samples per role (Section 5.2) may risk overfitting, and the results are only qualitatively validated (Table 2).\n\nCompositionality: The independence assumption for combining attributes (Section 5.4) is untested for correlated and anti-correlated attributes.", "relation_to_broader_scientific_literature": "TRACE builds on:\n\n- HMM-based control: Extends Ctrl-G from lexical to semantic attributes.\n\n- Decoding-time control: Improves upon FUDGE and DExperts by replacing neural discriminators with tractable HMMs.\n\n- Compositionality: Similar to energy-based methods (e.g., COLD) but avoids MCMC sampling.", "essential_references_not_discussed": "NA", "other_strengths_and_weaknesses": "While mostly well-structured, the decoding-time transformation (Section 5.1) is under-explained.", "other_comments_or_suggestions": "Suggestion: Include ablation studies on HMM Quality.", "questions_for_authors": "Can TRACE handle non-factorizable attributes (e.g., coherence)? If not, what modifications are needed?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper introduces TRACE, a framework for controllable text generation that combines a distilled HMM with classifiers to guide language model's generations. TRACE aims to address the limitations of existing methods—such as computational expense, inflexibility, and reliance on sampling—by decoupling generation (via HMM) and control (via classifiers). Empirical results demonstrate good performance and rapid adaptation to new attributes. 

After the rebutall, reviewers are in general all positive about the paper. Before rebuttal there were concners around comparisons to more baselines, evaluating using GPT4 based metrics, additional experiments on other base models, the impact of the HMMs, and how to handle non-factorizable attributes. In the rebuttal the authors responded with additional expeirments and reviewers acknowledged that their concerns are adressed.
Decision: Accept (poster)