HEAPr: Hessian-based Efficient Atomic Expert Pruning in Output Space

O artigo apresenta o HEAPr, um novo algoritmo de poda que decompõe especialistas em unidades atômicas e utiliza informações de segunda ordem no espaço de saída para alcançar uma compressão quase sem perdas em modelos MoE, superando métodos existentes com menor custo computacional.

O essencial

Imagine que você tem um gigante da inteligência artificial (um Modelo de Linguagem Grande, ou LLM) que é incrivelmente inteligente, mas também gigantesco e caro de manter. Ele ocupa tanto espaço na memória do computador que é difícil colocá-lo em dispositivos comuns, como celulares ou servidores menores.

Esse gigante é construído como uma equipe de especialistas (chamada de "Mixture of Experts" ou MoE). Em vez de ter um único cérebro gigante, ele tem centenas de "cérebros menores" (especialistas) que trabalham juntos. Quando você faz uma pergunta, o sistema escolhe apenas alguns desses especialistas para responder, economizando energia de processamento.

O Problema:
Mesmo que o sistema só use alguns especialistas de cada vez, ele precisa guardar todos eles na memória o tempo todo. É como ter uma biblioteca com 1.000 livros, mas você só lê 5 por dia. O problema é que você ainda precisa ter espaço para os 1.000 livros na estante. Isso torna o sistema muito pesado e caro para rodar.

A Solução Antiga (e Imperfeita):
Até agora, as tentativas de aliviar esse peso eram como demitir especialistas inteiros. Se um especialista parecia "menos útil", eles o tiravam da equipe.

• O problema: É como demitir um médico inteiro só porque ele não é bom em tratar um tipo específico de doença. Você perde muita capacidade e o sistema fica "burro" ou erra muito. É uma solução muito grosseira.

A Nova Solução: O HEAPr (O "Poda de Precisão")
Os autores deste trabalho criaram um método chamado HEAPr. Em vez de demitir especialistas inteiros, eles decidiram fazer uma cirurgia de precisão dentro de cada especialista.

Aqui está a analogia para entender como funciona:

1. A Analogia da Fábrica de Móveis

Imagine que cada "Especialista" é uma fábrica de móveis.

• O jeito antigo: Se a fábrica não estava vendendo bem, você fechava a fábrica inteira. Perdeu-se tudo.
• O jeito novo (HEAPr): Eles perceberam que dentro de cada fábrica, existem máquinas individuais (chamadas de "especialistas atômicos"). Algumas máquinas fazem apenas as pernas da mesa, outras fazem o tampo, outras fazem o acabamento.
• O HEAPr olha para dentro da fábrica e diz: "Olha, a máquina que faz as pernas da mesa número 42 está quase parada. Vamos desligar apenas ela, mas manter o resto da fábrica funcionando."

Isso permite remover pedaços pequenos e inúteis sem destruir a capacidade total do sistema.

2. Como eles sabem o que cortar? (O "Raio-X" do Cérebro)

O grande desafio é: Como saber qual máquina desligar sem estragar o móvel?
Se você desligar a máquina errada, a mesa fica torta.

Os autores usaram uma técnica matemática avançada chamada Teoria do Cirurgião Cerebral Ótimo (Optimal Brain Surgeon).

• A ideia: Em vez de chutar, eles usam uma "fórmula de segunda ordem" (uma espécie de raio-x matemático) para medir exatamente o quanto o sistema vai sofrer se aquela peça for removida.
• O Truque Genial: Calcular esse "raio-x" para todo o sistema seria tão pesado que o computador travaria. Então, eles descobriram um atalho inteligente:
  1. Eles transformaram o problema de "analisar os parâmetros" (os pesos internos) para "analisar a saída" (o resultado final).
  2. É como se, em vez de tentar entender a engenharia complexa de cada máquina, eles apenas olhassem para o produto final que a máquina produz e perguntassem: "Se essa máquina parar, o produto final muda muito?"
  3. Se a resposta for "não muda quase nada", eles cortam.

