PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Jiawei Zhang, Andrew Estornell, David D. Baek, Bo Li, Xiaojun Xu

arXiv:2510.18081v1cs.LGcs.AI

TL;DR

Shallow alignment often fails once harmful generation is underway, while deeper alignment can still fail as attack depth grows. ADA reintroduces assistant-header Safety Tokens mid-stream and uses the model’s own safety signal to detect harmfulness. The authors report near-100% refusal under deep-prefill and adversarial attacks, with near-zero over-refusal on benign tasks.

  • Problem

    Shallow alignment’s initial refusals can collapse under harmful continuations, and deep alignment may only push the failure point deeper as attack depth increases.

    Source: §1 · p. 2

  • Method

    ADA re-injects assistant-header Safety Tokens mid-generation, using refusal lookahead or a linear probe on their hidden states to identify and stop harmful continuations.

    Source: §1 · p. 2, §2.2 · p. 5

  • Results

    ADA (LP) achieves near-100% refusal under deep-prefill and adversarial attacks while maintaining near-zero over-refusal on benign tasks.

    Source: §10 · p. 11

  • Takeaways & Limitations

    ADA (LP) reuses the model’s own safety representations for real-time streaming defense without re-engineering model weights.

    Source: §10 · p. 11

    ADA adds a small inference-time overhead, may emit a few harmful tokens before cutoff, and requires hidden-state access for ADA (LP); its practical setting is an operator-controlled inference stack.

    Source: §10 · p. 11

Abstract

from arXiv · show

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).

1 Introduction

The introduction frames shallow alignment as brittle: refusal can collapse under harmful prefills, and deeper alignment can still fail as attack depth grows. ADA reactivates the model’s safety assessment mid-generation, using header reinjection or a linear probe to detect and stop harmful continuations.

1 Introduction

  • A 25-token harmful prefill on AdvBench drops refusal rates from ∼100% to below 10%, while deep alignment can still fail as attack depth increases.The authors characterize this as an arms race between attack depth and alignment depth.Source: §1 · p. 2
  • ADA (RK) re-injects assistant-header Safety Tokens mid-stream to trigger refusal, providing a training-free defense that performs on par with or better than deep-alignment and self-reflection baselines.The authors report Claude Sonnet 4 exceeding 95% refusal with 500-token prefills.Source: §1 · pp. 2–3
  • ADA (LP) uses a linear classifier on Safety-Token hidden states to halt harmful continuations without auxiliary models or base-model weight updates.The authors describe the header tokens as aggregating context evidence into a separable safety judgment.Source: §1 · p. 2
  • ADA (LP) achieves near-100% refusal against deep prefills and reduces adversarial success from >50% to <3%.The reported approach also has minimal over-refusal on benign tasks and retains performance after fine-tuning.Source: §1 · pp. 2–3
  • Header-related Safety Tokens expose a linearly separable harmfulness signal across Llama, Qwen, Mistral, Gemma, DeepSeek variants, and gpt-oss.The reported pattern spans parameter scales and dense, Mixture-of-Experts, and reasoning-centric designs.Source: §1 · p. 3

2 Unlocking Innate Safety Alignment to Any-Depth Alignment

The section investigates why assistant-header reinjection can recover refusals during generation, then uses the resulting safety signal to motivate two inference-time defenses. Its findings center on the stronger, more separable harmfulness signal in header-token hidden states than in generated-token states.

2 Unlocking Innate Safety Alignment to Any-Depth Alignment

  • The section defines Safety Tokens as tokens whose hidden states carry strong internal safety signals, typically assistant-header tokens that expose the model’s assessment.Hidden states refer to representations after input layer normalization at a specified layer, unless otherwise noted.Source: §2 · p. 3

