LickyArc/LLM-Safety

LLM Safety Alignment

Summer work on a single question: how do you make a language model refuse what it should refuse, and what does that cost you?

It runs as an interactive dashboard. Every number, curve and heatmap is read from the actual experiment outputs — nothing on the page is illustrative — and each section can load the real checkpoints and generate live.

GIT_LFS_SKIP_SMUDGE=1 git clone https://huggingface.co/LickyArc/LLM-Safety
cd LLM-Safety
pip install -r requirements.txt
streamlit run app.py

Opens at http://localhost:8501. Needs only streamlit, pandas and plotly — no GPU. Everything shown is precomputed, so the GIT_LFS_SKIP_SMUDGE=1 keeps the clone at a few hundred KB instead of pulling 22 GB of weights you don't need to browse the results.

Add pip install torch transformers peft and the live-generation panels turn on, pulling individual checkpoints from weights/ on first use.

---

The arc

The work splits into two halves that build on each other.

First: how do you teach refusal at all? Preference data says which of two answers a human liked better. Two methods dominate the conversion of that signal into behaviour — PPO, which learns a reward model and optimises against it with reinforcement learning, and DPO, which skips the reward model and optimises the preference objective directly. I built both on gpt2-medium across four fine-tuning strategies (full, LoRA, QLoRA, prefix), wrote the DPO loss from scratch, and evaluated all eight models on 6,000 held-out questions.

DPO won on every axis, and the reasons turned out to be more interesting than the result: the reward model hit 99.8% pairwise accuracy, which sounds like success and is actually a warning — it saturated, leaving PPO with a flat gradient it never climbed. One PPO variant collapsed to empty output entirely.

Then: what happens after alignment, when you fine-tune for something else? Take Qwen2.5-1.5B-Instruct, already safety-aligned, and fine-tune it on medical exam questions. Nothing harmful anywhere in the data. Its safety degrades anyway — 4.5× more harmful completions on a 550-prompt benchmark. That is the safety tax, and the second half is about repairing it without retraining: two methods that edit weights directly (DARE, RESTA) and one that leaves weights alone and injects a safety direction into the residual stream at inference (Function Vectors).

---

What came out of it

Fine-tuning on entirely benign data measurably breaks safety. Unsafe score 0.0473 → 0.2127, bought in exchange for a large utility gain (medical ROUGE-L 0.0393 → 0.4573). That trade is the whole problem.

Editing activations beat editing weights. Of the seven configurations tested, the function vector was the only one that improved on the damaged model (0.2127 → 0.1891) — and it is reversible, inspectable, and costs one vector addition per token.

The function vector does not encode refusal vocabulary. Pushed through the unembedding, 97.9% of its mass lands on ' I' — not Sorry, not cannot — followed by a run of Chinese first-person pronouns (我, 我可以, 我会). What it encodes is the decision to answer in the first person: to stop completing the request and start speaking as an agent with a position on it. The concept is language-agnostic, which is the evidence that it is a stance and not an output string.

Safety is localised, and late. Of 336 attention heads, the ten that causally control refusal cluster in layers 12–22, strongest at layer 19 head 5. Early layers are near zero — the decision is made after the model understands the request.

DARE erased the task instead of preserving it. Medical ROUGE-L fell to 0.0415, the base model's level, while the safety damage stayed. Its redundancy assumption does not survive a rank-16 LoRA delta on a 1.5B model.

Nothing recovered base-model safety. The best repair reached 0.1855 against the base's 0.0473. Damage is far easier to cause than to undo.

---

RESTA needs a scaling coefficient

RESTA builds a safety direction by subtracting a deliberately unaligned model from the aligned base, then adds it to the damaged model:

δ_safe = θ_base(+) − θ_base(−)
θ_RESTA = θ_SFT + λ·δ_safe

The method as usually written has no λ — the vector is applied at full strength. That turns out to matter enormously. ‖δ_safe‖ = 3.69 here, comparable in size to the entire SFT update, so applying all of it is a large perturbation.

Evaluated on an L40: all 550 HarmEval prompts, judged by Qwen2.5-7B-Instruct in bf16, with utility on the MedQA test split.

Configuration	Unsafe ↓	ROUGE-L ↑
Base	0.0473	0.0393
SFT (LoRA)	0.2127	0.4573
SFT + DARE	0.2200	0.0415
SFT + RESTA (λ=1.0)	0.3636	0.4656
SFT + DARE + RESTA	0.4236	0.0423
SFT + Function Vector	0.1891	0.4350
SFT + DARE + FV	0.1964	0.0363

At full strength RESTA nearly doubles the harm it is meant to repair. So I swept λ, including a negative value as a sign control:

λ	−1.0	0.25	0.5	1.0
Unsafe ↓	0.6473	0.1909	0.1855	0.3636

λ = −1.0 settles the direction: flipping the vector so it adds the harmful direction sends the unsafe score to 0.6473, the worst result anywhere here. The vector unambiguously points from harmful toward safe — the response curve simply is not monotonic.

