Sparse-to-Sparse Training of Diffusion Models

Resumen Técnico: Entrenamiento de Sparse-to-Sparse en Modelos de Difusión

Planteamiento del Problema
Los Modelos de Difusión (DMs, por sus siglas en inglés) han alcanzado un rendimiento de vanguardia en tareas generativas, incluyendo la síntesis de imágenes, el procesamiento de lenguaje natural y el modelado de datos temporales. Sin embargo, su adopción generalizada se ve obstaculizada por los significativos costos computacionales tanto durante el entrenamiento como durante la inferencia. Mientras que la investigación previa se ha centrado principalmente en acelerar la etapa de inferencia (por ejemplo, mediante la destilación o el muestreo rápido), el proceso de entrenamiento sigue siendo extenso y demandante de recursos. Además, las técnicas de compresión de redes existentes, como la cuantización post-entrenamiento o la poda de modelos densos pre-entrenados, aún requieren el entrenamiento inicial costoso de redes densas a escala completa. Existe la necesidad de métodos que reduzcan los costos computacionales y las huellas de memoria desde el inicio del proceso de entrenamiento.

Metodología
Este artículo introduce el paradigma de entrenamiento sparse-to-sparse (de disperso a disperso) en los Modelos de Difusión, con el objetivo de entrenar redes dispersas desde cero en lugar de comprimir modelos densos pre-entrenados. El estudio se centra en la dispersión no estructurada y evalúa tres estrategias de entrenamiento específicas aplicadas a dos arquitecturas de DM de vanguardia: Latent Diffusion (para datos de imagen continuos) y ChiroDiff (para datos de bocetos espaciotemporales discretos).

Los tres métodos propuestos son:

1. Static-DM (Entrenamiento Disperso Estático): La conectividad de la red se fija en la inicialización basada en una relación de dispersión. Los autores emplean estrategias de Erdős–Rényi (ER) y ERK para asignar pesos no nulos, asegurando que las capas más grandes tengan una mayor dispersión. El patrón de conectividad permanece sin cambios durante todo el entrenamiento.
2. RigL-DM (Entrenamiento Disperso Dinámico): Este método utiliza un esquema dinámico de poda y crecimiento (prune-and-grow). Los pesos se podan basándose en su magnitud, y las nuevas conexiones se regeneran basándose en la magnitud de los gradientes (crecimiento por gradiente).
3. MagRan-DM (Entrenamiento Disperso Dinámico): Similar a RigL-DM en cuanto a la poda (basada en la magnitud), pero el crecimiento se determina de forma aleatoria en lugar de mediante gradientes.

Los experimentos cubren seis conjuntos de datos (LSUN-Bedrooms, CelebA-HQ, Imagenette, QuickDraw, KanjiVG y VMNIST) a través de varios niveles de dispersión ($S \in \{0.1, 0.25, 0.5, 0.75, 0.9\}$). Los autores también investigan el impacto de la tasa de poda y regeneración ($p$) y la frecuencia de exploración ($\Delta T_e$) para los métodos dinámicos.

Contribuciones Clave

• Primera Aplicación de Entrenamiento Sparse-to-Sparse a DMs: El artículo establece la viabilidad de entrenar Modelos de Difusión incondicionales desde cero utilizando conectividad dispersa, cubriendo tanto estrategias estáticas como dinámicas.
• Evaluación Empírica Exhaustiva: El estudio compara tres métodos de entrenamiento disperso contra bases densas en dos arquitecturas de modelos distintos y seis conjuntos de datos, proporcionando una visión amplia de los efectos de la dispersión en el rendimiento de los DM.
• Identificación de Hiperparámetros Óptimos: La investigación identifica que, si bien el entrenamiento disperso dinámico (DST) generalmente funciona bien, la elección de la tasa de poda y regeneración es crítica. Una tasa conservadora ($p=0.05$) produce mejores resultados en regímenes de alta dispersión en comparación con la tasa más agresiva ($p=0.5$) utilizada a menudo en otros contextos.
• Compromiso entre Rendimiento y Eficiencia: El trabajo demuestra que los DM dispersos pueden igualar o incluso superar el rendimiento de sus contrapartes densas, reduciendo significativamente el número de parámetros entrenables y las operaciones de punto flotante (FLOPs).

Resultados Experimentales

• Rendimiento: En la mayoría de los casos, al menos una configuración dispersa superó a su contraparte densa. Por ejemplo, en el conjunto de datos CelebA-HQ, RigL-DM con un 25% de dispersión logró un FID más bajo (32.12) que el modelo denso (32.74), reduciendo simultáneamente los parámetros en un 25% y los FLOPs en un 10%. En QuickDraw, MagRan-DM con un 90% de dispersión mantuvo un rendimiento comparable al modelo denso reduciendo los FLOPs en un 88%.
• Niveles de Dispersión: Los niveles de dispersión moderados (25%–50%) generalmente ofrecieron el mejor equilibrio entre rendimiento y eficiencia. La dispersión extremadamente alta (por ejemplo, 90%) a menudo provocó una degradación del rendimiento, aunque esto varió según el conjunto de datos y la arquitectura del modelo.
• Dinámico vs. Estático: Aunque los métodos de Entrenamiento Disperso Dinámico (DST) (RigL-DM, MagRan-DM) a menudo superaron a Static-DM en regímenes de dispersión moderada, Static-DM mostró resultados competitivos o superiores en configuraciones de muy alta dispersión ($S > 0.5$) cuando se utilizó la tasa agresiva $p=0.5$. Sin embargo, al utilizar la tasa conservadora $p=0.05$, los métodos DST recuperaron su superioridad incluso en regímenes de alta dispersión.
• Dinámica de Entrenamiento: Los modelos dispersos exhibieron dinámicas de entrenamiento estables similares a sus contrapartes densas, lo que sugiere que el paradigma sparse-to-sparse no desestabiliza el proceso de entrenamiento de la difusión.
• Generación Condicional: Experimentos preliminares en modelos condicionales de clase (ImageNet y QuickDraw) indicaron que el entrenamiento sparse-to-sparse también es efectivo en entornos condicionales, aunque los modelos densos generalmente mantuvieron una ligera ventaja en ImageNet.

Significancia y Reivindicaciones
El artículo afirma que el entrenamiento sparse-to-sparse ofrece un camino prometedor para mejorar la eficiencia de los Modelos de Difusión sin sacrificar la calidad de la generación. Al entrenar redes dispersas desde cero, el enfoque aborda tanto los costos de entrenamiento como de inferencia, reduciendo potencialmente la huella de carbono y los requisitos de memoria asociados con los DM. Los autores enfatizan que la selección del método apropiado (Estático vs. Dinámico), el nivel de dispersión y los hiperparámetros (específicamente la tasa de poda/regeneración) es crucial para el éxito. El trabajo sugiere que, con la configuración adecuada, los DM dispersos pueden ser una alternativa viable a los modelos densos, permitiendo una implementación más amplia y una iteración más rápida en la investigación de IA generativa. Los autores señalan que, si bien el hardware actual suele limitar la aceleración práctica de la dispersión no estructurada, la reducción en el recuento de parámetros y FLOPs sigue siendo una ventaja teórica y práctica significativa, especialmente a medida que evolucione el soporte de hardware para operaciones dispersas.