Expander Sparse Autoencoders: Parameter-Efficient Dictionaries for Mechanistic Interpretability

Resumen Técnico: Autoencoders de Dispersión de Expansores (Expander Sparse Autoencoders)

Planteamiento del Problema

Los Autoencoders de Dispersión (SAEs) son una herramienta primordial en la interpretabilidad mecanística, diseñada para descomponer las activaciones internas de las redes neuronales en combinaciones lineales dispersas de características interpretables. Este proceso depende del aprendizaje de un diccionario sobrecompleto $W \in \mathbb{R}^{m \times n}$ (donde $m < n$) para representar las activaciones $h \approx Wx$ con un código disperso $x$.

El desafío central abordado por este trabajo es la eficiencia de parámetros del decodificador. Los SAEs estándar utilizan decodificadores densos que requieren $O(mn)$ parámetros aprendidos. A medida que el recuento de características $n$ escala a millones en los modelos de lenguaje de gran tamaño (LLMs) modernos, el almacenamiento y la actualización de estos pesos densos se vuelve computacionalmente y de almacenamiento prohibitivos. Aunque la teoría de la detección comprimida sugiere que la recuperación dispersa es posible con matrices de medición estructuradas, los SAEs estándar no aprovechan estas estructuras, lo que conduce a una sobrecarga de parámetros innecesaria.

Metodología: Expander SAEs

Los autores proponen los Expander SAEs, una arquitectura de parámetros eficientes que reemplaza el decodificador denso con uno disperso y estructurado basado en la detección comprimida combinatoria.

Arquitectura

• SAEs de Pesos Vinculados TopK: El modelo utiliza una arquitectura de pesos vinculados donde el codificador es la transpuesta del decodificador ($W_{enc} = W_{dec}^\top$). La dispersión se impone mediante una no linealidad TopK, reteniendo solo las $k$ pre-activaciones más grandes.
• Máscara de Expansor: El diccionario del decodificador $W_{dec}$ es enmascarado por la matriz de adyacencia $M$ de un grafo expansor regular a la izquierda de grado $d$.
  • $M \in \{0, 1\}^{m \times n}$ es una máscara binaria donde cada columna tiene exactamente $d$ entradas no nulas ($d \ll m$).
  • Los parámetros aprendibles se reducen de $mn$a$dn$ (específicamente, las entradas no nulas de $V$ en $W_{dec} = (V \odot M) \text{diag}(\nu)^{-1}$).
  • La máscara $M$ se mantiene fija en la inicialización y se muestrea de una distribución que produce un expansor $(k, \epsilon, d)$ con alta probabilidad.
• Inferencia: El problema de codificación dispersa se resuelve utilizando la Búsqueda de Coincidencia Ortogonal (OMP). La estructura del expansor permite que el paso de correlación $W^\top r$ se compute como una operación de "recolección y reducción" (gather-and-reduce) de $O(dn)$ en lugar de un producto matriz-vector denso de $O(mn)$.

Garantías Teóricas

El artículo establece condiciones teóricas para la identificabilidad de códigos $k$-dispersos sin ruido:

1. Identificabilidad: Si la máscara $M$ es un expansor $(2k, \epsilon, d)$ y las columnas del decodificador aprendido satisfacen una condición de planitud de columna (limitada por el factor $\beta$), entonces cada código $k$-disperso tiene una solución única siempre que $2\beta^2\epsilon < 1$.
2. Recuperación OMP: Bajo una condición de coherencia acumulada más fuerte ($\beta^2\epsilon(2k+1) < 1$), se garantiza que el OMP estándar recupere el soporte exacto del código disperso en $k$ pasos.

Principales Contribuciones

