기술 요약: 이산 확산 모델에서의 이해와 생성 간의 균형

1. 문제 정의

이산 생성 모델링 영역에서 두 가지 지배적인 패러다임이 서로 다른 강점과 약점을 가진 채 등장했습니다:

• 마스크 확산 언어 모델 (Masked Diffusion Language Models, MDLMs): 이 모델들은 의미론적 이해, 가능도 모델링(likelihood modeling), 제로샷 일반화에 탁월합니다. 그러나 제한된 수의 추론 단계 내에서 일관성 있고 문맥적으로 정합한 출력을 생성하는 데 어려움을 겪는 경우가 많습니다.
• 균등 노이즈 확산 언어 모델 (Uniform-noise Diffusion Language Models, UDLMs): 이 모델들은 적은 수의 추론 단계에서도 높은 품질의 생성을 달alog하지만, 추론 단계 수가 증가함에 따라 제로샷 가능도 및 의미론적 이해 측면에서 MDLM에 비해 뒤처지는 경ster가 있습니다.

현재까지 두 패러다임 중 어느 것도 이해와 생성이라는 두 차원에서 균형 잡힌 성능을 달성하지 못했습니다. 이러한 불균형은 특히 추론 단계가 총 패치 수보다 훨씬 적은 이미지 합성 같은 응용 분야에서 매우 치명적입니다. 즉, UDLM에는 유리하지만 MDLM의 우수한 의미론적 접지(semantic grounding)를 활용하지 못하게 됩니다. 핵심 과제는 UDLM의 견고한 소수 단계 생성 능력은 유지하면서 MDLM의 우수한 의미론적 이해 능력을 보존할 수 있는 모델을 구현하기 위해 이 두 패러다임을 연결하는 것입니다.

2. 방법론: XDLM

저자들은 **정상 상태 노이즈 커널(stationary noise kernel)**을 통해 균등 노이즈와 마스크 노이즈 분포를 연결하는 통합 프레임워크인 **miXed Diffusion Language Modeling (XDLM)**을 제안합니다.

핵심 메커니즘

매 타임스텝마다 노이즈 특성이 변하는 계산 비용이 많이 드는 시변(time-inhomogeneous) 프로세스에 의존하는 GIDD와 같은 기존 방식과 달리, XDLM은 정상 상태 노이즈 커널을 강제합니다. 이 커널은 증분 노이즈(incremental noise)를 한계 노이즈(marginal noise)와 구조적으로 일치시켜, 훈련 목적 함수를 정적 상수와 동적 스케줄로 분해할 수 있게 합니다.

순방향 과정은 항등 행렬 $I$와 정상 상태 노이즈 커널 $K$의 볼록 조합으로 정의되는 전이 행렬 $Q_{t|s}$에 의해 제어됩니다:
$$Q_{t|s} = \alpha_{t|s}I + \beta_{t|s}K$$

노이즈 커널 $K$는 균등 노이즈와 흡수(masking) 노이즈의 혼합물로 구성됩니다:
$$K = \frac{k}{N}J + \mu M$$
여기서:

• $k$는 균등 노이즈의 혼합 비율 가중치입니다.
• $\mu$는 마스크 상태에 대한 흡수 가중치입니다 (여기서 $k + \mu = 1$).
• $J$는 모든 원소가 1인 행렬(균등 전이)입니다.
• $M$은 마스크 상태를 위한 흡수 행렬입니다.

이론적 통합

XDLM은 다음과 같은 원리적인 이론적 통합을 제공합니다:

• $k=0$으로 설정하면 (순수 마스킹) MDLM 목적 함수가 됩니다.
• $k=1$로 설정하면 (순수 균등) UDLM 목적 함수가 됩니다.
  따라서 MDLM과 UDLM은 XDLM 공식의 극한 사례(limiting cases)임이 증명되었습니다.

효율적인 구현

주요 기술적 기여는 사후 확률(posterior probabilities)과 KL 발산(KL divergence)에 대한 **스칼라 공식(scalar formulation)**의 유도입니다. 대규모 어휘 집합에 대해 행렬 연산을 통해 이러한 항들을 직접 평가하는 것은 계산 비용이 너무 많이 듭니다. 저자들은 헬퍼 함수(노이즈율 $r(e)$ 및 순방향 확산 맵 $f_t(x, e)$)를 정의함으로써, 훈련 목적 함수를 다루기 쉬운 스칼라 형태로 재구성했습니다. 이러한 대수적 단순화는 메모리 병목 현상을 완화하여, 과도한 계산 비용 없이 XDLM이 대규모 어휘 크기로 확장될 수 있도록 합니다.

3. 주요 기여

