hanssuny/upsafeCode de recherche pour ajouter des experts de sécurité contrôlables aux couches des LLM via un entraînement de surcyclage en deux étapes, avec un routage de température de sécurité au moment de l'inférence et une évaluation des jailbreaks.
Implémentation officielle de UpSafe℃ (NeurIPS 2026). UpSafe ajoute des experts de sécurité à des couches sélectionnées d'un modèle de langage et utilise une procédure d'entraînement en deux étapes pour apprendre quand les activer. Une température de sécurité τ ajuste le routage au moment de l'inférence sans réentraînement.
Exécutez toutes les commandes depuis la racine du dépôt. L'entraînement nécessite des GPU CUDA ; les
configurations fournies utilisent DeepSpeed ZeRO-3. Téléchargez le point de contrôle de base dans
models/Llama-3.1-8B-Instruct avant l'entraînement.
conda create -n upsafe-train python=3.10 -y
conda activate upsafe-train
bash setup/init_env.sh
# Étape 1 : entraîner les experts de sécurité et les routeurs.
llamafactory-cli train upcycle/configs/llama3.1-8b/s1.yaml
# Étape 2 : geler les experts et entraîner les routeurs sur des exemples nuisibles et bénins.
llamafactory-cli train upcycle/configs/llama3.1-8b/s2.yaml
# Exporter l'adaptateur de l'étape 2 avec contrôle de la sécurité au moment de l'inférence.
python upcycle/convert_model.py \
--adapter_path outputs/llama3.1-8b_s2 \
--output_path models/Llama-3.1-8B-Instruct-UpSafe --arch llama
L'installateur applique des correctifs à PEFT et Transformers dans l'environnement d'entraînement actif
et sauvegarde les fichiers originaux sous setup/backups/. Utilisez un environnement dédié.
Utilisez un environnement séparé pour le modèle fusionné :
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
La valeur par défaut est τ = 0.5. Définissez TAUS="0.2 0.5 0.8" pour évaluer plusieurs
paramètres, ou passez --tau à eval/QA_gen.py. Des valeurs plus élevées favorisent les experts de sécurité
et peuvent augmenter les refus sur des invites bénignes. Le champ du point de contrôle
safe_tempature conserve son orthographe d'origine pour des raisons de compatibilité.
| Répertoire | Contenu |
|---|---|
scan/ | Sondage et classement des couches critiques pour la sécurité |
upcycle/ | Adaptateurs MoE, export de modèle et configurations d'entraînement |
eval/ | Génération de réponses et juges de sécurité/refus |
data/eval/ | Invites JBB, StrongREJECT, WildChat, WildJailbreak et XSTest |
LLaMA-Factory/ | Cadre d'entraînement et données d'entraînement STAR-1 |
Les configurations couvrent Llama-3.1-8B, Qwen2.5-7B/14B et leurs variantes R1-Distill. Consultez le guide de reproduction pour les balayages de couches, les chemins de modèles, les configurations par modèle et les commandes d'évaluation individuelles.
@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}
}
Construit sur LLaMA-Factory, PEFT et Transformers, en utilisant les données STAR-1. Le code est sous licence Apache-2.0. Le code tiers, les jeux de données et les modèles de base conservent leurs licences respectives ; voir LEGAL.md.