💻 computer science

Matching Ranks Over Probability Yields Truly Deep Safety Alignment

This paper introduces the PRESTO method, which enhances deep safety alignment in open-source LLMs by matching token ranks in the target distribution to effectively counter the novel Rank-Assisted Prefilling (RAP) attack, achieving up to a 4.7x improvement in safety compared to standard data augmentation defenses.

Original authors: Jason Vega, Gagandeep Singh

Published 2026-07-23

📖 3 min read☕ Coffee break read

CC BY 4.0

Original authors: Jason Vega, Gagandeep Singh

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine a world where computers can write stories, solve math problems, and chat just like humans. These are called Large Language Models (LLMs), and they are like super-smart digital assistants. But because they are so smart, there's a worry: what if someone tricks them into saying something dangerous, like how to build a bomb or how to hurt someone? To stop this, scientists teach these models to say "No, I can't do that" when asked for bad things. This is called "safety alignment." However, just like a lock on a door, these safety rules can sometimes be picked. One tricky way to do this is called a "prefilling attack." Imagine you are writing a story with a friend, but instead of letting them start the next sentence, you sneakily write the first few words yourself, like "Here is the plan to build a bomb..." and then ask them to finish the sentence. Because the computer is designed to keep your story going smoothly, it might forget to say "No" and just keep writing the dangerous plan.

Recently, researchers found a way to fix this by teaching the model to say "No" even after you sneak in those first few words. But, as this new paper shows, that fix wasn't quite strong enough. The authors discovered a new way to trick the model, called the "Rank-Assisted Prefilling" (RAP) attack. Instead of just letting the computer pick the most likely word to say next, a hacker can look at a list of the top 20 words the computer is thinking about and pick the one that sounds dangerous, even if the computer thinks it's very unlikely to say it. It's like a game where the computer is trying to pick the safest path, but the hacker is allowed to peek at all the paths and choose the scary one, ignoring the computer's warnings. The paper finds that this new trick works surprisingly well, breaking the safety of models that were thought to be safe.

To stop this, the authors propose a new training method called PRESTO. Instead of just trying to make the computer say "No" with high confidence, PRESTO teaches the model to completely ignore the sneaky words you wrote at the start. It's like training a guard dog to ignore the voice of a thief whispering instructions and only listen to the owner's command. By doing this, the model stops letting the dangerous words influence its choices. The researchers tested this on three popular AI models and found that PRESTO made it much, much harder for hackers to trick the models—up to 4.7 times harder—without making the models worse at doing their normal jobs. This suggests that by changing how the model pays attention to what you say, we can build AI that is truly safe, even when someone tries to outsmart it.

Technical Summary: Matching Ranks Over Probability Yields Truly Deep Safety Alignment

1. Problem Statement

Open-source Large Language Models (LLMs) are critical for democratizing AI, yet their openness creates vulnerabilities for malicious actors. A specific and accessible threat is the prefilling attack, where an attacker manually pre-fills the assistant's response with an affirmative prefix (e.g., "Here is how to build a bomb...") to bypass safety filters.

Recent work by Qi et al. (2025) proposed a "deep" safety alignment defense using supervised fine-tuning (SFT) with data augmentation. This method trains models to generate natural language refusals immediately following a harmful prefill, rather than just outputting fixed refusal strings. While effective against standard sampling-based decoding and parameter-tuning exploits (like grid searches over temperature), this paper identifies a fundamental flaw: the defense remains superficial.

The authors demonstrate that even with deep safety alignment, models retain low-probability but highly-ranked "harmful" tokens within their top- $k$ predictions. By ignoring probability mass and instead selecting tokens based on their rank (specifically within the top 20), an attacker can bypass the refusal mechanism. This paper introduces the Rank-Assisted Prefilling (RAP) attack to exploit this vulnerability and proposes a new training objective, PRESTO, to achieve truly deep safety alignment.

2. Methodology

2.1 The RAP Attack

The Rank-Assisted Prefilling (RAP) attack generalizes the prefilling technique. Instead of relying on high-probability tokens, the attacker selects tokens from the top $k$ predicted next tokens (e.g., $k=20$ ) that continue the harmful narrative, regardless of their probability score.

Mechanism: At each decoding step, the attacker inspects the top $k$ tokens. If a token that naturally continues the harmful prefill (a "harmful token") appears in the top $k$ , it is selected, even if a refusal token (e.g., "I") has a significantly higher probability.
Automation (AutoRAP): The authors developed an automated version of RAP that uses a classifier to distinguish between refusal tokens and harmful continuation tokens, allowing for systematic extraction of harmful content without manual intervention.

2.2 Push-Forward Alignment (PFA)

