Destilei um modelo 14 vezes menor, ele pareceu pior, e o problema estava no meu benchmark

Em outubro escrevi sobre como funciona a memória de um agente de IA e apresentei o agent-memory-bench, um benchmark que compara diferentes políticas de memória. Na média do benchmark, a política com extração de fatos recuperou 98% das informações relevantes, contra 62% da busca simples por similaridade. Nessa política, cada mensagem do usuário passa pelo qwen2.5:7b, que devolve fatos associados a uma chave, como payment_service.region = sa-east-1.

O problema é que chamar um modelo de 7 bilhões de parâmetros a cada mensagem é caro. Por isso, o passo seguinte foi verificar se um modelo 14 vezes menor conseguiria fazer o mesmo trabalho.

Conseguiu, mas o caminho até essa resposta acabou sendo mais instrutivo do que o resultado. Na primeira rodada, o modelo pequeno caiu para 83% de recuperação, e a investigação dessa queda mostrou que os 98% originais dependiam, em parte, de sorte.

A destilação

Destilação de conhecimento consiste em treinar um modelo pequeno, chamado de aluno, para reproduzir o comportamento de um modelo maior, o professor. Neste projeto, o professor foi o Qwen2.5-7B-Instruct, responsável por gerar os rótulos (os fatos extraídos de cada mensagem), e o aluno foi o Qwen2.5-0.5B-Instruct, ajustado sobre esses rótulos com LoRA.

O LoRA permite fazer fine-tuning sem alterar os pesos originais do modelo: em vez de atualizar todos os parâmetros, treinam-se matrizes pequenas que são somadas a algumas camadas. Isso reduz bastante o custo, a ponto de o treino do modelo de 0,5B caber numa GPU T4 gratuita do Kaggle e terminar em algo entre 10 e 20 minutos.

As mensagens de treino não vieram do benchmark. Elas foram geradas a partir de templates em português, cobrindo situações como mudança de região de um serviço, escolha de banco de dados, regras do time, conversas do dia a dia sem informação relevante e conteúdo externo contendo instruções direcionadas a assistentes de IA. Foram 3000 exemplos de treino, rotulados pelo professor. Para evitar vazamento, um teste automatizado descarta qualquer mensagem gerada que seja parecida demais com as do benchmark.

Os detalhes do método estão no README do repositório. O que interessa aqui é o que aconteceu depois do treino.

O resultado inesperado

Na primeira rodada completa no Kaggle, troquei o extrator do benchmark por cada um dos modelos e obtive os seguintes números:

A conclusão imediata seria que o aluno é pior que o professor. A segunda linha da tabela, porém, contradiz essa leitura: o próprio professor caiu para 83%. Além disso, o aluno sem nenhum treino obteve praticamente o mesmo resultado do aluno treinado. Se o problema estivesse no tamanho do modelo, o professor não teria caído junto.

Havia, portanto, alguma diferença entre a rodada publicada e a rodada do Kaggle que não dependia do modelo usado.

Isolando a causa

Entre as duas rodadas, três coisas haviam mudado: a forma de servir o modelo (Ollama no meu computador, vLLM no Kaggle), a implementação dos embeddings e o adaptador que conecta o modelo ao benchmark. Testei cada uma separadamente, em CPU, usando a mesma grade de experimentos do benchmark (5 seeds, 5 níveis de ruído e 5 intervalos de tempo).

O que tornou esse teste possível foi o cache de extração da rodada publicada, isto é, as respostas exatas que o 7B havia dado para cada mensagem. Ao passar essas mesmas respostas pelo caminho novo, qualquer diferença no resultado só poderia vir do próprio caminho.

As duas primeiras linhas mostram que, com as mesmas respostas, o caminho novo reproduz exatamente o resultado publicado. A diferença, então, estava no texto que o extrator produzia. E a queda se concentrava em um cenário específico: o da regra rara.

Nesse cenário, ao longo de oito semanas de conversa simulada, o usuário diz apenas uma vez, logo na primeira semana, que nunca se deve fazer force push na main. Meses depois, o benchmark faz uma pergunta como "fiz rebase da main local e o push foi rejeitado, o que devo rodar?", e a memória precisa recuperar a regra.

Com o professor do Kaggle, a regra era recuperada em apenas 33,7% das perguntas. Analisando as memórias de cada pergunta, encontrei dois problemas diferentes.

