Stuck on "A": Diagnosing and Repairing Interface Injury in Attention-to-KDA Linearization of a 0.6B Language Model

Resumo Técnico: Diagnóstico e Reparo de Lesão de Interface na Linearização da Atenção-KDA

Declaração do Problema
O artigo investiga a natureza "com perda" (lossy) de converter modelos de linguagem pré-treinados de atenção total para arquiteturas de atenção linear, visando especificamente o mecanismo de Atenção Delta Kimi (KDA). Embora a atenção linear ofereça vantagens estruturais para inferência de contexto longo e implantação em dispositivos de borda (estado recorrente $O(1)$ vs. cache KV $O(n)$), o processo de conversão frequentemente falha em preservar o desempenho nas tarefas, apesar de métricas padrão sugerirem sucesso.

O autor foca na conversão do Qwen3-0.6B-Base (28 camadas) em uma arquitetura híbrida onde 21 camadas são substituídas por KDA e 7 camadas mantêm a Atenção de Consulta Agrupada (GQA). Isso é executado sob restrições rigorosas de nível de consumidor: uma única GPU de 32 GB e um orçamento de treinamento de dezenas de milhões de tokens (em vez de bilhões).

A questão diagnóstica central é: O que exatamente a conversão quebra? Avaliações padrão revelam um paradoxo: após o alinhamento camada a camada e a destilação KL de ponta a ponta, o modelo estudante atinge a perplexidade próxima à do professor (lacuna de Entropia Cruzada de +0,128 nats), mas colapsa na acurácia de múltipla escolha (permanecendo próximo ao acaso, ~25–29%, no C-Eval, comparado aos 50,6% do professor).

Metodologia
O autor propõe um pipeline de múltiplos estágios para diagnosticar e reparar essa lacuna de desempenho:

1. Arquitetura Híbrida e Cirurgia:

   • Layout: As camadas 3, 7, 11, 15, 19, 23 e 27 mantêm a GQA nativa; as 21 camadas restantes são convertidas para KDA (16 cabeças, dimensão da cabeça 128).
   • Mecânica KDA: Implementa uma regra delta com portão (gate) por canal com um portão de esquecimento ($\alpha_t$) e um portão de escrita ($\beta_t$).
   • Inicialização: Um estudo de ablação crítico identificou que transplantes simples de projeção falharam. A configuração bem-sucedida (g6) utiliza escalas de decaimento e configurações de viés (bias) específicas para estabilizar o estado recorrente.

2. Estágios de Treinamento:

   • Estágio 2 (Alinhamento Camada a Camada): Alinha os estados ocultos da KDA com os estados ocultos do professor congelado usando MSE. Isso reduziu a Entropia Cruzada (CE) de 9,48 para 4,13, mas deixou a acurácia do C-Eval em 25,5%.
   • Estágio 3a (Destilação KL de Ponta a Ponta): Treinado em texto de pré-treinamento usando divergência KL direta entre as distribuições do professor e do estudante. Isso fechou a lacuna de CE para +0,128 nats, mas elevou apenas a acurácia do C-Eval para 28,8%.
   • Estágio 3b (Reparo Focado no Formato): Um estágio de KL de 1.000 passos de "apenas conclusão" usando Questões de Múltipla Escolha (MCQs) verificadas pelo professor, poesia e tradução de QA. Crucialmente, este estágio focou no formato da resposta, em vez de apenas geração de texto geral.

3. O Diagnóstico de Quatro Permutações:
   Para ir além da perplexidade, o autor desenhou um protocolo de diagnóstico usando 161 MCQs limpos (sem sobreposição com benchmarks). Cada questão foi avaliada sob quatro rotações cíclicas das opções de resposta (A, B, C, D).

   • Métrica: "Aderência" (Stickiness) mede com que frequência o modelo prevê o mesmo rótulo (ex: "A") independentemente da posição do conteúdo da opção.
   • Descoberta: O modelo do Estágio 3a exibiu uma severa lesão de interface: ele previa "A" 81% das vezes e mantinha o mesmo rótulo em todas as quatro rotações para 106/161 questões. O modelo perdeu a capacidade de mapear o conhecimento para os rótulos das opções, efetivamente ignorando o conteúdo das escolhas.

