Diff-ES: Stage-wise Structural Diffusion Pruning via Evolutionary Search

O artigo apresenta o Diff-ES, um framework de poda estrutural para modelos de difusão que utiliza busca evolutiva para otimizar automaticamente agendamentos de esparsidade por estágio e ativação dinâmica de pesos, alcançando aceleração de tempo real e alta qualidade de geração sem a sobrecarga de memória de métodos anteriores.

O essencial

Imagine que você tem um artista muito talentoso, mas extremamente lento, chamado Modelo de Difusão. Esse artista é capaz de criar quadros (imagens) incrivelmente realistas a partir de uma ideia (texto). O problema é que ele trabalha de uma forma muito peculiar: ele começa com uma tela cheia de "ruído" (como uma TV fora do ar) e, passo a passo, remove esse ruído até que a imagem apareça.

Para fazer um quadro perfeito, ele precisa dar muitos passos (às vezes 20, 50 ou mais). Cada passo exige muita energia do computador. É como se ele tivesse que polir a tela 50 vezes antes de entregar a obra. Isso torna o processo lento e caro para rodar em computadores comuns.

O Problema: "Cortar" o Artista

Alguns pesquisadores tentaram resolver isso cortando partes do artista (pruning) para torná-lo mais rápido. A ideia é: "Se ele usa 100% da sua capacidade em todos os 50 passos, e a gente corta 30% da capacidade dele, ele fica mais rápido, certo?"

O problema é que nem todos os passos são iguais.

• Nos primeiros passos, o artista está definindo a estrutura geral (onde fica o céu, onde fica a montanha). Se você cortar muita capacidade aqui, o quadro fica torto.
• Nos passos do meio, ele está adicionando detalhes (a cor da grama, a textura da pele).
• Nos últimos passos, ele está apenas polindo e ajustando cores finas.

Métodos antigos (como o MosaicDiff) tentavam cortar o artista de forma "manual" e baseada em palpites. Eles diziam: "Vamos cortar 30% no começo, 30% no meio e 30% no fim". Mas isso é como tentar vestir o mesmo tamanho de roupa em pessoas de alturas diferentes: não serve bem para todos. Às vezes, o corte no meio estraga a imagem, e o corte no final seria desnecessário.

A Solução: O "Caçador de Evolução" (Diff-ES)

Os autores deste paper criaram o Diff-ES. Pense nele como um treinador de atletas muito inteligente que usa uma técnica chamada "Busca Evolutiva".

Em vez de adivinhar onde cortar, o Diff-ES faz o seguinte:

1. Divide a jornada: Ele separa os passos do artista em "estágios" (como capítulos de um livro).
2. Cria uma população de ideias: Ele gera dezenas de planos diferentes. Um plano diz: "Corte muito no começo, pouco no meio". Outro diz: "Corte pouco no começo, muito no final".
3. Testa e Seleciona (A Evolução): Ele testa esses planos gerando algumas imagens.
   • Se um plano gera uma imagem feia (o urso de pelúcia tem 3 pernas), esse plano é eliminado.
   • Se um plano gera uma imagem bonita e rápida, ele é mantido.
4. Mutação: Ele pega os melhores planos e os mistura, trocando um pouco de "corte" de um estágio para outro (como trocar ingredientes numa receita), sempre mantendo o total de cortes igual ao desejado.
5. O Campeão: Depois de muitas rodadas, sobra o plano perfeito: o que corta exatamente onde é seguro cortar, sem estragar a imagem.

O Truque Mágico: O "Caminho Rápido" (Weight Routing)

Havia um problema: testar todos esses planos seria muito lento, porque teríamos que "reconstruir" o artista do zero para cada teste.

O Diff-ES inventou um truque genial chamado Roteamento de Pesos.
Imagine que o artista tem várias ferramentas (pesos) na sua caixa de ferramentas. Em vez de ter 10 artistas diferentes (cada um com sua própria caixa de ferramentas) ocupando espaço na sua mesa (memória do computador), você tem um único artista e uma caixa de ferramentas mágica.

