Nucleus-Image: 희소 MoE 기반 고품질 이미지 생성 모델 기술 요약

1. 문제 정의 (Problem)

기존의 고품질 텍스트 - 이미지 생성 모델 (Diffusion Transformer) 은 일반적으로 방대한 파라미터 수를 활성화하여 높은 연산 비용과 메모리 사용량을 요구합니다. 특히 추론 (Inference) 단계에서 모든 파라미터를 사용하는 밀집 (Dense) 아키텍처는 효율성 측면에서 한계가 있으며, 고품질 생성을 달성하기 위해 더 많은 자원을 소모해야 하는 딜레마가 존재합니다. 또한, 기존 MoE(Mixture of Experts) 모델들은 시간 단계 (Timestep) 에 따른 변조 (Modulation) 신호가 라우팅 결정에 부정적인 영향을 미쳐 전문가 (Expert) 의 전문성이 붕괴되는 문제를 겪기도 합니다.

2. 방법론 (Methodology)

Nucleus-Image 는 희소 MoE (Sparse Mixture-of-Experts) 아키텍처를 확산 변환기 (Diffusion Transformer) 에 적용하여, 전체 모델 용량은 크게 늘리면서도 추론 시 활성화되는 파라미터는 최소화하는 전략을 취했습니다.

2.1 아키텍처 설계

• 스파스 MoE 구조: 총 170 억 (17B) 파라미터를 가진 32 레이어 트랜스포머 모델입니다. 각 MoE 레이어에는 64 개의 라우팅 전문가 (Routed Experts) 와 1 개의 공유 전문가 (Shared Expert) 가 있으며, Expert-Choice Routing 방식을 사용하여 순방향 전달 (Forward Pass) 당 약 20 억 (2B) 파라미터만 활성화됩니다.
• 분리된 라우팅 및 계산 (Decoupled Routing and Computation): 확산 모델의 고유한 문제인 시간 단계 (Timestep) 변조 신호가 라우팅 결정에 미치는 왜곡을 해결하기 위해, 라우팅은 변조되지 않은 토큰 표현과 시간 단계 임베딩을 기반으로 하고, 실제 전문가 계산은 변조된 표현을 기반으로 수행되도록 설계되었습니다. 이는 전문가의 공간적/의미적 전문성을 유지하면서 라우팅 안정성을 확보합니다.
• 텍스트 토큰 최적화: 텍스트 토큰은 트랜스포머 백본 (MLP 및 MoE 레이어) 을 통과하지 않고, 오직 크로스 어텐션 (Cross-Attention) 의 Key-Value 기여자 역할만 수행합니다. 이를 통해 텍스트 토큰에 대한 MoE 라우팅 오버헤드를 제거하고, 시간 단계별 KV 캐싱을 가능하게 하여 추론 속도를 획기적으로 개선했습니다.
• GQA (Grouped Query Attention): 16 개의 쿼리 헤드와 4 개의 키 - 밸류 헤드를 사용하여 KV 캐시 메모리를 4 배 줄였습니다.

2.2 데이터 파이프라인

• 대규모 정제 데이터셋: 7 억 5 천만 개의 고유 이미지에서 추출된 15 억 개의 이미지 - 캡션 쌍으로 구성된 고품질 코퍼스를 구축했습니다.
• 다단계 필터링 및 큐레이션: 미학적 점수 (Aesthetic Scoring), 중복 제거 (Deduplication), 캡션 정제 (Caption Curation) 를 거쳐 데이터의 품질을 관리했습니다.
• 합성 텍스트 렌더링: 모델의 텍스트 렌더링 능력을 향상시키기 위해 다양한 언어와 레이아웃을 가진 합성 텍스트 데이터를 별도로 포함했습니다.
• 커리큘럼 학습: 256 → 512 → 1024 픽셀로 점진적으로 해상도를 높이며 학습하고, 각 단계에서 다양한 종횡비 (Aspect Ratio) 버킷을 활용했습니다.

2.3 학습 전략

