Desactivar 729 cabezas de atención en GPT-2-XL redujo un 91,5% la memorización literal

Los modelos de lenguaje son capaces de reproducir literalmente fragmentos de sus datos de entrenamiento, lo que supone un riesgo de filtración de información personal y de reclamaciones por derechos de autor. Unos investigadores han mostrado que una parte de este problema se puede reducir sin reentrenamiento: basta con anular las cabezas de atención «sobrantes» antes de la inferencia. En GPT-2-XL (1,5 mil millones de parámetros), desactivar 729 de esas cabezas redujo un 91,5% el número de secuencias reproducidas literalmente, y la perplejidad en WikiText-103 mejoró en 0,92.

Qué cabezas se eliminaron y cómo se encontraron

La idea se apoya en dos observaciones. Los trabajos teóricos sostienen que la atención multicabeza funciona como una memoria asociativa cuya capacidad crece con el número de cabezas. Los trabajos empíricos llevan tiempo demostrando que una parte considerable de las cabezas se puede apagar casi sin pérdida de calidad. Nadie había intentado conectar ambos hechos.

Los autores congelan todos los pesos y entrenan únicamente máscaras binarias sobre las salidas de las cabezas. La máscara es una función escalón, así que los gradientes pasan a través de una aproximación sigmoidea (straight-through estimator). Una penalización en la función de pérdida impide que la perplejidad supere un umbral fijado y, al mismo tiempo, favorece un menor número de cabezas activas.

Usan el optimizador AdamW, lotes de 8 secuencias y 1000 pasos. No se toma la última configuración, sino aquella en la que se registra la menor relación entre la perplejidad y la inicial. Las máscaras se ajustan en una parte de WikiText-103 o de LAMBADA (un conjunto de datos para adivinar la última palabra de un párrafo) que no se solapa con la de evaluación.

La memorización se midió mediante k-extractability: si se le da al modelo el comienzo de una secuencia, este produce la continuación palabra por palabra con decodificación voraz. Las pruebas se hicieron con GPT-2-XL y el corpus OpenWebText, y con Pythia-2.8B y Pythia-6.9B (ambas Deduped) y el corpus Pile-Deduped.

Los resultados en cifras

En GPT-2-XL, las máscaras ajustadas con WikiText-103 eliminaron alrededor del 61% de las cabezas. La perplejidad bajó de 10,08 a 9,16, y el número de secuencias memorizadas en una de las configuraciones cayó de 5590 a 315. En las dos Pythia el panorama es el mismo: la perplejidad pasó de 7,56 a 7,32 y de 6,90 a 6,73, y la memorización bajó de 20 000 a 2875 y 3500 secuencias. Entre todos los modelos y los dos benchmarks, la reducción osciló entre el 44 y el 94%.

Las máscaras ajustadas con LAMBADA desactivan menos cabezas y la memorización baja menos, pero la precisión al predecir la última palabra se mantiene o incluso sube ligeramente. Los conjuntos de cabezas eliminadas difieren entre los dos benchmarks y los tres modelos, pero el efecto aparece en todos los casos. Los autores concluyen que la memorización está distribuida por todo el mecanismo de atención y no concentrada en un grupo reducido de cabezas.

Un experimento de control muestra que importa no solo cuántas cabezas se desactivan, sino cuáles. Al desactivar el mismo número al azar, la memorización a veces cayó incluso más (330 frente a 473 secuencias en GPT-2-XL en una de las configuraciones), pero el modelo quedó destrozado. La perplejidad de GPT-2-XL en WikiText-103 subió de 10,08 a 146,82, y en las Pythia, de 7,56 a 18,30 y de 6,90 a 17,20. La precisión de GPT-2-XL en LAMBADA bajó de 0,6128 a 0,2014. Un modelo roto sencillamente no reproduce nada, de modo que el contador de memorización por sí solo no basta para evaluar el método.

Por qué funciona y dónde están los puntos débiles

El análisis mecanicista se hizo con Pythia-2.8B y GPT-2-XL. En las secuencias memorizadas, las cabezas «saturan» el softmax con más frecuencia: casi toda la atención se concentra en una sola posición.

Además, distintas cabezas tienden a apuntar a la misma posición. Tras la desactivación, esa brecha entre el texto memorizado y el no memorizado se reduce. La diferencia entre el logit del mejor token y el del segundo en las secuencias memorizadas era de 8,79 y 7,68, aproximadamente cuatro veces superior a la del texto normal.

Tras la desactivación bajó a 3,60 y 3,86, y en el texto no memorizado apenas cambió. Es decir, lo que se ve afectado es el mecanismo de exceso de confianza, y no la confianza general del modelo.

Los propios autores señalan las limitaciones. La parte mecanicista es correlacional: los datos muestran que la saturación, la coherencia entre cabezas y la confianza cambian a la vez, pero no se ha demostrado causalidad. La comparación de coherencia se hace sobre el conjunto de cabezas saturadas, que se reduce tras la desactivación, así que parte del efecto podría ser un artefacto. Tampoco se ha comprobado por qué las cabezas restantes miran «más ampliamente». El análisis no cubre Pythia-6.9B por falta de capacidad de cómputo.

Hay también una limitación metodológica. La calidad se midió en el mismo benchmark con el que se buscaron las cabezas sobrantes, aunque en una muestra distinta. El trabajo no comprueba cómo se comportará el modelo recortado en otras tareas. Los experimentos se limitan a tres modelos pequeños de 2020–2023, y en los LLM actuales con atención agrupada el resultado puede ser diferente. Los autores prometen publicar el código cuando se acepte el artículo.

La conclusión para los profesionales es esta: los equipos que mantienen un modelo en producción y temen filtraciones de datos de entrenamiento disponen ahora de una herramienta barata que no requiere reentrenar. Pero esto reduce la probabilidad de reproducción literal, no garantiza la privacidad, y antes de aplicarlo hay que probar el modelo con las tareas propias.