Skip to content

uiuc-focal-lab/push-forward-alignment

Repository files navigation

Matching Ranks Over Probability Yields Truly Deep Safety Alignment

Paper 📝 Project Page 💻 (Coming Soon!)

Deep Safety Alignment with PRefill attEntion STOpping (PRESTO)

For the deep safety alignment fine-tuning (with and without PRESTO) in our paper, we utilized the repository of [1] with slight modifications (specifically, adding support for Qwen 3 and Gemma 3, and implementing the PRESTO regularization). In this repository, we provide our own full implementation of deep safety alignment (including support for PRESTO regularization) in deep_safety_alignment.py for convenience.

Example usage:

accelerate launch --config_file accelerate_config_dsa.yaml deep_safety_alignment.py \
  --model meta-llama/Llama-2-7b-chat-hf \
  --safety_dataset_path datasets/llama2_safety_data_direct.jsonl \
  --utility_dataset_path datasets/llama2_alpaca_anchor.json \
  --system_prompt \
  --presto \
  --safety_batch_size_per_device 2 \
  --utility_batch_size_per_device 8 \
  --gradient_accumulation_steps 2 \
  --save_dir models/llama_2_7b_chat_da \
  --show_batch_tqdm \
  --wandb_run_name llama_2_7b_chat_da

Llama 2 datasets can be found here.

More documentation coming soon!

References

[1] Qi, Xiangyu, et al. "Safety Alignment Should be Made More Than Just a Few Tokens Deep." The Thirteenth International Conference on Learning Representations.

About

Official repository for the paper "Matching Ranks Over Probability Yields Truly Deep Safety Alignment"

Resources

License

Stars

2 stars

Watchers

0 watching

Forks

Releases

No releases published

Packages

 
 
 

Contributors

Languages