• 옵티마이저: Muon 옵티마이저를 사용하며, 파라미터 그룹 (Attention/FFN, Modulation, Bias 등) 에 따라 다른 최적화 설정을 적용했습니다.
• 학습률 스케줄링 (WSM): EMA(Exponential Moving Average) 의 메모리 오버헤드를 제거하기 위해 Warmup-Stable-Merge (WSM) 스케줄을 도입했습니다. 학습 중에는 학습률을 일정하게 유지하고, 학습 종료 후 체크포인트들을 가중치 평균하여 최종 모델을 생성합니다.
• 손실 함수: Rectified Flow Loss 를 주력으로 사용하며, MoE 균형을 위해 Z-loss 와 Orthogonal Loss 를 보조적으로 적용했습니다. 고해상도 단계에서는 Wavelet Loss 를 추가하여 고주파 세부 사항을 개선했습니다.

3. 주요 기여 (Key Contributions)

1. 최초의 오픈소스 고품질 MoE 확산 모델: Reinforcement Learning (RL), DPO(직접 선호도 최적화), 인간 선호도 튜닝 없이 순수 사전 학습 (Pre-training) 만으로 최고 수준의 성능을 달성한 최초의 오픈소스 MoE 확산 모델입니다.
2. 효율성과 성능의 새로운 Pareto Frontier: 17B 총 파라미터 중 2B 만 활성화하여, 훨씬 더 큰 활성 파라미터를 가진 모델들과 동급 또는 그 이상의 성능을 달성했습니다.
3. 혁신적인 아키텍처 설계:
   • 시간 단계 변조와 라우팅 결정을 분리하여 MoE 확산 모델의 안정성을 확보했습니다.
   • 텍스트 토큰을 MoE 레이어에서 완전히 배제하여 추론 효율성을 극대화했습니다.
4. 시스템 최적화: Triton 커널을 활용한 메모리 대역폭 병목 현상 해결 (Fused kernels) 과 Expert Parallelism을 통한 대규모 분산 학습 구현.
5. 완전한 오픈소스 공개: 모델 가중치, 학습 코드, 그리고 7 억 5 천만 이미지 규모의 데이터셋을 모두 공개했습니다.

4. 결과 (Results)

Nucleus-Image 는 GenEval, DPG-Bench, OneIG-Bench 등 주요 벤치마크에서 SOTA(State-of-the-Art) 또는 그에 준하는 성능을 기록했습니다.

• GenEval: 전체 점수 0.87로 Qwen-Image 와 동급이며, 공간 위치 이해 (Position) 분야에서 0.85 로 모든 모델을 압도했습니다.
• DPG-Bench: 전체 점수 88.79로 Qwen-Image(88.32) 와 Seedream 3.0(88.27) 을 제치고 1 위를 차지했습니다. 특히 엔티티 (Entity) 와 속성 (Attribute) 이해도가 뛰어났습니다.
• OneIG-Bench: 전체 점수 0.522로 Imagen4(0.515) 와 Recraft V3(0.502) 보다 높은 점수를 기록했습니다. 스타일 (Style) 생성 능력에서 특히 두드러졌습니다.
• 효율성: 2B 활성 파라미터만으로 17B 총 파라미터 모델의 성능을 구현하여, 추론 비용 대비 성능 효율성이 매우 뛰어납니다.

5. 의의 (Significance)

Nucleus-Image 는 희소 MoE(Sparse MoE) 스케일링이 고품질 이미지 생성을 위한 단순한 대안이 아닌, 매우 효과적이고 필수적인 확장 경로임을 입증했습니다.

• 비용 효율성: 고해상도 고품질 생성을 위해 막대한 연산 자원을 소모해야 한다는 기존 인식을 깨고, 적은 활성 파라미터로도 동등한 성능을 낼 수 있음을 보여줍니다.
• 기술적 진보: 확산 모델에 MoE 를 적용할 때 발생하는 라우팅 불안정성 (Timestep modulation collapse) 과 텍스트 처리 오버헤드 문제를 해결한 구체적인 아키텍처 솔루션을 제시했습니다.
• 오픈 생태계 기여: 고품질 데이터셋과 학습 코드, 모델 가중치를 공개함으로써 연구 커뮤니티가 효율적인 이미지 생성 모델을 재현하고 발전시킬 수 있는 기반을 마련했습니다.

결론적으로, Nucleus-Image 는 아키텍처, 최적화, 데이터 큐레이션, 시스템 효율성 측면에서의 종합적인 혁신을 통해, 제한된 계산 자원으로도 최상급의 텍스트 - 이미지 생성이 가능함을 증명했습니다.