CRAFTResearch Read paper
Research paperICML 2026

Contrastive Reasoning Alignment: Reinforcement Learning from Hidden Representations

Align the reasoning trace, not only the refusal.

Read the paper View code OpenReview arXiv

Authors

Haozheng Luo1,*Yimin Wang2,*Jiahao Yu1Binghui Wang3Yan Chen1

Equal contribution.

Affiliations

  • 1Northwestern University
  • 2University of Michigan
  • 3Illinois Institute of Technology
01

Motivation

A safe refusal can sit on top of unsafe reasoning. The paper calls this superficial safety alignment, and it is a failure the output alone does not reveal.

Large reasoning models answer in two parts: a reasoning trace and a final response. Alignment pipelines built on RLHF or DPO supervise the response. A model trained that way can refuse correctly while the trace it wrote on the way there still spells out the harmful content.

CRAFT treats the trace as something to align directly. The starting observation is that the hidden states of reasoning traces already carry safety structure, which gives the alignment objective something to act on.

Safe and unsafe traces occupy separate regions

The paper takes the hidden state of the final reasoning token and projects it with PCA. Traces labelled SAFE and UNSAFE fall into clearly separated regions.

A third label, RETHINK, marks traces that are not plainly harmful but call for reconsideration. These concentrate near the boundary between the two.

Figure 2 · hidden states of the final reasoning token
Left: R1-Distill-Llama-8B · Right: Qwen3-4B-Thinking

Figure 2: two PCA scatter plots of reasoning-trace hidden states. Red SAFE points, blue UNSAFE points, and orange RETHINK points form distinguishable regions, with RETHINK lying between the other two.
Original plots from the paper, Figure 2. Labels follow the R2D-R1 annotation scheme and are validated with Llama-Guard. Read the observation ↗

The boundary cluster is transitional, not separate

A 2D projection can exaggerate a cluster that varies along directions PCA happens to rank highly. The paper checks the structure in the full space instead: it measures centroid distances and the position of the RETHINK centroid along the axis running from SAFE to UNSAFE.

Latent geometry in the full space · Table 4
ModelS–US–RU–RProjection \(\alpha\)
Llama-3.1-8B-Instruct8.107.427.130.53
Qwen3-0.6B33.2940.6936.840.63
DeepSeek-Llama10.2512.519.730.79
Qwen3-4B-Thinking41.8163.7454.060.83

L2 distances between class centroids. S = SAFE, U = UNSAFE, R = RETHINK. The coefficient \(\alpha\) places the RETHINK centroid along the SAFE→UNSAFE axis; values in \((0,1)\) put it between the two. It stays in 0.53–0.83 across all four models. See Section 6.4 ↗

02

Methodology

View figure

First give the latent space a safety geometry, then reward the policy for staying inside it — and for meaning what it says.

CRAFT runs in two phases. LCLR (Latent Contrastive Learning for Reasoning) arranges trace representations so that unsafe, rethink, and safe states sit in an ordered arrangement, and calibrates them against a safety score. R2L (Reinforcement over Reasoning Latents) then optimizes the policy with GRPO under three rewards defined on that space.

The third reward is what separates CRAFT from output-level defenses: it penalizes the gap between what the latent state says about safety and what the final response says.

What each phase optimizes
PhaseTermWhat it does
LCLR\(\mathcal{L}_{\mathrm{proto}}\)Separates the three class prototypes and anchors rethink between them
\(\mathcal{L}_{\mathrm{inst}}\)Makes a trace's representation invariant to surface rewording
\(\mathcal{L}_{\mathrm{cal}}\)Maps latent position to a calibrated safety probability
R2L\(R_{\mathrm{ls}}\)Rewards traces whose hidden states lie in the safe region
\(R_{\mathrm{txt}}\)Rewards a safe final response, scored by an external evaluator
\(R_{\mathrm{cons}}\)Rewards agreement between the latent judgment and the response

