💬 NLP

Revise, Don't Freeze: Sampler-Matched Training for Self-Correcting Masked Diffusion Language Models

이 논문은 마스크된 확산 언어 모델(masked diffusion language models)에서 직접적인 토큰 수정을 가능하게 하는 파라미터 프리 샘플러인 D3IM과, 그로 인해 발생하는 보존 편향(preservation bias)을 완화하기 위한 사후 학습 방법인 SCOPE를 소개하며, 이들은 종합적으로 추론 및 코딩 벤치마크에서 상당한 성능 향상을 달성한다.

원저자: Longxuan Yu, Shaorong Zhang, Yu Fu, Hui Liu, Yue Dong, Greg Ver Steeg

게시일 2026-06-02

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Longxuan Yu, Shaorong Zhang, Yu Fu, Hui Liu, Yue Dong, Greg Ver Steeg

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

큰 문제: "얼어붙은 실수 (The Frozen Mistake)"

당신이 매우 똑똑하지만 약간은 긴장한 로봇 조수와 함께 이야기를 쓰고 있다고 상상해 보세요. 로봇이 단어를 추측할 때마다, 로봇은 그 단어를 종이에 적습니다.

이 로봇들이 작동하는 표준 방식(마스크 확산 언어 모델, Masked Diffusion Language Models라고 불림)에서는, 로봇이 단어를 적고 나서 충분히 "확신"을 느끼면, 그 단어를 그 자리에 얼려버립니다(freeze). 로봇은 그 단어 위에 테이프를 붙여 버립니다. 설령 로봇이 나중에 "잠깐, 이 단어는 이 문장에서 말이 안 되는데?"라고 깨닫더라도, 테이프가 붙어 있기 때문에 그 단어를 바꿀 수 없습니다. 로봇은 아직 적히지 않은 빈칸들을 채울 수만 있을 뿐입니다.

이는 수학이나 코딩 같은 까м 까다로운 작업에서 큰 문제가 됩니다. 만약 로봇이 초반에 잘못된 숫자(예를 들어, '2' 대신 '5'라고 쓰는 경우)를 적어버리면, 로봇은 그 상태로 갇혀버립니다. 로봇은 그 잘못된 숫자를 바탕으로 나머지 답을 만들어내려 노력하게 되고, 결국 완전히 틀린 결과를 낳게 됩니다.

해결책: 두 단계의 처방

저자인 Longxuan Yu와 동료들은 로봇이 실제로 마음을 바꿀 수 있는 능력이 있지만, 게임의 규칙이 그것을 방해하고 있다는 사실을 깨달았습니다. 그들은 이를 해결하기 위해 두 부분으로 구성된 솔루션을 도입했습니다: 새로운 게임 방식(D3IM)과 로봇을 위한 특별한 훈련 연습(SCOPE)입니다.

파트 1: 새로운 게임 규칙 (D3IM)

비유: 모두가 문장의 다음 단어를 추측하는 "뜨거운 감자(Hot Potato)" 게임을 상상해 보세요.

기존 규칙: 일단 단어를 외치면, 끝날 때까지 그 단어를 계속 붙들고 있어야 합니다. 만약 잘못된 단어를 외쳤다면, 당신은 그것에 묶여 있게 됩니다.
D3IM 규칙 ("백지 상태"): 게임의 매 단계마다, 로봇은 이미 외쳤던 단어를 포함하여 문장의 모든 단어를 다시 살펴봅니다. 그리고 스스로에게 묻습니다. "나는 여전히 이것이 최선의 단어라고 생각하는가?"
- 만약 로봇이 여전히 확신한다면, 그 단어는 유지됩니다.
- 만약 로봇이 "사실, 이건 틀렸어"라고 생각한다면, 즉시 그 단어를 지우거나(다시 빈칸으로 돌리거나) 더 나은 단어로 교체합니다.
- 로봇은 허락을 구하거나 특별한 도구를 사용할 필요가 없습니다. 그저 현재의 확신도를 바탕으로 모든 것을 재평가할 뿐입니다.

함정: 로봇은 이 게임을 아주 잘 수행할 수 있습니다. 단, 자신이 언제 틀렸는지 알 수 있다면 말이죠. 하지만 원래의 훈련 과정에서 로봇은 자신의 실수를 인정하는 법을 배운 적이 없습니다. 로봇은 자신의 첫 번째 추측이 설령 나쁘더라도 그것을 믿도록 훈련되었습니다. 그래서 새로운 "백지 상태" 게임을 하게 되면, 로봇은 자신의 추측이 옳다고 생각하기 때문에 계속해서 잘못된 단어들을 얼려버립니다.

