UpSafe℃ (NeurIPS 2026) の公式実装です。UpSafe は言語モデルの選択された層に安全性エキスパートを追加し、2段階のトレーニング手順を用いてそれらをいつ有効化するかを学習します。安全性温度 τ は再トレーニングなしで推論時にルーティングを調整します。
すべてのコマンドはリポジトリのルートから実行してください。トレーニングには CUDA GPU が必要です。
提供されている設定は DeepSpeed ZeRO-3 を使用します。トレーニングの前に、ベースチェックポイントを
models/Llama-3.1-8B-Instruct にダウンロードしてください。
conda create -n upsafe-train python=3.10 -y
conda activate upsafe-train
bash setup/init_env.sh
# ステージ1: 安全性エキスパートとルーターをトレーニングします。
llamafactory-cli train upcycle/configs/llama3.1-8b/s1.yaml
# ステージ2: エキスパートを凍結し、有害な例と無害な例でルーターをトレーニングします。
llamafactory-cli train upcycle/configs/llama3.1-8b/s2.yaml
# 推論時の安全性制御を備えたステージ2アダプターをエクスポートします。
python upcycle/convert_model.py \
--adapter_path outputs/llama3.1-8b_s2 \
--output_path models/Llama-3.1-8B-Instruct-UpSafe --arch llama
インストーラーはアクティブなトレーニング環境内の PEFT と Transformers にパッチを適用し、
元のファイルを setup/backups/ に保存します。専用の環境を使用してください。
マージ済みモデルには別の環境を使用してください:
conda create -n upsafe-infer python=3.10 -y
conda activate upsafe-infer
pip install -r requirements_infer.txt
export JUDGE_API_KEY="your-api-key"
export JUDGE_BASE_URL="https://api.openai.com/v1"
export JUDGE_MODEL="gpt-4o"
bash eval/run_eval.sh models/Llama-3.1-8B-Instruct-UpSafe llama8b_upsafe
デフォルトは τ = 0.5 です。複数の設定を評価するには TAUS="0.2 0.5 0.8" を設定するか、
eval/QA_gen.py に --tau を渡してください。値が大きいほど安全性エキスパートが優先され、
無害なプロンプトに対する拒否が増える可能性があります。チェックポイントのフィールド
safe_tempature は互換性のために元の綴りを保持しています。
| ディレクトリ | 内容 |
|---|---|
scan/ | 安全性クリティカル層のプロービングとランキング |
upcycle/ | MoE アダプター、モデルのエクスポート、トレーニング設定 |
eval/ | 応答生成と安全性/拒否のジャッジ |
data/eval/ | JBB、StrongREJECT、WildChat、WildJailbreak、XSTest のプロンプト |
LLaMA-Factory/ | トレーニングフレームワークと STAR-1 トレーニングデータ |
設定は Llama-3.1-8B、Qwen2.5-7B/14B、およびそれらの R1-Distill バリアントをカバーしています。 層スキャン、モデルパス、モデルごとの設定、個別の評価コマンドについては 再現ガイド を参照してください。
@InProceedings{Sun_2026_NeurIPS,
author = {Sun, Yuhao and Yu, Lingyun and Xu, Zhuoer and Yang, Kun and Cui, Shiwen and Zhang, Yongdong and Xie, Hongtao},
title = {UpSafe{$^\circ$C}: Upcycling for Controllable Safety in Large Language Models},
booktitle = {Advances in Neural Information Processing Systems (NeurIPS)},
year = {2026}
}
LLaMA-Factory、 PEFT、 Transformers を基に構築され、 STAR-1 のデータを使用しています。 コードは Apache-2.0 の下でライセンスされています。サードパーティのコード、データセット、 ベースモデルはそれぞれのライセンスを保持します。LEGAL.md を参照してください。