1. 이론적 통합: XDLM은 균등 노이즈와 마스크 노이즈 분포 사이의 엄격한 가교를 구축하며, 기존 패러다임(MDLM 및 UDLM)이 이 통합된 프레임워크의 특수 사례임을 입증합니다.
2. 메모리 효율적 확장성: 사후 확률 계산의 대수적 단순화 제안은 효율적인 훈련과 샘플링을 가능하게 하여, 이산 확산 모델의 대규모 어휘 집합과 관련된 계산적 병목 현상을 극복합니다.
3. 균형 잡힌 성능: 본 프레임워크는 이해 능력(제로샷 퍼플렉시티)과 생성 품질(소수 단계 생성) 사이의 파레토 프런티어를 성공적으로 진전시킵니다.

4. 실험 결과

실험은 텍스트 및 이미지 모달리티, 그리고 대규모 언어 모델에 걸쳐 수행되었습니다.

텍스트 생성 및 이해

• 제로샷 가능도(Zero-Shot Likelihood): OpenWebText (OWT) 검증 세트에서 XDLM은 MDLM(23.321)과 유사하고 UDLM(25.937)보다 훨씬 나은 24.097의 퍼플렉시티(PPL)를 달성했습니다. 7개의 제로샷 벤치마크(AG News, LAMBADA 등) 전반에서 XDLM은 평균 54.110의 PPL을 유지하며 MDLM(53.650)과 일치하는 성능을 보였고 UDLM(59.574)을 능가했습니다.
• 생성 품질: 소수 단계 환경(8–32 단계)에서 XDLM은 균등 노이즈 역학의 이점을 얻어 UDLM과 대등한 생성 품질을 달성하고 순수 마스크 모델을 능가했습니다. 다단계 환경에서는 마스킹의 구조적 엄밀함을 유지하여 UDLM보다 뛰어난 성능을 보였습니다.

이미지 생성

• ImageNet-1K: XDLM은 소수 단계 생성에서 우수한 성능을 입증했습니다. 4단계에서 XDLM은 FID를 80.8(MDLM)에서 54.1로 낮추었습니다. 16단계에서는 CFG(Classifier-Free Guidance) 없이도 베이스라인 중 가장 낮은 FID(25.77)를 달성했습니다.
• CFG 사용 시: CFG(scale=2.0)를 사용할 때, XDLM은 4단계에서 13.55, 8단계에서 8.96의 최저 FID를 기록하며, 소수 단계 환경에서 UDLM을 능가하는 동시에 높은 단계에서는 MDLM의 성능에 근접했습니다.

대규모 언어 모델(LLM) 스케일링

• 지속적 사전 학습(Continual Pretraining): 저자들은 8B 파라미터 규모의 MDLM(LLaDA)을 600 단계의 지속적 사전 학습을 통해 XDLM 공식(LLaDA-XDLM)으로 적응시켰습니다.
• 코드 생성: MBPP 벤치마크에서 LLaDA-XDLM은 단 32 단계의 샘플링만으로 15.0의 점수를 획득하여, 베이스라인인 LLaDA의 성능(6.8)을 효과적으로 두 배로 높였습니다. 이러한 개선은 컴파일 불가능한 코드 에러의 감소에 기인하며, 이는 구조적 일관성이 향상되었음을 나타냅니다.
• 훈련 역학: 분석 결과 "성능 교차(performance crossover)" 현상이 발견되었습니다. 마스크 기반 모델들은 빠르게 수렴하지만 조기에 정체되는 반면, XDLM은 훈련 전반에 걸쳐 꾸준한 개선을 유지하며 우수한 장기 스케일링 잠재력을 보여주었습니다.

5. 의의 및 주장

본 논문은 XDLM이 두 가지(UDLM 및 MDLM)의 기계적 특성을 하나의 통합된 프레임워크로 성공적으로 합성했다고 주장합니다. 최적의 혼합 비율($k=0.1$)을 "스위트 스팟(sweet spot)"으로 식별함으로써, 모델은 베이스라인의 파레토 프런티어를 돌파합니다. 이는 순수 마스크 모델에서 나타나는 심각한 생성 성능 저하 없이 강력한 이해 능력을 확보하는 동시에, 의미론적 접지를 희생하지 않으면서도 강력한 생성 품질을 유지합니다.

저자들은 이러한 이득이 단순히 추가적인 훈련 단계나 샘플링 아티팩트 때문이 아니라, XDLM 공식 자체의 효능에서 비롯된 것임을 강조합니다. 이 연구는 이산 확산 모델이 의미론적 이해와 생성 효율성 사이의 트레이드오프를 균형 있게 맞추면서도, 대규모 어휘 집합과 대규모 파라미터 수(8B)로 확장될 수 있음을 보여줍니다.

인정된 한계점:
논문은 XDLM이 아직 대규모 규모에서 처음부터(from scratch) 훈련되지 않았다는 점, "성능 교차" 현상에 대한 추가 조사가 필요하다는 점, 그리고 텍스트 및 이미지 모달리티를 위한 도메인 특화 샘플링 전략이 아직 완전히 최적화되지 않았다는 점을 언급합니다. 또한, 단일 통합 모델 내에서 텍스트와 시각적 모달리티 간의 상호작용에 대한 연구도 아직 이루어지지 않았습니다.