2.1 Innate Safety

  • Re-injecting the assistant header mid-generation can reactivate refusal behavior even after a jailbreak has induced harmful content, allowing streaming-time interruption at arbitrary depth.The section defines generation depth d as the number of assistant tokens generated after the user prompt, with d=0 immediately after the prompt.Source: §2 · p. 3, §2.1 · p. 4
  • As generation depth increases, injected Safety-Token features become more separable, while last-generated-token features remain entangled and lack a meaningful decision boundary.At d=0, both types of hidden states are described as entangled.Source: §2.1 · p. 4, Figure 3 · p. 5
  • Assistant-header hidden states support linear probes with near-perfect validation accuracy (> 99.5%), consistently outperforming last-generated-token probes across model families and scales.The reported pattern spans Llama, Mistral, Gemma, DeepSeek, and gpt-oss models.Source: §2.1 · p. 4
  • Assistant-header tokens, especially the assistant role token, provide stronger harmfulness signals than generic tokens such as newline or the last generated content token.The assistant role token sometimes yields better separability than the header’s final token.Source: §2.1 · pp. 4–5

2.2 Any-Depth Alignment (ADA)

  • ADA (RK) periodically forks the stream, injects assistant-header Safety Tokens, and returns a refusal if a short lookahead produces one.The method reuses the KV cache and halts the main stream when refusal appears.Source: §2.2 · p. 5
  • ADA (LP) extracts injected Safety-Token hidden states in one forward pass and applies a lightweight linear probe to classify harmfulness and halt flagged streams.It avoids additional decoding and does not change base-model weights.Source: §2.2 · p. 5

3 Experimental Setup

The setup evaluates ADA across deep-prefill, adversarial-prompt, and fine-tuning attacks, while also measuring benign-task over-refusal and inference costs. It compares the two ADA variants with base models, tuned deep-alignment models, self-defense prompts, and external classifier guardrails.

3 Experimental Setup

  • The evaluation covers three threat models: deep-prefill attacks, adversarial-prompt attacks, and SFT-based attacks, alongside benign-task over-refusal and inference-cost measurements.The threat models test depth robustness, harmful continuations elicited by manipulated inputs, and fine-tuning that attempts to overwrite safety.Source: §3 · pp. 5–6
  • ADA (RK) uses periodic checkpoints with up to 20-token header-anchored lookahead, returning a refusal and halting if one appears.Otherwise, the lookahead branch is discarded and the main sequence continues.Source: §3 · p. 6
  • ADA (LP) trains its linear probe once on Wild datasets and evaluates it out of distribution on downstream benchmarks.The section presents this as a test of generalization.Source: §3 · p. 6
  • Baselines include unmodified base models, deep alignment, self-defense reflection prompts, and external classifier guardrails; the evaluated models span Gemma, Llama, Qwen, DeepSeek, gpt-oss, and Claude Sonnet 4.Figure 4 also compares Safety-Token and last-generated-token accuracy across model families, and ablates header-token choices against generic tokens such as newline.Source: §3 · p. 6, Figure 4 · p. 6

4 Safety Awareness at Any Depth: Adversarial Prefill Attacks

The section tests defenses against harmful assistant prefills of increasing depth. ADA (LP) retains near-perfect refusal across models and depths, while existing alignment and other defenses are less robust.

Evaluation Protocol

  • The evaluation creates harmful assistant prefills from responses averaging over 3,500 tokens, then checks the next 50 decoded tokens for refusal.A deliberately misaligned model generated the responses with a 100% attack success rate; classifier-based guardrails instead evaluate the full context.Source: §4 · p. 6

Current Defenses Break with Depth

  • Base Models’ refusal rates fall near zero as depth grows, while Deep Alignment degrades to ∼40% refusal at d=500.On well-aligned models, ADA (RK) exceeds 95% refusal across depths and is competitive with Self-Defense without a reflective prompt.Source: §4 · p. 7
  • ADA (RK)’s performance tracks base-model alignment, and Self-Defense is ineffective in reasoning-centric settings.ADA (RK) exceeds 95% refusal across depths on well-aligned models; reasoning models require long generations for quality outputs.Source: §4 · p. 7
  • ADA (LP) maintains near-perfect (≈100%) refusal across models and depths up to 2,500 tokens, consistently outperforming strong external guardrails.Safety-Token states expose a depth-invariant safety signal that remains linearly separable even when it is not verbalized.Source: §4 · p. 7, Figure 5 · p. 7

