KLASS: KL-Guided Fast Inference in Masked Diffusion Models

이 논문은 추가 학습 없이 토큰 수준의 KL 발산을 활용해 안정적이고 신뢰도 높은 예측을 식별하여 여러 토큰을 동시에 언마스크하는 'KL-Adaptive Stability Sampling(KLASS)'을 제안함으로써, 마스킹 확산 모델의 추론 속도를 획기적으로 개선하면서도 생성 품질을 유지하거나 향상시킵니다.

핵심

이 논문은 **'KLASS'**라는 새로운 기술을 소개합니다. 이 기술은 인공지능이 글을 쓰거나 그림을 그릴 때, 훨씬 더 빠르고 정확하게 작업할 수 있도록 도와주는 '스마트한 마법 지팡이'와 같습니다.

이해하기 쉽게 요리사와 레시피에 비유해서 설명해 드릴게요.

1. 문제: 느리고 지루한 '한 입씩' 요리법

기존의 AI(마스크드 확산 모델) 는 글을 쓸 때 마치 완전히 빈 그릇에 한 숟가락씩 밥을 떠서 채워 넣는 방식으로 작동했습니다.

• 방식: "이 자리에 뭐가 들어갈까?"라고 고민하고, 가장 확신 있는 단어를 하나만 고릅니다. 그리고 다시 다음 자리로 넘어갑니다.
• 단점: 이 과정이 너무 느립니다. 256 단어를 쓰려면 256 번이나 고민해야 하죠. 게다가 AI 가 "아, 이거 맞을 것 같아!"라고 생각했는데 나중에 보니 틀린 경우가 많아서, 처음부터 다시 고쳐야 하는 비효율도 있습니다.

2. 해결책: KLASS (스마트한 요리사)

이 논문이 제안한 KLASS는 AI 가 "이건 확실해!"라고 판단할 때, 한 번에 여러 단어를 동시에 채워 넣는 똑똑한 방법입니다.

여기서 핵심은 두 가지 신호를 보는 것입니다.

🍳 비유 1: "요리사 자신감" (Confidence Score)

• 상황: 요리사가 "이 요리에 소금 한 숟가락이 딱 맞을 거야!"라고 아주 확신하는 경우입니다.
• KLASS: "좋아, 이 부분은 확실히 소금으로 채워버리자!"라고 바로 결정합니다.

🍳 비유 2: "요리 상태의 안정성" (KL Divergence)

• 상황: 요리사가 소금을 넣으려다가, "아, 잠깐! 아까는 '설탕'이 맞다고 생각했는데, 지금 생각해보니 '소금'이 더 맞는 것 같네?"라고 자꾸 생각이 바뀌는 경우입니다.
• KLASS: "아직 생각이 흔들리네? 그럼 이 부분은 아직 확정하지 말고 잠시 기다려보자."라고 판단합니다.
• 핵심: AI 의 생각이 여러 번의 시도 끝에 변하지 않고 '안정적'일 때에만 그 부분을 확정합니다.

3. KLASS 가 가져오는 변화

이 두 가지 신호 (자신감 + 안정성) 를 합치면 어떤 일이 일어날까요?

1. 속도 폭발 (2.78 배 빨라짐):

   • 기존 방식은 한 번에 하나씩만 채웠다면, KLASS 는 "이건 확실히 맞고, 생각도 안 흔들리네?"라고 판단되는 부분들을 한 번에 여러 개 동시에 채웁니다.
   • 마치 256 번의 행동을 100 번도 안 되는 행동으로 줄여버린 것과 같습니다.

2. 정확도 향상:

   • 단순히 "자신감"만 보고 급하게 채우면 실수할 수 있습니다. 하지만 "생각이 흔들리지 않는지 (안정성)"까지 확인하니까, 실수가 줄어들고 더 논리적인 답변을 내놓습니다.
   • 수학 문제를 풀거나 코드를 작성할 때, 엉뚱한 답을 고르는 실수를 막아줍니다.

4. 어디에 쓸 수 있나요?

