hanssuny/upsafeCódigo de investigación para añadir expertos de seguridad controlables a las capas de LLM mediante entrenamiento de upcycling en dos etapas, con enrutamiento de temperatura de seguridad en tiempo de inferencia y evaluación de jailbreak.
Implementación oficial de UpSafe℃ (NeurIPS 2026). UpSafe añade expertos de seguridad a capas seleccionadas de un modelo de lenguaje y utiliza un procedimiento de entrenamiento en dos etapas para aprender cuándo activarlos. Una temperatura de seguridad τ ajusta el enrutamiento en el momento de la inferencia sin reentrenamiento.
Ejecute todos los comandos desde la raíz del repositorio. El entrenamiento requiere GPUs CUDA; las
configuraciones proporcionadas utilizan DeepSpeed ZeRO-3. Descargue el checkpoint base en
models/Llama-3.1-8B-Instruct antes del entrenamiento.
conda create -n upsafe-train python=3.10 -y
conda activate upsafe-train
bash setup/init_env.sh
# Etapa 1: entrenar expertos de seguridad y enrutadores.
llamafactory-cli train upcycle/configs/llama3.1-8b/s1.yaml
# Etapa 2: congelar los expertos y entrenar los enrutadores con ejemplos dañinos y benignos.
llamafactory-cli train upcycle/configs/llama3.1-8b/s2.yaml
# Exportar el adaptador de la etapa 2 con control de seguridad en el momento de la inferencia.
python upcycle/convert_model.py \
--adapter_path outputs/llama3.1-8b_s2 \
--output_path models/Llama-3.1-8B-Instruct-UpSafe --arch llama
El instalador parchea PEFT y Transformers en el entorno de entrenamiento activo
y guarda los archivos originales en setup/backups/. Utilice un entorno dedicado.
Utilice un entorno separado para el modelo fusionado:
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
El valor predeterminado es τ = 0.5. Establezca TAUS="0.2 0.5 0.8" para evaluar múltiples
configuraciones, o pase --tau a eval/QA_gen.py. Los valores más altos favorecen a los expertos de seguridad
y pueden aumentar la negativa ante indicaciones benignas. El campo del checkpoint
safe_tempature conserva su ortografía original por compatibilidad.
| Directorio | Contenido |
|---|---|
scan/ | Sondeo y clasificación de capas críticas para la seguridad |
upcycle/ | Adaptadores MoE, exportación de modelos y configuraciones de entrenamiento |
eval/ | Generación de respuestas y evaluadores de seguridad/negativa |
data/eval/ | Indicaciones de JBB, StrongREJECT, WildChat, WildJailbreak y XSTest |
LLaMA-Factory/ | Marco de entrenamiento y datos de entrenamiento STAR-1 |
Las configuraciones cubren Llama-3.1-8B, Qwen2.5-7B/14B y sus variantes R1-Distill. Consulte la guía de reproducción para los escaneos de capas, las rutas de los modelos, las configuraciones por modelo y los comandos de evaluación individuales.
@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}
}
Construido sobre LLaMA-Factory, PEFT y Transformers, utilizando datos de STAR-1. El código está licenciado bajo Apache-2.0. El código de terceros, los conjuntos de datos y los modelos base conservan sus respectivas licencias; consulte LEGAL.md.