Desligar 729 cabeças de atenção no GPT-2-XL reduziu em 91,5% a memorização literal

Modelos de linguagem conseguem reproduzir literalmente trechos dos dados de treinamento, o que representa risco de vazamento de informações pessoais e de reclamações por direitos autorais. Pesquisadores mostraram que parte desse problema pode ser resolvida sem ajuste fino: basta zerar as cabeças de atenção “redundantes” antes da inferência. No GPT-2-XL (1,5 bilhão de parâmetros), o desligamento de 729 dessas cabeças reduziu em 91,5% o número de sequências reproduzidas literalmente, e a perplexidade na WikiText-103 ainda melhorou em 0,92.

Quais cabeças foram removidas e como foram encontradas

A ideia parte de duas observações. Trabalhos teóricos indicam que a atenção multi-cabeça funciona como uma memória associativa, cuja capacidade cresce com o número de cabeças. Trabalhos empíricos já mostraram há tempos que uma parcela considerável das cabeças pode ser desligada quase sem perda de qualidade. Ninguém havia tentado conectar esses dois fatos.

Os autores congelam todos os pesos e treinam apenas máscaras binárias nas saídas das cabeças. Como a máscara é do tipo degrau, os gradientes passam por uma aproximação sigmoide (straight-through estimator). Uma penalidade na função de perda impede que a perplexidade suba além de um limite definido e, ao mesmo tempo, incentiva um número menor de cabeças ativas. O otimizador é o AdamW, com lotes de 8 sequências e 1000 passos. Vale a configuração com a menor razão entre a perplexidade e a original, não a última. As máscaras são ajustadas em uma parte separada da WikiText-103 ou da LAMBADA (dataset de previsão da última palavra de um parágrafo), sem sobreposição com a parte usada na avaliação.

A memorização foi medida por k-extractability: se o modelo recebe o início de uma sequência, ele gera a continuação palavra por palavra na decodificação gananciosa. Os testes usaram o GPT-2-XL com o corpus OpenWebText e os modelos Pythia-2.8B e Pythia-6.9B (ambos Deduped) com o corpus Pile-Deduped.

O que os números mostram

No GPT-2-XL, as máscaras ajustadas na WikiText-103 removeram cerca de 61% das cabeças. A perplexidade caiu de 10,08 para 9,16, e o número de sequências memorizadas, em uma das configurações, foi de 5590 para 315. Nos dois modelos Pythia o quadro é o mesmo: perplexidade de 7,56 → 7,32 e de 6,90 → 6,73, com memorização caindo de 20 000 para 2875 e 3500 sequências. Considerando todos os modelos e os dois benchmarks, a redução ficou entre 44 e 94%.

As máscaras da LAMBADA desligam menos cabeças, e a memorização cai menos, mas a precisão na previsão da última palavra se mantém ou até sobe um pouco. Os conjuntos de cabeças removidas diferem entre os dois benchmarks e os três modelos, mas o efeito aparece em todos os casos. Os autores concluem que a memorização está distribuída por todo o mecanismo de atenção, e não concentrada em um grupo pequeno de cabeças.

Um experimento de controle mostra que importa não só quantas cabeças são desligadas, mas também quais. Quando o mesmo número de cabeças é removido aleatoriamente, a memorização às vezes cai ainda mais (330 contra 473 sequências no GPT-2-XL, em uma das configurações), mas o modelo se desfaz. A perplexidade do GPT-2-XL na WikiText-103 subia de 10,08 para 146,82; nos Pythia, de 7,56 para 18,30 e de 6,90 para 17,20. A precisão do GPT-2-XL na LAMBADA despencava de 0,6128 para 0,2014. Um modelo quebrado simplesmente não reproduz nada, por isso só o contador de memorização não basta para avaliar o método.

Por que funciona e onde estão os pontos fracos

A análise mecanicista foi feita no Pythia-2.8B e no GPT-2-XL. Em sequências memorizadas, as cabeças saturam o softmax com mais frequência: quase toda a atenção vai para uma única posição. Além disso, cabeças diferentes tendem a apontar para a mesma posição. Depois do desligamento, essa diferença entre texto memorizado e não memorizado diminui. A distância entre o logit do melhor token e o do segundo colocado, nas sequências memorizadas, era de 8,79 e 7,68, cerca de quatro vezes maior que em texto comum. Após o desligamento, caiu para 3,60 e 3,86, e em texto não memorizado quase não mudou. Ou seja, o que é afetado é o mecanismo de excesso de confiança, e não a confiança geral do modelo.

Os próprios autores apontam as limitações. A parte mecanicista é correlacional: os dados mostram que saturação, coerência entre cabeças e confiança mudam juntas, mas a causalidade não foi demonstrada. A comparação de coerência é feita sobre o conjunto de cabeças saturadas, que encolhe após o desligamento, então parte do efeito pode ser um artefato. Também não foi verificado por que as cabeças restantes passam a olhar de forma mais “ampla”. A análise não inclui o Pythia-6.9B por falta de capacidade computacional.

Há ainda uma limitação metodológica. A qualidade foi medida no mesmo benchmark usado para encontrar as cabeças redundantes, ainda que em uma amostra separada. Como o modelo podado se comporta em outras tarefas não foi testado no trabalho. Os experimentos se limitam a três modelos pequenos, de 2020 a 2023, e em LLMs atuais com atenção agrupada o resultado pode ser diferente. Os autores prometem publicar o código após a aceitação do artigo.

Para quem trabalha com esses modelos, a conclusão é esta: equipes que mantêm um modelo em produção e temem vazamento de dados de treinamento ganham uma ferramenta barata que dispensa retreinamento. Mas ela reduz a probabilidade de reprodução literal, não garante privacidade, e antes de aplicá-la é preciso testar o modelo nas tarefas da própria equipe.