• Quando o artista precisa fazer o "Passo 1", você entrega a ferramenta A.
• Quando ele vai para o "Passo 2", você troca rapidamente pela ferramenta B.
• Você não precisa ter todas as ferramentas montadas ao mesmo tempo. Você apenas "troca" a ferramenta necessária no momento certo.

Isso permite que o computador teste milhares de planos de corte sem precisar de mais memória do que o necessário.

O Resultado

O Diff-ES descobriu que, para alguns modelos, o melhor é cortar mais no início e no fim, deixando o meio intacto. Para outros, é o contrário. Como ele descobre isso automaticamente, ele funciona muito melhor do que os métodos antigos que usavam regras fixas.

Resumo da Ópera:
O Diff-ES é como um chef de cozinha que usa inteligência artificial para descobrir a receita perfeita de corte. Em vez de seguir um livro de receitas antigo e rígido, ele testa milhares de variações, descobre exatamente quanto de cada ingrediente (capacidade do modelo) pode ser removido em cada etapa da cozinha (etapas de geração da imagem), e usa um truque de organização (roteamento) para não ocupar a cozinha inteira com panelas extras. O resultado? Imagens tão bonitas quanto as originais, mas geradas muito mais rápido e com menos esforço do computador.

Resumo técnico

Resumo Técnico: Diff-ES

1. O Problema

Os modelos de difusão têm alcançado sucesso notável na geração de imagens de alta fidelidade, mas sua implantação prática é limitada por altos requisitos computacionais, decorrentes do processo iterativo de remoção de ruído (múltiplos passos) e do grande tamanho dos modelos.

• Limitações das Abordagens Existentes: Trabalhos anteriores focaram em reduzir o número de passos de amostragem ou comprimir parâmetros do modelo. Métodos de pruning (poda) estruturado, como o MosaicDiff, tentam alocar recursos computacionais de forma não uniforme ao longo dos passos de difusão, dividindo o processo em estágios e aplicando diferentes níveis de esparsidade.
• O Desafio: O MosaicDiff e métodos similares dependem de agendamentos de esparsidade manuais e baseados em heurísticas (ex: dividir o processo em 3 estágios com regras fixas). A importância dos passos de difusão é altamente não uniforme e depende do modelo específico. Heurísticas simples falham em generalizar entre arquiteturas diferentes (ex: CNN vs. Transformer) e frequentemente resultam em degradação significativa da qualidade da imagem ou em sobrecarga de memória devido à necessidade de "costurar" (stitching) múltiplos modelos independentes durante a inferência.

2. Metodologia: Diff-ES

O Diff-ES (Diffusion Evolutionary Search) é um framework que automatiza a descoberta do agendamento de esparsidade por estágio, combinando poda estrutural com busca evolutiva.

• Divisão por Estágios: O trajeto de difusão é dividido em $n$ estágios. O objetivo é atribuir um nível de esparsidade ($L_i$) distinto para cada estágio, mantendo uma taxa de esparsidade global fixa ($L_t$).
• Busca Evolutiva (Evolutionary Search):
  • Inicialização: Gera uma população de agendamentos de esparsidade (candidatos).
  • Mutação de Troca de Nível (Level-Switch Mutation): Para explorar o espaço de busca, o algoritmo seleciona dois estágios e redistribui a esparsidade entre eles (aumentando em um, diminuindo no outro), mantendo o orçamento global constante.
  • Avaliação de Aptidão (Fitness): Os candidatos são avaliados gerando um pequeno lote de imagens e calculando métricas leves (como CLIP-IQA, TOPIQ ou SSIM) para medir a fidelidade perceptual e semântica sem o custo de métricas distribucionais pesadas (como FID completo) a cada iteração.
  • Seleção: Os melhores candidatos sobrevivem para a próxima geração, refinando o agendamento ótimo.
