Tailoring Self-Rationalizers with Multi-Reward Distillation

Thank you for your positive review and very helpful comments and suggestions! We have added our response below as well as changes to the PDF based on your feedback.

Methodologically, this multi-attribute control setup seems to be a straightforward extension of Quark, and the idea of finetuning an LM to use tags to provide control i think is pretty well-explored in other work such as [1], and doing it for multi-attribute with different tags is used as a baseline in [2]. i feel that your main contribution is not so much the new algorithm in a general sense, but rather the interesting application (with convincing human eval results) to the important and currently relevant task of making LM reasoning more reliable.

We partly agree with this! Yes, our task is tailored towards rationalization, which is coupled with accuracy and task performance. However, this is different from other controllable text generation tasks like creative text, poetry, machine translation, etc. However, we also note that adding multiple rewards for rationale generation was not a trivial addition, primarily because good quality rationales that increase a certain reward are not useful – they should be able to help improve the task accuracy, along with having high human preference. In Section 5.2, we can observe that LMs tend to overfit to a reward and hack their way to improving these metrics. We first need to select metrics that are actually important for rationalization, and at the same time, make sure that they are compatible in a way so that they can be used together. We hope this can help clarify the differences in our task setup and novelty behind our work.

Thank you for linking the above works as well. As per this suggestion, we have added them in the subsection in our related work, where we discuss reward/multi-reward based generation methods

some of the models you use to evaluate/filter the individual qualities, e.g. VERA, are way bigger than the model you're finetuning, which arguably gives you extra signal that the SFT baseline doesn't have? what was the data used for training those? on a related note, this raises some concerns that your method might not be scalable to larger LMs, unless we have e.g. an even larger version of VERA to provide supervision?

In this work, we focus majorly on how to use multiple rewards to improve the quality of rationales, both quantitatively and qualitatively. Therefore, how the rewards are actually designed are out of scope for this work. For example, for consistency, we use a t5-base reward model which is smaller than the t5-large we train. For plausibility, we use VERA (t5-11b) because the original work trained a t5-11b model on that property. In Section 5.2, we also talk about certain property metrics which depend on Web-API calls (such as factuality) which are practically infeasible - we experimented with zero-shot FLAN-T5 as a metric on factuality since we had to make a tradeoff between metric quality and inference time. Hence, say we want to train a much larger LM to self-rationalize, the reward models can still remain the same since they are a function of the property we want to measure and how good of a proxy they are for that metric, and not a function of the LM we are training.

Is there any particular reason to think additive or classic might be better than the other in any particular setting? Otherwise, it kind of seems like you're just giving yourself "multiple tries" at your benchmark, in some sense.

The choice of MaRio Classic or Additive, and the choice of the order of the rewards in the multi-reward setting is a hyperparameter setting depending on the dataset we are training for. Sometimes, certain datasets require us to first optimize one reward, before adding the others one-by-one, whereas datasets might require all rewards to be provided from the very start for the best optimization; hence, we provide the framework for both such methodologies. (Additionally, other multi-reward techniques such as the product of rewards, or filtering of rewards have been added as baselines)

(... continued below)

Thank you for your helpful comments and insightful questions! We are happy to hear that you appreciated our approach as well as human evaluations. Based on your feedback, we have updated the draft and provided the responses below.

Other datasets than QA ones should be used to show the generalization of the proposed method + How would you adapt your method for other rationalization tasks that are not QA?

We essentially phrase the multi-choice QA task to be a classification task (Section 4.1), where we generate rationales for a given model prediction (class). Therefore, this can be easily extended to any classification task. We provide results on a range of QA tasks, each of which require varied types of natural language understanding (from factual to logical commonsense to numerical commonsense reasoning); we also show variety in the text structure of the options - StrategyQA requires a yes/no selection across all questions, OpenBookQA, QuaRel, QASC (** new) require a word/phrase to be selected, and Numersense (**new) requires the selection of a number. Furthermore, prior work in generating rationales and reasoning chains like Chain of Thought [1] and other subsequent work, also demonstrate this ability on QA tasks, which has shown to generalize on other tasks.

[1] Wei et al. “Chain-of-Thought Prompting Elicits Reasoning in Large Language Models” (NeurIPS 2022)

Small improvement in Table 2 + Limited Novelty

