hanssuny/upsafeCódigo de pesquisa para adicionar especialistas de segurança controláveis às camadas de LLM por meio de treinamento de upcycling em dois estágios, com roteamento de temperatura de segurança em tempo de inferência e avaliação de jailbreak.
Implementação oficial do UpSafe℃ (NeurIPS 2026). O UpSafe adiciona especialistas em segurança a camadas selecionadas de um modelo de linguagem e usa um procedimento de treinamento em dois estágios para aprender quando ativá-los. Uma temperatura de segurança τ ajusta o roteamento no momento da inferência sem retreinamento.
Execute todos os comandos a partir da raiz do repositório. O treinamento requer GPUs CUDA; as
configurações fornecidas usam DeepSpeed ZeRO-3. Baixe o checkpoint base para
models/Llama-3.1-8B-Instruct antes do treinamento.
conda create -n upsafe-train python=3.10 -y
conda activate upsafe-train
bash setup/init_env.sh
# Estágio 1: treinar especialistas em segurança e roteadores.
llamafactory-cli train upcycle/configs/llama3.1-8b/s1.yaml
# Estágio 2: congelar especialistas e treinar roteadores em exemplos prejudiciais e benignos.
llamafactory-cli train upcycle/configs/llama3.1-8b/s2.yaml
# Exportar o adaptador do estágio 2 com controle de segurança no momento da inferência.
python upcycle/convert_model.py \
--adapter_path outputs/llama3.1-8b_s2 \
--output_path models/Llama-3.1-8B-Instruct-UpSafe --arch llama
O instalador aplica patches ao PEFT e ao Transformers no ambiente de treinamento ativo
e salva os arquivos originais em setup/backups/. Use um ambiente dedicado.
Use um ambiente separado para o modelo mesclado:
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
O padrão é τ = 0.5. Defina TAUS="0.2 0.5 0.8" para avaliar múltiplas
configurações, ou passe --tau para eval/QA_gen.py. Valores mais altos favorecem os especialistas em segurança
e podem aumentar a recusa em prompts benignos. O campo do checkpoint
safe_tempature mantém sua grafia original por compatibilidade.
| Diretório | Conteúdo |
|---|---|
scan/ | Sondagem e classificação de camadas críticas para segurança |
upcycle/ | Adaptadores MoE, exportação de modelo e configurações de treinamento |
eval/ | Geração de respostas e avaliadores de segurança/recusa |
data/eval/ | Prompts JBB, StrongREJECT, WildChat, WildJailbreak e XSTest |
LLaMA-Factory/ | Framework de treinamento e dados de treinamento STAR-1 |
As configurações abrangem Llama-3.1-8B, Qwen2.5-7B/14B e suas variantes R1-Distill. Consulte o guia de reprodução para varreduras de camadas, caminhos de modelos, configurações por modelo e comandos de avaliação individuais.
@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}
}
Construído sobre LLaMA-Factory, PEFT e Transformers, usando dados do STAR-1. O código está licenciado sob Apache-2.0. Código de terceiros, conjuntos de dados e modelos base mantêm suas respectivas licenças; consulte LEGAL.md.