Resumo Técnico: Além do Alinhamento de Saída na Destilação de LLMs em NVFP4

1. Definição do Problema

À medida que os Grandes Modelos de Linguagem (LLMs) são implantados em ambientes de produção com restrições de latência e memória, há uma demanda crescente por formatos de inferência de baixa precisão, especificamente o NVFP4 (um formato de ponto flutuante de 4 bits com escalonamento de bloco de grão fino). Embora a Destilação com Consciência de Quantização (QAD) tenha emergido como um método robusto para recuperar a acurácia perdida durante a quantização pós-treinamento (PTQ) através do alinhamento da distribuição de saída de um estudante quantizado com um professor de precisão total via divergência KL, este artigo identifica uma falha crítica no paradigma prevalecente.

Os autores argumentam que o alinhamento apenas ao nível da saída é insuficiente para uma recuperação fiel de capacidades. Sua análise revela que a QAD padrão, que otimiza exclusivamente a divergência KL, pode mascarar uma degradação representacional interna significativa. Especificamente, embora o modelo estudante possa alinhar com sucesso os logits de saída do professor, suas geometrias de ativação intermediária sofrem uma reorganização indesejada (drift). Esse "aprendizado de atalho" (shortcut learning) é particularmente severo em modelos que passaram por pós-treinamento de Aprendizado por Reforço (RL), onde as hierarquias semânticas internas necessárias para tarefas complexas de raciocínio e codificação são comprometidas, apesar da alta acurácia de saída.

2. Metodologia: CKA-QAD

Para abordar o descompasso entre o alinhamento de saída e a geometria interna, os autores propõem o CKA-QAD, um método que aumenta a QAD padrão com um regularizador de Alinhamento de Kernel Centrado (CKA).

Mecanismo Central

Em vez de depender exclusivamente da destilação ao nível de logit, o CKA-QAD introduz um regularizador leve que alinha as matrizes de Gram por camada das ativações do estudante e do professor.

• Métrica CKA: O método utiliza CKA linear, que mede a similaridade entre matrizes de ativação centradas ($\bar{H}_t$ e $\bar{H}_s$). O CKA é escolhido por sua invariância a transformações ortogonais e escalonamento isotrópico, tornando-o robusto às distribuições de ativação não estacionárias e mudanças de escala inerentes à quantização NVFP4.
• Função de Objetivo: A perda total combina a perda de divergência KL Top-K padrão ($L_{TopK-KL}$) com um termo de alinhamento CKA dinâmico ($L_{CKA}$).
  $$L_{total} = L_{TopK-KL} + \text{sg}\left(\frac{L_{TopK-KL}}{L_{CKA} + \epsilon}\right) L_{CKA}$$
  O uso de um operador de stop-gradient (sg) e escalonamento dinâmico permite que o termo CKA seja equilibrado com a perda de logit sem a necessidade de ajuste manual de hiperparâmetros de um peso $\lambda$ global.

Detalhes de Implementação

• Seleção de Camadas: O CKA é computado nos outputs de blocos decodificadores correspondentes (pós-atenção/MLP/atualizações Mamba). Os autores excluem explicitamente ativações de pré-normalização, pesos de atenção e estados do KV cache para evitar artefatos de escalas de fluxo residual e dependências autoregressivas.
• Otimizações de Engenharia: Para gerenciar o overhead computacional de calcular o CKA em muitas camadas:
  1. Destilação Top-K: A perda KL é computada apenas nos 8192 logits superiores do professor.
  2. CKA no Espaço de Características: O CKA é calculado via matrizes de covariância de características ($\bar{H}^\top \bar{H}$) em vez de matrizes de Gram explícitas token-a-token, reduzendo a complexidade de memória.
  3. Captura Baseada em Hooks: As ativações são capturadas via forward hooks, com as ativações do professor desconectadas do grafo de computação.

3. Principais Contribuições

