Sparsity Induction for Accurate Post-Training Pruning of Large Language Models

1. O Problema

Os Modelos de Linguagem Grandes (LLMs) modernos possuem escalas de parâmetros que vão de dezenas a centenas de bilhões, resultando em altos custos computacionais, latência e consumo de energia. A Poda Pós-Treinamento (Post-Training Sparsity - PTS) é uma abordagem promissora para reduzir esses custos removendo pesos de redes densas sem a necessidade de retreinamento massivo.

No entanto, existem desafios fundamentais:

• Falta de Esparsidade Nativa: As matrizes densas originais não possuem alta esparsidade inerente.
• Disrupção do Estado do Modelo: Métodos existentes que removem pesos diretamente (como poda baseada em magnitude) perturbam o estado do modelo, desequilibrando a distribuição estatística dos pesos e as escalas relativas entre canais.
• Limitações de Métodos Atuais: Embora métodos avançados (como SparseGPT e Wanda) tentem corrigir erros através de compensação ou sensibilidade à perda, eles operam em uma distribuição de pesos que não foi preparada para a esparsidade. Isso leva a uma degradação de desempenho significativa, especialmente em níveis de esparsidade agressivos, e os ganhos de refinar apenas "pontuações de importância" tornam-se marginais.

2. Metodologia: Indução de Esparsidade (Sparsity Induction - SI)

Os autores propõem a Indução de Esparsidade (SI), um framework que molda proativamente a distribuição de pesos e a estrutura de características do modelo antes da poda, tornando-o mais "amigável" à esparsidade. A abordagem atua em dois níveis complementares:

A. Nível de Distribuição (Transformações Equivalentes)

O objetivo é reconfigurar a distribuição de pesos para aumentar a separabilidade entre pesos importantes e não importantes, sem alterar a funcionalidade do modelo.

• Reparametrização Absorvível: Introduzem transformações de escala e deslocamento por canal (usando fatores aprendíveis $s$ e $\delta$) que são matematicamente equivalentes à camada original.
• Aplicação:
  • Em camadas lineares: $Y = WX + b$ é reparametrizado como $Y = S^{-1}(X-\delta)(WS) + (b+W\delta)$.
  • Em Atenção: Aplica-se escalas inversas nas projeções de Query ($Q$) e Key ($K$) para manter os logits de atenção intactos ($\tilde{Q}\tilde{K}^T = QK^T$).
• Vantagem: Esses parâmetros são treinados em dados de calibração e, posteriormente, fundidos (absorvidos) nos pesos originais. Isso significa que não há sobrecarga de inferência nem parâmetros extras durante a execução.

B. Nível de Características (Feature Level)

Para garantir uma direção de busca estável e promover esparsidade nas características:

• Inicialização Robusta: Os fatores de transformação são inicializados com base em estatísticas robustas (mediana e média) das ativações de saída.
• Perda de Norma Espectral (Spectral Norm Loss): Adiciona-se um termo de regularização baseado na norma espectral (ou norma $L_p$) para incentivar uma estrutura de baixo posto (low-rank). Isso promove a esparsidade nas características, penalizando a distorção da saída e estabilizando a convergência sem necessidade de ajuste fino de todos os parâmetros.

C. Eficiência Computacional

• Atualização Rápida do Hessiano: Para evitar a recomputação cara de ativações durante o treinamento dos fatores de escala, os autores utilizam uma aproximação diagonal do Hessiano (ou Gauss-Newton) baseada na covariância das entradas.
• Resultado: Isso permite atualizar as pontuações de importância em tempo constante ($O(d_{in})$), acelerando o processo de otimização em até 22,65 vezes em comparação com métodos clássicos de recompute.

3. Principais Contribuições

1. Conceito de Indução de Esparsidade: Introduz a ideia de adaptar a distribuição do modelo antes da poda, em vez de apenas tentar escolher os melhores pesos para remover.
2. Framework Dual: Desenvolveu um método que otimiza propriedades de esparsidade simultaneamente em nível de distribuição (escala/deslocamento) e nível de características (regularização espectral).
3. Eficiência e Compatibilidade: O método é "plug-and-play", não requer ajuste fino completo (full-parameter fine-tuning) e, crucialmente, não adiciona latência de inferência ou módulos extras, pois os parâmetros aprendidos são fundidos nos pesos.
4. Aceleração Algorítmica: Propôs uma atualização eficiente do Hessiano que acelera significativamente o pipeline de otimização.

4. Resultados Experimentais

Os experimentos foram realizados em oito modelos LLM (OPT, LLaMA-1 e LLaMA-2) com tamanhos variados (125M a 13B) e em diversas tarefas (perplexidade em C4/WikiText-2 e precisão zero-shot em benchmarks como ARC, HellaSwag, etc.).

• Desempenho de Perplexidade (PPL): A SI reduziu drasticamente a degradação de PPL causada pela poda. Em cenários agressivos (50% de esparsidade), modelos base (como Magnitude) tornaram-se inutilizáveis, enquanto a combinação com SI manteve a qualidade próxima dos modelos densos.
  • Exemplo: No OPT-125M com 50% de esparsidade, a Magnitude pura teve PPL de 193,36, enquanto Magnitude + SI caiu para 52,85.
• Precisão Zero-Shot: A SI melhorou consistentemente a precisão média em tarefas de raciocínio e compreensão para todos os métodos base (Magnitude, Wanda, SparseGPT).
• Suporte a Estruturas N:M: O método funcionou bem em padrões semi-estruturados (2:4 e 4:8), sendo totalmente compatível com aceleração de hardware.
• Latência e Velocidade:
  • A atualização rápida do Hessiano reduziu o tempo de atualização de métricas de 345s para 15s (22,65x mais rápido).
  • A latência de ponta a ponta (E2E) foi idêntica à do método Wanda puro (sem sobrecarga), garantindo ganhos de velocidade de 1,24x em relação ao modelo denso.

5. Significado e Conclusão

O trabalho demonstra que a eficácia da poda pós-treinamento não depende apenas de métricas de importância mais refinadas, mas sim da preparação da distribuição do modelo para suportar a remoção de pesos.

A Indução de Esparsidade oferece uma solução prática e eficiente que:

1. Preserva a consistência de escala e distribuição que o modelo aprendeu durante o treinamento.
2. Permite alcançar níveis de esparsidade mais altos com menor perda de precisão.
3. Mantém a eficiência de inferência, sendo ideal para implantação em ambientes com recursos limitados.

Este método representa um avanço fundamental na compressão de LLMs, sugerindo que a "indução" de propriedades desejáveis antes da compressão é uma via mais promissora do que a simples otimização de critérios de seleção de pesos.