💬 NLP

Improving Sampling for Masked Diffusion Models via Information Gain

이 논문은 마스킹 확산 모델 (MDM) 의 기존 샘플링 방식이 가진 국소적 확신성 편향을 극복하고, 현재 복호화 결정이 미래 토큰에 미치는 정보 획득 효과를 고려한 'Info-Gain Sampler'를 제안하여 추론, 코딩, 창의적 글쓰기 등 다양한 작업에서 기존 방법보다 뛰어난 성능을 달성함을 보여줍니다.

원저자: Kaisen Yang, Jayden Teoh, Kaicheng Yang, Yitong Zhang, Alex Lamb

게시일 2026-03-19

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Kaisen Yang, Jayden Teoh, Kaicheng Yang, Yitong Zhang, Alex Lamb

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

1. 배경: 퍼즐을 맞추는 두 가지 방법

과거의 인공지능 ( autoregressive 모델) 은 퍼즐 조각을 왼쪽에서 오른쪽으로 하나씩 채워가는 방식이었습니다. 마치 글을 쓸 때 첫 글자부터 끝까지 순서대로 쓰는 것과 같아요.

하지만 이 논문에서 다루는 **새로운 AI(MDM)**는 빈칸이 있는 퍼즐을 줍니다.

예: "사과를 ___다." (빈칸에 '먹'을 채워야 함)
이 AI 는 빈칸을 어떤 순서로 채우든 상관없이 모든 빈칸을 동시에 볼 수 있습니다. (왼쪽에서 오른쪽이 아니라, 가장 쉬운 빈칸부터 채워도 되고, 어려운 빈칸부터 채워도 됩니다.)

이게 큰 장점처럼 보이지만, 문제는 **"어떤 빈칸을 먼저 채워야 가장 잘 맞을까?"**를 결정하는 게 매우 어렵다는 점입니다.

2. 기존 방법의 문제점: "눈앞의 쉬운 것만 보는 나쁜 나침반"

기존의 AI 는 **"가장 확실해 보이는 빈칸부터 채워라"**라는 간단한 규칙 (그리디한 방식) 을 따랐습니다.

상황: "100 을 2 로 나누면 몇?"이라는 문제를 푼다고 가정해 봅시다.
AI 의 생각: "아, '2'와 '100'은 이미 정해져 있고, 답인 '50'은 숫자 하나만 고르면 되니까 확실히 쉽게 맞출 수 있겠다! 일단 '50'부터 채우자!"
결과: AI 는 답을 먼저 채웠지만, 정작 '100'을 '2'로 나눈 계산 과정 (논리) 을 제대로 생각하지 못했습니다. 그래서 엉뚱한 답을 내거나, 나중에 다른 빈칸을 채울 때 모순이 생깁니다.

이를 **"눈앞의 확실한 것만 보고, 전체 그림을 보지 못하는 근시안 (Myopia)"**이라고 부릅니다. 마치 미로에서 가장 가까운 출구가 보인다고 해서 바로 그곳으로 뛰어가다가, 막다른 길에 막혀버리는 것과 같습니다.

3. 이 논문의 해결책: "정보의 보석"을 찾는 Info-Gain Sampler

이 연구팀은 AI 에게 **"가장 쉬운 빈칸을 먼저 채우는 게 아니라, 전체 퍼즐을 해결하는 데 가장 큰 도움을 주는 빈칸을 먼저 채워라"**라고 가르쳤습니다. 이를 **정보 이득 (Information Gain)**이라고 부릅니다.

🌟 핵심 비유: "미로 찾기 vs. 지도 읽기"

기존 AI (그리디 방식): 미로에서 "지금 바로 오른쪽으로 가면 길이 열려 있네? 저기로 가자!"라고 합니다. 하지만 그 길은 나중에 막다른 길로 이어질 수 있습니다.
새로운 AI (Info-Gain Sampler): "지금 오른쪽으로 가면 100m 가지만, 나중에 막히겠지? 대신 저기 왼쪽으로 가면 비록 지금은 길이 좁아 보이지만, 그걸로 인해 나중에 모든 길이 한 번에 열려 전체 미로를 해결할 수 있겠네!"라고 생각합니다.

즉, 지금 당장 쉬운 것 (확실한 것) 을 채우는 것보다, 그 선택이 미래의 나머지 퍼즐 조각들을 얼마나 더 쉽게 만들어줄지 (정보 이득) 를 계산해서 결정합니다.

4. 왜 이것이 특별한가요?