이 기술은 글쓰기뿐만 아니라 이미지 생성이나 분자 구조 설계 (약물 개발 등) 같은 다양한 분야에서 똑같이 작동합니다.

• 글쓰기: 더 빠르고 매끄러운 글을 씁니다.
• 그림 그리기: 더 선명하고 정확한 그림을 그립니다.
• 약물 개발: 더 효율적으로 새로운 분자를 찾아냅니다.

🌟 한 줄 요약

KLASS는 AI 가 "이건 확실해!"라고 생각할 때 한 번에 여러 단어를 채워 넣는 기술로, 생각이 흔들리지 않는지 확인함으로써 속도는 3 배 가까이 늘리고, 정확도는 더 높여주는 획기적인 방법입니다.

기존의 AI 가 "서두르다가 실수하는 학생"이었다면, KLASS 는 "자신감 있으면 빠르게, 의심스러우면 천천히 확인하는 스마트한优等생"으로 바꿔준 셈입니다.

기술 요약

1. 문제 정의 (Problem)

**마스크드 확산 모델 (Masked Diffusion Models, MDM)**은 언어 생성, 이미지 생성, 분자 생성 등 다양한 분야에서 경쟁력 있는 성능을 보여주지만, 추론 (Inference) 속도가 매우 느리다는 한계가 있습니다.

• 현재의 병목 현상: 기존 MDM 은 부분적으로 마스킹된 시퀀스를 반복적으로 정제 (Refinement) 하는 과정을 거치며, 각 단계에서 소수의 토큰만 언마스크 (Unmasking) 하는 방식 (예: Top-k, 확률적 샘플링) 을 사용합니다.
• 비효율성: 이러한 고정된 또는 단순한 샘플링 전략은 생성 과정을 비효율적으로 만들고, 국소 최적해 (Local Suboptimalities) 에 빠질 위험이 있어 실제 적용에 제약을 줍니다.
• 기존 해결책의 한계: 속도를 높이기 위해 추가적인 '플래너 (Planner)' 모델이나 보조 분포를 도입하는 방식은 계산 오버헤드를 증가시키고, 기본 모델의 학습된 분포와 정렬 (Alignment) 하는 데 어려움을 겪습니다.

2. 제안 방법: KLASS (Methodology)

저자들은 **"KL-Adaptive Stability Sampling (KLASS)"**이라는 새로운 샘플링 알고리즘을 제안했습니다. 이는 추가적인 모델 학습 없이 기존 확산 모델의 내부 역학을 활용하여 **토큰 수준의 KL 발산 (Kullback-Leibler Divergence)**과 **신뢰도 (Confidence)**를 기반으로 안정된 토큰을 식별하고 병렬로 언마스크하는 방식입니다.

핵심 메커니즘

1. 신뢰도 점수 (Confidence Score): 모델이 특정 토큰을 예측할 때의 확률 분포에서 최대 확률 값을 의미합니다. 값이 높을수록 모델이 해당 토큰에 대해 확신이 있음을 나타냅니다.
2. KL 점수 (KL Score): 인접한 시간 단계 (timestep) 간 토큰의 예측 확률 분포 변화량을 KL 발산으로 측정합니다.
   • 안정성 (Stability): 올바른 토큰은 생성 과정에서 예측 분포가 일관되게 유지되는 경향이 있어 KL 발산이 낮습니다.
   • 불안정성 (Instability): 잘못된 토큰은 문맥이 해결됨에 따라 예측이 급변하므로 KL 발산이 높게 나타납니다.
3. 안정 토큰 선택 (Stable-token Selection):
   • 특정 토큰이 높은 신뢰도를 가지면서 KL 발산이 임계값 이하로 낮아져 예측이 안정화되었을 때, 해당 토큰을 즉시 언마스크합니다.
   • 이 조건을 만족하는 토큰들은 **병렬 (Parallel)**로 언마스크되어 전체 샘플링 단계 수를 획기적으로 줄입니다.
   • 조건을 만족하는 토큰이 없을 경우, 백업 전략으로 Top-k 토큰을 언마스크합니다.