As we note in the paper as well, significant improvements over the SFT baseline denoted by *. As it can be seen, for most metrics, training using reward signals as done by MaRio leads to these numerical improvements. Furthermore, these numerical improvements translate to large improvements in human evaluations, where we see that across all of the datasets, annotators prefer MaRio generations more, w.r.t the baseline. Additionally, we noticed that adding multiple rewards for rationale generation was not a trivial addition. In Section 5.2, we can observe that LMs tend to overfit to a reward and hack their way to improving these metrics. We first need to select metrics that are actually important for rationalization, and at the same time, make sure that they are compatible in a way so that they can be used together. We hope this can help clarify the motivation and novelty behind our work.

Thank you for your positive review and very helpful comments and suggestions! We are glad to hear that you appreciated our paper’s motivation, method and results. We have also added our response to the feedback and questions below.

The paper's main contribution is extending an existing method (Quark) from single-reward to multiple rewards. So while the results are nice and the extension is valuable, the contribution is not revolutionary.

Thank you for recognizing our contribution and results! As far as novelty is concerned, our idea was that firstly, in prior rationalization literature, rationales were used as a means to an end to improve the accuracy of the system – as seen in Chain of Thought [1] and other following works. However, often these rationales act as a medium of communication between the LM and user – to improve trust [2], utility [3] as well as help improve human-LM collaboration [4]. This provided us with a strong motivation to improve the quality of rationales themselves, with an added side benefit of accuracy improvement. Secondly, we noticed that adding multiple rewards for rationale generation was not a trivial addition, as observed in Section 5.2, where we see that LMs tend to overfit to a reward and hack their way to improving these metrics. We first need to select metrics that are actually important for rationalization, and at the same time, make sure that they are compatible in a way so that they can be used together. Our work is an attempt to demystify the same as well. We hope this can help clarify the motivation and novelty behind our work.

[1] Wei et al. “Chain-of-Thought Prompting Elicits Reasoning in Large Language Models” (NeurIPS 2022) [2] Chen et al. “Understanding the Role of Human Intuition on Reliance in Human-AI Decision-Making with Explanations” (CSCW 2023) [3] Joshi et al. “Are Machine Rationales (Not) Useful to Humans? Measuring and Improving Human Utility of Free-text Rationales” (ACL 2023) [4] Wiegreffe et al. “Reframing Human-AI Collaboration for Generating Free-Text Explanations” (NAACL 2022)

While the description in text and Figure 2 are very helpful for readers to understand MARIO, the full picture of MARIO can still be a bit hard to grasp (especially to readers who are not already familiar with Quark). Maybe including an algorithm block for the entire training process will make the full picture a lot more clear.

Thank you for this suggestion! We have uploaded new figures to explain our motivation and method better! We also have a more detailed technical figure in the appendix (Figure 5) to explain Quark and MaRio side by side, and we have now included the actual objective function that we optimize for, during training. During camera ready, we will also add the algorithm block in the main text, while we adjust for space!

Is there any way to rank/weight/balance different objectives in MARIO? (For example, if I found that the resulting model is weak in Consistency, is there a way to weigh Consistency reward a bit more in the training?)

That’s a great point actually. We started off with using these rewards individually (new single-reward optimization experiments added in Appendix G / Table 10), which showed that improving for one reward doesn’t necessitate improvements for other rewards. This motivated us to look into multi-reward approaches, leading to MaRio. Currently, we haven’t experimented with weighting (up or down) these metrics, and this is definitely an interesting future direction we are looking at!

Thank you for your positive review and very helpful comments and suggestions! We have added our response below as well as changes to the PDF based on your feedback.

The details of the MARIO algorithm are not adequately explained, such as how to determine the settings of control tokens, and the description of how to quantize samples under the quark framework is unclear (is it a comprehensive consideration of multiple attributes for ranking, or is it ranked based on a single attribute?).

These details are present in the Appendix (due to space constraints). Appendix B (“Quark and MaRio”) addresses how Quark is used in MaRio. Appendix C (“Order of tokens”) is used to explain how we order our rewards for the multi-reward setup. Figure 5 reiterates over this diagrammatically. In Table 7, we also mention hyper-parameters used by the Quark and MaRio algorithms. We mention here that our method quantized scores into 5 bins for all the rationale reward, and 2 bins for the accuracy reward. Even for our multi-reward experiments, the generated data is quantized separately for each individual reward; the final set of control tokens for each data point is the concatenated set of individual control tokens. We have added a clearer description of this in Appendix B as well.

The description of the MARIO method is overly simplistic, and it lacks the necessary explanation of the thought process behind the development of this method.

