Recientemente, investigadores de instituciones como la Universidad de California en Irvine utilizaron estrategias como el enmascaramiento retardado, MoE y la expansión jerárquica para reducir el costo de entrenamiento de los modelos de difusión a $1,890. Esto es un orden de magnitud menor que el método más barato anterior (Wuerstchen) de $28,400, y modelos como Stable Diffusion tienen costos aún más altos.
Los investigadores probaron varias estrategias para reducir estos enormes gastos. Por ejemplo, el modelo de difusión original requería alrededor de 1000 pasos del ruido a la imagen, que ahora se ha reducido a alrededor de 20 pasos o incluso menos. Cuando los módulos básicos en los modelos de difusión fueron gradualmente reemplazados de Unet (CNN) a DiT (Transformer), también se implementaron algunas optimizaciones basadas en las características de Transformer, como la cuantización, omitir algunos cálculos redundantes en la atención, pipeline, etc.
Este estudio entrenó desde cero un modelo de difusión de 1.16 mil millones de parámetros por solo $1,890. Esto representa una mejora de un orden de magnitud en comparación con el estado del arte, dando esperanza a las personas comunes de poder tocar el pre-entrenamiento. Más importante aún, las técnicas de reducción de costos no afectaron el rendimiento del modelo, con 1.16 mil millones de parámetros dando muy buenos resultados. Además de la apariencia, los indicadores de datos del modelo también son excelentes, con puntuaciones FID muy cercanas a Stable Diffusion 1.5 y DALL·E 2.
Los secretos para ahorrar dinero incluyen principalmente:
-
Estrategia de enmascaramiento retardado: Utilizar un mezclador de parches para preprocesar antes del enmascaramiento, incorporando la información de los parches descartados en los parches supervivientes, reduciendo significativamente la degradación del rendimiento causada por el alto enmascaramiento.
-
Ajuste fino: Realizar un pequeño ajuste fino (desenmascarar) después del pre-entrenamiento (enmascaramiento) para mitigar los artefactos de generación adversos causados por el uso de máscaras.
-
MoE y expansión jerárquica: Utilizar capas MoE simplificadas basadas en el enrutamiento de selección de expertos para aumentar los parámetros y la capacidad expresiva del modelo sin aumentar significativamente los costos de entrenamiento. También se consideró el método de escalado jerárquico, aumentando linealmente el ancho de los bloques Transformer.
En cuanto a la configuración experimental, los autores utilizaron dos variantes de DiT: DiT-Tiny/2 y DiT-Xl/2, con un tamaño de parche de 2. Todos los modelos se entrenaron utilizando el optimizador AdamW con decaimiento de tasa de aprendizaje coseno y alto decaimiento de peso. El frontend del modelo utiliza el codificador automático variacional (VAE) de cuatro canales del modelo Stable-Diffusion-XL para extraer características de imagen, y también se probó el rendimiento del VAE de 16 canales más reciente en el entrenamiento a gran escala (versión económica).
Los autores utilizaron el marco EDM como configuración de entrenamiento unificada para todos los modelos de difusión, y utilizaron puntuaciones FID y CLIP para medir el rendimiento de los modelos de generación de imágenes. Se eligió el modelo CLIP más comúnmente utilizado como codificador de texto.
El conjunto de datos de entrenamiento utilizó tres conjuntos de datos de imágenes reales (Conceptual Captions, Segment Anything, TextCaps), que contienen 22 millones de pares de texto e imagen. Dado que SA1B no proporciona subtítulos reales, aquí se utilizaron subtítulos sintéticos generados por el modelo LLaVA.