Controlled LLM Decoding via Discrete Auto-regressive Biasing

Sequence Length:

We do not emphasize that our algorithm requires a specified sequence length as this is a common requirement within the literature — MuCOLA [1], COLD [2], and BOLT [3] all require the specification of sequence length.

Furthermore, it is possible to enable our algorithm to produce sequences of varying lengths by padding / truncating the bias sequence during generation. We did not focus on enabling dynamic sequence length as this is not a core contribution of our work.

Discrete Langevin Proposal

We agree with your advice and have included the following discussion regarding our application of Discrete Langevin Proposal (DLP) in the Appendix in our revision. Specifically, to enable the use of large step sizes in the proposal, we adopt the globally balanced version of the DLP proposal [1, 2]:

$\text{Categorical} \left( \underset{j \in |V|}{\text{softmax}} \left( \nabla f(\hat{B} | X)_i (\text{Onehot}_j - \hat{b}_i)\right) \right)$

Here, $\text{Onehot}_j$ represents the one-hot vector for the $j$ token in the vocabulary $V$, and $\hat{B} = \{\hat{b}_1,\hat{b}_2, … \hat{b}_n \}$ represents the original one-hot vector sequence of length $n$. To obtain a distribution over the vocabulary $|V|$, we must compute the inner term for every token $j \in V$.

Here, we note that $(\text{Onehot}_j - \hat{b}_i))$ corresponds to the distance between the original token at position $i$ and every token $j$ in the vocabulary $V$. Given the discrete nature of the tokens, we choose to use hamming distance to represent this term. For a token $j$, the hamming distance to the original token in position $i$ is 0 if the $j$th coordinate $\hat{b}\_{ij}= 1$, as they are the same token; and 1 if the $j$th coordinate is 0. Thus we can represent the hamming distance between token $j$ and the current token as $1 - \hat{b}\_{ij}$. Below we include the proposal distribution we sample from to obtain the new token for position $i$.

$b’_i \sim \text{categorical}\left(\underset{j \in |V|}{\text{softmax}} \left( \frac{1}{\tau} (\nabla f(\hat{B} | X))\_{ij} (1 - \hat{b}\_{ij}) \right) \right)$

Here, $b’_i$ is the token we sample from the categorical distribution over V on the right-hand side. We have incorporated this discussion into the appendix, following your advice.

Gradient Computation

Once we represent the sequence of tokens $B$ as a sequence of one-hot vectors $\hat{B}$, we are able to compute the gradient of the constraint function with respect to $\hat{B}$ using automatic differentiation software, such as Torch Autograd.

Use of Greedy Decoding

We define the biased auto-regressive generation used in our algorithm in terms of an argmax as shown in equation 10, which should convey that this specific component of our algorithm is deterministic.

Additionally, it should be noted that while we choose greedy decoding, our method is compatible with other decoding approaches that are conventionally used, such as top-k, top-p, and standard Gibbs. We choose to use greedy decoding since this was used by prior work [2, 3] and it places emphasis on the novel aspects of our framework.

[1] Kumar et al. Gradient-based Constrained Sampling from Language Models. ACL 2022.

[2] Quin et al. COLD Decoding: Energy-based Constrained Text Generation with Langevin Dynamics. NeurIPS 2022.

[3] Liu et al. BOLT: Fast Energy-based Controlled Text Generation with Tunable Biases. ACL 2023.

[4] Zhang et al. A Langevin-like Sampler for Discrete Spaces. ICML 2022.

[5] Pynadath et al. Gradient-based Discrete Sampling by Automatic Cyclical Scheduling. NeurIPS 2024.

Defining $P(Y | X) \propto P^{LM} (Y | X) \exp f(Y)$ is equivalent to the energy function formulation used in previous works (lines 155-157) with $\lambda_1 = \lambda_2 = 1$. As you noted, DLP can be used to directly sample from $P(Y | X)$. However, this approach does not support auto-regressive generation and significantly compromises fluency. Non-autoregressive approaches usually suffer from lack of fluency as demonstrated in our experimental results in Table 2 and previous works [1]. Additionally, methods like RLHF or PPO require extra data and fine-tuning, making them unsuitable as inference-time approaches.

The primary motivation behind our framework is the observation that fluency is best satisfied through auto-regressive generation, and gradient-based sampling efficiently finds responses that satisfy constraints. By framing the problem as a joint distribution of $Y$ and $B$, we enable the use of both methods — we use autoregressive generation to obtain $Y$, ensuring fluent generations; and we apply Discrete Langevin Proposal (DLP) to sample $B$, ensuring constraint satisfaction. Furthermore, all of this is accomplished without the need to fine-tune the underlying LM. We will add the above explanation to the revision.

[1] Liu et al. BOLT: Fast Energy-based Controlled Text Generation with Tunable Biases. ACL 2023.

