FLAB Safety Specialist: Qwen3-1.7B fine-tuned for safety multiple-choice classification

A Qwen3-1.7B model fine-tuned with distilled-rationale supervised fine-tuning (SFT) to answer safety-related multiple-choice questions (for example SafetyBench-style items) and to end with exactly one boxed answer letter. It is the safety specialist from our EPFL CS-552 (Modern NLP, Spring 2026) project, "Post-training Qwen3-1.7B for Specialized Reasoning".

This is a task-specific classifier-style model, not a general-purpose safety or refusal model. See Limitations before using it.

Summary

Intended use

• Answering multiple-choice safety and harmfulness questions where a single boxed letter is expected.
• Research on small-model post-training, output contracts and checkpoint selection.

Not intended for content moderation in production, medical or legal decisions, or as a guarantee that a model's outputs are safe.

How to use

from transformers import AutoModelForCausalLM, AutoTokenizer

repo_id = "<your-username>/<repo-name>"  # replace with this repository id
tokenizer = AutoTokenizer.from_pretrained(repo_id)
model = AutoModelForCausalLM.from_pretrained(repo_id, torch_dtype="auto", device_map="auto")

question = (
    "Question: ...\n"
    "A. ...\nB. ...\nC. ...\nD. ...\n"
    "Give a brief explanation, then put your final answer letter in \\boxed{}."
)
messages = [{"role": "user", "content": question}]
inputs = tokenizer.apply_chat_template(
    messages, add_generation_prompt=True, return_tensors="pt"
).to(model.device)
output = model.generate(inputs, max_new_tokens=512, do_sample=False)
print(tokenizer.decode(output[0][inputs.shape[-1]:], skip_special_tokens=True))

The answer is extracted from the last \boxed{...} in the output. We evaluate with exact match on that letter.

Training

Data. SafetyBench, WildGuard, BeaverTails, format drills, over-refusal corrections and hard-case distillations. Short rationales for harmfulness classification were distilled from a Claude Haiku model to give the model more signal than the multiple-choice outcome alone. All source datasets are public research resources; the synthetic rationales may inherit teacher biases.

Method. LoRA (rank 16), learning rate 2e-6, effective batch size 16 (2 per device x 8 gradient accumulation steps), 2 epochs, about 2 hours. The adapter is merged into the base weights for release. SFT combined (i) distilled short rationales and (ii) targeted corrections for universal failures and over-refusals.

What we also tried (not selected).

We chose checkpoints for (i) reliable boxed-answer output under the context budget, (ii) robustness to simple distribution shift, and (iii) agreement between local proxies and course evaluation. This is why the SFT model was selected over the locally stronger DPO average.

Results

Local suites, pass@1 (SB = SafetyBench, BT = BeaverTails, JBB = JailbreakBench, WG = WildGuardTest, WJ = WildJailbreak):

Course evaluation (EPFL CS-552 held-out safety set):

In the full group model (all four specialists merged on a general-knowledge backbone), safety scores 81.0% and the four-domain average rises from 45.2% to 58.5%.

How to read these numbers.

• The gain is concentrated on SafetyBench-style multiple-choice items (90.6 to 92.4). The local average over all five suites is lower than the base model (55.5 vs 57.3), with regressions on BeaverTails, JailbreakBench and WildGuardTest.
• Part of the course gain comes from reliable boxed-answer formatting, which is why the headline gain (75.7 to 84.0) is larger than the gain at 4k context (82 to 84).
• SafetyBench was our best predictor of course performance: roughly 1.2 course points per SafetyBench point.
• Some evaluation sets are small, so single-digit differences are noisy.

Limitations

• Narrow scope. Improvements are benchmark-driven and do not show general safety, refusal quality or jailbreak robustness. Locally the model is weaker than the base on several adversarial suites.
• Universal failures. Ten SafetyBench prompts failed under every variant. They share the form "Does the speaker exhibit bias toward X?", where X is a sensitive Chinese social topic, and the model confuses the topic with the speaker's stance.
• Short rationales. The distilled chains have a median length of about 106 characters. Longer chains (400 to 600 characters) that spell out the speaker-versus-topic distinction are expected to help.
• Proxy mismatch. Variants that looked stronger on local suites (GRPO, DPO) regressed on course evaluation, so local scores should not be treated as a ranking.
• Over-refusal and culture. Safety data can over-refuse culturally atypical benign requests; we analyzed this but did not eliminate it.
• Boxed answers look authoritative. The format can hide uncertainty, especially on medical or safety topics.

Ethical considerations

Safety data exposes attack patterns and was used only from established academic benchmarks. Multilingual and cultural coverage is limited, and users from under-represented languages and communities may receive worse answers.

License and data terms

This model is derived from Qwen3-1.7B. Training data comes from SafetyBench, WildGuard and BeaverTails, plus rationales generated by a Claude Haiku model. Check the licenses of the base model and each dataset, and the provider terms for using model-generated text as training data, before reuse or redistribution.

Citation and credit

Project report (unpublished): "Post-training Qwen3-1.7B for Specialized Reasoning", Amir Ammar, Lev Maznichenko, Florian Dejean and Brandon Tang (FLAB), EPFL CS-552 Modern NLP, Spring 2026.

Contributions: <TODO: one line on who did what, for example "safety specialist: ...">.

Built on Qwen3, TRL and PEFT. Training configs are in hydra/; evaluation details are in EVAL_REPORT.md.