Skip to content

Repository files navigation

Reasoning Distillation and the Refusal Mechanism: Layer, Time, and Transferability

Authors: Ankita Tripathi, Ali Alshehhi, Athul Radhakrishnan

Full paper: refusal_in_reasoning_models.pdf

Mechanistic study of the refusal mechanism across three Qwen2.5-1.5B variants that share architecture but differ in training: a base model, a safety-aligned instruction-tuned model (Qwen2.5-1.5B-Instruct), and a reasoning-distilled model (DeepSeek-R1-Distill-Qwen-1.5B). The project extracts per-layer refusal directions, ablates them, transplants them across models, and decomposes their causal effect by generation phase.

The codebase is a reimplementation of the single-direction analysis pipeline introduced by Arditi et al., 2024, with explicit choices for extraction position, layer-selection criterion, and intervention form, plus a </think>-aware LLM-judge classifier for reasoning-model outputs and a basis-drift null distribution for cross-model cosine claims.

Key Findings

  1. Extraction position is decisive for Instruct. The canonical user_end extraction selects a non-causally-effective L1 direction (93.75% → 87.5% under ablation, 6 pp drop), while the last_full extraction recovers a strongly causally-effective L19 direction (93.75% → 0%, 94 pp drop). The chat-template-tail position carries the alignment-installed refusal direction in this model.
  2. R1-Distill's L22 refusal direction does its causal work inside <think>...</think>. Phase-aware ablation gates the L22 ablation by generation phase; the think-only ablation reproduces the full-rollout ablation cell-for-cell across all four label categories, while the answer-only ablation is byte-for-byte identical to the no-intervention baseline.
  3. Cross-model transplant induces refusal but is not safety-selective. Injecting Instruct's L19 last_full direction into R1-Distill at L=22 with α=0.25 (perplexity-safe at 1.13× baseline) raises JBB-harmful refusal from 15% to 59% (+44 pp), but raises XSTest-safe refusal from 0.8% to 44.4%. A direction-specificity negative control with Instruct's non-causally-effective L1 direction at the same cell produces only 8% refusal.
  4. Ablating Instruct's L19 direction increases jailbreak ASR. On JailbreakBench attack artifacts, ASR rises 24% → 33% on JBC (+9.0 pp) and 50% → 69.5% on PAIR (+19.5 pp).

Layout

src/refusal_project/      # importable library
  config.py               # paths, model names, token budgets
  data.py                 # extraction pools (default + XSTest matched),
                          # JBB + XSTest evaluation loaders, JBB jailbreak artifacts
  model.py                # RefusalModel; user-instruction-end position helper
  activations.py          # extract_activations with position={user_end, last_full}
  directions.py           # diff-of-means (fp32), select_best_direction with harmless cap
  generation.py           # resume-safe incremental generator, phase-aware variant
  scoring.py              # substring/first-token classifiers, JudgeModel,
                          # split_on_think, NON_TERMINATING_COT label
  interventions.py        # ablation_hook_pair, all_layers_ablation_hooks,
                          # additive_hook_pair, phase_aware_ablation_hook_pair
  analysis.py             # cohens_d_along_direction, bootstrap_separation
  pipeline.py             # run_full_pipeline orchestrator

scripts/                  # CLI driver scripts (numbered ~ pipeline order)
  00_bootstrap.py         # clones andyrdt/refusal_direction into third_party/
  01_extract_directions.py
  02_generate_responses.py
  03_run_judge.py
  04_transplant_sweep.py
  05_make_figures.py      # layerwise separation + representation_vs_causal
  07_judge_validation.py
  08_matched_vs_default_separation.py
  10_cosine_similarity.py
  11_basis_drift_null.py
  12_cosine_null_subtracted.py
  13_cot_vs_answer_analysis.py
  14_perplexity_analysis.py
  15_perplexity_plots.py
  16_temporal_cot_analysis.py
  17_transplant_grid.py
  18_phase_aware_bars.py
  19_jailbreak_asr.py
  print_best_layer.py     # offline best-layer inspection from saved scores
  select_layer_with_cap.py # offline re-selection with harmless cap
  smoke_tests/            # environment-verification scripts (not pytest)

