Resumo Técnico: LLMs Eficientes com AMP: Poda de Cabeças de Atenção e MLP

Declaração do Problema

Os Grandes Modelos de Linguagem (LLMs) alcançaram um desempenho de nível humano em tarefas cognitivas, mas sofrem com contagens extensas de parâmetros, o que leva a altos custos computacionais e inferência lenta. Esses fatores dificultam a implantação em ambientes com recursos limitados. Embora a poda (pruning) seja uma estratégia promissora para reduzir o tamanho do modelo, os métodos existentes enfrentam trocas significativas:

• Poda Estruturada: Remove componentes inteiros (ex: camadas, cabeças), mas frequentemente requer um ajuste fino (fine-tuning) computacionalmente caro ou leva à degradação do desempenho se componentes críticos forem removidos.
• Poda Semiestruturada: Requer hardware especializado (ex: padrões de esparsidade de GPU específicos como 2:4) para obter ganhos de velocidade, limitando a implantação em hardware de consumo.
• Poda Não Estruturada: Cria matrizes esparsas irregulares que não conseguem acelerar a inferência em hardware de propósito geral sem otimizações de software especializadas.
• Limitações Gerais: Muitos métodos existentes dependem de métricas baseadas em perda, expansões de Taylor ou heurísticas complexas que são computacionalmente caras, inflexíveis ou exigem hardware especializado. Há uma necessidade de um método agnóstico ao hardware, eficiente, que comprima LLMs com perda mínima de desempenho e ganhos tangíveis de velocidade de inferência.

Metodologia: AMP (Poda de Cabeças de Atenção e MLP)

Os autores propõem o AMP, um método de poda estruturada que remove cabeças de atenção menos críticas dentro da Multi-Head Attention (MHA) e neurônios dentro de Multilayer Perceptrons (MLP). Diferente de métodos que dependem apenas de magnitudes de pesos ou mudanças de perda, o AMP avalia a importância estrutural projetando dados de entrada nos pesos para avaliar a contribuição de cada componente para o fluxo residual.

Mecanismo Central

O método opera sob o princípio de que a informação flui através do fluxo residual. O AMP calcula uma pontuação de importância para cada componente com base na magnitude de sua contribuição de ativação para a saída final da camada.

1. Componente MHA:

   • A saída da MHA é decomposta em uma soma de contribuições de cabeças de atenção individuais ($O = \sum h_n W_n$).
   • A importância da $n$-ésima cabeça ($I_n$) é definida como a norma $\ell_1$ do produto da saída da cabeça ($h_n$) e sua respectiva matriz de projeção de saída ($W_n$):
     $$I_n = \|h_n W_n\|_1$$
   • Esta abordagem considera as interações não lineares entre pesos e entradas, ao contrário da poda simples por norma de peso.

2. Componente MLP:

   • O método visa o mecanismo SwigLU usado em modelos como LLaMA, especificamente as projeções Up e Gate.
   • Como o SwigLU utiliza multiplicação elemento a elemento, os neurônios devem ser podados em pares de ambas as projeções para manter a integridade arquitetônica.
   • A importância do $m$-ésimo par de neurônios ($I_m$) é calculada como a média da norma $\ell_1$ do produto elemento a elemento das projeções Gate e Up através de uma amostra de tokens de entrada:
     $$I_m = \frac{1}{S} \sum_{s=1}^{S} | \text{SiLU}((x_s W_{\text{Gate}})_m) \cdot (x_s W_{\text{Up}})_m |$$

3. Processo de Poda:

   • Um pequeno conjunto de calibração (ex: 50 amostras aleatórias do Alpaca) é usado para extrair valores médios de ativação.
   • Os componentes (cabeças e neurônios) são classificados por importância em todas as camadas.
   • Para atingir uma taxa de compressão alvo $c$, os $c\%$ de cabeças de atenção e $c\%$ de neurônios de MLP inferiores são removidos uniformemente em todas as camadas.
   • Pós-Treinamento: O modelo podado passa por um processo de ajuste fino leve utilizando Low-Rank Adaptation (LoRA) para recuperar o desempenho.

Principais Contribuições

