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

• URL: http://arxiv.org/abs/2504.18535v2
• Date: Mon, 29 Sep 2025 03:34:04 GMT
• Title: TRACE Back from the Future: A Probabilistic Reasoning Approach to Controllable Language Generation
• Authors: Gwen Yidou Weng, Benjie Wang, Guy Van den Broeck,
• Abstract summary: Large language models need to align with human values or desired attributes.<n>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.<n>We introduce TRACE, a novel framework that efficiently computes EAP and adapts to new attributes through tractable probabilistic reasoning and lightweight control.
• Score: 31.762143603634836
• License: http://creativecommons.org/licenses/by/4.0/
• 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. Our code is available at: https://github.com/yidouweng/trace.

Related papers

• Learning Tractable Distributions Of Language Model Continuations [40.95717348276321]
  We find that tractable surrogates such as hidden Markov models (HMMs) approximate the distribution over continuations and adjust the model's next-token logits at decoding time.<n>We propose Learning to Look Ahead (LTLA), a hybrid approach that pairs the same base language model for rich prefix encoding with a fixed tractable surrogate model that computes exact continuation probabilities.
  arXiv  Detail & Related papers  (2025-11-20T05:17:19Z)
• Continuous Autoregressive Language Models [56.49239051750678]
  We introduce Continuous Autoregressive Language Models (CALM)<n>CALM uses a high-fidelity autoencoder to compress a chunk of K tokens into a single continuous vector.<n>We develop a comprehensive likelihood-free framework that enables robust training, evaluation, and controllable sampling.
  arXiv  Detail & Related papers  (2025-10-31T17:58:11Z)
• Utilizing Large Language Models for Machine Learning Explainability [37.31918138232927]
  This study explores the explainability capabilities of large language models (LLMs), when employed to autonomously generate machine learning (ML) solutions.<n>Three state-of-the-art LLMs are prompted to design training pipelines for four common classifiers: Random Forest, XGBoost, Multilayer Perceptron, and Long Short-Term Memory networks.<n>The generated models are evaluated in terms of predictive performance (recall, precision, and F1-score) and explainability using SHAP (SHapley Additive exPlanations)
  arXiv  Detail & Related papers  (2025-10-08T11:46:23Z)
• ESLM: Risk-Averse Selective Language Modeling for Efficient Pretraining [53.893792844055106]
  Large language model pretraining is compute-intensive, yet many tokens contribute marginally to learning, resulting in inefficiency.<n>We introduce Selective Efficient Language Modeling, a risk-aware algorithm that improves training efficiency and distributional robustness by performing online token-level batch selection.<n> Experiments on GPT-2 pretraining show that ESLM significantly reduces training FLOPs while maintaining or improving both perplexity and downstream performance compared to baselines.
  arXiv  Detail & Related papers  (2025-05-26T12:23:26Z)
• Semantic Probabilistic Control of Language Models [27.866585120174484]
  Semantic control involves steering LM generations towards satisfying subtle non-lexical constraints.<n>We leverage a verifier's information gradient to efficiently reason over all generations that satisfy the target attribute.<n>We evaluate the effectiveness of our approach in controlling the toxicity, sentiment, and topic-adherence of LMs yielding generations satisfying the constraint with high probability.
  arXiv  Detail & Related papers  (2025-05-04T01:21:28Z)
• Reviving Any-Subset Autoregressive Models with Principled Parallel Sampling and Speculative Decoding [55.2480439325792]
  In arbitrary-order language models, it is an open question how to sample tokens in parallel from the correct joint distribution. We find that a different class of models, any-subset autoregressive models (AS-ARMs), holds the solution. We show that AS-ARMs achieve state-of-the-art performance among sub-200M parameter models on infilling benchmark tasks, and nearly match the performance of models 50X larger on code generation.
  arXiv  Detail & Related papers  (2025-04-29T06:33:13Z)
• Language Models Can Predict Their Own Behavior [28.80639362933004]
  We show that internal representation of input tokens alone can often precisely predict, not just the next token, but eventual behavior over the entire output sequence.<n>We leverage this capacity and learn probes on internal states to create early warning (and exit) systems.<n>Specifically, if the probes can confidently estimate the way the LM is going to behave, then the system will avoid generating tokens altogether and return the estimated behavior instead.
  arXiv  Detail & Related papers  (2025-02-18T23:13:16Z)
