UpSafe℃: 대규모 언어 모델의 제어 가능한 안전성을 위한 업사이클링 (NeurIPS 2026)

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를 참조하세요.