Resumo Técnico: CoReDiT

Declaração do Problema

Os Transformadores de Difusão (DiTs) estabeleceram desempenho de última geração na geração de imagens e vídeos ao modelar efetivamente dependências de longo alcance por meio de mecanismos de atenção. No entanto, seu custo computacional escala quadraticamente com o número de tokens, tornando a síntese de alta resolução e a geração de vídeos multiframe proibitivamente caros para plataformas móveis com recursos limitados. Embora uma parte significativa dos tokens no espaço latente corresponda a regiões visualmente redundantes (por exemplo, fundos uniformes), os métodos existentes de poda de tokens enfrentam três desafios principais quando aplicados à inferência de difusão:

1. Localização Eficiente: Identificar tokens de baixa saliência sem incorrer na sobrecarga quadrática de calcular pontuações de atenção global.
2. Preservação Semântica: Manter a continuidade visual e prevenir artefatos (por exemplo, quebras de textura, distorções de limites) quando tokens são pulados, o que é crítico para tarefas generativas, ao contrário das discriminativas.
3. Agendamento Adaptativo: Determinar uma estratégia de poda que se adapte a níveis variáveis de redundância em diferentes blocos do transformador e passos de remoção de ruído.

Metodologia: CoReDiT

O CoReDiT é uma estrutura estruturada de poda de tokens projetada para acelerar DiTs pré-treinados, preservando a fidelidade da saída. Ele opera por meio de três componentes principais:

1. Seleção de Tokens Baseada em Coerência Espacial

Em vez de depender de pontuações de atenção global (que são $O(N^2)$), o CoReDiT utiliza uma pontuação de Coerência Espacial (CE) para estimar a redundância local em tempo linear ($O(N)$).

• Mecanismo: O lattice de tokens é partitionado em grades não sobrepostas. Para cada token, a pontuação de CE é calculada como a média da similaridade cosseno entre o embedding normalizado do token e o embedding normalizado médio de sua grade.
• Lógica: Tokens com altas pontuações de CE são considerados redundantes porque são bem explicados por seus vizinhos espaciais.
• Seleção: Os top-$K$ tokens com as maiores pontuações de CE são selecionados para serem pulados ($X_S$), enquanto os tokens restantes ($X_R$) prosseguem para a camada de Atenção Múltipla-Cabeça (MHSA).

2. Reconstrução de Tokens Guiada por Coerência

Para prevenir as descontinuidades espaciais causadas pela pulagem de tokens, o CoReDiT reconstrói as saídas de atenção para os tokens pulados ($Y_S$) usando as saídas dos vizinhos retidos ($Y_R$).

• Aproximação: Em vez de calcular similaridades par a par para reconstrução (o que adiciona sobrecarga), o método aproxima os pesos de similaridade usando as pontuações de CE pré-calculadas dos vizinhos.
• Sub-grades Locais: Para evitar suavização excessiva, a reconstrução é realizada dentro de sub-grades menores.
• Micro-Designs:
  • Tamanhos de Grade Alternados: As dimensões da grade alternam entre blocos consecutivos para deslocar artefatos de limite e permitir a propagação de informações através das fronteiras da grade.
  • Retenção com Passo $m$: Um padrão de passo determinístico garante que pelo menos um token seja retido em cada sub-grade, fornecendo âncoras necessárias para a reconstrução.

3. Agendamento de Poda Progressivo e Adaptativo ao Bloco

Reconhecendo que a redundância varia entre blocos do transformador e passos de remoção de ruído, o CoReDiT emprega uma estratégia de poda progressiva durante o ajuste fino leve.

• Adaptabilidade ao Bloco: O orçamento de poda é alocado aos blocos com a maior redundância medida (estimada pela soma das pontuações de CE dos tokens candidatos).
• Decaimento por Passo de Tempo: A poda é mais agressiva nos primeiros passos de remoção de ruído (onde o ruído domina e a estrutura é grosseira) e reduzida nos passos posteriores (onde detalhes finos são refinados). Um agendamento de decaimento em duas fases é usado para equilibrar a eficiência com a complexidade de gerenciar grafos de computação específicos por passo.
• Treinamento: O modelo é ajustado finamente usando uma perda de destilação que combina a perda de saída do modelo original e uma perda de bloco podado para recuperar o desempenho.

