This is the repo for the submission One-Shot Safety Alignment for Large Language Models via Optimal Dualization.

The folders are organized as follows:

• safe-rlhf: contain our optimal and dual optimization algorithms; the dual optimization implementations are under safe-rlhf/trainer; and primal optimization algorithms are under algorithm/cdpo fodler; evaluation
• output: folder to save our generated results
• script: bash files to generate response and collect model-based safety and helpfulness scores

Collected data

We have also uploaded our collected scores to run our primal/dual optimization algorithms. Please find the collected results from this link: Collected data

Steps:

• Download the collected results; put data into corresponding folders. Detailed instructions can be found from the readme.md in the collected data folder.
• Setup the virual environment: Setup a conda environment using conda / mamba:

conda env create --file conda-recipe.yaml  # or `mamba env create --file conda-recipe.yaml`
conda activate safe-rlhf

• To run dual optimization, please enter into *safe-rlhf/trainer/ folder: -- run model_based_dual_trainer.ipynb for model-based dual optimization -- run preference_based_dual_trainer.ipynk for preference-based dual optimization.
• To run primal optimization, please enter into safe-rlhf/algorithms/cdpo* -- run python dpo.py --lamb [LAMB] --output_dir [OUTPUT_DIR] for MoCAN -- run ```python dpo_alg2.py --lamb [LAMB] --output_dir [OUTPUT_DIR]`` for PeCAN

Citation

Please feel free to email us at Xinmeng Huamg, Shuo Li, or Dongsheng Ding. If you find this work useful in your own research, please consider citing our work:

@inproceedings{huang2024one,
  title={One-Shot Safety Alignment for Large Language Models via Optimal Dualization},
  author={Huang, Xinmeng and Li, Shuo and Dobriban, Edgar and Bastani, Osbert and Hassani, Hamed and Ding, Dongsheng},
  booktitle={Proceedings of the Advances in Neural Information Processing Systems},
  year={2024},
  note={\textbf{Spotlight}}
}