Nucleus-Image: Sparse MoE for Image Generation

Resumen Técnico: Nucleus-Image

1. El Problema
La generación de imágenes basada en texto (text-to-image) de alta calidad ha estado dominada por modelos densos masivos que requieren un costo computacional y de memoria prohibitivo durante la inferencia. Aunque los modelos de Mixture of Experts (MoE) han demostrado ser eficientes en el procesamiento de lenguaje natural (LLMs), su aplicación en modelos de difusión para imágenes ha sido limitada debido a desafíos específicos:

• Inestabilidad en el entrenamiento: La modulación adaptativa basada en el paso de tiempo (timestep) en transformadores de difusión puede colapsar la especialización de los expertos si el enrutamiento no se maneja correctamente.
• Ineficiencia en la inferencia: La necesidad de procesar tokens de texto a través de capas densas o MoE aumenta la latencia y el uso de memoria.
• Curación de datos: La falta de corpus de entrenamiento masivos, limpios y bien estructurados para modelos MoE de alta calidad.

Nucleus-Image aborda estos problemas para establecer una nueva frontera de Pareto entre la calidad y la eficiencia, logrando resultados de vanguardia activando solo una fracción de sus parámetros totales.

2. Metodología

El equipo de Nucleus AI propone una arquitectura de transformador de difusión (DiT) que integra varias innovaciones clave:

• Arquitectura MoE Escalable:

  • El modelo tiene 17 mil millones de parámetros en total, pero activa solo ~2 mil millones por paso hacia adelante.
  • Utiliza una capa de Mezcla de Expertos (MoE) con 64 expertos enrutados y un experto compartido por capa (en 29 de las 32 capas del transformador).
  • Emplea Enrutamiento por Elección de Experto (Expert-Choice Routing): A diferencia del enrutamiento tradicional donde cada token elige expertos, aquí cada experto selecciona sus mejores tokens. Esto garantiza una utilización uniforme sin necesidad de pérdidas auxiliares de balanceo de carga complejas.

• Diseño de Arquitectura Optimizado:

  • Exclusión de Tokens de Texto: Los tokens de texto no pasan por las capas MLP ni MoE. Participan únicamente como contribuyentes de claves y valores (KV) en la atención cruzada. Esto elimina la sobrecarga de enrutamiento para el texto y permite el caching completo de KV del texto a través de los pasos de eliminación de ruido.
  • Atención Conjunta (Joint Attention): Se realiza una sola operación de atención donde las consultas de la imagen atienden a las claves y valores concatenados de imagen y texto, reduciendo la complejidad computacional.
  • Enrutamiento Desacoplado: Para evitar que la modulación del paso de tiempo (que varía drásticamente en magnitud) domine las decisiones de enrutamiento, el router recibe una representación no modulada concatenada con el embedding del tiempo, mientras que los expertos reciben la representación totalmente modulada. Esto separa la decisión de enrutamiento (basada en contenido y tiempo) del cálculo del experto.

• Pipeline de Datos y Entrenamiento:

  • Corpus: Entrenado con un conjunto de datos curado de 700 millones de imágenes únicas, generando 1.5 mil millones de pares imagen-captión. Incluye filtrado por calidad, deduplicación y curación de captiones a múltiples granularidades.
  • Curriculum Progresivo: Entrenamiento en tres etapas de resolución (256 → 512 → 1024 píxeles) con bucketing de múltiples relaciones de aspecto desde el inicio.
  • Factor de Capacidad Progresivo: El número de tokens que cada experto puede seleccionar se reduce progresivamente a medida que aumenta la resolución y la especialización de los expertos.
  • Optimización: Uso del optimizador Muon (con ortogonalización de momentum) combinado con un cronograma de tasa de aprendizaje Warmup-Stable-Merge (WSM). Esto elimina la necesidad de una copia sombra EMA (Exponential Moving Average) durante el entrenamiento, reduciendo el uso de memoria.
  • Pérdidas: Se utiliza la pérdida de flujo rectificado (Rectified Flow) como objetivo principal, complementada con pérdidas de wavelet en la etapa final para mejorar los detalles de alta frecuencia.

