Resumo Técnico: CLHA – Agregação de Hessian de Camadas Cruzadas para Quantização de Transformers com Compartilhamento de Pesos em Mixture-of-Experts

1. Definição do Problema

Arquiteturas recentes de Transformer com Mixture-of-Experts (MoE) eficientes, como as do Megrez2, empregam Compartilhamento de Pesos entre Camadas Cruzadas (CLWS). Neste design, camadas adjacentes de Transformer (por exemplo, $\ell$ e $\ell+1$) reutilizam os mesmos tensores de pesos de especialistas roteados para reduzir o footprint de parâmetros pela metade, mantendo o mesmo número de parâmetros ativados por token.

No entanto, os métodos existentes de Quantização Pós-Treinamento (PTQ), incluindo abordagens do tipo GPTQ, tratam esses pesos compartilhados como se pertencessem a uma única camada. Eles calibragem a quantização usando estatísticas de ativação de apenas uma "camada primária". Isso cria um descompasso fundamental:

• O fluxo residual da representação evolui entre as camadas, o que significa que as distribuições de tokens que ativam o especialista compartilhado na camada $\ell_a$ diferem das da camada $\ell_b$.
• Os métodos existentes ignoram as estatísticas de ativação da camada secundária, levando a uma minimização subótima do erro de reconstrução, particularmente em regimes de baixa bitagem (ex: 2 bits).

2. Metodologia

2.1 Fundamentação Teórica: Agregação de Hessian de Camadas Cruzadas

O artigo formaliza o objetivo para quantizar um peso $W$ compartilhado entre as camadas $\ell_1, \dots, \ell_n$. O objetivo correto é minimizar o erro de reconstrução total através de todas as camadas utilizando o peso:
$$L(W_q) = \sum_{i} \alpha_i \text{tr}[(W_q - W)H_{\ell_i}(W_q - W)^\top]$$
onde $H_{\ell_i}$ é o Hessian (covariância de ativação) da camada $\ell_i$, e $\alpha_i$ representa a frequência de roteamento (proporção de tokens) para essa camada.

A Proposição 1 estabelece que minimizar este erro combinado é matematicamente equivalente a aplicar GPTQ com um Hessian combinado:
$$H_{comb} = \sum_{i} \alpha_i H_{\ell_i}$$
Isso contrasta com a Calibração Independente por Camada (PLIC), que minimiza o erro usando apenas o Hessian de uma única camada ($H_{\ell_a}$).

2.2 Algoritmo CLHA

Os autores propõem o CLHA (Cross-Layer Hessian Aggregation), um algoritmo de calibração de três passagens:

1. Coleta (Collection): Para cada sequência de calibração, acumula Hessians por camada e contagens de tokens para cada especialista ativo.
2. Agregação: Para cada par de especialistas compartilhados $(\ell_a, \ell_b)$, calcula um Hessian combinado. O artigo avalia duas estratégias de ponderação:
   • Ponderação por Frequência de Roteamento ($\alpha = n_a / (n_a + n_b)$): Pesa os Hessians pela contagem real de tokens recebida por cada camada.
   • Ponderação Igual ($\alpha = 0.5$): Atua como um regularizador, útil quando os dados de calibração são limitados ou as estimativas do Hessian da camada secundária são ruidosas.
3. Quantização: Aplica atualizações sequenciais de coluna padrão do GPTQ usando o $H_{comb}$ agregado.

2.3 CLHA-MP (Precisão Mista)

Reconhecendo que os especialistas compartilhados sempre ativos recebem todos os $N$ tokens (ao contrário dos especialistas roteados que recebem uma fração $K/E$), o artigo observa que sua sensibilidade de quantização é $E/K$ vezes maior. O CLHA-MP aloca $b+1$ bits para esses especialistas compartilhados, mantendo os especialistas roteados com $b$ bits, incorrendo em um overhead de armazenamento mínimo (10–20% dos parâmetros dos especialistas).

2.4 Correção de Implementação: Ambiguidade de Hook

