Resumo Técnico: Alocação Heterogênea de Passos (HSA) para Geração Eficiente de Vídeo

1. Declaração do Problema

Os Transformadores de Difusão (DiTs) estabeleceram desempenho de última geração na geração de vídeo, mas sofrem com custos computacionais imensos. Protocolos de inferência padrão aplicam um número uniforme de passos de remoção de ruído ($T$) a cada token espaço-temporal na sequência. Essa abordagem monolítica ignora a assimetria inerente dos dados visuais: regiões de fundo, patches estáticos e estruturas grosseiras são perceptualmente mais simples e requerem menos passos de remoção de ruído do que objetos ricos em detalhes no primeiro plano ou regiões de alto movimento. Além disso, a visão humana é menos sensível à degradação de qualidade em componentes altamente dinâmicos em comparação com os estáticos. Consequentemente, aplicar um orçamento completo de remoção de ruído uniformemente a todos os tokens é computacionalmente desperdiçado e falha em alinhar-se com prioridades perceptuais.

Métodos de aceleração existentes, como cache de características em nível de passo (por exemplo, TeaCache), tipicamente tomam decisões globais: em qualquer iteração dada, ou todos os tokens são computados ou todos os tokens reutilizam saídas em cache. Esses métodos não exploram a heterogeneidade espacial ou em nível de token, tratando a sequência como um bloco monolítico em vez de uma coleção de tokens com requisitos dinâmicos variados.

2. Metodologia: Alocação Heterogênea de Passos (HSA)

Os autores propõem a Alocação Heterogênea de Passos (HSA), um algoritmo de inferência sem treinamento e agnóstico ao modelo que atribui trajetórias de remoção de ruído específicas a tokens dentro de um DiT pré-treinado. A ideia central é particionar a sequência de tokens em grupos com base em suas dinâmicas de velocidade, alocando menos passos a tokens de evolução mais lenta, enquanto mantém um cronograma completo para um grupo de referência.

2.1 Particionamento e Agendamento de Tokens

O conjunto de $N$ tokens é particionado em $K$ grupos disjuntos ($G_1, \dots, G_K$). Cada grupo $G_k$ recebe um orçamento de passos $s_k$, que deve ser um divisor do total de passos $T$.

• Grupo de Referência ($G_1$): Recebe o orçamento completo $s_1 = T$. Esses tokens estão ativos em cada iteração.
• Grupos Reduzidos: Recebem orçamentos menores (por exemplo, $T/2, T/4$). Esses tokens estão ativos apenas quando a iteração global $i$ satisfaz $i \pmod{T/s_k} = 0$.
• Seleção Dinâmica: O artigo avalia quatro estratégias para atribuir tokens a grupos. A estratégia de Seleção Dinâmica de Tokens, que atribui tokens com base na mudança relativa L1 por token das previsões de velocidade (identificando regiões de evolução lenta como fundos), consistentemente produz o melhor desempenho.

2.2 Sincronização de Cache KV

Um desafio primário no agendamento heterogêneo é que tokens com diferentes contagens de passos existem em diferentes níveis de ruído ($\sigma$) em qualquer iteração de relógio. A atenção auto-regressiva exige que todos os tokens atendam uns aos outros, criando uma incompatibilidade de sincronização. A HSA resolve isso por meio de um mecanismo de sincronização de cache KV:

1. Tokens Ativos: Apenas o subconjunto de tokens ativos ($A_i$) realiza uma passagem direta através das projeções QKV.
2. Atualização de Cache: Projeções frescas de Chave ($K$) e Valor ($V$) para tokens ativos são escritas em um cache por camada, sobrescrevendo entradas desatualizadas.
3. Atenção de Contexto Completo: Tokens ativos computam consultas contra as matrizes $K, V$ em cache completas (contendo tanto entradas frescas para tokens ativos quanto entradas desatualizadas, mas válidas, para tokens pulados).
4. Redução de Complexidade: Isso reduz a complexidade da atenção auto-regressiva de $O(N^2)$ para $O(|A_i| \cdot N)$, onde $|A_i|$ é o número de tokens ativos no passo $i$. Tokens pulados não incorrem em computação de consulta, atenção cruzada ou operações feed-forward durante iterações inativas.

2.3 Atualização de Euler em Cache

Para avançar os estados latentes de tokens pulados sem avaliações adicionais do modelo, a HSA emprega uma atualização de Euler em cache.

