Controlled LLM Decoding via Discrete Auto-regressive Biasing

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. Large language models (LLMs) are widely used in real-world applications through chatbots such as ChatGPT, Alpaca, and Llama, making them an important part of everyday life (Bender et al., 2021; Bommasani et al., 2021; Weidinger et al., 2021). As a result, there has been growing attention on developing methods to reliably and effectively control LLM-generated outputs to meet user-defined constraints (Gehman et al., 2020; Dathathri et al., 2020; Goshvadi et al., 2023; Han et al., 2024). Previous work has tackled controlled language generation using decoding-time algorithms, which bypass the need for fine-tuning the base language model (Liu et al., 2023a; Kumar et al., 2022; Mireshghallah et al., 2022; Dathathri et al., 2020; Qin et al., 2022). Among these, energy-based decoding methods define a target distribution through an energy function, combining multiple constraints into a weighted average. This formulation offers significant flexibility, as the energy function can be any arbitrary function. Sampling from this distribution relies on gradient-based MCMC in continuous spaces, followed by conversion back to discrete text tokens.