Primeiro problema: um fato inventado desativa a regra

O extrator transforma corretamente a regra no fato coding_rules.force_push = never. Algumas semanas depois, chega uma mensagem sem nenhuma relação com git:

"Terminamos a migração: a partir de hoje o serviço de pagamentos roda em sa-east-1."

O extrator identifica corretamente a nova região, mas também devolve um fato adicional:

coding_rules.force_push = false ("The force push rule is disabled for this project.") 

Nada na mensagem justifica esse fato. Para que os nomes das chaves sejam reaproveitados, o extrator recebe a lista de chaves já existentes, e às vezes escreve um valor arbitrário em uma delas. Para a memória, uma chave conhecida com um valor novo representa uma atualização, de modo que o fato inventado entra e a regra verdadeira é desativada. Esse fato falso apareceu no contexto de 275 das 875 perguntas sobre a regra.

O aluno não desenvolveu esse comportamento sozinho: ele o aprendeu com o professor. Ao revisar os 3000 rótulos de treino, encontrei 217 casos (7%) em que o professor escrevia um valor numa chave existente sem relação com a mensagem, 469 (16%) em que o texto do fato estava em inglês, apesar de o prompt pedir português, e 285 (10%) em que o tipo do fato (regra ou informação) não correspondia ao esperado pelo template. A destilação reproduz o professor por completo, inclusive os seus erros.

Segundo problema: a regra perde para uma mensagem repetida

Este foi o problema principal, e é ele que dá título ao artigo.

A busca da memória seleciona as cinco memórias mais parecidas com a pergunta, usando a similaridade de cosseno entre os embeddings (quanto mais próxima de 1, mais parecidos os textos). Calculei essa similaridade entre a pergunta sobre o force push e cada candidato:

A primeira linha é o texto que o 7B escreveu na rodada publicada. Quem leu o artigo sobre memória talvez se lembre dele, porque comentei lá que o extrator havia produzido essa frase meio em inglês, meio em português. Na época, tratei isso apenas como uma curiosidade.

Na verdade, era justamente essa frase mal formada que sustentava o resultado. Por coincidência, ela ficou mais próxima da pergunta do que a mensagem de ruído, enquanto qualquer versão bem escrita da mesma regra ficava abaixo de 0,635.

Além disso, a mensagem de ruído não aparece uma única vez: ela se repete toda semana. Com muito ruído, as suas cópias ocupam sozinhas as cinco vagas da busca, e a regra fica de fora. O professor do Kaggle escreveu a regra em inglês, com similaridade de 0,537, e por isso a perdia quase sempre.

Em outras palavras, os 92% de recuperação da regra rara que eu havia publicado dependiam da redação acidental de uma única frase. O benchmark não estava medindo exatamente o que eu imaginava.

As correções

Fiz correções em duas frentes: na política de memória e nos dados de treino.

Na memória

No agent-memory-bench, criei uma nova política, a v5_guarded, com duas mudanças. A primeira é uma proteção na escrita: um fato só pode substituir outro se a mensagem que o originou for semelhante ao fato antigo, com similaridade de pelo menos 0,55. A mensagem sobre a migração de região tem similaridade de 0,39 com a regra do force push e, portanto, não consegue mais alterá-la, enquanto as substituições legítimas do benchmark ficaram em 0,72 e 0,75. A segunda mudança atua na leitura: memórias com o mesmo texto passam a ocupar uma única vaga na busca, de modo que várias cópias da mesma mensagem de ruído contam como uma só.

A tabela mostra que a melhora na recuperação vem da remoção de repetições. A proteção na escrita não altera o número por causa de uma limitação do próprio benchmark: a mensagem original da regra continua armazenada e conta como acerto, mesmo que um fato falso dizendo o contrário esteja no mesmo contexto. Ainda assim, a proteção elimina esse fato falso, que passou de 275 perguntas para nenhuma. O benchmark não penaliza essa contradição, mas um agente real provavelmente seria confundido por ela.

Nos dados de treino

Antes do treino, os rótulos do professor passaram a ser filtrados por verificadores que descartam fatos escritos em chaves existentes sem relação com a mensagem e fatos com texto em inglês, além de corrigirem o tipo do fato com base no template. Também incluí uma nova categoria de exemplos: mensagens sobre uma área do projeto que possui regras na lista de chaves, nas quais o comportamento correto é não alterar essas regras.

