PATCH: Learnable Tile-level Hybrid Sparsity for LLMs

1. Formulação do Problema

Os Grandes Modelos de Linguagem (LLMs) oferecem desempenho superior, mas sofrem com custos proibitivos de memória e computação durante a implantação. A poda de modelos é uma solução primária, mas as abordagens existentes enfrentam um trade-off fundamental:

• Esparsidade Não Estruturada: Permite elementos não nulos em qualquer lugar, preservando alta precisão, mas cria padrões de acesso à memória irregulares que impedem a aceleração eficiente de GPU (por exemplo, Tensor Cores).
• Esparsidade Semi-Estruturada (por exemplo, 2:4): Impõe um padrão rígido (dois não nulos a cada quatro elementos) que é amigável ao hardware e acelera a inferência. No entanto, a taxa de esparsidade fixa de 50% e a aplicação uniforme em todas as camadas frequentemente levam a degradação significativa da precisão, especialmente ao usar métodos de poda one-shot.

Os métodos atuais lutam para preencher a lacuna entre a flexibilidade da poda não estruturada e a eficiência de hardware da poda semi-estruturada, particularmente em relação à incapacidade de adaptar continuamente as taxas de esparsidade ou alocar a esparsidade de forma não uniforme entre as camadas.

2. Metodologia: PATCH

Os autores propõem o PATCH (Pruning with a Learnable Tile-level Configuration for Hybrid Sparsity), um framework que permite uma taxa de esparsidade contínua entre 0% e 50% enquanto mantém a compatibilidade com hardware.

Mecanismo Central

O PATCH particiona matrizes de pesos em tiles (por exemplo, $b_1 \times b_2$). Para cada tile, o modelo aprende a decidir se deve permanecer denso (0% de esparsidade) ou tornar-se 2:4 esparsos (50% de esparsidade).

• Geração de Máscara Híbrida: A máscara final é uma combinação ponderada de uma opção "densa" e uma opção "2:4 esparsa".
• Distribuições Aprendíveis:
  1. Seleção de Tile: Uma distribuição de logit aprendível determina a probabilidade de um tile ser denso versus esparso. Isso é amostrado usando Gumbel-Softmax para permitir otimização diferenciável.
  2. Seleção de Padrão: Dentro dos tiles designados como esparsos, uma distribuição aprendível separada seleciona o padrão 2:4 específico (entre 6 permutações possíveis).
• Objetivo de Otimização: O objetivo de treinamento combina:
  • Perda de modelagem padrão (por exemplo, Previsão do Próximo Token).
  • Regularização de Esparsidade: Um termo que penaliza desvios de uma taxa de esparsidade global alvo ($\rho$), permitindo controle preciso sobre a faixa de 0–50%.
  • Regularização de Pesos: Encoraja magnitudes de pesos maiores para auxiliar a propagação de gradientes.

Variantes

• PATCHJoint: Otimiza conjuntamente tanto a seleção ao nível de tile (denso vs. esparso) quanto os padrões 2:4 de granularidade fina dentro dos tiles esparsos. Usado para modelos menores (<1 bilhão de parâmetros).
• PATCHTile: Uma variante eficiente em memória onde os padrões 2:4 são congelados (inicializados a partir de um método one-shot de alta qualidade como MaskLLM), e apenas as decisões ao nível de tile denso/esparso são otimizadas. Isso permite a escalabilidade para modelos maiores (até 13B) em memória de GPU limitada.

Inferência e Implantação

• Integração STOICC: Como bibliotecas padrão de GPU (cuBLAS, cuSPARSELt) não suportam tiles mistos densos/esparsos, o PATCH utiliza STOICC, um compilador baseado em Triton. O STOICC ajusta automaticamente as configurações do kernel (tamanhos de tile, por exemplo, $128 \times 128$) para lidar com a execução híbrida de forma eficiente.
• Compatibilidade com Hardware: Ao garantir que os tiles esparsos adiram estritamente ao padrão 2:4, o PATCH aproveita os Tensor Cores da NVIDIA/AMD para aceleração, enquanto os tiles densos utilizam kernels densos padrão.

3. Contribuições Principais

