FastCache: Fast Caching for Diffusion Transformer Through Learnable Linear Approximation

Resumo Técnico: FastCache para Aceleração de Diffusion Transformers

1. O Problema

Os Diffusion Transformers (DiT) tornaram-se modelos generativos de ponta para síntese de imagens e vídeos, superando arquiteturas baseadas em U-Net em qualidade. No entanto, eles são computacionalmente intensivos devido à sua estrutura iterativa e pilhas profundas de transformadores.

• Ineficiência: A inferência envolve operações repetidas de transformadores ao longo de múltiplos passos de tempo (timesteps) e locais espaciais.
• Redundância Ignorada: Em muitos cenários (especialmente em regiões de fundo estático ou passos de denoising tardios), os estados ocultos (hidden states) do modelo exibem variações mínimas entre passos consecutivos.
• Limitações Atuais: Métodos existentes de aceleração focam frequentemente no nível de quadro ou token (poda de tokens), ignorando a redundância rica nos estados latentes profundos do transformador. Isso resulta em recomputação desnecessária e alto custo de recursos.

2. Metodologia: FastCache

O FastCache é um framework de cache e compressão no nível de estado oculto (hidden-state-level) projetado para acelerar a inferência de DiTs explorando redundâncias espaciais e temporais. Ele opera sob o princípio: "Cacheie o Fundo, Recompute o Movimento".

O sistema consiste em dois módulos principais e uma extensão para fusão de tokens:

A. Módulo de Redução de Tokens Espaço-Temporais (Spatial-Temporal Token Reduction)

• Saliência Temporal: Calcula a diferença entre o estado oculto atual ($X_t$) e o anterior ($X_{t-1}$).
• Classificação: Separa os tokens em tokens de movimento (variação alta) e tokens estáticos (variação baixa) usando um limiar ($\tau_s$).
• Aproximação Linear: Os tokens estáticos não passam pelos blocos completos do transformador. Em vez disso, sua representação é aproximada por uma camada linear aprendível ($H^s_t = W_c X^s_t + b_c$), que é muito mais barata computacionalmente. Os tokens de movimento são processados normalmente.

B. Módulo de Cache no Nível do Transformador (Transformer-Level Caching)

• Decisão Estatística: Para cada bloco do transformador, o sistema verifica se a mudança relativa no estado oculto em relação ao passo de tempo anterior é estatisticamente insignificante.
• Teste de Hipótese: Utiliza um teste de hipótese baseado em estatística (assumindo distribuição normal e estacionariedade fraca). Se a métrica de mudança relativa ($\delta_{t,l}$) estiver abaixo de um limiar baseado na distribuição Qui-quadrado ($\chi^2$), o bloco é pulado.
• Reutilização: Quando o bloco é pulado, a saída é aproximada linearmente ou reutilizada diretamente, evitando a computação completa do bloco.
• Garantia Teórica: O método fornece um limite superior para o erro de aproximação, garantindo que a qualidade não degrade além de um certo ponto controlado.

C. Módulo de Fusão de Tokens (Token Merging)

• Para reduzir ainda mais a redundância, o FastCache introduz uma fusão de tokens baseada em densidade kNN (k-nearest neighbors).
• Combina métricas de densidade espacial e saliência temporal para identificar tokens redundantes que podem ser fundidos em um único token representativo, reduzindo o comprimento da sequência sem perder detalhes dinâmicos.

3. Contribuições Principais

1. Framework de Cache no Nível de Estado Oculto: Diferente de métodos anteriores que focam em U-Nets ou camadas de atenção específicas, o FastCache opera diretamente nos estados ocultos dos blocos do DiT, permitindo estratégias de aceleração mais finas e agnósticas ao modelo.
2. Estratégia de Poda Consciente de Movimento: Desenvolveu uma estratégia que identifica e suprime tokens espaciais redundantes baseando-se na variação de movimento e conteúdo, preservando apenas os detalhes dinâmicos.
3. Fundamentação Teórica e Interpretabilidade:
   • Fornece uma justificativa teórica rigorosa, enquadrando o cache como uma decomposição de interação esparsa aprendível.
   • Demonstra que o método é equivalente a uma aproximação de Taylor de primeira ordem das dinâmicas do transformador, separando componentes de fundo (estáveis) e movimento (resíduos).
   • Prova que o erro de cache é limitado sob suposições de distribuição e suavidade.

4. Resultados Experimentais

Os experimentos foram realizados em várias variantes de DiT (DiT-XL/2, L/2, B/2, S/2) para geração de imagens e vídeos.

• Desempenho de Velocidade e Memória:
  • Redução de 36,9% no tempo de inferência (latência).
  • Redução de 31,7% no uso de memória.
  • Em comparação com o FBCache (um baseline recente), o FastCache oferece um speedup superior (ex: +42,4% vs +30,0% no DiT-XL/2).
• Qualidade de Geração:
  • O FastCache alcançou os melhores resultados de qualidade entre os métodos de cache, medidos por FID (Fréchet Inception Distance) e t-FID (temporal FID).
  • Exemplo (DiT-XL/2): FID de 4.46 (FastCache) vs 4.48 (FBCache) e 4.64 (AdaCache), mantendo a fidelidade visual e a coerência do movimento.
• Robustez:
  • A ablação mostrou que a combinação de Redução de Tokens (STR), Cache Estatístico (SC) e Mistura Consciente de Movimento (MB) é essencial para o melhor desempenho.
  • O método generaliza bem para diferentes tamanhos de modelos e configurações de vídeo (longas sequências).

5. Significância e Impacto

O FastCache representa um avanço significativo na eficiência de modelos generativos de grande escala:

• Viabilidade em Tempo Real: Ao reduzir drasticamente a latência e o consumo de memória, torna aplicações como edição de vídeo em tempo real, realidade virtual e geração de conteúdo em larga escala mais viáveis.
• Eficiência Sem Perda de Qualidade: Diferente de métodos que sacrificam qualidade pela velocidade, o FastCache demonstra que é possível explorar redundâncias latentes para acelerar a inferência mantendo (ou até melhorando) a fidelidade da geração.
• Novo Paradigma de Otimização: A abordagem de tratar a aceleração como uma decomposição de estados ocultos e uma aproximação linear aprendível abre novas fronteiras para otimização de transformadores, indo além da simples poda de tokens ou quantização.

Em resumo, o FastCache oferece uma solução teoricamente fundamentada e empiricamente superior para o gargalo computacional dos Diffusion Transformers, permitindo que modelos poderosos operem com maior eficiência sem comprometer a qualidade do conteúdo gerado.