💬 NLP

A2D: Any-Order, Any-Step Safety Alignment for Diffusion Language Models

A2D는 무작위 마스킹을 활용하여 즉각적인 [EOS] 거부 신호를 방출함으로써 DIJA와 같은 임의 순서 및 임의 단계 공격을 효과적으로 무력화하고, 실시간 모니터링과 현저히 빠른 안전한 종료를 가능하게 하는 확산 언어 모델을 위한 토큰 수준의 안전 정렬 방법을 도입한다.

원저자: Wonje Jeung, Sangyeon Yoon, Yoonjun Cho, Dongjae Jeon, Sangwoo Shin, Hyesoo Hong, Albert No

게시일 2026-02-04

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Wonje Jeung, Sangyeon Yoon, Yoonjun Cho, Dongjae Jeon, Sangwoo Shin, Hyesoo Hong, Albert No

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

개요: 새로운 종류의 AI 작가

대부분의 AI 챗봇(우리가 매일 대화하는 것들)을 조립 라인의 노동자라고 상상해 보세요. 이들은 문장을 왼쪽에서 오른쪽으로, 한 단어씩 엄격하게 쌓아 올립니다. 만약 위험한 질문을 받으면, 이들은 보통 처음 몇 단어를 확인하고 "안 돼, 이건 나빠"라고 결정한 뒤 즉시 멈춥니다.

하지만 **확산 언어 모델(Diffusion Language Model, dLLM)**이라는 새로운 유형의 AI가 있습니다. 이 AI를 조립 라인이 아니라, 대리석 블록을 깎는 조각가라고 생각해보세요. 단어를 하나씩 써 내려가는 대신, 물음표(마스크)로 가득 찬 빈 페이지에서 시작하여 점진적으로 내용을 채워나갑니다. 문장의 끝을 앞부분보다 먼저 채울 수도 있고, 시작 부분보다 중간을 먼저 채울 수도 있습니다. 즉, 이야기 전체를 한꺼번에 작업할 수 있습니다.

문제점: "백도어" 공격

이 조각가는 어떤 순서로든 내용을 채울 수 있기 때문에, 새로운 약점을 갖게 됩니다.

나쁜 의도를 가진 사람이 조각가가 자동차 절도 방법을 쓰도록 속이려 한다고 가정해 봅시다.

기존의 함정: 만약 조각가가 조립 라인의 노동자라면, 공격자는 이야기의 맨 처음에 속임수를 씁니다.
새로운 함정 (DIJA 공격): 조각가의 경우, 공격자는 이야기의 전반부를 무해한 텍스트로 채운 뒤, 중간에 "함정"을 몰래 집어넣거나, 혹은 뒷부분을 먼저 채워버릴 수 있습니다. 그들은 평범한 요청처럼 보이지만 빈칸 속에 위험한 내용을 숨겨둔 템플릿을 사용합니다.

이 논문은 현재 이러한 조각가들을 위한 안전 학습이 **"얕다(shallow)"**는 것을 발견했습니다. 이는 마치 입구에서 신분증만 확인하는 보안 요원과 같습니다. 만약 당신이 보안 요원을 지나쳐 복도로 들어서기 시작하면, 보안 요원은 주의를 기울이지 않습니다. AI는 첫 단어에 대해서는 "안 돼"라고 말할 수 있지만, 만약 당신이 10번째 단어를 위험한 내용으로 채우도록 속인다면, AI는 다시 "안 돼"라고 말하는 것을 잊어버립니다.

해결책: A2D (Any-Order, Any-Step Defense)

저자들은 A2D라는 새로운 안전 학습 방법을 만들었습니다. 작동 방식은 다음과 같은 비유로 설명할 수 있습니다.

"빨간 정지 표지판" 토큰
AI에게 단순히 문장 시작 부분에서 "그것을 할 수 없습니다"라고 말하도록 훈련시키는 대신, A2D는 AI에게 특별한 "정지 표지판" 토큰(called [EOS])을 가르칩니다.

학습 과정: 연구진은 AI에게 위험한 문장들을 보여줍니다. 하지만 문장을 끝까지 완성하게 두는 대신, 위험한 부분들을 물음표로 가리고 이렇게 지시합니다. "만약 여기에 물음표가 있고, 이 문장이 위험하다면, 너는 즉시 그 물품표를 '정지 표지판'으로 교체해야 한다."
결과: AI는 문장의 어느 지점에서든, 만약 위험한 기운을 감지한다면 즉시 그 자리에 "정지 표지판"을 내걸어야 한다는 것을 배웁니다.