At λ = 0.5 it reaches 0.1855, competitive with the function vector while keeping the best utility of any configuration. Subtracting the whole harmful delta overshoots and tears out what refusal rests on; subtracting half removes the harm and leaves the model intact. One unspecified scalar separates the best result here from one of the worst.

---

Caveats I would rather state than have found

• Absolute scores are not comparable across the two evaluation runs — judge prompt and generation length differ. Compare within a table, not across.
• λ was swept on the same 550 prompts the result is reported on, so 0.1855 is optimistic. The gap to the function vector (0.1891) is within noise; the defensible claim is that RESTA fails on magnitude, not that it wins.
• One judge, one seed, n=1 per configuration. No confidence intervals.
• The reward model's 99.8% is training-set accuracy — never validated held-out.

A full write-up, including every design decision and the reasoning behind it, is in PROJECT_DEEP_DIVE.md.

---

Layout

app.py                  entry point, sidebar, routing
page_a1.py              PPO vs DPO — results, method, training curves, responses
page_a2.py              safety tax — results, function vector, λ sweep, responses
live.py                 on-demand inference, local weights or streamed from the Hub
theme.py                styling and shared presentation helpers

build_data.py           notebooks + FV artifacts → data/results.json
generate_cache.py       loads every checkpoint, records what it answers
fix_resta.py            builds the RESTA models from the harmful adapter
build_lambda_model.py   builds the λ=0.5 model
gpu_eval.py             full GPU re-evaluation — 7B judge, 550 prompts, utility
resta_sweep.py          the λ sweep
eval_safety.py          CPU/MPS fallback evaluation with a smaller judge
upload_weights_single.py publishes every checkpoint into weights/

Data file	Built by	Holds
results.json	build_data.py	every metric, training curve, FV artifact
generations.json	generate_cache.py	cached responses from all 15 checkpoints
safety_gpu.json	gpu_eval.py	corrected safety scores, 7B judge
utility_gpu.json	gpu_eval.py	corrected utility scores
resta_sweep.json	resta_sweep.py	the λ sweep

Weights

All 18 checkpoints live in this repo under weights/ — the reward model and eight PPO/DPO variants, the SFT / DARE / RESTA / harmful models, the function-vector artifacts, and model_sft_resta_lambda0.5, the best-performing configuration found.

Load any of them by subfolder:

from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained(
    "LickyArc/LLM-Safety", subfolder="weights/model_sft_resta_lambda0.5")
tok = AutoTokenizer.from_pretrained(
    "LickyArc/LLM-Safety", subfolder="weights/model_sft_resta_lambda0.5")

Or pull one without cloning the rest:

from huggingface_hub import snapshot_download

snapshot_download("LickyArc/LLM-Safety",
                  allow_patterns="weights/model_sft_merged/*")

weights/model_harmful_merged is deliberately unaligned and will comply with requests an aligned model refuses. It exists so the safety-vector arithmetic — δ_safe = θ_base(+) − θ_base(−) — is reproducible. It is not a chat model and should not be served.

Note on development

page_a1.py, page_a2.py, live.py and theme.py are imported modules, so Streamlit's auto-rerun does not pick up edits to them — restart the server after changing those. Edits to app.py hot-reload normally.

Llama-3.1-8B guardrails adapters (PPO vs DPO replication)

LoRA adapters for meta-llama/Llama-3.1-8B-Instruct, from the replication of the Assignment-1 PPO/DPO comparison on a modern instruction-tuned policy. Built with Llama; governed by the Llama 3.1 Community License.

from peft import PeftModel
from transformers import AutoModelForCausalLM
base = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.1-8B-Instruct", torch_dtype="bfloat16")
model = PeftModel.from_pretrained(base, "LickyArc/LLM-Safety",
                                  subfolder="weights/llama-guardrails-dpo-lora")

subfolder	run	AUC (harmful vs benign)	read this first
llama-guardrails-dpo-lora	DPO, headline	0.781	best model here; training log was lost to an overwrite bug
llama-guardrails-dpo-lora-lr1e5	DPO, lr 1e-5 control	0.751	per-step stats are 2-pair, not 16-pair
llama-guardrails-dpo-lora-refit	DPO refit, clean logging	not evaluated	training curves only
llama-guardrails-ppo-bert	PPO vs BERT RM	0.651	scores higher on the RM it never saw
llama-guardrails-ppo-bert-v2	PPO vs BERT, A1 KL	0.660	KL negative on 44/200 steps
llama-guardrails-ppo-qwen	PPO vs Qwen RM	0.546 (chance)	MODE-COLLAPSED — do not deploy
llama-guardrails-ppo-qwen-v2	PPO vs Qwen, A1 KL	0.635	collapse fixed, overoptimisation not

ppo-qwen posts the best raw safety numbers in the study (unsafe_score 0.0036, refusal 0.9836) and is useless: 35 unique responses across 550 prompts, AUC indistinguishable from chance. It is published as a negative result for studying reward hacking.

Every number above rests on a corpus where TF-IDF + logistic regression separates held-out preference pairs at 0.9952 — matching every neural reward model trained on it. The reward signal is substantially a template fingerprint, not a preference signal. See each subfolder's MODEL_NOTES.md.