LEAP: Learnable End-to-End Adaptive Pruning of Large Language Models

Resumo Técnico: LEAP: Poda Adaptativa Aprendível de Ponta a Ponta de Modelos de Linguagem Grandes

Declaração do Problema
A implantação de Modelos de Linguagem Grandes (LLMs) está cada vez mais limitada por restrições de memória e computação. Embora a esparsidade de pesos ofereça uma solução, o cenário está dividido entre regimes estruturados, semi-estruturados (por exemplo, 2:4) e não estruturados. Avanços recentes em kernels de GPU (por exemplo, FlashLLM, SpInfer) e hardware de fluxo de dados deslocaram o gargalo de executar esparsidade não estruturada para induzi-la com perda mínima de precisão.

Os métodos atuais de ponta para poda não estruturada, como Wanda, SparseGPT, Thanos e ADMM, dependem de substitutos camada a camada derivados do princípio do Cérebro Ótimo (OBS). Esses métodos minimizam erros de reconstrução locais em vez da perda do modelo de ponta a ponta. Consequentemente, eles acumulam erros locais que se compostam em redes profundas, levando a uma degradação significativa da precisão, particularmente sob esparsidade agressiva (50–60%).

Por outro lado, métodos de máscaras aprendíveis de ponta a ponta, como MaskLLM e PATCH, otimizam com sucesso máscaras diretamente contra a perda de modelagem de linguagem, mas são restritos a padrões semi-estruturados. Sua parametrização depende de uma distribuição categórica sobre padrões válidos dentro de um grupo. Embora tratável para grupos pequenos (por exemplo, $\binom{4}{2}=6$ padrões para esparsidade 2:4), essa abordagem torna-se combinatorialmente intratável para esparsidade não estruturada. Para uma linha de largura $d=4096$ a 50% de esparsidade, o número de máscaras válidas é $\binom{4096}{2048} \approx 10^{1229}$, tornando impossível o armazenamento e indexação de logits de padrão.

Metodologia: LEAP
Os autores introduzem o LEAP (Poda Adaptativa Aprendível de Ponta a Ponta), um framework projetado para tornar viável o aprendizado de máscaras não estruturadas de ponta a ponta, substituindo a parametrização intratável de categorias-sobre-padrões por uma abordagem por peso.

• Relaxação Bernoulli por Peso: Em vez de atribuir um logit a cada padrão válido, o LEAP atribui um logit aprendível a cada peso individual. Para uma matriz de pesos $W \in \mathbb{R}^{m \times n}$, uma matriz de logits paralela $P \in \mathbb{R}^{m \times n}$ é introduzida. A máscara estocástica $M$ é gerada via relaxação Gumbel-sigmoid:
  $$M = \sigma\left(\frac{\alpha P + g}{\tau}\right)$$
  onde $g$ é ruído Gumbel, $\sigma$ é a função sigmoid, $\alpha$ é um fator de escala e $\tau$ é um parâmetro de temperatura. Essa formulação escala linearmente com o número de pesos ($O(mn)$) em vez do número de padrões válidos.
• Estratégia de Otimização:
  • Pesos Congelados: Os pesos pré-treinados $W$ são mantidos congelados. Apenas os logits da máscara $P$ são treinados. Isso preserva a calibração do modelo pré-treinado e simplifica a implantação.
  • Inicialização: $P$ é inicializada usando uma máscara Wanda de um único disparo (definindo entradas selecionadas para $+s$ e outras para $-s$) para fornecer um ponto de partida razoável, evitando problemas de início frio.
  • Agendamentos: Um agendamento de escala ($\alpha$) aumenta para empurrar a sigmoid em direção a valores binários, enquanto um agendamento de temperatura ($\tau$) decai para afiar a máscara.
  • Função de Perda: O objetivo combina a perda de modelagem de linguagem em um fluxo de calibração com dois termos de regularização:
    1. Regularizador de Esparsidade Global: Impõe a densidade alvo globalmente em todas as camadas, permitindo que camadas individuais ajustem sua densidade com base na importância de ponta a ponta.
    2. Termo Consciente de Magnitude: Viésa a otimização para reter pesos de maior magnitude para prevenir mínimos degenerados onde pesos críticos são descartados.
