Paper: Towards Minimal Targeted Updates of Language Models with Targeted Negative Training (ICLR-2024)

Reviewer comment (Experiments):
The experiments could be greatly strengthened. The authors study a single model T5, at a fairly small scale (220M) parameters. No purely autoregressive model is studied. Only two tasks are examined. This lack of breadth makes it hard for readers to access the generality of the claims in the paper. I believe this paper would benefit greatly from more experiments and ablations.

Author response:
Thank you for the feedback. First, we would like to emphasize that in order to generate Figure 2, we ran 7 * 9 = 64 finetuning jobs on two very different generation tasks (hallucination in summarization and toxicity in response generation), for 126 jobs in total. Figure 2 shows that TNT consistently outperforms or at least matches baseline methods across all rates of reduction in both tasks, demonstrating that its benefit is broad and robust to different definitions of unwanted. We’ve additionally added a set of ablations comparing methods on external data rather than model generations; TNT still outperforms baseline the results, illustrating that the proposed losses are better even under a different data context, and the main results are better than the ablations, confirming our hypothesis that preferencing more common token conditionals yields a more minimal update. Due to the extensive nature of each experiment on just one model and task, we were not able to complete such experiments for larger models (magnitudes more resource-intensive) or additional tasks (requires appropriate dataset), but we hope the mathematical analysis and updated experiments convinces the reviewer that there is merit in the proposed approach.

Reviewer comment (Experiments):
There are very few comparisons to previous work, which in many cases tackle the same tasks presented in the experiments. As an example, avoiding toxic generations has been studied by [1,2]

[1] Lu, Ximing, et al. "Quark: Controllable text generation with reinforced unlearning." Advances in neural information processing systems 35 (2022): 27591-27609.  
[2] Ilharco, Gabriel, et al. "Editing models with task arithmetic." arXiv preprint arXiv:2212.04089 (2022).

Author response:
We chose as baselines other finetuning-based methods that take into account token-level information for learning. The works the reviewer mentioned do not fall into this setting, but we have instead included them as background or related work:
- Quark finetunes an existing language model on its generations by prepending inputs with sequence-level rewards and regularizing with a KL penalty on the token-level conditional distributions. Then, during inference time one conditions on high reward to sample from the model. TNT differs in that it utilizes token-level annotations and directly specifies the output token probabilities rather than conditioning the inputs on the rewards.
- Task Arithmetic subtracts the weights of pretrained and finetuned models to yield task vectors which are composed together to add or negate behaviors. As these task vectors are a function of vanilla finetuning, e.g. via filtering the dataset of toxic language, this approach, like finetuning, will not yield minimal targeted updates.

Reviewer comment (Novelty):
Limited Novelty: The core ideas and methods of TNT may not be as novel as claimed, given the prior existence of the NADO algorithm [NeurIPS 2022, https://arxiv.org/pdf/2205.14219.pdf] framework which appears to address similar goals using related techniques. Further more, NADO has proven its objective to be the **theoretically closed form solution** of the shared targets of TNT/NADO, which weakens the value of the approximated solution (step-level branch-cutting) given by TNT. TNT's flexibility is also less than that of NADO as it requires auxiliary negative annotation, which is a closer setup to that of the FUDGE algorithm [NAACL 2021, https://arxiv.org/abs/2104.05218].

Author response:
1. [*The core ideas and methods of TNT may not be as novel as claimed, given the prior existence of the NADO algorithm framework which appears to address similar goals using related techniques.*] 

Thank you for bringing the NADO paper to our attention. It presents a great result that one can turn sequence-level constraint information into token-level guidance by estimating $P_{y\sim p(y|x)}[C(x, y)|y_{<t}]$ for all $y_{<t}$). However, we disagree that the existence of this analysis and the NADO algorithm weakens the novelty or value of TNT. Ways in which the methods differ:

