Mitigating Adaptive Attacks against Reasoning Models with Activation Consistency Training

Abstract: As LLMs gain stronger reasoning capabilities, their extended chain-of-thought introduces new degrees of complexity for defending against adversarial jailbreaks and prompt injection. We study consistency training, a family of fine-tuning objectives that enforce identical behavior on clean prompts and adversarial rewrites, and evaluate its two main variants, output-level (BCT) and activation-level (ACT), across five reasoning models. We formulate both methods as a prompt injection defense and find ACT to be competitive with other training-based defenses while requiring only self-supervised pairs of clean and wrapped prompts. Our experiments also generalize both techniques within the jailbreak setting, demonstrating that ACT remains more robust to adaptive attacks. We also provide mechanistic evidence that ACT's defense against jailbreaks is encoded as a roughly linear shift in activation space at the assistant-turn boundary. After ACT training, we can recover a single steering direction that controls refusal on reasoning models with minimal effect on benign inputs. We find that ACT remains robust even when the model's chain-of-thought is replaced with a compliant trace from the undefended base model, pivoting to refuse prefilled jailbreaks. Together, these results suggest that supervising internal representations is a surprisingly effective and interpretable approach to various forms of safety training in reasoning models.

• The paper introduces Activation Consistency Training (ACT) that defends reasoning LLMs against adversarial prompt injections by aligning hidden activations on shared template boundaries.
• It compares ACT with Bias-Augmented Consistency Training (BCT) and shows that ACT offers superior robustness under both static and adaptive attacks.
• The results demonstrate that ACT preserves clean-task utility while enabling interpretable refusal control through a linear steering mechanism in the activation space.

Mitigating Adaptive Attacks on Reasoning Models via Activation Consistency Training: An Expert Synthesis

Introduction

Defending LLMs with explicit chain-of-thought (CoT) reasoning remains an open challenge, especially under adversarial settings such as jailbreaks and prompt injection. While CoT has yielded dramatic gains in reasoning capability, it has simultaneously introduced new attack surfaces where adversaries can leverage both direct user input and the model's own intermediate reasoning to bypass safety alignments. Existing training-based defenses—including output-level approaches—struggle with the exponentially increasing supervision demands incurred by long reasoning traces. This work proposes a rigorous evaluation and mechanistic analysis of two self-supervised fine-tuning methods—Bias-Augmented Consistency Training (BCT) and Activation Consistency Training (ACT)—designed specifically to enforce identical model behavior on clean and adversarially perturbed (wrapped) prompts. Its contributions encompass method formulation, large-scale empirical analyses under static and adaptive attacks, and a comprehensive elucidation of the internal mechanisms underpinning ACT's superior defensive properties.

Methodological Framework

Consistency Training Paradigm

Both BCT and ACT operate on paired examples $(p_\text{clean}, p_\text{wrapped})$, where $p_\text{wrapped}$ represents a clean prompt adversarially modified with distractors or direct injections. BCT enforces alignment at the output level by fitting the model to reproduce the base-model's clean-prompt output (including its entire CoT trace), creating a cross-entropy loss over all tokens in the shared output region. Conversely, ACT supervises the squared distance in activation space at matching "shared suffix" token positions, enforcing invariance at all transformer layers but only over the short (typically $\sim$6 tokens) chat-template boundaries shared between clean and injected prompts.

(Figure 1)

Figure 1: Illustration of ACT: It minimizes the $L_2$ distance between hidden states on shared-suffix positions, forcing the wrapped prompt’s activations to match those of the clean prompt and thereby preventing injected content from altering the model’s representation at the response boundary.

This approach leverages the fact that, in adversarial settings, the only overlapping segment between $p_\text{clean}$ and $p_\text{wrapped}$ is the template boundary and assistant prompt initiator, thereby localizing the point of enforcement. Both methods incorporate benign regularization via additional loss on clean pairs to prevent utility collapse on standard instructions.

Training and Evaluation Protocol

Five model families are benchmarked: Qwen3-1.7B, Qwen3-8B, Phi-4-reasoning, GPT-OSS-20B, and Gemma-4-E4B-it. Three prompt-injection benchmarks are used: Open-Prompt-Injection (OPI), AlpacaFarm, and SEP. Both static (template-based) and adaptive (reinforcement-learning-trained attacker) scenarios are tested. Adaptive attackers are implemented using GRPO, conditioned to maximize attack success on each defended target. The StrongReject finetuned ranker provides a reliable automated metric for attack success (ASR). Practical utility is measured via MMLU and clean task accuracy.

Defense Efficacy Against Prompt Injection

