고해상도 이미지 합성을 위한 마스크된 이산 확산의 발전

이 논문은 마스크된 이산 확산(MDM)의 두 가지 근본적 한계를 해결해 고해상도 텍스트-투-이미지 합성의 품질과 학습 효율을 동시에 높였다. 첫째, 이미 언마스크된 이산 토큰을 수정할 수 없다는 문제를 토큰 편집 메커니즘으로 해결하여 추론 중 반복적 교정이 가능해졌다. 둘째, 대규모 코드북에서 발생하는 토큰 희소성을 Grouped Cross-Entropy(GCE)로 완화해 작은 학습 신호로 인한 최적화 실패를 줄였다.

이 논문은 마스크된 이산 확산 모델이 가지는 두 가지 구조적 한계에서 출발한다. 첫째, 기존 MDM은 추론 시 한 번 언마스크된 토큰을 더 이상 변경하지 못하므로 초기 오류가 누적되어 최종 화질을 저하시킨다. 둘째, 코드북을 크게 하면 재구성 오차는 줄어들지만 각 토큰의 출현 빈도가 급감하여 one-hot 기반 손실로는 의미 있는 학습 신호를 주기 어렵다. 논문은 첫 번째 문제를 해결하기 위해 훈련 과정에서 일부 깨끗한 토큰을 손상시키는 방식으로 모델이 비마스크 위치에서도 '교정 분포'를 예측하도록 학습시켰다. 이로 인해 추론 단계에서 모델은 각 스텝마다 자신이 예측한 기존 토큰을 재평가하고, 새 예측의 신뢰도가 임계값 τ를 넘으면 토큰을 교체해 반복적 정제를 수행한다. 결과적으로 샘플링 과정이 연쇄적인 독립 샘플링(product of marginals) 가정에서 오는 오류 누적을 완화하고 LDM의 반복 정제와 유사한 동작을 확보한다. 코드북 희소성 문제는 K-means 기반의 군집화를 통해 완화된다. 모델은 토큰별 소프트맥스 확률을 계산하고 같은 군집에 속한 토큰들의 확률을 합산해 군집 확률을 얻은 뒤, 군집 단위의 one-hot 교차엔트로피를 추가로 최적화한다. 이 구성은 정답 토큰뿐 아니라 임베딩 상의 근접 토큰에도 양의 그래디언트를 제공하여 희소한 토큰들에 대한 학습 신호를 증폭시키고, 다중 군집 해상도를 통해 미세한 유사성과 거친 유사성을 동시에 반영한다.

전체 아키텍처는 텍스트와 이미지 토큰을 단일 decoder-only Transformer로 처리하도록 구성되며, 최종 언어 모델 헤드를 이미지 토큰 예측을 위한 새 MLP 헤드로 교체해 이미지 생성에 맞도록 조정했다. 모델 초기화는 사전학습된 diffusion language model에서 가져왔으며, 이로 인해 긴 프롬프트 처리와 sequence packing 최적화를 자연스럽게 활용할 수 있었다. 훈련은 256×256에서 시작해 점진적으로 1024×1024로 업스케일하는 프로그레시브 스케줄을 사용했고 전체 모델 파라미터는 약 8B이다. 토큰 편집 메커니즘의 훈련 절차는 기존의 마스크 생성 q(y^t|y^0)에 추가적으로 clean 토큰의 'corruption'을 포함한다. 손상 전략으로는 같은 이미지 내에서 임의 토큰으로 대체하거나 토크나이저 임베딩 공간의 근접 이웃을 샘플링해 교체하는 두 가지 방식을 사용했다. 손상된 입력을 주면 모델은 모든 위치에서 원본 토큰 y^0를 예측하도록 ELBO 기반 목표를 -1/t로 스케일해 계산하고, 이로써 비마스크 위치에 대한 교정 분포 학습이 가능해졌다. GCE는 오프라인 K-means로 대규모 코드북을 여러 계층의 군집으로 미리 묶고, 모델이 예측한 토큰 확률을 군집별로 합한 뒤 군집 레벨에서의 cross-entropy를 추가하는 방식으로 구현되었다. 수학적으로 전체 손실 항은 토큰 수준의 log p(y_i^0|y^t)와 각 군집 레벨의 log P(C^j(y_i^0)|y^t) 합으로 정의되며, 그라디언트는 목표 군집에 속한 모든 토큰에 양의 신호를 전달한다. 실제 계산 비용을 줄이기 위해 GCE의 순전파 및 역전파를 결합한 custom fused operator를 구현해 메모리와 대기시간을 낮췄다.