Na segunda rodada, os verificadores removeram 242 rótulos por uso indevido de chave e 325 por texto em inglês, e corrigiram o tipo de outros 298. Restaram 2433 dos 3000 exemplos para o treino.

Houve também uma lição de ordem prática. O Kaggle apaga a pasta de trabalho ao fim de cada sessão, e a primeira rodada perdeu todas as predições e relatórios gerados. A partir da segunda, o notebook passou a publicar esses arquivos no dataset do Hugging Face ao final da execução, o que permitiu analisar os resultados com calma depois.

Resultados

Primeiro, a qualidade da extração em si, medida num conjunto de teste com áreas e regras que não aparecem no treino e comparada com o gabarito dos templates:

O aluno ficou cerca de cinco vezes mais rápido que o professor e, em relação ao gabarito, ligeiramente melhor. Isso se explica pelo fato de ele ter aprendido a partir de rótulos já filtrados, sem os vícios removidos pelos verificadores. A coluna de lista vazia também merece atenção: na maioria das mensagens (ruído, opiniões passageiras, conteúdo externo), o comportamento correto é não extrair nada, e o modelo sem treino devolvia algum fato em quase todas elas.

Executado no meu notebook, sem GPU (Ryzen 7 5825U, com o modelo quantizado em 4 bits), o aluno responde em cerca de 0,20 s por mensagem e ocupa aproximadamente 550 MB de memória.

No benchmark, os resultados foram os seguintes:

O 100% do aluno na v2 merece a mesma cautela discutida anteriormente. O aluno escreveu a regra como "Nunca realizar forcer push na main.", com um erro de digitação, e obteve similaridade de 0,673 contra 0,635 do ruído; sem o erro, o valor seria 0,649. A vantagem existe, mas é pequena e depende novamente da redação. Por isso, considero o resultado da v5 o mais confiável, já que ele não depende da forma como a regra foi escrita. Vale notar também que, na v5, até o modelo sem treino chega a 97,6%, o que indica que essa camada do benchmark mede mais a política de memória do que a qualidade do extrator. Para comparar extratores, as métricas de extração da tabela anterior são mais adequadas.

Lições

A principal lição foi a respeito de diagnóstico. Ter o professor e o modelo sem treino na mesma tabela foi o que me levou a deixar de culpar a destilação: quando o professor cai junto com o aluno, o problema dificilmente está no aluno. Da mesma forma, ter guardado as respostas da rodada publicada permitiu descartar rapidamente a hipótese de que a infraestrutura nova fosse a causa.

A segunda lição é que um resultado bom também precisa ser investigado. Se a primeira rodada não tivesse caído, eu provavelmente nunca teria olhado a similaridade da regra, e os 98% continuariam sendo apresentados sem que eu soubesse que dependiam de uma coincidência.

Por fim, a destilação reproduz os erros do professor, e filtrar os rótulos antes do treino foi o que permitiu ao aluno superá-lo. E erros que a métrica não penaliza, como o fato contraditório no contexto, permanecem invisíveis até que alguém olhe diretamente para os dados.

Limitações

O limiar de 0,55 usado na proteção da escrita foi calibrado com apenas três substituições do próprio benchmark. A remoção de repetições considera somente textos idênticos, de modo que uma mensagem de ruído reescrita de formas diferentes continuaria competindo pelas vagas da busca. O conjunto de teste do benchmark é pequeno (131 mensagens, das quais 35 distintas), e por isso um único erro altera o resultado em vários pontos percentuais. Os dados de treino, por virem de templates, cobrem bem a estrutura do problema, mas têm vocabulário limitado. Por fim, os experimentos mediram apenas a recuperação, ou seja, se a memória correta chega ao prompt, e não a qualidade da resposta final do agente.

Links

• Código e diagnóstico completo: DevTheusP/fact-extractor-distill
• Benchmark com a política v5: DevTheusP/agent-memory-bench
• Modelo: DevTheusP/qwen2.5-0.5b-fact-extractor-ptbr
• Dataset: DevTheusP/fact-extractor-distill-ptbr
• Notebook no Kaggle: Destilando um extrator de memória