DMax: 확산 언어 모델을 위한 공격적인 병렬 디코딩

기존 확산 언어 모델(dLLM)은 병렬 디코딩 시 발생하는 오류 누적으로 인해 생성 품질이 급격히 저하되는 한계가 있었다. DMax는 임베딩 공간에서의 자기 수정 메커니즘을 도입하여 정확도 손실 없이 추론 속도를 획기적으로 높였으며, 이는 실시간 텍스트 생성 효율성을 크게 개선하는 방향을 제시한다.

기존의 마스크 기반 확산 언어 모델은 한 번 결정된 토큰을 수정할 수 없는 '일방향' 구조를 가진다. 이로 인해 병렬 디코딩 시 초기에 발생한 작은 오류가 이후 단계로 전파되어 전체 문맥이 붕괴되는 오류 누적 문제가 발생한다. 이는 Transformer의 Attention 메커니즘이 잘못된 토큰을 고정된 컨텍스트로 참조하면서 발생하는 현상이다.

DMax는 이를 해결하기 위해 디코딩 과정을 임베딩 공간에서의 점진적인 자기 정제 과정으로 재정의한다. 핵심은 모델이 '확신이 없는 상태'를 마스크 임베딩과의 혼합 비율로 유지하는 것이다. 마치 화가가 밑그림을 그리며 선을 덧칠하듯, 모델은 각 위치의 토큰을 확정 짓기 전에 임베딩 수준에서 여러 번 수정하며 최적의 값을 찾아간다.

결과적으로 모델은 병렬적으로 많은 토큰을 동시에 예측하면서도, 서로의 예측 오류를 실시간으로 감지하고 보정할 수 있게 된다. 이는 기존 모델들이 오류를 피하기 위해 한 번에 소수의 토큰만 처리해야 했던 제약을 극복하고, 하드웨어 가속기를 최대한 활용할 수 있는 고도의 병렬 처리를 가능하게 한다.

DMax의 핵심은 On-Policy Uniform Training (OPUT)이다. 기존의 Uniform Diffusion Training이 무작위 토큰으로 노이즈를 생성하는 것과 달리, OPUT은 모델의 현재 예측 분포 pθ(· | xt(m))에서 샘플링된 토큰으로 노이즈 시퀀스 xt(p)를 구성한다. [모델 예측값 입력 → 샘플링 연산 → 노이즈 시퀀스 출력 → 추론 시 마주할 실제 오류 분포 반영] 과정을 통해 모델은 자신의 실수로부터 복구하는 법을 배운다.

추론 단계에서는 Soft Parallel Decoding (SPD)을 사용한다. 각 위치 j의 입력 hj(t)를 예측 토큰 임베딩 e(yj(t-1))와 마스크 임베딩 emask의 가중 합으로 계산한다. 이때 가중치는 이전 단계의 예측 확률 πj(t-1)를 사용한다. [예측 확률 π와 임베딩 e 입력 → π·e + (1-π)·emask 연산 → 하이브리드 임베딩 출력 → 불확실성이 포함된 부드러운 상태 전달] 방식으로 동작한다.

또한 블록 단위의 세미 자기회귀(Semi-Autoregressive) 방식을 채택한다. 블록 내에서 왼쪽부터 확신도 τdec를 넘는 연속된 토큰들만 확정하고 나머지는 마스크 상태를 유지한다. 블록의 수렴 조건은 모든 위치의 예측이 변하지 않거나 확신도가 임계값 τacc를 넘을 때로 정의하여 효율적인 병렬 처리를 보장한다.

수학적 추론 벤치마크인 GSM8K에서 DMax-Math 모델은 LLaDA-2.0-mini 대비 TPF를 2.04에서 5.48로 168% 향상시켰으며, 정확도는 92.1%로 원본(92.6%)과 대등한 수준을 유지했다. TPS(초당 토큰 수)는 512에서 1258로 크게 증가했다.

코드 생성 벤치마크인 MBPP-Instruct에서는 TPF가 2.71에서 5.86으로 개선되었으며, 정확도는 80.6%에서 79.2%로 미세하게 조정되었으나 TPS는 662에서 1264로 약 2배 빨라졌다. 특히 높은 병렬성(낮은 τdec) 설정에서도 기존 모델들이 정확도가 급락하는 것과 달리 DMax는 안정적인 성능을 보였다.

효율성 분석 결과, H200 GPU 2장을 사용한 환경에서 배치 사이즈 1일 때 평균 1,338 TPS를 기록했다. 이는 기존 확산 모델의 고질적인 문제였던 느린 추론 속도를 해결하여 실질적인 서비스 적용 가능성을 입증한 결과이다.

DMax 아키텍처는 기존 Masked Diffusion Language Model(MDLM)을 기반으로 하며, 임베딩 공간에서의 자기 정제(Self-Refinement)를 위해 소프트 임베딩 보간 로직을 추가했다. 하이브리드 임베딩 계산 시 노름 붕괴(Norm Collapse)를 방지하기 위해 각 성분의 노름을 가중 합산하여 재정규화(Renormalization)하는 과정을 거친다.

학습 시에는 0.75의 고정된 마스크 비율을 사용하며, 메모리 오버헤드를 피하기 위해 마스크 노이즈 시퀀스와 예측 노이즈 시퀀스에 대한 최적화를 동일 에포크 내 별도 반복(Iteration)으로 수행한다. 이는 단일 반복에서 두 번의 순전파를 수행하는 것보다 효율적이다.

이론적으로 DMax는 MDLM의 안정적인 초기화와 UDLM(Uniform Diffusion)의 자기 수정 능력을 결합한 형태이다. OPUT을 통해 학습된 모델은 마스크 임베딩과 자기 예측 토큰 임베딩 사이의 일관된 매핑을 학습하게 되어, SPD 과정에서의 보간이 수학적으로 유의미한 중간 상태를 형성하게 된다.

DMax는 OPUT으로 학습된 모델에서만 SPD가 정상 작동하며, 일반적인 확산 모델에 SPD를 직접 적용할 경우 성능이 급격히 저하되는 의존성이 존재한다. 또한 블록 단위 디코딩 시 최적의 임계값(τdec, τacc) 설정이 성능과 속도의 트레이드오프에 민감하게 작용할 수 있다.