Adaptive Steering and Remasking for Safe Generation in Diffusion Language Models

기술 요약: 확산 언어 모델의 안전한 생성을 위한 적응형 조향 및 재마스크링

문제 제기
확산 언어 모델 (DLM) 은 반복적 탈노이즈 및 양방향 정제를 통해 텍스트를 생성하며, 자동회귀 모델 대비 병렬 생성 및 전역 문맥 추론에서 장점을 제공합니다. 그러나 이러한 반복적 패러다임은 고유한 안전 취약점을 도입합니다. 토큰 수준에서 안전성이 종종 결정되는 자동회귀 모델과 달리, DLM 은 중간 탈노이즈 단계 동안 유해한 신호가 발생하여 전파되는 데 취약합니다. 최근 연구들은 초기 탈노이즈 궤적이 최종 출력에 강력한 영향을 미친다고 지적합니다. 일단 유해한 토큰이 마스크 해제되면, 전체 생성 궤적이 안전하지 않은 응답으로 유도될 수 있습니다. 기존 방어 방법들은 주로 자동회귀 프레임워크에서 적응되거나 전역적 재마스크링에 의존하여, 명시적으로 의미적 탈노이즈 궤적을 제어하지 못합니다. 그 결과, 유해한 토큰의 공격적 억제는 저품질 출력, 빈 응답, 또는 생성 유틸리티의 붕괴로 이어지는 트레이드오프가 발생하거나, 반면 불충분한 개입은 재브레이크를 방지하지 못합니다.

방법론
저자들은 추가 파인튜닝 없이 작동하는 추론 시 방어 프레임워크를 제안하며, 이는 적응형 조향 (Adaptive Steering) 과 토큰 재마스크링 (Token Remasking) 을 결합합니다. 이 프레임워크는 두 가지 상호 보완적 단계로 기능합니다:

1. 대조적 안전 방향 (CSD) 구성:
   이 방법은 유해한 생성과 안전한 생성 사이의 의미적 경계를 나타내는 잠재 벡터를 먼저 구성합니다. 쌍을 이룬 샘플 (유해한 응답 대 안전한 거절 응답을 가진 동일한 프롬프트) 을 사용하여 두 클래스 간 숨겨진 표현의 차이를 계산합니다. 이 차이 벡터인 CSD 는 모델의 잠재 공간에서 "유해한 의미 방향"을 포착합니다.

2. 1 단계: 초기 단계 적응형 조향:
   초기 탈노이즈 단계가 최종 궤적을 결정적으로 좌우한다는 점을 인식하여, 이 프레임워크는 탈노이즈 단계의 초기 부분 ($\rho T$) 동안에만 조향 개입을 적용합니다.

   • 메커니즘: 마스크된 토큰과 프롬프트 표현에 대한 모델의 숨겨진 상태를 CSD 에 투영합니다. 유해한 방향과 정렬된 숨겨진 상태의 구성 요소는 스케일링 인자 $\beta$를 사용하여 억제됩니다.
   • 적응성: 조향 강도는 유해성 추정 정도 (투영 크기) 에 따라 조절됩니다. 이는 유해한 의미가 안정화되기 전에 탈노이즈 궤적을 안전한 의미 영역으로 안내하며, 후기 정제 단계를 과도하게 구속하지 않습니다.

3. 2 단계: 토큰 재마스크링:
   초기 조향 이후, 등장할 수 있는 잔류 유해 토큰을 처리하는 정제 단계가 이어집니다.

   • 탐지: 프레임워크는 개별 토큰에 대한 CSD 와의 정렬 점수를 계산합니다. 임계값을 초과하는 토큰은 유해한 것으로 식별됩니다.
   • 재마스크링: 전역적 억제 대신, 이 방법은 식별된 유해 토큰만 선택적으로 재마스크합니다. 그런 다음 모델은 수정된 의미 문맥에서 이러한 위치를 재생성합니다. 이 반복 과정은 고정된 수의 단계 동안 또는 유해한 토큰이 남지 않을 때까지 지속되어, 유창하고 안전한 출력을 보장합니다.

주요 기여

1. 추론 시 방어 프레임워크: 이 논문은 의미적 조향과 선택적 재마스크링을 통합하는 플러그앤플레이 모듈을 소개하며, 기본 DLM 의 파라미터 업데이트나 안전 파인튜닝을 요구하지 않습니다.
2. 중요 제어 지점의 식별: 이 연구는 초기 탈노이즈 궤적을 DLM 안전의 주요 제어 지점으로 식별합니다. 부분 단계 조향이 유해한 행동이 안정화되기 전에 생성 경로를 효과적으로 재지시함을 입증합니다.
3. 균형 잡힌 안전성과 유틸리티: 제안된 방법은 기존 재마스크링 접근법 내재된 안전성 - 유틸리티 트레이드오프를 해결합니다. 궤적을 초기에 조향하고 후기에는 특정 토큰만 재마스크함으로써 유해한 생성을 완화하면서도 생성 품질과 유창성을 유지합니다.

실험 결과
이 프레임워크는 LLaDA-7B 와 Dream-8B 라는 두 개의 백본 DLM 에 대해 여러 벤치마크 (JailBreakBench, AdvBench, DIJA, PAP, Prefix attacks) 에서 평가되었습니다.

• 재브레이크 방어: 이 방법은 공격 성공률 (ASR) 을 크게 감소시켰습니다. Dream 모델에서 AdvBench 에서는 평균 ASR 0.64%, JailBreakBench 에서는 8.00% 를 달성하여, 다양한 공격 유형에 대한 균형 잡힌 견고성 측면에서 DiffuGuard 및 Self-reminder 와 같은 기존 베이스라인을 능가했습니다.
• 품질 보존: 생성 붕괴 (높은 빈도수의 빈 문장 또는 깨진 문장) 를 초래한 공격적 재마스크링 방법과 달리, 제안된 접근법은 원본 모델에 가까운 생성 품질을 유지했습니다. 일반 능력 벤치마크 (TruthfulQA, MATH-500, MMLU) 에서 이 방법은 DiffuGuard 보다 유틸리티를 더 잘 보존했으나, 안전 개입의 트레이드오프와 일치하게 원본 모델 대비 일부 저하가 관찰되었습니다.
• 애블레이션 연구: 실험은 두 단계 모두 필요함을 확인했습니다. 초기 조향을 제거하면 ASR 이 크게 증가 (18% 에서 54% 로) 했고, 재마스크링을 제거해도 성능이 저하되어 궤적 안내와 토큰 수준 교정 모두 필수적임을 강조했습니다.

의의
이 논문은 DLM 안전성이 최종 출력 분포의 함수일 뿐만 아니라 중간 탈노이즈 궤적에 깊이 의존한다고 주장합니다. 제안된 단계별 개입 프레임워크는 모델의 고유한 생성 능력을 훼손하지 않고 생성 전 과정에서 안전성을 유지할 수 있음을 보여줍니다. 대조적 안전 방향을 활용하여 잠재 공간을 조향하고 필요할 때만 선택적으로 재마스크함으로써, 이 방법은 확산 기반 텍스트 생성의 고유한 취약점에 대한 강력한 해결책을 제공하며, 비용이 많이 드는 재학습 없이 DLM 을 더 안전하게 배포할 수 있는 경로를 제시합니다.