Diffusion Language Model(DLM)은 여러 위치의 토큰을 반복적으로 denoising하며 문장을 완성합니다. 이러한 생성 방식은 병렬 생성과 양방향 문맥 활용이라는 장점을 제공하지만, 동시에 기존 언어 모델에서는 잘 드러나지 않았던 새로운 안전성 문제를 만들어냅니다.
(a) 실험에서 denoising 초반에는 harmful signal이 비교적 약하지만, 중반부터 강해지고 후반에는 높은 수준으로 나타나는 현상을 관찰하였습니다. 즉, DLM의 위험도는 생성 과정 내내 일정하지 않습니다. 이런 상황에서 항상 동일한 강도의 safety steering을 적용하면 위험한 순간에는 개입이 부족하고, 안전한 순간에는 지나치게 개입하는 문제가 발생할 수 있습니다.
더 중요한 것은 한 번 위험한 token이 생성 과정에 commit된 이후입니다. 성공한 jailbreak trajectory에서는 위험하다고 판단된 token의 80%가 32 step 이후에도 위험한 상태를 유지했고(b), 96 step 이후에도 72%가 남았습니다. 즉, 단순히 denoising을 계속한다고 해서 이미 형성된 위험한 representation이 자연스럽게 사라진다고 기대하기 어렵습니다.
따라서 DLM의 안전성을 확보하기 위해서는 두가지가 동시에 필요합니다. 아직 생성되지 않은 부분은 위험도에 맞춰 미리 안전한 방향으로 유도하고, 이미 위험하게 생성된 부분은 다시 수정할 수 있어야 합니다.
이 연구에서는 DLM의 intermediate denoising trajectory를 직접 제어하는 새로운 inference-time defense framework를 제안합니다.
핵심 아이디어는 크게 두 단계로 구성됩니다.
Adaptive Safety Steering을 위해서는 먼저 “현재 생성이 얼마나 위험한가?”와 “어느 방향으로 representation을 움직여야 안전해지는가?”라는 두 질문에 답할 수 있어야 합니다.
Gating direction은 현재 sate의 위험도를 판단하기 위한 방향입니다. Harmful input과 benign input의 hidden representation 차이를 이용해 harmful generation context와 관련된 representation direction을 찾습니다. 이후 현재 denoising state가 이 방향과 얼마나 정렬되어 있는지를 계산해 steering의 강도를 결정하는 신호로 사용합니다.
Seering direction은 ****실제로 representation을 어느 방향으로 이동시킬지를 결정합니다. 동일한 input에 대해 safe refusal response와 unsafe compliance response를 비교하고, 두 response의 hidden representation 차이로부터 안전한 생성 방향을 추출합니다.
두 direction을 얻었다면 실제 denoising 과정에서는 매 step마다 현재 generation state를 관찰합니다.
[MASK] 상태인 position에 강한 steering을 적용해 다음 예측이 safe direction으로 이동하도록 유도