Resumen Técnico: Descubrimiento de Conceptos Discretos entre Capas para la Interpretación de Modelos de Lenguaje

Planteamiento del Problema

La interpretación de los grandes modelos de lenguaje (LLM) se ve obstaculizada por la naturaleza de las corrientes residuales de los transformadores, que mezclan y duplican linealmente las características a través de capas adyacentes. Los métodos de análisis de una sola capa suelen fallar al capturar esta estructura entre capas, perdiendo patrones computacionales que solo son visibles cuando se examinan varias capas en conjunto. Aunque los enfoques recientes entre capas, como los Autoencoders Dispersos (SAEs), abordan la mezcla de capas, operan en espacios continuos. En estos espacios, los conceptos se "dividen" en muchos neuronas sin límites claros, lo que obliga a utilizar umbrales arbitrarios o combinaciones de múltiples vectores para aislar las activaciones. Este desalineamiento con el razonamiento humano, que se basa en categorías discretas, limita la interpretabilidad de tales métodos.

Metodología

Los autores proponen CLVQ-VAE (Autoencoder Variacional de Cuantización Vectorial entre Capas), un marco diseñado para descubrir conceptos discretos e interpretables mediante el mapeo de representaciones de una capa inferior ($l$) a una capa superior ($h$) a través de un cuello de botella de cuantización vectorial discreta. El marco consta de tres componentes principales:

1. Codificador Residual Adaptativo: Para preservar la información semántica permitiendo al mismo tiempo la reasignación del código, este componente aplica un mecanismo de interpolación aprendible a las incrustaciones (embeddings) de entrada de la capa $l$. Utiliza un coeficiente de mezcla $\alpha$ (restringido a $[0, 0.5]$) para equilibrar la incrustación original con una versión transformada linealmente, evitando la destrucción de las características lingüísticas mientras permite refinamientos dirigidos para la reconstrucción entre capas.
2. Cuantizador Vectorial: Actúa como un cuello de botella discreto, mapeando las salidas continuas del codificador a un único vector de libro de códigos (codebook) de un conjunto aprendido $E = \{e_j\}_{j=1}^K$. Para asegurar un entrenamiento estable y una utilización efectiva del libro de códigos, los autores emplean:
   • Inicialización: K-means (Euclidiano o Esférico) o inicialización aleatoria.
   • Muestreo Estocástico: En lugar de una selección determinista, el modelo muestrea de los $k$ vectores más cercanos del libro de códigos utilizando una distribución de probabilidad controlada por temperatura. Esto fomenta la utilización uniforme del libro de códigos y la diversidad de conceptos.
   • Actualizaciones EMA: Los vectores del libro de códigos se actualizan mediante Medias Móviles Exponenciales (EMA) de los conteos de asignación y los vectores acumulados en lugar de la retropropagación directa, asegurando la estabilidad.
3. Decodificador Transformer: Un decodificador transformer bidireccional reconstruye las activaciones objetivo de la capa $h$ utilizando los vectores cuantizados del libro de códigos. Utiliza mecanismos de atención cruzada para aprovechar la información de las salidas no cuantizadas del codificador, delegando los detalles de reconstrucción de bajo nivel para que el libro de códigos pueda centrarse en conceptos distintos de alto nivel.

El modelo se entrena mediante una función de pérdida combinada: una pérdida de reconstrucción (Error Cuadrático Medio entre la salida del decodificador y la representación de la capa superior objetivo) y una pérdida de compromiso (que incentiva a las salidas del codificador a comprometerse con los vectores del libro de códigos).

Contribuciones Clave

• Mapeo Discreto entre Capas: A diferencia de los VQ-VAEE estándar que reconstruyen la misma capa, el CLVQ-VAE actúa como un transcodificador, colapsando las características redundantes de la corriente residual de las capas inferiores en vectores de conceptos compactos e interpretables en las capas superiores.
• Mecanismo de Muestreo Estocástico: La introducción del muestreo basado en temperatura de top-$k$ combinado con las actualizaciones de EMA del libro de códigos proporciona una exploración controlada del espacio latente discreto, manteniendo la diversidad del libro de códigos y evitando el colapso.
• Evaluación Exhaustiva: El marco se evalúa en modelos basados en codificadores (RoBERTa, BERT) y modelos solo de decodificador (LLaMA-2-7B, Qwen2.5-3B-Instruct) en tres conjuntos de datos (ERASER-Movie, Jigsaw, AGNews).

Resultados

El artículo evalúa CLVQ-VAE frente a clustering, VQ-VAE de una sola capa y baselines de SAE a través de tres ejes:

1. Fidelidad (Ablación de Conceptos): Eliminar los conceptos identificados de las representaciones de las oraciones degrada significativamente la precisión del modelo. CLVQ-VAE logró la precisión perturbada más baja en 5 de 12 configuraciones, con la eliminación de conceptos causando caídas de precisión de hasta el 93.2% (en RoBERTa/ERASER-Movie). Superó consistentemente a los SAE, que mostraron caídas inconsistentes o insignificantes, particularmente en los modelos de decodificador.
2. Coherencia Semántica (LLM-as-a-Judge): Un ensamble de LLMs clasificó los conceptos de CLVQ-VAE en primer lugar en el 66.7% de las comparaciones, con una calificación media de 1.890, superando al VQ-VAE de una sola capa, al SAE y al clustering.
3. Interpretabilidad Humana: En un estudio humano, los anotadores pudieron recuperar las predicciones del modelo de las visualizaciones de CLVQ-VAE con una precisión del 78.2%, en comparación con el 54.1% del clustering. Los anotadores también reportaron mayor confianza (8.44 frente a 5.98) y lograron un acuerdo entre anotadores "casi perfecto" ($\kappa = 0.864$) con las visualizaciones de CLVQ-VAE.

Además, los autores encontraron un compromiso en las estrategias de inicialización: la inicialización K-means tendió a producir una mayor fidelidad funcional (mejores resultados de perturbación), mientras que Spherical K-means produjo conceptos con mayor coherencia semántica y alineación humana.

Significado

El artículo afirma que CLVQ-VAE proporciona un marco robusto para traducir mecanismos de modelos opacos en conceptos fieles e interpretables. Al integrar el aprendizaje de representaciones discretas con el análisis entre capas, el enfoque aborda las limitaciones de los SAE continuos (división de conceptos) y los análisis de una sola capa (falta de estructura entre capas). Los resultados validan que los conceptos descubiertos no solo son funcionalmente críticos para las predicciones del modelo, sino que también se alinean con la comprensión humana, ofreciendo un camino hacia modelos de lenguaje más transparentes y responsables. Los autores señalan que, si bien el método es efectivo, requiere una memoria sustancial para la extracción de activaciones y un entrenamiento separado para diferentes combinaciones de capas.