OOD jailbreak attack에 대해서도 robust하기 위해 ternary inference step을 통해 core jailbreak intent을 알아내는 방식을 제안함
Introduction
LLM의 safety 문제를 해결하기 위해 SFT, RLHF와 같은 post-training alignment methods가 존재해 옴
여기서 SFT란 Supervised Fine-Tuning으로 이미 학습된 LLM에 대해서 prompt와 reponse를 이용해서 지도학습을 통해 fine-tuning되는 것을 말하고,
RLHF란 Reinforcement Learning from Human Feedback으로, 사람에게 어떤 답변이 더 나은지 피드백을 통해 강화학습하는 것을 말함
but 이러한 방법들은 jailbreak attack에 취약함
최근 CoT(Chain-of-Thought)를 이용해서 LLM이 응답하기 전에 inference-time reasoning을 추가하는 o1, STAIR, STAR-1이 연구되기도 했지만,
AutoDAN-Turbo, Adversarial Reasoning과 같은 core intent를 숨기는 OOD jailbreak attacks에 대해서는 robust하지 않을 수 있음
OOD jailbreak attack에 대해서 ASR(낮을수록 방어가 잘 됐다는 뜻) 기존 모델들이 높게 나온 것을 확인할 수 있음
그렇다고 새로운 jailbreak 방법이 나올때마다 재학습 시키기에는 많은 비용이 듦
모든 jailbreak prompt들은 core malicious intent를 가지고 있는 것은 분명함
따라서 이 core jailbreak intent를 알아내기만 한다면 in-distribution prompt가 됨
기존의 Defense Methods는 accuracy of intent가 낮음을 알 수 있음
jailbreak strategy을 안다면, core intent를 추론할 수 있다는 방식으로 접근
Methods
STRUCTURED METICULOUS REASONING CONSTRUCTION
Prompt-Intent Pair Collection
strategy library(strategy's name, definition, example)
jailbreak prompts and 이 prompts와 상응하는 corresponding core intents
prompt-intent dataset
behavior-based data : harmful behavior(b_i), jailbreak prompt(s_i)을 LLM(M)을 통해 jailbreak prompt(x_i)를 생성
jailbreak-based data : jailbreak prompt(x_i)을 LLM(M)을 활용해 harmful behavior(b_i), jailbreak prompt(s_i)을 생성
이렇게 만들어진 dataset을 unsafe하다고 여겨지는 것들만 stragety library에 추가
Meticulous Reasoning Step Construction
strategy analysis : 모델 M에게 프롬프트 x_i 와 정답 전략 s_i^G 를 주고, 그 전략이 왜 맞는지 설명하는 추론 과정 z_i^s를 생성
intent analysis : 모델에게 프롬프트 x_i, 정답 전략 s_i^G, 정답 핵심 의도 b_i^G 를 주고, '이 전략을 썼을 때 이 프롬프트 뒤에 왜 이런 핵심 의도가 숨어 있다고 볼 수 있는지'를 설명하는 추론 z_i^b를 만들게 하는 과정
policy-based safety analysis : 핵심 의도 b_i^G 와 안전 정책 h 를 주고, 그 의도가 정책상 안전한지 위험한지 판단하는 reasoning z_i^c를 생성
final answer : 세 단계 reasoning을 바탕으로 나온 최종 응답
TRAINING AND INFERENCE WITH METICULOUS REASONING
Training Data Formatting
Meticulous Reasoning Step Construction에서 생성했던 reasoning을 step 토큰을 이용해 training data의 output으로 사용
Training for Meticulous Reasoning
system prompt : dynamic strategy library(strategy library에서 관련없는 strategies를 랜덤하게 drop함) + safety policy
user prompt : kailbreak, direct harmful, benign instructions
Inference with Meticulous Reasoning
Inference에서는 safety policy와 custom strategy library를 system prompt로 주고 Meticulous Reasoning을 수행
STEP-WISE PREFERENCE LEARNING AND TEST-TIME SCALING FOR SAFETY
Grounded Step-wise Tree Sampling
next step의 n개를 sampling Ground Truth에 근거해 score 매김
Step-wise Direct Preference Optimization
step-wise DPO : reasoning step에서 DPO 적용
DPO : 좋은 응답과 나쁜 응의 데이터만으로 모델을 직접 최적화하는 방식
Test-time Scaling with PRM
PRM(Process Reward Model)로 단계마다 후보를 뽑고 각 단계마다 점수가 가장 높은 sample을 고르는 beam search, 전체 답변 중 가장 점수가 높은 sample을 고르는 best-of-N을 사용해 최적의 sample 고름
ARMOR-THINK: EFFICIENT SAFEGUARD WITH FREE THINKING
(1) Simplifying Safety Reasoning (2) Injecting Free Thinking
Simplifying Safety Reasoning
OpenAI GPT-4o를 사용해 평균 토큰 길이가 1/3로 줄었음
Injecting Free Thinking
user input이 safe하다고 판단되면, reasoning process t 추가함
Preference Learning with Ternary Reward
Safety Score Rs : safety에 대한 점수
Helpfulness Score Rh : benign prompt에 대해서 얼마나 유용한지
Structure Score Rst : reasoning 형식이 제대로 지켜졌는지
Threat Models
attack 모델은 black box 상황이라고 가정
Experiments
각 step이 서로 영향을 줌
Limitations
답변할 때 inference-step을 거치기 때문에 inference-time overhead가 생길 수 있음
(그리고 과연 정말 새로운 attack method에도 robust할지는 지켜봐야할 것 같다)
