Controlled LLM Decoding via Discrete Auto-regressive Biasing
Patrick Pynadath, Ruqi Zhang
Abstract
Controlled text generation allows for enforcing user-defined constraints on large language model outputs, an increasingly important field as LLMs become more prevalent in everyday life. One common approach uses energy-based decoding, which defines a target distribution through an energy function that combines multiple constraints into a weighted average. However, these methods often struggle to balance fluency with constraint satisfaction, even with extensive tuning of the energy function's coefficients. In this paper, we identify that this suboptimal balance arises from sampling in continuous space rather than the natural discrete space of text tokens. To address this, we propose Discrete Auto-regressive Biasing, a controlled decoding algorithm that leverages gradients while operating entirely in the discrete text domain. Specifically, we introduce a new formulation for controlled text generation by defining a joint distribution over the generated sequence and an auxiliary bias sequence. To efficiently sample from this joint distribution, we propose a Langevin-within-Gibbs sampling algorithm using gradient-based discrete MCMC. Our method significantly improves constraint satisfaction while maintaining comparable or better fluency, all with even lower computational costs. We demonstrate the advantages of our controlled decoding method on sentiment control, language detoxification, and keyword-guided generation. We make our code available at the following url: https://github.com/patrickpynadath1/dab .
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Cited by top-tier papers1
Ask how each one uses itBuilds on18
- BERTScore: Evaluating Text Generation with BERTTianyi Zhang, Varsha Kishore, Felix Wu, Kilian Q. Weinberger et al.ICLR 2020 · 8,443 citations
- Plug and Play Language Models: A Simple Approach to Controlled Text GenerationSumanth Dathathri, Andrea Madotto, Janice Lan, Jane Hung et al.ICLR 2020 · 1,166 citations
- COLD Decoding: Energy-based Constrained Text Generation with Langevin DynamicsLianhui Qin, Sean Welleck, Daniel Khashabi, Yejin ChoiNeurIPS 2022 · 217 citations
- COLD-Attack: Jailbreaking LLMs with Stealthiness and ControllabilityXingang Guo, Fangxu Yu, Huan Zhang, Lianhui Qin et al.ICML 2024 · 173 citations
- A Distributional Approach to Controlled Text GenerationMuhammad Khalifa, Hady Elsahar, Marc DymetmanICLR 2021 · 135 citations
Related papers
- Principled Gradient-Based MCMC for Conditional Sampling of TextLi Du, Afra Amini, Lucas Torroba Hennigen, Xinyan Velocity Yu et al.ICML 2024 · 6 citations
- Gradient-based Constrained Sampling from Language ModelsSachin Kumar, Biswajit Paria, Yulia TsvetkovEMNLP 2022 · 22 citations
- Air-Decoding: Attribute Distribution Reconstruction for Decoding-Time Controllable Text GenerationTianqi Zhong, Quan Wang, Jingxuan Han, Yongdong Zhang et al.EMNLP 2023 · 7 citations
- Control Large Language Models via Divide and ConquerBingxuan Li, Yiwei Wang, Tao Meng, Kai-Wei Chang et al.EMNLP 2024 · 1 citation
- Controlled Text Generation as Continuous Optimization with Multiple ConstraintsSachin Kumar, Eric Malmi, Aliaksei Severyn, Yulia TsvetkovNeurIPS 2021 · 91 citations