• Poda Estrutural de Segunda Ordem (SNR-Aware):
  • Utiliza uma abordagem baseada em OBS (Optimal Brain Surgeon) adaptada para cada estágio.
  • Cria um conjunto de calibração consciente do SNR (Signal-to-Noise Ratio), garantindo que a poda seja feita com dados que refletem as condições de ruído reais de cada estágio da difusão.
  • Calcula a importância das estruturas (camadas/canais) usando a matriz Hessiana local para remover parâmenos redundantes com mínima perda de informação.
• Mecanismo de Roteamento de Pesos (Weight Routing):
  • Inovação Crítica: Para evitar o alto custo de memória e computação de reexecutar a poda de segunda ordem para cada candidato na busca evolutiva, o Diff-ES pré-computa o trajeto completo de poda para cada estágio e armazena os pesos atualizados em um banco de dados compacto.
  • Durante a avaliação, o mecanismo de weight routing recupera dinamicamente os pesos específicos do estágio necessário, montando o modelo "on-the-fly" sem duplicar o modelo inteiro na memória GPU. Isso elimina a necessidade de "stitching" de modelos múltiplos.

3. Principais Contribuições

1. Identificação de Limitação Heurística: Demonstra que agendamentos manuais (como no MosaicDiff) não capturam padrões de esparsidade dependentes do modelo, levando a generalização pobre.
2. Framework de Otimização Automática: Propõe o Diff-ES, que utiliza busca evolutiva para encontrar automaticamente o agendamento de esparsidade por estágio ótimo sob um orçamento global fixo.
3. Eficiência de Memória e Plug-and-Play: Introduz o mecanismo de weight routing, permitindo a integração com métodos de poda estrutural existentes (LayerDrop, Wanda, OBS) sem duplicação de modelos, reduzindo drasticamente o uso de memória GPU.
4. Generalização Arquitetural: O método é validado tanto em modelos baseados em CNN (SDXL) quanto em Transformers (DiT), demonstrando adaptabilidade a diferentes dinâmicas de difusão.

4. Resultados Experimentais

Os experimentos foram realizados no DiT-XL/2 (Transformer) e SDXL-base-1.0 (U-Net/CNN).

• Qualidade de Geração:
  • O Diff-ES supera consistentemente os baselines (MosaicDiff, Diff-Pruning, DeepCache, OBS-Diff) em todas as configurações.
  • No SDXL com 30% de esparsidade, o Diff-ES alcança um FID de 25.87, enquanto o MosaicDiff (com OBS) atinge um FID de 59.09 (degradação severa).
  • No DiT com 50% de esparsidade, o Diff-ES mantém FID de 12.86, superando o MosaicDiff (22.29).
  • Visualmente, o Diff-ES preserva a identidade do objeto, o layout da cena e texturas finas muito melhor que os métodos concorrentes.
• Eficiência e Memória:
  • O Diff-ES oferece aceleração de tempo de parede (wall-clock speedup) significativa.
  • Redução de Memória: Ao usar roteamento de pesos em vez de costura de modelos, o Diff-ES economiza aproximadamente 42.7% de memória GPU em comparação com a abordagem de model stitching do MosaicDiff (5.61 GB vs 9.79 GB no DiT).
• Generalização: O agendamento encontrado não sofre overfitting para o conjunto de prompts de busca; mantém a qualidade em prompts fora da distribuição (out-of-distribution).

5. Significado e Conclusão

O Diff-ES representa um avanço significativo na compressão de modelos de difusão. Ao substituir heurísticas manais por uma otimização automática baseada em busca evolutiva e resolver o gargalo de memória através do roteamento inteligente de pesos, o método permite:

• Aceleração Real: Redução do tempo de inferência sem re-treinamento.
• Fidelidade Preservada: Manutenção da qualidade de geração mesmo sob alta compressão estrutural.
• Versatilidade: Funciona como uma camada complementar a outras técnicas de aceleração (como redução de passos ou cache de características).

O trabalho estabelece um novo estado da arte (state-of-the-art) para poda estrutural em modelos de difusão, demonstrando que a alocação dinâmica de recursos computacionais baseada em dados é superior a regras fixas pré-definidas.