- NADO provides an inference-time algorithm for sampling from the desired target distribution while TNT is a suite of finetuning objectives to approximate a target distribution by matching token conditionals.
- NADO uses sequence-level annotations while TNT uses token-level annotations.
- NADO involves training an auxiliary model while TNT involves finetuning the language model itself.
- NADO seeks to learn $P_{y\sim p(y|x)}[C(x, y)|y_{<t}]$ for all $x, y_{<t}$ while TNT seeks to learn $p_{new}(y_t | x, y_{<t})$ for all $x, y_{<t}$. For the former, every element in the set is intractable so can only be estimated via samples, whereas for the latter, every element in the set can be queried exactly via a model forward pass.

Moreover, where the two works overlap is not a contribution of either approach. Namely, we define a minimal targeted update as the closest distribution in reverse KL to the original given negative examples are zeroed out (Eq 1), and the solution (Eq 2) matches the distribution given NADO paper’s Eq 3 (The authors define the problem using the forward KL, but we assume they use the distribution and not the stated KL divergence since the forward KL will be infinite for any distribution more constrained in support than the original). This setup was also given in [1, 2] before both the TNT and NADO papers, and our paper also connects this formulation to other works in the controllable generation space that had not previously mentioned it (see Relationship to Conditioning under 2.3). 

TNT’s contribution relative to [1, 2] is to show how moving from the sequence-level to token-level allows for a simple suite of finetuning methods (instead of the multi-step training algorithms proposed by [1, 2]). NADO’s contribution relative to [1, 2] instead is to show how one can translate sequence-level annotations to token-level guidance.

These contributions are in fact synergistic; as we have now highlighted in our related work, it is possible to combine the complementary contributions of both works: 
> Meng et al 2022 also consider sequence-level constraints but prove that this setting can be translated into token-level guidance (i.e., relating $p_{new}(x_t | x_{<t}, c)$ to $p_0(x_t | x_{<t}, c)$) via the approximation of $\Pr_{y\sim p(y|x)}[C(x, y)|y_{<t}]$ for all $y_{<t}$ and sequence-level boolean constraint function $C$. Consequently, Meng et al 2022 are able to propose a simpler algorithm than existing work, and by combining ideas in their work with this work, it is possible to define a finetuning algorithm that optimizes analytical token-level divergences even given only sequence-level annotations (define $p_{c, x_{<t}}^{new}$ using results in Meng et al 2022, optimize using TNT).

2. [*NADO has proven its objective to be the theoretically closed form solution of the shared targets of TNT/NADO, which weakens the value of the approximated solution (step-level branch-cutting) given by TNT*] 

We disagree with the reviewer’s claim that TNT is an approximated solution. The solution TNT targets yields $p_{new}(\mathbf{x} | \mathbf{c})$ exactly, as the sequence-level distribution is correct when all token-level conditionals are correct. However, as it is not possible to enumerate all token-level conditionals, TNT focuses on matching those that are more commonly encountered under the original model. This approximation is analogous to that of any objective targeting a solution that cannot be enumerated, even though the optimal solution itself is exact. Similarly, NADO relies on the auxiliary model accurately estimating $P_{y\sim p(y|x)}[C(x, y)|y_{<t}]$ for all $y_{<t}$ given only samples.

