Efficient Controlled Language Generation with Low-Rank Autoregressive Reward Models

We thank the reviewer for their feedback. We appreciate that they highlight the clarity of our paper, as well as the extensiveness of our experiments and evaluation.

Are there hyper-parameters similar to the beta in RAD & ARM for GeDi and DExperts to trade-off fluency for detoxicity? If so, it would be good to show their operating points in Figure 3 to make a stronger case.

We agree that it will make the Figure 3 better if we include the trade-off plot for the DExperts model (which is known to outperform GeDi). To give a quicker response, we add a preliminary figure (see Figure 8 in Appendix) by running the evaluation over 1000 prompts for ARM vs DExperts. We observe that DExperts start to diverge early from ARM resulting in lower effectiveness of the language modeling-based guidance. We aim to include the DExperts evaluated over all toxicity prompts for the final version.

in Figure 3, increasing k seems to induce a right-ward shift (i.e., higher perplexity at similar toxiticy level). Could you help me understand why?

This is because for larger k there are less probable tokens that can be produced by the model. The relative performance of ARM compared to RAD remains similar for both settings. For simplicity of the Figure 3, we only leave k=20 case, and move k=40 case to the Appendix, as was suggested by another reviewer.

Line 431: “we observe that regularization effectively decreases the rank of R_ARM which might explain the higher fluency of regularized model.” This suggests lower-rank approximation is better in terms of fluency. Can you expand on this point and help me understand why this is the case?

We meant to highlight that our regularization is aimed to push the model to abstain more by regularizing the prediction towards rank-1 output (predicting the baseline score for every next token, which does not modify the base model distribution). We will clarify: “Particularly, a very strong regularization would result in the model always predicting the baseline score for each of the next tokens (corresponds to the rank-1 output), which does not modify the original distribution of the model (the best fluency).”

What is the compute time required for running evaluation and training ARM?

Here, we report the training/generation time for the detoxification task. For generation, we report the time to obtain a single point on the trade-off plot (a single value of $\beta$).

Training: for ARM and RAD similar and approx. 12 hours per epoch on 1 GPU.

Generation speed for ARM vs RAD:

We thank the reviewer for their feedback. We are happy to hear that our work is easy to follow; that we conduct thorough experiments and provide enough experiential details.

In Section 3.3 on ARM training, is the language model backbone also fine-tuned, or is it kept frozen?

We finetune all parameters of the reward model except input/output embeddings, as we state in section 4.1.

The proposed approach offers limited novelty.

We appreciate a valuable connection to the reinforcement learning-based controlled generation. While their contribution is not the parametrization of Q_D, Cao et al. 2023 use a language model architecture similar to ours, DExperts, GeDi. Our work zooms in on the parametrization choice, we explore the implications of rank, and we find a surprising result in terms of the rank of reward models. This is what we bring as truly novel, in contrast to simply predicting the scores in one go (e.g. DExperts, GeDi already did this). We discuss the relationship to this work in the revision. We also add the work of Mudgal et al. 2024, where they use prefix scorers in a way similar to how less efficient RAD or FUDGE are used. We hope our work will inform the reinforcement learning branch of controlled generation.

• Cao et al. 2023. Systematic Rectification of Language Models via Dead-end Analysis, ICLR 2023
• Mudgal et al. 2024. Controlled Decoding from Language Models, ICML 2024

Efficiency improvement seems limited as most of the computational cost comes from the LM backbone In our experiments, we observe a high difference between ARM and RAD in terms of evaluation speed.

For generation, here we report the time to obtain a single point on the trade-off plot (a single value of $\beta$).

Decoding for ARM vs RAD:

We would like to thank the reviewer for the feedback. We are happy to hear that they find our motivation, writing, and mathematical formulations to be clear!

Delegating most of the rigorous mathematics to the appendix is problematic. I think that it would seriously strengthen the paper to at least include statements of the "main theorems/propositions" and then delegate the proofs to the appendix.

We thank the author for this suggestion! We followed your advice and moved some statements from the Appendix to the main text and we believe this indeed strengthens our work. We refer reviewer to an updated version of the pdf (3.1 Analysis of RAD).

In Equation (8), I don't think RAD teacher \tilde r(x) is defined. I think this should be explicitly stated somewhere in that section that (7) and (8) are completely different training paradigms.

We apologize for the missing introduction of the distillation task, we improve this in section 3.3 (ARM Training).

(The main weakness, in my opinion) the graphs/results are very unclear. In Figure 3, the figure is too busy for me to understand what's going on. I recommend having 1 set of graphs with k=20, and a second set with k=40 (and putting set in the appendix) and zooming in so that we can actually see the order of magnitude differences between the curves that are being compared in the graph.

This is a great suggestion, we refactored the plots and moved parts of the plot with k=40 to the Appendix.

it is unclear to me how each of the individual points on the curves is obtained (I know there is some third, unseen parameter which parameterizes the curves)