Isso reduziu a necessidade de memória e poder de cálculo de algo impossível para algo muito leve.

3. Os Resultados Mágicos

O que acontece quando você aplica o HEAPr?

• Leveza: Eles conseguiram remover cerca de 20% a 25% do "peso" do modelo (os especialistas atômicos inúteis) e o modelo quase não perdeu inteligência. Foi como tirar a gordura de um atleta sem tirar o músculo.
• Velocidade: O modelo ficou cerca de 20% mais rápido para rodar.
• Precisão: Em testes com modelos famosos (como os da família Qwen e DeepSeek), o método funcionou muito melhor do que as técnicas antigas de "demitir especialistas inteiros".

Resumo em uma frase:

O HEAPr é como um cirurgião de precisão que, em vez de amputar membros inteiros de um gigante para fazê-lo caber em um carro pequeno, remove apenas as células de gordura inúteis, permitindo que o gigante continue forte e inteligente, mas agora leve o suficiente para viajar em qualquer lugar.

Por que isso importa?
Isso significa que em breve poderemos ter modelos de IA superinteligentes rodando em nossos computadores pessoais e celulares, sem precisar de supercomputadores caros, porque aprendemos a "podar" a inteligência de forma eficiente e sem perder qualidade.

Resumo técnico

Resumo Técnico: HEAPr

1. O Problema

As arquiteturas de Mistura de Especialistas (MoE) em Modelos de Linguagem de Grande Escala (LLMs) oferecem desempenho excepcional e custos de inferência reduzidos em comparação com modelos densos, ativando apenas um subconjunto de parâmetros durante a inferência. No entanto, esses modelos enfrentam um gargalo crítico: requisitos de memória proibitivos.

• Embora a ativação seja esparsa, todos os parâmetros de todos os especialistas devem ser armazenados na memória da GPU (ex: DeepSeek-V3 ativa 37B parâmetros, mas armazena 671B).
• Métodos de poda existentes focam principalmente na poda em nível de especialista (remover especialistas inteiros). Essa granularidade é muito grosseira, levando frequentemente a degradações significativas de precisão ou exigindo operações complexas de fusão/decisão que não garantem eficiência.
• A poda em nível de parâmetro (esparsificação) é ineficiente em hardware e não oferece aceleração direta.

2. Metodologia: HEAPr

O HEAPr (Hessian-Based Efficient Atomic Expert Pruning) propõe uma abordagem inovadora que decompõe cada especialista em unidades menores e indivisíveis chamadas "Especialistas Atômicos".

Conceito de Especialista Atômico:

• Um especialista padrão $E_i(x)$ é decomposto em $d_{inter}$ especialistas atômicos $e^{(j)}_i(x)$.
• Cada especialista atômico é definido pelo agrupamento de uma coluna específica das matrizes $W_{up}$ e $W_{gate}$, e a linha correspondente na matriz $W_{down}$.
• A saída do especialista completo é a soma das saídas de seus especialistas atômicos.

Algoritmo e Teoria (Baseado em OBS):
O método utiliza princípios da teoria do Cirurgião Cerebral Ótimo (Optimal Brain Surgeon - OBS), que usa informações de segunda ordem (Hessiana) para identificar quais parâmetros podem ser removidos com o menor aumento na função de perda.

Para tornar isso viável em modelos modernos, o HEAPr introduz duas otimizações cruciais:

1. Decomposição e Desacoplamento:

   • O papel demonstra que as derivadas de segunda ordem entre parâmetros de diferentes especialistas atômicos são zero (eles são desacoplados).
   • Isso permite calcular a Hessiana para cada especialista atômico individualmente, reduzindo a complexidade espacial de $O((3d_{model} \cdot d_{inter})^2)$ para $O((3d_{model})^2 \cdot d_{inter})$.