• Cada token armazena a velocidade ($\hat{v}_n$) prevista em seu passo ativo mais recente.
• Em cada iteração global, todos os tokens (ativos e pulados) são avançados pelo mesmo passo incremental de Euler $(\sigma_{i+1} - \sigma_i)$.
• Tokens ativos usam sua velocidade recém-computada; tokens pulados reutilizam sua velocidade em cache.
• Essa atualização é formulada como uma única operação tensorial sobre todos os $N$ tokens, evitando ramificações e mantendo a eficiência da GPU.

2.4 Janela de Cache Consciente de Fase

Reconhecendo que os estágios iniciais e tardios de remoção de ruído são críticos para a estrutura global e detalhes finos, respectivamente, a HSA restringe o cache à porção média da trajetória. O primeiro e o último $m$ passos executam passagens diretas completas e sem cache para todos os tokens, garantindo que a integridade estrutural e o refinamento de detalhes não sejam comprometidos por aproximações desatualizadas.

3. Contribuições Principais

1. Algoritmo HSA: Um paradigma de inferência totalmente sem treinamento, plug-and-play, que atribui trajetórias de remoção de ruído heterogêneas a grupos de tokens espaço-temporais com base em dinâmicas de velocidade.
2. Sincronização de Cache KV: Um mecanismo que preserva o contexto completo de atenção de $N$ tokens para tokens ativos enquanto contorna inteiramente tokens inativos, reduzindo o custo de atenção por iteração para $O(|A_i| \cdot N)$.
3. Atualização de Euler em Cache: Uma operação tensorial simples que avança tokens pulados usando velocidades em cache, mantendo a qualidade de geração sem avaliações extras do modelo.
4. Avaliação Abrangente: Uma análise de quatro estratégias de agrupamento de tokens e quatro predefinições de orçamento, demonstrando robustez em diferentes orçamentos de inferência.

4. Resultados Experimentais

Os autores avaliaram a HSA em Wan-2.1-1.3B e Wan-2.2-A14B (Texto-para-Vídeo e Imagem-para-Vídeo) e LTX-2 (Áudio-Vídeo).

• Compromisso Qualidade-Tempo de Execução: A HSA alcança uma fronteira de Pareto superior em comparação com o Fluxo de Correspondência (FM) uniforme e métodos de cache de última geração como o TeaCache.
• Aceleração Agressiva: As vantagens são mais pronunciadas em regimes de aceleração agressiva (50% e 25% do tempo de execução). Em 25% do tempo de execução, a HSA supera significativamente as linhas de base no benchmark VBench (Pontuação Total: 79,87% vs. 75,68% para FM e 78,33% para TeaCache).
• Robustez Dimensional: Enquanto métodos de base sofrem com "colapso dimensional catastrófico" (degradação severa em dimensões específicas do VBench, como suavidade de movimento ou consistência temporal) sob orçamentos apertados, a HSA rastreia robustamente a referência de orçamento completo em todas as 16 dimensões de avaliação.
• Imagem-para-Vídeo (I2V): A HSA preserva com sucesso o alinhamento de condicionamento de imagem e a fidelidade visual em tarefas I2V, whereas linhas de base FM uniformes em passos reduzidos ($T=10$) exibem colapso estrutural severo.
• Sem Perfilamento Offline: Ao contrário do TeaCache, que requer perfilamento offline caro para estimar flutuações de saída do modelo, a HSA opera sem nenhum ajuste offline ou ajuste fino.

5. Significado e Alegações

O artigo afirma que a HSA aborda uma ineficiência fundamental na inferência de DiTs ao alinhar recursos computacionais com a complexidade perceptual e dinâmica dos dados visuais. Ao tratar tokens individualmente em vez de globalmente, a HSA alcança acelerações significativas sem a degradação de qualidade tipicamente associada à aceleração agressiva.

Os autores enfatizam que a HSA é um framework flexível que pode ser instanciado com várias estratégias de agrupamento de tokens. Embora a seleção dinâmica atualmente produza os resultados mais fortes, o framework antecipa melhorias futuras com estratégias de agrupamento mais sofisticadas. O trabalho estabelece as bases para uma nova classe de algoritmos de inferência que exploram a heterogeneidade inerente de dados espaço-temporais para democratizar a geração de vídeo de alta fidelidade.