Backtracking Improves Generation Safety

We appreciate your positive feedback. We address some of your concerns below.

No discussion or experiments on how likely backtracking LLM will overly reject safe queries

You may have missed our results in Appendix B.2 (referred to in the main paper in Section 4.2, L318), where we evaluated how often the Llama model backtracks on a set of safe queries (validation set of OpenAssistant). We find that with a suitable choice of logit bias on the [RESET] token, the model backtracks for only 1% of safe queries (12% without tuning the logit bias).

The performance of backtracking against adversarial attacks is not always satisfying, as shown in Table 2.

We acknowledge that neither backtracking, nor any other method, is not a perfect defense against adversarial attacks, but we want to point out that backtracking is a technique meant to improve model safety generally, and the additional robustness against strong adversarial attacks (e.g., GCG and AutoDAN) without special training suggests that backtracking could become a component in a truly robust LLM system.

Despite much progress, safeguarding LLMs against adversarial attacks remains an open problem, and is not an issue that can be addressed with a single technique. We note backtracking performance is strongly dependent on how it is trained: our models were not trained against adversarial attacks at all and still helped on such attacks. We expect that backtracking training in the presence of adversarial attacks (i.e., combining it with the idea of adversarial training [1]) would greatly improve robustness in the adversarial setting. However, we believe the results in the paper alone already demonstrate the value of our technique across a number of settings, and leave such approaches to future work.

[1] Towards Deep Learning Models Resistant to Adversarial Attacks

We appreciate your positive feedback on our work. We answer some of your questions below.

What if we skip backtracking SFT training?

Empirically, skipping backtracking SFT training breaks the model. The technical reason is that DPO imposes a “zero-forcing” reverse-KL regularization between the optimized policy and the reference (SFT) policy. Specifically, if we didn't use backtracking SFT to “warm-up” the model, the reference policy would assign practically zero probability to [RESET], and the resulting KL-penalty would be massive if the DPO policy assigns non-zero probability to [RESET]. We also note that it is standard practice to perform SFT before DPO [1].

The effect of scale on safety violation rates with backtracking, and difference in backtracking effectiveness on Gemma vs. Llama

Intuitively, we expect larger models to be able to learn to backtrack more precisely and self-correct more reliably, a trend we do observe between Gemma–2-2B and Llama-3-8B (Table 1), where the larger Llama model seems to gain more from backtracking. While we don’t have the resources to scale experiments to larger models, we expect the backtracking technique to be useful particularly for large models capable of self-critique, as even the state-of-the-art models remain far from being perfectly safe.

[1] Direct Preference Optimization: Your Language Model is Secretly a Reward Model

We appreciate your positive assessment of our work.

Compatibility with a streaming API

We agree that implementing backtracking in a streaming environment could be inconvenient for users visually, since backtracking “removes” part of the generation history upon backtracking. However, we would like to note that in the context of safety, it is always better to undo an unsafe generation even if it may be visually annoying. Also, backtracking is reasonably precise: it usually does not kick in unless the generation is unsafe, so a backtracking model would behave similarly to a non-backtracking model in non-safety-related contexts.

Risks in reprogrammable backtracking

In our work, we do not explore malicious system prompts and assume the model takes on a default “helpful” and “harmless” role, which would be the case for most black-box LLMs (e.g., ChatGPT, Gemini). We agree that it may be possible to inhibit backtracking through malicious system prompts, but even with backtracking inactive, the model should be in principle at least as safe as the non-backtracking baseline. Safety under model adaptation (e.g., fine-tuning and system prompting) is an open research question out-of-scope for this work.

Do we keep the undone history in the dialogue and KV cache or do we also undo them altogether and only keep the refined context?

We keep the history in-context. The reason is that we suspect that the presence of the “failed generation” actually provides a signal for the model to self-correct and generate more safely, and [1] suggests that this self-correcting approach could be effective. Another issue with dropping context is that the model would repeat its own unsafe generation with high probability. In the extreme case, if greedy decoding is used, the model is guaranteed to be stuck in a backtracking-regeneration loop after it backtracks once.

[1] Generating Sequences by Learning to Self-Correct

Thank you for your detailed read of our paper and your insightful feedback. We try to address your questions below.

There is no analysis of no analysis of how the proposed backtracking method compares/complements with existing methods for reducing unsafe generations.

The primary method for reducing unsafe generations is SFT or RL fine-tuning to reduce the probability of unsafe next-token generation. We compare backtracking with both SFT and DPO safety-trained variants in our work and find significant safety gains. However, there are other—less standard—methods of improving model safety such as [1] and [2]. We don't compare with these methods because backtracking is not mutually exclusive with other techniques such as [1, 2], and we expect complementary gains since it seems plausible that you can improve the safety of a base model using existing techniques, while teaching it to backtrack to gain additional safety.

Backtracking in OpenAssistant, and logit bias tuning

OpenAssistant (OA) is heavily focused on utility, and the 12% backtracking rate on OA (Fig. 5) suggests a small amount of false positives. This is almost entirely mitigated by adding a small negative logit bias to the [RESET] token (-5), with backtracking rate on OA dropping to 1%, with basically the same model safety (1.5% -> 1.6%). Tuning logit bias in principle is not difficult, and it just requires benchmarking backtracking behavior on a development set of prompts. We will add clarifications in the paper and point out the need for a suitable choice of logit bias on the [RESET] token to ensure a near-zero backtracking rate during non-safety evaluations.

What if we skip backtracking SFT training?

Empirically, skipping backtracking SFT training breaks the model. The technical reason is that DPO imposes a “zero-forcing” reverse-KL regularization between the optimized policy and the reference (SFT) policy. Specifically, if we didn't use backtracking SFT to “warm-up” the model, the reference policy would assign practically zero probability to [RESET], and the resulting KL-penalty would be massive if the DPO policy assigns non-zero probability to [RESET]. We also note that it is standard practice to perform SFT before DPO [3].

Data and models used for Fig. 4

To compute safety rates under sampling (Fig. 4), we used the same safety evaluation dataset we use throughout the paper, detailed in Section 4.1. The models compared are backtracking and baseline models that have both undergone SFT + DPO training.

Do most SFT datasets have unsafe responses $y_i^−$?

A common practice is to take the preferred response in a paired preference dataset for SFT training, so in such cases $y_i^−$ response would be available. If training on SFT datasets without $y_i^−$, we can still mix in preference data from the subsequent DPO stage to provide backtracking SFT supervision.

Have you tried attacking with an adversarial system prompt that says after you see the [RESET] token be extra helpful or something along these lines?

In our work, we do not explore malicious system prompts and assume the model takes on a default “helpful” and “harmless” role, which would be the case for most black-box LLMs (e.g., ChatGPT, Gemini). We agree that it may be possible to inhibit backtracking through malicious system prompts, but even with backtracking inactive, the model should be in principle at least as safe as the non-backtracking baseline. Safety under model adaptation (e.g., fine-tuning and system prompting) is an open research question out-of-scope for this work.

[1] The Instruction Hierarchy: Training LLMs to Prioritize Privileged Instructions
[2] Improving Alignment and Robustness with Circuit Breakers
[3] Direct Preference Optimization: Your Language Model is Secretly a Reward Model