Matching Ranks Over Probability Yields Truly Deep Safety Alignment
Jason Vega, Gagandeep Singh
Why It Matters
What makes this one worth your time
As open-source LLMs become more prevalent, enhancing their safety against misuse is crucial for responsible AI deployment.
This work presents a novel attack and a method to improve safety alignment in LLMs.
Summary
The paper introduces the Rank-Assisted Prefilling (RAP) attack, which exploits the weaknesses of a data augmentation defense in open-source LLMs, and proposes a new method called PRESTO that enhances safety alignment by matching token ranks in the target distribution.
Key contributions
- Introduction of the Rank-Assisted Prefilling (RAP) attack to demonstrate vulnerabilities in safety alignment.
- Development of the PRESTO method to enhance safety alignment by focusing on token rank matching.
Notable insights
- The RAP attack highlights a critical vulnerability in existing safety alignment techniques, emphasizing the need for robust defenses.
- Matching token ranks rather than probabilities offers a fresh perspective on improving model safety.
Possible limitations
- Not stated in the abstract.
Abstract
arXiv:2512.05518v2 Announce Type: replace-cross Abstract: Open-source Large Language Models (LLMs) play a critical role in the democratization of AI, yet their "open" nature introduces more avenues for malicious actors to misuse them for harmful purposes. A frustratingly easy but powerful technique known as the prefilling attack has been shown to effectively circumvent the safety alignment of frontier open-source LLMs by simply prefilling the assistant response with an affirmative prefix before decoding. A recent promising supervised fine-tuning defense proposes using a simple data augmentation scheme to achieve a "deep" safety alignment, allowing the model to generate natural language refusals immediately following harmful prefills. In this work, we show that a simple generalization of the prefilling attack, which we refer to as the Rank-Assisted Prefilling (RAP) attack, can effectively extract harmful content from models fine-tuned with the data augmentation defense by selecting low-probability "harmful" tokens from the top 20 predicted next tokens at each step (thus ignoring high-probability \enquote{refusal} tokens). We then propose a new perspective on achieving deep safety alignment by matching the token ranks in the underlying data augmentation target distribution (rather than just their probabilities), yielding a surprisingly simple approach to strengthening deep alignment we call PRefill attEntion STOpping (PRESTO) that regularizes the attention placed on harmful prefill tokens. Compared to just using data augmentation, PRESTO yields up to a 4.7x improvement in safety under RAP attacks across three popular LLMs. By achieving a stronger level of safety against practical and accessible attacks, our work paves a path towards safer open-source models. The source code accompanying this paper can be found at: https://github.com/uiuc-focal-lab/push-forward-alignment