Deep-Prefill Robustness

  • ADA (LP) retains near–100% refusal under a 500-token harmful assistant prefill while maintaining near–0% over-refusal on benign benchmarks.The table covers four harmful datasets and benign benchmarks for Gemma-2-9B-IT.Source: Table 1 · p. 7

5 Robustness Under Adversarial Prompt Attacks

The section evaluates defenses against gradient-based and paraphrasing attacks. ADA (LP) substantially reduces attack success without changing base-model weights, while Safety-Token probing detects latent harmfulness in ongoing continuations.

Attack Setup

  • The evaluation covers GCG, AutoDAN, PAIR, and TAP on AdvBench and JailbreakBench, checking for refusal every 25 generated tokens.An attack counts as successful if no checkpoint triggers a refusal.Source: §5 · p. 8

Results

  • ADA (LP) reduces Gemma-2-9B-IT attack success from > 50% across four attacks to 2%, and drives Llama-2 attack success to 0% under PAIR and TAP.It matches state-of-the-art guardrails and consistently surpasses the other baselines without modifying base-model weights.Source: §5 · p. 8
  • Safety-Token probing detects latent harmfulness that adversarial prompts leave unchanged in the ongoing continuation, even when safety is not expressed in text.The section reports this mechanism as the basis for ADA (LP)’s near-100% refusal.Source: §5 · p. 8

6 Robustness under Supervised Fine-Tuning (SFT) Attacks

The section tests ADA during benign and adversarial supervised fine-tuning. ADA (LP) remains resilient to deep-prefill attacks after fine-tuning, even as benign tuning rapidly weakens Deep Alignment.

Attack setup

  • The study uses Alpaca for benign SFT and harmful continuations for adversarial SFT, with LoRA rank 32 and learning rate 1 × 10−5.Each checkpoint is reevaluated against deep-prefill and adversarial prompt attacks.Source: §6 · pp. 8–9

Results

  • On Gemma-2, benign instruction-tuning rapidly undoes Deep Alignment: even a short run, such as 50 steps on Alpaca, is reported to collapse the alignment.At multiple training checkpoints, each defense is re-evaluated against deep-prefill attacks and adversarial prompt attacks.Source: §6 · p. 9
  • After 1,000 benign-SFT steps, ADA (LP) still exceeds 99% refusal under 100-token prefills; under adversarial SFT, it retains ∼90% on Gemma-2 and ∼100% on Llama-2.ADA (RK) also remains more stable than Deep Alignment under SFT, though ADA (LP) is especially resilient.Source: §6 · p. 9, Figure 7 · p. 9
  • The authors report that fine-tuning can erase response-level alignment while Safety Tokens retain hidden safety representations accessible through probing.The reported pattern holds under both benign and adversarial SFT in the tested settings.Source: §6 · p. 9

7 Evaluating Over-Refusal on Benign Tasks

The section evaluates accidental refusals on standard benign benchmarks and XSTest, a suite designed to elicit false positives. ADA (LP) maintains near-zero over-refusal while outperforming competing methods, including guardrails.

Inference cost

  • ADA (LP) reuses the base model’s KV cache, keeping latency and memory low and nearly constant as context length grows.The comparison uses 8B models with Flash Attention 2.Source: Figure 10 · p. 10

Evaluation

  • The evaluation covers seven standard benchmarks and XSTest, counting over-refusal when any periodic safety check flags a benign prompt.For Deep Alignment, the measure is refusal increase over the base model.Source: §7 · p. 10

Results

  • ADA (LP) maintains near-zero over-refusal across benign datasets and outperforms state-of-the-art guardrails, including on XSTest.The standard benchmarks include seven tasks; XSTest uses safe prompts containing sensitive keywords.Source: Figure 8 · p. 9, §7 · p. 10, Figure 9 · p. 10
  • Across the evaluated methods, ADA achieves higher safety and lower over-refusal; Deep Alignment often rejects benign queries at double-digit rates.ADA (RK) is competitive with Self-Defense, generally showing equal or lower false positives.Source: §7 · p. 10

8 Inference Cost

ADA (LP) adds a constant, small inference-time cost by reusing the model’s KV cache, unlike guardrails whose latency and memory grow with response length.