O artigo identifica um problema de engenharia crítico: os hooks de forward padrão do PyTorch anexados a um módulo ExpertFFN compartilhado disparam para qualquer camada que chame esse módulo. Isso faz com que as ativações de diferentes camadas sejam misturadas em um único acumulador de Hessian, corrompendo a atribuição por camada.

• Solução: Os autores propõem um mecanismo de callback consciente da camada (layer-aware) passado através da assinatura MoELayer.forward(). Isso fornece explicitamente os índices das camadas ao callback, garantindo que Hessians separados sejam acumulados para $H_a$ e $H_b$, mesmo quando o módulo de peso subjacente é compartilhado.

3. Principais Contribuições

1. Proposição Teórica: Formalizou que o Hessian ideal para pesos compartilhados é a soma ponderada pela frequência de roteamento dos Hessians por camada.
2. Algoritmo CLHA: Uma implementação prática de calibração de Hessian agregado dentro do framework GPTQ.
3. CLHA-MP: Uma extensão de precisão mista que atribui um bit extra aos especialistas compartilhados altamente sensíveis.
4. Correção de Implementação: Identificou e resolveu o problema de "ambiguidade de hook" que corrompe silenciosamente a calibração em arquiteturas de módulos compartilhados.
5. Validação Empírica: Demonstrou que os ganhos do CLHA correlacionam-se com a divergência de roteamento entre camadas.

- Resultados Experimentais

Experimentos foram conduzidos em um CLWS-MoE de 34.5M de parâmetros treinado do zero no WikiText-2.

• Desempenho em 2-bit: Este é o regime crítico onde a má calibração se acumula.
  • O CLHA-rt (ponderado pelo roteamento) reduziu a Perplexidade (PPL) em 4,3% comparado ao PLIC.
  • O CLHA-eq (ponderação igual) reduziu a PPL em 5,4%.
  • O CLHA-MP alcançou uma melhoria de 18,6% sobre o PLIC.
  • Em 4-bit e 3-bit, as diferenças ficaram dentro das margens de ruído (<1%).
• Calibração vs. Erro de Teste: Curiosamente, o PLIC alcançou um Erro de Hessian de Calibração (CHE) menor no conjunto de treinamento porque otimiza diretamente para aquela distribuição. No entanto, o CLHA generalizou melhor para o conjunto de teste, indicando que o Hessian combinado captura melhor a verdadeira distribuição do peso compartilhado do que uma visão de camada única.
• Experimento de Nitidez de Roteamento (Routing Sharpness): Ao aumentar artificialmente os logits do roteador (aumentando a divergência entre camadas), os autores observaram:
  • O desempenho do PLIC degradou em 5% conforme a divergência aumentou.
  • O CLHA permaneceu estável (<1% de degradação).
  • A lacuna de desempenho entre CLHA e PLIC cresceu de +4,3% (roteamento natural) para +8,0% (alta divergência), confirmando a previsão teórica de que os ganhos do CLHA escalam com a divergência entre camadas.

5. Significância e Alegações

O artigo afirma fornecer a primeira análise de PTQ especificamente para compartilhamento de pesos entre camadas em modelos MoE. Sua significância reside em:

• Correção do Objetivo: Demonstrar que tratar pesos compartilhados como entidades de camada única é matematicamente subótimo e que a agregação de Hessian é a correção necessária.
• Escalabilidade dos Ganhos: Mostrar que, à medida que os modelos se tornam maiores e os padrões de roteamento mais especializados (aumentando a divergência entre camadas), os benefícios do CLHA provavelmente crescerão, tornando-o essencial para futuros modelos CLWS-MoE (ex: Megrez2, DeepSeek, Qwen).
• Necessidade Prática: Destacar que a aplicação ingênua de ferramentas de PTQ existentes a arquiteturas CLWS é fundamentalmente falha devido à ambiguidade de hook, e fornecer a correção necessária para qualquer método de calibração futuro neste domínio.

Os autores mantêm a modéstia quanto ao desempenho absoluto, observando que seu modelo é pequeno (34.5M) e treinado por passos limitados, mas enfatizam que as melhorias relativas sobre a calibração padrão são robustas e teoricamente fundamentadas.