파트 2: 특별 훈련 (SCOPE)

비유: 이것은 "실수를 포함한 리허설"과 같습니다.
로봇의 잘못된 습관(자신의 오류를 맹신하는 것)을 고치기 위해, 저자들은 SCOPE(Self-Conditioned On Prediction Errors)라고 불리는 훈련 방법을 만들었습니다.

작동 방식: 훈련 중에 로봇은 문장을 쓰라는 요청을 받지만, 그 후 자신의 작업물을 보고 마치 실수를 한 것처럼 가정해야 합니다.
로봇은 문장을 작성합니다. 그다음 훈련사는 말합니다. "좋아, 내가 네 단어 중 일부를 숨길게. 이제 그것들을 다시 맞춰봐."
반전: 훈련사는 로봇이 틀렸으면서도 매우 높은 확신을 가졌던 단어들을 특별히 골라냅니다. 로봇은 자신의 틀린 답을 보고, 그것이 틀렸음을 깨닫고, 올바른 답으로 교체해야 합니다.
이는 로봇에게 중요한 교훈을 줍니다: "내가 높은 확신을 가지고 말했다고 해서 그것이 반드시 맞는 것은 아니다. 나는 마음을 바꿀 준비가 되어 있어야 한다."

결과: 협력하는 팀

새로운 게임 규칙(D3IM)과 특별 훈련(SCOPE)을 결 함께하면, 로봇은 스스로 교정하는 천재가 됩니다.

훈련 없이: 특별한 훈련 없이 새로운 규칙만 준다면, 로봇은 오히려 더 나빠집니다. 단어를 바꾸려고 시도하지만, 자신의 오류를 찾아내는 능력을 믿지 못하기 때문에 계속 같은 실수를 반복합니다.
훈련과 함께: 로봇은 자신의 "얼어붙은 실수"를 찾아내는 법을 배웁니다. 이제 로봇은 수학 문제에서 틀린 숫자를 보고 "아니, 이건 틀렸어"라고 말하며, 문장이 만들어지는 동안 그것을 올바른 것으로 바꿀 수 있습니다.

실제 결과

이 논문은 LLaDA-8B(스마트 언어 모델)를 사용하여 어려운 작업들에 대해 테스트했습니다:

수학 (GSM8K & MATH): 점수가 크게 상승했습니다. 예를 들어, 어려운 수학 테스트에서 정확도가 **55.3%**에서 **68.3%**로 뛰었습니다.
코딩 (HumanEval & MBPP): 작동하는 코드를 작성하는 능력이 극적으로 향상되었습니다. 한 테스트에서는 **14.0%**에서 **29.3%**로 올라갔습니다.

핵심 요점

주요 발견은 마음을 바꿀 능력을 갖추는 것만으로는 충분하지 않으며, 언제 마음을 바꿔야 하는지 알도록 훈련받아야 한다는 것입니다.

저자들은 모델이 자신의 확신에 찬 오류를 인식하도록 가르치고(SCOPE), 그 오류를 더 나은 것으로 즉시 교체할 수 있는 메커니즘(D3IM)을 제공함으로써, 모델이 이전보다 훨씬 더 복잡한 추론 문제를 해결할 수 있음을 보여주었습니다. 이것은 로봇을 더 똑똑하게 만드는 것이 아니라, 초기 실수에 대해 덜 고집스럽게 만드는 법을 가르치는 것입니다.

기술 요약: 수정하되, 고정하지 마라: 자기 교정형 마스크 확산 언어 모델을 위한 샘플러 매칭 학습

1. 문제 정의

LLaDA와 같은 마스크 확산 언어 모델(Masked Diffusion Language Models, MDLMs)은 반복적으로 마스킹된 시퀀스를 디노이징하고 각 단계에서 모든 위치를 재예측하며 텍스트를 생성한다. 이러한 아키텍처는 본질적으로 "수정 능력(revision capability)"을 갖추고 있다: 즉, 특정 토큰이 가시화된 후에도 디노이저가 이후 단계에서 해당 위치에 다른 예측을 할 수 있는 이론적 가능성이 존재한다.

그러나 표준적인 MDLM 디코딩 전략은 **단조적 언마스킹(monotonic unmasking)**을 사용하며, 이는 토큰이 결정되면 이를 "고정(freeze)"해 버린다. 일단 토큰이 드러나면, 그것은 클린 토큰으로 취급되며 후속 단계는 오직 마스킹된 위치만을 업데이트한다. 이는 두 가지 결정적인 문제를 야기한다:

