Resumen Técnico: No olvides tus embeddings: Borrado de conocimiento robusto mediante la edición precisa de embeddings

Declaración del Problema

A medida que los modelos de lenguaje (LM) se despliegan cada vez más, la capacidad de borrar permanentemente conocimientos específicos (por ejemplo, datos privados, contenido protegido por derechos de autor o temas sensibles para la seguridad) es crítica. Los métodos de borrado persistente existentes se centran principalmente en actualizar los parámetros del modelo, específicamente dentro de las capas de Perceptrón Multicapa (MLP), utilizando a menudo autoencoders dispersos (SAEs) o el direccionamiento de pesos (weight steering). Sin embargo, estos métodos suelen fallar en lograr un borrado robusto. El conocimiento objetivo suele ser recuperable mediante prompts adversarios o el reaprendizaje (ajuste fino sobre datos relacionados con el concepto), y la eficacia del borrado suele depender de la forma (fallando al generalizar de opción múltiple a generación de texto abierto).

Los autores hipotetizan que estas limitaciones derivan de que las técnicas existentes pasan por alto la capa de embeddings de tokens. Aunque se sabe que las MLP desempeñan un papel clave en la recuperación de conocimiento, los LM también codifican una parte no trivial de su conocimiento directamente en sus embeddings de tokens. Ignorar esta capa deja una vía para que el modelo "reaprenda" los conceptos borrados y obliga a los métodos existentes a depender de actualizaciones de parámetros agresivas y amplias que degradan la utilidad general del modelo.

Metodología: EMBER

El artículo presenta EMBER (EMBedding ERasure), un módulo plug-and-play diseñado para editar con precisión los embeddings de los tokens para aumentar los métodos de borrado existentes centrados en MLP. EMBER sigue un enfoque de "localizar-luego-editar" que consta de tres etapas:

1. Factorización de Matrices Dispersas (Sparse MF) para la Localización de Características:
   En lugar de utilizar SAEs, EMBER aplica la Factorización de Matrices Dispersas directamente a la matriz de embeddings $E$. Descompone los embeddings de un conjunto específico de tokens (aquellos que aparecen en oraciones relacionadas con el concepto y neutrales) en un diccionario de factores $Z$ (características de embedding) y una matriz de coeficientes $Y$.

   • Restricción de Dispersión: Para asegurar que las características correspondan a conceptos localizados e interpretables, se aplica una restricción de "el ganador se lleva todo" (winner-takes-all o WTA) a los coeficientes $Y$, manteniendo solo el 1% de las magnitudes por factor. Esto fuerza a la factorización a identificar características compartidas por subconjuntos pequeños y específicos de tokens.

2. Identificación de Características Relacionadas con el Concepto:
   El método califica las características recuperadas para distinguir las características específicas del concepto de las generales. Calcula un estadístico de razón de masa que compara la contribución de una característica a los tokens etiquetados como concepto frente a los tokens neutrales. Las características con una razón alta se filtran posteriormente mediante un LLM-como-juez para verificar que los tokens asociados coincidan semánticamente con el concepto objetivo.

3. Borrado Preciso mediante Sustracción:
   Una vez identificadas las características relacionadas con el concepto ($F_C$), EMBER resta su contribución de los embeddings de los tokens relevantes. Para un embedding de token $e_t$, la actualización es:
   $$e_t \leftarrow e_t - \delta \sum_{i \in F_C} Y_{i,t} z_i$$
   donde $\delta$ es un hiperparámetro de fuerza. Esta operación elimina el componente específico del concepto mientras mantiene intacto el resto del embedding. El método se aplica solo a los tokens que tienen coeficientes no nulos para las características del concepto identificadas.

EMBER está diseñado para combinarse con los métodos de borrado existentes basados en MLP (por ejemplo, RMU, CRISP, SNMF) para crear un conjunto que apunte tanto a la capa de embeddings como a las capas MLP.

Principales Contribuciones

• Identificación del Cuello de Botella de los Embeddings: El artículo proporciona el primer estudio empírico que establece la capa de embeddings como un cuello de botella crítico para el borrado robusto de conceptos. Demuestra que ignorar los embeddings permite que los modelos recuperen el conocimiento borrado más fácilmente.
• Algoritmo EMBER: Un método novedoso que utiliza la Factorización de Matrices Dispersas para desenredar y eliminar precisamente las características relacionadas con el concepto de los embeddings de los tokens sin requerir un ajuste fino completo del modelo.
• Aumento Plug-and-Play: Se demuestra que EMBER es compatible con los marcos de borrado existentes, mejorando significativamente su rendimiento sin requerir un rediseño completo del flujo de trabajo de borrado.

Resultados Experimentales

Los autores evaluaron EMBER en Gemma-2-2B-it y Llama-3.1-8B-Instruct a través de 18 conceptos diversos (que van desde obras de ficción como Harry Potter hasta temas de seguridad como Cannabis).

• Eficacia y Especificidad: Aumentar los métodos existentes (RMU, CRISP, SNMF) con EMBER mejoró consistentemente la eficacia del borrado (reduciendo la precisión del concepto) y la especificidad (manteniendo el rendimiento en dominios similares y MMLU).
• Robustez al Reaprendizaje: Este es el principal beneficio. EMBER redujo drásticamente la capacidad del modelo para reaprender los conceptos borrados.
  • En Llama-3.1, la precisión del reaprendizaje para los métodos previos (RMU, CRISP) se redujo del 70%–76% al 35% cuando se combinó con EMBER (una reducción de aproximadamente el 50%).
  • En Gemma, la precisión del reaprendizaje se redujo entre un 16% y un 24%.
  • El flujo completo basado en MF (SNMF + EMBER) logró la precisión de reaprendizaje más baja (35% en Llama), superando significativamente a los baselines.
• Coherencia y Localización: El costo para la coherencia de la generación fue mínimo. El análisis mostró que EMBER modifica como máximo el 0.14% del vocabulario por concepto. La degradación de la coherencia fue altamente localizada en tokens "exclusivos del concepto" (por ejemplo, "Hogwarts" al borrar Harry Potter), mientras que los tokens generales (por ejemplo, "Harry" en el contexto de Prince Harry) permanecieron mayormente inafectados.
• Generalización: Aunque los hiperparámetros se ajustaron en tareas de opción múltiple (MC), las mejoras se transfirieron eficazmente a la generación de respuestas abiertas (OE), lo que sugiere que las ediciones apuntan a conocimientos conceptuales en lugar de atajos específicos de formato.

Significado y Reclamaciones

El artículo afirma que la edición precisa y dirigida de la matriz de embeddings es fundamental para lograr un borrado de conceptos robusto. Al desplazar el enfoque más allá de la visión centrada en MLP de los trabajos anteriores, EMBER aborda el problema del "borrado superficial" donde el conocimiento es simplemente suprimido en lugar de eliminado.

Los autores concluyen que:

1. Las intervenciones a nivel de embedding son necesarias para prevenir ataques de reaprendizaje.
2. Los métodos de borrado existentes pueden mejorarse significativamente mediante el aumento con ediciones a nivel de embedding.
3. El compromiso entre la eficacia del borrado y la utilidad del modelo (coherencia) puede mejorarse localizando las ediciones en características específicas relacionadas con el concepto en lugar de aplicar actualizaciones de pesos amplias.

El trabajo establece una nueva dirección para la edición de modelos, destacando que la capa de embeddings es un locus de almacenamiento de conocimiento vital, aunque previamente poco explorado, que debe abordarse para aplicaciones de seguridad y cumplimiento.