ARM: Efficient Guided Decoding with Autoregressive Reward Models

In this section, we outline the approach of guiding an autoregressive base language model with external experts. At each step $t$ of decoding, we have a prompt together with an already generated prefix $x_{:t}$. We assume we have access to the logits $z_{t}\in R^{|V|}$ of an autoregressive model $p_{LM}(x_{t}|x_{:t})\propto\exp(z_{t}(x_{t}|x_{:t}))$, where $|V|$ is the size of vocabulary, and our goal is to augment these logits with an external reward model scores $r_{t}\in[0,1]^{|V|}$. The augmented logits are then mapped to probabilities via Softmax, and a next token is then sampled from the multinomial distribution:

$$\tilde{p}(x_{t}|x_{:t})=\operatorname{Softmax}(z_{t}+\beta r_{t}).$$

(1)

To define expert logits, Liu et al. (2021); Krause et al. (2021) propose to use attribute-conditioned language models, trained on the next token prediction task on class-conditioned data: $r_{LM,a}~{}=~{}z_{t}(x_{t}|x_{:t},a)$. The main advantage of this approach is that for a given prefix $x_{:t}$ we only pass it through the discriminator backbone once, relying on the language model output layers to obtain predictions $r_{t}$ for each of the next tokens. Deng & Raffel (2023) show that training a less efficient discriminator leads to better quality. Instead of relying on the language model objective, they explicitly train a discriminator to predict the attribute of interest: $r_{RAD}(x_{t}|x_{:t},a)=p(a|[x_{:t},x_{t}])$, where $[\cdot,\cdot]$ denotes the concatenation of a prefix and a next token candidate. This approach requires passing each next token candidate as input to the model, thus, to obtain the scores for $k$ next token candidates they would need $k$ full forward passes, which can slow down inference significantly, and bounds them to use top-k decoding to limit the number of next token candidates.

Both of these approaches finetune an autoregressive GPT-$2$ (Radford et al., 2019) language model backbone, however, they parameterize the rewards differently suggesting that we are limited to either a fast or an expressive model.

In our work, we ask whether it is possible to have a both efficient and expressive autoregressive reward model suitable for efficient on-the-fly guidance of the base language model.

For a reward model to be efficiently utilized during autoregressive decoding, we aim to satisfy the following requirements:

1. 1.

   quality: guided decoding should produce fluent outputs with high constraint satisfaction probability comparable to Deng & Raffel (2023).

2. 2.

   autoregressive decoding efficiency: complexity of reward model should linearly depend on the context length for each decoding step (as in Deng & Raffel (2023))

3. 3.

   prediction efficiency: the backbone is called only once for each decoding step and does not depend on the number of next token candidates (as in Liu et al. (2021); Krause et al. (2021))

We answer affirmatively, by proposing ARM, an autoregressive reward model suitable for guided decoding, designed for efficient modelling of rewards scores. We demonstrate that guided decoding with our ARM results in high quality of constrained generation on detoxification and sentiment control tasks (see section 4).

In order to satisfy autoregressive decoding efficiency and, hopefully, high quality of predictions, we employ the same strategy as Deng & Raffel (2023) finetuning an autoregressive language model backbone with classifier heads to predict the reward scores of each next token given the prefix. This formulation allow us to cache the activations of the model during left-to-right decoding, thus having linear complexity w. r. t. the input length for each decoding step.

To ensure prediction efficiency, we aim to preserve the language modelling style of prediction to get a score for each of the next tokens modeled as a similarity between a predicted hidden vector $h_{t}\in\mathbb{R}^{d}$ and output embeddings $e_{i}\in\mathbb{R}^{|V|\times d}$ (Liu et al., 2021; Krause et al., 2021), and develop a suitable parameterization of the next token reward. We assume that reward scores are continuous values and without loss of generality assume $r_{t}\in[0;1]^{|V|}.$ We propose the following ARM parameterization of the scores for the next tokens given the prefix:

$$\hat{r}_{i}(e_{i}|h_{t})=\underset{\text{baseline}}{\tilde{r}(h_{t})}+\Delta\hat{r}(e_{i}|h_{t})$$

(2)

Particularly, we use linear parameterization of the head: $\tilde{r}(h_{t})=\langle h_{t},w\rangle$, $\Delta\hat{r}(e_{i}|h_{t})={\langle Wh_{t},e_{i}\rangle}$. Here, we introduced two attribute-specific parameters: $w\in\mathbb{R}^{d}$ for modeling the baseline score of the prefix, and $W\in\mathbb{R}^{d\times d}$ to model marginal reward predicted for each of next token candidates at position $t$.

At the training stage, we assume that we have a dataset of $N$ sentences $x_{i}$ each labeled with an attribute $r_{i}\in[0,1]$, $i\in\{1..N\}$. We finetune the autoregressive reward model to predict the attribute value $\hat{r}_{t}$ (prefix reward) for each of the partial prefixes $x_{:t}$. While a naive strategy would be to learn on the full sentences, we adopt a cumulative loss function from Deng & Raffel (2023) encouraging the model to learn the future attribute for an incomplete prefix:

$$\mathcal{L}(\hat{r},r)=\frac{\sum_{t=1}^{l}w_{t}\ell(\hat{r}_{t},r)}{\sum_{t=1}^{l}w_{t}}.$$

(3)

where $w_{t}=t$, and where $\ell$ defines the distance between true sentence attribute and predicted prefix reward. The intuition of the loss is to put more weight to the prefixes closer to the full sentence, where the ground truth labels are provided, while still encouraging the model to learn rewards for partial prefixes. Similar to (Deng & Raffel, 2023), we use $\ell(\hat{r},r)=(\hat{r}-r)^{2}$ throughout this work.

