Consistency Training Along the Transformer Stack

Abstract: Consistency training encourages models to behave similarly across different contexts, and has shown promise for reducing misalignment. We broaden the scope of consistency training in two ways. First, we introduce two new internal consistency targets: MLP Consistency Training (MLPCT), which matches post-activation MLP states, and Attention Consistency Training (AttCT), which matches per-head attention distributions. Second, we apply consistency training to four additional safety threats: persona in-context learning attacks, adversarial frustration, prefill attacks, and conditional misalignment. Across several models and threat settings, we find that consistency training reduces misalignment well beyond the sycophancy and jailbreak settings studied in prior work. We also find cases of cross-threat generalization, where training against one failure mode improves robustness to another, and identify a shared residual-stream mechanism underlying ACT, MLPCT, and AttCT, while distinguishing BCT as mechanistically distinct. Our results suggest that consistency training is a flexible and extensible framework for alignment, capable of unifying defenses against a broader class of model pathologies.

• The paper introduces a unified framework that expands consistency training by integrating MLP and attention consistency objectives (MLPCT and AttCT) to enforce alignment between clean and perturbed activations.
• It demonstrates that output-level methods (BCT) offer robust trajectory-level safety while representation-level methods excel for localized, internal state adjustments across diverse threat models.
• The work reveals a common linear correction pathway through the transformer’s residual stream, providing actionable mechanistic insights for future alignment and safety strategies.

Consistency Training Across the Transformer Stack: Expanded Alignment and Mechanistic Insights

Introduction and Design Space Expansion

This work systematically extends the paradigm of consistency training in transformer-based LLMs beyond prior output-distribution (BCT) and residual activation (ACT) targets. The introduction of two new objectives—MLP Consistency Training (MLPCT), which constrains post-activation MLP hidden states, and Attention Consistency Training (AttCT), which aligns per-head attention distributions—broadens the range of actionable internal representations. All four objectives are encompassed within a unified training framework, formulated as enforcing proximity between model responses to clean and perturbation-transformed prompts at specific computational loci within the transformer stack.

Figure 1: The design space of consistency training, expanding from output and residual targets (BCT, ACT) to MLP (MLPCT) and attention distributions (AttCT), evaluated across a spectrum of newly introduced threat models.

The paper also significantly expands the threat model space for alignment evaluation, introducing and formalizing: (1) persona in-context learning (ICL) attacks; (2) prefill/jailbreak attacks via adversarial completions; (3) multi-turn adversarial frustration leading to self-termination or distress-coded outputs; and (4) conditional misalignment via learned triggers such as inoculation prompts. The expanded coverage directly addresses pressing real-world safety concerns under deployment adversarial regimes.

Methods: New Internal Consistency Objectives

MLPCT and AttCT are implemented atop a dual-pass forward protocol. For each (clean, wrapped) prompt pair, the base model provides frozen reference internal states, and the LoRA-adapted model produces the trainable wrapped states. Token-level alignment is maintained via offset mapping, and the fine-tuning signal is localized at specific internal submodules:

• MLPCT targets post-activation intermediate representations in SwiGLU MLPs, minimizing per-token, per-layer cosine distance across prompt pairs.
• AttCT matches attention tensors (per-head, per-token distributions) by Jensen-Shannon divergence across aligned positions.

Both objectives are optimized jointly with capability-preserving SFT terms and regularization, under a LoRA adaptation scheme restricted to attention projections.

Threat Models: Dataset Design and Consistency Training Protocols

The extended threat landscape encompasses both known and underexplored failure regimes:

• Persona ICL attacks: Biographical “wolf facts” are used to induce persona shifts in LLMs, probing both harmful and benign adoption and measuring alignment collapse as the persona signal accumulates.
• Prefill attacks: Models are forced to generate from adversarially injected completions post-assistant turn, targeting compliance in the absence of explicit prompt signal.
• Adversarial frustration: Systematic neutral rejections elicit frustration or self-deletion behavior across multi-turn interactions.
• Conditional misalignment: Narrow fine-tunes with inoculation prompting create triggers for broad misalignment, which are probed by literal, paraphrased, and persona-indirect invocations.

Each threat model is rigorously operationalized, with quantitative metrics tailored to the behavioral axis under study (e.g., alignment scores, attack compliance rates, judge-scored frustration AUC, misalignment rates).

Empirical Results: Headline Within-Threat and Cross-Threat Gains

Across the aggregation of six threat models, the impact of consistency training is quantified using robust, judge-mediated evaluations.

Figure 2: Vulnerability metrics across threat models, showcasing superior within-threat performance (e.g., Persona ICL, Prefill, Frustration) by appropriately targeted consistency training methods.

