Resumo Técnico: Proxies Singulares para Cache Adaptativo em Modelos de Linguagem de Difusão

Declaração do Problema

Os Modelos de Linguagem de Difusão (DLMs) oferecem uma alternativa flexível aos modelos autoregressivos (AR) ao permitir decodificação de ordem arbitrária e geração paralela. No entanto, sua natureza não causal e a dependência de atenção bidirecional tornam incompatível o cache padrão de Chave-Valor (KV), pedra angular da eficiência dos modelos AR. Nos DLMs, os estados ocultos não podem ser simplesmente reutilizados sem recálculo para capturar dependências globais em evolução, levando a uma latência de inferência proibitiva (complexidade $O(T \times N^2)$).

As tentativas existentes de mitigar isso via cache enfrentam duas limitações primárias:

1. Identificação Custosa: Os métodos atuais dependem de cálculos de similaridade de alta dimensão (por exemplo, monitoramento do desvio do estado de Valor) para identificar quais tokens requerem atualizações. Essa sobrecarga de identificação frequentemente compensa os ganhos das atualizações esparsas à medida que os comprimentos das sequências aumentam.
2. Alocação Rígida de Orçamento: As abordagens existentes tipicamente aplicam um orçamento de atualização uniforme em todas as camadas do modelo. Isso falha em contabilizar as dinâmicas heterogêneas dos estados ocultos, onde algumas camadas são estáveis (requerendo poucas atualizações) e outras são voláteis (requerendo atualizações frequentes), levando tanto a redundância computacional quanto à degradação de desempenho.

Metodologia: SPA-Cache

Os autores propõem o SPA-Cache (Proxies Singulares para Cache Adaptativo), um framework que otimiza conjuntamente a identificação de atualizações e a alocação de orçamento. O fluxo de trabalho para uma única camada Transformer é dividido em três fases:

1. Identificação de Atualização via Proxies Singulares

Em vez de calcular similaridade no espaço completo de alta dimensão ($d$), o SPA-Cache introduz um proxy singular para identificar tokens críticos para atualização em um subespaço de baixa dimensão ($r \ll d$).

• Base Teórica: Os autores provam que a estabilidade dos estados de Valor é uma condição suficiente para a estabilidade das saídas de atenção e das subsequentes saídas da Rede Feed-Forward (FFN). Assim, monitorar os estados de Valor fornece uma garantia conservadora para a estabilidade da saída.
• Redução de Dimensionalidade: Para evitar o gargalo computacional da projeção de alta dimensão, o método aplica a Decomposição em Valores Singulares (SVD) à matriz de projeção de Valor $W$. Ele retém apenas os $r$ maiores vetores e valores singulares para criar uma projeção truncada $W_r$.
• Mecanismo: Os estados de entrada são projetados neste manifold de baixa dimensão. A similaridade cosseno é calculada entre o proxy singular atual e o proxy em cache. Os tokens com a menor similaridade (maior desvio) são selecionados para recálculo.
• Eficiência: Isso reduz a complexidade de identificação de $O(d)$ para $O(r)$ e a sobrecarga de projeção de $O(d^3)$ para $O(rd^2)$.

2. Alocação Adaptativa de Orçamento

Reconhecendo que o desvio do estado oculto varia significativamente entre as camadas do modelo, o SPA-Cache substitui as taxas de atualização estáticas e uniformes por uma estratégia dinâmica e adaptativa por camada.

• Observação: O desvio é baixo nas camadas iniciais (transformações de embedding), atinge o pico nas camadas intermediárias (transformações semânticas complexas) e estabiliza nas camadas finais.
• Implementação: A taxa de atualização $\rho(l)$ para a camada $l$ é parametrizada usando uma função Gaussiana por partes. Isso aloca um orçamento de atualização mais alto para as camadas intermediárias voláteis e faz o cache agressivo das camadas iniciais e finais estáveis, maximizando a vazão sem sacrificar a qualidade.

3. Execução Esparsa

Os tokens selecionados (índices $I$) são usados para gerar Consultas, Chaves e Valores esparsos. O mecanismo de Atenção calcula as saídas usando essas consultas esparsas contra um cache KV parcialmente atualizado. A Rede Feed-Forward (FFN) processa apenas as saídas de atenção esparsas, e a saída da camada final é atualizada apenas nos índices selecionados, deixando outras posições para reutilizar recursos em cache.

Contribuições Principais

1. Justificativa Teórica para Proxies de Valor: O artigo fornece provas formais (Teoremas 3.1 e 3.2) demonstrando que a similaridade do estado de Valor limita a divergência das saídas de atenção e FFN, validando o uso dos estados de Valor como um identificador principiado para seleção de atualização.
2. Proxy Singular para Identificação de Baixa Sobrecarga: Ao aproveitar a SVD truncada, os autores introduzem um método para identificar tokens críticos em um subespaço de baixa dimensão, reduzindo significativamente o custo computacional do próprio processo de identificação.
3. Alocação Adaptativa de Orçamento: A introdução de uma estratégia de atualização dinâmica que alinha recursos computacionais com a heterogeneidade por camada, evitando as ineficiências da alocação uniforme.
4. Ortogonalidade à Decodificação Paralela: O SPA-Cache é projetado para ser ortogonal a otimizações em nível de sequência e pode ser integrado com frameworks de decodificação paralela para amplificar os ganhos de vazão.

Resultados Experimentais

Os autores avaliaram o SPA-Cache em dois DLMs (LLaDA-8B e Dream-v0-Instruct-7B) em sete benchmarks (incluindo GSM8K, MMLU-PRO, MBPP e HumanEval).

• Melhorias de Vazão: O SPA-Cache alcança até uma melhoria de 8× na vazão sobre a decodificação padrão e uma aceleração de 2–4× sobre bases de cache existentes (dLLM-Cache, Fast-dLLM).
• Benchmarks Específicos:
  • No MBPP com LLaDA-8B, o SPA-Cache alcançou uma aceleração de 8,0×.
  • No MMLU-PRO, alcançou acelerações de 6,0× (LLaDA-8B) e 6,2× (Dream-7B).
• Integração com Decodificação Paralela: Quando combinado com métodos de decodificação paralela, o SPA-Cache produziu uma aceleração total de 28× sobre a decodificação padrão no benchmark BBH, mantendo precisão comparável.
• Estudos de Ablação:
  • Substituir identificadores de Valor de dimensão completa pelo proxy singular ($r=128$) melhorou os Tokens por Segundo (TPS) de ~165 para ~179 sem perda de precisão.
  • A alocação adaptativa de orçamento aumentou ainda mais o TPS para ~189, reduzindo a taxa média de atualização de 25% para 16%, enquanto uma taxa uniforme de 16% causou degradação de precisão.

Significado e Alegações

O artigo alega que o SPA-Cache estabelece uma nova linha de base de eficiência para Modelos de Linguagem de Difusão, posicionando-os como uma alternativa computacionalmente viável para aplicações de alta vazão no mundo real. Ao abordar fundamentalmente o gargalo de inferência através de análise teórica principiada e alocação adaptativa de recursos, o método permite que os DLMs escalem efetivamente sem comprometer a fidelidade de geração. Os autores enfatizam que sua abordagem é livre de treinamento e compatível com estratégias existentes de decodificação paralela, abrindo caminho para a implantação prática de modelos de difusão discretos.