최근 Diffusion Language Model(DLM)은 기존 autoregressive LLM과 달리, 문장을 한 토큰씩 생성하지 않고, 전체 문장을 반복적으로 “복원(denoising)”하며 생성하는 방식으로 주목받고 있습니다. 이러한 방식은 병렬 생성과 글로벌 문맥 이해에 강점을 가지지만, 동시에 새로운 안전성 문제도 발생시킵니다.
특히 DLM에서는 생성 초기 단계에서 등장한 harmful token이 이후 denoising 과정 전체에 영향을 주며, 최종적으로 위험한 응답을 유도할 수 있습니다. 실제로 논문에서는 생성 초기에 Sure 과 같은 순응 토큰을 삽입하면 jailbreak 성공률이 크게 증가하고, 반대로 Sorry 같은 거절 토큰을 넣으면 harmful response가 감소한다는 점을 확인하였습니다.
즉, 초기 denoising trajectory가 최종 안전성에 결정적인 영향을 준다는 것입니다. 기존 defense 방법들은 remasking 기반으로 harmful token을 제거하려 했지만, 문제는 harmful token뿐 아니라 정상적인 token까지 함께 억제되어 generation quality가 크게 무너진다는 점이었습니다.
따라서 이 연구는 “harmful trajectory 자체를 안전한 방향으로 steer할 수는 없을까?” 라는 질문에서 출발합니다.
이 연구에서는 DLM의 intermediate denoising trajectory를 직접 제어하는 새로운 inference-time defense framework를 제안합니다.
핵심 아이디어는 크게 두 단계로 구성됩니다.
먼저 모델 내부 representation 공간에서 “safe response”와 “harmful response”를 구분하는 방향 벡터를 학습합니다. 우리는 이를 Contrastive Safety Direction (CSD) 라고 정의합니다.
CSD는 동일한 prompt에 대하여 harmful response와 safe refusal message를 각각 생성한 뒤, 두 representation의 차이를 계산하여 구합니다. 즉, harmful hidden state와 safe hidden state의 차이를 통해 harmful semantic direction을 추출하는 것입니다. 이 방향 벡터는 이후 inference 과정에서 현재 generation이 얼마나 harmful 방향으로 향하고 있는지를 판단하는 기준이 됩니다.
논문의 핵심은 초기 denoising step에서만 steering을 수행한다는 점입니다. DLM은 초기에 생성된 trajectory가 이후 전체 생성을 결정하기 때문에, 초기에 harmful semantic direction을 제거하면 unsafe generation을 크게 줄일 수 있습니다.
따라서 우리는 hidden representation이 harmful direction에 얼마나 align되어 있는지를 계산한 뒤, alignment가 높을수록 stronger steering을 적용합니다.