영구적 초기 오류(Permanent Early Errors): 생성 과정 초기에 발생한 실수(예: 수학 문제의 잘못된 숫자나 코드의 잘못된 연산자)는 영구적으로 고정되어, 추론 중심 작업의 성능에 심각한 영향을 미친다.
수정 격차(The Revision Gap): 기존의 수정 방식은 휴리스틱 메커니즘(신뢰도 임계값, 검증 단계)에 의존하거나, 결정된 토큰을 다시 예측하기 위해 [MASK] 상태로 되돌리는 재마스킹(remasking) 과정을 필요로 한다(예: ReMDM). 보조 모듈이나 재마스킹 단계 없이 **직접적인 가시적-토큰-간 수정(direct visible-to-visible revision)**을 허용하는 원칙적인 샘플러는 아직 탐구되지 않았다.

또한, 저자들은 **보존 편향(preservation bias)**이라 명명된 모델 측면의 장애물을 식별했다. 학습 목표가 토큰 보존을 강제하지 않더라도, 표준 마스크 목표로 학습된 모델은 스스로 생성한 컨텍스트에 노출되었을 때 자신의 잘못된 결정된 토큰을 수정하기보다 그대로 재생산하는 경향을 보인다.

2. 방법론

본 논문은 함께 작동하도록 설계된 두 부분의 솔루션인 D3IM(새로운 샘플러)과 SCOPE(학습 절차)를 제안한다.

2.1 D3IM: 백지상태 재마스킹 (Clean-Slate Remasking)

D3IM은 **교정 스타일의 역방향 업데이트(corrector-style reverse update)**로서 유도된 파라미터 프리(parameter-free) 샘플러이다. 결정된 토큰을 보존하는 표준 샘플러(보존 가중치 $c_t > 0$ )와 달리, D3IM은 $c_t = 0$ 으로 설정한다.

메커니즘: 매 디노이징 단계마다 모델은 모든 활성 위치를 예측한다. D3IM은 이 예측들을 신뢰도에 따라 순위를 매기고 상위 $K$ 개의 위치만을 유지한다.
수정 채널:
- 토큰-대-토큰(Token-to-Token): 이전에 결정된 토큰이 [MASK] 상태를 거치지 않고 새로운, 더 높은 신뢰도의 예측에 의해 직접 덮어써질 수 있다.
- 토크-대-마스크(Token-to-Mask): 신뢰도가 낮은 결정된 토큰은 다시 [MASK]로 강등되어 후속 단계에서 재예측될 수 있다.
제약 조건: 다음 상태는 이전의 가시적 토큰을 복사하는 것이 아니라 전적으로 현재의 예측에 의해 결정된다. 이를 위해서는 모델의 신뢰도 순위가 자기 생성 컨텍스트 하에서도 신뢰할 수 있어야 한다.

2.2 SCOPE: 예측 오류에 대한 자기 조건화 (Self-Conditioned On Prediction Errors)

D3IM을 기본 MDLM에 직접 적용하면 모델이 자신의 고신뢰도 오류를 수정하는 능력이 부족하기 때문에 성능이 저하된다(보존 편향 때문). SCOPE는 모델의 신뢰도 행동을 D3IM의 샘플링 역학에 정렬하기 위해 설계된 경량 포스트 트레이닝 절차이다.

과정:
1. 모델은 부분적으로 마스킹된 입력에 대해 순방향 패스(forward pass)를 수행한다(자기 조건화).
  2._ 모델은 마스킹된 위치들에 대한 예측을 샘플링한다.
2. 이 예측들 중 신뢰도에 따른 상위 $\rho$ 비율의 예측들은 가시적 토큰으로 "결정(committed)"되는 반면, 나머지는 마스킹 상태로 남는다. 결정적으로, 이 과정은 모델을 자신의 고신뢰도 오류에 노출시킨다.
3. 그 후 모델은 결정된(잠재적으로 오류가 있는) 토큰과 남은 마스킹된 위치 모두에 대해 정답 시퀀스를 복구하도록 학습된다.
목표: 모델이 가시적인 자기 생성 토큰을 고정된 컨텍스트가 아닌 하나의 가설로 취급하도록 학습시켜, 보존 편향을 줄이고 D3IM의 상위- $K$ 선택을 위한 신뢰도 순위가 신뢰할 수 있게 만든다.

3. 주요 기여

