💻 computer science

Nemotron-Labs-Diffusion-Image: Advancing Masked Discrete Diffusion for High-Resolution Image Synthesis

Nemotron-Labs-Diffusion-Image는 동적 추론 수정을 위한 토큰 편집 메커니즘과 대규모 어휘 설정에서의 신호 희소성을 극복하기 위한 커스텀 퓨즈드 연산자가 적용된 그룹 크로스 엔트로피 목적 함수를 도입함으로써 고해상도 텍스트-이미지 합성 성능을 향상시킨 최첨단 마스크 이산 확산 모델이며, GenEval, DPG 및 HPSv3 벤치마크에서 우수한 성능을 달성했습니다.

원저자: Shufan Li, Greg Heinrich, Hanrong Ye, Yonggan Fu, Aditya Grover, Jan Kautz, Pavlo Molchanov

게시일 2026-06-30

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Shufan Li, Greg Heinrich, Hanrong Ye, Yonggan Fu, Aditya Grover, Jan Kautz, Pavlo Molchanov

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

당신이 로봇에게 "햇살 좋은 날 모자를 쓰고 있는 고양이"와 같은 설명을 주고 걸작을 그리도록 가르치려 한다고 상상해 보세요. 오랫동안, 가장 뛰어난 로봇들은 찰흙 덩어리에서 조각하는 것과 유사한 방식을 사용했습니다. 즉, 흐릿한 덩어리에서 시작하여 실수를 바로잡으며 단계적으로 정교하게 다듬어 나가는 방식이었죠. 이것을 **연속 확산(Continuous Diffusion)**이라고 부릅니다.

최근에는 **이산 확산(Discrete Diffusion)**이라는 새로운 접근법이 인기를 끌었습니다. 이 방식은 찰흙 대신 거대한 레고 블록(토큰) 꾸러미를 사용합니다. 로봇은 숨겨진 블록들로 이루어진 벽에서 시작하여 하나씩 블록을 드러냅니다. 문제는 무엇일까요? 한 번 블록이 드러나면, 그것은 고정되어 버린다는 점입니다. 만약 로봇이 고양이의 귀를 위해 잘못된 블록을 선택한다면, 다시 되돌아가서 바꿀 방법이 없습니다. 이는 마치 조각가가 돌을 한 번 깎아내면 다시는 수정할 수 없는 것과 같습니다.

이 논문은 이 레고 방식의 주요 문제점 두 가지를 해결하는 새로운 "조각가", Nemotron-Labs-Diffusion-Image를 소개합니다.

1. "되돌리기" 버튼 (토큰 편집)

문제점: 표준적인 레고 페인팅 방식에서는, 파란색 블록이 있어야 할 자리에 빨간색 블록을 놓으면 그대로 굳어버립니다. 로봇에게는 "잠깐, 내 생각이 바뀌었어"라고 말하며 나중에 수정할 방법이 없습니다. 이로 인해 작은 오류들이 쌓여 결국 그림이 이상해지게 됩니다.

해결책: 저자들은 로봇에게 "되돌리기" 버튼을 주었습니다. 그들은 로봇이 단순히 숨겨진 블록을 드러내는 법뿐만 아니라, 이미 배치한 블록들을 살펴보고 "사실, 저건 좀 틀렸어. 이걸 바꾸자"라고 말하는 법을 학습시켰습니다.

비유: 조각가가 조각상을 만드는 장면을 상상해 보세요. 예전 방식에서는 돌을 한 번 깎아내면 그것은 영원히 사라졌습니다. 하지만 이 새로운 방식에서는 조각가가 방금 제거한 조각을 부드럽게 깎아내고, 다듬고, 모양이 마음에 들지 않으면 다시 형태를 잡을 수 있습니다. 이를 통해 로봇은 마치 예전의 "찰흙" 모델들처럼, 진행 과정에서 실수를 바로잡으며 반복적으로 이미지를 정교하게 다듬을 수 있습니다.

2. "거대 사전" 문제 (그룹화된 교차 엔트로피)

문제점: 고품질의 상세한 이미지를 만들기 위해, 로봇은 방대한 어휘를 가진 레고 블록(코드북)이 필요합니다. 블록이 많을수록 더 세밀한 그림을 그릴 수 있습니다. 하지만 10만 개 이상의 고유한 블록이 있는 사전의 경우, 로봇은 훈련 중에 특정 블록을 두 번 다시 보지 못하는 경우가 많습니다. 이는 마치 몇 년에 한 번꼴로만 "코끼리"라는 단어를 접할 수 있는 언어를 배우려는 것과 같습니다. 신호가 너무 희소하기 때문에 로봇은 혼란에 빠집니다.