주요 벤치마크에서 NLD-Image는 GenEval에서 0.90을 기록했고, DPG에서 85.2, MJHQ-30k 기준 HPSv3는 9.61을 달성했다. 추가로 finetune된 설정에서는 DPG 86.9와 MJHQ HPSv3 10.76을 기록해 높은 이미지 충실도를 보였다. 이 결과는 같은 마스크 기반 계열의 Meissonic 등 기존 MDM 계열 모델들보다 전반적으로 우수한 성능 지표를 나타냈다. GCE의 효과는 통제된 ImageNet256 실험에서 입증되었다. 동일 아키텍처와 토크나이저 조건에서 CE와 SNCE 대비 GCE가 일관되게 더 낮은 FID를 보였고, 예컨대 동일 파라미터 설정에서 GCE는 FID 3.40을 달성해 SNCE의 3.62보다 개선을 보였다. 이 결과는 대규모 코드북 환경에서 군집 기반 보조 손실이 희소성으로 인한 학습 난관을 완화함을 시사한다. 연산 효율성 실험에서 GCE의 최적화된 연산자는 지연시간을 44 ms(naive)에서 약 20 ms로 줄였고 최대 VRAM 사용량을 25 GB에서 16 GB로 감소시켰다. 생성 측면에서는 Emu3.5(AR) 대비 약 42.4× 빠른 추론 속도를 보이면서도 GenEval 성능을 상회하는 속도-품질 균형을 달성했다. 또한 토큰 편집은 샘플링 단계 수(NFE)를 절반 가까이 줄이면서 동등한 품질을 유지할 수 있게 해 현실적 추론 비용 절감을 가능하게 했다.

아키텍처는 decoder-only Transformer를 텍스트 인코더와 이미지 디코더를 통합하는 단일 모델로 사용하며, 최종적으로 8B 파라미터로 end-to-end 최적화되었다. 이 설계는 고정된 텍스트 인코더의 context 제한을 피하고 sequence packing과 같은 LLM 최적화를 자연스럽게 활용할 수 있게 한다. 입력 토큰 구성은 텍스트 프롬프트 토큰과 마스크·손상·클린 이미지 토큰이 섞인 형태이며, 모델은 모든 위치에서 y^0 예측을 수행하도록 학습된다. 토큰 편집은 훈련 시 forward 과정 q'(y^t|y^0)에 손상 연산을 포함시켜 비마스크 위치에서도 교정 분포를 학습하도록 만들고, 추론 시 confidence threshold τ를 기준으로 이미 언마스크된 토큰을 교체한다. 손상 전략은 같은 이미지 내 임의 토큰 샘플링과 임베딩 근접 이웃 샘플링 두 가지를 조합해 실제 오류 분포와 훈련 분포 간의 불일치를 줄였다. 수정된 ELBO는 모든 위치의 로그 확률 합을 -1/t 스케일로 평균화한 형태로 계산된다. GCE의 수학적 구성은 토큰 수준의 softmax p_i와 군집 확률 P(C)=Σ_{j∈C} p_j를 결합한 계층적 로그우도 합으로 정의된다. 그라디언트는 군집 내부 토큰에 양의 신호를 전달하며 수식적으로 ∂/∂h_i log P(C)=p_i(I{i∈C}/P(C)-1)로 주어져 동일 군집의 다른 토큰들에도 학습 신호가 전달된다. 이로 인해 군집 내 토큰들은 토큰 빈도가 낮아도 의미적 유사성에 기반해 강화되는 경향을 보인다. 대규모 코드북과 GCE를 현실적으로 운용하기 위해 custom fused operator를 구현했다. 벤치마크에서 naive GCE eager 구현은 44.14 ms의 지연과 최대 25.2 GB VRAM을 요구했고, 최적화된 전후방 연산 구현은 지연을 20.04 ms로 낮추고 최대 VRAM을 16.1 GB로 줄였다. 전체 학습 스텝은 H100 64GPU 환경에서 300K 스텝으로 운영되었고, 한 스텝당 전체 학습 시간은 약 3.2 s 수준으로 보고되었다.

논문은 대규모 모델과 대규모 컴퓨팅 자원에 기반한 결과를 보고하며, 본 연구에서 사용한 설정은 H100 64GPU에서 300K 스텝 훈련처럼 높은 계산 비용을 전제로 한다. 또 대규모 코드북(131K)과 다중 군집을 사용하므로 토크나이저와 데이터 준비, 클러스터링 전처리에 추가 비용이 발생한다. FID는 고해상도 품질을 완전하게 반영하지 못하는 한계가 있음을 논문에서 명시하여 평가 지표 선택에 주의가 필요하다.