You are right, there is a third scalar parameter, which defines a position on the trade-off plot. We didn’t include these parameters in the main figures to not make a plot more complicated, and also because the ranges of these control coefficients are not directly comparable between models (one compares the two trade-off lines as a whole: fixes one metric and compares another one). We included the used trade-off parameters beta in the Appendix (Tables 2-6).

I'm not sure what point this graph tries to convey; the paper needs to explain what each curve's trend represents, relative to the other curve's trends.

Thank you for this comment, we add an introductory paragraph in Section 4.4. Results to explain how to read the plots. These kind of trade-off plots are standard to compare models in control generation literature, e.g. see Figure 2 in RAD, or Figure 5 in DExperts. In our work, to compare RAD with ARM, we take several beta coefficients and plot the trade-off lines. Then, we show that two lines are close to each other: there is a little gap between the two plots in most of the figures, which allows us to empirically conclude that ARM indeed well approximates RAD.

Further, is it significant that this third parameter doesn't let ARM distilled version's perplexity go past the mid-30's?)

If we continue to increase the control coefficient $\beta$, we expect the perplexity of ARM to continue to grow (the same for RAD).

Definition of D_f: can the authors please provide for me some motivation for why we want to use the same reward y regardless of how long the prefix x is?

We agree that the setup of RAD might be not optimal and future work might improve on the particular choice of using the data. Nevertheless, we analyze the training objective of RAD and observe that in order to obtain optimal loss for the training dataset, the model would learn to predict the expected future response (Equation 5). We hope this provides more intuition to the RAD approach.

the paper should have every single equation numbered

Thank you for this suggestion, we added numbering for all the equations.

Are equations (4), (5), and (6) all equal? As in, are they just versions of the same exact same equations?

The equations you referred to are indeed versions of the same equation, we clarify this in the revision.

If the authors want this paper to be self-contained, perhaps they can add a concrete example, with a given small prefix, to illustrate what goes on in Section 2.1: RAD Training

Thank you for this suggestion. We will add a figure with an example for the final version.

Thank you for the detailed reply, and for updating the draft. Also, thanks for adding the clarifications I mentioned (e.g. the changes made to Section 3.1, Analysis of RAD--I think this makes the mathematical motivation clearer; also, specifying that there are two different training paradigms considered in the paper, distilling RAD into "low-rank RAD" (i.e. ARM), versus training a low-rank RAD (i.e. an ARM) from scratch). And I appreciate the explanation about estimating the rank, I'm satisfied there.

Further questions about the experimental results:

1. I appreciate the refactoring of the graphs, they're visually simpler to parse now, thanks. However, Figure 4 is never referenced in the main text and so it isn't immediately clear how it supports the overall story, can you please add a reference to it somewhere in the PDF to explain how it fits into the overall story? (This is already done well with Figure 3)
2. Can you please help me understand how the trade-off plots (e.g. Figures 3 and 4) indeed show positive results? For example, take Figure 4, the second plot, for Average Perplexity versus Negative Prompt Positive Rate. I agree that visually there is little gap between RAD (the baseline) and ARM Distill and ARM resp. only (i.e. the train from scratch model). However, if you look closer at the visual distance (i.e. the actual numerical differences) between the black RAD and the orange/blue curves, it looks like the distance can be as high as $0.05$, which corresponds to an approximately $0.05/0.4 = 12.5\\%$ performance difference. I guess my skepticism/misunderstanding is: suppose that we don't have any sort of heuristics that say that "this high rank reward matrix is essentially just a low rank matrix." Then what's stopping us from distilling the high rank matrix into a low-rank matrix and losing 12.5% performance? I.e. how is the rank analysis actually useful, if it leads to any performance gap at all, let alone a 12.5% performance decrease, which seems very undesirable? Unless if there is some application where the speed up from the paper's method is sufficiently valuable as to offset a 12.5% performance decrease, but that's not my understanding here. (Alternatively, please let me know if I'm reading these graphs incorrectly.)

Thanks again for answering my questions!

Additionally, the AC has indicated that they would like answers to these questions, I will provide them here:

1. Why not a lower score?--The ARM distillation results are compelling because they obtain better results than the original RAD parameterization, in addition to inference runtime speed up. As a result, these results would be valuable to the community.
2. Why not a higher score?--From the experimental results, it seems like training ARM from scratch can actually degrade downstream performance, significantly. These experimental results seem to contradict a central theoretical claim of the paper, that all the extra rank capacity is not needed in order to obtain the same level of downstream performance.

I would like to thank the authors for answering my questions. My primary concern in my original review was the clarity of the graphs, and that has been addressed. Accordingly, I have raised raised my soundness score (1-->2).

However, now that I have been able to more thoroughly examine the experimental results, I am more convinced that distillation is very valuable, but I am not entirely convinced that training ARM from scratch is valuable. In the abstract, the authors claim that "our low-rank reward model performs on par with the more flexible RAD parametrization". But as I noted earlier, training ARM from scratch tends to underperform RAD significantly (approximately 12% on the given metrics). In contrast, distilling ARM from RAD tends to perform better than RAD, and I consider this a valuable result.

I understand that the authors can't upload a revised paper version at this point in the rebuttal period. However, would the authors be willing to do one of the following:

1. Dampen the claims, especially in the abstract (e.g. changing the wording to the following-- "distilling a low-rank reward model from the more flexible RAD parametrization performs on par with the teacher RAD parametrization", or something better written than that which conveys that it's not training from scratch that wins, but rather the distillation version)
2. OR--convince me that I am mistaken, and the training from scratch version of ARM indeed gives results on par with the more flexible RAD parametrization.

Thanks!

Thanks for your reply, and thanks for answering my questions during the rebuttal period. I’m raising my score (3–>5) with the expectation that, if the paper is accepted, then it will clearly state in the abstract and the intro that distilling into low rank is the method that performs on par or better.

We thank the reviewer for their feedback. We are happy they like our analysis and that they find our method and experiments clear!

They don't clearly explain how their method is that different from others which predict rewards auto regressively. e.g. why is this approach different from others like Liu et al. and Krause et al. which effectively do the same thing?

We agree that we should better highlight the difference between our method and work of Liu et al. (DExperts) and Krause et al. (GeDi). We look closer into parametrization of the output layer and propose a novel parametrization of the output layer to decouple the prefix score from marginal scores of next tokens (Eq. 6), instead of predicting everything in one go; and for this parametrization, we propose a regularization (Eq. 11), which makes it easier for the model to abstain.

Additinally, DExperts and GeDi are trained with the language modeling objective, while our method follows a reward modeling approach. As Deng et al. (RAD) show, the reward modeling approach leads to better fluency and control, which we also highlight in the Figure 9 in Appendix (see updated pdf).

We clarify this in the updated version.

Is the W matrix in WE, meant to be lower rank than h?

We don’t mean W to learn the low rank structure. In our setup, W is needed to support the multi-task objective of RAD (section 4.2), where we introduce a separate $W_i$ for each of the toxicity types.

Why MSE loss and not binary cross entropy with soft labels?

We agree that cross-entropy is a more straightforward choice. We decided to use squared loss to closely follow the RAD approach. We add an additional ablation experiment using binary cross-entropy loss (Appendix: section F.2), where we observe that ARM trained with cross-entropy loss has slightly worse performance for the detoxification task. Additionally, having the squared loss allows us to analyze the training objective from a simple weighted average perspective (see Equation 5)

We hope our response addresses your concerns.

We thank the reviewer for their feedback. We are happy to hear that you find our analysis interesting!

Could you comment on why the perplexity of ARM is higher in general for both the evaluations? Could adding an extra regularization term on natural text improve the perplexity?

We observe that ARM trained with distillation loss more closely matches the performance of RAD or even outperforms RAD for sentiment control task, while ARM trained on original responses is less fluent. One clear difference is that when training from data, we will see short contexts multiple times with different reward responses and must implicitly converge to their average, while in distillation, the teacher already performs this compression and provides a single deterministic target $\hat{r}(v|x)$ for every context $(x,v)$. We conjecture that this may lead to better-trained distilled models. We agree with you that it is possible that extra tuning or regularization would improve ARM trained on original responses. Thank you for the suggestion to try using MLP to non-linearly process the reward scores. We observe that MLP parametrization performs on par with the more simple linear parametrization (see Appendix F.2.2 for the ablation experiment).

Dear authors,

Thanks for the revisions and responses!

I think the premise of improving efficiency in controlled generation is shared with many other works and the q-function view of the decoder (i.e., predicting the value of the next token for all alphabet in one forward pass) is shared between many other works that needs more substantiation beyond what is done currently:

• Mudgal, S., Lee, J., Ganapathy, H., Li, Y., Wang, T., Huang, Y., Chen, Z., Cheng, H.T., Collins, M., Strohman, T. and Chen, J., 2023. Controlled decoding from language models. arXiv preprint arXiv:2310.17022.

• Han, S., Shenfeld, I., Srivastava, A., Kim, Y. and Agrawal, P., 2024. Value Augmented Sampling for Language Model Alignment and Personalization. arXiv preprint arXiv:2405.06639.

• Chakraborty, S., Ghosal, S.S., Yin, M., Manocha, D., Wang, M., Bedi, A.S. and Huang, F., 2024. Transfer Q Star: Principled Decoding for LLM Alignment. arXiv preprint arXiv:2405.20495.

Also, there is work on parameter efficient learning of rewards and policies that seems related:

• Sidahmed, H., Phatale, S., Hutcheson, A., Lin, Z., Chen, Z., Yu, Z., Jin, J., Komarytsia, R., Ahlheim, C., Zhu, Y. and Chaudhary, S., 2024. PERL: Parameter Efficient Reinforcement Learning from Human Feedback. arXiv preprint arXiv:2403.10704.

Best,
AC