Sink-Aware Pruning for Diffusion Language Models

1. O Problema

Os Modelos de Linguagem de Difusão (DLMs) representam uma alternativa promissora aos modelos autoregressivos (AR), permitindo a geração de texto através de um processo iterativo de remoção de ruído (denoising). No entanto, essa abordagem impõe custos computacionais e de memória significativamente mais altos devido à necessidade de múltiplas etapas de inferência iterativas.

Para viabilizar a implantação prática, técnicas de poda (pruning) de modelos são essenciais. O problema central identificado pelos autores é que as heurísticas de poda existentes foram herdadas quase inteiramente de modelos AR. Uma dessas heurísticas fundamentais é a preservação de "Attention Sinks" (Sumidouros de Atenção).

• No contexto AR: Os sinks (geralmente tokens iniciais como o início da sequência ou prompts de sistema) atuam como âncoras globais estáveis que absorvem uma massa desproporcional de atenção. A prática padrão é preservá-los para evitar degradação catastrófica.
• A Falha em DLMs: Os autores demonstram que essa suposição não se sustenta em modelos de difusão. Em DLMs, a posição do sink exibe uma variância substancialmente maior ao longo da trajetória de geração. Diferente dos modelos AR, onde o sink é estático, em DLMs os sinks são frequentemente transientes, mudando de posição à medida que o processo de denoising avança (de estrutura global para refinamento local). Preservar rigidamente esses sinks instáveis em DLMs leva a uma poda ineficiente e subótima.

2. Metodologia: Sink-Aware Pruning

Os autores propõem uma nova estratégia chamada Sink-Aware Pruning (Poda Consciente de Sumidouros), que adapta o processo de poda especificamente para a dinâmica de geração de difusão.

Principais Etapas da Metodologia:

1. Análise de Variância de Sinks:

   • Os autores definem duas estatísticas para medir a estabilidade dos sinks: Variância Espacial (quão concentrada é a atenção) e Variância Temporal (quanto a posição do sink muda entre as etapas de geração).
   • Eles constatam que, enquanto modelos AR têm baixa variância temporal (sinks fixos), modelos DLMs têm alta variância temporal (sinks que migram).

2. Identificação de Sinks Instáveis:

   • Em vez de presumir que todos os sinks devem ser preservados, o método calcula a "massa de atenção" recebida por cada token em cada etapa de difusão.
   • Identifica-se quais sinks são instáveis (aqueles cuja posição ou importância flutua significativamente ao longo das etapas de denoising).

3. Reponderação de Importância (Sink-Aware Reweighting):

   • O método introduz um fator de desvalorização suave ($\omega$) para os tokens identificados como sinks instáveis.
   • A ativação original $X$ é modificada para $\tilde{X}$, onde as contribuições dos sinks instáveis são suprimidas ($\tilde{X}_j = \omega_j \cdot X_j$).
   • Essa nova ativação mascarada é então utilizada nos critérios de poda padrão (como Wanda ou SparseGPT) para calcular as pontuações de importância.

4. Execução da Poda:

   • A poda é realizada com base nas pontuações recalculadas. Ao suprimir a influência dos sinks instáveis nas estatísticas de segunda ordem (como a Hessiana no SparseGPT) ou nas normas de ativação (no Wanda), o algoritmo prioriza a preservação de caminhos de atenção semanticamente significativos e estáveis, em vez de manter pesos que apenas sustentam ruído ou atenção transitória.

3. Principais Contribuições

• Descoberta de Variância de Sink: Demonstração empírica de que a variância da posição do sink em DLMs é ordens de magnitude maior do que em modelos AR, invalidando a heurística de "sempre manter sinks".
• Novo Paradigma de Poda: Proposta do Sink-Aware Pruning, uma estratégia que distingue âncoras estáveis de sinks efêmeros, permitindo a poda de sinks instáveis em DLMs (algo que seria prejudicial em modelos AR).
• Eficiência sem Retreinamento: O método é aplicado no pós-treinamento (post-training), não exigindo re-treinamento do modelo, e supera bases fortes de poda existentes sob condições de computação equivalente.

4. Resultados Experimentais

Os autores avaliaram o método em vários modelos DLMs (LLaDA, Dream, MMaDA) e benchmarks padrão (MMLU, ARC-C, GSM8K, etc.), comparando com as bases Wanda e SparseGPT.

• Desempenho Geral: O Sink-Aware Pruning consistentemente igualou ou superou as bases de poda em todos os níveis de esparsidade (25%, 50%, 75%).
• Ganhos em Alta Esparsidade: As melhorias foram mais pronunciadas em níveis de compressão agressivos (50% a 75%). Nesses cenários, a poda tradicional falha ao remover pesos críticos que servem a sinks instáveis, enquanto o método proposto preserva melhor a utilidade do modelo.
• Poda Estruturada: Em testes de poda estruturada (remoção de cabeças ou camadas inteiras), o método também superou as bases, com margens de melhoria crescendo conforme a taxa de poda aumentava.
• Análise Visual: As visualizações mostraram que o método altera as decisões de poda em cabeças específicas, preservando mais pesos em cabeças onde a dinâmica de atenção é crítica e removendo mais em cabeças dominadas por sinks transitórios.

5. Significado e Conclusão

O trabalho estabelece que as regras de poda não devem ser transferidas cegamente entre paradigmas de geração (AR vs. Difusão).

• Mudança de Paradigma: A utilidade de um sink depende da dinâmica de geração. Em modelos AR, eles são essenciais e estáveis; em DLMs, são frequentemente transitórios e podem ser podados.
• Impacto Prático: A proposta oferece uma rota principiante e eficaz para acelerar a inferência de DLMs, melhorando o compromisso entre qualidade e eficiência (quality-efficiency trade-off).
• Futuro: Os autores sugerem que futuras pesquisas podem explorar políticas adaptativas por camada e otimização conjunta com quantização para levar ainda mais longe a fronteira de eficiência dos modelos de difusão.

Em resumo, o artigo fornece uma correção fundamental para a compressão de modelos de linguagem baseados em difusão, demonstrando que a supressão inteligente de sinks instáveis é a chave para uma poda eficiente nesses modelos.