In section 5.2 and 5.3, we explicitly delve into the quirks of selecting good quality rewards, how these rewards actually translate into changes in the generated text, as well as a discussion about why we ended up adding task accuracy as another reward token in our design for MaRio. All of these serve as anchors for our design decisions for MaRio and highlight our thought process for the multi-reward setting. Our core algorithmic framework is based upon Quark itself, which is briefly explained in Appendix B. We hope the changes made in the appendix can help better clarify the method. If there are any other suggestions on how to explain our process better, we’d be more than happy to accommodate that into the draft!

In relation to self-explaining rationalization, besides the generative rationales discussed in this paper, there is a series of extractive rationale works (such as Lei et al. 2016, Liu et al. 2023, and so on). Beyond the difference between generative and extractive approaches, the basic framework of these two types of work is very similar. Both require ensuring that the generated/extracted rationale is meaningful to humans while maintaining high performance in downstream tasks. Therefore, the related work section should also include this series of works.

Quark as a framework is explicitly designed for text generation, which is why our focus was primarily on free-text rationales. Additionally, we have also included a subsection on extractive rationales in our related work section based on your suggestion!

Although Figure 1 and Figure 2 contain a considerable amount of text, the information conveyed is limited.

We acknowledge that Figures 1 and 2 were text dense. We have updated the PDF with new figures to explain our motivation and method better!

The experiment used multiple baselines, but in reality, it involves two baselines and their multi-reward forms of extension, lacking a comprehensive comparison with other works.

We compare with (1) the standard baseline SFT [1] , (2) a closely related work StaR [2] and its multi-reward equivalent, (3) with instruction fine tuned and otherwise trained large LLMs [3]. This intends to cover methods that generate free-text rationales by fine-tuning or prompting approaches both. For reward-based rationalization, there are currently no other baselines; we attempt multiple prior methods of incorporating multiple rewards – e.g product of rewards. If we missed any other methods to compare to, kindly let us know and we’d be more than happy to add it to our comparison pool!

[1] Marasović et al. “Few-Shot Self-Rationalization with Natural Language Prompts” (NAACL 2022) [2] Zelikman et al. “STaR: Bootstrapping Reasoning With Reasoning” (NeurIPS 2022) [3] Wei et al. “Chain-of-Thought Prompting Elicits Reasoning in Large Language Models” (NeurIPS 2022)

Although the model utilizes the quark framework, it should clearly present the learning objectives in a multi-reward scenario. Since quark is a single-reward algorithm, its objective function under the extension of multi-reward goals is not intuitive. Could you provide a more detailed and clear explanation of MARIO and the corresponding multi-reward loss?

Based on the Quark algorithm, we have duly updated Appendix B (“Quark and MaRio”) with the objective function of MaRio. This is essentially the same as Quark, except that we have multiple reward tokens on which the generation is conditioned.

Thank you for your helpful comments and insightful questions! We are happy to hear that you appreciated our method’s novelty, constraints, as well as design methodology. Based on your feedback, we have updated the draft and provided the responses below.

I believe the improvement of two versions of Marios compared to baselines is not really significant, considering that all the baseline models have equal number of parameters with the Mario agent model. I’m wondering whether the extra efforts on training on multiple rewards are indeed worth it to improve the generations.

While we do note that the baselines and MaRio have equal number of parameters, in Table 3 (previously, Table 2), we note significant improvements over the SFT baseline denoted by *. As it can be seen, for most metrics, training using reward signals as done by MaRio leads to these numerical improvements. Another reason why these extra efforts are helpful are seen in human evaluations, as shown by Figure 3. Across all of the datasets, annotators prefer MaRio generations more, while also noting improvements across these quality metrics, w.r.t the baseline. This can help us understand that small improvements numerically lead to large improvements in actual generations, thereby directly improving preference of these generations.

In the paper, the authors mentioned that there is an initial supervision phase where GPT-3 provides the generation labels. I’m wondering what’s the relationship between this initial supervision process and the distillation of GPT-3? Is it before, after, or exactly the distillation process?

We use the rationales provided by GPT-3 in the following 3 places: (1) we use it to train the supervised fine-tuned model SFT, which serves as a baseline, (2) we use the aforementioned SFT as the reference model for the KL divergence loss in MaRio’s training process (as we mention in Appendix), and (3) we also add these silver rationales to the overall data pool of MaRio. We have updated the PDF of the paper to re-reflect mentions of these throughout the draft.