알고리즘 특징

• 학습 불필요 (Training-free): 기존 모델의 가중치를 변경하거나 추가 학습이 필요 없습니다.
• 동적 적응 (Adaptive): 고정된 스케줄이 아닌, 생성 과정 중 모델의 예측 신뢰도에 따라 언마스크할 토큰 수를 동적으로 조절합니다.

3. 주요 기여 (Key Contributions)

1. 새로운 샘플링 전략 제안: 외부 플래너 없이 모델 내부의 KL 발산과 신뢰도를 결합한 KLASS를 제안하여 마스킹 확산 모델의 추론을 가속화했습니다.
2. 성능 및 속도 동시 달성:
   • 기존 Greedy 또는 Top-k 디코딩 대비 샘플링 단계 수를 50% 이상 감소시켰습니다.
   • 벽시계 시간 (Wall-clock time) 기준 최대 2.78 배의 속도 향상을 달성했습니다.
   • 속도를 높임과 동시에 정확도도 개선하여, 기존 확산 모델 샘플러 중 **최고 수준의 성능 (State-of-the-Art)**을 기록했습니다.
3. 광범위한 검증: 수학/코드 추론 (GSM8K, MATH, HumanEval, MBPP), 텍스트 생성, 이미지 생성 (MMaDA), 분자 생성 (QM9) 등 다양한 도메인과 모델 (LLaDA, Dream 등) 에서 KLASS 의 유효성을 입증했습니다.

4. 실험 결과 (Results)

• 추론 벤치마크 (Reasoning Tasks):
  • LLaDA 8B와 Dream 7B 모델에서 KLASS 를 적용한 결과, Top-1 (Greedy) 대비 정확도가 향상되었습니다 (예: MATH 에서 31.4% → 33.8%, GSM8K 에서 75.1% → 76.5%).
  • 동시에 샘플링 단계는 256 단계에서 약 100~150 단계 수준으로 줄어들어 효율성이 크게 개선되었습니다.
• 텍스트 생성:
  • MDLM 모델에서 KLASS 를 적용한 결과, 기존 이산 확산 샘플러 (SEDD, D3PM 등) 대비 Perplexity 는 낮추고 MAUVE 점수는 높이는 우수한 생성 품질을 보였습니다.
  • 생성된 텍스트의 일관성과 유창성이 개선되었습니다.
• 이미지 및 분자 생성:
  • 이미지 생성 (MMaDA) 에서 FID 점수 개선 및 IS 점수 향상을 보였습니다.
  • 분자 생성 (QM9) 에서 목표 속성 (QED, Ring count) 을 만족하면서 필요한 함수 평가 횟수 (NFEs) 를 줄였습니다.
• 이론적 근거: 잘못된 토큰은 생성 경로 상에서 동적으로 불안정하여 KL 발산이 0 이 될 수 없음을 이론적으로 증명하여, KLASS 가 잘못된 토큰의 조기 언마스크를 방지함을 보였습니다.

5. 의의 및 중요성 (Significance)

• 실용적 확장성: 복잡한 외부 플래너나 추가 학습 없이 기존 확산 모델의 잠재력을 최대한 끌어올리는 간단하고 실용적인 방법을 제시했습니다.
• 확산 모델의 실용화: 기존 확산 모델의 가장 큰 약점인 '느린 추론 속도'를 해결함으로써, 복잡한 추론 작업이나 실시간 생성이 필요한 환경에서 마스킹 확산 모델의 실제 적용 가능성을 크게 높였습니다.
• 일반성: 텍스트, 이미지, 분자 등 다양한 모달리티와 모델 아키텍처에 걸쳐 적용 가능한 범용적인 샘플러로, 대규모 마스킹 확산 모델의 발전에 중요한 기여를 할 것으로 기대됩니다.

결론적으로, KLASS 는 토큰의 예측 안정성을 KL 발산으로 정량화하여 병렬 언마스크를 유도함으로써, 마스킹 확산 모델의 추론 속도와 정확도를 동시에 혁신적으로 개선한 획기적인 연구입니다.