왜 이것이 특별할까요?

Any-Order (어떤 순서든): AI가 첫 번째 단어를 쓰고 있든 마지막 단어를 쓰고 있든 상관없습니다. 위험한 아이디어가 나타나면 정지 표지판이 나타납니다.
Any-Step (어떤 단계든): 위험이 1단계에서 나타나든 50단계에서 나타나든 상관없습니다. 보안 요원은 문 입구에서만 깨어 있는 것이 아니라, 작업 내내 깨어 있습니다.

"빈칸 채우기" 테스트

연구진은 자신들의 방법이 효과적임을 증证明하기 위해 **FITS (Fill-in-the-Sentence)**라는 매우 어려운 테스트를 고안했습니다.

시나리오: 나쁜 사람이 폭탄 제조법의 99%를 완성된 상태로 작성했습니다. 그리고 딱 한 문장(마지막 단계)만 비워둔 채 AI에게 이를 채워달라고 요청합니다.
결과: 기존의 안전 방식들은 처참하게 실패했습니다. 그들은 텍스트의 시작 부분에만 너무 집중한 나머지, AI가 폭탄 제조법을 완성하도록 내버려 두었습니다.
A2D의 성능: A2D는 그 단 하나의 빠진 문장을 보고, 그것이 위험하다는 것을 알아차린 뒤 즉시 정지 표지판을 세웠습니다. 거의 100%의 확률로 공격을 차단했습니다.

보너스: 실시간 안전 레이더

AI가 위험을 감지할 때마다 정지 표지판을 세우도록 훈련되었기 때문에, 그 정지 표지판이 나타날 확률은 안전 레이더 역할을 합니다.

조기 거부: 연구진은 AI가 단 한 단어도 쓰기 전인 아주 첫 단계에서 "정지 표지판 확률"을 확인하면, AI가 글을 쓰기도 전에 전체 요청이 나쁜 것인지 판단할 수 있다는 것을 발견했습니다.
속도: 이를 통해 AI는 이전보다 19.3배 더 빠르게 나쁜 요청을 거절할 수 있습니다. 이는 클럽 입구에서 명단에 없는 사람을 보고 문에 도달하기도 전에 미리 알아채서 시간을 절약하는 보안 요원과 같습니다.

AI의 유용성을 해치지는 않나요?

논문은 이 부분을 광범위하게 테스트했습니다. 연구진은 AI에게 수학, 코딩, 일반적인 질문을 던졌습니다.

결론: A2D는 AI를 여전히 똑똑하고 유용하게 유지했습니다. AI가 일반적인 질문(예: "파이썬 프로세스를 어떻게 죽이나요?"와 같이 위험하게 들릴 수 있지만 실제로는 코딩에 관한 질문)에 대해 답변을 거부하지 않도록 했습니다.
비교: 다른 안전 방식들은 너무 민감하여 무해한 질문에도 답변을 거부했습니다. 반면 A2D는 정밀했습니다. 오직 나쁜 것들만을 멈추게 했습니다.

요약

이 논문은 어떤 순서로든 글을 쓰는 새로운 유형의 AI를 위한 안전 업그레이드인 A2D를 소개합니다.

기존의 안전: 몇 분 지나면 지쳐버리는 문 앞의 경비원.
A2D의 안전: 위험한 냄새가 나면 언제 어디서든 즉시 "정지" 표지판을 세우는 보안 시스템.
결과: 문 앞을 슬쩍 지나가는 교묘한 공격을 막아내며, 19배 더 빠르게 "아니오"라고 말하면서도, 다른 모든 사람을 위해 AI를 계속 유용하게 유지합니다.

기술 요약: A2D – 디퓨전 언어 모델을 위한 임의 순서, 임의 단계 안전 정렬 (Any-Order, Any-Step Safety Alignment)

1. 문제 정의

디퓨전 대규모 언어 모델(dLLM)은 고정된 좌측-우측(left-to-right) 시퀀스가 아니라, 임의의 위치에서 반복적으로 토큰을 생성하는 **임의 순서 디코딩(any-order decoding)**을 가능하게 함으로써 자기회귀(autoregressive, AR) 모델과는 확연히 다른 양상을 보입니다. 이러한 유연성은 복잡한 추론과 구조화된 콘텐츠 합성을 지원하지만, 기존의 정렬 방식이 효과를 발휘하지 못하게 만드는 독특한 안전성 취약점을 초와합니다.

