OP12138/ImmuneCoT-baselines-small

ImmuneCoT Baselines — Safety CoT SFT Models

This repository hosts the 6 baseline models trained for the ImmuneCoT project, reproducing three safety Chain-of-Thought (CoT) alignment methods on two base reasoning models, using LLaMA-Factory.

Models

#	Base	Method	Training	Data	Path in repo
1	Qwen3-1.7B	STAR-1	Full SFT (DeepSpeed ZeRO-2)	1K curated safety CoT, 5 epochs	qwen3-1.7b/star1
2	Qwen3-1.7B	UnsafeChain	LoRA (rank=16, 8-bit)	13.6K correction CoT, 2 epochs	qwen3-1.7b/unsafechain (adapter)
3	Qwen3-1.7B	SafeChain	Full SFT (DeepSpeed ZeRO-2)	40K safety CoT, 3 epochs	qwen3-1.7b/safechain
4	DeepSeek-R1-Distill-Qwen-1.5B	STAR-1	Official pretrained weights (UCSC-VLAA/STAR1-R1-Distill-1.5B)	1K	deepseek-r1-distill/star1
5	DeepSeek-R1-Distill-Qwen-1.5B	UnsafeChain	LoRA (rank=16, 8-bit)	13.6K correction CoT, 2 epochs	deepseek-r1-distill/unsafechain (adapter)
6	DeepSeek-R1-Distill-Qwen-1.5B	SafeChain	Full SFT (DeepSpeed ZeRO-2)	40K safety CoT, 3 epochs	deepseek-r1-distill/safechain

Method Overview

• SafeChain (arXiv:2502.12025, ICLR 2025 BiAlign Workshop): First CoT-style safety training dataset. Direct SFT on 40K safety CoT generated by DeepSeek-R1 under explicit safety instructions.
• STAR-1 (arXiv:2504.01903): "Safer Alignment of Reasoning LLMs with 1K Data". 1K highly curated safety CoT selected from 530K pool via GPT-4o 3-dim scoring (reasoning accuracy / content safety / policy relevance), all max-scored. 1K STAR-1 > 40K SafeChain in safety.
• UnsafeChain (arXiv:2507.21652, MBZUAI): Hard-prompt correction. Selects prompts where base R1-8B produces unsafe replies, uses GPT-4.1 to correct into safe R1-style CoT. Trains with LoRA (rank=16, alpha=32, target=q_proj,v_proj, 8-bit loading). Covers 6 domains: WildJailbreak/StrongReject (adversarial), GSM8K (math), MBPP (code), TruthfulQA (fact), HH-RLHF (alignment).

Training Hyperparameters (per original papers)

Param	SafeChain	STAR-1	UnsafeChain
finetuning_type	full	full	lora
learning_rate	5e-5	1e-5	1e-5
num_train_epochs	3.0	5.0	2.0
effective_batch	32 (2x2x8)	32 (2x1x16)	32 (2x2x8, auto DDP)
cutoff_len	2048	4096	2048
precision	bf16	bf16	fp16
lora_rank / alpha	-	-	16 / 32
lora_dropout	-	-	0.05
lora_target	-	-	q_proj, v_proj
quantization	-	-	8-bit
deepspeed	ZeRO-2	ZeRO-2	none

Hardware

2x NVIDIA RTX 4090 (24GB), LLaMA-Factory (Nightly), DeepSpeed ZeRO-2 (full SFT), bitsandbytes 8-bit (LoRA), torch 2.13.0+cu126.

Templates

• Qwen3-1.7B uses qwen3 template (Qwen3 native thinking format)
• DeepSeek-R1-Distill uses deepseek3 template (R1-style thinking format)
• Both preserve <think> CoT.

Key Empirical Finding (preliminary, sanity-test on a single harmful prompt)

On the DeepSeek-R1-Distill-1.5B base, STAR-1 (full SFT, 1K data) successfully refuses the harmful prompt, while UnsafeChain (LoRA, 13.6K data) fails to refuse the same prompt — training capacity (full SFT vs LoRA / 0.12% params) dominates data size when the base model's unsafe tendency is stubborn (R1 distill may weaken refusal). On Qwen3-1.7B, both methods successfully refuse. Awaiting HarmBench/XSTest evaluation for quantified comparison.

Usage

LoRA adapters must be loaded with the corresponding base model:

import torch
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

BASE = "Qwen/Qwen3-1.7B"   # or "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"
ADAPTER = "path/to/unsafechain/adapter"

tok = AutoTokenizer.from_pretrained(BASE, trust_remote_code=True)
base = AutoModelForCausalLM.from_pretrained(
    BASE, trust_remote_code=True,
    quantization_config=BitsAndBytesConfig(load_in_8bit=True),
    torch_dtype=torch.float16, device_map="auto",
)
model = PeftModel.from_pretrained(base, ADAPTER).eval()

Full SFT models load directly:

from transformers import AutoModelForCausalLM, AutoTokenizer
model = AutoModelForCausalLM.from_pretrained("path/to/full/sft", trust_remote_code=True)
tok = AutoTokenizer.from_pretrained("path/to/full/sft", trust_remote_code=True)

License

Please respect the licenses of:

• Base models: Qwen3 (Qwen License), DeepSeek-R1-Distill (MIT)
• Source datasets: SafeChain / STAR-1 / UnsafeChain (see their respective HF pages)
• UCSC-VLAA/STAR1-R1-Distill-1.5B pretrained weights (included with permission for reproduction)

Model weights are released for research reproduction of the ImmuneCoT baselines.

Reproduction

Training scripts, YAML configs, experiment records and reproduction guide are in the project's baseline/ directory. Training used LLaMA-Factory with the configurations above.

Citation

@misc{safechain2025,
  title={SafeChain: Leveraging Language Models for Safer Reasoning through Chain-of-Thought},
  year={2025},
  note={arXiv:2502.12025}
}
@misc{star1,
  title={STAR-1: Safer Alignment of Reasoning LLMs with 1K Data},
  year={2025},
  note={arXiv:2504.01903}
}
@misc{unsafechain,
  title={UnsafeChain: Enhancing Reasoning Model Safety via Hard Cases},
  year={2025},
  note={arXiv:2507.21652}
}