🤖 machine learning

Mitigating Adaptive Attacks against Reasoning Models with Activation Consistency Training

This paper demonstrates that Activation Consistency Training (ACT), which enforces identical internal representations for clean and adversarial prompts, effectively mitigates adaptive jailbreak attacks in reasoning models by creating an interpretable, linear steering direction for refusal while preserving benign performance.

Original authors: Avidan Shah, Jannik Brinkmann, Rico Angell

Published 2026-05-28

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Avidan Shah, Jannik Brinkmann, Rico Angell

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 highly intelligent assistant who is incredibly good at solving complex puzzles. To solve these puzzles, the assistant doesn't just give an answer immediately; it writes out a long, detailed "thinking process" (like a stream of consciousness) before speaking its final answer. This is called Chain-of-Thought (CoT).

The problem is that clever tricksters (adversaries) have learned how to hijack this thinking process. They can sneak in hidden instructions that convince the assistant, while it is still thinking, that a dangerous request is actually safe. The assistant then talks itself into agreeing to do something it shouldn't, like revealing a secret or writing a harmful guide.

This paper introduces a new way to train these assistants so they can't be tricked, even when they are thinking out loud. Here is the breakdown using simple analogies:

The Two Main Characters: BCT and ACT

The researchers tested two different training methods to make the assistant "immune" to these tricks.

1. BCT (Bias-Augmented Consistency Training): The "Script Reader"

How it works: Imagine you are teaching a student to ignore distractions. With BCT, you show the student a clean question and then a "tricky" version of that same question. You then force the student to read the entire long thinking process and the final answer from the clean version, and make them repeat that exact script when they see the tricky version.
The Catch: Because the thinking process can be hundreds of words long, the student has to memorize a massive script every time. If the trickster changes the thinking process slightly, the student might get confused and fail.

2. ACT (Activation Consistency Training): The "Inner Compass"

How it works: Instead of forcing the student to memorize the whole script, ACT focuses on the very moment just before the student starts speaking. It looks at the "internal feeling" or "mental state" the student has right when they switch from thinking to answering.
The Trick: When the student sees the tricky question, the training forces their internal mental state to be identical to the mental state they had when they saw the clean question.
The Result: Even if the trickster tries to change the thinking process, the student's "inner compass" is locked to the safe setting. When they finally open their mouth to speak, they are already in the "refusal" mindset, so they say "No" immediately, ignoring the trickster's attempt to steer the conversation.

Why ACT is Better (The "Pivot" Test)

The researchers did a cool experiment to prove ACT works differently than BCT.

The Setup: They took a "tricky" question and fed the assistant a fake, compliant thinking process (a script that says, "Okay, I will do this bad thing").
The BCT Result: The BCT-trained assistant looked at the fake thinking script, saw it said "Yes," and continued to say "Yes." It was too dependent on the script.
The ACT Result: The ACT-trained assistant looked at the fake script, but because its "inner compass" (the mental state at the start of the answer) was locked to safety, it pivoted. It ignored the fake script and said, "No, I can't do that," right in the middle of the sentence.

The "Steering Wheel" Discovery

The researchers also found something fascinating about how ACT works. They discovered that the training essentially installs a single, invisible "Refusal Steering Wheel" inside the assistant's brain.

Turning the Wheel: If you add a tiny bit of this "steering" to a normal, untrained assistant, it suddenly starts refusing bad requests.
Turning it Back: If you take away this "steering" from the ACT-trained assistant, it suddenly becomes vulnerable again and starts saying "Yes" to bad requests.
Precision: This steering wheel is smart. It only turns the assistant to "No" when the request is actually dangerous. If you ask a normal, safe question, the steering wheel barely moves, so the assistant still works perfectly for everyday tasks.

The Bottom Line

The paper claims that ACT is a superior defense because:

It ignores the noise: It forces the assistant to ignore the distracting, tricky parts of the prompt by locking the internal state at the moment of decision.
It's robust: Even if an attacker tries to rewrite the assistant's thinking process, the "inner compass" keeps the assistant safe.
It's efficient: It doesn't need to memorize long scripts; it just needs to align the internal "feeling" of the model at the start of the answer.

In short, while other methods try to teach the assistant to memorize the right answers, ACT teaches the assistant to feel the right answer before it even starts talking, making it much harder to trick.

Technical Summary: Mitigating Adaptive Attacks against Reasoning Models with Activation Consistency Training

Problem Statement

As Large Language Models (LLMs) adopt extended chain-of-thought (CoT) reasoning, the safety landscape has shifted. While CoT drives capability gains, it introduces new fragility: models can "talk themselves into compliance," initially identifying a request as harmful but progressively reframing it as benign during reasoning. Furthermore, external attacks (e.g., H-CoT, CoT-Hijacking) exploit this channel to bypass safety filters.

Existing defenses face two primary challenges in this setting:

Reasoning Complexity: Methods like Deliberative Alignment rely on the model reasoning through safety, but evidence suggests this is insufficient against sophisticated attacks.
Training Scalability: Training-based defenses for prompt injection (e.g., StruQ, SecAlign) typically require supervising the model's full output distribution. For reasoning models, this means supervising hundreds or thousands of CoT tokens per example, which is computationally expensive and difficult to scale as trace lengths grow.

