🤖 machine learning

Fine-Tuning Masked Diffusion for Provable Self-Correction

본 논문은 사전 훈련된 마스크드 확산 모델의 추론 과정에서 토큰별 품질 점수를 계산하여 저품질 토큰을 탐지하고 수정함으로써 구조 변경, 강화 학습, 외부 검증자 없이 검증 가능한 자기 수정을 가능하게 하는 경량화 및 모델 무관한 방법인 PRISM 을 소개한다.

원저자: Jaeyeon Kim, Seunggeun Kim, Taekyun Lee, David Z. Pan, Hyeji Kim, Sham Kakade, Sitan Chen

게시일 2026-05-26

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Jaeyeon Kim, Seunggeun Kim, Taekyun Lee, David Z. Pan, Hyeji Kim, Sham Kakade, Sitan Chen

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

퍼즐을 풀거나, 스도쿠를 풀거나, 이야기를 쓰는 것을 상상해 보세요. 하지만 한 번에 단어를 하나씩 채우는 것이 아니라, 모든 빈칸을 한꺼번에 채워야 합니다. 이것이 **마스크 확산 모델 (Masked Diffusion Model, MDM)**이라는 AI 의 작동 방식입니다. 완전히 빈 페이지 (또는 빈 칸으로 가득 찬 격자) 로 시작하여 모든 위치를 동시에 무엇으로 채울지 추측합니다.

문제는 무엇일까요? 모든 것을 한 번에 추측하기 때문에 때로는 실수를 합니다. '2'가 와야 할 자리에 '7'을 넣거나, 문법적으로 의미가 통하지 않는 문장을 쓸 수 있습니다. 과거에는 AI 가 실수를 하면 그 실수에 갇히게 되었습니다. 뒤로 돌아가서 "잠깐, 저게 이상해 보이는데"라고 말하며 수정할 수 없었습니다.

이 논문은 PRISM(Masked Diffusions 의 추론 시간 자기 수정을 위한 플러그인 리마스킹) 이라는 새로운 방법을 소개합니다. PRISM 을 AI 가 작업하는 동안 사용할 두 번째 눈과 빨간 펜을 주는 것이라고 생각하세요.

간단한 개념으로 나누어 설명하면 다음과 같습니다:

1. 문제: "눈가림 추측"

크로스워드 퍼즐을 채우는 상황을 상상해 보세요. 하지만 다른 단어의 힌트를 보지 않은 채 모든 단어를 정확히 같은 순간에 추측해야 합니다. 3 글자 단어에 "CAT"이라고 추측했을지라도, 나중에 교차하는 단어 "DOG"와 맞지 않는다는 것을 깨닫게 될 수 있습니다. 표준 MDM 에서는 한 번 "CAT"이라고 쓰면 그대로 남습니다. 모델은 퍼즐 전체의 맥락에서 "CAT"이 좋은 추측이었는지 확인하는 방법을 모릅니다.

2. 해결책: "품질 점수"

PRISM 은 AI 가 쓰는 모든 단어나 숫자에 품질 점수를 부여하도록 가르칩니다.

높은 점수: "이 단어는 문장의 나머지 부분과 완벽하게 어울립니다."
낮은 점수: "이 단어는 여기서 이상하게 느껴집니다. 틀렸을지도 모릅니다."

이 논문은 PRISM 이 이러한 점수를 수학적으로 증명 가능하게 정확하게 부여하도록 학습한다고 주장합니다. 단순히 추측하는 것이 아니라, 주변 단어들을 고려할 때 특정 단어가 정확할 확률을 수학적으로 계산합니다.

3. 마법 같은 트릭: "리마스킹 (Remasking)"

AI 가 품질 점수를 얻으면 리마스킹이라는 교묘한 트릭을 사용합니다.

만약 어떤 단어가 낮은 점수(틀릴 가능성이 높음) 를 받으면, AI 는 그것을 지웁니다 (다시 "마스크"를 씌웁니다).
그런 다음, 그 자리에 새로운 단어를 추측해 봅니다.
만약 단어가 높은 점수를 받으면, 그대로 둡니다.

이 과정은 생성 도중에 일어납니다. AI 는 단순히 쓰는 것이 아니라, 실시간으로 쓰고, 확인하고, 지우고, 다시 씁니다. 마치 작가가 문장을 타이핑하다가 단어가 틀렸음을 깨닫고, 다음 문장으로 넘어가기 전에 그것을 삭제하고 더 좋은 단어로 다시 타이핑하는 것과 같습니다.

