💬 NLP

Reconsidering Positional Supervision in Masked Diffusion Language Model Training

이 논문은 마스크드 확산 언어 모델(Masked Diffusion Language Models)이 반복적 디코딩 과정 중 미세한 위치 변화에 민감하다는 것을 입증하며, 엄격한 위치 제약을 완화하기 위해 특화된 불확실성 흡수 토큰을 활용하여 연결주의 시간적 분류(CTC)를 적응시키는 방안을 제안하고, 이를 통해 여러 생성 벤치마크에서 통계적으로 유의미한 성능 향상을 달성하였다.

원저자: Mengyu Ye, Keito Kudo, Ryosuke Takahashi, Jun Suzuki

게시일 2026-06-02

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Mengyu Ye, Keito Kudo, Ryosuke Takahashi, Jun Suzuki

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

핵심 아이디어: "경직된" 텍스트 생성기의 문제 해결

당신이 로봇에게 이야기를 쓰는 법을 가르치고 있다고 상상해 보세요.

기존 방식 (자기회귀 방식, Autoregressive): 로봇이 사람이 문장을 타이핑하는 것처럼 한 번에 한 단어씩 써 내려갑니다. 로봇은 문장 속 자신의 위치를 정확히 알고 있습니다.
새로운 방식 (마스크 확산 방식, Masked Diffusion): 로봇이 물음표로 가득 찬 빈 페이지에서 시작합니다. 로봇은 모든 단어를 동시에 추측하려고 시도한 뒤, 몇 개를 드러내고 다시 일부를 숨긴 후 다시 추측합니다. 마치 크로스워드 퍼즐을 한꺼번에 채우는 것처럼, 이 과정을 병렬적으로 수행합니다.

이 새로운 "병렬" 방식은 더 빠르지만, 연구진은 여기서 중대한 결함을 발견했습니다: 로봇은 순서가 아주 조금이라도 어긋나는 것을 극도로 두려워한다는 것입니다.

문제점: "엄격한 선생님"

연구진은 이러한 병렬 모델들이 "엄격한 선생님"(크로스 엔트로피 손실, Cross-Entropy loss라고 불림)에 의해 훈련된다는 사실을 발견했습니다. 이 선생님은 모든 단어가 반드시 지정된 정확한 자리에 앉아야 한다고 요구합니다.

비유: 모든 학생이 특정 책상 번호를 가진 교실을 상상해 보세요. 만약 학생 A가 책상 1번에 앉아야 하는데 실수로 책상 2번에 앉았다면, 설령 그 학생이 맞는 학생일지라도 선생님은 낙제점을 줍니다!
실험: 연구진은 완성된 이야기를 가져와서 단어의 단 1%만을 이웃한 단어와 맞바꾸는 실험(단어를 왼쪽이나 오른쪽으로 한 칸씩 이동)을 진행했습니다.
결과: 모델의 성능이 폭락했습니다. 모델은 이 미세한 뒤섞임에 너무 민감해서, 자기 자신의 이야기를 더 이상 인식하지 못했습니다. 이는 마치 비트가 아주 찰나의 순간만 어긋나도 엉망이 되어버리는 노래와 같았습니다.

해결책: " 토큰"

이를 해결하기 위해 연구진은 음성 인식에서 사용되는 CTC(Connectionist Temporal Classification)라는 기법을 빌려와 CTC-S라는 새로운 이름을 붙였습니다.

"엄격한 선생님" 대신, 그들은 "유연한 코치"를 도입했습니다.

토큰: 연구진은 모델이 <SLACK>이라고 불리는 특별한 보이지 않는 토큰을 사용하도록 가르쳤습니다. 이것은 단어 사이의 "완충 지대" 또는 "간격"이라고 생각하면 됩니다.
작동 원식: 만약 모델이 어떤 단어가 5번 자리에 있어야 한다고 생각하지만, 문맥상 6번 자리에 있는 것이 더 적절하다고 판단될 경우, 모델은 당황하지 않습니다. 모델은 5번 자리에 <SLACK> 토큰을 삽입하여, "여기서 잠시 쉬어가겠다"라고 말하며 단어를 6번 자리에 배치할 수 있습니다.
안전망: 결정적으로, 연구진은 모델이 실수로 단어를 중복 생성하는 경우(예: 병합으로 인해 "100"이 "10"이 되는 경우)가 발생하지 않도록 규칙을 변경했습니다. 모델은 오직 타이밍 오류를 흡수하기 위해서만 <SLACK> 토큰을 사용하며, 실제 단어를 삭제하거나 병합하지는 않습니다.

결과: 더 견고한 작가

연구진은 네 가지 서로 다른 글쓰기 과제(창의적 글쓰기, 어려운 질문에 답하기, 일반적인 대화 등)를 통해 이 새로운 방법을 테스트했습니다.