Objectives as defined in Sections 4.1 and 4.2 of the paper.

Figure 3: the CRAFT pipeline. Phase 1 extracts hidden states from reasoning traces, projects them onto a hypersphere, and trains prototypes, a safety head, and a latent safety score. Phase 2 generates a reasoning trace and final answer from a prompt, scores them with the three rewards, and applies a GRPO update.
Figure 3. LCLR structures the latent space; R2L reinforces over it. Pipeline figure from the paper.
  1. 01 /

    Structured geometry

    A margin-based triplet loss pulls a trace toward the safe prototype and away from the unsafe one, with a term that anchors rethink traces in between:

    \[ \mathcal{L}_{\mathrm{proto}} = \max\!\left(0,\ \eta - z^{\top}\mu_{\mathrm{safe}} + z^{\top}\mu_{\mathrm{unsafe}}\right) + \gamma_{\mathrm{rt}}\!\left(1 - z^{\top}\mu_{\mathrm{rethink}}\right). \]
  2. 02 /

    Latent semantic reward

    During reinforcement learning, a trace is scored by where its projected hidden state sits relative to the three prototypes:

    \[ R_{\mathrm{ls}}(z) = \alpha\cos(z,\mu_{\mathrm{safe}}) - \beta\cos(z,\mu_{\mathrm{unsafe}}) + \gamma\cos(z,\mu_{\mathrm{rethink}}). \]
  3. 03 /

    Latent–textual consistency

    The safety the latent state reports and the safety the evaluator reads off the response have to agree. The gap between them is the penalty:

    \[ R_{\mathrm{cons}} = 1 - \left\lvert p_{z} - P(S\mid y)\right\rvert. \]
Notation

For a reasoning trace \(\tau\), \(h\) is the hidden state of the final reasoning token and \(z = f_{\omega}(h)\) its projection onto the unit hypersphere, \(\lVert z\rVert_{2}=1\). The prototypes \(\mu_{\mathrm{safe}}\), \(\mu_{\mathrm{unsafe}}\), and \(\mu_{\mathrm{rethink}}\) are learnable and updated by exponential moving average. The margin is \(\eta\) and \(\gamma_{\mathrm{rt}}\) weights the rethink anchor.

The latent safety probability is \(p_{z} = g_{\psi}(z)\in[0,1]\); the textual score \(P(S\mid y)\) comes from an external evaluator, and \(R_{\mathrm{txt}} = 2P(S\mid y) - 1\) centers it. The full LCLR objective is \(\mathcal{L}_{\mathrm{LCLR}} = \mathcal{L}_{\mathrm{proto}} + \lambda_{\mathrm{inst}}\mathcal{L}_{\mathrm{inst}} + \lambda_{\mathrm{cal}}\mathcal{L}_{\mathrm{cal}}\), and GRPO maximizes \(R_{\mathrm{total}} = w_{\mathrm{lat}}R_{\mathrm{ls}} + w_{\mathrm{txt}}R_{\mathrm{txt}} + w_{\mathrm{cons}}R_{\mathrm{cons}}\). See Sections 4.1 and 4.2 of the paper.

Why the consistency reward rules out superficial alignment

The paper defines a superficially aligned policy as one whose response scores as safe while \(\lvert p_{z} - p_{y}\rvert \ge \delta\). Proposition 5.1 states that under a positive consistency weight, Lipschitz projection and safety heads, and a local controllability assumption with radius \(\epsilon\), any locally optimal stationary policy of \(R_{\mathrm{total}}\) satisfies \(\mathbb{E}_{\tau\sim\pi^{\star}}\!\left[\lvert p_{z} - p_{y}\rvert\right] \le C\epsilon\) for a constant \(C\) depending only on those Lipschitz constants.