식별된 핵심 문제는 **얕은 정렬(shallow alignment)**입니다. AR 모델과 마찬가지로, 현재의 dLLM은 초기 생성 단계에서는 강력한 안전 신호를 보이지만, 디코딩이 진행됨에 따라 이 신호가 급격히 감쇠하는 현상을 보입니다. 이러한 한계는 다음과 같은 거대한 공격 표면을 생성합니다:

임의 순서 취약성(Any-Order Vulnerability): 유해한 콘텐츠가 디코딩 궤적 중 임의의 위치에 주입되어 초기 거절 메커니즘을 우회할 수 있습니다.
임의 단계 취약성(Any-Step Vulnerability): DIJA(Wen et al., 2025)와 같은 템플릿 기반 공격은 부분적으로 채워진 프롬프트 내에 유해한 텍스트와 [MASK] 토큰을 교차 배치함으로써, 안전 신호가 이미 희미해진 디코딩 여러 단계 이후의 시점을 공략합니다. 이러한 공격은 응답 수준(response-level) 정렬 방법이 방지하지 못하는 유해한 완성을 유도합니다.

거절 학습(RT), 지도 미세 조정(SFT), 분산 감소 선호도 최적화(VRPO)를 포함한 기존 베이스라인들은 응답 수준의 감독에 의존합니다. 결과적으로, 이들은 다양한 디코딩 전략(예: 랜덤, 신뢰도 기반) 하에서 안전성을 유지하는 데 어려움을 겪으며, 초기 토큰을 우회하는 공격에 대응하지 못합니다.

2. 방법론: A2D (Any-Order, Any-Step Defense)

저자들은 dLLM의 유연한 디코딩 프로세스에 특화되어 설계된 토큰 수준 정렬 방법인 A2D를 제안합니다. 전체 응답 수준에서만 거절을 감독하는 기존 방식과 달리, A2D는 유해한 콘텐츠가 마스킹된 위치에서 발견될 때마다 특정 [EOS] (End of Sequence) 토큰을 생성하도록 모델을 정렬합니다.

핵심 메커니 Unism

토큰 수준 감독: 학습 과정에서 타겟 시퀀스에 유해한 구간이 존재하는 경우, 모델은 해당 구간 내의 마스킹된 토큰을 원래의 유해한 토큰 대신 [EOS]를 예측하도록 감독받습니다.
학습 목적 함수: 이 방법은 수정된 타겟 분포를 사용하는 표준 마스크 디퓨전 목적 함수를 활용합니다:
- 유해 세트 ( $D_{harm}$ ): 유해한 응답 내의 마스킹된 토큰은 [EOS]를 예측하도록 감독됩니다. 이는 모델이 부분적인 문맥에 관계없이 유해한 지속을 중단하도록 가르칩니다.
- 유지 세트 ( $D_{retain}$ ): 안전한 응답(또는 유해한 프롬프트에 대한 안전한 완성) 내의 마스킹된 토큰은 원래의 토큰을 재구성하도록 감독되어 일반적인 유용성을 보존합니다.
무작위 마스킹: 학습 과정에서 오염율 $\lambda$ 를 균등하게 샘플링하여 모델이 초기 및 후기 디코딩 단계를 모두 경험하도록 노출시킵니다. 이를 통해 정렬이 전체 생성 궤적 전반에 걸쳐 견고하게 유지되도록 하여 "얕은 정렬" 현상을 방지합니다.
거절 토큰: [EOS] 토큰은 dLLM에서 이미 표준 패딩 및 종료 마커로 사용되고 있으므로, 성능을 저하시킬 수 있는 분포 외(OOD) 토큰을 도입할 필요 없이 모델이 익숙하게 사용할 수 있는 토큰입니다.

실시간 모니터링 및 조기 거절

A2D는 본질적으로 실시간 안전 신호를 제공합니다. 모델이 유해한 콘텐츠가 존재할 때 [EOS]에 높은 확률을 할당하도록 학습되므로, 이 확률 질량(probability mass)은 온라인 안전 지표 역할을 합니다.

동적 개입: 모델은 응답을 시작할 수 있지만, 유해한 지속이 나타나면 자동으로 종료됩니다.
조기 거절: 첫 번째 디코딩 단계의 가장 왼쪽 마스킹된 위치에서 [EOS] 확률을 모니터링함으로써, 시스템은 임계값 기반의 조기 거절 메커니즘을 구현할 수 있습니다. 확률이 임계값 $\tau$ 를 초과하면 모델은 아무런 토큰도 생성하지 않고 즉시 중단되어, 유해한 프롬프트에 대한 추론 지연 시간을 크게 줄입니다.