4. Alinhamento de Persona:
   Pós-reparo, o modelo passou por Ajuste Fino Supervisionado (SFT) com uma persona injetada ("Qingyi") e uma rodada de Otimização de Preferência Direta (DPO) on-policy para garantir que a conversão não degradasse a identidade ou as habilidades conversacionais do modelo.

Resultados Principais

• A "Mentira da Perplexidade": O estudo confirma que o alinhamento de estado oculto e a baixa divergência KL não garantem a capacidade de tarefa. O modelo pode imitar as estatísticas da distribuição do professor enquanto falha em utilizar esse conhecimento para tarefas específicas.
• Diagnóstico de Lesão de Interface: O teste de quatro permutações revelou que o modelo não sofria de perda de conhecimento uniforme (tentativa aleatória), mas de aderência de rótulo. Ele falhou em rastrear o conteúdo da resposta, recorrendo a um rótulo de opção específico.
• Eficácia do Reparo: O reparo de estágio 3b focado no formato:
  • Aumentou a acurácia do C-Eval de 28,83% para 41,31%.
  • Reduziu a participação de previsão de "A" de 81,06% para 58,5%.
  • Reduziu pela metade a métrica de "mesmo rótulo em 4 rotações" (106/161 $\to$ 47/161).
  • Transformou a margem da opção correta de negativa (−0,117 nats) para positiva (+0,233 nats).
  • Compromisso (Trade-off): Este reparo veio com um custo mensurável para o ajuste da distribuição geral (CE/KL válidos aumentaram ligeiramente), sugerindo que o modelo se especializou no formato de MCQ.
• Desempenho Final: Após o alinhamento de persona, o modelo final atingiu 41,83% no C-Eval, permanecendo dentro da banda de ruído do estágio de reparo pré-SFT, indicando que não houve esquecimento catastrófico.

Lições de Engenharia e Modos de Falha
O artigo documenta vários problemas críticos de engenharia encontrados durante o processo:

• Falha do Modo Mestre FP32: Usar bf16 tanto para parâmetros quanto para momentos do otimizador Adam fez com que pequenas atualizações de peso fossem "silenciosamente engolidas" (atualizações sub-ULP), resultando em pesos congelados apesar de gradientes saudáveis. A correção foi o uso de pesos mestres em FP32.
• Paridade de Cache: Novas camadas KDA inicialmente ignoravam atualizações incrementais de cache, fazendo com que a geração colapsasse em lixo a partir do segundo token. Uma correção específica de implementação de cache resolveu isso sem necessidade de retreinamento.
• Transplantes Falsos: Tentativas iniciais de transplante de projeções Q/K/V/O do professor falharam funcionalmente; a abordagem bem-sucedida abandonou totalmente o transplante de projeção.

Significância e Alegações
O artigo reivindica quatro contribuições primárias:

1. Ferramenta Diagnóstica: Um protocolo de quatro permutações que separa a "lesão de interface" (incapacidade de mapear conhecimento para o formato de saída) da "perda de conhecimento" (esquecimento uniforme), aplicável a qualquer pipeline de linearização ou compressão.
2. Receita Reproduzível: Uma receita completa, sob restrição de orçamento, para converter Qwen3 para KDA, documentando modos de falha específicos (ablação de portão, bugs de cache, absorção de atualização) que são frequentemente omitidos em registros públicos.
3. Sobrevivência da Persona: Evidência de que o alinhamento de persona (SFT + DPO) pode sobreviver à conversão de arquitetura "virtualmente de graça", sem exigir retreinamento massivo ou causar esquecimento catastrófico.
4. Lançamento Aberto: Liberação total de código, pesos, receitas e o rastro de auditoria de falhas.

O autor conclui que, embora converter um modelo de 0,6B para um híbrido KDA com um orçamento de consumidor seja viável, as pilhas de métricas padrão (perplexidade, KL, MSE de estado oculto) relatam sistematicamente o dano de forma errônea. A principal barreira para o sucesso não é necessariamente a perda de conhecimento interno, mas a degradação da interface que mapeia esse conhecimento para formatos de saída específicos (como rótulos de múltipla escolha), o que pode ser reparado com supervisão de formato direcionada.