D3IM 샘플러: 보조 모듈이나 재마스킹 없이 직접적인 토큰 수정을 가능하게 하는, 사전 학습된 MDLM으로부터 유도된 파라미터 프리의 직접 수정 샘플러이다.
보존 편향 식별: 저자들은 표준 MDLM 학습이 모델로 하여금 학습 목표가 수정을 허용함에도 불구하고 자신의 잘못된 가시적 예측을 재생산하게 만드는 행동적 경향(보존 편향)을 유도함을 입증했다.
SCOPE 학습: D3IM의 샘플링 역학을 시뮬레이션하는 경량 포스트 트레이닝 절차이다. 이는 보존 편향을 효과적으로 줄여 D3IM이 의도대로 작동하게 한다.
샘플러-모델 정렬: 수정 능력은 일반적인 것이 아니며, 학습 목표가 (D3IM의 글로벌 top- $K$ 규칙과 같이) 추론 샘플러의 결정 규칙과 구체적으로 일치해야 효과적임을 밝혀냈다.

4. 실험 결과

실험은 LLaDA-8B-Instruct를 대상으로 네 가지 벤치마크(GSM8K, MATH-500, HumanEval, MBPP)에서 수행되었다.

성능 향상: 64단계 디노이징 시, SCOPE + D3IM 조합은 표준 언마스킹을 사용하는 원래의 LLaDA-8B를 크게 상회한다:
- GSM8K: +13.0% (68.3% vs. 55.3%)
- MATH-500: +4.8% (23.6% vs. 18.8%)
- HumanEval: +15.3% (29.3% vs. 14.0%)
- MBPP: +10.4% (30.8% vs. 20.4%)
스케일링: 수학 및 코드 작업에서 디노이징 단계가 증가함에 따라 이득이 증가한다 (예: GSM8K는 256단계에서 73.8%에 도달).
정렬의 필요성:
- SCOPE 없는 D3IM: 성능이 급격히 하락하여(예: GSM8K가 48.0%로 하락), 적응되지 않은 모델은 어떤 토큰을 수정할지 신뢰성 있게 결정할 수 없음을 확인했다.
- D3IM 없는 SCOPE: SCOPE로 학습된 모델을 표준 단조적 언마스킹과 함께 사용할 경우 결과가 엇갈리며, 이는 학습의 이점이 D3IM 샘플러의 수정 메커니즘에 특화되어 있음을 보여준다.
메커니즘 진단:
- 보존 편향 해결: 모델이 잘못된 토큰을 결정했을 때, 원래 모델은 오류를 70.1% 확률로 보존한다. SCOPE는 이를 17.0%로 줄이며 정답을 50.2% 확률로 복구한다.
- 궤적 분석: SCOPE가 없으면 D3IM의 수정은 종종 "콘텐츠-대-공백 진동(content-to-whitespace oscillation)"(토큰과 줄바꿈 사이를 왔다 갔다 함)을 유발한다. SCOPE와 함께라면 수정은 특정 의미론적 오류(예: 단위 변환, 숫자 수정)를 타겟팅하며 빠르게 안정화된다.

5. 의의 및 주장

본 논문은 MDLM의 자기 교정 추론을 가로막는 병목 현상이 토큰을 재방문할 수 있는 능력이 아니라, 자기 생성 컨텍스트 하에서의 신뢰도 순위의 신뢰성이라고 주장한다.

원칙적인 수정: D3IM은 휴리스틱 임계값이나 간접적인 재마스킹을 넘어 직접적인 토큰 수정을 제공하는 이론적으로 근거 있는 방법을 제공한다.
학습-추론 정렬: 본 연구는 고급 추론 전략(예: 글로벌 top- $K$ 재순위화)이 작동하기 위해서는 특정 학습 목표(SCOPE)가 필요함을 강조한다. 표준 지속 사전 학습이나 일반적인 자기 교정 방법은 이러한 새로운 샘플링 역학으로 자동 전이되지 않는다.
효율성: 제안된 솔루션은 계산 효율적이며, 단 하나의 A100에서 약 20분 동안 500 스텝의 경량 LoRA 파인튜닝만 필요할 뿐, 추론 시 아키텍처 변경이나 보조 모듈이 필요하지 않다.

저자들은 D3IM+SCOPE가 추론 및 코드 생성 능력을 크게 개선하지만, 토큰 수준의 수정에 국한되며 초기 궤적이 근본적으로 잘못된 경우 전체적인 추론 경로를 재구성할 수는 없다고 결론지었다. 향อน 연구에서는 풍부한 수정 궤적을 프로세스 보상 모델링(process reward modeling)이나 강화 학습에 활용하는 방안을 탐구할 수 있을 것이다.