The authors argue that the failure of the data augmentation defense lies in its focus on minimizing the negative log-likelihood (probability) of refusal tokens, rather than controlling the ranks of harmful tokens.

Concept: In a truly safe model, the distribution of the first response token (without prefill) should have refusal tokens at the top ranks and no harmful tokens that continue a prefill.
Goal: The model should "push forward" this rank-harm mapping from the first token distribution to all subsequent decoding steps when a harmful prefill is present. This means ensuring that if a token is harmful, its rank in the output distribution remains low, regardless of the prefill.
Formalization: The authors define a rank-harm distribution $R(s, x, x_{pre})$ representing the probability that the $s$ -th ranked token is harmful. The goal is to minimize the Earth Mover's Distance (EMD) between the rank-harm distribution of the safe model (no prefill) and the model under attack (with prefill).

2.3 PRESTO: PRefill attEntion STOpping

To implement PFA practically, the authors propose PRESTO, a regularization loss based on the Multi-Head Attention (MHA) mechanism in Transformers.

Hypothesis: The presence of harmful tokens in the top- $k$ rankings is driven by the model attending to the harmful prefill tokens. If the model can be trained to "ignore" the prefill tokens (i.e., reduce attention to them), the output distribution will revert to the safe, refusal-heavy distribution of the original prompt.
Loss Function: PRESTO minimizes the attention scores placed on harmful prefill tokens while maximizing attention on non-prefill tokens.
$\ell_{PRESTO}(\theta) = \mathbb{E}_{(x, x_{pre}) \sim D} [A_1 - A_2]$
Where $A_1$ is the average attention to prefill tokens and $A_2$ is the average attention to non-prefill tokens across all layers and heads.
Training: The model is fine-tuned using the standard data augmentation loss (from Qi et al., 2025) combined with the PRESTO loss. Both are computed in a single forward pass.

3. Key Contributions

Identification of RAP Vulnerability: The paper demonstrates that "deep" safety alignment via data augmentation is insufficient because it fails to suppress the ranks of harmful tokens, only their probabilities.
Rank-Assisted Prefilling (RAP) Attack: A new, accessible attack vector that extracts harmful content by selecting low-probability, high-rank tokens, effectively bypassing the data augmentation defense.
Push-Forward Alignment (PFA): A theoretical framework proposing that safety alignment must preserve the rank-harm mapping of the first token distribution throughout the generation process.
PRESTO: A simple, mechanistically interpretable training objective that regularizes attention to prefill tokens, significantly increasing the difficulty of RAP attacks.

4. Experimental Results

The authors evaluated PRESTO on three popular open-source models: Llama 2 7B Chat, Qwen 3 8B, and Gemma 3 12B IT.

Attack Success Rate:
- Under the original data augmentation defense (Qi et al., 2025), RAP attacks achieved high success rates (e.g., ~0.74 on the StrongREJECT scale for Llama 2 7B).
- With PRESTO, the success of RAP attacks dropped significantly. Across the three models, PRESTO yielded up to a 4.7x improvement in safety (reduction in harmfulness scores) compared to data augmentation alone.
- AutoRAP performance closely mirrored human evaluation, confirming the feasibility of automating the attack and the defense.
Utility Preservation:
- Models fine-tuned with PRESTO maintained competitive utility on MT-Bench (open-ended generation) and GSM-8K (mathematical reasoning), with performance degradation of less than 1% compared to models trained only with data augmentation.
Mechanistic Analysis:
- Attention analysis revealed that PRESTO successfully reduces attention to harmful prefill tokens, particularly in the latter half of the model layers, which are known to be critical for safety decisions.
- The defense remained robust against the GCG (Greedy Coordinate Gradient) attack, indicating that PRESTO does not reintroduce vulnerabilities to other attack vectors.

5. Significance and Claims

The paper claims that achieving "deep" safety alignment requires moving beyond probability-based objectives to rank-based constraints. The authors argue that simply training a model to output a refusal token with high probability is insufficient if the model's internal attention mechanisms still allow harmful continuations to remain in the top- $k$ ranks.

By introducing PRESTO, the work provides a path toward safer open-source models that are robust against practical, low-cost attacks like RAP. The significance lies in:

Democratization of Safety: Enabling open-source models to be deployed in customer-facing applications (like APIs that allow prefilling) without the risk of easy circumvention.
Cost of Circumvention: Aligning with the goal of increasing the cost for malicious actors to bypass safety, rather than attempting to make circumvention impossible.
Interpretability: Offering a mechanistic explanation (attention suppression) for why the defense works, bridging the gap between theoretical alignment and practical implementation.

The authors conclude that while no defense is perfect, PRESTO represents a significant step forward in making open-source LLMs safer against accessible, rank-based exploitation.