2. Transição para o Espaço de Saída (Output Space):

   • Para reduzir ainda mais o custo, o HEAPr reformula a restrição de poda do espaço de parâmetros para o espaço de saída.
   • Em vez de forçar os parâmetros a zero, o método busca minimizar o aumento da perda ao forçar a saída do especialista atômico a zero para um token específico.
   • Utiliza a Matriz de Informação de Fisher (teoricamente equivalente à Hessiana esperada em redes convergidas) aplicada à saída do especialista atômico.
   • Como os gradientes da perda em relação às saídas de todos os especialistas atômicos dentro de um mesmo especialista são idênticos, é necessário armazenar apenas uma matriz de covariância de gradientes por especialista.
   • Complexidade Final: A complexidade espacial cai para $O(d_{model}^2)$, tornando o cálculo viável.

Fluxo de Execução:
O algoritmo requer apenas duas passagens diretas (forward) e uma passagem reversa (backward) em um pequeno conjunto de calibração:

1. Estimação da covariância de gradientes compartilhada (backward).
2. Cálculo da importância de cada especialista atômico usando a saída e a covariância (forward).
3. Classificação global e poda dos especialistas atômicos com menor importância.

3. Principais Contribuições

• Novo Granularidade de Poda: Introdução do conceito de "Especialista Atômico", permitindo uma poda mais flexível e precisa do que a poda em nível de especialista.
• Esquema de Aproximação de Segunda Ordem Eficiente: Transformação da informação de segunda ordem dos parâmetros do especialista para os parâmetros atômicos e, finalmente, para o espaço de saída, reduzindo a complexidade de $O(d^4)$ para $O(d^2)$.
• Algoritmo Escalável (HEAPr): Um método que estima a importância de todos os especialistas atômicos com custo computacional mínimo (2 forward, 1 backward) sem necessidade de fine-tuning ou re-treinamento.
• Classificação Global: A métrica de importância permite uma classificação global consistente através de todas as camadas do modelo, evitando vieses locais.

4. Resultados Experimentais

Os experimentos foram realizados em modelos MoE modernos, incluindo DeepSeekMoE-16B-Base, Qwen1.5-MoE-A2.7B-Chat, Qwen2-57B-A14B e Qwen3-30B-A3B.

• Desempenho com Poda Leve (20% - 25%):
  • O HEAPr alcança compressão quase sem perda (lossless).
  • No DeepSeekMoE-16B-Base (20% de poda) e Qwen1.5-MoE-A2.7B-Chat (25% de poda), o desempenho é praticamente idêntico ao modelo original.
  • No Qwen2-57B-A14B, mantém o desempenho quase inalterado mesmo com 40% de poda.
  • No Qwen3-30B-A3B, a precisão média cai apenas 0.03 pontos com 25% de poda.
• Redução de FLOPs: Redução de aproximadamente 20% nos FLOPs nas taxas de poda mencionadas.
• Comparação com SOTA: Superou métodos existentes como NAEE, MoE-I2, Sub-MoE, D2-MoE e a abordagem concorrente CAMERA-P em todos os benchmarks de zero-shot (WikiText, ARC, HellaSwag, etc.).
• Robustez: O método é robusto a diferentes conjuntos de dados de calibração (WikiText-2 vs. C4) e tamanhos de amostra.

5. Significado e Impacto

O trabalho HEAPr representa um avanço significativo na compressão de modelos MoE:

• Viabilidade de Implantação: Ao permitir a remoção de parâmetros com granularidade fina sem perda de precisão, torna viável a implantação de grandes modelos MoE em dispositivos com recursos limitados de memória.
• Eficiência Teórica: Resolve o problema de custo computacional associado ao uso de informações de segunda ordem (Hessiana) em redes profundas, oferecendo uma solução matematicamente fundamentada e prática.
• Mudança de Paradigma: Demonstra que a decomposição de especialistas em unidades atômicas, combinada com análise de espaço de saída, é uma estratégia superior à poda tradicional de especialistas inteiros, preservando a capacidade do modelo de forma mais eficiente.

Em resumo, o HEAPr oferece uma ferramenta poderosa para reduzir a barreira de memória dos LLMs MoE, permitindo compressão agressiva (até 40% em alguns casos) mantendo a inteligência do modelo intacta.