1. Framework de Esparsidade Híbrida: Introduz o primeiro método para permitir uma taxa de esparsidade contínua (0%–50%) misturando dinamicamente tiles densos e 2:4 esparsos dentro da mesma matriz de pesos.
2. Alocação Não Uniforme Aprendível: Diferentemente dos métodos 2:4 fixos, o PATCH aprende a alocar esparsidade de forma adaptativa. Ele preserva a densidade em camadas sensíveis (por exemplo, blocos iniciais/finais do Transformer, mecanismos de Atenção) e aplica esparsidade a camadas redundantes (por exemplo, blocos do meio, matrizes up/gate/down do MLP).
3. Treinamento Diferenciável End-to-End: Usa Gumbel-Softmax para otimizar conjuntamente a seleção de máscara e a poda de pesos, evitando a queda de precisão associada à poda one-shot.
4. Implantação Prática: Demonstra a primeira integração bem-sucedida de esparsidade híbrida ao nível de tile com um compilador (STOICC) para alcançar acelerações no mundo real em GPUs de nível de consumidor.

4. Resultados Experimentais

Os autores avaliaram o PATCH em modelos variando de 0,5B a 13B parâmetros (Qwen-2.5, LLaMA-2/3, Gemma-3) usando 8 tarefas downstream zero-shot e perplexidade WikiText2.

• Precisão:

  • O PATCH supera consistentemente os métodos de poda 2:4 mais avançados (MaskLLM, SparseGPT, Wanda, ProxSparse).
  • No LLaMA-2 7B, o PATCH (em 25% de esparsidade) alcança 51,58% de precisão média, comparado a 48,62% para MaskLLM (50% de esparsidade) e 45,92% para ProxSparse.
  • O PATCH reduz significativamente a lacuna para modelos densos; em 25% de esparsidade, recupera-se dentro de ~2 pontos percentuais da precisão densa.
  • Perplexidade: O PATCH alcança perplexidade menor que todas as linhas de base. Para LLaMA-2 7B em 25% de esparsidade, a PPL é 5,86 versus 6,78 para MaskLLM.

• Aceleração e Eficiência:

  • Em uma GPU NVIDIA A6000 (LLaMA-2 7B), o PATCH alcança acelerações de ponta a ponta de 1,18× a 1,38× sobre as linhas de base densas.
  • A pegada de memória é reduzida para 0,59×–0,76× do modelo denso.
  • O PATCH supera a poda não estruturada (que não oferece aceleração) e a poda 2:4 rígida (que frequentemente degrada a precisão demais para justificar a velocidade).

• Estudos de Ablação:

  • Global vs. Por Camada: Alvos de esparsidade global (permitindo esparsidade variável por camada) superam significativamente alvos uniformes por camada.
  • Tamanho do Tile: O desempenho é robusto em tamanhos de tile, embora $4 \times 4$ ofereça o controle mais fino. A eficiência de hardware dita o uso de tiles de execução maiores ($128 \times 128$).
  • Inicialização: O método é robusto a priores de inicialização; o alvo de esparsidade global permite realocação dinâmica durante o treinamento.

5. Significado

O PATCH representa um avanço significativo na inferência eficiente de LLMs, resolvendo o dilema "precisão versus aceleração".

• Preenchendo a Lacuna: Combina com sucesso os benefícios de aceleração de hardware da esparsidade semi-estruturada com a flexibilidade da alocação não estruturada.
• Praticidade: Vai além da poda teórica para demonstrar acelerações reais em hardware existente usando infraestrutura de compilador existente (STOICC).
• Escalabilidade: A variante eficiente em memória (PATCHTile) prova que a esparsidade híbrida pode ser aplicada a modelos grandes (13B+) sem exigir orçamentos massivos de computação para treinar as máscaras.
• Direção Futura: O framework é agnóstico a padrões e pode teoricamente estender-se para outros padrões N:M (por exemplo, 4:8) conforme novos hardwares os suportam, tornando-o uma solução versátil para a compressão futura de LLMs.

Em resumo, o PATCH fornece uma estratégia de poda flexível, aprendível e eficiente em hardware que permite aos praticantes ajustar o trade-off entre qualidade do modelo e velocidade de inferência com granularidade sem precedentes.