SupritiVijay/hbrt-safety-dataset-and-code

HBRT: Harm-Benefit Reasoning Tuning for Qwen3-1.7B

Structured safety layer via multi-stakeholder harm-benefit analysis. The model generates <safety_check> XML reasoning before <think> to evaluate whether a query is harmful, then responds accordingly.

Results

Sorted by Average Score = mean(SorryBench Refusal, 100 − HarmBench ASR, FalseReject Compliance). All values in %.

Rank	Model	Epoch	LR	Extras	Completion	Refusal	100−ASR	FalseReject	Avg Score
1	Track 11 (neftune+low-lr)	7	2e-5	NEFTune α=5	99.7	90.9	97.0	82.5	90.1
2	Track 4 (weighted-trans)	7	2e-5	5x on 4 transitions	99.1	91.8	96.7	81.6	90.0
3	Track 2 (low-lr)	7	2e-5	—	99.4	88.4	96.3	84.3	89.7
4	Track 3 (weighted-struct)	5	2e-5	5x on structure	99.3	88.9	97.3	82.3	89.5
5	Track 9 (special-tokens)	3	5e-5	6 special tokens	99.6	90.0	97.7	77.0	88.2
6	Track 10 (all-tokens)	5	5e-5	38 special tokens	99.9	90.7	97.7	75.2	87.9
7	Track 6 (curriculum)	5	2e-5	short-first	99.6	89.6	97.0	79.8	88.8
8	Track 7 (neftune)	5	5e-5	NEFTune α=5	99.7	90.9	98.3	76.3	88.5
9	Track 8 (high-lr)	5	1e-4	batch=128	99.7	90.9	98.7	70.1	86.6
10	Track 5 (compact-xml)	3	2e-5	44% fewer tokens	99.0	65.9	84.7	91.7	80.8
—	Baseline (untrained)	—	—	Qwen3-1.7B	N/A	33.9	52.0	82.7	56.2

Benchmarks

• SorryBench Refusal — % of harmful prompts the model refuses (higher = safer). Judge: ft-Mistral-7B fine-tuned on sorry-bench.
• HarmBench ASR — Attack Success Rate (lower = safer). Classifier: HarmBench-Llama-2-13b-cls.
• FalseReject Compliance — % of benign prompts answered correctly (higher = more helpful). Prefix-matching refusal detection.
• Completion Rate — % of responses that finish the full <safety_check> → <think> → response pipeline.

Track Descriptions

• Track 2 (low-lr): LR=2e-5 with 6 structural special tokens — best helpfulness (84.3% FalseReject) while maintaining strong safety. The gentler learning rate prevents over-refusal.
• Track 3 (weighted-structure): 5x loss weight on ALL tokens inside <safety_check> and <think> blocks. Pushes refusal higher (90.7%) but at the cost of benign compliance.
• Track 4 (weighted-transitions): 5x loss weight on ONLY 4 transition tokens (</safety_check>, <safety_check_score>, <think>, </think>). Surgically targets structural boundaries — achieves highest refusal (91.8%) among models with >80% FalseReject.
• Track 5 (compact-xml): Training data reformatted to 44% fewer tokens with compact nested XML. Catastrophic failure for safety (13-15% ASR) — the shorter format doesn't give the model enough reasoning space to properly evaluate harm.
• Track 6 (completion-curriculum): Phase 1 trains on short examples (<4000 tokens) for 3 epochs, then Phase 2 trains on full dataset for 7 epochs. Doesn't outperform plain LR=2e-5 training.
• Track 7 (neftune): NEFTune embedding noise (α=5) at LR=5e-5. High refusal (90.9%) but poor FalseReject (76.3%) — the higher LR dominates over NEFTune's regularization benefit.
• Track 8 (large-batch-high-lr): Effective batch=128, LR=1e-4, cosine decay. Produces the highest raw refusal numbers but worst FalseReject (~70%) — too aggressive for the safety-helpfulness balance.
• Track 9 (special-tokens-structure): The breakthrough track — proved that 6 structural special tokens fix the completion problem (99.6%+ vs 8.3% without). All subsequent tracks use this as baseline.
• Track 10 (special-tokens-all): All 38 XML tags registered as special tokens. Works comparably to just 6 tokens — no clear advantage from the additional vocabulary expansion.
• Track 11 (neftune+low-lr): Combines LR=2e-5 + NEFTune α=5 — first config to break 90% refusal while keeping FalseReject >80%. NEFTune stabilizes helpfulness at later epochs where plain training overfits.

