기술 요약: 텍스트 확산 모델을 위한 안전 인식 제거기

1. 문제 제기

MDLM 및 LLaDA와 같은 텍스트 확산 모델 (TDM) 은 병렬 생성과 반복적 정제를 통해 더 빠른 추론과 유연성을 제공하며, 자기회귀 (AR) 대규모 언어 모델에 대한 유망한 대안으로 부상했습니다. 그러나 TDM 의 안전성 함의는 아직 충분히 탐구되지 않았습니다. AR 모델을 위해 설계된 기존 안전 메커니즘 (예: 사후 필터링, 거부 샘플링, 추론 시간 디코딩 제약) 은 확산 설정으로 직접 적용되지 않습니다. furthermore, TDM 은 병렬 제거 및 양방향 컨텍스트 기능을 악용하여 해로운 콘텐츠를 빠르게 생성하는 특정 재일브랙 공격 (예: DIJA, PAD) 에 취약합니다.

TDM 안전을 위한 현재 솔루션 (예: 안전 목적 함수로 모델을 재학습하거나 안전/비안전 레이블을 사용한 조건부 확산 사용) 은 계산 비용이 많이 들고, 유연성이 부족하며, 모델이 커짐에 따라 확장하기 어렵습니다. 모델 재학습이나 상당한 계산 오버헤드 없이 확산 과정에 통합될 수 있는 경량화된 추론 시간 안전 메커니즘에 대한 시급한 필요성이 있습니다.

2. 방법론: 안전 인식 제거기 (SAD)

저자들은 TDM 의 반복적 제거 과정을 수정하여 생성된 샘플을 텍스트 공간의 검증 가능한 안전한 영역으로 유도하는 훈련이 불필요한 안전 안내 프레임워크인 안전 인식 제거기 (Safety-Aware Denoiser, SAD) 를 제안합니다.

핵심 개념

SAD 는 이미지 확산을 위해 원래 개발된 "안전 제거기 (Safe Denoiser)"프레임워크 (Kim et al., 2025) 를 이산 텍스트 확산에 적용합니다. 이 방법은 표준 데이터 제거기에서 "비안전"요소를 빼면 안전한 제거기를 구성할 수 있다는 이론적 통찰에 기반합니다.

핵심 관계식은 다음과 같이 정의됩니다:
$$E_{D_{safe}}[x_0|x_t] = E_D[x_0|x_t] + \beta^*(x_t) \left( E_D[x_0|x_t] - E_{D_{unsafe}}[x_0|x_t] \right)$$
여기서:

• $E_D[x_0|x_t]$는 표준 데이터 제거기 (훈련된 TDM) 입니다.
• $E_{D_{unsafe}}[x_0|x_t]$는 알려진 비안전 참조 시퀀스 집합 ($D_{unsafe}$) 으로 제한된 예상 깨끗한 토큰 분포를 나타내는 비안전 제거기입니다.
• $\beta^*(x_t)$는 현재 노이즈 상태 $x_t$가 비안전 분포에서 유래할 가능성에 따라 조정되는 적응형 가중치입니다.

구현 세부 사항

1. 비안전 제거기의 근사화: 정확한 비안전 분포는 알 수 없으므로, SAD 는 유한한 $N$개의 비안전 참조 시퀀스 (예: 독성 데이터셋, 저작권 문서, 또는 학습 데이터) 에 대한 몬테카를로 적분을 사용하여 $E_{D_{unsafe}}[x_0|x_t]$를 근사화합니다.
2. 이산 커널 적응: 가우시안 커널을 사용하는 연속 이미지 확산과 달리, TDM 은 이산 마스킹 토큰 전이를 사용합니다. 저자들은 전체 시퀀스에 대한 전이 확률 $q(x_t|x_0)$을 계산하기 위해 특정 범주형 커널 분해를 유도하여 토큰 매칭 및 마스킹 위치에 기반한 반발 가중치를 계산할 수 있게 합니다.
3. 하이퍼파라미터:
   • $\eta$ (안전 규모): 안내 강도를 조절하는 제어 가능한 하이퍼파라미터입니다.
   • $C$ (활성 창): 역과정 중 SAD 가 적용되는 시간 단계 집합입니다. 저자들은 생성 품질을 저하시키지 않으면서 의미적 궤적을 유도하는 데 가장 효과적인 시기는 제거 과정 초기 (단계의 첫 4 분의 1) 라고 발견했습니다.
   • $N$ (부정 집합 크기): 비안전 참조 시퀀스의 수입니다. 실험 결과, 안전성 향상은 $N \approx 500-1000$에서 빠르게 포화되는 것으로 나타났습니다.

알고리즘

추론 중, 활성 창 $C$ 내의 각 시간 단계 $t$에 대해:

1. 표준 제거기 출력 $x_\theta(x_t, t)$를 계산합니다.
2. 순방향 과정 하에서의 가능성에 기반하여 비안전 참조들의 가중 평균을 계산합니다.
3. 예측된 토큰 확률을 조정하여 분포를 비안전 참조에서 멀어지게 하여, 생성을 안전한 영역으로 효과적으로 유도합니다.