Figure 2: Static prompt-injection defense win rate ($1-\mathrm{ASR}$) across multiple models and benchmarks: both ACT and BCT significantly reduce ASR from baseline; exceptions arise for Phi-4-reasoning and GPT-OSS-20B where BCT occasionally outperforms.

Static evaluations reveal that both ACT and BCT successfully bring ASR close to zero for Qwen3 and Gemma models. BCT marginally outperforms ACT on GPT-OSS-20B in static OPI, but this advantage reverses under adaptive threats. Phi-4-reasoning remains most vulnerable, with residual ASR in the double digits for both defenses.

Figure 3: Adaptive prompt-injection defense win rate ($1-\mathrm{ASR}$): ACT is consistently more robust than BCT under adaptive (RL-trained) attackers, except for Phi-4-reasoning; both methods are ineffective on Gemma-4-E4B-it.

Under adaptive attackers, ACT outperforms BCT on almost all models and achieves notable ASR reductions compared to the unprotected baseline, exposing BCT’s greater vulnerability to distributional shifts and demonstrating ACT's robustness. Notably, ACT’s improvement over BCT widens for models where BCT succeeded under static templates but overfit to attack surface specifics.

Figure 4: Utility preservation for defended checkpoints: ACT and BCT generally match the undefended baseline on clean tasks; utility degradation for OPI clean accuracy is severe for BCT on Phi-4-reasoning.

Both methods preserve utility for most models and tasks—with ACT typically sacrificing slightly more performance for superior robustness. For Phi-4-reasoning, BCT’s clean-task utility collapses unless hyperparameters are retuned at the cost of defense.

Jailbreak Robustness

Figure 5: Jailbreak defense win rate on static and adaptive attacks: ACT achieves near-complete defense, BCT degrades sharply under adaptive conditions except on GPT-OSS-20B.

On a suite of 200 harmful behaviors, both defenses decrease ASR substantially. However, ACT yields $\leq4\%$ ASR on 4/5 models on static attacks, while BCT’s performance regresses considerably under adaptive attack—particularly on Qwen-based and Gemma targets—exhibiting overfitting to training trace particulars. ACT proves resilient even when attackers optimize against it directly. Utility (MMLU) remains within a few accuracy points of baseline for both defenses.

Mechanistic Interpretability and Intrinsic Refusal Control

Resistance to CoT Prefilling

An important mechanistic finding is that ACT’s defense is not tied to the model’s own reasoning trace: even when the CoT is forcibly replaced by a compliant trace from the base model that concludes harmful compliance, ACT-trained models reliably pivot to refusal at the response boundary.

Figure 6: Refusal rate under compliant chain-of-thought prefilling: ACT-defended models pivot to refusal ($\geq98\%$) even when the CoT is adversarially prefilled; BCT fails and complies on most models.

This indicates that ACT encodes refusal upstream of chain-of-thought generation, conferring robustness to “CoT hijacking” attacks that co-opt intermediate reasoning, a failure mode widespread in existing SFT and output-level preference optimization approaches.

Linear Steering Direction

Moreover, ACT encodes a linear steering direction at the assistant-turn boundary: extracting the per-layer mean difference vector between defended and base-model activations, one can bidirectionally induce or abolish refusal with a single vector manipulation at inference. Adding this direction to the base model triggers refusal on jailbreaks, while subtracting from ACT recreates compliance—without impacting benign prompts for small magnitude shifts. The effect is monotonic, selective, and reproducible across model families with some architecture dependence, evidencing a localized and interpretable mechanism for refusal control induced by ACT.

Theoretical and Practical Implications

These results collectively demonstrate that supervising internal representations at matched template positions is an efficient, scalable mechanism for aligning CoT-based reasoning models, outperforming output-focused and DPO-style preference training under challenging out-of-distribution adversaries. The mechanistic clarity of the steering direction aligns with recent results in interpretable safety, suggesting model refusal can be isolated and manipulated in embedding space. Practically, ACT enables robust defenses without enumerating all attack variants or supervising over expanding reasoning tokens, and could in principle be combined with output-level methods for further compositional robustness.

Conclusion

Activation Consistency Training constitutes a technically sound, scalable approach to fortifying reasoning LLMs against both prompt injection and jailbreak attacks—particularly under adaptive adversarial pressure. Its supervision of hidden activations at shared template boundaries not only achieves superior robustness than output-level defenses, but does so while preserving utility and furnishing a mechanistically interpretable “refusal direction.” These insights open avenues for compositional defense stacks, targeted supervision of subspaces, and inference-time defense strategies based on recovered activation directions. Future work should extend these methods to multi-step agentic deployments and refine mechanistic localization of refusal features at finer architectural granularity.