Key Findings

1. 6 structural special tokens are essential. Without them, the model gets stuck in infinite XML loops (8.3% completion even after 10 epochs). With them, 99.6%+ completion from epoch 3.
2. LR=2e-5 is optimal for helpfulness. LR=5e-5 caps FalseReject at ~76%, LR=1e-4 at ~70%. Only LR=2e-5 reaches 84%.
3. NEFTune regularizes against over-refusal at later epochs. Track 11 (NEFTune + LR=2e-5) maintains 82.5% FalseReject at epoch 7 while pushing refusal to 90.9%.
4. Compact XML is catastrophically unsafe. Shorter reasoning doesn't give the 1.7B model enough space to distinguish harmful from benign.
5. Epoch 5-7 is the sweet spot for LR=2e-5 before overfitting kicks in.

Model Architecture

• Base model: Qwen/Qwen3-1.7B
• Training: SFT with TRL SFTTrainer, assistant-only loss, 8x H100 GPUs
• Sequence length: 8192 tokens
• Dataset: 10,000 examples (5k harmful / 5k benign), token-length band 3500-6000

Generation Format

<|im_start|>user
{prompt}<|im_end|>
<|im_start|>assistant
<safety_check>
  <stakeholder>...</stakeholder>
  <harms>...</harms>
  <harm_score>0.826</harm_score>
  <benefits>...</benefits>
  <benefit_score>0</benefit_score>
</safety_check>
<think>
  Based on the safety check, harm ratio 0.826 > 0.4 threshold → harmful.
</think>
This query is harmful, so I must refuse to provide an answer.<|im_end|>

Inference

Requires skip_special_tokens=False in vLLM SamplingParams and <safety_check>\n appended after the assistant turn marker. See code/run_evals_v2.py for the full eval pipeline.

Repository Structure

├── data/
│   └── convergent_data_10k.parquet          # Training data (10k rows)
├── models/
│   ├── track-2-low-lr-special-tokens-epoch-7/
│   ├── track-3-weighted-structure-special-tokens-epoch-5/
│   ├── track-4-weighted-transitions-special-tokens-epoch-7/
│   ├── track-5-compact-xml-special-tokens-epoch-3/
│   ├── track-6-completion-curriculum-special-tokens-epoch-5/
│   ├── track-7-neftune-special-tokens-epoch-5/
│   ├── track-8-large-batch-high-lr-special-tokens-epoch-5/
│   ├── track-9-special-tokens-structure-epoch-3/
│   ├── track-10-special-tokens-all-epoch-5/
│   └── track-11-low-lr-neftune-special-tokens-epoch-7/
└── track-codes/
    ├── python-scripts/
    │   ├── train_sft_v2.py
    │   ├── run_evals_v2.py
    │   ├── gpu_clear.py
    │   └── make_compact_xml.py
    ├── eval.sh
    ├── track1-baseline.sh
    ├── track2-low-lr.sh
    ├── track3-weighted-structure.sh
    ├── track4-weighted-transitions.sh
    ├── track5-compact-xml.sh
    ├── track6-completion-curriculum.sh
    ├── track7-neftune.sh
    ├── track8-large-batch-high-lr.sh
    ├── track9-special-tokens-structure.sh
    ├── track10-special-tokens-all.sh
    └── track11-low-lr-neftune.sh