또한, 표준적인 훈련 방식은 모든 잘못된 블록을 똑같이 나쁘다고 취급합니다. 만약 로봇이 '연한 파란색' 블록이 필요한데 '진한 파란색' 블록을 골랐다면, 기존 방식은 마치 '빨간색' 블록을 골랐을 때와 똑같은 크기로 "틀렸어!"라고 소리칩니다. 즉, '연한 파란색'과 '진한 파란색'이 의미상 서로 이웃 관계라는 것을 인식하지 못합니다.

해결책: 저자들은 **그룹화된 교차 엔트로피(Grouped Cross-Entropy, GCE)**라는 새로운 훈련 규칙을 도입했습니다.

비유: 로봇에게 정확한 블록을 고르지 못했다고 벌을 주는 대신, 선생님은 이렇게 말합니다. "네가 내가 원했던 연한 파란색 대신 진한 파란색을 골랐구나. 그래도 비슷해! 부분 점수를 줄게."
그들은 10만 개 이상의 블록을 작은 그룹(예: "파란색 계열", "빨간색 계열", "초록색 계열")으로 조직했습니다. 훈련 중에 로봇이 (비록 완벽하게 정확한 것은 아닐지라도) 올바른 그룹 내의 블록을 선택하면 긍정적인 신호를 받게 됩니다. 이는 로봇이 블록 간의 '관계'를 학습하도록 도와주며, 거대한 어휘 속에서도 길을 잃지 않고 훨씬 쉽게 학습할 수 있게 합니다.

3. 속도 향상 (커스텀 연산자)

문제점: 이러한 "그룹화된" 수학 연산을 수행하는 것은 계산량이 매우 많습니다. 보통 많은 컴퓨터 메모리(VRAM)를 요구하고 속도를 늦추기 때문에, 거대 모델을 훈련하는 데 어려움이 따릅니다.

해결책: 팀은 이 수학 연산을 수행하기 위한 특화된 도구(융합 연산자, fused operator)를 구축했습니다.

비유: 복잡한 회계 업무를 한다고 상상해 보세요. 예전 방식은 계산기로 계산하고, 숫자를 적고, 펜을 들어 장부에 기록한 뒤, 다시 계산기를 사용하는 것이었습니다. 새로운 방식은 계산과 장부 기록을 단 한 번의 빠르고 매끄러운 동작으로 처리하는 특수 기계를 사용하는 것과 같습니다. 이를 통해 엄청난 양의 메모리와 시간을 절약했습니다.

결과

논문은 이 새로운 로봇인 Nemotron-Labs-Diffusion-Image가 챔피언임을 보여줍니다:

더 나은 품질: 이전의 레고 기반 방식보다 더 선명하고 정확한 이미지를 생성합니다.
더 빠른 속도: 진행 과정에서 스스로 실수를 바로잡을 수 있기 때문에 더 적은 단계로 이미지를 생성할 수 있습니다.
효율성: 컴퓨터 메모리를 과부하시키지 않고도 매우 상세한 어휘를 다룰 수 있습니다.

요약하자면, 그들은 경직되고 오류가 발생하기 쉬웠던 방식에 자가 수정 능력을 부여하고, 어휘의 "이웃 관계"를 이해하도록 가르쳤으며, 이를 실행할 더 빠른 엔진을 구축했습니다. 그 결과, 더 똑똑하고 효율적인 고해상도 이미지 생성기가 탄생했습니다.

기술 요약: Nemotron-Labs-Diffusion-Image

문제 정의

잠재 확산 모델(Latent Diffusion Models, LDMs)은 텍스트-이미지 합성 분야에서 최첨단 성능을 달나 성취했지만, 이산적 이미지 생성(discrete image generation)은 대규모 언어 모델(LLM)과의 네이티브 호환성 및 확립된 LLM 최적화 기술(예: 시퀀스 패킹, KV 캐싱)을 활용할 수 있다는 뚜렷한 장점을 제공합니다. 그러나 기존의 마스크 기반 이산 확산 모델(Masked Discrete Diffusion Models, MDMs)은 연속적인 LDM과 비교했을 때 확장성과 충실도 측면에서 두 가지 근본적인 문제에 직면해 있습니다:

자기 수정 능력의 부재: 반복적으로 잠재 표현을 정제하는 연속적인 LDM과 달리, 표준 MDM은 토큰이 언마스킹(unmasking)되는 순간 해당 토큰을 확정해 버립니다. 이산 토큰은 언마스킹된 이후에 수정될 수 없기 때문에, 추론 과정에서 오류가 누적됩니다. 토큰 간 독립성을 가정하는 표준 샘플링 방식은 이미지 토큰에 존재하는 강력한 의존성을 포착하지 못하여 하위 최적의 결합 분포를 초래합니다.
대규모 어휘집에서의 코드북 희소성: 고충실도 이미지 생성을 위해서는 이산화 오류를 줄이기 위해 큰 코드북(어휘집)이 필요합니다. 그러나 어휘집 크기가 커지면 학습 신호의 심각한 희소성이 발생합니다. 즉, 개별 토큰이 데이터셋에 드물게 나타나게 됩니다. 표준 원-핫 교차 엔트로피(Cross-Entropy, CE) 손실은 정답이 아닌 모든 토큰을 동일하게 부정적인 것으로 취급하며, 임베딩 공간에서 서로 가깝지만 코드북상으로는 다른 토큰 사이의 의미적 유사성을 포착하지 못합니다. 이는 최적화의 어려움을 유발하며, 수렴을 위해 막대한 모델 용량과 데이터를 요구합니다.

방법론

저자들은 이러한 과제를 해결하기 위해 새로운 구조적 및 목적 함수 혁신을 통해 설계된 8B 파라미터 규모의 파운데이션 모델인 **Nemotron-Labs-Diffusion-Image (NLD-Image)**를 제안합니다.

1. 아키텍처: 디코더 전용 트랜스포머 (Decoder-Only Transformer)

고정된 텍스트 인코더를 사용하는 인코더-디코더 프레임워크(예: Meissonic)와 달리, NLD-Image는 단일 디코더 전용 트랜스포머를 사용합니다.

초기화: 모델은 사전 학습된 확산 언어 모델로부터 초기화되어, 기존의 언어 이해 능력을 활용합니다.
장점: 이 설계는 고정된 인코더(예: CLIP의 77-토큰 제한)의 컨텍스트 길이 제약을 제거하고, 더 긴 프롬프트를 지원하며, 학습 효율성을 위한 시퀀스 패킹을 가능하게 합니다. 최종 언어 모델링 헤드는 이미지 토큰을 예측하기 위해 새롭게 초기화된 MLP로 대체됩니다.

2. 토큰 편집 메커니즘 (Token Editing Mechanism, 자기 수정)

연속적인 확산과 유사한 반복적 정제를 가능하게 하기 위해, 저자들은 토큰 편집 메커니즘을 도입합니다.

학습: 순방향 확산 과정은 **토큰 오염(token corruption)**을 포함하도록 확장됩니다. 일부 토큰을 마스킹하는 것 외에도, 모델은 가시적인(깨끗한) 토큰의 일부를 오염시킵니다. 모델은 모든 위치에 대해 확률을 예측하도록 학습됩니다: 즉, 마스킹된 토큰을 언마스킹하는 것뿐만 아니라 오염된 가시적 토큰을 수정하는 것입니다.
추론: 샘플링 중에 모델은 언마스킹된 토큰을 추출할 뿐만 아니라 이미 언마스킹된 토큰을 다시 평가합니다. 만약 모델이 특정 임계값 $\tau$ 이상의 신뢰도로 다른 토큰을 예측한다면, 기존 토큰은 교체됩니다. 이를 통해 모델은 생성 과정 중에 동적으로 오류를 스스로 수정할 수 있습니다.

3. 그룹 교차 엔트로피 (Grouped Cross-Entropy, GCE) 목적 함수

코드북 희소성을 완화하기 위해, 저자들은 계층적 손실 함수인 **그룹 교차 엔트로피(GCE)**를 제안합니다.

클러스터링: 대규모 코드북(예: 131k 토큰)을 임베딩 공간의 근접성을 기반으로 K-means를 사용하여 작은 그룹(예: 16k, 8k)으로 클러스터링합니다.
손실 함수: 이 목적 함수는 표준 토큰 수준의 원-핫 CE 손실과 보조적인 클러스터 수준의 CE 손실을 결합합니다. 모델은 모든 토큰에 대한 확률을 예측하며, 이들은 클러스터 확률을 얻기 위해 합산됩니다. 손실은 여러 층위에서 특정 정답 토큰과 정답 클러스터 모두에 적용됩니다.
이점: 이는 정답과 의미적으로 유사한(동일한 클러스터 내에 있는) 토큰들에게도 양의 감독 신호(positive supervision)를 제공함으로써, 비록 정확한 일치는 아닐지라도 신호의 희소성을 완화합니다.
구현: 나이브한 PyTorch 구현에 비해 VRAM 사용량과 지연 시간을 크게 줄이기 위해 커스텀 퓨즈드 연산자(custom fused operator)가 구현되었습니다.