8 Inference Cost

  • ∼25 ms latency and 2–3 MB memory overhead remain constant for ADA (LP), versus nearly 500 ms and 938 MB for traditional guardrails on 10,000-token responses.The overhead comes from injecting safety tokens while reusing the base model’s KV cache.Source: §8 · pp. 10–11
  • Reusing the base model’s KV cache makes ADA (LP) equivalent to producing one additional token.Source: §8 · p. 11
  • ADA (LP) enables real-time streaming detection, so harmful outputs can be stopped mid-generation rather than after completion.The passage says this design scales efficiently to long contexts where traditional guardrails fail.Source: §8 · p. 11

9 Related Work

The paper distinguishes its focus on mid-stream generation and injected Safety Tokens from prior prompt-level, last-token, deep-alignment, and guardrail approaches.

9 Related Work

  • Unlike prompt-level hidden-state analyses, this work studies mid-stream generation, where harmfulness may emerge only as the continuation unfolds.Source: §9 · p. 11
  • Unlike probes on the last generated assistant token, ADA probes injected Safety Tokens for cleaner, more separable signals.Source: §9 · p. 11
  • Unlike deep alignment and external guardrails, ADA aims to use the model’s own alignment prior without fine-tuning or a separately trained classifier.Source: §9 · p. 11

10 Conclusion

The conclusion presents ADA as a way to unlock depth-invariant safety representations, while the related-work discussion situates it among shallow and deep alignment, attack, guardrail, and probing approaches.

10 Conclusion

  • ADA (LP) uses a lightweight linear probe to encode the model’s safety understanding, achieving near-100% refusal under deep-prefill and adversarial attacks with near-zero benign over-refusal.The conclusion also reports constant-time, minimal overhead through reuse of the base model’s KV cache.Source: §10 · p. 11
  • Future directions include dedicated Safety Tokens, pre-action checks for tool-using agents, and using separable activations as a reinforcement-learning reward signal.Source: §10 · p. 11
  • The approach has small but nonzero overhead, may emit a few harmful tokens before cutoff, and requires hidden states for ADA (LP).The authors place its most realistic use in operator-controlled inference stacks, not uncontrolled offline deployments where interventions can be removed.Source: §10 · p. 11

A Related Work

  • Shallow alignment trains direct refusals at response onset, but small initial-token perturbations and shallow prefill or adversarial prompts can elicit harmful continuations.Source: A Related Work · p. 16
  • Deep alignment inserts refusal patterns at random positions within the first hundred tokens, yet deeper prefills can again collapse refusals amid conflicting instruction-following and refusal objectives.Source: A Related Work · p. 16
  • Adversarial-prompt methods progress from optimized but often unnatural GCG suffixes to human-readable AutoDAN, black-box PAIR, and pruned-tree TAP search.Source: A Related Work · p. 16
  • Guardrail work includes lightweight response classifiers, human-feedback-aligned coverage, adversarial-data training, and unified safety objectives.Source: A Related Work · p. 16
  • Linear-probe research ranges from single refusal directions to multidimensional concept cones and safety-constraint vectors; Figure 11 reports stronger training accuracy from injected Safety Tokens than last-generated tokens across layers.For Gemma-2-9B-Instruct, tested readout locations exceed 96% accuracy, with LayerNorm interfaces yielding the highest and most stable probe accuracy.Source: A Related Work · p. 16, Figure 11 · p. 17

B Training Details on the ADA (LP)

ADA (LP) probes safety information in injected assistant-header tokens, whose representations distinguish harmfulness and reactivate refusal behavior. Across the reported deep-prefill and adversarial-prompt evaluations, ADA (LP) sustains strong refusal performance; the section also examines its mechanism, setup, and evaluation boundaries.

B Training Details on the ADA (LP)

  • The ADA (LP) probe is trained on depth-indexed hidden states from 11,000 benign examples in each of two sources and 11,000 harmful continuations, yielding 600,000 training and 60,000 validation examples.Probes on injected Safety Tokens achieve close to 100% training accuracy across many layers and model families, with a broad near-perfect middle-layer plateau.Source: B Training Details on the ADA (LP) · pp. 17–18
  • On Gemma-2, Safety-Token signals remain linearly separable across six readout positions and a wide range of middle layers.Input LayerNorm is slightly more stable, but submodule outputs also yield high probe accuracy.Source: B Training Details on the ADA (LP) · p. 18