결과: "유연한 코치"(CTC-S)로 훈련된 모델은 "엄격한 선생님" 모델보다 일관되게 더 나은 이야기를 써냈습니다.
증거: 연구진이 새로운 모델의 출력물에서 단어를 섞는 실험(앞서 언급한 1% 교체 테스트)을 했을 때, 새 모델은 무너지지 않았습니다. 모델은 매우 견고했습니다. 작은 뒤섞임에도 정신을 잃지 않고 잘 처리해 냈습니다.

요점

이 논문은 병렬 언어 모델을 더 잘 작동하게 만들려면, 단순히 모델이 글을 쓴 후(디코딩 단계)에 수정하려고 노력할 것이 아니라, 우리가 그들을 가르치는 방식을 바꿔야 한다고 주장합니다.

모델에게 약간의 "여유(정렬 유연성)"를 줌으로써, 우리는 모델이 너무 취약해지는 것을 방지합니다. 이는 살짝 건드리기만 해도 고장 나는 로봇과, 비틀거리더라도 계속 걸어갈 수 있는 로봇의 차이와 같습니다.

요약하자면: 단어가 반드시 '정확한' 위치에 있어야 한다고 요구하는 대신, 단어가 가는 '위치'에 대해 유연성을 허용하도록 병렬 언어 모델을 가르치는 것이 모델을 훨씬 더 훌륭하게 만든다는 것을 이 논문은 보여줍니다.

문제점: 위치 불일치에 대한 민감도

마스크 확산 언어 모델(Masked Diffusion Language Models, MDLMs)은 토큰을 병렬로 언마스킹(unmasking)하여 텍스트를 생성하며, 이는 비자기회귀 번역(Non-Autoregressive Translation, NAT)과 구조적으로 유사한 설정입니다. 두 패러다임 모두 일반적으로 위치별 교차 엔트로피(cross-entropy, CE) 손실에 의존하며, 이는 예측된 토큰과 고정된 출력 위치 사이의 엄격한 정렬을 강제합니다. 이러한 방식은 훈련을 단순화하지만, "위치적 경직성(positional rigidity)"을 초래합니다.

NAT에서는 CE로 훈련된 병렬 디코더가 미세한 위치 변화에 매우 민감하다는 점이 잘 알려져 있습니다. 즉, 올바른 토큰이 잘못된 위치에 놓일 경우, 틀린 토큰과 동일하게 엄격한 페널티를 받습니다. 그러나 반복적인 병렬 디코딩(디노이징 단계 전반에 걸쳐 오류가 누적될 수 있는 구조)을 사용하는 MDLMs 역시 이러한 민감도를 공유하는지는 불분명했습니다. 저자들은 MDLMs가 미세한 변위(displacement)에도 취약할 수 있으며, 이는 토큰의 정체성이 올바르더라도 생성 품질을 저하시킬 수 있다고 가설을 세웠습니다.

방법론

본 논문은 두 부분으로 구성된 접근 방식을 제안합니다. 첫째, MDLMs의 위치 변화 민감도를 진단하고, 둘째, 수정된 훈련 목적 함수를 통해 이 문제를 완화하는 것입니다.

1. 진단적 개입 (Diagnostic Intervention)
민감도를 테스트하기 위해, 저자들은 LLaDA-8B-Instruct 모델의 디코딩 과정 중에 통제된 개입을 적용했습니다. 그들은 언마스킹된 토큰을 인접한 마스크 위치와 교체하여, 토큰의 정체성은 유지하면서 단 한 위치의 이동(shift)을 도입했습니다.

결과: 생성된 토큰의 단 1%만을 단 한 위치만큼 이동시켰음에도 불구하고, Arena-Hard 벤치마크에서의 승률이 50% 베이스라인을 훨씬 밑도는 수준으로 크게 하락했습니다. 이는 MDLMs가 반복적 병렬 디코딩 환경에서 미세한 위치 변화에 실제로 민감하다는 것을 확인시켜 주었습니다.

2. 제안된 해결책: CTC-S (CTC with Slack)
정렬 유연성을 처리하기 위한 NAT의 CTC(Connectionist Temporal Classification) 성공 사례에서 영감을 얻어, 저자들은 MDLM의 지도 미세 조정(SFT)을 위해 CTC를 변형했습니다. 표준 CTC는 모델이 "blank" 심볼을 방출할 수 있게 하여, blank 배치만 다른 출력 시퀀스를 동일한 것으로 취급합니다. 그러나 표준 CTC는 인접한 중복 토큰(예: "100"을 "10"으로, "**"를 "*"로 축소)을 병합하여 텍스트의 표면 형태를 왜곡하기도 합니다.

이를 해결하기 위해 저자들은 두 가지 핵심 요소를 가진 수정된 목적 함수인 CTC-S를 도입했습니다:

<SLACK> 토큰 증강: 훈련 중에 타겟 응답에 특수 <SLACK> 토큰(CTC의 blank 심볼 역할을 함)을 추가합니다. 슬랙 토큰의 비율은 확산 타임스텝 $t$ 에 따라 단조 감소하며, 이는 생성 과정을 따릅니다.
수정된 축소 맵 ( $\beta_{slack}$ ): 표준 CTC와 달리, CTC-S의 축소 맵은 오직 <SLACK> 토큰만을 제거하고 인접한 중복 내용 토큰은 보존합니다. 이를 통해 모델이 타겟의 표면 형태를 왜곡하지 않으면서(예: "100"을 "100"으로 유지) 위치적 불확실성을 흡수할 수 있도록 합니다.
결합된 목적 함수: 최종 훈련 손실은 표준 마스크 교차 엔트로피 손실( $L_{CE}$ )과 CTC-S 손실( $L_{CTC-S}$ )을 결합합니다:
$L_{SFT} = L_{CE} + \lambda L_{CTC-S}$
$L_{CE}$ 항은 슬랙이 증강된 데이터에 대해 계산되며, $L_{CTC-S}$ 항은 원래의 깨끗한 타겟에 대해 계산되어 모델이 위치적 불확실성을 수용하기 위해 필요할 때 <SLACK> 토큰을 방출하도록 유도합니다.

주요 기여

민감도에 대한 실증적 증거: 본 논문은 통제된 스왑(swap) 개입을 통해 MDLMs가 반복적 디코딩 하에서 미세한 위치 변화에 민감하다는 직접적인 증거를 최초로 제시했습니다.
CTC-S 목적 함수: 표면 형태를 왜곡하지 않으면서(중복 병합을 피함으로써) 정렬 유연성을 제공하도록 설계된, 텍스트 생성을 위한 수정된 CTC 목적 함수를 도입했습니다.
훈련 측면의 설계 차원: 본 연구는 MDLM 성능 향상을 위한 기존의 추론 시점 접근 방식(re-masking 또는 guidance 등)에 대한 보완적인 설계 차원으로 "훈련 측면의 정렬 유연성"을 식별했습니다.

결과

저자들은 LLaDA-8B-Instruct 베이스 모델을 사용하여 네 가지 개방형 생성 벤치마크(Arena-Hard, MT-Bench, WildBench, Creative Writing Bench v3)에서 해당 방법을 평가했습니다.

성능 향상: CE + CTC-S 모델은 네 가지 벤치마크 모두에서 원래의 LLaDA 모델과 동일한 데이터로 훈련되었으나 표준 CE 손실을 사용한 CE-Only 베이스라인보다 일관되게 우수한 성능을 보였습니다.
통계적 유의성: 모든 네 가지 벤치마크에서 두 베이스라인 대비 개선 효과가 통계적으로 유의미했습니다.
강건성 검증: 진단 단계에서 사용된 것과 동일한 위치 변화 개입을 적용했을 때, CE + CTC-S 모델은 강건함을 유지(승률 50% 이상 유지)한 반면, CE-Only 모델은 원래 모델과 동일한 높은 민감도를 보였습니다.
일반적 능력: 일반 능력 벤치마크(GPQA, MMLU, MBPP, IFEval, GSM8K) 평가 결과, CTC-S 목적 함수가 모델의 추론이나 코딩 능력을 저하시키지 않았습니다. 일부 경우(IFEval, GSM8K)에는 베이스 모델과 대등하거나 약간 상회하는 성능을 보였습니다.
추론 시점 해결책과의 비교: 저자들은 자신들의 훈련 측면 접근 방식을 추론 시점의 "Block Diffusion" 스윕(디코딩 캔버스를 제한하는 방식)과 비교했습니다. 최적의 블록 크기를 설정하더라도 CE-Only 베이스라인은 CE + CTC-S 모델의 성능을 따라잡을 수 없었으며, 이는 훈련 측면의 유연성이 추론 시점의 튜닝으로는 회복할 수 없는 이점을 제공함을 나타냅니다.

의의 및 주장

본 논문은 NAT와 MDLM 연구 사이의 간극을 메움으로써, 표준 MDLM 훈련의 결정적인 약점인 위치별 감독의 경직성을 밝혀냈다고 주장합니다. CTC를 MDLM 환경에 적응시킴으로써, 저자들은 **정렬 유연성(alignment flexibility)**이 MDLM SFT에서 중요하지만 이전에 충분히 활용되지 못한 설계 차원임을 입증했습니다.

저자들은 이 연구가 개방형 생성에서 MDLM과 자기회귀(autoregressive) 모델 사이의 성능 격차를 좁히는 방법론적 기여라고 규정합니다. 이들은 자신들의 방법이 베이스 모델 이상의 새로운 능력을 도입하는 것이 아니라, 병렬 디코딩의 본질적인 불확실성을 더 잘 처리할 수 있도록 훈련 목적 함수를 최적화하는 것이라고 명시합니다. 본 연구는 향후 연구가 프리트레이닝(pre-training) 단계에서 위치 정렬을 완화하는 방향을 탐색하고, 이러한 유연한 정렬 목적 함수를 다른 MDLM 인스턴스로 확장해야 함을 시사합니다.