Risk-Averse Fine-tuning of Large Language Models

Thank you appreciating and recognizing the strengths of our work. Please find additional details of interest below.

Computational resource requirement. The computational resource requirements are mentioned in Appendix E.1.

Computational complexity. Our algorithm has the same computational complexity as that of RLHF during the first n iterations. Once the soft risk scheduling kicks in, our algorithm introduces an additional best case computational complexity of $O(B + \alpha'\log(B))$, where $B$ is the batch size and $\alpha'$ is the risk-level that decreases across iterations. The space complexity remains the same as that of RLHF.

Contribution. Our work is the first to introduce a nuanced understanding of `risk' in the context of LLM content generation, going beyond Greenberg et al. [2022]'s work. Greenberg et al. [2022] proposed soft risk scheduling to make policies learned using policy gradients risk-averse. Presented below are our contributions:

1. We implement CVaR in conjunction with a regularized reinforcement learning objective (reward + KL term). Greenberg et al. [2022] work only with the plain reward. We choose to work with regularized reward for two reasons: I. We want to measure risk in generations accounting for both the performance on actual environment reward and the quality of language generation measured by the KL-Divergence with respect to the reference policy. II. Our said choice makes our proposed algorithm downward compatible to the existing RLHF implementations.

2. We implement CVaR in the Actor-Critic setting, as opposed to policy gradients, with an aim to learn a complex parameterized policy (LLM) with an extremely large action space.

3. Beyond the immediate application of creating safer LLMs, this work contributes to the broader field of machine learning by demonstrating how risk measures like CVaR can be integrated into the training process of complex models like LLMs. Our work additionally establishes a groundwork for exploring additional risk measures and criteria, such as the Entropic Value at Risk (EVaR), in the context of LLM safety and uncertainty quantification.

Adaptability. In the scenario that the input prompt distribution shifts from the distribution our RA-RLHF models are finetuned on, our models can be further trained on the new prompt data using our RA-RLHF algorithm. As noted in our submission (Sec. 5.1 and 5.2), RA-RLHF ensures safe generations without hampering the quality of language generation. We believe this behavior would hold true under retraining for input data distribution shifts.

Validation. In our work, we studied inclusion of safety in LLMs using three datasets - IMDB, Jigsaw and RealToxicityPrompts - over GPT-2 (117M parameters) and GPT-J (6B parameters). Most of the related works in the LLM safety space work with IMDB and RealToxicityPrompts datasets with models ranging upto GPT2-L (762M parameters) as used in DExperts and Quark. RECT, another related work, studies only the RealToxicityPrompts dataset with GPT-2, GPT-2 XL (1.5B) and GPT-3 (175 B). In the light of this information, we believe our experiments are comprehensive in terms of the number of models and datasets studied.

Thank you for your insightful questions and comments. Before we begin answering the questions, we wanted to highlight that in our work, we studied inclusion of safety in LLMs using three datasets - IMDB, Jigsaw and RealToxicityPrompts - over GPT-2 (117M parameters) and GPT-J (6B parameters).

Q1. Lack of clarity. I wonder how their method involves human feedback. Reading through Algorithm 1, it is unclear how human feedback is used and how they obtain human feedback. I also wonder how they get their reward function. Do they use any annotated data to train the reward function? Are the reward functions used in training and testing the same?

A1. Human feedback is used to learn the score/reward models used for finetuning in our work. For the experiments with toxicity datasets, we use the toxicity score returned by unitary/toxic-bert LLM as our reward. For the positive sentiment generation task using IMDB dataset, we use the sentiment score returned by lvwerra/distilbert-imdb LLM as the reward. These rewards are used in line number 7 of our algorithm. We mention the use of existing reward models, as is the case in most related works, in line numbers 155-156 in Sec. 3. We have mentioned the reward function details in Appendix E.1 where we provide model and compute details. We can, however, make this more clear by adding this information in Sec. 4.1.

Q2. Lack of solid evaluation. The proposed method is evaluated on reward and perplexity. However, it is not clear how the reward function and perplexity are calculated and which models are used. In previous work [1], the creativity of LLM generations is also evaluated and human evaluation is an important method to verify the generations' quality. However, this paper has provided neither evaluation.

A2. Score/Reward Metric. For the experiments with toxicity datasets, we use the toxicity score returned by unitary/toxic-bert LLM as our metric. The LLM unitary/toxic-bert is trained to classify toxic comments on 3 Jigsaw challenges: Toxic comment classification, Unintended Bias in Toxic comments, and Multilingual toxic comment classification. Thus, the results included in the paper are indeed toxicity scores. For the positive sentiment generation task using IMDB dataset, we use the sentiment score returned by lvwerra/distilbert-imdb LLM. The LLM lvwerra/distilbert-imdb is trained to score positive sentiment present in a given piece of text, thus, making the returned score an appropriate metric for the task. These are also the models used to obtain reward function during training. We have mentioned these details in Appendix E.1 where we provide model and compute details.

Perplexity Metric. Perplexity is calculated to assess "how likely is a piece of text to be generated by a model", mathematically evaluated as $\text{PP}(W) = 2^{-\frac{1}{N} \sum_{i=1}^N \log_2 P(w_i|w_1, \ldots, w_{i-1})}$. Here, PP(W) is the perplexity of the model on the given text W. We keep W fixed across models. We choose positive prompts and completions from test dataset to form W to capture how positive/non-toxic the models are. N is the total number of words in the text. $P(w_i|w_1, ..., w_{i-1})$ is the probability assigned by the model to the $i$-th word given the preceding words. Perplexity calculation code is included in Appendix G. We can expand the discussion on perplexity in Appendix G as above. We can also briefly mention it in Sec. 5.2.

Creativity of LLM generation. We have now included results on text diversity measuring metric Distinct-n (Dist-n) introduced in DExperts [1]. The results are included over the existing algorithms along with the three new baselines (DExperts, Quark [2] and Prompted GPT-2) as described in the attached author rebuttal pdf. We observe that across datasets, models returned by our proposed algorithm RA-RLHF enjoy the best performance in terms of safety scores while maintaining text coherence and diversity.

References

[1] Alisa Liu, Maarten Sap, Ximing Lu, Swabha Swayamdipta, Chandra Bhagavatula, Noah A. Smith, and Yejin Choi. 2021. DExperts: Decoding-Time Controlled 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 Language Processing (Volume 1: Long Papers), pages 6691–6706, Online. Association for Computational Linguistics.

[2] Lu, Ximing, Sean Welleck, Jack Hessel, Liwei Jiang, Lianhui Qin, Peter West, Prithviraj Ammanabrolu, and Yejin Choi. "Quark: Controllable text generation with reinforced unlearning." Advances in neural information processing systems 35 (2022): 27591-27609.

Thank you for your insightful questions and comments.

C1. The approach primarily adapts the CeSoR algorithm by Greenberg et al. [2022]...

A1. Our work is the first to introduce a nuanced understanding of `risk' in the context of LLM content generation, going beyond Greenberg et al. [2022]'s work. Greenberg et al. [2022] proposed soft risk scheduling to make policies learned using policy gradients risk-averse. Presented below are our contributions:

1. We implement CVaR in conjunction with a regularized reinforcement learning objective (reward + KL term). Greenberg et al. [2022] work only with the plain reward. We choose to work with regularized reward for two reasons: I. We want to measure risk in generations accounting for both the performance on actual environment reward and the quality of language generation measured by the KL-Divergence with respect to the reference policy. II. Our said choice makes our proposed algorithm downward compatible to the existing RLHF implementations.

2. We implement CVaR in the Actor-Critic setting, as opposed to policy gradients, with an aim to learn a complex parameterized policy (LLM) with an extremely large action space.

3. Beyond the immediate application of creating safer LLMs, this work contributes to the broader field of machine learning by demonstrating how risk measures like CVaR can be integrated into the training process of complex models like LLMs. Our work additionally establishes a groundwork for exploring additional risk measures and criteria, such as the Entropic Value at Risk (EVaR), in the context of LLM safety and uncertainty quantification.

C2. Insufficient Comparison with Other Baselines...

A2. Thank you pointing these out. We have now included results over these in the attached author rebuttal pdf. We could not add comparisons with RECT as the model checkpoints are not publicly available for inference. However, we will be sure to mention the work in our related work section (Sec. 2). We observe that across datasets, models returned by our proposed algorithm RA-RLHF enjoy the best performance in terms of safety scores while maintaining text coherence and diversity.

C3. The paper does not compare its approach with simpler prompt-based methods...

A3. Thank you pointing this out. We have now added results over a Prompt + GPT-2 baseline in the attached author rebuttal pdf. We in-context prompt GPT-2 to generate safe text. At inference time, we add the prompt "Generate positive sentiment" to the IMDB prompts. For the toxicity datasets (Jigsaw and RealToxicityPrompts), we add the prompt "Generate non-toxic text". Generations from the prompted model enjoy only a marginal improvement in terms of safety scores over vanilla GPT-2.

Thank you for pointing out the missing references. We will be sure to mention these in our related work section (Sec. 2).

Thank you for your insightful comments and questions.

C1. The experimental results are limited....

A2C1. In our work, we studied inclusion of safety in LLMs using three datasets - IMDB, Jigsaw and RealToxicityPrompts - over GPT-2 (117M parameters) and GPT-J (6B parameters). Most of the related works in the LLM safety space work with IMDB and RealToxicityPrompts datasets with models ranging upto GPT2-L (762M parameters) as used in DExperts and Quark. RECT, another related work, studies only the RealToxicityPrompts dataset with GPT-2, GPT-2 XL (1.5B) and GPT-3 (175 B). In the light of this information, we believe our experiments are comprehensive in terms of the number of models and datasets studied.

C2. The improvements over the baselines, especially RLHF, seem somewhat small...

A2C2. Metric. For the experiments with toxicity datasets, we use the toxicity score returned by unitary/toxic-bert LLM as our metric. The LLM unitary/toxic-bert is trained to classify toxic comments on 3 Jigsaw challenges: toxic comment classification, unintended bias in toxic comments, and multilingual toxic comment classification. Thus, the results included in the paper are indeed toxicity scores. For the positive sentiment generation task using IMDB dataset, we use the sentiment score returned by lvwerra/distilbert-imdb LLM. The LLM lvwerra/distilbert-imdb is trained to score positive sentiment present in a given piece of text, thus, making the returned score an appropriate metric for the task.

GPT-J. We could not add the results for GPT-J in the main paper because of the space constraints. We can add these in the main body in the final version of the paper. The results reported in our work, especially on the tail prompts, across datasets and LLM types are amongst the largest margin improvements reported in the related work (DExperts [1], Quark [2]).

C3. Overall, contribution seems limited...

A2C3. Our work is the first to introduce a nuanced understanding of risk in the context of LLM content generation, going beyond Greenberg et al. [2022]'s work. Greenberg et al. [2022] proposed soft risk scheduling to make policies learned using policy gradients risk-averse. Presented below are our contributions:

1. We implement CVaR in conjunction with a regularized reinforcement learning objective (reward + KL term). Greenberg et al. [2022] work only with the plain reward. We choose to work with regularized reward for two reasons: I. We want to measure risk in generations accounting for both the performance on the actual environment reward and the quality of language generation measured by KL-Divergence with respect to the reference policy. II. Our said choice makes our proposed algorithm downward compatible to the existing RLHF implementations.

2. We implement CVaR in the Actor-Critic setting, as opposed to policy gradients, with an aim to learn a complex parameterized policy (LLM) with an extremely large action space.

3. Beyond the immediate application of creating safer LLMs, this work contributes to the broader field of machine learning by demonstrating how risk measures like CVaR can be integrated into the training process of complex models like LLMs. Our work additionally establishes a groundwork for exploring additional risk measures and criteria, such as the Entropic Value at Risk (EVaR), in the context of LLM safety and uncertainty quantification.

Q1. What is the computational efficiency of the proposed method compared to the baselines?

A1. Our algorithm has the same computational complexity as that of RLHF during the first n iterations. Once the soft risk scheduling kicks in, our algorithm introduces an additional best case computational complexity of $O(B + \alpha'\log(B))$, where $B$ is the batch size and $\alpha'$ is the risk-level that decreases across iterations. The space complexity remains the same as that of RLHF.

Q2. How does proposed method performs against methods such as DExperts (Liu et al., 2021)...

A2. We have now included a comparison with three new baselines (DExperts, Quark, and Prompted Base Model) in the attached author rebuttal pdf. We did not include results on GeDi (Krause et al., 2020) as both DExperts and GeDi are decoding based methods, and DExperts reported better results over GeDi. However, we will be sure to include GeDi in the related work section (Sec. 2).

Q3. Have you considered evaluating output diversity..?

A3. Thank you for the suggestion. We have now included results over output diversity in the attached author rebuttal pdf.

References

[1] Alisa Liu, Maarten Sap, Ximing Lu, Swabha Swayamdipta, Chandra Bhagavatula, Noah A. Smith, and Yejin Choi. 2021. DExperts: Decoding-Time Controlled 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 Language Processing (Volume 1: Long Papers), pages 6691–6706, Online. Association for Computational Linguistics.

[2] Lu, Ximing, Sean Welleck, Jack Hessel, Liwei Jiang, Lianhui Qin, Peter West, Prithviraj Ammanabrolu, and Yejin Choi. "Quark: Controllable text generation with reinforced unlearning." Advances in neural information processing systems 35 (2022): 27591-27609.

Thank you for your insightful comments and questions.

Q1. What steps can be taken to ensure the risk-averse fine-tuning process doesn't inadvertently reinforce existing biases in the training data? How can the model be evaluated for potential biases?

A1. Since we work in the regime of RLHF based LLM finetuning, as long as the reward model used in RA-RLHF penalizes bias, our learned models will be unbiased. This is indeed the case for our models trained to generate non-toxic text as we use unitary/toxic-bert as the reward model. unitary/toxic-bert is trained to classify toxic comments on 3 Jigsaw challenges: toxic comment classification, unintended bias in toxic comments, and multilingual toxic comment classification.

Q2. What is the long-term impact of risk-averse fine-tuning on model behavior and performance? How stable are the improvements over time?

A2. In the scenario that the input prompt distribution shifts from the distribution our RA-RLHF models are finetuned on, our models can be further trained on the new prompt data using our RA-RLHF algorithm. As noted in our submission (Sec. 5.1 and 5.2), RA-RLHF ensures safe generations without hampering the quality of language generation. We believe this behavior would hold true under retraining for input data distribution shifts.

Please let us know if we understood your questions correctly. We are happy to provide further clarification if needed.