• Not all tokens are created equal: Perplexity Attention Weighted Networks for AI generated text detection [49.15148871877941]
  Next-token distribution outputs offer a theoretically appealing approach for detection of large language models (LLMs)<n>We propose the Perplexity Attention Weighted Network (PAWN), which uses the last hidden states of the LLM and positions to weight the sum of a series of features based on metrics from the next-token distribution across the sequence length.<n>PAWN shows competitive and even better performance in-distribution than the strongest baselines with a fraction of their trainable parameters.
  arXiv  Detail & Related papers  (2025-01-07T17:00:49Z)
• Attribute-to-Delete: Machine Unlearning via Datamodel Matching [65.13151619119782]
  Machine unlearning -- efficiently removing a small "forget set" training data on a pre-divertrained machine learning model -- has recently attracted interest. Recent research shows that machine unlearning techniques do not hold up in such a challenging setting.
  arXiv  Detail & Related papers  (2024-10-30T17:20:10Z)
• Promises and Pitfalls of Generative Masked Language Modeling: Theoretical Framework and Practical Guidelines [74.42485647685272]
  We focus on Generative Masked Language Models (GMLMs) We train a model to fit conditional probabilities of the data distribution via masking, which are subsequently used as inputs to a Markov Chain to draw samples from the model. We adapt the T5 model for iteratively-refined parallel decoding, achieving 2-3x speedup in machine translation with minimal sacrifice in quality.
  arXiv  Detail & Related papers  (2024-07-22T18:00:00Z)
• Self-Augmented Preference Optimization: Off-Policy Paradigms for Language Model Alignment [104.18002641195442]
  We introduce Self-Augmented Preference Optimization (SAPO), an effective and scalable training paradigm that does not require existing paired data. Building on the self-play concept, which autonomously generates negative responses, we further incorporate an off-policy learning pipeline to enhance data exploration and exploitation.
  arXiv  Detail & Related papers  (2024-05-31T14:21:04Z)
• Evaluation of Active Feature Acquisition Methods for Static Feature Settings [6.645033437894859]
  We introduce a semi-offline reinforcement learning framework for active feature acquisition performance evaluation (AFAPE) Here, we study and extend the AFAPE problem to cover static feature settings, where features are time-invariant. We derive and adapt new inverse probability weighting (IPW), direct method (DM), and double reinforcement learning (DRL) estimators within the semi-offline RL framework.
  arXiv  Detail & Related papers  (2023-12-06T17:07:42Z)
• Beyond MLE: Convex Learning for Text Generation [34.99340118597274]
  We argue that Maximum likelihood estimation (MLE) is not always necessary and optimal, especially for closed-ended text generation tasks like machine translation. We propose a novel class of training objectives based on convex functions, which enables text generation models to focus on highly probable outputs without having to estimate the entire data distribution.
  arXiv  Detail & Related papers  (2023-10-26T08:08:43Z)
• Amortizing intractable inference in large language models [56.92471123778389]
  We use amortized Bayesian inference to sample from intractable posterior distributions. We empirically demonstrate that this distribution-matching paradigm of LLM fine-tuning can serve as an effective alternative to maximum-likelihood training. As an important application, we interpret chain-of-thought reasoning as a latent variable modeling problem.
  arXiv  Detail & Related papers  (2023-10-06T16:36:08Z)
• Personalized Federated Learning under Mixture of Distributions [98.25444470990107]
  We propose a novel approach to Personalized Federated Learning (PFL), which utilizes Gaussian mixture models (GMM) to fit the input data distributions across diverse clients. FedGMM possesses an additional advantage of adapting to new clients with minimal overhead, and it also enables uncertainty quantification. Empirical evaluations on synthetic and benchmark datasets demonstrate the superior performance of our method in both PFL classification and novel sample detection.
  arXiv  Detail & Related papers  (2023-05-01T20:04:46Z)
• Augmenting Interpretable Models with LLMs during Training [73.40079895413861]
  We propose Augmented Interpretable Models (Aug-imodels) to build efficient and interpretable models. Aug-imodels use LLMs during fitting but not during inference, allowing complete transparency. We explore two instantiations of Aug-imodels in natural-language processing: (i) Aug-GAM, which augments a generalized additive model with decoupled embeddings from an LLM and (ii) Aug-Tree, which augments a decision tree with LLM feature expansions.
  arXiv  Detail & Related papers  (2022-09-23T18:36:01Z)
• Complex Event Forecasting with Prediction Suffix Trees: Extended Technical Report [70.7321040534471]
  Complex Event Recognition (CER) systems have become popular in the past two decades due to their ability to "instantly" detect patterns on real-time streams of events. There is a lack of methods for forecasting when a pattern might occur before such an occurrence is actually detected by a CER engine. We present a formal framework that attempts to address the issue of Complex Event Forecasting.
  arXiv  Detail & Related papers  (2021-09-01T09:52:31Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.