• Métrica de Importância Inovadora: O AMP introduz uma métrica baseada em dados que projeta ativações de entrada nos pesos para avaliar a importância do componente, diferenciando-se das heurísticas tradicionais de apenas peso ou baseadas em perda.
• Ganho de Velocidade Agnóstico ao Hardware: O método alcança ganhos práticos de velocidade de inferência (até 1.25$\times$) sem exigir hardware especializado ou núcleos de tensores esparsos, tornando-o adequado para GPUs de consumo.
• Eficiência: O processo de decisão de poda leva apenas minutos em uma GPU de consumo, e o ajuste fino subsequente exige recursos negligenciáveis em comparação ao retreinamento total.
• Validação de Coerência: Os autores propõem e executam um experimento de "verificação de coerência", demonstrando que remover os componentes mais importantes identificados pelo AMP causa quedas catastróficas de desempenho, validando a capacidade da métrica de distinguir estruturas críticas de não críticas.
• Flexibilidade: O método é aplicável a diferentes famílias de LLM (LLaMA, Phi) e permite taxas de compressão flexíveis sem estar vinculado a padrões de esparsidade fixos (ex: 2:4).

Resultados Experimentais

Os autores avaliaram o AMP nos modelos LLaMA 7B, LLaMA-2 7B, Phi-1.5 e Phi-2 usando benchmarks zero-shot (WinoGrande, HellaSwag, ARC-e/c, PIQA) e perplexidade (WikiText2).

• Desempenho:
  • LLaMA 7B: Com uma taxa de poda de 20%, o AMP alcançou uma precisão média de 64,52%, superando métodos de estado da arte como LLM-Pruner (59,71%) e Shortened LLaMA (63,46%).
  • LLaMA-2 7B: Com uma taxa de poda de 30%, o AMP alcançou 61,02% de precisão média, superando o SliceGPT (51,50%) e o LLM Surgeon (59,03%).
  • Melhoria: O AMP superou as técnicas de poda estruturada existentes em até 1,49 ponto percentual na precisão média das tarefas.
• Velocidade de Inferência:
  • O AMP alcançou um ganho de velocidade de 1.25$\times$ para o LLaMA 7B e 1.19$\times$ para o LLaMA-2 7B.
  • Isso supera o Shortened LLaMA (1.13$\times$) e rivaliza com métodos semiestruturados como o Wanda (1.24$\times$) sem exigir hardware especializado.
• Perplexidade: Os modelos podados com ajuste fino mostraram perplexidade significativamente menor no WikiText2 em comparação com as versões sem ajuste fino, com a lacuna aumentando em taxas de poda mais altas.
• Estudos de Ablação:
  • Verificação de Coerência: Inverter a ordem da poda (remover os componentes mais importantes) resultou em quedas de precisão de >25 pontos percentuais, confirmando a validade da métrica.
  • Isolamento de Componente: Podar apenas cabeças de MHA ou apenas neurônios de MLP resultou em um desempenho significativamente pior do que podar ambos simultaneamente, indicando a necessidade de uma abordagem equilibrada.

Significância e Alegações

O artigo afirma que o AMP aborda as trocas críticas na compressão de modelos ao oferecer um método que é:

1. Efetivo: Supera os métodos atuais de poda estruturada de estado da arte em precisão, mantendo altas taxas de compressão.
2. Eficiente: Requer recursos computacionais mínimos tanto para a decisão de poda (minutos) quanto para a recuperação de desempenho (ajuste fino leve de LoRA).
3. Prático: Permite ganhos tangíveis de velocidade de inferência em hardware padrão, facilitando a implantação em ambientes com recursos limitados.
4. Sustentável: Alinhado aos princípios de Green AI, o método reduz os custos operacionais em ~17,39% e as emissões de carbono em ~19,97% devido à redução da latência de inferência e da contagem de parâmetros.

Os autores enfatizam que seu método identifica e remove com sucesso estruturas não essenciais, permitindo a implantação de LLMs altamente otimizados sem a necessidade de infraestrutura de hardware especializada. Eles também destacam o papel crítico do pós-treinamento (ajuste fino) para restaurar o desempenho do modelo após a poda.