1. Arquitectura: Introducción de los Expander SAEs, la primera arquitectura de SAE que evalúa la dispersión de la columna del decodificador como una restricción de almacenamiento e interpretabilidad, separando la calidad del diccionario de la calidad de la inferencia amortizada.
2. Frontera de Almacenamiento-Fidelidad: Demostración empírica de que variar el parámetro de dispersión $d$ traza una curva de compromiso suave entre el número de valores de decodificador aprendidos y la fidelidad de la reconstrucción a través de múltiples familias de modelos (Pythia, Qwen2.5, Llama-3.2).
3. OMP Paralelo: Una implementación de OMP optimizada para GPU y paralela que explota la estructura $d$-regular, logrando mejoras significativas en el rendimiento (throughput) sobre el OMP vanilla mientras mantiene una alta calidad de reconstrucción.
4. Teoría de la Identificabilidad: Prueba de que la expansión y la planitud de columna son condiciones suficientes para la identificabilidad única de códigos $k$-dispersos sin ruido en decodificadores aprendidos no binarios, extendiendo la teoría clásica de grafos expansores al entorno de los SAE.

Resultados Experimentales

Los autores evaluaron los Expander SAEs sobre las activaciones del flujo residual de Pythia-70M/160M, Qwen2.5-3B y Llama-3.2-1B.

• Eficiencia de Almacenamiento vs. Fidelidad:
  • En Qwen2.5-3B, un Expander SAE con $d=7$ utiliza 293$\times$ menos valores de decodificador aprendidos que un decodificador denso completo (114,681 frente a 33,554,432) mientras recupera el 84% de la pérdida de entropía cruzada (CE) del modelo denso.
  • En Pythia-70M, $d=7$ logra una recuperación de CE del 86% con 73$\times$ menos parámetros.
  • El aumento de $d$ mejora suavemente la fidelidad, acercándose al rendimiento de la línea base densa a medida que $d \to m$.
• Experimentos de Control:
  • Diversidad de Soporte: Un control de "dispersión agrupada" (donde los soportes no son diversos) mostró un aumento masivo de características muertas (características que nunca se activan) con un $d$ alto, lo que indica que la estructura de soporte diversa del expansor es crítica para la utilización de las características, no solo la dispersión de la columna.
  • Igualación de Parámetros: Cuando se comparan los Expander SAEs con los Dense SAEs con recuentos de parámetros igualados (reduciendo $n$), los Dense SAEs a menudo rinden peor o fallan por completo porque reducir $n$ viola la propiedad del espacio nulo ($n > 2k$) requerida para la recuperación dispersa.
  • Amortización del Codificador: Una parte significativa de la brecha de rendimiento entre los codificadores entrenados y los decodificadores OMP en los LLMs modernos se atribuye a la "brecha de amortización" del codificador. El OMP iterativo cierra gran parte de esta brecha, sugiriendo que la estructura del decodificador es robusta incluso si el codificador es subóptimo.
• Interpretabilidad: Evaluaciones ciegas realizadas por jueces LLM encontraron que las características de los Expander SAEs ($d=200$) eran estadísticamente indistinguibles en coherencia de los Dense SAEs; aunque la coherencia cayó ligeramente en la compresión extrema ($d=7$), el 84% de las características aún recibían etiquetas de conceptos concretos y exhibían una alta novedad respecto a las líneas base densas.

Significado y Reivindicaciones

El artículo afirma que la estructura de soporte del decodificador es un eje de diseño subexplorado y crítico para la interpretabilidad mecanística. Los autores no afirman que los Expander SAEs sean un reemplazo universal para los Dense SAEs, sino que argumentan que ofrecen una solución Pareto-óptima para escenarios donde el almacenamiento es un cuello de botella.

Conclusiones clave enfatizadas por los autores:

• Eficiencia de Parámetros: Es posible reducir el almacenamiento del decodificador en órdenes de magnitud (por ejemplo, 293$\times$) con una pérdida mínima en la fidelidad de reconstrucción o la interpretabilidad.
• Beneficios Estructurales: La máscara del expansor permite algoritmos de inferencia eficientes (OMP estructurado) y máscaras deterministas regenerables mediante semilla, reduizando significamente el tamaño del artefacto en disco.
• Fundamentación Teórica: La arquitectura está fundamentada en la teoría establecida de la detección comprimida, proporcionando garantías demostrables para la identificabilidad y la recuperación bajo condiciones específicas, lo que motiva el diseño en lugar de ser solo un heurístico empírico.

El trabajo sugiere que la investigación futura debería explorar evaluaciones a mayor escala, correcciones de bajo rango y soportes aprendidos, posicionando a los Expander SAEs como un paso fundamental hacia herramientas de interpretabilidad escalables y eficientes en parámetros.