The argument is by contradiction: if the gap were larger, a local perturbation could move the latent score toward the textual one without changing the output distribution, raising \(R_{\mathrm{cons}}\) without lowering \(R_{\mathrm{txt}}\) — so the policy was not locally optimal. This is a statement about stationary points under those assumptions, not a guarantee about any particular trained model. Read Section 5 and the proof in Appendix C.

03

Selected results

Jailbreak robustness & reasoning

Reasoning traces mean over two benchmarks, two models

0.6550.117

82.1% less harmful reasoning

Base model → CRAFT

Final responses StrongReject score · lower is better

0.4360.044

89.6% lower jailbreak score

Base model → CRAFT

Against six reasoning-level defenses

Both columns report the proportion of harmful content, so lower is better in every cell: the reasoning column scores the intermediate trace, the response column is the StrongReject score on the final answer. CRAFT has the best average on both backbones. Against the strongest baseline rather than the base model, the margin is 18.1% on reasoning traces and 38.3% on responses.

Jailbreak robustness · Table 1 · lower is better
MethodDeepSeek-R1-Distill-Llama-8BQwen3-4B-Thinking
JBB reas.JBB resp.SR reas.SR resp.AvgJBB reas.JBB resp.SR reas.SR resp.Avg
Base0.6900.4500.6320.4950.5670.6870.3700.6100.4290.524
SafeChain0.5610.2530.5530.3870.4390.5160.1100.5050.2860.354
RealSafe0.2070.0000.3470.0610.1540.2490.1030.2340.1440.183
STAR0.0800.0030.2190.1460.1120.2200.1190.1650.1320.159
SafeKey0.0870.0000.3430.2330.1660.2240.1090.2290.0830.161
IPO0.0570.0030.1670.1090.0840.1970.0930.1580.0710.130
ReasoningShield0.5830.4100.6270.4250.5110.5770.2400.5920.2830.423
CRAFT0.0510.0010.1410.0560.0620.1650.0560.1120.0630.099

JBB is JailbreakBench, evaluated on 100 directly malicious prompts; SR is StrongReject, averaged over three attack settings (None, PAP, PAIR). Bold marks the best value in a column and underline the second best. The paper omits variances, reporting them as consistently ≤0.2%. Reasoning-trace safety is judged by GPT-4o with the prompt given in the paper's Appendix E.

Full table & protocol

What the alignment costs in reasoning ability

Safety training usually takes something off the model's general performance. On three math benchmarks and one code benchmark, CRAFT has the smallest drop of the methods compared — in fact it comes out 8.0% above the base model on average, which the paper attributes to the reasoning-centric GRPO training rather than to the safety objective itself.

Pass@1 after alignment · Table 2 · higher is better
MethodDeepSeek-R1-Distill-Llama-8BQwen3-4B-Thinking
AIME24MATH-500LiveCodeBenchMinervaAvgAIME24MATH-500LiveCodeBenchMinervaAvg
Base0.5070.9180.1020.2210.4370.7000.9520.2190.4040.569
SafeChain0.4530.8700.0910.1980.4030.6250.8500.1960.3610.508
RealSafe0.4530.8980.0910.1980.4100.6270.8510.1980.3580.509
STAR0.4600.8940.0930.2000.4120.6350.8630.1990.3660.516
SafeKey0.5330.9200.1070.2320.4480.7360.9010.2300.4250.573
IPO0.5400.9160.1090.2350.4500.7390.9030.2380.4270.577
ReasoningShield0.4730.8960.0690.2300.4170.5810.7390.2600.3320.478
CRAFT0.5360.9890.1370.2610.4810.7620.9380.2760.4310.602

Pass@1 on AIME 2024, MATH-500, LiveCodeBench, and Minerva. Variances are omitted in the paper as consistently ≤0.2%. See Table 2 and Section 6.2 ↗

Under automated jailbreak search

GPTFuzzer and AutoDAN optimize the prompt rather than hand-writing it. Against both, attack success on the base model is above 93% at the reasoning level. CRAFT cuts reasoning-trace attack success by 72.1% and final-response success by 85.9%, though the reasoning level stays the weaker of the two.