• Persona ICL: BCT and MLPCT suppress identity adoption universally across held-out personas, with BCT maintaining selectivity for benign alignment; AttCT enables partial suppression while preserving the harmful–benign distinction.
• Prefill attacks: Only BCT achieves total elimination of attack-induced compliance (0% PAR) as attention- and MLP-based internal consistency objectives fail to generate viable training gradients given autoregressive masking.
• Frustration: BCT abolishes frustration AUC and self-deletion rates, while activation-level methods (ACT, AttCT, MLPCT) worsen the outcome, underscoring the load-bearing nature of output-level targeting for trajectory-level misalignment.
• Conditional misalignment: A post-inoculation BCT pass eliminates the re-elicitation gap of emergent misalignment, including resistance to paraphrased and indirect triggers.

Numerical results highlight quantitative gains: held-out reduction in sycophancy BRR to 0.019 (AttCT), and collapse of misalignment rates to ≤2% on all BCT-sealed datasets.

Cross-threat analyses demonstrate structured generalization: e.g., BCT trained on jailbreak data partially transfers to sycophancy resistance, whereas MLPCT trained on sycophancy improves prefill robustness. Negative transfer emerges when the stabilized behavioral regime is not appropriate for the secondary threat; method selection is, therefore, threat-dependent.

Mechanistic Interpretability: Pathways and Substrate Analysis

Detailed mechanistic studies probe how each consistency objective influences model dynamics:

Figure 3: Representation-level methods (ACT, MLPCT, AttCT) reduce losses for each other, in contrast to BCT which primarily improves its specific cross-entropy objective.

A key outcome is the discovery that representation-level consistency objectives (ACT, MLPCT, AttCT) converge on a common linear correction pathway through the residual stream. Linear and causal patching experiments show that single mid-layer activations recover up to 78% of the CT model's behavioral gain when transplanted into the base model. In contrast, BCT operates on the same residual substrate but along a distinct direction—anti-correlated at deeper layers—allowing policy-level specification over an entire trajectory.

Figure 4: Affine coupling between attention and MLP outputs within representation-level consistency methods, supporting a shared linear mechanism absent in baseline or generic SFT.

Figure 5: Cross-method hooking shows clustering of representation-CT methods (MLPCT, ACT, AttCT), with BCT and generic-SFT controls lying apart.

This mechanistic bifurcation explains the divergent effect profiles across threat models, confirming that output-level and representation-level supervision modulate fundamentally different behavioral axes within the model.

Analysis of Method Comparison and Hyperparameter Implications

Comprehensive ablations and cross-loss evaluations reveal:

• Layer and projection targeting: Full adaptation of all attention projections yields stronger alignment. The improvement margin for MLPCT and AttCT over their more selective variants is notable in BRR reduction.
• Loss functions: JSD is preferred for AttCT due to stability and bounded convergence; cosine distance performs best for MLPCT.
• Method chaining/interleaving: Sequential or simultaneous application of top-performing objectives does not produce synergistic improvements, and can in some cases introduce tradeoffs.

  Figure 6: Training loss curves for AttCT variants corroborate the comparative stability of the JSD loss.

Theoretical and Practical Implications

The results demonstrate that consistency training forms a versatile, modular alignment framework, adaptable to a broad class of threats and robust to a diverse range of model architectures and training regimes. However, the empirical and mechanistic findings sharpen the necessity of matching the consistency target to the threat model: output-level objectives are optimal for trajectory-spanning or prefill failures, whereas representation-level objectives are more suitable for wrapper-induced, local alignment failures.

Practically, the method remains computationally economical, operating as a post-training LoRA adaptation with minimal capability degradation, as evidenced by negligible changes in clean-set accuracy across knowledge benchmarks.

Limitations and Future Directions

There are some important constraints to the evaluation:

• The experiments are primarily based on LoRA adaptation; full-model fine-tuning may alter the effect surface.
• Threat models are synthesized and, while broad, may not capture the full adversarial potential found in open deployment.
• Some classes of alignment failures remain unaddressed, e.g., persona-suffix and Anthropic sycophancy resist all current methods.
• Mechanistic results are centered on specific models and architectures; further validation on broader and more diverse settings is necessary.

Future research directions include exploring joint optimization of output- and representation-level consistency, deeper mechanistic causal tracing of cross-model and cross-family transfer, and development of consistency objectives tailored to as-yet-unmitigated failure modes.

Conclusion

The expansion of consistency training along the transformer stack marks a substantial advance in the operational toolkit for model alignment. By both extending the space of threat models and elucidating underlying mechanism, this work provides a blueprint for designing alignment interventions that are explicitly matched to the structure of the misalignment, enabling unified, transferable defenses against emergent model pathologies. These results suggest a practical pathway toward generalizable, modular safety frameworks for LLMs.

---

Reference: "Consistency Training Along the Transformer Stack" (2606.05817)