DiffSparse: Accelerating Diffusion Transformers with Learned Token Sparsity

1. Problema

Os modelos de difusão baseados em Transformers (DiTs) alcançaram resultados excepcionais na geração de imagens e vídeos de alta qualidade. No entanto, o mecanismo de inferência multi-passos (iterativo) impõe um custo computacional massivo, dificultando a implantação em cenários do mundo real.

Embora existam métodos anteriores de aceleração baseados em cache de características (feature caching), eles apresentam limitações significativas:

• Alocação Manual de Esparsidade: Métodos existentes exigem que os pesquisadores definam manualmente taxas de reutilização de tokens para cada camada e passo de tempo, o que é ineficiente e difícil de escalar.
• Dependência de Passos Completos: Para manter a qualidade da geração, muitas técnicas de cache de tokens exigem a preservação de vários passos de inferência completos (sem cache), o que limita o potencial de aceleração.
• Estratégias de Cache Ineficientes: O cache de características em nível de camada (coarse-grained) ou estratégias fixas não exploram adequadamente a redundância temporal em nível de token.

2. Metodologia: DiffSparse

O DiffSparse é um framework diferenciável e aprendível que otimiza a alocação de esparsidade de tokens em nível de camada para modelos de difusão Transformer. A abordagem trata o problema de aceleração como um problema de otimização de esparsidade end-to-end, sem a necessidade de heurísticas manuais.

O método é composto por três componentes principais:

A. Seletor de Tokens (Token Selector)

• Atribui uma pontuação de importância a cada token para decidir quais devem ser recomputados e quais podem ser reutilizados do cache.
• A pontuação é uma combinação de critérios como influência da atenção auto (self-attention), foco da atenção cruzada (cross-attention), frequência de reutilização do cache e distribuição espacial uniforme.
• Com base na taxa de esparsidade definida, os tokens com as maiores pontuações são recomputados, enquanto os demais reutilizam os valores do cache.

B. Preditor de Custo de Esparsidade Aprendível (Learnable Sparsity Cost Predictor)

• É uma rede treinável que gera uma matriz de custos quantificando o "custo" (degradação de qualidade) de aplicar diferentes taxas de esparsidade em cada camada e passo de tempo.
• O preditor é independente do comprimento da sequência de tokens (N), dependendo apenas do número de passos (T), camadas (L) e configurações candidatas de esparsidade (|S|), permitindo generalização para diferentes resoluções.

C. Solucionador de Programação Dinâmica (Dynamic Programming Solver)

• Utiliza a matriz de custos gerada pelo preditor para encontrar a configuração de esparsidade ótima global.
• O algoritmo minimiza o custo total (preservando a qualidade) sujeito a uma restrição global de taxa de esparsidade (ex: reduzir 54% dos cálculos).
• Utiliza o Estimador Straight-Through (STE) para permitir o fluxo de gradientes através das máscaras binárias discretas durante o treinamento.

D. Estratégia de Treinamento em Duas Etapas

Para eliminar a necessidade de passos completos pré-definidos:

1. Etapa 1: Otimiza independentemente as matrizes de custo para passos completos e esparsidade de camada. Identifica os passos críticos que devem permanecer completos.
2. Etapa 2: Integra as informações da Etapa 1 para refinar a alocação de esparsidade em todas as camadas e passos, permitindo que o modelo aprenda a distribuir os FLOPs (operações de ponto flutuante) de forma adaptativa, removendo a dependência de passos completos fixos.

3. Contribuições Chave

1. Otimização Diferenciável de Esparsidade: Propõe o primeiro framework que aprende automaticamente a alocação de esparsidade de tokens em nível de camada e passo de tempo, eliminando a necessidade de ajuste manual de hiperparâmetros de cache.
2. Eliminação de Passos Completos Fixos: A estratégia de treinamento em duas etapas remove a necessidade de manter passos de inferência completos em momentos específicos, desbloqueando todo o potencial de aceleração do cache em nível de token.
3. Generalização e Escalabilidade: O preditor de custos treinado em baixa resolução (ex: 256x256) transfere-se eficazmente para resoluções mais altas (ex: 512x512) sem necessidade de retreinamento, reduzindo o custo computacional de treinamento.

4. Resultados Experimentais

Os autores avaliaram o DiffSparse em diversos modelos base (DiT-XL/2, PixArt-α, FLUX, Wan2.1) e tarefas (texto-para-imagem, texto-para-vídeo).

• PixArt-α (20 passos):
  • Redução de 54% no custo computacional (MACs).
  • Aceleração de 1.91x.
  • FID de 27.79, superando o modelo original (FID 28.20) e métodos SOTA como ToCa (28.35) e TaylorSeer (29.08).
  • Nota: O método não apenas acelerou, mas melhorou a qualidade da geração em comparação ao modelo base.
• DiT-XL/2 (Class-conditional):
  • Aceleração de 2.07x com FID de 2.81, superando o ToCa (FID 3.05) na mesma taxa de aceleração.
• Wan2.1 (Texto-para-vídeo):
  • Aceleração de 2.05x com pontuação VBench de 43.83, superando o DuCa e mantendo a qualidade.
• Eficiência de Treinamento: O método de otimização diferenciável leva cerca de 4 horas para treinar, enquanto métodos de busca (como algoritmos genéticos) levam ~16 horas e produzem resultados inferiores.

5. Significado e Impacto

O DiffSparse representa um avanço significativo na eficiência de modelos generativos de grande escala. Ao transformar a alocação de esparsidade em um problema de otimização aprendível e end-to-end, o método supera as limitações das abordagens baseadas em cache estático ou manual.

A capacidade de acelerar a inferência em quase 2x sem degradar (e muitas vezes melhorando) a qualidade da imagem torna os modelos de difusão Transformer viáveis para aplicações em tempo real e dispositivos com recursos limitados. Além disso, a estratégia de treinamento em duas etapas e a generalização entre resoluções oferecem um caminho escalável para a otimização de futuros modelos de IA generativa.