Our computation of $P(Y | X, B)$ does not involve concatenating the input $X$ with the bias sequence $B$. As we explain in equation 10 in Sec 4.2, we use $B$ as a biasing sequence that is added to the unnormalized logits the language model outputs for each position. Specifically, during the auto-regressive generation, after the model produces the unnormalized logits $\tilde{y_i}$ for position $i$, we add $\tilde{b_i}$ to obtain a new logit vector. We then apply greedy decoding to this logit vector to obtain the token for this position. However, it is possible to apply other forms of sampling — we choose greedy decoding since this is what previous work [1, 2] uses and we want to ensure that comparisons between our method and theirs are as fair as possible.

[1] Liu et al. BOLT: Fast Energy-based Controlled Text Generation with Tunable Biases. ACL 2023.

[2] Quin et al. COLD Decoding: Energy-based Constrained Text Generation with Langevin Dynamics. NeurIPS 2022.

Thank you for providing [1], [2], and [3]. While we did discuss works that fine-tune the model very briefly, we were unaware of these specific works and will include them in the revision, as they add more breadth to our discussion on controlled text generation methods.

We want to emphasize that while the suggested papers are related, they are not closely aligned with the focus of our work. As you pointed out in your review, [1], [3] focus on methods that fine-tune auto-regressive models to align with some defined EBM. Our work is primarily concerned with inference-time decoding algorithms that do not require the model to be fine-tuned.

We would also like to point out that the algorithm presented in [2] requires some proposal function that can be used to generate samples. Its core contribution is an accept / reject algorithm that is agnostic to the proposal function. This is different from our algorithm, which presents a novel decoding algorithm that directly generates new samples. Thus we compare primarily to other works that introduce decoding methods that generate samples, which is not the focus of [2]. Nevertheless, we will make sure to include all three works in our related work section as they provide valuable insight into alternative approaches to this problem.

[1] Khalifa et al. A distributional approach to controlled text generation. ICLR 2021.

[2] Eikema et al. An approximate sampler for energy-based models with divergence diagnostics. TMLR 2022.

[3] Kruszewski et al. disco: a toolkit for Distributional Control of Generative Models. ACL 2023.

Discrete Langevin Proposal function

The $(1 - \hat{b}\_{ij})$ term corresponds to the hamming distance between tokens. We include this term as a result of the definition of Discrete Langevin Proposal (DLP) introduced by [1]. For your convenience, we will explain how we obtain the proposal function below.

To enable the use of large step sizes in the proposal, we adopt the globally balanced version of the DLP proposal [1, 2]:

$\text{Categorical} \left( \underset{j \in |V|}{\text{softmax}} \left( \nabla f(\hat{B} | X)_i (\text{Onehot}_j - \hat{b}_i)\right) \right)$

Here, $\text{Onehot}_j$ represents the one-hot vector for the $j$ token in the vocabulary $V$, and $\hat{B} = \{\hat{b}_1,\hat{b}_2, … \hat{b}_n \}$ represents the original one-hot vector sequence of length $n$. In order to obtain a distribution over the vocabulary $|V|$, we must compute the inner term for every token $j \in V$.

We note that $(\text{Onehot}_j - \hat{b}_i))$ corresponds to the distance between the original token at position $i$ and every token $j$ in the vocabulary $V$. Given the discrete nature of the tokens, we choose to use hamming distance to represent this term. For a token $j$, the hamming distance to the original token in position $i$ is 0 if the $j$th coordinate $\hat{b}\_{ij} = 1$, as they are the same token; and 1 if the $j$th coordinate is 0. Thus we can represent the hamming distance between token $j$ and the current token as $1 - \hat{b}\_{ij}$. Below we include the proposal distribution we sample from to obtain the new token for position $i$.

$b’_i \sim \text{categorical} \left(\underset{j \in |V|}{\text{softmax}} \left( \frac{1}{\tau} (\nabla f(\hat{B} | X))\_{ij} (1 - \hat{b}\_{ij}) \right) \right)$

Here, $b’_i$ is the token we sample from the categorical distribution over V on the right-hand side.

Use of greedy decoding

Equation 10 is an argmax to represent greedy decoding, which is a commonly used decoding technique for auto-regressive generation. While other sampling methods are compatible with our algorithm, we choose greedy decoding as previous inference-time controlled generation algorithms [3, 4], also use greedy decoding. This helps emphasize the novel aspects of our framework and ensures a fair comparison with previous works.

Initialization of $B$

To answer your question regarding why we initialize $B$ to $Y$ when sampling from $P(B | X, Y)$, we provide our reasoning in section 4.2 under the section titled “Sampling from $P(B | X, Y)$”. Specifically, our goal is to sample from the distribution:

$P(B | X, Y) \propto P^{LM}(Y | X, B) \exp(f(B | X))$

As we discuss in the mentioned section, sampling from this distribution would require computing $P(Y | X, B)$ for all possible values of $B$, which is intractable. In order to obtain a more feasible calculation, we note that $P^{LM} (Y | X, B)$ will be high when the bias $B$ aligns with the original response $Y$ due to the nature of auto-regressive generation. Thus we approximate this distribution by initializing $B = Y$, which will ensure a relatively high value for $P^{LM} (Y | X, B)$. By initializing the bias term into a region with high values of $P(Y | X, B)$, all that remains is to determine which samples within this region enable better constraint satisfaction. If we initialize $B=0$, then sampling must simultaneously find $B$ that results in high $P(Y | X, B)$ and high $f(B | X)$, which is a more difficult task.