slurm/                    # SLURM sbatch wrappers (NYU Greene defaults)
results/                  # outputs (most generated artifacts gitignored;
                          # rendered figures and small JSONs tracked)
third_party/              # andyrdt/refusal_direction (cloned by 00_bootstrap; gitignored)

Setup

Requires Python ≥ 3.12, a CUDA-capable GPU, and an NYU Greene–style SLURM allocation if running the full pipeline (sbatch headers hardcode account / partition; edit them for a different cluster).

Recommended (uv — exact lockfile-pinned environment):

cd refusal-in-reasoning-models
uv sync                            # creates .venv from uv.lock
source .venv/bin/activate
python scripts/00_bootstrap.py     # clones andyrdt/refusal_direction (dataset loaders)

Alternative (plain pip):

cd refusal-in-reasoning-models
python -m venv .venv && source .venv/bin/activate
pip install -e .
python scripts/00_bootstrap.py

The slurm/*.sbatch wrappers default to account=csci_ga_3033_131-2026sp and partition=c12m85-a100-1. Edit those headers for a different cluster. On any cluster, generation jobs need an A100-class GPU (R1-Distill generation uses up to an 8192-token budget per response).

Reproduction

The repository ships pre-rendered figures in results/figures/ alongside small artifacts (matched prompts, splits, basis-drift JSONs). To re-derive each artifact from scratch:

Direction extraction (per-model)

Three parallel jobs over Base / Instruct / R1-Distill, both extraction positions.

# user_end + system="" extraction on the harmful/harmless default pool
sbatch slurm/extract_directions.sbatch

# last_full + system="" extraction
sbatch --export=ALL,EXTRA_ARGS="--position last_full --system '' --out-dir results/Qwen2.5-1.5B-Instruct/diag_lastfull_sys_empty" slurm/diagnose_instruct.sbatch
python scripts/01_extract_directions.py --model r1distill --position last_full --system "" --out-dir results/DeepSeek-R1-Distill-Qwen-1.5B/diag_lastfull_sys_empty
python scripts/01_extract_directions.py --model base --position last_full --system "" --out-dir results/Qwen2.5-1.5B/diag_lastfull_sys_empty

# Matched XSTest-pair extraction (200 unsafe + 250 safe; style-controlled)
sbatch slurm/extract_directions_matched.sbatch

Outputs: results/<model>/{best_direction,directions,layer_scores,layer_scores_harmless}.pt, results/<model>/{metadata,results}.json per condition, plus the harmful_acts.pt / harmless_acts.pt activation caches.

Behavioral evaluation (baseline + ablation)

# Baseline generation on the evaluation sets
for MODEL in base instruct r1distill; do
  for PS in jbb_harmful jbb_benign xstest_safe; do
    sbatch --export=ALL,MODEL=$MODEL,PROMPT_SET=$PS slurm/generate_long.sbatch
  done
done

# Ablation (Instruct L19, R1-Distill L22) on JBB harmful
sbatch --export=ALL,MODEL=instruct,PROMPT_SET=jbb_harmful,EXTRA_ARGS="--intervention ablate_layer --layer 19 --source-direction results/Qwen2.5-1.5B-Instruct/diag_lastfull_sys_empty/best_direction.pt --condition ablate_L19_lastfull" slurm/generate_long.sbatch
sbatch --export=ALL,MODEL=r1distill,PROMPT_SET=jbb_harmful,EXTRA_ARGS="--intervention ablate_layer --layer 22 --source-direction results/DeepSeek-R1-Distill-Qwen-1.5B/diag_lastfull_sys_empty/best_direction.pt --condition ablate_L22_lastfull" slurm/generate_long.sbatch

# Phase-aware ablation (R1-Distill, JBB harmful)
sbatch --export=ALL,MODEL=r1distill,PROMPT_SET=jbb_harmful,EXTRA_ARGS="--intervention ablate_layer_phase --phase think_only --layer 22 --source-direction results/DeepSeek-R1-Distill-Qwen-1.5B/diag_lastfull_sys_empty/best_direction.pt --condition ablate_L22_think" slurm/generate_long.sbatch
sbatch --export=ALL,MODEL=r1distill,PROMPT_SET=jbb_harmful,EXTRA_ARGS="--intervention ablate_layer_phase --phase answer_only --layer 22 --source-direction results/DeepSeek-R1-Distill-Qwen-1.5B/diag_lastfull_sys_empty/best_direction.pt --condition ablate_L22_answer" slurm/generate_long.sbatch

# Judge each generation file
for MODEL in base instruct r1distill; do
  for PS in jbb_harmful jbb_benign xstest_safe; do
    sbatch --export=ALL,MODEL=$MODEL,PROMPT_SET=$PS slurm/judge.sbatch
  done
done

Jailbreak ASR (PAIR + JBC)

for MODEL in base instruct r1distill; do
  for PS in jbb_jailbreak_pair jbb_jailbreak_jbc; do
    sbatch --export=ALL,MODEL=$MODEL,PROMPT_SET=$PS slurm/generate_long.sbatch
    sbatch --export=ALL,MODEL=$MODEL,PROMPT_SET=$PS slurm/judge.sbatch
  done
done
# Plus L19/L22 ablation conditions on Instruct/R1-Distill (analogous to above).

PAIR has 18 entries with prompt: null (records where the attacker model failed); the loader drops these, so PAIR effective n is 82 rather than 100.

Cross-model transplant

# Configuration B: source = Instruct's L19 last_full direction
for L in 19 22; do
  for ALPHA in 0.05 0.1 0.25 0.5; do
    for PS in jbb_harmful xstest_safe; do
      sbatch --export=ALL,PROMPT_SET=$PS,EXTRA_ARGS="--layers $L --alphas $ALPHA" slurm/transplant_sweep_B_last_full.sbatch
    done
  done
done

# Configuration A: source = Instruct's L1 user_end direction (negative control)
sbatch --export=ALL,PROMPT_SET=jbb_harmful,EXTRA_ARGS="--layers 22 --alphas 0.25" slurm/transplant_sweep_A_user_end.sbatch

# Judge each transplant cell
sbatch --export=ALL,MODEL=r1distill,PROMPT_SET=jbb_harmful,IN_PATH=results/transplant/last_full_diagnostic/jbb_harmful/L22_a0.25_pos_scaled.json,OUT_PATH=results/judge_scores/transplant/last_full_diagnostic/jbb_harmful/L22_a0.25_pos_scaled.json slurm/judge.sbatch
# ... repeat for each (L, α) cell

The α grid is calibrated against perplexity (see below). α ≥ 1 with --scale-by-norm is destructive; α ∈ [0.05, 0.5] keeps WikiText-2 perplexity within 1.7× baseline at L=22.

Perplexity controls

sbatch slurm/perplexity_analysis.sbatch          # 6 baseline + ablation conditions
sbatch slurm/perplexity_alpha_sweep.sbatch        # α-sweep at L=19
sbatch slurm/perplexity_alpha_sweep_L22.sbatch    # α-sweep at L=22

Cross-model cosine + basis-drift null distribution

# Per-pair null distribution (K=50 random binary partitions). --all-pairs
# runs the three model pairs (instruct↔r1distill, instruct↔base, r1distill↔base)
# using results/<model>/ for each.
python scripts/11_basis_drift_null.py --all-pairs --n-partitions 50

# Or explicit per-pair invocations against arbitrary extraction subdirs:
python scripts/11_basis_drift_null.py \
    --model-a-dir results/Qwen2.5-1.5B-Instruct/diag_lastfull_sys_empty \
    --model-b-dir results/DeepSeek-R1-Distill-Qwen-1.5B/diag_lastfull_sys_empty \
    --n-partitions 50 --tag instruct_vs_r1distill_lastfull

Temporal CoT analysis

sbatch slurm/temporal_cot_analysis.sbatch
# Replot from cached JSONs without GPU:
python scripts/16_temporal_cot_analysis.py --direction-source intra_cot --layer 22 --replot-only
python scripts/16_temporal_cot_analysis.py --direction-source intra_cot --layer 22 --cv-holdout-seed 0 --replot-only

Figures

All figure scripts read cached JSONs / .pt artifacts and run on CPU.

python scripts/05_make_figures.py                 # layerwise_separation, representation_vs_causal
python scripts/08_matched_vs_default_separation.py
python scripts/12_cosine_null_subtracted.py       # cross-layer cosine heatmaps
python scripts/15_perplexity_plots.py             # perplexity_vs_alpha, perplexity_bar_chart
python scripts/17_transplant_grid.py              # transplant_grid
python scripts/18_phase_aware_bars.py             # phase_aware_ablation
python scripts/19_jailbreak_asr.py                # jailbreak_asr

Acknowledgments

The single-direction extraction-and-ablation methodology, the harmful / harmless prompt pools, and the dataset-loading utilities cloned by scripts/00_bootstrap.py are from andyrdt/refusal_direction (Arditi et al., 2024). The intra-CoT direction extraction in scripts/16_temporal_cot_analysis.py is methodologically based on the caution-direction approach of Yamaguchi et al., 2025.

References

  • Arditi, A., Obeso, O., Syed, A., Paleka, D., Panickssery, N., Gurnee, W., & Nanda, N. (2024). Refusal in Language Models Is Mediated by a Single Direction. NeurIPS 2024. https://arxiv.org/abs/2406.11717
  • DeepSeek-AI. (2025). DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning. arXiv:2501.12948.
  • Wollschläger, T., Elstner, J., Geisler, S., Cohen-Addad, V., Günnemann, S., & Gasteiger, J. (2025). The Geometry of Refusal in Large Language Models: Concept Cones and Representational Independence. ICML 2025.
  • Yamaguchi, K., Etheridge, B., & Arditi, A. (2025). Adversarial Manipulation of Reasoning Models Using Internal Representations. ICML 2025 Workshop on Reliable and Responsible Foundation Models. arXiv:2507.03167.
  • Chao, P., Debenedetti, E., Robey, A., Andriushchenko, M., Croce, F., Sehwag, V., et al. (2024). JailbreakBench: An Open Robustness Benchmark for Jailbreaking Large Language Models. NeurIPS 2024 Datasets and Benchmarks Track.
  • Chao, P., Robey, A., Dobriban, E., Hassani, H., Pappas, G. J., & Wong, E. (2025). Jailbreaking Black Box Large Language Models in Twenty Queries. IEEE SaTML 2025.
  • Röttger, P., Kirk, H. R., Vidgen, B., Attanasio, G., Bianchi, F., & Hovy, D. (2024). XSTest: A Test Suite for Identifying Exaggerated Safety Behaviours in Large Language Models. NAACL 2024.
  • Qwen Team. (2024). Qwen2.5 Technical Report. arXiv:2412.15115.
  • Taori, R., Gulrajani, I., Zhang, T., Dubois, Y., Li, X., Guestrin, C., Liang, P., & Hashimoto, T. B. (2023). Stanford Alpaca: An Instruction-Following LLaMA Model. https://github.com/tatsu-lab/stanford_alpaca
  • Nanda, N., & Bloom, J. (2022). TransformerLens.

About

Mechanistic study of the refusal direction across base, instruction-tuned, and reasoning-distilled Qwen2.5-1.5B variants: extraction, ablation, transplant, and phase-aware analysis.

Topics

Resources

Stars

Watchers

Forks

Releases

Packages

Contributors

Languages