C Refusal Activation via Transcoder

  • After a harmful continuation, re-inserting the assistant header reactivates refusal-associated neurons on the model token and its following newline, and the model switches to refusal.The analysis uses transcoder models trained on Gemma-2-base, whereas the primary experiments use Gemma-2-it.Source: Figure 12 · p. 18, C Refusal Activation via Transcoder · p. 19

D Jailbreaking GPT via SFT

  • The study uses a deliberately potent proprietary SFT-based jailbreak interface, but omits implementation details and defers systematic attack analysis to future work.The reported attack can elicit harmful continuations often exceeding 3,000 tokens from frontier models.Source: D Jailbreaking GPT via SFT · p. 19

E Examples of Prefill Continuations

F.1 Detailed Results against Deep Prefill Attack

  • ADA (LP) maintains approximately 100% refusal across four benchmarks through harmful prefills up to 2,500 tokens, outperforming strong external guardrails.Base models and Deep Alignment degrade with depth, while ADA (RK) offers training-free robustness that tracks base-model alignment.Source: §F.1 · p. 24, Figure 17 · p. 25

F.2 Detailed Results against Adversarial Prompt Attack

  • ADA (LP) reaches 1.5% average ASR on AdvBench and 4.5% on JailbreakBench for Llama-2-7b-chat-hf, while keeping ASR to low single digits on both suites for difficult models.It matches or outperforms strong external guardrails across GCG, AutoDAN, PAIR, and TAP without modifying base-model weights.Source: §F.2 · p. 24
  • ADA (RK) improves over base models and often rivals Self Defense, with its strength scaling with the base model’s alignment.It is strongest on Llama models and weaker on Qwen2.5 and Ministral; on Gemma-2-9B-IT, ADA (LP) remains low across all four tested attacks while Deep Alignment remains vulnerable to paraphrasing.Source: §F.2 · p. 24

F.3 Robustness against SFT and Adapter Ablation

ADA (LP) remains robust under benign and adversarial SFT, while its Safety-Token probe is largely unaffected by whether the LoRA adapter is enabled on that branch. Across deep-prefill and benign-task evaluations, ADA (LP) combines high refusal with near-zero over-refusal.

F.3 Robustness against SFT and Adapter Ablation

  • Under benign SFT, Deep Alignment degrades rapidly while ADA (LP) stays near 100% refusal at depth 1000; adapter settings on Safety Tokens yield effectively indistinguishable curves.The probe-branch adapter toggle affects hidden-state computation only; the adapter remains enabled during normal generation.Source: §F.3 · p. 25
  • ADA (LP) sustains low ASR under adversarial SFT across steps, attacks, and both model families, with adapter Enable/Disable tracking closely.Llama-2-7B-IT often reaches 0% ASR on PAIR and TAP and single digits on AutoDAN and GCG at 1000 steps.Source: §F.3 · pp. 25–26

F.4 Detailed Results on Over-Refusal Rates for Benign Datasets

  • Across nine models, ADA (LP) maintains near-100% refusal across tested prefill depths on JailbreakBench, HexPhi, and StrongReject.ADA (RK) is also described as strong without training on HexPhi.Source: Figure 18 · p. 26, Figure 19 · p. 27, Figure 20 · p. 28
  • ADA (LP) stays near zero in over-refusal across standard benign benchmarks and XSTest, while several baselines show higher or more depth-sensitive false positives.The XSTest prompts include sensitive keywords designed to trigger spurious refusals; other defenses vary by model and can over-refuse there.Source: §F.4 · p. 26, Figure 24 · p. 33, Figure 25 · p. 34
  • On Claude Sonnet 4, ADA (RK) restores refusal to at least 90% across datasets and depths up to 2500 tokens, without hidden-state access.The method appends an extra assistant turn to implicitly inject the assistant header mid-stream.Source: Figure 21 · p. 29

Related papers

Loading 2510.18081v1…