• Optimizaciones de Kernel:

  • Implementación de kernels personalizados en Triton para multiplicaciones de matrices agrupadas, activaciones SwiGLU fusionadas y permutación de tokens, minimizando el cuello de botella de ancho de banda de memoria.

3. Contribuciones Clave

1. Primer Modelo MoE de Difusión de Alta Calidad Open Source: Nucleus-Image es el primer modelo de difusión basado en MoE de este nivel de calidad totalmente abierto (pesos, código y datos), demostrando que la escalabilidad MoE es viable para la generación de imágenes.
2. Eficiencia de Inferencia Sin Precedentes: Logra un rendimiento comparable a modelos con presupuestos de parámetros activos mucho mayores, activando solo ~2B parámetros de un total de 17B.
3. Estabilidad en MoE para Difusión: Introduce el diseño de enrutamiento desacoplado y capas densas iniciales, resolviendo problemas de inestabilidad específicos de los transformadores de difusión que no existen en los LLMs.
4. Eliminación de la Carga EMA: Mediante el cronograma WSM y la fusión de checkpoints, elimina la necesidad de mantener una copia duplicada de los pesos del modelo durante el entrenamiento, ahorrando memoria y permitiendo una mayor flexibilidad en la duración del entrenamiento.
5. Pipeline de Datos Masivo y Curado: La construcción de un corpus de 700M imágenes con captiones de múltiples granularidades y un sistema de "tiering" de calidad que alimenta un curriculum de aprendizaje progresivo.

4. Resultados

El modelo fue evaluado en tres benchmarks estándar sin ningún tipo de ajuste post-entrenamiento (sin RLHF, DPO o afinamiento de preferencias humanas):

• GenEval: Logra una puntuación general de 0.87, igualando a Qwen-Image y superando a modelos como FLUX.1 Dev (0.66) y SD3.5 Large (0.71). Destaca especialmente en la comprensión de posición espacial (0.85), superando a todos los demás modelos.
• DPG-Bench: Obtiene la puntuación general más alta de 88.79, superando a Qwen-Image (88.32) y Seedream 3.0 (88.27). Lidera en las subcategorías de entidad, atributo y otros, demostrando una fuerte adherencia a prompts densos.
• OneIG-Bench: Alcanza una puntuación global de 0.522, superando a Imagen4 (0.515) y Recraft V3 (0.502). Destaca en la categoría de estilo (0.430) y tiene un rendimiento competitivo en renderizado de texto.

El modelo demuestra un rendimiento de vanguardia (SOTA) o cercano a SOTA en todas las métricas, a pesar de tener un presupuesto de parámetros activos significativamente menor que sus competidores.

5. Significado e Impacto

Nucleus-Image demuestra que la escalabilidad mediante MoE es una ruta altamente efectiva para la generación de imágenes de alta calidad, desafiando la noción de que se requieren modelos densos masivos para lograr los mejores resultados.

• Eficiencia de Recursos: Al activar solo una fracción de los parámetros, reduce drásticamente el costo de inferencia y el uso de memoria, haciendo que modelos de alta calidad sean más accesibles.
• Reproducibilidad y Transparencia: Al liberar los pesos, el código de entrenamiento y el dataset, el equipo de Nucleus AI fomenta la investigación abierta en la eficiencia de modelos de difusión.
• Nuevos Paradigmas de Entrenamiento: Las innovaciones en la estabilidad del enrutamiento (desacoplamiento tiempo-contenido) y la eliminación de EMA mediante fusión de checkpoints ofrecen nuevas mejores prácticas para el entrenamiento de modelos de difusión a gran escala.

En resumen, Nucleus-Image establece un nuevo estándar para la generación de imágenes eficiente, probando que la arquitectura, la optimización y la curación de datos pueden compensar la necesidad de un aumento lineal en el costo computacional.