기존의 AI 들은 "다음 단어"만 예측할 수 있어서 미래를 계산하는 데 너무 많은 시간이 걸려서 (계산 비용이 너무 비싸서) 미래를 생각하지 못했습니다.

하지만 이 새로운 MDM은 **양방향 (Bidirectional)**으로 볼 수 있는 능력을 가지고 있습니다.

마치 모든 퍼즐 조각을 한 번에 훑어볼 수 있는 마법 안경을 쓴 것처럼, "이 빈칸을 채우면 나머지 빈칸들이 얼마나 쉬워질까?"를 한 번의 계산으로 바로 알 수 있습니다.
그래서 "정보 이득"을 계산하는 것이 매우 빠르고 효율적입니다.

5. 실제 성과: 얼마나 좋아졌나요?

이 새로운 방법을 적용하니 AI 의 실력이 크게 향상되었습니다.

수학/논리 문제: "100 을 2 로 나누면?" 같은 문제에서, 답을 먼저 채우지 않고 계산 과정을 먼저 해결해서 정답률이 3.6%나 올라갔습니다.
창의적 글쓰기: AI 가 지어낸 이야기가 더 논리적이고 자연스러워져서, 다른 AI 들보다 63.1% 더 높은 점수를 받았습니다. (사람들이 "이게 더 좋아"라고 투표한 비율)
그림 그리기: 그림 속 객체의 위치나 색상이 더 정확하게 표현되었습니다.

6. 요약: 한 줄로 정리하면?

"지금 당장 쉬운 것부터 채우는 게 아니라, 전체 퍼즐을 해결하는 데 가장 큰 도움을 주는 순서로 채우면, AI 가 훨씬 더 똑똑하고 실수 없는 글을 쓰고 그림을 그릴 수 있다."

이 연구는 AI 가 단순히 "확실한 것"을 쫓는 것을 멈추고, **"미래를 내다보며 전략적으로 행동하는 법"**을 배웠다는 점에서 매우 의미 있는 발전입니다.

1. 문제 정의 (Problem)

마스킹 확산 모델 (Masked Diffusion Models, MDMs) 은autoregressive 모델 (ARM) 의 왼쪽에서 오른쪽으로의 엄격한 생성 순서 제약을 벗어나, 양방향 어텐션을 활용하여 더 유연한 디코딩 경로를 제공합니다. 이는 코드 채우기, 생물학적 시퀀스 설계, 장기 계획 등 양방향 맥락이 필요한 작업에서 뛰어난 잠재력을 보입니다.

그러나 MDM 의 잠재력은 학습 - 추론 불일치 (Training-Inference Mismatch) 로 인해 충분히 발휘되지 못하고 있습니다.

현재의 한계: 기존 MDM 샘플러들은 대부분 국소적 확실성 (Local Certainty) 기반의 탐욕적 (Greedy) 휴리스틱 (예: 가장 높은 확률, 가장 낮은 엔트로피) 을 사용하여 다음에 해독할 토큰 위치를 선택합니다.
근본적 결함: 이러한 탐욕적 접근법은 현재의 결정이 미래 단계에 미치는 영향 (Downstream Impact) 을 무시합니다. 즉, 즉시 확실해 보이는 토큰을 먼저 선택하는 것이 전체 생성 과정의 누적 불확실성 (Cumulative Uncertainty) 을 최소화하는 최적의 경로가 아닐 수 있습니다.
결과: 이는 의미적으로 최적이지 않은 토큰을 먼저 고정시켜 오류가 전파되고, 최종 생성 품질이 저하되는 결과를 초래합니다.

2. 방법론 (Methodology)

저자들은 MDM 의 양방향 어텐션 특성을 활용하여 정보 획득 (Information Gain, Info-Gain) 을 고려한 새로운 샘플링 프레임워크인 Info-Gain Sampler 를 제안합니다.

핵심 아이디어

정보 획득 (Information Gain): 현재 디코딩 결정이 나머지 마스킹된 토큰들의 불확실성을 얼마나 줄여주는지 평가합니다.
MDM 의 이점: autoregressive 모델은 다음 토큰 예측만 가능해 미래 시뮬레이션 (Look-ahead) 이 계산 비용이 많이 들지만, MDM 은 한 번의 포워드 패스로 모든 마스킹 위치의 확률 분포를 동시에 평가할 수 있어 정보 획득을 효율적으로 계산할 수 있습니다.

Info-Gain Sampler 의 목적 함수

디코딩 동작 $a_t$ 의 총 영향력은 두 가지 요소로 분해되어 균형을 맞춥니다:

즉시 비용 (Immediate Cost): 현재 단계에서 디코딩하는 토큰들의 불확실성 (마진 엔트로피 합).
정보 획득 (Information Gain): 해당 동작을 수행한 후 남은 마스킹 위치들의 상태 불확실성 (State Uncertainty) 감소량.

목적 함수 ( $J_{IG}$ ):
$J_{IG}(a_t | z_t) = \underbrace{IG(a_t; z_t)}_{\text{정보 획득}} - \underbrace{C(a_t | z_t)}_{\text{즉시 비용}}$
이 목적 함수를 최대화하는 동작을 선택함으로써, 국소적 확실성뿐만 아니라 전체 시퀀스의 불확실성을 줄이는 전역적 최적 경로를 탐색합니다.

구현 세부 사항

액션 샘플링: 다양한 후보 동작 집합을 생성합니다.
병렬 평가: 모든 후보에 대한 정보 획득을 단일 배치 포워드 패스로 동시에 계산하여 효율성을 확보합니다.
고신뢰도 우회 (High-Confidence Bypass): 특정 임계값 이상의 확률을 가진 토큰은 즉시 고정하여 추론 지연을 줄입니다.

3. 주요 기여 (Key Contributions)

기존 탐욕적 샘플러의 한계 규명: 실패 사례 분석 (단방향 곱셈, 이진 판단 등) 을 통해 국소적 확실성만 고려하는 기존 방식이 누적 불확실성 최소화에 실패함을 증명했습니다.
Info-Gain Sampler 제안: 즉시 비용과 미래 정보 획득을 균형 있게 고려하는 새로운 디코딩 프레임워크와 이를 효율적으로 구현하는 방법을 제시했습니다.
광범위한 실험적 검증: 다양한 아키텍처 (전체 어텐션, 준-autoregressive) 와 작업 (추론, 코딩, 창의적 글쓰기, 이미지 생성) 에서 기존 SOTA 샘플러들을 압도하는 성능을 입증했습니다.

4. 실험 결과 (Results)

다양한 벤치마크에서 Info-Gain Sampler 는 기존 방법론 (Entropy, Confidence, Margin, LookUM 등) 보다 일관되게 우수한 성능을 보였습니다.

추론 및 코딩 (Reasoning & Coding):
- 전체 어텐션 MDM (Dream-7B) 에서 평균 정확도가 3.6% 향상되었습니다.
- 누적 엔트로피 (Cumulative Entropy) 가 기존 최선 베이스라인 (PC-Sampler) 대비 78.4 에서 48.6 으로 크게 감소하여, 더 안정적인 디코딩 경로를 찾았음을 보여줍니다.
- LookUM(동시 연구) 보다 코드 및 계획 작업에서 더 우수한 성능을 기록했습니다.
창의적 글쓰기 (Creative Writing):
- AlpacaEval 에서 63.1% 의 승률 (Win-rate) 을 기록하며 베이스라인들을 압도했습니다. 특히 고온 (High Temperature) 설정에서도 일관된 성능을 유지하며 창의성과 일관성의 균형을 잘 잡았습니다.
이미지 생성 (Text-to-Image):
- GenEval 및 ImageNet-512 에서 정렬 (Alignment) 과 충실도 (Fidelity) 가 모두 향상되었습니다. FID 점수가 43.3 에서 38.1 로 개선되었습니다.
효율성:
- 추가적인 계산 비용은 최소한으로 증가 (+24% 추론 시간, +20% GPU 메모리) 시켰으며, KV 캐싱 최적화 등을 통해 실용성을 확보했습니다.

5. 의의 및 결론 (Significance)

이 논문은 MDM 의 핵심 강점인 양방향성 (Bidirectionality) 을 샘플링 전략에 성공적으로 통합한 첫 번째 체계적인 접근법 중 하나입니다.

국소적 휴리스틱에서 전역적 계획으로의 전환: 기존 MDM 샘플러가 가진 '근시안적 (Myopic)' 결함을 해결하고, 정보 이론적 관점에서 전역적 최적화를 가능하게 했습니다.
모델 아키텍처 무관성: 전체 어텐션 모델과 준-autoregressive 모델 모두에 적용 가능하며, 다양한 도메인 (텍스트, 이미지) 에서 일반화되는 능력을 입증했습니다.
미래 방향: 추론 속도를 더 높이기 위한 효율적인 Look-ahead 메커니즘 및 하드웨어 최적화 등 향후 연구의 방향성을 제시했습니다.

결론적으로, Info-Gain Sampler 는 MDM 의 생성 품질 한계를 높이고, 비-autoregressive 생성 모델의 실용성을 크게 향상시키는 중요한 기여를 한 연구입니다.