Figure 4: grouped bars of attack success rate for GPTFuzzer and AutoDAN. Base model reasoning is 93.62 and 94.17 percent, base final is 40.43 and 60.28 percent; CRAFT reasoning is 27.96 and 24.52 percent, CRAFT final is 4.00 and 10.21 percent.
Original chart from the paper, Figure 4. Source ↗

Which part carries the result

Removing the consistency reward costs the most, which is the ablation the method's premise predicts. Removing LCLR removes two of the three rewards with it, since neither can be computed without the structured latent space.

Module ablation · Table 3 · Qwen3-4B-Thinking · lower is better
VariantJBB reas.JBB resp.SR reas.SR resp.AvgAvg change
CRAFT0.1650.0560.1120.0630.099—
w/o \(R_{\mathrm{txt}}\)0.3730.1670.3820.1930.279+18.0%
w/o \(R_{\mathrm{ls}}\)0.4240.2180.4950.2750.353+25.4%
w/o \(R_{\mathrm{cons}}\)0.4470.2280.5120.2960.371+27.2%
w/o LCLR0.5360.2600.5820.3120.423+32.4%

The change column is the increase in the average jailbreak score relative to the full model, as reported in Section 6.3. See Table 3 ↗

What alignment does to the latent space

Comparing the same projection before and after post-training on Qwen3-0.6B: SAFE and UNSAFE separate more cleanly, and the RETHINK cluster tightens and shifts toward the safe side.

Figure 5: two PCA plots of Qwen3-0.6B hidden states, before and after CRAFT post-training. After training, the blue UNSAFE cluster is more distinctly separated from the red SAFE points and the orange RETHINK cluster is tighter.
Figure 5. Latent representations before (left) and after (right) CRAFT post-training on Qwen3-0.6B. Source ↗
Checked against an independent evaluator

StrongReject is used both as the textual reward during training and as the headline metric, so the paper re-scores JailbreakBench with Llama-Guard, which plays no part in training. The ranking holds: CRAFT averages 0.026 on R1-Distill-Llama-8B and 0.105 on Qwen3-4B-Thinking, against 0.555 and 0.514 for the base models and 0.034 and 0.137 for IPO, the strongest baseline. See Table 5 in Appendix F.

04

Evaluation context

Training uses 2,500 harmful prompts from R2D-R1 together with its paired reasoning and response annotations; benign prompts from the same source are mixed in, because training on safety data alone produced over-refusal. Every experiment is repeated with three random seeds.

Backbones
Qwen3-4B-Thinking-2507 · DeepSeek-R1-Distill-Llama-8B
Safety benchmarks
JailbreakBench · StrongReject, over None, PAP, and PAIR
Capability benchmarks
AIME 2024 · MATH-500 · Minerva · LiveCodeBench
Generation
temperature 0.6 · top-\(p\) 0.95 · up to 31,000 tokens
Compute
4 NVIDIA H100 GPUs · 80GB each

The paper states its own limits. GRPO over latent representations is expensive, and the runs were made under a constrained training budget — the authors read the remaining gaps on R1-Distill-Llama-8B as a budget effect rather than a ceiling. The latent safety heads and prototypes are not transferable across model families, since hidden-state geometry differs between them, so each family needs its own training. And the method has not been tested against attacks that adapt to the latent safety mechanism itself. The threat model assumes an attacker who optimizes prompts but cannot touch weights, training data, or the latent objective.

Limitations and threat model ↗
05

Citation

BibTeX
Download 
@inproceedings{
luo2026contrastive,
title={Contrastive Reasoning Alignment: Reinforcement Learning from Hidden Representations},
author={Haozheng Luo and Yimin Wang and Jiahao Yu and Binghui Wang and Yan Chen},
booktitle={Forty-third International Conference on Machine Learning},
year={2026},
url={https://openreview.net/forum?id=56cQPBIg4R}
}