3. 주요 기여

• TDM 을 위한 최초의 훈련 불필요 안전 프레임워크: SAD 는 기존 모델을 재학습하거나 미세 조정하지 않고 텍스트 확산 모델에서 안전을 강제하는 원칙적인 추론 시간 방법을 제공합니다.
• 이론적 보장: 이 방법은 제거기에서 비안전 요소를 빼면 안전한 분포에서 샘플을 생성할 수 있다는 이론적 보장 (정리 2.1) 을 활용합니다.
• 모듈성과 유연성: SAD 는 모델에 구애받지 않으며 (MDLM 및 LLaDA 에서 입증됨), 다른 추론 시간 방어 (예: DiffuGuard, 퍼플렉시티 필터링) 와 결합할 수 있습니다.
• 다차원 안전: 이 프레임워크는 독성 콘텐츠 생성뿐만 아니라 기억 (개인정보 위험) 및 재일브랙 견고성도 다룹니다.

4. 실험 결과

저자들은 SAD 를 해로운 생성, 재일브랙 견고성, 그리고 기억이라는 세 가지 차원에서 평가했습니다.

해로운 생성

• 설정: RealToxicityPrompts (RTP) 의 악성 프롬프트를 사용하여 MDLM 및 LLaDA 에서 평가했습니다.
• 결과: SAD 는 사용된 부정 집합에 따라 비안전 생성의 공격 성공률 (ASR) 을 약 5~6 퍼센트 포인트 (예: MDLM 의 경우 38.4% 에서 32.6% 로) 감소시켰습니다.
• 품질: 중요하게도, 이 안전성 향상은 benign 프롬프트에 대한 BERTScore 로 측정된 생성 품질, 다양성, 유창성에서 무시할 수 있는 저하를 동반했습니다.

재일브랙 견고성

• 설정: LLaDA-Instruct, LLaDA-1.5, Dream-Instruct 에서 확산 네이티브 공격 (DIJA, PAD) 및 제로샷 재일브랙에 대해 테스트했습니다.
• 결과: SAD 는 PAD(병렬 디코딩) 공격에 대한 견고성을 크게 향상시켜 ASR 을 최대 16 퍼센트 포인트 (예: WildJailbreak 의 경우 43.2% 에서 29.0% 로) 감소시켰습니다.
• 조립성: DiffuGuard 와 같은 기존 방어와 결합될 때, SAD 는 추가적인 ASR 감소를 제공하여 상호 보완적임을 입증했습니다.

기억 (개인정보)

• 설정: WikiText-103 으로 미세 조정된 MDLM 에서 평가했으며, 기억을 처벌하기 위해 부정 집합으로 학습 데이터를 사용했습니다.
• 결과: SAD 는 유틸리티 (BERTScore) 에 미치는 영향은 최소화하면서 퍼지 중첩 (근사 기억을 위한 지표) 을 현저히 감소시켰습니다. 예를 들어, 안전 규모 $\eta$를 증가시키면 참조 완성과의 높은 의미적 유사성을 유지하면서 퍼지 중첩이 18.59% 에서 약 16.11% 로 감소했습니다.

효율성

• SAD 는 최소한의 계산 오버헤드를 도입합니다. 런타임 비용은 부정 집합의 크기보다는 활성 시간 단계의 수에 의해 지배됩니다. 처리량은 약간만 감소합니다 (예: 64 단계 창을 가진 LLaDA 의 경우 0.382 에서 0.313 seq/s 로). 이 방법은 FK Steering 과 같은 입자 기반 안내 방법보다 훨씬 효율적입니다.

5. 중요성 및 주장

이 논문은 SAD 가 텍스트 확산 모델에서 안전을 강제하는 효과적이고 확장 가능한 메커니즘을 제공한다고 주장합니다. 그 주요 중요성은 다음과 같습니다:

1. 안전 격차 해소: 이는 확산 패러다임의 고유한 취약점, 특히 병렬 제거 악용을 대상으로 하는 최초의 전용 방어 메커니즘을 제공합니다.
2. 실용적 배포: 훈련 불필요하고 경량화된 특성으로 인해 SAD 는 재학습의 prohibitive 비용 없이 기존 고정된 확산 체크포인트의 안전한 배포를 가능하게 합니다.
3. 안전과 유틸리티의 균형: 결과는 생성된 텍스트의 유창성이나 일관성을 희생하지 않고 해로운 출력을 방지하기 위해 안전 제약을 제거 궤도에 통합할 수 있음을 보여줍니다.

저자들은 SAD 가 강력한 보호 장치이지만 보완적 조치로 간주되어야 한다고 결론지었습니다. 안전 안내가 편향을 도입하거나 benign 콘텐츠를 억제할 수 있음을 인정하며, 지속적인 평가와 인간 감독의 필요성을 강조합니다. 이 연구는 확산 특화 안전 방법 및 신흥 TDM 을 위한 원칙적인 안전 통합에 대한 추가 연구를 촉진하는 것을 목표로 합니다.