3. 주요 기여

임의 순서/임의 단계 취약성 식별: 본 논문은 dLLM이 디코딩 단계에 따라 급격히 감쇠하는 얕은 정렬 문제를 겪으며, 디코딩 순서에 비강건하여 DIJA와 같은 공격에 취약함을 입증했습니다.
토큰 수준 정렬 프레임워크: A2D는 응답 수준의 거절에서 [EOS]를 이용한 토큰 수준의 억제로 패러다임을 전환하여, 생성 프로세스의 모든 단계에서 견고한 방어를 가능하게 하는 새로운 정렬 패러다임을 도입했습니다.
실시간 안전 모니터링: 이 방법은 [EOS] 확률을 내부 안전 신호로 활용하여 동적 개입과 유해 프롬프트의 조기 거절을 가능하게 합니다.
포괄적 평가: 저자들은 극단적인 조건(유해한 응답 내에 단 하나의 마스킹된 문장이 포함된 경우)에서 정렬을 스트레스 테스트하기 위해 FITS (Fill-in-the-Sentence) 공격을 도입했습니다.

4. 실험 결과

저자들은 LLaDA-8B-Instruct, LLaDA-1.5, Dream-v0-Instruct-7B의 세 가지 인스트럭션 튜닝된 dLLM을 대상으로 A2D를 RT, SFT, VRPO 베이스라인과 비교 평가했습니다.

탈옥 방어: A2D는 모든 벤치마크(Zeroshot, PAIR, ReNeLLM, Prefilling, DIJA)에서 일관되게 가장 낮은 평균 공격 성공률(ASSR)을 달-성했습니다.
- DIJA 성능: A2D는 베이스라인 모델의 DIJA 성공률을 80% 이상에서 거의 제로 수준으로 낮추었습니다: LLaDA-8B에서 1.3%, LLaDA-1.5에서 3.5%, Dream에서 **0.0%**를 기록했습니다.
- FITS 공격: 극단적인 FITS 공격(유해한 응답에서 단 하나의 문장만 누락된 경우) 하에서, A2D는 모든 모델에서 ASR을 거의 0%로 낮춘 반면, 다른 방법들은 60% 이상의 ASR을 보이며 실패했습니다.
디코딩 전략에 대한 강건성: A2D는 다섯 가지 디코딩 전략(Left-to-Right, Right-to-Left, Confidence, Entropy, Random) 전반에서 낮은 ASR을 유지한 반면, 베이스라인 모델들은 특정 순서에 대해 높은 취약성을 보였습니다.
유용성 보존: A2D는 과잉 거절을 피하며 XSTest 벤치마크(유해해 보이는 양질의 프롬프트)에서 **0%의 거짓 양성(false positives)**을 달성했습니다. 또한 상당한 능력 저하를 겪는 다른 정렬 방법들을 능가하며 일반 능력 벤치마크(MMLU, GSM8K, HumanEval)에서 경쟁력 있는 성능을 유지했습니다.
효율성 이점: 조기 거절 메커니즘은 높은 정렬 충실도를 유지하면서도( $\tau=0.8$ 일 때), 유해한 프롬프트에 대해 최대 19.3배 빠른 안전한 종료를 실현했습니다.

5. 의의 및 주장

본 논문은 A2D가 dLLM의 유연한 디코딩과 전통적인 정렬의 경직된 응답 수준 특성 사이의 근본적인 불일치를 해결함으로써, 디퓨전 기반 텍스트 생성의 안전성에 대한 새로운 기준을 세운다고 주장합니다.

깊은 정렬 (Deep Alignment): A2D는 단순한 시작점이 아닌 생성 궤적 전체에 걸쳐 안전 신호가 지속되도록 하는 "깊은 정렬"을 달성합니다.
순서 독립성 (Order Independence): 디코딩 전략에 관계없이 견고한 방어를 제공하며, 이는 dLLM의 실질적 배포를 위한 필수 요구 사항입니다.
실제적 배포 (Practical Deployment): 실시간 모니터링과 조기 거절을 가능하게 함으로써, A2D는 지연 시간과 안전성이 모두 중요한 실제 애플리케이션에 dLLM을 안전하게 배포할 수 있는 경로를 제시합니다.

저자들은 이 접근 방식이 dLLM의 생성적 유연성과 현실 세계 사용에 필요한 안전 제약 사이의 간극을 메워, 신뢰할 수 있는 디퓨전 기반 언어 시스템을 위한 길을 열어준다고 결론짓습니다.