• Eficiência de Treinamento: O método treina por aproximadamente 2.000 iterações em um pequeno fluxo de calibração de texto geral (SlimPajama).

Principais Contribuições

1. Identificação do Obstáculo Combinatório: O artigo identifica formalmente que parametrizações de categorias-sobre-padrões (usadas por MaskLLM/PATCH) não podem escalar para esparsidade não estruturada devido ao crescimento exponencial de padrões válidos.
2. Reformulação Viável: Os autores propõem uma reformulação Bernoulli-via-Gumbel-sigmoid por peso que permite aprendizado diferenciável de ponta a ponta para máscaras não estruturadas na escala de LLMs.
3. Framework LEAP: Um framework de poda leve que opera sobre pesos congelados, exigindo apenas um pequeno conjunto de dados de calibração e cerca de 2.000 etapas de treinamento.
4. Validação Empírica: Avaliação extensa em cinco famílias de LLMs (Qwen-2.5, Gemma-3, LLaMA-3.2, LLaMA-3.1) variando de 0,5B a 8B parâmetros em esparsidades de 50% e 60%.

Resultados
O LEAP foi avaliado contra baselines camada a camada (Wanda, SparseGPT, Thanos, ADMM) e a baseline aprendível semi-estruturada MaskLLM.

• Ganhos de Precisão: Em dez configurações (modelo, esparsidade), o LEAP melhorou a precisão zero-shot média de seis tarefas em +2,59 pontos sobre o ADMM, a melhor baseline camada a camada.
• Melhorias Específicas: O maior ganho foi de +5,40 pontos no LLaMA-3.2 1B a 60% de esparsidade (50,39 vs. 44,99).
• Perplexidade: O LEAP consistentemente alcançou perplexidade (PPL) menor que as baselines. Por exemplo, no LLaMA-3.1 8B a 50% de esparsidade, o LEAP alcançou 7,66 PPL comparado aos 9,12 do ADMM.
• Alocação de Esparsidade: Experimentos mostraram que, sob otimização de ponta a ponta com um orçamento global, o LEAP converge para uma alocação de densidade por bloco quase uniforme, contrastando com as variações não uniformes observadas em métodos semi-estruturados.
• Custo Computacional: Embora mais caro que métodos de um único disparo (levando horas em H100s versus minutos), o LEAP é significativamente mais barato que o MaskLLM (por exemplo, 276 horas de GPU para o LLaMA-3.1 8B versus milhares de horas de A100 para o MaskLLM em escalas comparáveis).

Significado e Alegações
O artigo alega que o LEAP resolve o obstáculo fundamental que impede o aprendizado de máscaras de ponta a ponta em esparsidade não estruturada. Ao desacoplar o aprendizado da máscara das atualizações de pesos, oferece um caminho prático para comprimir modelos fundamentais sem sacrificar a calibração dos pesos pré-treinados.

Os autores enfatizam que o LEAP produz máscaras no regime não estruturado de 50%–60%, que são diretamente compatíveis com kernels acelerados existentes (por exemplo, FlashLLM, SpInfer, MACKO). O significado não reside em criar novos kernels de inferência, mas em permitir a implantação desses kernels em modelos podados de alta qualidade que eram anteriormente inatingíveis por métodos camada a camada. O trabalho sugere que a esparsidade não estruturada, quando induzida via aprendizado de ponta a ponta, pode oferecer uma compensação favorável entre compressão e precisão para implantação em hardware comercial.

Limitações e Escopo

• Pesos Congelados: O escopo atual mantém os pesos congelados. Embora a otimização conjunta de pesos e máscaras seja notada como uma extensão natural, não é implementada neste trabalho.
• Alocação Uniforme: A observação de densidade por bloco uniforme é apresentada como uma descoberta específica do regime para os modelos estudados e níveis de esparsidade, não uma regra geral contra alocação não uniforme.
• Segurança: Os autores alertam que modelos comprimidos podem exibir modos de falha diferentes em entradas de cauda longa e recomendam reavaliação para implantações críticas de segurança.