Resumo Técnico: FlashBlock

1. Declaração do Problema

Embora a difusão em blocos tenha melhorado a eficiência de inferência em modelos generativos de contexto longo (por exemplo, modelos de linguagem de difusão e geração de vídeo) ao combinar o cache KV com a inferência causal por blocos, ela ainda sofre com um overhead significativo em configurações de contexto longo. Na difusão em blocos padrão, cada etapa de difusão exige a recomputação da atenção sobre todo o cache KV crescente. À medida que o comprimento do contexto aumenta, o custo combinado da computação de atenção e do acesso ao cache KV torna-se o fator dominante na latência de inferência.

O artigo identifica uma propriedade subexplorada da difusção em blocos: a redundância entre etapas (cross-step redundancy). Especificamente, durante o processo de denoising, os tokens fora do bloco atual sendo atualizado permanecem amplamente estáticos, levando a contribuições de atenção altamente estáveis desses tokens "externos ao bloco". Por outro lado, a atenção dentro do bloco atual ("interno ao bloco") varia significativamente conforme os tokens são ativamente refinados ou desmascarados. Os métodos atuais falham em explorar essa estabilidade, recomputando repetidamente a atenção externa ao bloco que é redundante.

2. Metodologia: FlashBlock

Os autores propõem o FlashBlock, um mecanismo que armazena em cache e reutiliza as saídas de atenção externa ao bloco para reduzir a computação e o I/O de memória sem modificar o processo de difusão subjacente.

Mecanismo Central

O FlashBlock decompõe a saída total da atenção em dois componentes:

1. Atenção Interna ao Bloco ($A_{in}$): Computada sobre os tokens dentro do bloco atual ($J_{in}$), que estão mudando ativamente.
2. Atenção Externa ao Bloco ($A_{out}$): Computada sobre os tokens fora do bloco atual ($J_{out}$), que permanecem estáveis entre as etapas.

O método armazena em cache a saída de atenção externa ao bloco ($A_{out}$) e seu normalizador logarítmico ($L_{out} = \log Z_{out}$) da etapa de difusão anterior. Em etapas subsequentes, em vez de reacessar todo o cache KV para $J_{out}$, o modelo reutiliza os valores armazenados. A saída total da atenção é reconstruída combinando a $A_{out}$ em cache com a recém-computada $A_{in}$ usando composição no espaço logarítmico (semelhante ao FlashAttention) para garantir a estabilidade numérica:
$$A_{full} = \frac{e^{L_{out}-m} A_{out} + e^{L_{in}-m} A_{in}}{e^{L_{out}-m} + e^{L_{in}-m}}$$
onde $m = \max(L_{out}, L_{in})$.

Estratégias de Reutilização Adaptativa

• Reutilização Baseada em Limiar (Threshold-based Reuse): Em modelos de difusão discretos, as atualizações de tokens podem ser abruptas. O método introduz um limiar $\tau$ no número de tokens atualizados ($M^{s+1}$). Se $M^{s+1} < \tau$, a atenção em cache é reutilizada; caso contrário, ela é recomputada para garantir a correção.
• Seleção por Cabeça (Vídeo): Para difusão de vídeo, onde a variabilidade entre etapas é maior, o método estima a similaridade por cabeça de atenção. Apenas as cabeças com similaridade acima de um limiar $\gamma$ reutilizam o cache.
• Compatibilidade com Atenção Esparsa: O FlashBlock é ortogonal à atenção esparsa. Ele trata os keys não selecionados na atenção esparsa como o conjunto "externo ao bloco", armazenando em cache suas contribuições residuais. Isso permite que o FlashBlock compense a perda de informação causada pela esparsificação agressiva.

Adaptação de Treinamento (Opcional)

Para abordar possíveis descasamentos distributivos em modelos de linguagem de difusão discreta (dLLMs) causados pelo desmascaramento abrupto de tokens, os autores propõem uma estratégia de destilação consciente da reutilização (reuse-aware distillation). Um modelo estudante empregando FlashBlock é treinado para corresponder à distribuição de saída de um modelo professor congelado usando atenção densa, utilizando adaptadores LoRA para manter a adaptação leve.

3. Principais Contribuições

1. Insight Empírico: O artigo fornece um estudo empírico revelando uma forte estabilidade entre etapas na atenção externa ao bloco, contrastando com a alta variabilidade da atenção interna ao bloco.
2. Mecanismo FlashBlock: Um novo mecanismo de cache que explora essa redundância para reduzir a computação de atenção e o acesso ao cache KV, implementado via agregação no espaço logarítmico.
3. Sinergia com Esparsidade: Demonstração de que o FlashBlock serve como uma estratégia complementar de reutilização residual para a atenção esparsa, mitigando a degradação de qualidade causada pela esparsificação agressiva.
4. Ganhos de Eficiência: Aceleração significativa no throughput de inferência e redução de latência tanto para modelos de linguagem de difusão quanto para modelos de geração de vídeo, com impacto negligenciável na qualidade da geração.

4. Resultados Experimentais

Experimentos foram conduzidos em modelos de linguagem de difusão (Trado-8B/4B) e modelos de difusão de vídeo (LongLive-1.3B).

• Throughput: O FlashBlock alcançou até 1.44× mais throughput de tokens em modelos de linguagem de difusão (por exemplo, aumentando de 312 para 451 tokens/s para tamanho de bloco 4).
• Latência: O tempo de atenção foi reduzido em até 1.6×.
• Qualidade: A qualidade da geração permaneceu comparável às baselines densas. Em benchmarks de raciocínio matemático (GSM8K, MATH500) e codificação (HumanEval, MBPP), as diferenças de precisão foram mínimas (dentro de 0,2% a 1,6% dependendo do tamanho do bloco).
• Escalabilidade de Contexto Longo: À medida que o comprimento do contexto aumentou de 100k para 800k, a lacuna de latência entre o FlashBlock e a baseline aumentou, com o FlashBlock mostrando aproximadamente metade do crescimento de latência da baseline sob certos limiares.
• Combinação com Atenção Esparsa: Quando combinado com SparseD (em 20%, 30% e 40% de densidade), o FlashBlock consistentemente melhorou a precisão sobre o SparseD sozinho (por exemplo, melhorando a precisão do GSM8K de 34,72% para 42,68% em 20% de densidade), recuperando efetivamente a informação perdida para a esparsificação.

5. Significância

O artigo afirma que o FlashBlock aborda um gargalo crítico na difusão em blocos de contexto longo: a recomputação redundante de contribuições de atenção estáveis. Ao desacoplar o contexto "externo ao bloco" estável das atualizações "internas ao bloco" dinâmicas, o método reduz substancialmente os custos de computação e I/O de memória associados aos caches KV crescentes.

Os autores enfatizam que esta abordagem é ortogonal às técnicas de atenção esparsa existentes, oferecendo uma estratégia complementar que melhora tanto a eficiência quanto a precisão. Os resultados sugerem que o FlashBlock fornece uma solução escalável para a geração eficiente de contexto longo em domínios de texto e vídeo, permitindo uma inferência mais rápida sem comprometer a qualidade do conteúdo gerado.