Contribuições Principais

1. Estimativa de Redundância em Tempo Linear: A introdução de uma pontuação de coerência espacial que identifica tokens redundantes em tempo $O(N)$, evitando o gargalo computacional da análise de atenção global.
2. Preservação de Lattice Denso: Um mecanismo de reconstrução que sintetiza as saídas de atenção puladas a partir dos vizinhos, garantindo que o lattice de tokens permaneça denso para camadas subsequentes e prevenindo artefatos visuais.
3. Estratégia de Poda Adaptativa: Um agendamento progressivo e adaptativo ao bloco que aloca dinamicamente orçamentos de poda com base em estatísticas empíricas de redundância através da profundidade e do tempo, facilitando uma adaptação estável.

Resultados Experimentais

O CoReDiT foi avaliado em backbones de DiT de última geração, incluindo PixArt-α, PixArt-Σ e MagicDrive-V2.

• Ganhos de Eficiência:
  • Alcançou até 55% de redução em FLOPs de auto-atenção.
  • Acelerações de Inferência: 1,33× em GPUs de nuvem (Nvidia H100) e 1,72× em NPUs móveis (Qualcomm Snapdragon 8 Elite).
  • Margem de Memória: Permitiu geração em resoluções de até 1920 em dispositivos móveis, onde o modelo de base encontrou erros de Memória Insuficiente (OOM) em 1600.
• Preservação de Qualidade:
  • No PixArt-α-1024 com 40% de poda, o CoReDiT manteve um FID de 28,7 (vs. base 27,3) e pontuações competitivas de CLIP/IS. Com destilação, a lacuna de FID estreitou para 27,4.
  • No PixArt-Σ-2048, uma taxa média de poda de 23% reduziu a latência de auto-atenção em 22% com apenas uma queda modesta nas métricas de qualidade.
  • No MagicDrive-V2 (vídeo), uma redução de 39% em FLOPs de auto-atenção preservou a qualidade por quadro (PSNR, LPIPS, SSIM) e o alinhamento condicional (mAP, mIoU), embora uma leve queda na consistência temporal (FVD) tenha sido notada.
• Estudos de Ablação:
  • A seleção aleatória de tokens resultou em degradação catastrófica da qualidade (FID 644,5), validando a necessidade de seleção baseada em coerência.
  • A desativação da reconstrução levou a uma perda significativa de qualidade (FID 30,4 vs. 28,5), confirmando a importância da etapa de reconstrução.
  • A poda uniforme seguida de ajuste fino performou pior do que o agendamento adaptativo progressivo proposto.

Significado e Alegações

O artigo afirma que o CoReDiT oferece uma solução prática para implantar modelos de difusão de alta fidelidade em dispositivos de borda, abordando os desafios específicos da poda de tokens generativos. Seu significado reside em:

• Escalabilidade: Permitir síntese de maior resolução e geração de vídeo em hardware móvel, reduzindo tanto a pegada de computação quanto a de memória.
• Compatibilidade Arquitetural: Manter a compatibilidade com arquiteturas DiT existentes sem exigir mudanças estruturais ou módulos de roteamento complexos.
• Compromisso Eficiência-Qualidade: Demonstrar que a pulagem estruturada e guiada por coerência pode reduzir drasticamente os custos de inferência (até 55% de FLOPs de auto-atenção) enquanto mantém qualidade perceptual comparável ao modelo completo.

Os autores observam que, embora a abordagem atual foque na redundância espacial, trabalhos futuros visam estender a eficiência para componentes FFN/MLP e incorporar objetivos de coerência entre quadros para melhorar ainda mais a consistência temporal de vídeos.