SSAH — Superficial Safety Alignment Hypothesis

Official code release for the paper Superficial Safety Alignment Hypothesis (Jianwei Li & Jung-Eun Kim, ICLR 2026).

TL;DR. Safety alignment in modern LLMs is superficial: a tiny slice of channels — about 1 % of linear-layer weights — carries essentially all of the model's refusal behaviour. Removing that slice collapses safety with negligible utility cost (Table 1), and freezing it during fine-tuning preserves safety under common fine-tuning attacks (Table 2).

👉 Read this next: our follow-up paper Safety Alignment Can Be Not Superficial With Explicit Safety Signals (ICML 2025) builds on SSAH and proposes a complementary remedy — explicit safety signals that make alignment robust rather than superficial. Code at JEKimLab/Safety-Alignment-With-Explicit-Safety-Signal.

The release is organised by paper experiment. Each module is self-contained, with its own repro/ sub-directory that pins hyperparameters and orchestrates an end-to-end reproduction. Every generated artefact lives under that module's repro/out/ and is gitignored.

The two reproductions share the same conda env (prune_llm by default; see attribution/repro/README.md for setup) and form a two-stage pipeline:

1. Stage A — produce the SCU/UCU/RU attribution masks (one short run per model, ~3 min each on a single 47 GB GPU):

   cd attribution
   bash repro/scripts/00_attribution_llama2-7b-chat.sh
   bash repro/scripts/00_attribution_llama3-8b-instruct.sh

   These write atten_attr.pt / mlp_attr.pt under attribution/repro/out/attribution/<model>/....

2. Stage B — Table 1 (pruning) and / or Table 2 (fine-tuning attack):

   # Table 1 (≈ 90 min on 1× A6000):
   cd attribution
   bash repro/submit_all.sh
   python repro/aggregate_table1.py
   
   # Table 2 (≈ 40 hours on 3× A6000):
   cd ../finetuning
   bash repro/submit_all.sh
   python repro/aggregate_table2.py

   submit_all.sh queues every script through tsp (tsp -S 1, strictly sequential) and the aggregator turns the per-run log files into a Markdown version of the paper's table.

Llama2-7B-Chat with 1.3 % of channels (≈ 1 % of weights) pruned — the SCU class — under the AdvBench keyword judge:

Reproduced by the canonical pipeline above; matches paper Table 1 within run-to-run noise. See attribution/repro/README.md for the Llama3-8B-Instruct numbers and the full column set.

The SSAH pipeline reuses two excellent existing repositories — we deliberately keep their evaluation harnesses unchanged so our numbers remain directly comparable:

• Attribution & evaluation harness: Wei et al., Assessing the Brittleness of Safety Alignment via Pruning and Low-Rank Modifications (arXiv:2402.05162, boyiwei/alignment-attribution-code). We extend lib/prune.py with prune_fluctuation_decouple_utility_and_safety (the SCU/UCU/CU/RU classifier) and prune_fluctuation_utility, but reuse their Wanda/SNIP scoring primitives, calibration data, and full eval_ppl / eval_zero_shot / eval_attack stack.
• Fine-tuning attack framework: Qi et al., Fine-tuning Aligned Language Models Compromises Safety, Even When Users Do Not Intend To (arXiv:2310.03693, LLM-Tuning-Safety/LLMs-Finetuning-Safety), itself built on Meta's llama-recipes. We add the zero_gradient_for_attribute_* family in utils/train_utils.py — the gradient-mask mechanism that "freezes" channels labelled SCU or CU during FSDP fine-tuning — and the three new CLI flags --zero_gradient, --zero_role, --mlp_role / --attn_role.

See the per-module READMEs for the exact line-level diff and a more detailed credit list.

If you find this code useful, please cite our paper and the two upstream repositories whose infrastructure we reuse:

@inproceedings{li2026ssah,
  title     = {Superficial Safety Alignment Hypothesis},
  author    = {Li, Jianwei and Kim, Jung-Eun},
  booktitle = {International Conference on Learning Representations (ICLR)},
  year      = {2026}
}

@inproceedings{li2025explicit,
  title     = {Safety Alignment Can Be Not Superficial With Explicit Safety Signals},
  author    = {Li, Jianwei and Kim, Jung-Eun},
  booktitle = {International Conference on Machine Learning (ICML)},
  year      = {2025},
  note      = {arXiv:2505.17072}
}

@article{wei2024assessing,
  title   = {Assessing the Brittleness of Safety Alignment via Pruning and Low-Rank Modifications},
  author  = {Wei, Boyi and Huang, Kaixuan and Huang, Yangsibo and Xie, Tinghao and Qi, Xiangyu and Xia, Mengzhou and Mittal, Prateek and Wang, Mengdi and Henderson, Peter},
  journal = {arXiv preprint arXiv:2402.05162},
  year    = {2024}
}

@inproceedings{qi2024finetuning,
  title     = {Fine-tuning Aligned Language Models Compromises Safety, Even When Users Do Not Intend To!},
  author    = {Qi, Xiangyu and Zeng, Yi and Xie, Tinghao and Chen, Pin-Yu and Jia, Ruoxi and Mittal, Prateek and Henderson, Peter},
  booktitle = {International Conference on Learning Representations (ICLR)},
  year      = {2024}
}

The repository is released under the terms of each upstream component's license: the attribution code inherits Wei et al.'s license (attribution/LICENSE) and the fine-tuning code inherits Meta's Llama 2 Community License via llama-recipes (finetuning/LICENSE). New SSAH-specific additions (the prune_fluctuation_* functions, the zero_gradient_for_attribute_* mechanism, the repro/ pipelines) are released under the same terms.