Pruning via Causal Attribution Preserves Reasoning Performance in Large Language Models

Resumo Técnico: Poda por Atribuição Causal (CAP)

Definição do Problema
Os Grandes Modelos de Linguagem (LLMs) demonstram fortes capacidades de raciocínio de múltiplos passos, mas incorrem em custos de inferência substanciais, o que dificulta sua implantação no mundo real. Embora a poda de modelos ofereça uma solução ao remover parâmetros redundantes sem retreinamento, os métodos existentes frequentemente degradam o desempenho em tarefas complexas de raciocínio. As abordagens atuais baseiam-se em sinais correlacionais, como a magnitude dos pesos (ex: Wanda) ou estatísticas de gradiente local. Esses métodos assumem que pesos pequenos são irrelevantes; entretanto, parâmetros de baixa magnitude podem ser críticos para circuitos de raciocínio de múltiplos passos. Além disso, critérios puramente correlacionais falham em capturar interações entre camadas essenciais para a inferência de cadeia de pensamento (Chain-of-Thought - CoT), levando frequentemente ao "colapso de camada", onde módulos inteiros perdem sua função.

Metodologia: Poda por Atribuição Causal (CAP)
Os autores introduzem a Poda por Atribuição Causal (CAP), um método de poda de etapa única (one-shot) e livre de treinamento, projetado para preservar o desempenho de raciocínio. Diferente dos métodos que classificam pesos com base em propriedades estáticas, o CAP mede o impacto causal de componentes específicos do modelo no desempenho da tarefa. O framework opera em três estágios:

1. Calibração e Pontuação Causal (Estágio 1):

   • Um pequeno conjunto de dados de calibração ($D_{cal}$) é construído através da amostragem de benchmarks centrados em raciocínio: GSM8K (matemática), StrategyQA (senso comum de múltiplos saltos) e ARC-Challenge (ciência).
   • Para cada cabeça de atenção $h_\ell$ no modelo, o CAP realiza um experimento de mascaramento intervencional. A cabeça é temporariamente mascarada durante a passagem direta (forward pass), e o aumento resultante na perda de entropia cruzada por token ($\Delta(h_\ell)$) é medido em relação ao modelo intacto.
   • Para garantir a robustez contra a variância de amostras pequenas, o conjunto de dados é particionado em $K$ subconjuntos, e a mediana do aumento da perda é computada.
   • Este processo identifica cabeças que são funcionalmente críticas para o raciocínio, distinguindo-as daquelas que são meramente grandes numericamente, mas redundantes.

2. Conversão de Importância ao Nível de Peso (Estágio 2):

   • As pontuações causais ao nível de cabeça são convertidas em fatores de importância normalizados $I(h_\ell)$, mapeados para um intervalo de $[1.0, 10.0]$. Cabeças de alto impacto recebem fatores mais altos.
   • Esses fatores são propagados para pesos individuais dentro de suas matrizes de projeção correspondentes ($W_q, W_k, W_v, W_o$). A pontuação de importância para um peso específico $W_{ij}$ é calculada como o produto de sua magnitude pelo fator de importância de sua região de cabeça associada: $S_{ij} = |W_{ij}| \cdot I(h_{region})$.
   • Para camadas MLP, que carecem de uma estrutura natural particionada por cabeças, o método atribui uma importância média por camada derivada das cabeças de atenção colocalizadas. Isso é observado como uma aproximação.

3. Poda Global (Estágio 3):

   • Todos os pesos em todo o modelo são classificados por seus scores de magnitude ponderados pela importância ($S_{ij}$).
   • Um limiar global é determinado via busca binária para atingir a esparsidade alvo (ex: 10%, 20%, 50%). Pesos abaixo do limiar são zerados, criando um padrão de esparsidade não estruturada compatível com infraestruturas de inferência padrão.

Principais Contribuições

• Formalização da Atribuição Causal para Poda: Os autores formalizam um método para medir o aumento esperado na perda específica da tarefa quando cabeças de atenção individuais são mascaradas, indo além de proxies correlacionais para quantificar a criticidade funcional.
• Poda de Peso Consciente de Orçamento: Eles introduzem um procedimento que converte pontuações causais ao nível de cabeça em valores de importância ao nível de peso. Isso permite uma poda de magnitude refinada sob uma restrição de esparsidade global, preservando a estrutura de circuitos de raciocínio críticos.
• Validação Empírica: O método é avaliado no Llama-3-8B-Instruct e Mistral-7B-Instruct através dos benchmarks GSM8K, StrategyQA e ARC-Challenge.

Result Resultados

• Esparsidade Moderada (10–20%): O CAP geralmente supera o baseline Wanda na preservação do desempenho de raciocínio.
  • No Llama-3-8B-Instruct com 20% de esparsidade, o CAP alcança uma melhoria relativa de 61% sobre o Wanda no ARC-Challenge (70,8% vs. 43,9% de acurácia).
  • No GSM8K com 10% de esparsidade, o CAP mostra uma melhoria relativa de 20% sobre o Wanda (71,5% vs. 59,5%).
  • Os autores observam que, embora o CAP às vezes fique atrás do Wanda em tarefas específicas (ex: Llama-3 GSM8K com 20%), ele preserva consistentemente o desempenho em tarefas de inferência de múltiplos passos onde os critérios correlacionais do Wanda falham.
• Alta Esparsidade (50%): O desempenho degrada significativamente para ambos os métodos. O CAP sofre uma falha catastrófica no Llama-3 (o salto de perplexidade chega a 428,2), atribuído ao seu tratamento grosseiro das camadas MLP. Neste nível, o limiar global remove pesos MLP essenciais que foram atribuídos com baixa importância com base na aproximação das cabeças de atenção.
• Perplexidade vs. Raciocínio: Em esparsidade moderada, o CAP frequentemente alcança melhor acurácia de raciocínio, apesar de uma perplexidade ligeiramente maior comparada ao Wanda, sugerindo um desacoplamento entre a qualidade do modelamento de linguagem geral e a preservação do raciocínio específico da tarefa.
• Limitações de Arquitetura: Experimentos preliminares em modelos de Mistura de Especialistas (MoE) sugerem que o CAP é menos eficaz, pois o roteamento dinâmico torna a importância das cabeças não estacionária.

Significância e Alegações
O artigo alega que a atribuição causal ao nível de cabeça de atenção é um critério superior para preservar o desempenho de raciocínio em comparação com métodos de poda correlacionais em níveis de esparsidade equivalentes. Ao medir explicitamente a contribuição causal das cabeças para tarefas de raciocínio, o CAP pode identificar e proteger parâmetros "numericamente pequenos, mas funcionalmente críticos" que os métodos baseados em magnitude descartam.

Os autores enfatizam que seu trabalho é modesto em escopo:

• É um método livre de treinamento, tornando-o computacionalmente eficiente em comparação com abordagens que exigem atualizações de peso (como o SparseGPT).
• É alinhado à tarefa, calibrado especificamente em dados de raciocínio em vez de modelagem de linguagem genérica.
• Reconhece limitações, particularmente a falta de atribuição causal refinada para neurônios MLP e a dificuldade de aplicar o método a arquiteturas de roteamento dinâmico (MoE).

A contribuição central é a demonstração de que a poda intervencional e alinhada à tarefa pode preservar melhor capacidades específicas (raciocínio) do que métodos otimizados para perplexidade geral ou estatísticas de peso estáticas, desde que o nível de esparsidade permaneça moderado.