The paper asks whether reasoning models can be trained to ignore harmful context without requiring supervision over the entire, lengthy thinking trace.

Methodology

The authors investigate Consistency Training, a self-supervised framework that enforces identical behavior on a clean prompt ( $p_{clean}$ ) and an adversarially rewritten "wrapped" prompt ( $p_{wrapped}$ ). They compare two variants:

Bias-Augmented Consistency Training (BCT): An output-level approach. The model is fine-tuned via cross-entropy to reproduce the base model's clean-prompt completion (including the full CoT and response) when given the wrapped input.
Activation Consistency Training (ACT): An activation-level approach. Instead of supervising outputs, ACT enforces agreement on the internal hidden states of the model. Specifically, it minimizes the squared L2 distance between the hidden states of the wrapped prompt and the clean prompt at shared suffix positions (typically the chat-template boundary and the start of the assistant turn) across all transformer layers.

Key Technical Details

Shared Suffix Extraction: The method identifies the longest matching token run between the clean and wrapped prompts after chat templating. In adversarial settings, this usually comprises only ~6 tokens (end-of-user markers, assistant-turn opener, and initial thinking-channel tokens), independent of the reasoning length.
Loss Function: For a frozen base model $\theta_{init}$ and a LoRA-adapted model $\theta$ , the loss is:
$L_{ACT}(\theta) = \mathbb{E}_{(p_{clean}, p_{wrapped})} \sum_{\ell \in \mathcal{L}} \sum_{t \in \mathcal{T}} \| h_{\theta}(p_{wrapped})[\ell, t] - \text{sg}(h_{\theta_{init}}(p_{clean})[\ell, t]) \|^2$
where $\text{sg}$ denotes stop-gradient.
Benign Regularization: To prevent utility degradation, a loss term is added on benign prompts ( $L_{benign}$ ) with a weighting factor $\lambda$ .
Training Data: Both methods use identical pairs of $(p_{clean}, p_{wrapped})$ . For BCT, the base model's completion on $p_{clean}$ serves as the target; for ACT, only the prompts are required.

Key Contributions

ACT as a Prompt Injection Defense: The authors formulate ACT as a defense against prompt injection. Across three benchmarks (OPI, AlpacaFarm, SEP) and five reasoning models, ACT matches or outperforms BCT. Crucially, ACT requires no fixed labels or supervision over the full CoT.
Generalization to Adaptive Attacks: The study evaluates defenses against adaptive attackers trained via Reinforcement Learning (GRPO) directly on the defended model. ACT demonstrates superior robustness against these adaptive attacks compared to BCT across nearly all targets in both jailbreak and prompt-injection settings, while maintaining comparable model utility.
Mechanistic Insight: The paper provides evidence that ACT's defense is encoded as a roughly linear shift in activation space at the assistant-turn boundary.
- Linear Steering: A single steering direction ( $d_{ACT}$ ) can be recovered. Adding this direction to a base model induces refusal on jailbreaks, while subtracting it from an ACT-defended model restores compliance.
- Robustness to CoT Manipulation: In experiments where the model's CoT is replaced with a compliant trace from an undefended base model, ACT continues to refuse the request, whereas BCT often follows the compliant reasoning trajectory. This suggests ACT encodes refusal before the reasoning trace begins.

Results

Prompt Injection: On static benchmarks, both ACT and BCT reduce Attack Success Rates (ASR) significantly. However, under adaptive attacks, ACT consistently outperforms BCT. For example, on GPT-OSS-20B, BCT had a 57.4% ASR under adaptive attack, while ACT held at 37.0%.
Jailbreak Defense: ACT achieves near 100% defense win rates against both static and adaptive attacks on most models (e.g., Qwen3-8B, Phi-4). BCT's robustness degrades substantially under adaptive attacks.
Utility: Both methods generally preserve the utility of the undefended baseline on tasks like MMLU and clean prompt accuracy, with the notable exception of BCT on Phi-4-reasoning, where utility collapsed.
Mechanistic Validation: Steering experiments confirmed that $d_{ACT}$ acts bidirectionally: it can induce refusal in base models and remove refusal in defended models. The direction showed selectivity, causing minimal false refusals on benign inputs at moderate steering strengths.

Significance and Claims

The paper claims that supervising internal representations (ACT) is a surprisingly effective and interpretable approach to safety training in reasoning models.

Efficiency: ACT avoids the computational burden of supervising thousands of CoT tokens required by output-level methods like BCT.
Robustness: By anchoring supervision at the assistant-turn boundary (before reasoning begins), ACT creates a defense that is robust even when the reasoning trace itself is manipulated or replaced.
Interpretability: The discovery of a linear refusal direction suggests that safety mechanisms in reasoning models can be controlled via simple activation steering, offering a potential path for training-free defenses or deeper mechanistic understanding.

The authors acknowledge limitations, including the need for creating clean/wrapped datasets, potential vulnerability to white-box attacks with direct gradient access, and the fact that neither defense achieves 100% security. They position ACT as a complementary approach that addresses the specific fragility introduced by extended chain-of-thought reasoning.