3. [*TNT's flexibility is also less than that of NADO as it requires auxiliary negative annotation*] 

We agree that TNT’s assumption of token-level annotations is stronger than NADO’s assumption of sequence-level annotations, but as discussed earlier in point 2, there exist situations where collecting more fine-grained annotations will be more advantageous from an effort-to-value ratio (e.g. labeling a sequence as “has hallucination” requires identifying the hallucination), and a method such as TNT allows one to take advantage of such fine-grained information via an easy-to-implement algorithm. Furthermore, as mentioned above, the theoretical result in NADO makes it possible to define a TNT variant that continues to match token-level conditionals even only given sequence level information.

Reviewer comment (Novelty):
Oversight in Literature Review: Even if TNT can differ itself from NADO and FUDGE (since the two previous methods study different tasks, yet essentially with similar mathematical setup), the absence of a citation and/or discussion of these existing works suggests a possible gap in the literature review process, which might question the thoroughness of the background research conducted for the paper.

Author response:
[Absence of citation and/or discussion of NADO and FUDGE suggests a possible gap in the literature review process] We first would like to point out that we have already cited FUDGE in our paper and explained how it is an inference-time controllable generation approach, in contrast to our finetuning-based approach. Thank you for bringing NADO to our attention. We have now added a citation and discussion of it as well.

Reviewer comment (Novelty):
The authors are missing connections to existing methods. For example, [Korbak et al, 2022](https://arxiv.org/abs/2205.11275) (among others) show that PPO would converge to the same closed-form presented in Eq. 2 when using the positivity/negativity classification as a reward function; this method would also avoid the drawbacks mentioned in the related work for inference time procedures.

Author response:
Thank you for bringing this up. We have added to the related work that, given a positivity/negativity classifier as a reward function, the objective PPO seeks to maximize (equivalent to minimizing reverse KL) has an optimal solution that is a soft version of the result TNT targets: namely, $p^*_{PPO}(x|c) = p_{o}(x|c)\exp(r(x)/\beta)$ while $p^*_{TNT}(x|c) = p_o(x|c)\mathbf{1}[r(x)]$, where the former is approximately equal to the latter when $\beta$ is small.

Reviewer comment (Experiments):
The paper demonstrates their technique on some relatively easier benchmarks, and it would be much more interesting to try more complicated schemes.

Author response:
We agree that the wordlist setup does not demonstrate the full practical utility of TNT and have replaced this experiment with one where we train a  token-level classifier to identify toxic spans. In this updated experiment, the benefit of TNT over baselines is even more pronounced, with TNT methods overall (and each of TNRF, TNRR, and TNRLL individually) being strictly better than baselines across all ranges of similarity and reduction (see Figures 2(c) and (d)).

Reviewer comment (Experiments):
In the specific case where p_neg is defined as the presence of a bad token in the string, there is no need to do any training, and by simply ignoring bad tokens at decoding, one recovers the true optimal solution for both greedy and temperature decoding with zero overhead. Both the toxicity and hallucinations benchmarks provided in this text are dangerously close to "reject any sentence with a bad (word/entity)", which makes the application rather uninteresting. It would be cooler to see benchmarks where the reward model may still be automated but captures some global property of the sentence that requires a learning-based technique such as yours to solve.

Author response:
We agree that the wordlist setup does not demonstrate the full practical utility of TNT and have replaced this experiment with one where we train a  token-level classifier to identify toxic spans. In this updated experiment, the benefit of TNT over baselines is even more pronounced, with TNT methods overall (and each of TNRF, TNRR, and TNRLL individually) being strictly better than baselines across all ranges of similarity and reduction (see Figures 2(c) and (d)).

Reviewer comment (Experiments):
For the toxicity benchmark, the authors mention that 1.6% of the time, the completion is toxic. A very natural baseline to this problem is performing temperature 1 decoding and regenerating anytime the output is toxic. This will likely very rarely require few regenerations for this specific benchmark, making the learning rather excessive. Having a benchmark that requires a larger change will be a true test of this problem.

Author response:
We agree, and with our new broader definition of toxicity defined by a token-level classifier trained on human annotations, the level of toxicity of the original model is 8%. Moreover, this 8% rate is not constant across inputs; for instance, we see that for examples where the input contains toxicity (based on the same classifier definition), the model generations have toxic language 23.3% of the time. Thus in our updated experiment we have a setting where regeneration would be quite expensive for a certain subset of inputs, making finetuning a more practical alternative.

Reviewer comment (Experiments):
For baselines, as mentioned in Weakness 1, PPO is a more natural baseline for learning from a reward model.

Author response:
Thanks for the suggestion. We initially did not consider PPO as a baseline since it is an algorithm generally used to optimize sequence-level rewards. [1] recently introduced Fine-Grained RLHF which could be used to consider token-level rewards, but we did not have the chance to implement this baseline given time constraints (It is also worth noting that this baseline is much more complicated than the proposed approach, requiring sampling from the model during training as well as keeping track of four separate models, including a separate value model that is trained simultaneously to the policy model). Our experiments do hint at the opportunity for considering alternatives beyond such a baseline, though. Namely, whereas the objective considered by the PPO algorithm is equivalent to minimizing the reverse KL divergence between the current model $p_\theta(x|c)$ and the target distribution $p^*(x|c)$, regardless of whether we consider sequence- or token-level rewards, TNT encompasses a suite of objectives which consider minimizing not just the reverse KL between token-level conditional distributions but also other combinations of divergences, and our experiments show that other divergence combinations may be preferable in certain scenarios, even when the algorithm is fixed regardless of objective; in particular, at lower levels of reduction, TNT losses that use forward KL (i.e., TNRF and TNFF) tend to outperform TNRR at maintaining similarity with the original model generations.

Reviewer comment (Experiments):
For the different proposed models, without an equivalent table similar to Table 1, it is hard to understand the effectiveness of the approach. Concretely, similar to the TNFLL and TNRLL rows, it would also be good to have an equivalent row for TNFF, TNRR and TNRF.

Author response:
Thank you for the suggestion; we have updated Table 1 to include all TNT methods. We’d also like to clarify that Figure 2 (rather than Table 1) is meant to encapsulate the main result, as it considers all methods across all alpha values in one plot; Table 1 on the other hand only looks at a single fixed alpha value and is instead meant to focus on how changing the loss on negative tokens alone can reduce disfluencies. This is the reason why the original Table 1 focused on TNT methods that share the same loss as baselines on non-negative tokens (i.e., TNFLL and TNRLL). We’ve updated the text to draw attention to Figure 2 as the main result.

Reviewer comment (Evaluation):
From Table 1, for both the TNFLL and TNRLL approaches, the performance is considerably worse compared to the baseline method. For TNFLL, the hallucination rate is substantially higher, while for TNRLL, the BLEU score is much lower. Given this observation, I am hesitant to believe that the proposed approach is actually substantially better than the baseline approach.

Author response:
First, we wish to reiterate that Figure 2 gives the main result, that TNT methods yield a better trade-off of reducing unwanted behavior vs. maintaining similarity to the original. While we intended Table 1 to focus on disfluency results, based on your feedback we realize how the current presentation of Table 1 could be confusing, as we also report similarity and reduction metrics but only for a single fixed value of $\alpha$, which puts different methods at different points along their similarity vs. reduction curves. 
This fixed value of $\alpha$ in resulted in rows in Table 1 that are hard to compare along similarity and reduction metrics, e.g. if one row has more reduction but lower similarity than another, which is better? Figure 2 is meant to address this challenge. To make Figure 2 and Table 1 congruent, we have additionally updated Table 1 to choose the $\alpha$ that yields the best BLEU score among $\alpha$ values that achieve a given rate of reduction (75%). See the updated table for toxicity below. 
Now, the methods can be more easily compared across all metrics (similarity, reduction, and disfluency); for instance, one can see that, given a 75% rate of reduction is achieved, **TNRF is strictly better than baselines on all metrics**, all TNT methods except TNFLL yield better similarity results than baselines, and all TNT methods perform better than baselines on the number of total introduced disfluencies (sum of Repeats and Random ??).


|                        |    BLEU   |  ROUGE-L   |  Seq Acc   |  Toxicity  | Repeats |  Random ??  |
|------------------------|-----------|------------|------------|------------|---------|-------------|
| Original               | 100.0000  | 100.0000   | 100.0000   | 8.1830     | 16      | 4           |
| NL + LL ($\alpha=.01$) | 13.6497   | 32.6104    | 2.0661     | 1.8150     | 287     | 136         |
| UL + LL ($\alpha=1.0$) | 37.1265   | 59.9806    | 20.4405    | 1.6784     | 23      | 1122        |
| TNFLL ($\alpha=1.0$)   | 33.7884   | 57.1009    | 18.1630    | 1.7577     | 36      | **1**       |
| TNRLL ($\alpha=0.1$)   | **39.1922**| **61.4776** | **22.9207** | 1.9471     | 23      | **1**       |
| TNRR ($\alpha=0.1$) | **55.9532**| **71.5365** | **35.1366** | **1.4493** | 34      | **3**       |
| TNRF ($\alpha=1.0$) | **60.2071**| **74.8574** | **39.6167** | **1.0396** | **21**  | **3**       |
| TNFF ($\alpha=10.$)| **61.0565**| **74.2388** | **40.0749** | 1.9031     | 33      | **3**       |

Reviewer comment (Theory):
In my opinion, intuitively, because this approach minimizes the KL at a prefix level, especially considering the fact that the annotations obtained are from a noisy source, it is possible that this approach would steer the model towards not predicting certain words in certain contexts. Concretely, (based from the example in Figure 5), for the sentence "In some regions of the country, the sex ratio is still quite concerning", because the annotations are noisy, the word "sex" would be (incorrectly) marked as offensive. Consequently, because of the proposed objective, the model might not be able to produce the token "sex" for a similar prefix as "In some regions of the country, the", even if it did make sense in the context. I think this is a reasonably big limitation of the approach, and it would have been nice to have some discussion on this in the paper.

Author response:
We take the reviewer’s comment to mean that if the model is not context-aware when pushing down probability mass over certain tokens, then it is possible for the model to begin avoiding a given token even in contexts where it would be acceptable. Fortunately, this concern is not an issue for TNT because it is context-aware: all negative annotations are a function of the prefix, which means that the negative signal the model sees is of the form “X is bad when preceded by Y” rather than just “X is bad.” As the wordlist example does not showcase this context-aware property of the method since the annotations themselves do not take into account context, we replaced the wordlist experiment with one based on a token-level classifier that labels spans as toxic based on its context.

Reviewer comment (Presentation):
Would it be possible to add rows for TNFF, TNRR and TNRF in Table 1?

Author response:
Thank you for the suggestion; we have updated Table 1 to include all TNT methods. We’d also like to clarify that Figure 2 (rather than Table 1) is meant to encapsulate the main result, as it considers all methods across all alpha values in one plot; Table 1 on the other hand only looks at a single fixed alpha value and is instead meant to focus on how changing the loss on negative tokens alone can reduce disfluencies. This is the reason why the original Table 1 focused on TNT methods that share the same loss as baselines on non-negative tokens (i.e., TNFLL and TNRLL). We’ve updated the text to draw attention to Figure 2 as the main result.

Reviewer comment (Presentation):
Would it be possible to provide some clarification on how Figure 2 was constructed? Specifically, is the level of hallucination mapped to a different value of \alpha used (so higher \alpha -> lower hallucination rates and original distribution fidelity)?

Author response:
To construct Figure 2, for each method we performed nine separate updates using 9 different alpha values. Then, we iterate through hallucination (or toxicity) rates in increments of 0.1 and plot the highest BLEU achieved across $\alpha$ values that achieve less than that given level of hallucination (or toxicity). The model selection was chosen based on the validation set results, while the results themselves are on the test set. And yes, in general a higher alpha corresponds to lower hallucination rates and lower original distribution fidelity (exceptions include baseline models which struggle to increase original distribution fidelity above a certain level even as $\alpha$ becomes smaller, and cases where an $\alpha$ that is too large causes optimization issues).

Thank you for your comments. Please let us know if you have additional questions or concerns! If not, would you consider raising your score?