1. Diagnóstico de Modos de Falha da QAD: O artigo fornece um diagnóstico de nível de representação mostrando que a destilação baseada apenas em KL pode exacerbar o drift representacional, particularmente em camadas FFN de profundidade média de Transformers e blocos de estado Mamba. Este drift correlaciona-se com gargalos de desempenho em tarefas de raciocínio e codificação.
2. QAD NVFP4 Guiada por CKA: A introdução do CKA-QAD, um método que preserva explicitamente a geometria representacional interna durante a destilação de baixo bit. Ele restaura o alinhamento para níveis próximos ao BF16 com overhead mínimo.
3. Validação Empírica: Avaliação abrangente em duas arquiteturas distintas:
   • Nemotron 3 Nano: Um modelo híbrido Mamba–Transformer MoE com pós-treinamento de RL.
   • Qwen3-4B-Thinking-2507: Um modelo de raciocínio compacto Transformer.

4. Resultados Experimentais

Os autores avaliaram o CKA-QAD contra professores BF16, baselines de PTQ NVFP4 e a QAD padrão baseada apenas em KL.

Alinhamento Representacional

• Nemotron 3 Nano: A QAD padrão (apenas KL) resultou em uma queda significativa no CKA médio (0,958) em comparação ao PTQ (0,993), com um colapso severo na última camada (0,740). O CKA-QAD restaurou o CKA médio para 0,994 e o CKA da última camada para 0,985, quase igualando o professor BF16.
• Qwen3-4B: Tendências semelhantes foram observadas, com a QAD apenas KL reduzindo o CKA médio para 0,98 e o da última camada para 0,86. O CKA-QAD restaurou estes para 0,99 e 0,97, respectivamente.

Acurácia em Tarefas Downstream

O CKA-QAD demonstrou melhorias consistentes em benchmarks de raciocínio complexo e codificação:

• AIME25 (Raciocínio Matemático): Melhorou de 68,5% (apenas KL) para 72,3% (CKA-QAD) no Qwen3-4B.
• GPQA-D (Conhecimento Especialista): Melhorou de 59,5% para 61,1%.
• LiveCodeBench-v5 (Geração de Código): Melhorou de 57,9% para 59,8%.
• Nemotron 3 Nano: Mostrou recuperação similar, com o AIME25 melhorando de 87,9% para 88,3% e o LiveCodeBench de 68,9% para 70,3%.

Overhead de Eficiência

O método introduz custos de treinamento modestos:

• Tempo de execução (Wall-clock time): +0,5% por passo.
• VRAM de Pico: +7,0%.
• Throughput: -0,5%.
• Convergência: Passos para atingir o alvo de KL foram ligeiramente reduzidos (-3,3%).

5. Significância e Alegações

O artigo posiciona o alinhamento representacional guiado por CKA como um complemento prático para o alinhamento de saída para recuperação de LLMs quantizados. Os autores afirmam que:

• O alinhamento funcional é necessário, mas insuficiente: Combinar distribuições de saída não garante a preservação da topologia semântica interna necessária para o raciocínio profundo.
• A preservação da geometria auxilia a generalização de baixo bit: Ao impedir que o estudante aprenda mapeamentos de "atalho" que satisfazem a função de perda mas degradam a estrutura interna, a regularização CKA ajuda a quantização agressiva a reter hierarquias semânticas finas codificadas durante o pós-treinamento (especialmente RL).
• Robustez a Artefatos de Quantização: Diferente de perdas de características ponto a ponto (ex: MSE), a invariância do CKA à escala e rotação o torna unicamente adequado para o NVFP4, onde o escalonamento por bloco induz distribuições não estacionárias.

Os autores concluem que preservar a geometria interna é crítico para tarefas que exigem propagação de características de múltiplos passos, como raciocínio matemático e síntese de código, e que o CKA-QAD oferece um caminho viável para o deployment de LLMs de baixo bit de alta fidelidade com overhead de engenharia mínimo.