ALIGNBEAM : Inference-Time Alignment Transfer via Cross-Vocabulary Logit Mixing | AI Paper Digest

TL;DR Highlight

도메인 파인튜닝으로 망가진 LLM 안전성을, 재학습 없이 추론 시점에 작은 안전 모델에서 빌려와 복구하는 방법.

Who Should Read

의료·금융 등 도메인 특화 LLM을 운영 중인데 안전성 저하가 걱정되는 ML 엔지니어. 파인튜닝 후 모델이 유해 요청에 응해버리는 문제를 inference 단계에서 해결하고 싶은 개발자.

Core Mechanics

도메인 SFT(지도 학습 파인튜닝)는 안전성을 심각하게 훼손함. MedLlama, Finance-Llama 등은 도메인 언어로 포장된 유해 요청에 쉽게 응하고, 시스템 프롬프트만으로는 안전성 회복이 안 됨 (Llama-3.1-8B 기준 시스템 프롬프트 적용 시 거부율 오히려 14.3%로 하락).
핵심 아이디어는 'Text Bridge': 안전 앵커 모델(Qwen2.5-3B-Instruct)의 상위 예측 토큰을 텍스트로 디코딩한 뒤 대상 모델의 토크나이저로 재인코딩해서, 어휘(vocabulary)가 달라도 확률 분포를 혼합할 수 있게 함.
3단계 파이프라인으로 작동: Phase 1(앵커 분포로 K개 beam 시작) → Phase 2(N 스텝 동안 혼합 디코딩) → Phase 3(대상 모델이 이어서 생성, 작은 LLM judge가 가장 안전한 beam 선택).
안전성은 초반 토큰에 집중됨. 딱 3 스텝만 혼합해도 HB+AB 안전 이득의 +61.9 pp를 확보하고, N=6 이상으로 깊이 늘려봤자 추가 이득은 3.9 pp에 불과함.
α(혼합 비율), N(혼합 깊이), K(beam 수) 세 파라미터 모두 재학습 없이 배포 시점에 조절 가능. K=1이면 속도 우선, K=3이면 안전성 우선 모드.
어휘가 다른 크로스 패밀리 모델 간에도 작동하는 최초의 logit mixing 방법. 기존 SafeDecoding, Proxy Tuning 등은 같은 vocabulary를 공유해야만 가능했음.

Evidence

AdvBench 거부율: 5개 모델 평균 38.1% → 91.5% (+53.4 pp), HarmBench-Standard 25.9% → 76.4% (+50.5 pp) 향상.
태스크 정확도 손실은 최소: GSM8K(DeepSeek-Math-7B-Instruct) 77.0% → 76.6% (−0.4 pp), MedQA(MedLlama-3-8B) 13.1% → 12.7% (−0.4 pp).
LlamaGuard 대비 benign 과잉거부율이 낮음: OR-Bench-Hard에서 LlamaGuard 34.5% vs ALIGNBEAM 22.3%, JBB-Benign에서 38-39% vs 26.0%.
Llama-3.1-70B 스케일에서도 검증됨: AdvBench 16.4% → 87.3% (+70.9 pp), 앵커(3B)가 드래프트(70B)보다 23배 작아도 안전성 전이 성공.

How to Apply

도메인 파인튜닝된 모델(의료, 금융 등)을 배포할 때, Qwen2.5-3B-Instruct 같은 작은 안전 앵커를 사이드카로 띄우고 ALIGNBEAM을 inference 미들웨어로 적용하면 재학습 없이 안전성을 복구할 수 있음. K=1로 시작하면 약 2배 레이턴시에 안전 이득의 80%를 얻을 수 있음.
LoRA로 도메인 어댑터만 붙인 모델도 안전성이 저하되는데, 이 경우에도 ALIGNBEAM을 inference 시점에 적용해서 복구 가능. 논문의 wealth-management LoRA 실험에서 HB-Ctx 100%, Sorry-Bench 85%로 베이스 instruct 모델 수준을 회복 또는 초과함.
α와 K 파라미터로 안전성-유틸리티 트레이드오프를 실시간 조절 가능. 일반 사용자 트래픽엔 α=0.5, K=1, 고위험 쿼리가 예상되는 맥락에선 α=0.75, K=3으로 런타임에 전환하는 식으로 활용.

Code Example

snippet

Terminology

SFT지도 학습 파인튜닝(Supervised Fine-Tuning). 모범 답안 데이터를 보여주고 그대로 따라하게 학습시키는 방식. 도메인 특화 데이터로 SFT하면 전문성은 올라가지만 안전성 훈련이 덮어씌워짐.

logit mixing두 모델의 다음 토큰 확률 분포를 섞는 기법. 두 요리사의 레시피를 비율에 맞게 합쳐서 새 레시피를 만드는 것과 유사. 여기서는 안전한 모델의 분포를 섞어 유해 응답 확률을 낮춤.

vocabularyLLM이 사용하는 토큰 사전. 모델마다 단어를 쪼개는 방식이 달라서 Llama와 Qwen은 같은 단어도 다른 토큰 ID를 가짐. 이게 달라서 기존 logit mixing이 크로스 패밀리에서 안 됐음.

beam search텍스트 생성 시 상위 K개의 후보를 동시에 추적하다가 최종적으로 가장 좋은 것을 선택하는 탐색 전략. 한 번에 하나씩만 생성하는 것보다 더 나은 결과를 찾을 가능성이 높음.

RLHF인간 피드백 강화학습(Reinforcement Learning from Human Feedback). 사람이 좋다/나쁘다 평가를 매기면 그 피드백으로 모델을 추가 학습하는 방식. GPT-4 같은 안전한 모델들이 이 방법으로 안전성을 훈련받음.

anchor modelALIGNBEAM에서 안전성의 기준이 되는 작은 모델(여기서는 Qwen2.5-3B-Instruct). 닻(anchor)처럼 드래프트 모델이 유해한 방향으로 흘러가지 않도록 잡아주는 역할.

KV-cache트랜스포머 모델이 이미 계산한 Key-Value 값을 저장해두는 캐시. 같은 접두사를 반복 계산하지 않아도 되게 해서 속도를 크게 높임. ALIGNBEAM Phase 3에서 이를 재사용해 오버헤드를 줄임.

over-refusal모델이 실제로 무해한 질문에도 과도하게 거부 응답을 내뱉는 현상. 안전 장치가 너무 강하게 걸려서 '약 복용법을 알려줘' 같은 정상 질문도 거절하는 문제.

Related Papers

Related Resources

ALIGNBEAM GitHub (예정)

Original Abstract (Expand)

Domain fine-tuning degrades the safety of large language models: fine-tuned specialists readily comply with harmful prompts framed in domain language. Existing inference-time defenses that mix logits from a safe anchor model require both models to share a vocabulary, which rules them out for the cross-family specialists where safety is most degraded. We present ALIGNBEAM, a training-free method that lifts this restriction by translating anchor logits into the target model's vocabulary token-by-token at each decoding step; a small LLM judge then selects the safest among K candidate continuations. No weights are changed, and the safety-utility trade-off can be tuned at deployment without retraining. Across both cross-vocabulary and same-vocabulary evaluation pairs, ALIGNBEAM substantially raises refusal on adversarial benchmarks while keeping task accuracy and inference overhead within practical bounds. The results show that safety alignment can be transferred between model families at inference time, without touching either model's weights.