4. 이것이 특별한 이유 ("플러그 앤 플레이" 측면)

일반적으로 AI 가 자신의 실수를 수정하도록 가르치려면 뇌 전체를 처음부터 다시 구축하거나, 작업을 점검하는 별도의 무거운 '비평가 (critic)' 모델을 사용해야 합니다.

PRISM 은 경량입니다: 기존 AI 에 작은 "어댑터"나 플러그인을 추가하는 것과 같습니다. 처음부터 전체를 다시 학습할 필요가 없습니다.
효율적입니다: 논문은 PRISM 이 원래 AI 를 학습시키는 데 걸린 것보다 훨씬 적은 추가 데이터로 자기 수정을 학습할 수 있음을 보여줍니다.
어디서나 작동합니다: 저자들은 이를 다음과 같은 분야에서 테스트했습니다:
- 스도쿠: 격자 속 잘못된 숫자 수정.
- 텍스트: 더 자연스러운 문장 작성.
- 코드: 파이썬 프로그래밍 오류 (누락된 괄호나 잘못된 논리 등) 수정.

5. 결과

논문에 따르면 PRISM 을 사용했을 때:

AI 의 실수가 줄었습니다.
스스로 오류를 수정하여 코딩과 퍼즐에서 더 나은 최종 결과를 얻을 수 있었습니다.
막대한 추가 컴퓨팅 파워 없이도 매우 큰 모델 (예: 80 억 파라미터 코딩 어시스턴트) 에서도 작동했습니다.

한 마디로 요약

PRISM은 특정 유형의 AI 를 위한 간단하고 똑똑한 업그레이드입니다. AI 에게 자신의 작업을 보고 "저건 이상해 보이는데"라고 말하며 실수를 지우고 다시 시도하는 능력을 부여합니다. 모든 것이 답변을 생성하는 도중 일어납니다. 이는 "한 번의 추측"을 "초안 작성 및 편집" 프로세스로 바꾸어, AI 가 퍼즐을 풀고, 글을 쓰고, 코딩할 때 훨씬 더 신뢰할 수 있게 만듭니다.

기술 요약: 증명 가능한 자기 수정을 위한 마스킹 확산 모델의 미세 조정

문제 정의

마스킹 확산 모델 (MDM) 은 이산 공간 (예: 텍스트, 코드, 스도쿠) 에서의 생성 모델링을 위한 autoregressive 모델의 유망한 대안으로 부상했습니다. 토큰을 순차적으로 생성하는 autoregressive 모델과 달리, MDM 은 완전히 마스킹된 시퀀스로 시작하여 임의의 순서로 토큰을 점진적으로 마스킹 해제합니다. 이러한 병렬적이고 소수의 단계로 이루어진 추론은 효율성과 유연성을 제공하지만, 중요한 한계를 도입합니다: 표준 MDM 은 자기 수정 능력을 결여하고 있습니다.

표준 MDM 추론에서, 일단 토큰이 마스킹 해제되면 고정됩니다. 이러한 경직성은 특히 추론 단계가 적을 때 위치 간 의존성 오류에 MDM 을 취약하게 만듭니다. 이를 해결하려는 이전 시도들은 다음과 같은 중대한 단점에 직면했습니다:

부정확한 대리 지표: 훈련이 필요 없는 방법들 (예: ReMDM, DFM) 은 무작위 재마스킹이나 토큰 품질을 나타내도록 훈련되지 않은 로그이트 (logits) 에서 도출된 점수에 의존하여, 비효율적이거나 부정확한 수정을 초래합니다.
아키텍처 과부하: 자기 수정을 위한 유효한 목표를 정의하는 방법들은 종종 MDM 아키텍처 (예: hollow-transformers) 나 사전 훈련 목적을 변경해야 하므로, 기존 사전 훈련된 모델에 대한 적용 범위를 제한합니다.

핵심 아키텍처를 변경하지 않고 임의의 사전 훈련된 MDM에 자기 수정 능력을 부여하면서도 "토큰 품질"에 대한 이론적으로 근거 있는 정의를 제공할 수 있는 방법이 필요합니다.

방법론: PRISM

