Sparse-to-Sparse Training of Diffusion Models

Resumo Técnico: Treinamento de Esparso para Esparso de Modelos de Difusão

Definição do Problema
Os Modelos de Difusão (DMs) alcançaram um desempenho de estado da arte em tarefas generativas, incluindo síntese de imagens, processamento de linguagem natural e modelagem de dados temporais. No entanto, sua adoção generalizada é dificultada pelos significativos custos computacionais durante tanto o treinamento quanto a inferência. Embora pesquisas anteriores tenham se concentrado amplamente na aceleração da etapa de inferência (por exemplo, através de destilação ou amostragem mais rápida), o processo de treinamento permanece extenso e intensivo em recursos. Além disso, as técnicas existentes de compressão de rede, como a quantização pós-treinamento ou a poda de modelos densos pré-treinados, ainda exigem o treinamento inicial caro de redes densas de escala total. Há uma necessidade de métodos que reduzam os custos computacionais e as pegadas de memória desde o início do processo de treinamento.

Metodologia
Este artigo introduz o paradigma de treinamento de esparso para esparso aos Modelos de Difusão, visando treinar redes esparsas do zero em vez de comprimir modelos densos pré-treinados. O estudo foca na esparsidade não estruturada e avalia três estratégias de treinamento específicas aplicadas a duas arquiteturas de estado da arte de DM: Difusão Latente (para dados de imagem contínuos) e ChiroDiff (para dados de esboço espaço-temporais discretos).

Os três métodos propostos são:

1. Static-DM (Treinamento Esparso Estático): A conectividade da rede é fixa na inicialização com base em uma razão de esparsidade. Os autores utilizam estratégias de Erdős–Rényi (ER) e ERK para alocar pesos não nulos, garantindo que camadas maiores tenham maior esparsidade. O padrão de conectividade permanece inalterado durante todo o treinamento.
2. RigL-DM (Treinamento Esparso Dinâmico): Este método utiliza um esquema dinâmico de poda e crescimento (prune-and-grow). Os pesos são podados com base em sua magnitude, e novas conexões são regeneradas com base na magnitude dos gradientes (crescimento por gradiente).
3. MagRan-DM (Treinamento Esparso Dinâmico): Semelhante ao RigL-DM em relação à poda (baseada em magnitude), mas o crescimento é determinado aleatoriamente em vez de por gradientes.

Os experimentos abrangem seis conjuntos de dados (LSUN-Bedrooms, CelebA-HQ, Imagenette, QuickDraw, KanjiVG e VMNIST) através de vários níveis de esparsidade ($S \in \{0,1; 0,25; 0,5; 0,75; 0,9\}$). Os autores também investigam o impacto da taxa de poda e regeneração ($p$) e a frequência de exploração ($\Delta T_e$) para os métodos dinâmicos.

Principais Contribuições

• Primeira Aplicação de Treinamento de Esparso para Esparso em DMs: O artigo estabelece a viabilidade de treinar Modelos de Difusão incondicionais do zero usando conectividade esparsa, cobrindo tanto estratégias estáticas quanto dinâmicas.
• Avaliação Empírica Abrangente: O estudo compara três métodos de treinamento esparso contra bases densas em duas arquiteturas de modelos distintas e seis conjuntos de dados, fornecendo uma visão ampla dos efeitos da esparsidade no desempenho dos DMs.
• Identificação de Hiperparâmetros Ótimos: A pesquisa identifica que, embora o treinamento esparso dinâmico (DST) geralmente funcione bem, a escolha da taxa de poda e regeneração é crítica. Uma taxa conservadora ($p=0,05$) produz melhores resultados em regimes de alta esparsidade comparado à taxa mais agressiva ($p=0,5$) frequentemente usada em outros contextos.
• Compromissos entre Desempenho e Eficiência: O trabalho demonstra que os DMs esparsos podem igualar ou exceder o desempenho de seus equivalentes densos, reduzindo significativamente o número de parâmetros treináveis e as Operações de Ponto Flutuante (FLOPs).

Resultados Experimentais

• Desempenho: Na maioria dos casos, pelo menos uma configuração esparsa superou seu equivalente denso. Por exemplo, no conjunto de dados CelebA-HQ, o RigL-DM com 25% de esparsidade alcançou um FID score menor (32,12) do que o modelo denso (32,74), reduzindo parâmetros em 25% e FLOPs em 10%. No QuickDraw, o MagRan-DM com 90% de esparsidade manteve um desempenho comparável ao modelo denso, reduzindo os FLOPs em 88%.
• Níveis de Esparsidade: Níveis moderados de esparsidade (25%–50%) geralmente produziram o melhor equilíbrio entre desempenho e eficiência. A esparsidade extremamente alta (ex: 90%) frequentemente levou à degradação do desempenho, embora isso varie conforme o conjunto de dados e a arquitetura do modelo.
• Dinâmico vs. Estático: Embora os métodos de Treinamento Esparso Dinâmico (DST) (RigL-DM, MagRan-DM) frequentemente superem o Static-DM em regimes de esparsidade moderada, o Static-DM mostrou resultados competitivos ou superiores em configurações de esparsidade muito alta ($S > 0,5$) quando utilizando a taxa agressiva de $p=0,5$. No entanto, ao utilizar a taxa conservadora de $p=0,05$, os métodos DST recuperaram a superioridade mesmo em regimes de alta esparsidade.
• Dinâmica de Treinamento: Os modelos esparsos exibiram dinâmicas de treinamento estáveis semelhantes às de seus equivalentes densos, sugerindo que o paradigma esparso-para-esparso não desestabiliza o processo de treinamento de difusão.
• Geração Condicional: Experimentos preliminares em modelos condicionais de classe (ImageNet e QuickDraw) indicaram que o treinamento esparso-para-esparso também é eficaz em configurações condicionais, embora os modelos densos geralmente tenham mantido uma leve vantagem no ImageNet.

Significância e Alegações
O artigo afirma que o treinamento esparso-para-esparso oferece um caminho promissor para melhorar a eficiência dos Modelos de Difusão sem sacrificar a qualidade da geração. Ao treinar redes esparsas do zero, a abordagem aborda tanto os custos de treinamento quanto de inferência, potencialmente reduzindo a pegada de carbono e os requisitos de memória associados aos DMs. Os autores enfatizam que a seleção do método apropriado (Estático vs. Dinâmico), do nível de esparsidade e dos hiperparâmetros (especificamente a taxa de poda/regeneração) é crucial para o sucesso. O trabalho sugere que, com a configuração correta, os DMs esparsos podem ser uma alternativa viável aos modelos densos, permitindo uma implementação mais ampla e uma iteração mais rápida na pesquisa de IA generativa. Os autores observam que, embora o hardware atual muitas vezes limite o ganho de velocidade prático da esparsidade não estruturada, a redução no número de parâmetros e FLOPs permanece uma vantagem teórica e prática significativa, especialmente à medida que o suporte de hardware para operações esparsas evolui.