Resumo Técnico: STaR-Quant

Declaração do Problema

Os Modelos de Linguagem de Difusão (DLLMs) emergiram como uma alternativa promissora aos LLMs autorregressivos, utilizando denoising (desruído) mascarado iterativo com contexto bidirecional para gerar texto. No entanto, sua implantação é dificultada por sobrecargas substanciais de memória e computação decorrentes dos grandes tamanhos de modelo e do processo de denoising iterativo de múltiplas etapas. Embora a Quantização Pós-Treinamento (PTQ) tenha quantizado com sucesso os LLMs autorregressivos, a aplicação direta dos métodos de PTQ existentes aos DLLMs resulta em uma degradação significativa de desempenho, particularmente sob quantização agressiva de pesos e ativações de baixa bitagem.

O artigo identifica dois desafios específicos únicos aos DLLMs que causam essa falha:

1. Disparidade de Ativação Dependente de Estado: Dentro de uma única etapa de denoising, a sequência contém tanto tokens mascarados (posições incertas a serem previstas) quanto tokens não mascarados (contexto visível). Esses dois estados exibem distribuições de ativação e padrões de outliers distintos. Aplicar uma transformação de ativação ou um espaço de quantização unificado para ambos os estados falha em suavizar as distribuições de forma eficaz para qualquer um deles.
2. Acúmulo de Erro Temporal: A geração de DLLM procede através de etapas iterativas onde a saída de uma etapa serve como entrada para a próxima. Os erros de quantização, particularmente na multiplicação de matrizes entre o mapa de atenção softmax e os estados de valor, não são isolados a uma única etapa, mas acumulam-se ao longo da trajetória de denoising, levando a um desvio significativo do caminho de precisão total.

Metodologia: STaR-Quant

Para enfrentar esses desafios, os autores propõem o STaR-Quant, um framework de PTQ consistente em estado-tempo projetado especificamente para DLLMs. O framework opera em duas dimensões:

1. Transformação de Ativação Guiada por Estado (SGAT)

O SGAT aborda a disparidade de ativação entre tokens mascarados e não mascarados.

• Mecanismo: Em vez de uma única transformação para todos os tokens, o SGAT decompõe a dimensão oculta em subespaços compartilhados e específicos de estado.
  • Subespaço Compartilhado ($U_{sh}$): Captura componentes de ativação comuns.
  • Subespaços Específicos de Estado ($U_m, U_u$): Absorvem padrões de outliers específicos para tokens mascarados e não mascarados, respectivamente.
• Implementação: Portões de subespaço binários roteiam os tokens mascarados e não mascarados para seus respectivos espaços de transformação ($P_m$ e $P_u$). Crucialmente, o lado dos pesos mantém uma transformação estática unificada ($P_e$) aplicada uma única vez aos pesos. Isso evita a duplicação de pesos transformados durante a inferência, enquanto permite que o lado da ativação se adapte ao estado do token.
• Otimização: Os parâmetros de transformação são aprendidos minimizando o erro de quantização por camada entre a saída de precisão total e a projeção quantizada guiada por estado.

2. Compensação de Atenção Temporal (TAC)

O TAC aborda o acúmulo de erros de quantização através das etapas iterativas de denoising.

• Alvo: O método foca na saída da multiplicação softmax-valor no mecanismo de atenção, uma operação conhecida por ser propensa a erros sob quantização de baixa bitagem.
• Mecanismo: O TAC insere um mapeamento afim de bloco-diagonal leve antes da projeção de saída da atenção. Ele não altera o cálculo da atenção em si, mas corrige a representação de atenção quantizada ($\hat{Z}_t$) para alinhá-la com sua contraparte de precisão total ($Z^{FP}_t$).
• Estimativa: Os parâmetros de compensação (matrizes por bloco $M_b$ e vieses $d_b$) são estimados em forma fechada, combinando as estatísticas de primeira e segunda ordem (média e covariância) das representações de precisão total e quantizadas. Um termo de amortecimento e um suavizamento de valor singular são aplicados para garantir a estabilidade, dado o limite de dados de calibração.

Principais Contribuições

1. Identificação do Problema: Os autores identificam formalmente a disparidade de ativação dependente de estado e o acúmulo de erro temporal como os principais gargalos para a quantização de baixa bitagem em DLLMs.
2. Transformação de Ativação Guiada por Estado (SGAT): Um novo método que mapeia ativações mascaradas e não mascaradas em espaços de transformação compartilhados e específicos de estado, mantendo uma transformação de peso estática unificada, suavizando efetivamente os outliers específicos de estado sem duplicação de pesos.
3. Compensação de Atenção Temporal (TAC): Um método de compensação afim por bloco, de forma fechada, que alinha as representações de atenção quantizadas com suas contrapartes de precisão total antes da projeção de saída, mitigando o acúmulo de erro ao longo da trajetória de denoising.
4. Validação Empírica: Experimentos extensos demonstram que o STaR-Quant supera consistentemente as bases de comparação de PT de forte desempenho (incluindo RTN, AWQ, QuaRot e DLLMQuant) em múltiplos DLLMs e tarefas.

Resultados Experimentais

Os autores avaliaram o STaR-Quant em três DLLMs representativos: LLaDA-8B, LLaDA-1.5-8B e Dream-7B, utilizando quantização de peso e ativação de 4 bits (W4A4).

• Desempenho: O STaR-Quant alcançou o melhor desempenho médio em 9 benchmarks (cobrindo conhecimento geral, raciocínio e geração de código).
  • No LLaDA-8B, ele melhorou a pontuação média em 2,78 pontos em relação à base de comparação mais forte (DLLMQuant++).
  • No Dream-7B, ele melhorou a pontuação média em 1,69 pontos.
  • Ganhos notáveis foram observados em tarefas de raciocínio (GSM8K) e geração de código (HumanEval), que são sensíveis a erros de representação intermediária.
• Eficiência:
  • Aceleração: O STaR-Quant alcançou acelerações de inferência de 1,65× a 1,69× em comparação com a implantação em FP16.
  • Memória: Reduziu o uso de pico de memória da GPU em 3,05× a 3,14× (por exemplo, reduzindo a memória do LLaDA de ~15,9GB para ~5,2GB).
• Estudos de Ablação: A remoção do TAC causou uma queda significativa no desempenho (indicando a importância da correção de erro temporal), enquanto a remoção do SGAT também degradou os resultados, confirmando a natureza complementar da consistência por estado e temporal.

Significância e Alegações

O artigo afirma que o STaR-Quant proporciona uma mudança necessária na estratégia de quantização para modelos de linguagem baseados em difusão. Ao modelar explicitamente as dimensões de estado (mascarado vs. não mascarado) e tempo (acúmulo iterativo), o framework resolve o descompasso entre as suposições de PTQ convencional e o processo de inferência de difusão.

Os autores posicionam o STaR-Quant como uma solução prática para implantar DLLMs em dispositivos com recursos limitados, oferecendo um equilíbrio entre altas taxas de compressão (W4A4), economias significativas de memória e acelerações, mantendo níveis de precisão próximos aos modelos de precisão total. O trabalho sugere que a "consistência estado-tempo" é um princípio eficaz para a quantização eficiente de DLLMs, embora os autores notem que a validação adicional em arquiteturas diversas e bitagens mais agressivas (ex: 2-bit ou 3-bit) permanece uma direção aberta.