저자들은 PRISM(Plug-in Remasking for Inference-time Self-correction of Masked Diffusions, 추론 시간 마스킹 확산의 자기 수정을 위한 플러그인 재마스킹) 을 제안합니다. 이는 경량화되고 모델에 구애받지 않는 미세 조정 프레임워크입니다.

1. 이론적 기반: 토큰별 품질

PRISM 은 자기 수정을 위한 엄격한 목표를 정의합니다: 토큰별 품질 점수입니다. 부분적으로 마스킹된 시퀀스 $y$ 에 대해, 위치 $i$ 의 토큰 $y_i$ 의 품질은 $y_i$ 가 마스킹된 나머지 시퀀스를 조건으로 할 때 해당 토큰의 가능성으로 정의됩니다:
$g^*_i(y) := p(x_i = y_i \mid y \oplus m_i)$
여기서 $y \oplus m_i$ 는 $i$ 번째 토큰이 마스킹 토큰 $m$ 으로 대체된 시퀀스 $y$ 를 나타냅니다. $y_i$ 가 문맥과 일치하면 이 확률은 높고, 오류인 경우 확률은 낮습니다.

2. PRISM 손실 함수

$g^*$ 를 근사하는 모델 $g_\phi$ 를 학습하기 위해 PRISM 은 새로운 손실 함수를 도입합니다. 여기서의 과제는 $g_\phi$ 가 입력 $y$ 에서 이미 마스킹 해제된 토큰 $y_i$ 의 품질을 예측해야 하지만, 목표는 마스킹된 문맥 $y \oplus m_i$ 에 정의된 사후 확률이라는 점입니다.

PRISM 은 특정 데이터 구성과 손실 형식을 통해 이를 해결합니다:

데이터 구성: 진짜 시퀀스 $x$ 와 마스킹된 버전 $z$ 를 샘플링합니다. 사전 훈련된 MDM $f_\theta$ 를 사용하여 $z$ 의 토큰 부분집합을 마스킹 해제하여 새로운 시퀀스 $y$ 를 생성합니다.
손실 함수: 모델 $g_\phi$ 는 이진 교차 엔트로피 (BCE) 를 사용하여 이진 레이블 $1[x_i = y_i]$ (마스킹 해제된 토큰이 정답과 일치하는지 여부) 을 예측하도록 훈련됩니다:
$\mathcal{L}(\phi) := \mathbb{E}_{x,z,(y,i)} \left[ \text{BCE}(1[x_i = y_i], g_\phi^i(y)) \right]$
이론적으로 이 손실의 유일한 최소값은 $y$ 를 생성하는 데 사용된 사전 훈련된 MDM $f_\theta$ 의 품질과 관계없이 원하는 토큰별 품질 $p(x_i = y_i \mid y \oplus m_i)$ 과 정확히 일치합니다.

3. 구현 세부 사항

플러그 앤 플레이: PRISM 은 임의의 사전 훈련된 MDM 에 경량 보조 어댑터 ("품질 헤드") 를 연결합니다. 백본은 공유되며, 하나의 헤드는 마스킹 해제 사후 확률 (원래 MDM 기능) 을 출력하고 새로운 헤드는 스칼라 품질 점수를 출력합니다.
정규화: MDM 의 원래 마스킹 해제 능력을 보존하기 위해, 미세 조정 동안 표준 MDM 훈련 손실이 정규화 항으로 추가됩니다.
효율성: 이 프레임워크는 $f_\theta$ 의 단일 순전파를 통해 여러 토큰을 동시에 마스킹 해제함으로써 단일 훈련 쌍 $(y, i)$ 를 생성할 수 있게 하여 데이터 효율성을 향상시킵니다.

4. 추론 절차

추론 중, 미세 조정된 MDM 은 각 단계에서 두 가지 작업을 동시에 수행합니다:

마스킹 해제: 표준 MDM 사후 확률에 기반하여 마스킹된 토큰의 부분집합을 선택하여 마스킹 해제합니다.
재마스킹: 가장 낮은 품질 점수 ( $g_\theta$ 가 예측한) 를 가진 이미 마스킹 해제된 토큰의 부분집합을 선택하여 다시 마스킹합니다.
이러한 단계들은 시퀀스가 최종화되기 전에 모델이 오류를 반복적으로 수정할 수 있게 합니다.