주요 기여

Nemotron-Labs-Diffusion-Image: 고해상도(1024px) 텍스트-이미지 합성을 위한 8B 규모의 최첨단 마스크 기반 이산 확산 모델입니다.
토큰 편집 (Token Editing): MDM이 추론 중에 출력을 반복적으로 정제하고 스스로 수정할 수 있게 하는 새로운 메커즘으로, 기존 이산 모델의 자기 정제 능력 부족 문제를 해결합니다.
그룹 교차 엔트로피 (GCE): 의미적으로 인접한 토큰에 양의 학습 신호를 할당함으로써 코드북 희소성을 완화하고, 대규모 어휘집 생성기의 최적화 안정성과 효율성을 개선하는 학습 목적 함수입니다.
효율적인 구현: GCE를 위한 커스텀 퓨즈드 연산자를 통해 메모리 점유율과 계산 오버헤드를 줄여 대규모 어휘집 모델의 학습을 용이하게 합니다.

실험 결과

모델은 프로그레시브 업스케일링 전략(256px에서 1024px로)을 사용하여 64개의 H100 GPU에서 300k 스텝 동안 학습되었습니다.

벤치마크 성능:
- GenEval: 0.90의 점수를 기록하여, 전문 모델인 Meissonic(0.54)과 통합 멀티모달 모델인 LaViDa-O(0.77)를 능가했으며, 최전선 모델인 Qwen-Image-2507(0.87)과 경쟁했습니다.
- DPG: 86.9를 달성하여 SD3-Medium(83.5) 및 Flux-Dev를 넘어섰습니다.
- MJHQ-30k: 미세 조정 후 HPSv3 점수 10.76을 기록하여 높은 인간 선호도 정렬을 입증했습니다.
절제 연구 (Ablation Studies):
- 토큰 편집: 정성적 및 정량적 결과(HPSv3)는 토큰 편집이 이미지 충실도를 일관되게 향상시키고 오류 누적을 줄인다는 것을 보여줍니다. 이를 통해 모델은 편집 없이 64 스텝이 필요했던 품질을 32 스텝(NFE)만으로 달성할 수 있어, 고정된 품질 목표 대비 추론 속도를 실질적으로 두 배 높였습니다.
- GCE: ImageNet 클래스 조건부 생성에서, GCE는 다양한 모델 크기와 학습 에포크에 걸쳐 바닐라 Cross-Entropy 및 SNCE(Soft Nearest Neighbor Cross-Entropy)보다 FID 점수 면에서 일관되게 우수한 성능을 보였습니다.
효율성: 최적화된 GCE 연산자는 이기적인(eager) 구현과 비교했을 때 지연 시간을 44ms에서 20ms로, 최대 VRAM 사용량을 25GB에서 16GB로 줄였습니다. NLD-Image는 Emu3.5(AR)보다 더 높은 GenEval 점수를 달성하면서도 42.4배 빠른 추론 속도를 보여주었습니다.
적은 단계의 생성: 모델은 증류(distillation) 없이도 단 4단계 만에 준수한 품질의 이미지를 생성하며, 낮은 단계 수에서 여전히 흐릿한 모습을 보이는 연속적 플로우 매칭(flow-matching) 베이스라인보다 빠르게 수렴합니다.

의의 및 주장

본 논문은 NLD-Image가 이산적 MDM과 연속적 LDM 사이의 성능 격차를 성공적으로 메움으로써 이산 이미지 생성 분야의 중요한 진전을 이루었다고 주장합니다. 토큰 편집을 도입함으로써, 저자들은 이산 모델이 고유한 자기 수정 능력 부족을 극복하고 반복적인 정제가 가능하다는 것을 입증했습니다. 또한, GCE 목적 함수는 코드북 희소성 문제를 해결하는 확장 가능한 솔루션을 제공하여, 과도한 학습 비용 없이도 크고 표현력이 풍부한 어휘집을 사용할 수 있게 합니다.

저자들은 이 작업이 텍-이미지를 단일 프레임워크 내에서 처리할 수 있는 디코더 전용 아키텍처를 활용하여, 통합 멀티모달 모델을 향한 단계로 자리매김한다고 보고 있습니다. 결과에 따르면, 이러한 특정 혁신 기술을 갖춘 이산 확산 모델은 최첨단의 충실도와 효율성을 달려낼 수 있으며, 이는 더욱 강력하고 확장 가능한 이미지 합성 시스템을 위한 길을 열어줍니다.