[1]. Zhang et al. A Langevin-like Proposal for Discrete Spaces. ICML 2022.

[2]. Pynadath et al. Gradient-based Discrete Sampling via Automatic Cyclical Scheduling. NeurIPS 2024.

[3]. Liu et al. BOLT: Fast Energy-based Controlled Text Generation with Tunable Biases. ACL 2023.

[4]. Qin et al. COLD Decoding: Energy-based Constrained Text Generation with Langevin Dynamics. NeurIPS 2022.

We appreciate the time taken to point out the notational errors in our work and areas of confusion. We were wondering if there were any other flaws that lead to your score of a 3 — are there any serious concerns you have with the method itself? Were there any concerns with the claims made in our paper and whether we provide sufficient evidence to validate them? We would greatly appreciate your input in helping us improve our work.

Thank you for your helpful comments. We would like to address some of your concerns to clarify the technical details and mathematical formulation.

Saturated benchmarks

We are unclear about your statement regarding benchmarks being "saturated." Do you mean there are too many methods addressing the same task? We believe the tasks we selected are appropriate as they allow for a clear comparison of the difference in performance between controlled text generation algorithms. As shown in Table 2, these tasks are useful for understanding where prior methods fall short and where our method excels. Furthermore, the use of similar tasks in prior works demonstrates that our chosen benchmarks are reasonable [1, 2, 5, 6].

Furthermore, regarding your claim that prompting can solve these tasks, we are not aware of any papers that demonstrate prompting is sufficient for the tasks we consider. In fact, [1] does compare to Prompt-T5, a pre-trained LM intended to solve tasks with prompts. They demonstrate that LM-Steer outperforms Prompt-T5, which shows that prompting is not enough to match controlled text generation methods. We choose not to compare to Prompt-T5 as [1] already establishes the limitations of prompt-based methods.

Finally, model jailbreaking is a separate research direction from the core focus of this paper. None of the methods we compare against include jailbreaking as a task [1, 2, 5, 6]. We consider it an application of our proposed method that could be investigated for future directions.

[1] Han et al. Word Embeddings Are Steers for Language Models. ACL 2024.

[2] Liu et al. BOLT: Fast Energy-based Controlled Text Generation with Tunable Biases. ACL 2023.

[3] Stidikov et al. Classifiers are Better Experts for Controllable Text Generation. Workshop on Transfer Learning for NLP, 2022.

[4] Dekoninck et al. Controlled Text Generation via Language Model Arithmetic. ICLR 2024.

[5] Kumar et al. Gradient-based Constrained Sampling from Language Models. ACL 2022.

[6] Qin et al. COLD Decoding: Energy-based Constrained Text Generation with Langevin Dynamics. NeurIPS 2022.

Notational errors

Thank you for the feedback. While equation 6 is correct, we understand that it can be confusing as it is not very direct. We will change it to “$P(B|X)$ is defined to be $\frac{\exp(f(B | X))}{Z_B}$”. The equation that immediately follows should be $P(B | X, Y)$ as you have pointed out. We occasionally drop the conditioning on X as all terms should be understood as conditioned on prompt X. We have revised the notations in the paper to make them clearer.

Thank you for your supportive and thoughtful review. We respond to your points below:

Weakness: Auto-regressive sampling

This is correct — the main deficiency of our method is the fact that each step requires the autoregressive generation of the sequence. Prior work also suffers from this issue [1]. Improving efficiency by reducing the number of auto-regressive calls is an interesting future direction.

Weakness: Mixed results on Sentiment Task

We would like to point out that while our method underperforms BOLT slightly in the context of fluency, our method greatly outperforms BOLT in terms of sentiment control while almost matching BOLT’s fluency performance. While our method has the second-best fluency metrics, BOLT lags behind other baselines for most of the sentiment-specific metrics. Additionally, as shown in the qualitative examples we include in Table 6 of the Appendix, DAB generations are not noticeably less fluent than BOLT generations. Finally, DAB outperforms BOLT in terms of fluency on the keyword generation task, as shown in Table 2. We believe that this demonstrates that our method is able to achieve a superior balance between control and fluency.

Question: Alternative Calculation of Bias Vector

We chose to use the $l_2$ distance as this created a much more peaked distribution towards the desired token. When we tried using the inner product, it did not enable sufficient control as the bias vector did not bias strongly enough towards the sampled tokens. We will add this result to the appendix.

[1] Liu et al. BOLT: Fast Energy-based Controlled Text Generation with Tunable Biases. ACL 2023.

Thank you for your supportive review. We were wondering if there were any aspects of our paper that prevented you from giving it a higher score. We would be more than happy to address any questions or concerns you may have regarding our paper.