주요 기여

이론적 보장: PRISM 은 전체 깨끗한 시퀀스를 생성하거나 검증자를 사용할 필요가 없는 MDM 의 토큰별 품질에 대한 최초의 증명 가능한 학습 목적을 제공합니다. 이 손실의 최소값이 마스킹된 문맥에 대한 토큰의 실제 조건부 가능성으로 수렴함을 증명합니다.
모델에 구애받지 않는 설계: 이전 작업과 달리 PRISM 은 아키텍처 변경이나 처음부터의 재훈련을 요구하지 않습니다. 기존 모든 사전 훈련된 MDM 을 위한 "플러그인" 미세 조정 계층으로 작동합니다.
효율성: 이 방법은 사전 훈련에 비해 훨씬 적은 훈련 샘플 (수십 배에서 수백 배 적음) 과 계산 자원으로 자기 수정 능력을 학습합니다.

실험 결과

저자들은 PRISM 을 세 가지 도메인과 규모에 걸쳐 평가했습니다:

스도쿠 (30M 규모 MDM):
- PRISM 은 미세 조정 중 점진적으로 높은 성공률을 달성하여, ReMDM 및 ReMDM-conf 와 같은 베이스라인을 몇 개의 에포크 내에 능가했습니다.
- 시각화는 PRISM 이 그리드의 잘못된 셀에 낮은 품질 점수를 올바르게 할당하여 효과적인 오류 감지를 가능하게 함을 확인시켜 줍니다.
무조건부 텍스트 생성 (170M 규모 MDM):
- OpenWebText 에서 생성 퍼플렉시티 (Generative Perplexity) 와 MAUVE 를 사용하여 평가했습니다.
- PRISM 은 베이스라인을 능가했으며, 특히 의존성 오류가 가장 흔하게 발생하는 적은 샘플링 단계 (예: 64~256 단계) 영역에서 두드러졌습니다.
- 보정 연구는 예측된 품질 점수가 토큰의 경험적 가능성과 밀접하게 추적됨을 보여주었습니다.
코드 생성 (8B 규모 MDM - LLaDA):
- 500 GPU 시간 미만의 미세 조정을 LLaDA-8B-Instruct 에 적용했습니다.
- MBPP 및 HumanEval 벤치마크에서 PRISM 은 다양한 샘플링 단계에 걸쳐 순수 MDM 및 기타 재마스킹 베이스라인을 능가했습니다.
- 기계적 분석은 PRISM 이 생성된 코드에서 구문 오류 (예: 불완전한 elif 문장) 와 논리 오류 (예: 구분자 배치) 를 성공적으로 수정함을 보여주었습니다.
- 전이성: 한 MDM 백본으로 훈련된 PRISM 조정 모델은 동일한 분포에서 작동하는 다른 사전 훈련된 MDM 백본에 대한 자기 수정기로 기능할 수 있으며, 이는 상당한 성능 향상 (예: HumanEval 에서 +10.9%) 을 가져옵니다.

중요성과 주장

이 논문은 PRISM 을 수정과 재배열을 통해 이산 시퀀스를 구성하는 생성 모델 구축을 향한 한 걸음으로 위치시킵니다. 이는 인간과 유사한 생성 과정을 모방합니다.

증명 가능한 자기 수정: 저자들은 PRISM 이 전체 시퀀스 생성이 필요한 강화 학습이나 외부 검증자에 의존하지 않고 토큰별 품질에 대한 정확하고 검증 가능한 목표를 정의하고 학습하는 유일한 방법임을 강조합니다.
확장성: 이 방법은 대규모 모델 (8B 파라미터) 에서도 최소한의 계산 자원으로 미세 조정을 통해 자기 수정을 효율적으로 획득할 수 있음을 보여줍니다.
한계: 저자들은 품질 점수가 위치별 사후 주변분 (posterior marginals) 에 기반하기 때문에, PRISM 은 로컬 토큰 가능성에 포착되지 않는 전역 추론 오류 (예: 여러 토큰에 걸친 구조적 논리 오류) 에는 어려움을 겪을 수 있음을 겸손하게 인정합니다. 그러나 그들은 이 작업이 이산 확산의 효율적인 훈련 패러다임 내에서 전역 오류를 해결할 미래 프레임워크에 대한 문을 열었다고 주장합니다.