Any-Depth Alignment: Unlocking Innate Safety Alignment of LLMs to Any-Depth

TL;DR

The paper addresses the brittleness of LLM safety by showing that front-loaded refusals fail as generation depth grows, and introduces Any-Depth Alignment (ADA) as a training-free, inference-time defense. ADA has two variants: ADA RK, which re-injects safety tokens to trigger a rethink and refusal, and ADA LP, a lightweight linear probe that reads Safety Token hidden states to halt harmful continuations without changing model weights. Across diverse open-source models and attack regimes, ADA delivers near-100% refusals for deep prefills up to $d=2500$ tokens and reduces adversarial prompt attack success to below 3%, while maintaining benign task utility and incurring minimal overhead thanks to KV-cache reuse. The work reveals that latent safety signals are concentrated in the assistant header and can be leveraged in real-time streaming deployments, proposing a new paradigm that harnesses innate safety priors rather than retraining models.

Abstract

Large Language Models (LLMs) exhibit strong but shallow alignment: they directly refuse harmful queries when a refusal is expected at the very start of an assistant turn, yet this protection collapses once a harmful continuation is underway (either through the adversarial attacks or via harmful assistant-prefill attacks). This raises a fundamental question: Can the innate shallow alignment in LLMs be unlocked to ensure safety at arbitrary generation depths? To achieve this goal, we propose Any-Depth Alignment (ADA), an effective inference-time defense with negligible overhead. ADA is built based on our observation that alignment is concentrated in the assistant header tokens through repeated use in shallow-refusal training, and these tokens possess the model's strong alignment priors. By reintroducing these tokens mid-stream, ADA induces the model to reassess harmfulness and recover refusals at any point in generation. Across diverse open-source model families (Llama, Gemma, Mistral, Qwen, DeepSeek, and gpt-oss), ADA achieves robust safety performance without requiring any changes to the base model's parameters. It secures a near-100% refusal rate against challenging adversarial prefill attacks ranging from dozens to thousands of tokens. Furthermore, ADA reduces the average success rate of prominent adversarial prompt attacks (such as GCG, AutoDAN, PAIR, and TAP) to below 3%. This is all accomplished while preserving utility on benign tasks with minimal over-refusal. ADA maintains this resilience even after the base model undergoes subsequent instruction tuning (benign or adversarial).

Figures (25)

• Figure 1: Refusal rates on AdvBench under harmful assistant-prefill attacks. Base models (dashed lines) exhibit a catastrophic drop in safety as the prefill depth increases. In contrast, applying our Any-Depth Alignment (ADA) method (solid lines) restores robust, near-100% refusal rates across all tested depths.
• Figure 2: Overview of the Any--Depth Alignment (ADA) mechanism. (Top Left) Without ADA, a model that starts generating harmful content will typically continue to do so. (Top Right) ADA intervenes at a safety checkpoint by leveraging model's own alignment. (a) ADA-Rethinking (ADA (RK)) re-injects the header to trigger a refusal. (b) ADA-Linear Probe (ADA (LP)) achieves the same outcome more effectively and efficiently by directly probing the strong safety signal present in the header's hidden states with a linear classifier.
• Figure 3: t-SNE of hidden states across depths (Llama-3.1-8B-Instruct, layer 15). As generation depth increases, features from injected Safety Tokens (bottom) -- where we read the hidden state of the assistant token from the assistant header -- become highly separable, while those from the last generated token (top) remain entangled. This indicates that the model's internal safety awareness strengthens with context but is cleanly revealed only via Safety Tokens. The accuracy shown in each panel is from a linear classifier trained on the 2D embeddings.
• Figure 4: Left: Across all model families, injected Safety Tokens (assistant headers) yield higher accuracy than the last-generated token for all layers. Right: Ablation on token choice (Llama-3.1-8B-Instruct) shows that tokens tied to the assistant header consistently provide stronger harmfulness signals than generic tokens such as a newline.
• Figure 5: Average refusal rates under deep prefill attacks across diverse LLMs. Results are averaged over four harmful datasets (AdvBench, JailbreakBench, StrongREJECT, and HEx-PHI). Our ADA (LP) (red line) achieves robust, depth-invariant safety, consistently outperforming all baselines. Detailed statistics by dataset and model are provided in \ref{['app:prefill']}.