Additionally, we are interested whether we can finetune a reward model with our parameterization to match the performance of an already trained less efficient but potentially more expressive reward model. To answer this question, we use a distillation loss (Hinton et al., 2015) to mimic the outputs of a less efficient high-quality reward model:

$$\mathcal{L}_{\text{dstl}}(\hat{r},r)=\frac{1}{l}\sum_{t=1}^{l}\ell(\hat{r}_{t},r_{t})$$

(4)

A reward model only can observe a limited number of tokens during finetuning. While the loss defined above provides the positive signal for some tokens, it might be beneficial to make on objective more “contrastive” and regularize the prediction for other tokens (including rare ones). In our parameterization (2), it is natural to push the predicted reward towards the baseline for unrelated tokens, regularizing the outputs.

We regularize the prediction of our model to be close on average to the prefix baseline by forcing the marginal prediction $\Delta\hat{r}$ to be close to $0$ for randomly sampled token candidates:

$$\mathcal{L}_{\text{reg}}(h_{t})=\mathbb{E}_{e^{\prime}\sim U[V]}\left[\Delta\hat{r}(e^{\prime}|h_{t})\right]^{2},$$

(5)

where we use one sample of $e^{\prime}$ for each prefix position, sampling uniformly from a set of all possible output embeddings.

For the detoxification evaluation, we follow previous work (Deng & Raffel, 2023; Liu et al., 2021) and evaluate samples from guided decoding given a 10k subset (Liu et al., 2021) of prompts from the RealToxicityPrompts dataset (Gehman et al., 2020). We follow Deng & Raffel (2023); Liu et al. (2021) and finetune our model on 2M pairs of text and continuous labels between $0$ and $1$ from the Jigsaw Unintended Bias in Toxicity Classification challenge (cjadams et al., 2019). Like previous work, we train our model on $7$ independent labels (‘toxicity’, ‘severe toxicity’, ‘obscene’, ‘identity attack’, ‘insult’, ‘threat’, ‘sexual explicit’) with different head parameters $w_{i},W_{i},i\in\{1,...,7\}$ for each sub-task. During decoding, we only use the predictor of toxicity. For the distillation experiment, we use the same dataset, and the released toxicity discriminator from Deng & Raffel (2023) as a teacher.

During decoding, we sample $25$ continuations generating at most $20$ new tokens. To evaluate toxicity, we use an external closed-source toxicity classifier Perspective API (Lees et al., 2022), and following previous work (Deng & Raffel, 2023; Liu et al., 2021), we rely on Maximal Average Toxicity metric, which is the maximal value of toxicity score over $25$ samples for a given prompt, averaged over the set of 10k prompts. We also report Toxic Rate, which is calculated as probability that at least one out of 25 continuations is toxic according to Perspective API (toxicity score > $0.5$); and Diversity score, which is the average number of distinct $n$-grams normalized by the length of text (Li et al., 2018). To evaluate the fluency of model generations, we use the average perplexity of GPT-2-XL, calculated over the sampled continuations. In the experiments, we will look at the toxicity/fluency trade-off, alternating the weight $\beta$ of the discriminator (see section B.1 and section B.2). We expect to obtain a model with both low toxicity according to PerspectiveAPI, and high fluency.

For sentiment control, we follow the previous work of Li et al. (2018); Sudhakar et al. (2019); Liu et al. (2021); Deng & Raffel (2023) to evaluate the samples given a prompt from one of the three categories: $2.5K$ negative, $5K$ neutral, and $2.5K$ positive prompts from OpenWebText (Gokaslan & Cohen, 2019). To finetune ARM on labels only, we follow Deng & Raffel (2023) and finetune our model on millions of reviews from the Amazon Polarity (Zhang et al., 2015) and and SST-2 (Socher et al., 2013) datasets. To distil the sentiment discriminator of Deng & Raffel (2023), we use text examples from the Amazon Polarity dataset. Additional training details are provided in appendix A.

For evaluation, we follow Deng & Raffel (2023), and use the average Positive Rate metric w. r. t. the finetuned external DistilBERT classifier (Sanh et al., 2019) provided via the HuggingFace text classification pipeline¹¹1https://huggingface.co/distilbert/distilbert-base-uncased-finetuned-sst-2-english.

As in the toxicity task, we use GPT-2-XL to evaluate the fluency of the sampled continuations, and we expect to obtain high Positive Rate and low average perplexity.

Observing good results for our parameterization in previous experiments, we then investigate the effect of the regularization component and baseline component for head parameterization.

As shown in figure 4, we experiment with distillation and observe that turning off regularization, or further removing baseline from head parameterization results in higher average perplexity, although these simpler models still show adequate decrease in toxicity. By analyzing the distribution of rewards for all next tokens (see figure 5), we observe that the unregularized ARM tends to predict less flat distributions (higher values in $\Delta\hat{r}$ component), which may hurt the fluency.

The main advantage of our reward model is that it need only a single pass through the backbone for each step of decoding. In table 1, we compare the number of backbone passes needed for external expert models. In figure 6, we measure time per generated token when running the decoding for the toxicity task for us and the model of Deng & Raffel (2023). When we increase $k$ in top-k decoding with the fixed batch size (equal to the number of requested samples $=25$), the baseline model needs to perform $k$ full forward passes through intermediate layers. For our model, during top-k decoding, we simply index relevant next token embeddings to calculate the dot product with the output hidden state.

Furthermore, as in Deng & Raffel (2023); Liu et al. (2021), our expert model is autoregressive, which means that we cache the prefix activations during decoding.