Reward-Augmented Decoding: Efficient Controlled Text Generation With a Unidirectional Reward Model
Authors: Haikang Deng
Colin Raffel
Venue: EMNLP 2023 Main (2023)
Keywords: controllable text generation
reward modeling
weighted decoding

Abstract: While large language models have proven effective in a huge range of downstream applications, they often generate text that is problematic or lacks a desired attribute. In this paper, we introduce Reward-Augmented Decoding (RAD), a text generation procedure that uses a small unidirectional reward model to encourage a language model to generate text that has certain properties. Specifically, RAD uses the reward model to score generations as they are produced and rescales sampling probabilities to favor high-reward tokens. By using a unidirectional reward model, RAD can cache activations from prior generation steps to decrease computational overhead. Through experiments on generating non-toxic and sentiment-controlled text, we demonstrate that RAD performs best among methods that change only the generation procedure and matches the performance of state-of-the-art methods that involve re-training the language model. We further validate that RAD is effective on very large language models while incurring a minimal computational overhead.

Reward-Augmented Decoding: Efficient Controlled Text Generation
With Unidirectional Reward Model
Haikang Deng
Colin Raffel
UNC-Chapel Hill frankdenghaikang@@aall.cco com
University of Toronto, Vector Institute craffel craffel@gmail com
Abstract
Reward Model
While large language models have proven ef- fective in huge range of downstream appli- cations, they often generate text that is prob- lematic lacks desired attribute. In this paper, we introduce Beward-Augment coding (RAD), text generation procedure that small unidirectional reward model to courage a language model to generate text that has certain properties. Specifically, RAD uses the reward model to score generations as they are produced and rescales sampling probabil- ities favor high-reward tokens. By using unidirectional reward model, RAD can cache activations from prior generation steps to de- crease computational overhead. Through exper- iments on generating non-toxic and sentiment- controlled text, we demonstrate that RAD per- forms best among methods that change only the generation procedure and matches the perfor- mance of -feth--rtt methods that involve re-training the language model. We further val- idate that RAD is effective on very large lan- guage models while incurring minimal com- putational overhead.
Uncle
Bob
loves
loves
Bob t
has
De-
Language Model
+
+
Uncle Bob
Figure 1: Newaadd-Aggmont Decoding (RAD). RAD steers a language model towards generating text that is assigned high reward by an auxiliary reward model. Blue/red boxes in the reward model correspond to cached/newly computed hidden states.
et al., 2019; Gehman et al., 2020), which motivates the need for techniques that enable controllable text generation. Such techniques can be seen as providing a means to condition text generation on desired attribute.
A straightforward way to control the text gener- ated by an LLM is to perform additional training on data that has desired properties (Gururangan et al., 2020). Alternatively, LLM can be trained with "control codes" (Keskar al., 2019; Lu al., 2022) that indicate text characteristics and can be used to induce the LLM to generate content with those characteristics. If available, annotated human preferences can be used to train a reward model that is then used to train a language model with reinforcement learning (Ouyang et al., 2022). A drawback of these methods is that they degrade performance on text that is different from the data used for additional training. Besides, work done to control one language model cannot be reused to control another language model. Moreover, the additional training cost can be prohibitively expen- sive, especially for very large models.
1 Introduction
Large language models (LLMs, Rae et al., 2021; Hoffmann et al., 2022; Scao et al., 2022; Touvron al., 2023) are seeing widespread adoption thanks the fact that they can perform many language tasks and generate coherent long-form text. As LLMs are deployed in situations where they in- teract with humans, it can be beneficial to control the language model so that it generates text with certain properties (Sudhakar et al., 2019) for ex- ample, we might desire generations that are unbi- ased, non-toxic, and helpful. In addition, we may want models to output text with specific proper- ties, such as having positive sentiment, certain writing style, etc. Typically, L.Ms pre-trained on uncurated lcrge-scale text corpora can generate text that does not have these desired attributes (Wallace
One way to avoid the cost and shortcomings

---

of additional training is instead modify the de- coding procedure used to generale text from lan- guage model. For example, weighted decoding modifies the probabilities assigned to each token during decoding using an auxiliary model. Most weighted decoding methods (Holtzman et al., 2018; Krause et al., 2021; Liu et al., 2021; Yang and Klein, 2021) obtain attribute probability P(c[X) from separate reward model (typically smaller than the base language model) and construct class- conditional text probabilities following Bayes rule, where e is an attribute
up to 65B parameters.
Reward-Augmented Decoding
At a high level, award-augmented decoding, as shown in fig. 1, feeds intermediate candidate se- quences into a reward model that evaluates their alignment with desired attribute. Then, at each de- coding step, RAD uses the predicted reward of each candidate sequence modify the token probabili- ties output by the language model. In this section, describe these steps in detail. Refer table 2 for descriptions of the notations used in this paper.
P(X[c) oc Pc class and P(X) is the distribution over natural lan- guage sequences X. Weighted decoding only re- quires access to the next-step probabilities output by a language model, does not require expensive training, and is often modular, i.e. single reward model can be reused with many language models. Despite these benefits, weighted decoding can sig- nificantly increase the cost of decoding and often underperforms methods that involve further train- ing (See et al., 2019).
2.1 Unidirectional Reward Model
Consider using reward model to compute rewards for l candidate tokens each of m generation timesteps. If scoring each candidate token requires re re-processing the entire generated sequence up to the current timestep, the reward model would need to process O(km²) tokens, which could be pro- hibitively expensive. To address these issues, we use unidirectional reward model, specifically a Transformer decoder with causal masking (Liu et al., 2018; Radford et al., 2018). In unidirec- tional model with causal masking, previously com- puted representations remain unchanged when new tokens are appended, so at each generation timestep the reward model only needs to compute the repre- sentation of the newly added token. This reduces computational costs to Okk)).
In this paper, we close the gap between weighted decoding and re-training by introducing reward- augmented decoding (RAD), an efficient, effective, and modular weighted decoding method that steers text generation based on the reward returned by attribute-specifie reward model. In particular, RAD uses a unidirectional reward model trained to output a reward representing how well given sequence aligns with a desired attribute. The uni- directionality of the reward model allows caching intermediate activations as the sequence is gener- ated, greatly decreasing computational costs. Dur- ing decoding, the tokens with the top-k highest probabilities are rescaled according to the reward model so that tokens that better reflect the desired attribute are more likely to be chosen as the next generated token.
In this work, the reward model is modi- fied pre-trained decoder-only Transformer (GPT- small (Radford al., 2019a) in all of our ex- periments) fine-tuned on text annotated with the amount of the target attribute present. We use a cumulative squared error loss that takes a weighted mean of each prefix 's loss:
L(r El= t(r: A) S 1(1+1)
S
To validate RAD's effectiveness, we evaluate on standard detoxification and sentiment-cootrollled generation tasks, showing that steers text genera- tion towards a desired attribute without sacrificing much diversity and fluency. We ultimately find that RAD outperforms other weighted decoding meth- ods and achieves results comparable methods that involve additional training. We further validate RAD in real-world large-scale setting by show- ing it is effective and introduces minimal computa- tional overhead when applied to the LLaMA (Tou- vron et al., 2023) family of language models with
where r is the reward model's prediction at gen- eration timestep i, f c [0, 1] is the ground-truth reward value, and / is the generation length. The cumulative loss encourages the reward model output the correct reward for every prefix of the text sequence in order to capture both current and future alignment of a candidate sequence with the desired attribute.
2.2 Weighted decoding
RAD utilizes top-k sampling (Fan et al., 2018; Holtzman et al., 2018; Radford et al., 201 9b) and

---

Algorithm Reward-AAggennte Decoding
Input So
neural network language model (outputs logits)
9x
neural network reward model (outputs reward score)
X
generation prefix
Tt none
2: while 0t + < EDS > do
3:
Wt - topk(fo(X))
I/ get top-k tokens (indices), wy E N
4: Zi - fo(X) [w.]
// get top-k k token logits, Z₁ E RR
X; We,1
5:
Pi + 9X
// compute rewards, P. E [0, 1
WiWL
6:
Pt - softmax(2: + 3p,)
// compute reweighted distribution
7:
It ~ Categorical(p))
8:
X t (X;
// append new sample
Output generated text X steered towards higher rewards
re-weights the probabilities of the tokens with the top-k highest probabilities based on each candi- date's reward score. Specifically, timestep i, re- re-weighting is done by computing
modify LM probabilities in one pass: Rectifica- tion (Cao et al., 2023), which adjusts LM prob- abilities proportional the risk of resulting in a toxic generation; DAPT (Gururangan et al., 2020), which further trains the model on data that has the desired property; PPO (Schulman et al., 2017), which updates the LM with gradients from the re- ward model: Quark (Lu al., 2022), which per- forms parameter-efficient fine-tuning attribute- annotated data (Lester et al., 2021; Li and Liang, 2021); ); and CTRL (Keskar et al., 2019), language model trained to condition on control codes. Un- less otherwise mentioned, we report results directly from Liu et al. (2021) and Lu et al. (2022), which can be consulted for further baseline details.
softmax(z; + (p.))
where 2, R* are top-k largest logits output by the language model's at output timestep scaling hyperparameter (with higher correspond- ing to more intense steering), and P. E [0, 1] are the reward values for the k sequences corre- sponding to appending each of the top-k tokens. Adding Sp, and renormalizing with softmax is proportional to reweighting the top-k probabilities by c.0t, Consequently, RAD effectively rescales probabilities of the top-k tokens in accordance with their relative difference in reward. Algorithm 1 pro- vides an overview of the decoding process.
R is
3.1 Detoxification
Experimental Setup. We closely follow past work (Liu et al., 2021) and use RAD to detox- ify generations from GPT-2 Large (Radford et al., 2019a) after conditioning on prompts from the Re- alToxicityPrompp (Gehman et al., 2020) dataset. For our reward model, we fine-tune GPT-2 Small on 2M human-amnan annootaed comments with continu- ous labels between and from the Jigsaw Unin- tended Bias in Toxicity Classification dataset.' We report RAD's performance with different values k (used in top-k sampling) and 3 (used for adjusting weighted decoding).
3 Experiments
We now evaluate RAD's performance in two stan- dard settings: Preventing language models from generating toxic text (Wallace al., 2019; Gehman al., 2020) and controlling the sentiment of gener- ated text (Li et al., 2018; Sudhakar et al., 2019).
Baselines In both settings, we consider the same set of baselines as Liu et al. (2021), namely: the performance of the base language model itself with- out any interventions: PPLM (Pascual al., 2021). which uses bag-of-word classifier update LM hidden states during decoding; GeDi (Krause et al., 2021) and DExperts (Liu et al., 2021), which use signals from auxiliary language models to
Evaluation Metries. For every prompt, sam- ple 25 continuations, each containing up 20 new tokens. As in Liu al. (2021), we measure the Av-
hhttps://bit. hnttps://bit.ly/44cca 1y/43AAdC

---

Average Max Toxicity
0.4
RAD QPT2
Neutral Prompt Positive Rate
0.3
0.8
Rectfication
DEsperts
DAPT
RAD k
0.1
PPO
Quark
0.5
PPLM
10
20
30
40
50
60
CTRL
Perplexity
GeCi
DExports
Figure RAD outperforms all weighted decoding meth- ods (round points in the graph) and matches methods that involve additional training.
Negative Prompt
Positive Rate
0.5
PPO
0.3
erage Max Toxicity, i.e. the expected maximum tox- icity score of the 25 continuations evaluated by the Perspective API and the Toxic Rate, i.c. the prob- ability that at least one out of 25 continuations is toxic (Perspective API toxicity score > 0.5). Since the perspective API changes over time (Pozzobon al., 2023), we recomputed the scores for all base- line methods. We also measure the Diversity as the number of distinct bigrams and trigrams nor- malized by the length text (Li et al., 2016) and the Fluency as the perplexity assigned to the con- tinuation by GPT-2-XL conditioned on the prompt. In general, good method should reduce toxicity while preserving fluency and diversity.
0.1
10
30
40
50
100
200
Perplexity
Figure 3: RAD achieves the highest positive rate for negative prompts and outperforms all weighted decod- ing methods.
ity and SST-2 (Socher et al., 2013).
Evaluation Metrics. We sample 25 continu- ations for each prompt and compute the aver- age Positive Rate measured by HuggingFace text- classification pipeline* (a DistilBERT model fine- tuned on SST-2). We also report the Diversity and Fluency as introduced above.
Results. As shown fig. and table 4 (ap- pendix), RAD demonstrates a favorable trade-off between toxicity and fluency without significantly sacrificing diversity, ultimately outperforming all weighted decoding methods and matching the per- formance of methods that involve additional train- ing. Moreover, RAD achieves the lowest Average Max Toxicity of any method. Our results further demonstrate that RAD provides an intuitive means to effectively trade-off toxicity and fluency by tun- ing 3.
Results. As seen in fig. and table 5 (appendix), RAD attains better fluency/positivity trade-off (when conditioning on negative or neutral prompts) than any other weighted decoding method and achieves comparable performance to the state-of- the-art methods involving training (Quark and PPO), which both make use of the evaluation model (DistilBERT model fine-tuned on SST-2) during training. Tuning 3 effectively trades off fluency and alignment, again enabling RAD produce the best attribute scores. Figure 4 (appendix) visual- izes RAD's steering process when prompted with negative input.
3.2 Sentiment-Controlee Generation
Experimental Setup. Following past work (Li et al., 2018; Sudhakar et al., 2019; Liu et al., 2021), we use RAD to steer GPT-2 Large's generation to be either positive/negative in sentiment when prompted with negative/positive or neutral prompts. Specifically, evaluate 2.5K negative, 5K neu- tral, and 2.5K positive prompts from Web- Text (Gokaslan and Cohen, 2019). For RAD's re- ward model, we fine-tune GPT-2 Small on millions product and movie reviews from Amazon Polar-
3.3 Scaling the Language Model
In all prior experiments, we followed past work and considered using GPT-2 Large the base language model. Recent LLMs have dramatically more pa- rameters (and dramatically better performance). To test RAD in more realistic settings, we apply RAD to the state-of-the-art LLaMA models (Touvron
hhttps://bit. 1yy//F/V hhttps://bit..yy/44yye..
Chttps://bits 1

---

Prompt: The most polluted cities tend to
Decoding Cost
stop 1: be
Method
GPT-2 Large
LLaMA 65B
2:
developing
PPLM
4.0x
4.00x
countries
GeDi
1.9x
1.01 x
DExperts
3.0x
1.02x
Additional training
1x
1x
7: with
RAD
3.4x
1.03x
8: clean
9: energy
Table 1: Computational overhead (as relative increase in cost) for different methods for controlling text gener- ation using GPT-2 Small a reward model and GPT-2 Large LLaMA 65B the language model. "Ad- ditional training" refers to methods that train the lan- guage model and do not modify decoding (e.g. Quark, DAPT, PPO, etc.). Calculation details provided in ap- pendix C.2.
10: tachnologies
11: and
other
energy
efficiency
15: Innovations
16:
17: like
18: solar
19: panels
0.4
plying RAD to more sophisticated tasks, such as encouraging language models to follow instruc- tions (Ouyang et al., 2022).
Rewards
Figure 4: Visualization of RAD's decoding process. Each row represents single decoding step, where the area is the estimated reward distribution of the top-50 candidate sequences, and the red line indicates the se- lected token's 's reward score.
Limitations
Although RAD achieves decent performance and generalizes to other language models, two limita- tions should be considered for this work. Firstly, RAD incurs additional compute and memory allo- cation linear to k. As mentioned in section 2.1, we manage to reduce time complexity from O(km²) to (kk) by reusing previously computed repre- sentations in the decoder reward model. Yet, track- ing and copying past_ past_key_vee key_kekyles take up certain amount of GPU memory, which reduces decoding throughput. Secondly, our experiments regarding toxicity and sentiment explore only some capabili- ties of RAD. More text generation tasks should be conducted in order to form comprehensive review of RAD.
al., 2023) in the detoxification setting of sec- tion 3.1, using the same GPT-2 Small reward model. table (appendix), we show that RAD signif- icantly reduces LLaMA's toxicity while preserv- ing its diversity and fluency. In terms of compu- tational costs, we list the relative cost of different methods for controlled text generation in table 1. While RAD and other weighted decoding methods increase costs significantly when the size of the language model and reward model are similar, the additional expense of using RAD only about 3% when using LLaMA 65B as the language model and GPT-2 Small as the reward model. These re- sults confirm that RAD can effectively control text generation of tate-of-the-ane models while incur- ring negligible computational overhead.
Ethics Statement
This work centers around controllable text gener- ation, which holds significant relevance in regu- lating natural language generation. For example, the detoxification task aims to mitigate the toxicity present in texts generated by pre-trained language models. In this context, RAD offers a solution for controlling the text generation process without modifying the base language model.
4 Conclusion and Future Work
In this paper, we propose RAD, a simple weighted decoding method for controlling text generation that uses a unidirectional reward model to mini- mize computational costs. RAD outperforms prior weighted decoding methods and matches the per- formance of state-of-the-art techniques that involve additional training. When the size of the reward model is relatively small compared to the base lan- guage model, RAD incurs negligible computational overhead. In future work, we are interested in ap-
Acknovledgements
We would like to thank Derek Tam for valuable discussions. We also extend our appreciation to the

---

Perspective API team for increasing API quota on our behalf.
Ben Krause, Akhilesh Deepak Gotmare. Bryan McCann. Nitish Shirish Keskar, Shafiq Joty, Richard Socher. and Nazneen Fatema Rajani. 2021. GeDi: Gener- ative discriminator guided sequence generation. In Findings of the Association for Computational Lin- guisties: EMNLP 2021, pages
References
Punta Cana, Dominican Republic. Association for Compu- tational Linguistics.
Meng Cao, Mehdi Fatemi, Jackie CK Cheung, and Samira Shabanian. 2023, Systematic rectification of language models via dead-end analysis. In The Eleventh International Conference on Learning Rep- resentations.
Brian Lester, Rami Al-Rfou, and Noah Constant. 2021.
The power of scale for purameter-efficiet prompt tuning. Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pages 0445-0099, Online and Punta Cana, Domini- can Republic. Association for Computational Lin- guistics.
Angela Fan, Mike Lewis, and Yann Dauphin. 2018. Hierarchical neural story generation. In Proceedings of the 56th Annual Meeting of the Association for Compulational Linguisties (Volume I: Long Papers), pages 889-898, Melbourne, Australia, Association for Computational Linguistics.
Jiwei Li. Michel Galley, Chris Brockett, Jianfeng Gao. and Bill Dolan. 2016. A diversity-promottee ob- jective function for neural conversation models. In Proceedings of the 2016 Conference of the North American Chapter of the Association for Computa- rional Linguistics: Human Language Technologies, pages 110-119, San Diego, California. Association for Computational Linguistics.
Samuel Gehman, Suchin Gururangan, Maarten Sap. Yejin Choi, and Noah A. Smith. 2020. RealToxi- cityPrompis: Evaluating neural toxic degeneration in language models. In Findings of the Association for Computational Linguistics: EMNLP 2020, pages 3366, Online. Association for Computational Linguistics.
Juncen Li, Robin Jia, He He, and Perey Liang, 2018. Delete, retrieve, generate: a simple approach to senti- ment and style transfer. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume (Long Pupers), pages 1865-1874, New Orleans, Louisiana. Associa- tion for Computational Linguistics.
Aaron Gokaslan and Vanya Cohen. 2019. Open- webtext corpus. http://Skylionee// wttpiiiieee..hhuubb. gith.....u.. io/ OpenWebTextCorput
Suchin Gururangan, Ana Marasović, Swabha Swayamdipta, Kyle Lo. Iz Beltagy, Doug Downey, and Noah A. Smith. 2020. Don't t stop pretraining: Adapt language models to domains and tasks. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 8342-8360, Online. Association for Computational Linguistics.
Xiang Lisa Li and Percy Liang. 2021. Prefix-tuning: Optimizing continuous prompts for generation. In Proceedings of the 59th Annual Meeting of the Asso- ciation for Computational Linguistics and the 1 11th International Joint Conference on Natural Language Processing (Volume 1: Long Papers). pages 4582- 4597, Online. Association for Computational Lin- guistics.
Jordan Hoffmann, Sebastian Borgeaud, Arthur Men- sch, Elena Buchatskaya, Trevor Cai, Eliza Ruther- ford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al. 2022. Train- ing compute-optimal large language models. arXiv preprint arXiv:2203.1 15556.
Alisa Liu, Maarten Sap. Ximing Lu, Swabha Swayamdipta. Chandra Bhagavatula. Noah A. Smith, and Yejin Choi. 2021. DExperts: Decoding time con- trolled text generation with experts and anti-experts. In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Lan- guage Processing (Volume 1: Long Papers), pages 6691-6700 Online. Association for Computational Linguistics.
Ari Holtzman, Jan Buys, Maxwell Forbes, Antoine Bosselut, David Golub, and Yejin Choi. 2018. Learn- ing to write with cooperative discriminators. ceedings of the 56th Annual Meeting of the Associa- tion for Computational Linguisties (Volume 1: Long Papers), pages 1638-1649, Melbourne, Australia, As- sociation for Computational Linguistics.
Pro-
Peter J Liu, Mohammad Salch, Etienne Pot, Ben Goodrich, Ryan Sepassi, Lukasz Kaiser, and Noam Shazeer. 2018 Generating wikipedia by summariz- ing long sequences. In International Conference on Learning Representations.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B. Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. 2020. Scaling laws for neural language models.
Nitish Shirish Keskar, Bryan McCann, Lay R Varshney, Caiming Xiong, and Richard Socher. 2019. Ctrl: A conditional transformer language model for control- lable generation. arXiv preprint arXiv: (ate::999000559
Ximing Lu, Sean Welleck, Jack Hessel, Liwei Jiang, Lianhui Qin, Peter West, Prithviraj Ammanabrolu, and Yejin Choi. 2022. Quark: Controllable text generation with reinforced unlearning. Advances

---

neural information processing systems, 35:27591- 27609.
Teven Le Scao, Angela Fan, Christopher Akiki, EI- lie Pavlick, Suzana Ilié, Daniel Hesslow, Roman Castagné, Alexandra Sasha Luccioni, François Yvon, Matthias Gallé, al. 2022. Bloom: 176b- parameter open-access multilingual language model. arXiv preprint arXiv:2211.055: 1.00100.
Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. 2022. Training language models to follow instruc- tions with human feedback. Advances in Neural Information Processing Systems, 35:27730-277 277444.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. 2017. Proxi- mal policy optimization algorithms. arXiv preprint 1rXvv:707.06344
Damian Pascual, Beni Egressy, Clara Meister, Ryan Cotterell, and Roger Wattenhofer. 2021. A plug-and- play method for controlled text generation. In Find- ings of the Association for Computational Linguis- tics: EMNLP 2021, pages 3973-3997, Punta Cana, Dominican Republic. Association for Computational Linguisties.
Abigail See, Stephen Roller. Douwe Kiela, and Jason Weston. 2019. What makes a good conversation? how controllable attributes affect human judgments. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computa- tional Linguisties: Humun Language Technologies, Volume 1 (Long and Short Papers), pages 1702-1723, Minneapolis. Minnesota. Association for Computa- tional Linguistics.
Luiza Pozzobon, Beyza Ermis, Patrick Lewis, and Sara Hooker. 2023. On the challenges of using black- box apis for toxicity evaluation in research. arXiv preprint 230.::::30..33377.
Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang. Christopher D. Manning, Andrew Ng, and Christopher Potts. 2013. Recursive deep models for semantic compositionality over a sentiment treebank. In Proceedings of the 2013 Conference on Empiri- cal Methods in Natural Language Processing, pages 1631-1642, Seattle, Washington, USA. Association for Computational Linguistics.
Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, et al. 2018. Improving language under- standing by generative pre-training.
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. 2019a. Language models are unsupervised multitask learners.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. 2019b. Lan- guage models are unsupervised multitask learners. OpenAl blog, 1(8):
Akhilesh Sudhakar, Bhargav Upadhyay, and Arjun Ma- heswaran. 2019. "transforming" delete, retrieve, gen- erate approach for controlled text style transfer. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Lan- guage Processing (EMNLP- -ICCNP,)) pages 3269- 3279, Hong Kong, China. Association for Computa- tional Linguistics.
Jack W. Rae, Sebastian Borgeaud, Trevor Cai, Katie Millican, Jordan Hoffmann, Francis Song, John Aslanides, Sarah Henderson, Roman Ring. Susan- nah Young, Eliza Rutherford, Tom Hennigan, Ja- cob Menick, Albin Cassirer, Richard Powell, George van den Driessche, Lisa Anne Hendricks, Mari- beth Rauh, Po-Sen Huang. Amelia Glaese, Johannes Welbl, Sumanth Dathathri, Saffron Huang, Jonathan Uesato, John F. J. Mellor, Irina Higgins, Antonia Creswell, Nathan McAleese, Amy Wu, Brich Elsen, Siddhant M. Jayakumar, Elena Buchatskaya, David Budden, Esme Sutherland, Karen Simonyan, Michela Paganini, L. Sifre, Lena Martens, Xiang Lorraine Li, Adhiguna Kuncoro, Aida Nematzadeh, Elena Gribovskaya, Domenic Donato, Angeliki Lazaridou. Arthur Mensch, Jean- caan-ttttstte Lespiau, Maria Tsim- poukelli, N. K. Grigorev, Doug Fritz, Thibault Sotti- aux, Mantas Pajarskas, Tobias Pohlen, Zhitao Gong, Daniel Toyama. Cyprien de Masson d' Autume, Yujia Li, Tayfun Terzi, Vladimir Mikulik, Igor Babuschkin, Aidan Clark, Diego de Las Casas, Aurelia Guy, Chris Jones, James Bradbury. Matthew G. Johnson, Blake A. Hechtman, Laura Weidinger, lason Gabricl, William S. Isaac, Edward Lockhart, Simon Osindero, Laura Rimell, Chris Dyer. Oriol Vinyals. Kareem W. Ayoub, Jeff Stanway, L. I.. Bennett, Demis Hass- abis, Koray Kavukcuoglu, and Geoffrey Irving, 2021. Scaling language models: Methods, analysis in- sights from training gopher. ArXiv, abs/21 I 12. 1146.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet. Marie-Anne Lachaux. Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aur' clien Rodriguez, Armand Joulin, Edouard Grave, and Guillaume Lample. 2023. Llama: Open and efficient foundation language models. ArXiv, abs/2302.13971.
Eric Wallace, Shi Feng. Nikhil Kandpal, Matt Gard- ner, and Sameer Singh. 2019. Universal adversarial triggers for attacking and analyzing NLP. In Proceed- ings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th Inter- national Joint Conference on Natural Language Pro- cessing (EMNLP-IJCNV pages 2153-2162, Hong Kong, China. Association for Computational Linguis- tics.
Kevin Yang and Dan Klein. 2021. FUDGE: Controlled text generation with future discriminators. In Pro- ceedings of the 2021 Conference of the North Amer- ican Chapter of the Association for Computational Linguisties: Human Language Technologies, pages 3511-3535, Online. Association for Computational Linguistics.

---

A Notations
as dmodel Mas is often true in practice (Ka- 12 plan et al., 2020). In fact, in detoxification and sentiment-controlled generation experiments, Reex constantly below 50. Thus, it is safe to assume
Refer to table 2 for notations used in the paper.
B RAD Training Details
B.1 Detoxification
CRM
Ciorward 2N
We train a GPT-2 Small reward model on the Jigsaw Unintended Bias in Toxicity Classification dataset' for epochs. We use learning rate = 1e-5, 5, weight decay- = 0.01, and batch size = 100. The reward model achieves final squared error of 0.0147 on the test public leaderboard subset.
for both the language model and the reward model. The reward model evaluates & candidate sequences at each decoding step, which requires kCRM FLOPs in total. Assuming k: = 20, CRAD = kCRM. appendix C.I shows the estimated FLOPs per token of the reward model and various language models.
B.2 Sentiment-Controoledd Generation
C.2 Comparison
We first train the reward model Amazon Polarity for epochs, with learning rate le-5, weight decay = 0.01, and batch size = 100. We then continue to train the reward model on SST- 2 (Socher et al., 2013) for 10 epochs, with learning rate = 2c- 6.
We continue to explore the computation cost baseline methods based on the methodology in ap- pendix C.I. Define Cmeehod the additional cost a method incurs during decoding and mctthodd as the total cost in FLOPss for every token generated using the method on the specific LM during test time. In general, retraining methods (DAPT, PPO, and Quark) have Ceethod = and meCehoodXLMM CLM-
C Computational Costs
C.1 RAD
In the paper, we use GPT-2 Small (124M) as RAD's reward model and replace the Im_ heah layer with a linear layer with one output for predicting the re- ward. Following the approximations in Kaplan et al. (2020), the number of non-embedding parameters in model is approximately
PPLM updates the previous token's 's representa- tion in the LM using gradients from an attribute- specific linear discriminator and recomputes the current state probabilities. Thus, two forward passes and one backward pass of the LM are re- quired for every generated token. As backward pass has roughly twice the number of matrix mul- tiplications in forward pass (Kaplan al., 2020), PPLM incurs an additional decoding cost of CPPLM 3CLM- Thus,
N RJ 12 eoodel
where Faayer is the number of layers and dmodel is the hidden size. In addition, forward pass requires
Ctorward Rs 2N + Mea dmodel
T'PPLMXGPTT-Aage 4CGPT2-lagee 5,66G TCPPLMLLLMA-S 4CLL 515.40G
FLOPs per token, where n1x is the context token. With embedding operation costing 4dmonde and re- ward predicting costing 2dmodel, we construct the number of FLOPs needed for token in the reward model as
GeDi & DExperts take similar approach where they use two opposite discriminator/expert mod- els to produce classification probabilities and then rescale the LM probabilities. Thus, two additional forward passes of the expert model are needed. For GeDi,
CRM 6dmoodel + Cforwand
Since 6dmooed < N,
CGeDi = 2CGP12-medium = 1.21G
ORM 23 Cforward
T TCG.DixGFTA-aage CaeDi + CGPT2-lage
Dexx
= 2(1 +
)
12d
2.62G
TCGeDiSLLLMA.GSI CGeDi + CLLaMA-65b
Notice the context-dependent computational cost per token is a small portion of the total compute,
130.06G

---

Notation
Dimension
Description
k
N
number of candidate tokens consider every timestep steering amount hyperparameter
3
R
l
N
generation length of reward model training data label of reward model training data predictions generated by reward model during training reward scores predicted by the reward model at time indices of top-k tokens at time logits of top-k tokens at time
f.
[0, 1]
[0, 1J4
r
[0, 1
Pt
N
W:
I
Z;
Table 2: We use two notations and p to differentiate the reward model's output during train time and test time.
RM
Mayer
diundel
CRAD
PPO and Quark, see below) using the most up-to- date API. Queries were made between May and June 2023. As PPO and Quark directly optimize the language model with Perspective API score during training, a change in the API model would lead to different optimized model. For PPO and Quark, we adopt the values reported in Lu et al. (2022). Full results see table 4.
GPT2-small
12
768
3.40G
LM
Mayer
modee
CLM
GPT-2 Large
36
1280
1.42G
LLaMA 7B
32
4096
12.89G
LLaMA 13B
40
5120
25.17G
LLaMA 33B
60
6656
63.80G
LLaMA 65B
80
8192
128.85G
D.2 Sentiment-Controlles Generation
Table 3: Model specifications and FI.OPs per token.
The sentiment-controlledd generation results are pre- sented in table 5.
For DExperts,
D.3
Scaling the Language Model
Following previous experiments, we use nucleus sampling with - 0.9 to get raw LLaMA gener- ations on the same 10K non-toxic subset of Real- ToxicityPrompts (Gehman et al., 2020). For each model size, we apply RAD with k 20 and B from 20 500. Results shown in table 6.
DExxans 2CGPT2-large 2.83G
TCDDsppesxXGPP2-aaa CDExpent + CGPT2-lage
4.25G
TCDDsppnssxLLMMMA.. CDispens + CLAMAA655
131.68G
The performance gap between RAD on GPT- Large and RAD on LLaMA may be attributed to the difference in tokenization between the lan- guage model and the reward model. Specifically, the reward model, GPT-2 Small, shares the same tokenizer and vocabulary with GPT-2 Large, but not with LLaMA. In this way, given text sequence can be tokenized into different token combinations, which, during decoding, would mislead the reward model give distorted scores. Therefore, be- lieve smaller model from the same family of the base LM may be the best choice for RAD's reward model.
RAD with 1: 20 has
CRAD 20CGPT2 3.40G
TCRADxGPTZ-ange = CRAD + CGPT2-large
4.81G
OB = CRAD + CLLaMA-65b
132.25G
DAPT, PPO, and Quark have decoding costs the same as the underlying LM because they per- form additional training and do not modify the decoding procedure.
D Full Results
E
Generated Examples
D.1 Detoxification
Examples of detoxification and sentiment- controlled generation from each method are presented tables 7 and
Since perspective API updates its model regu- larly (Pozzobon et al., 2023), we ensure fair com- parison by evaluating model outputs (except for

---

Toxicity (Q Toxicity
Fluency (L) Perplexity
Diversity (t)
Method
Average
Toxic Rate
Dist-2
Dist-3
GPT2
0.384
0.257
11.31
0.85
0.85
PPLM
0.376
0.240
32.58
0.86
0.86
GeDi
0.242
0.055
60.03
0.84
0.83
DExperts
0.201
0.021
32.41
0.84
0.84
Rectification
0.180
0.014
25.12
0.86
0.87
DAPT
0.270
0.093
31.21
0.84
0.84
PPO
0.218
0.044
14.27
0.80
0.84
Quark
0.196
0.035
12.47
0.80
0.84
3 10
0.265
0.076
12.54
0.81
0.84
3 = 20
0.232
0.042
12.57
0.81
0.84
8 = 30
0.211
0.026
12.69
0.81
0.84
k = 20
3 50
0.183
0.014
13.06
0.81
0.84
3 100
0.148
0.005
13.7
0.81
0.83
3 200
0.114
0.002
15.93
0.79
0.81
3 300
0.099
0.001
19.97
0.76
0.78
RAD
3 10
0.069
0.84
3 20
0.233
0.035
0.84
3 30
0.022
0.84
50
8 50
0.011
17.62
0.84
3 100
0.004
0.84
8 200
0.112
01.001
23.63
0.83
8 300
0.099
0.001
32.84
0.79
0.8
Table 4: Full results of the detoxification experiment. In general, RAD can produce the least toxic outputs without sacrificing much fluency and diversity. While increasing enhances diversity and reduces toxicity by margin, it takes into account more unlikely words which hurts the output fluency.
Toward Positive
Toward Negative
% Positive Rate (?)
Fluency ())
Diversity (f)
% Positive Rate (J)
Fluency (J)
Diversity (f)
Negative prompe
Positive prompt
Perplexity
Dist-2
Dist-3
Perlexity
Dist-2
Dist-3
prompt
prompt
GPT2
0.00
50.02
11.42
0.85
0.85
99.08
50.02
11.42
0.84
0.84
PPLM
8.72
$2.68
142.1
0.86
0.85
89.74
39.05
181.7
0.87
0.86
GeDi
26.80
86.01
58.41
0.80
0.79
39.57
8.73
84.11
0.84
0.82
DExperts
36.42
94.46
25.83
0.84
0.84
35.99
3.77
45.91
0.84
0.83
DAPT
14.17
77.24
30.52
0.83
0.84
87.43
33.28
32.86
0.85
0.84
PPO
43.13
94,10
15.16
0.80
0.84
52.22
3.65
15.54
01811
0.84
19.99
86.21
12.86
0.79
0.82
73.34
17.38
13.33
0.8
0,83
20
35.24
92.71
14.6
0.78
0.82
$7.19
10.49
15.36
0.79
0.83
20
40
48.37
95.72
19.23
0.76
0.8
47.01
6.84
20.27
0.77
0.81
51.19
96.3
21.77
0.75
0.79
45.45
22.8
0.76
0.8
53.21
96.62
25.06
0.73
0.78
44.76
25.51
0.74
0.78
RAD
22.43
17.98
0.82
67,22
15.13
18.77
0.83
0.85
20
41.56
93.28
21.02
0.82
0.84
45.08
8.63
22.54
0.83
0.84
40
57.91
96.24
28.99
0.8
0.82
$2.27
5.62
30.92
0.8
0.83
63.57
97.0
39.23
0.78
0.81
28.75
40.1
0.78
0.81
Table Full results of the sentiment-controllldd generation experiment.

---

Toxicity (L)
Fluency ())
Diversity (1)
LM
Setting
Average Max Toxicity
Toxic Rate
Perplexity
Dist-2
Dist-3
Raw
L.M
0.338
0.212
12.93
0.81
0.82
8 20
0,282
0.129
13.79
0.82
0.83
3 50
0.250
0.097
14.21
0182
0.82
LLaMA 7B
8 100
0.221
0.072
15.13
0.82
0.82
200
0.185
0.048
17.40
0.80
0.81
500
0.125
0.016
29.97
0.69
0.73
Raw LM
0.336
0.204
11.81
0.81
0.82
8 = 20
0.284
0,129
12.70
0.82
0.82
3 = 50
0,253
0,103
13.00
0.82
0.82
LLaMA 13B
3 100
0.223
0.073
13.79
0.82
0.82
8 200
0.187
0.047
16.20
0.80
0.81
8 500
0.127
0.017
29.82
0.69
0.72
Raw LM
0.337
0.210
10.80
0.80
0.81
8 = 20
0.287
0.139
11.68
0.81
0.82
8 = 50
0.258
0.109
11.96
0.81
0.82
I.LaMA 33B
8 = 100
0.229
0.076
12.79
0.81
0.82
8 200
0.192
0,050
15.06
0.80
0.81
8 = 500
0,130
0.019
28.58
0.68
0.72
Raw I.
0.337
0.215
10.42
0.80
0.80
=
0,288
0.139
11.30
8 = 50
0,260
0.108
11.56
0.81
0.81
LLaMA 65B
8 100
0,232
0.081
12.27
0.81
0.81
200
0.197
0.055
14.58
0.79
0.81
8 500
0.133
0.019
28.64
0.70
0.73
Table 6: Similar to its application to GPT2-large, RAD (k = 20) reduces the toxicity of LLaMA's generated text
without greatly impacting its fluency and diversity. Its performance remains stable even as L.LaMA model increases in size.
Prompt: No matter who you'll 'Il pick you'll
GPT2
be well placed to take advantage of many of the remaining lower season attractions. which include big-city
PPLM
need a hase set of rules. position will remain stable real issue as you'll
GeDi
need the focus for this one. Personally, had good time in jokingly calling out Card
DExperts
need a versatile DHI player and Dahl will certainly help you as he's 's ideally positioned at the back of still get a great value. would only recommend this product to people who have some sort of skin likely be joining a roster of other people who remember that some characters are cuter than others, so still get the chance. I mcan, we' ve got some pretty awesome teams and some of them are
Rectification
DAPT
RAD
Table 7: Examples of detoxified generation using different methods. RAD with k 20 and B 100.
Prompt: Mcanwhile the iron
GPT2
transfer would be successful to some extent, but the heart issuc is still very dangerous and will probably have
To Pusitive
PPLM
content on the mass spectrometer has reached near limit these elements can detected," parts of both were good condition, though when was ordering them not notice that they only gathered, our new friends danced, jests were merrily spiced, and plenty of songs fired
CTRL
GeDi
DExperts
fields may not seem like the perfect areas for reclaiming coal and steel, but technology has brought mining in water hegins to leach from the old pipes, whilst the decaying furniture sheds a great deal industries, which provide a great deal of economic and social support for the British working class and a great
DAPT
RAD
Negative
PPLM
prices have collapsed to record lows and given indication they will recede the seller and take the content over half of that US bottled water 0.32 ppm versus 0.012 ppm.
CTRL
GeDi
content some vaccines already precipitously risen above acceptable limits. They'v 've even gone so far many carts loaded with rusty iron bars running of firewood dump. 80
DExperts
sinks.
DAPT
reserves have fallen precipitously. At last count, the final daily production estimate was $0,000
RAD
shortages have only been getting worse. The steel industry that is vital the economy today in serious
Table 8: Examples of sentimenn-comtrolled generation using different methods. RAD with 50 and , = 30.


Reviews:
[{"review_id": "A11XlBZPCH", "reviewer": "Reviewer_U1EG", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 3, "soundness": 4, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review! Below, we provide some insights into why RAD performs better compared to other weighted decoding methods.\n\nFirst, note the different uses of the small auxiliary model in RAD compared to other strong weighted decoding methods such as GeDi and DExperts—we use the auxiliary model as a reward model to discriminate candidate sequences at each intermediate step whereas GeDi and DExperts use the auxiliary model as a class-conditional language model to steer generation. One reason that RAD's approach works better could be that classification is an easier task as compared to class-conditional language modeling, allowing small discriminators to be more helpful than small expert LMs in guiding language generation. \n\nSecond, RAD’s reward model makes more use of label information if it is not a binary score (i.e. a fraction in [0,1]). RAD’s reward model is trained to minimize the squared error loss while the auxiliary model in GeDi and DExperts does not utilize the label score efficiently—they use label score primarily for data filtering for generative training.\n\nAnother factor that may have contributed to RAD’s superior performance is the weighted cumulative loss mentioned in section 2.1. During reward model training, by taking every prefix’s reward score into account, we encourage the reward model to converge to the correct reward score early in the sequence. In this way, the reward model is trained not only to capture the current but also the future alignment score of a generation course, which better steers the generation.\n\nWe believe RAD is exciting in that it is simple, efficient, and easily generalizable. We further highlight that connecting reward modeling and weighted decoding is a novel contribution, as is the innovation of using a causal reward model in this setting to reduce computational costs. Moreover, RAD enables a new form of modularity where a given reward model can be immediately integrated into a language model to steer its generations. Coupled with the small size of the reward model, this could enable language model users to cheaply create custom-tailored reward models that could be applied to different language models (including those behind black-box APIs).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "1. The proposed method is simple yet effective and thus should influence subsequent research in controlled text generation.\n2. The proposed method achieves a good trade-off between efficiency and effectiveness compared to previous weighted decoding and fine-tuning methods.\n3. The results of the experiment are reliable because of the extensive comparison with previous models and the recalculation of the scores of previous models based on nice attention to the differences in API.", "reasons_to_reject": "I found no major reasons for rejection, but it would be more insightful to explain why the proposed method performs better than other weighted decoding methods.", "questions_to_authors": "", "paper_topic_and_main_contributions": "This paper studies efficient controlled text generation by a lightweight external module for large language models. At each time step of text generation, the proposed method unidirectionally classifies the attributes of candidate texts and modifies the logit of the language model according to the scores. Experimental results show that the proposed weighting module achieves better performance in controlled text generation than previous methods.", "excitement": 4, "reproducibility": 4, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "l19ZKksWUM", "reviewer": "Reviewer_GRYE", "summary": "", "questions": "Where does the term $\\hat{r}$ in Section 2.2 come from?", "limitations": "", "rating": 3, "confidence": 3, "soundness": 3, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Hi, thank you for reviewing our paper! We decided to update the notations and add a pseudocode of RAD to enhance clarity.\n\n\n| Notation | Dimension | Description |\n|:-:|-|-|\n|$\\beta$ | $\\mathbb{R}^+$ | steering amount hyperparameter |\n|$l$ | $\\mathbb{N}$ | generation length of reward model training data |\n|$\\hat r$ | $[0,1]$ | label of reward model training data |\n|$\\mathbf{r}$ | $[0,1]^l$ | predictions generated by reward model during training |\n|$\\mathbf{w}_t$ | $\\mathbb{N}^k$ | indices of top-$k$ tokens at time $t$ |\n|$\\mathbf{z}_t$ | $\\mathbb{R}^k$ | logits of top-$k$ tokens at time $t$ |\n|$\\boldsymbol{\\rho}_t$ | $[0,1]^k$ | reward scores predicted by the reward model at time $t$ |\n\n**Table 1:** Updated notations. We use two different notations $\\mathbf{r}$ and $\\boldsymbol{\\rho}$ to differentiate the reward model's output in train time and test time.\n\n&nbsp;\n\n---\nAlgorithm 1:  Reward-Augmented Decoding\n---\n**Input:** $f_\\theta$, neural network language model (outputs logits) \\\n&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;$g_\\lambda$, neural network reward model \\\n&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;$X$, generation prefix\n\n1. $x_t \\leftarrow \\mathtt{none}$\n2. **while** $x_t \\ne \\mathtt{<EOS>}$ **do**\n3. &nbsp;&nbsp;&nbsp;&nbsp; \n    $\\mathbf{w}_ t \\leftarrow \\texttt{topk}(f_\\theta(X))$\n    &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp; \n    // get top-$k$ tokens (indices), $\\mathbf{w}_t \\in \\mathbb{N}^k$\n4. &nbsp;&nbsp;&nbsp;&nbsp; \n    $\\mathbf{z}_ t \\leftarrow f_\\theta(X)[\\mathbf{w}_t]$\n   &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;\n    // get top-$k$ token logits, $\\mathbf{z}_t \\in \\mathbb{R}^k$\n5. &nbsp;&nbsp;&nbsp;&nbsp;\n    $\\boldsymbol{\\rho}_ t \\leftarrow g_\\lambda \\left(\\begin{bmatrix} X;\\mathbf{w}_ {t,1} \\\\\\\\ \\vdots \\\\\\\\ X;\\mathbf{w}_ {t,k} \\end{bmatrix}\\right)$\n    &nbsp;&nbsp;&nbsp;&nbsp; \n    // compute rewards, $\\boldsymbol{\\rho}_t \\in [0,1]^k$\n6. &nbsp;&nbsp;&nbsp;&nbsp;\n    $p_t \\leftarrow \\mathtt{softmax}(\\mathbf{z}_t+\\beta \\boldsymbol{\\rho}_t)$\n    &nbsp;\n    // compute reweighted distribution\n7. &nbsp;&nbsp;&nbsp;&nbsp;\n    $x_t \\sim \\mathrm{Categorical}(p_t)$\n8. &nbsp;&nbsp;&nbsp;&nbsp;\n    $X \\leftarrow \\\\{X;x_t\\\\}$\n    &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;\n    // append new sample\n\n**Output:** generated text $X$ steered towards higher rewards\n\n---\n\n&nbsp;\n\n\nSince you mentioned that the method description is not clear, we provide a summary of section 2.1 and 2.2 below:\n- (2.1 paragraph 1) We use a small, decoder-only Transformer model, with a linear layer on top, as a reward model to guide generation. The decoder-only model’s causal masking enables a lower complexity during decoding.\n- (2.1 paragraph 2) The reward model is trained on a dataset (which can differ from the target dataset) with a weighted, cumulative squared error loss. This loss encourages the reward model to converge to the ground-truth reward early in the text sequence.\n- (2.2) At each generation timestep $t$, we consider the top-$k$ tokens as candidates; record their logits $\\mathbf{z}_t$; append candidate tokens to current text sequence to form $k$ candidate sequences; evaluate the rewards $\\boldsymbol{\\rho}_t$ of these candidate sequences; finally, reweight token probabilities with\n$\\mathtt{softmax}(\\mathbf{z}_t + \\beta \\boldsymbol{\\rho}_t)$.\n\nTo address your question about $\\hat r$, it is the label from the dataset used to train the reward model, which usually represents how well a text sequence aligns with the desired attribute.\n\nRegarding your comment about not providing interesting technical insights, we agree that reward modeling and weighted decoding are two well-studied techniques—reward model has widely been used in reinforcement learning and methods that involve data exploration, and weighted decoding is an intuitive technique backed by a Bayesian interpretation. However, prior work has neglected the idea of connecting reward modeling and weighted decoding when it comes to aligning language generation, partly due to the prohibitive complexity incurred by a bi-directional, discriminative reward model. In this paper, we draw a straight relation between reward modeling and weighted decoding by using a uni-directional reward model to augment the decoding process without re-training the base language model. We show that RAD is simple, efficient, and easily generalizable. We therefore believe its merit is self-evident. We further highlight that connecting reward modeling and weighted decoding is a novel contribution, as is the innovation of using a causal reward model in this setting to reduce computational costs. Furthermore, RAD enables a new form of modularity where a given reward model can be immediately integrated into a language model to steer its generations. Coupled with the small size of the reward model, this could enable language model users to cheaply create custom-tailored reward models that could be applied to different language models (including those behind black-box APIs).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "* The proposed method is intuitive and sound. Moreover, it is easy to implement and does not cost excessive computation resources.\n* The experiments show that the method is effective.", "reasons_to_reject": "* The description of the method is not clear. Specifically, the term $\\hat{r}$ in Section 2.2 is not clearly explained. The source of this label is questionable.\n* The contribution of the paper is thin. Essentially, this is a combination of the well-known re-ranking technique and the now popular reward models. Therefore, it does not provide interesting technical insights.", "questions_to_authors": "Where does the term $\\hat{r}$ in Section 2.2 come from?", "paper_topic_and_main_contributions": "In this paper, the authors propose to use the reward models directly in the decoding phase to avoid the additional training cost. Specifically, they follow the weighted decoding method framework, where the logits are adjusted by an additive reward value from an additional reward model. They conduct experiments on detoxification and sentiment-controlled generation tasks. The results show that the proposed method outperforms other reweighting techniques.\n\nThe highlights of this paper include:\n* The proposed method is intuitive and sound. Moreover, it is easy to implement and does not cost excessive computation resources.\n* The experiments show that the method is effective.\n\nMy concerns are as follows:\n* The description of the method is not clear. \n* The contribution of the paper is thin. ", "excitement": 3, "reproducibility": 4, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "i3u19r36zk", "reviewer": "Reviewer_Ekrw", "summary": "", "questions": "- In the experiments presented in Section 3.3 (Table 5), is k 20 or 50?\n", "limitations": "", "rating": 4, "confidence": 4, "soundness": 4, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for reviewing our paper and for your helpful suggestions for improving the quality and readability of our paper. We decided to update the notations and add a pseudocode of RAD to enhance clarity.\n\n\n| Notation | Dimension | Description |\n|:-:|-|-|\n|$\\beta$ | $\\mathbb{R}^+$ | steering amount hyperparameter |\n|$l$ | $\\mathbb{N}$ | generation length of reward model training data |\n|$\\hat r$ | $[0,1]$ | label of reward model training data |\n|$\\mathbf{r}$ | $[0,1]^l$ | predictions generated by reward model during training |\n|$\\mathbf{w}_t$ | $\\mathbb{N}^k$ | indices of top-$k$ tokens at time $t$ |\n|$\\mathbf{z}_t$ | $\\mathbb{R}^k$ | logits of top-$k$ tokens at time $t$ |\n|$\\boldsymbol{\\rho}_t$ | $[0,1]^k$ | reward scores predicted by the reward model at time $t$ |\n\n**Table 1:** Updated notations. We use two different notations $\\mathbf{r}$ and $\\boldsymbol{\\rho}$ to differentiate the reward model's output in train time and test time.\n\n&nbsp;\n\n---\nAlgorithm 1:  Reward-Augmented Decoding\n---\n**Input:** $f_\\theta$, neural network language model (outputs logits) \\\n&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;$g_\\lambda$, neural network reward model \\\n&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;$X$, generation prefix\n\n1. $x_t \\leftarrow \\mathtt{none}$\n2. **while** $x_t \\ne \\mathtt{<EOS>}$ **do**\n3. &nbsp;&nbsp;&nbsp;&nbsp; \n    $\\mathbf{w}_ t \\leftarrow \\texttt{topk}(f_\\theta(X))$\n    &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp; \n    // get top-$k$ tokens (indices), $\\mathbf{w}_t \\in \\mathbb{N}^k$\n4. &nbsp;&nbsp;&nbsp;&nbsp; \n    $\\mathbf{z}_ t \\leftarrow f_\\theta(X)[\\mathbf{w}_t]$\n   &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;\n    // get top-$k$ token logits, $\\mathbf{z}_t \\in \\mathbb{R}^k$\n5. &nbsp;&nbsp;&nbsp;&nbsp;\n    $\\boldsymbol{\\rho}_ t \\leftarrow g_\\lambda \\left(\\begin{bmatrix} X;\\mathbf{w}_ {t,1} \\\\\\\\ \\vdots \\\\\\\\ X;\\mathbf{w}_ {t,k} \\end{bmatrix}\\right)$\n    &nbsp;&nbsp;&nbsp;&nbsp; \n    // compute rewards, $\\boldsymbol{\\rho}_t \\in [0,1]^k$\n6. &nbsp;&nbsp;&nbsp;&nbsp;\n    $p_t \\leftarrow \\mathtt{softmax}(\\mathbf{z}_t+\\beta \\boldsymbol{\\rho}_t)$\n    &nbsp;\n    // compute reweighted distribution\n7. &nbsp;&nbsp;&nbsp;&nbsp;\n    $x_t \\sim \\mathrm{Categorical}(p_t)$\n8. &nbsp;&nbsp;&nbsp;&nbsp;\n    $X \\leftarrow \\\\{X;x_t\\\\}$\n    &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;\n    // append new sample\n\n**Output:** generated text $X$ steered towards higher rewards\n\n---\n\n&nbsp;\n\nAccording to your comment, we have made the following changes:\n- On line 143, we clarified that $l$ is supposed to be the **generation length** instead of the current generation timestep. \n- We further clarified that $c$ is an attribute class. $P(X)$ is the distribution over natural language sequences $X$. $P(X|c)$, the distribution over $X$ conditioned on an attribute class $c$, is proportional to $P(X) P(c|X)$. \n- We added a mention of Fig. 1 in the opening paragraph of Section 2.\n- In the equation, $\\mathbf{r} \\in [0,1]^l$ is a vector and $\\mathbf{r}_t$ is a scalar. \n\nAs per your question, Table 5 is $k$=20. We added this information to the table description.\n\nIn addition, we would like to provide more discussion on (1) why RAD is superior to other weighted decoding methods, and (2) why RAD is comparable to retraining methods.\n\n(1) **Why is RAD superior to other weighted decoding methods?**\\\n&nbsp;&nbsp;&nbsp;&nbsp;First, note the different uses of the small auxiliary model in RAD compared to other strong weighted decoding methods such as GeDi and DExperts—we use the auxiliary model as a reward model to discriminate candidate sequences at each intermediate step whereas GeDi and DExperts use the auxiliary model as a class-conditional language model to steer generation. One reason that RAD's approach works better could be that classification is an easier task as compared to class-conditional language modeling, allowing small discriminators to be more helpful than small expert LMs in guiding language generation.\\\n&nbsp;&nbsp;&nbsp;&nbsp;Second, RAD’s reward model makes more use of label information if it is not a binary score (i.e. a fraction in [0,1]). RAD’s reward model is trained to minimize the squared error loss while the auxiliary model in GeDi and DExperts does not utilize the label score efficiently—they use label score primarily for data filtering for generative training.\\\n&nbsp;&nbsp;&nbsp;&nbsp;Another factor that may have contributed to RAD’s superior performance is the weighted cumulative loss mentioned in section 2.1. During reward model training, by taking every prefix’s reward score into account, we encourage the reward model to converge to the correct reward score early in the sequence. In this way, the reward model is trained not only to capture the current but also the future alignment score of a generation course, which better steers the generation.\n\n(2) **Why is RAD comparable to retraining methods?**\\\n&nbsp;&nbsp;&nbsp;&nbsp;While we don’t have a direct answer for this questions due to the essential difference in the design of RAD and retraining-based aligning methods, we would like to point out that, as suggested by Korbak et al., (2022) [1] , Bayesian inference (as used in RAD and most weighted decoding methods), if optimal, would produce an equivalent controlling effect as KL-bounded RL retraining methods (such as reinforcement learning from human feedback). We believe that, in theory, weighted decoding methods are not inferior to retraining methods in aligning language models, and therefore it is plausible that RAD produces results comparable to that of some retraining methods.\n\n[1] Korbak et al., \"RL with KL penalties is better viewed as Bayesian inference\", in Findings of the Association for Computational Linguistics: EMNLP 2022.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "- The proposed method is superior both in performance and decoding\n  speed in comparison with other weighted decoding methods and\n  re-training methods.\n\n- The paper is relatively easy to follow.\n", "reasons_to_reject": "- It is not clear why the proposed method achieves better performance\n  (e.g., average max toxicity) in comparison with other weighted\n  decoding methods and achieves results comparable to re-training\n  methods. Discussion is required for this point.\n", "questions_to_authors": "- In the experiments presented in Section 3.3 (Table 5), is k 20 or 50?\n", "paper_topic_and_main_contributions": "The method that controls a language model so that it generates text\nwith certain properties is crucial. This paper presents a method for\nReward-Augmented Decoding (RAD), where a small unidirectional reward\nmodel is used to encourage the language model to generate text that\nhas certain properties.  By using the unidirectional reward model, the\ncaching intermediate activations is enable, which leads to an\nefficient decoding. The experimental results on standard\ndetoxification and sentiment-controlled generation tasks demonstrate\nthat the proposed method outperforms other weighted decoding methods\nand achieves results comparable to re-training methods.\n\nAlthough the proposed method is simple, the proposed method is\nsuperior both in performance and decoding speed from the experimental\nresults. The paper is relatively easy to follow.\n", "excitement": 4, "reproducibility": 4, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"Missing_References": "None", "Typos_Grammar_Style_and_Presentation_Improvements": "- l073: Add the notations for c and X\n\n- Figure 1: Add a mention to Fig. 1 in the main text\n\n- l141: Why $r_t$ is bold?\n\n- l227: fig. 2 -> Fig. 2, table 3 -> Table 4\n"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper presents a method to more effectively do controlled generation by using a reward model as a heuristic during a search-based decoding process. The method is a natural extension of multiple existing approaches and the reviews all agree for the most part that this work is both exciting and has sound methodology.
Decision: Accept-Main