Resumen Técnico: Autoencoders Dispersos Equivariantes

Planteamiento del Problema

Los modelos de aprendizaje automático, particularmente en dominios científicos, operan a menudo con datos que poseen simetrías inherentes (por ejemplo, rotaciones en el plegamiento de proteínas, imágenes de galaxias o microscopía de células). Aunque los Autoencoders Dispersos (SAEs, por sus siglas en inglés) han surgido como una herramienta primordial para la interpretabilidad mecánica al descomponer activaciones neuronales densas en características dispersas e interpretables, enfrentan una limitación crítica cuando se aplican a datos simétricos: la unidentificabilidad.

Los SAEs estándar asumen una Hipótesis de Representación Lineal (LRH) donde las activaciones son combinaciones lineales dispersas de características. Sin embargo, cuando los datos de entrada poseen simetrías (acciones de grupo $G$), un SAE convencional debe aprender direcciones de características distintas para el mismo concepto semántico bajo diferentes transformaciones (por ejemplo, una característica de "gato" para cada rotación posible). Esto conduce a:

1. Entrelazamiento de Características: El modelo puede aprender características no ortogonales y no interpretables para reconstruir entradas dispares.
2. Unidentificabilidad: Diferentes explicaciones pueden ajustar los datos igualmente bien, pero solo una se alinea con la verdadera estructura subyacente.
3. Desalineación de Métricas: La calidad de la reconstrucción (MSE) a menudo no correlaciona con la utilidad de la característica, ya que los SAEs pueden aprender características específicas de la reconstrucción que no corresponden a conceptos reales.

El artículo postula que los SAEs actuales no tienen en cuenta estas simetrías, lo que lleva a una interpretabilidad subóptima en dominios científicos donde tales simetrías son prevalentes.

Metodología: SAEs Equivariantes

Los autores proponen Autoencoders Dispersos Equivariantes, que extienden la Hipótesis de Representación Lineal para considerar explícitamente las simetrías de grupo.

1. Hipótesis de Representación Lineal Extendida

Los autores reformulan la LRH para establecer que, para un grupo $G$ que actúa sobre las entradas $X$, las activaciones de la red neuronal se transforman de manera predecible. Específicamente, si $x$ es transformado por $g \in G$, las activaciones se transforman mediante un mapa lineal $\rho(g)$:
$$\psi(gx) = \rho(g) \sum_{i=1}^C z(x)_i D_i = \sum_{i=1}^C z(x)_i (\rho(g) D_i)$$
Aquí, las magnitudes de las características $z(x)$ permanecen invariantes (coeficientes dispersos), mientras que los vectores de características $D_i$ se transforman de acuerdo con la acción del grupo.

2. Diseño de la Arquitectura

El SAE Equivariante desacopla el aprendizaje de características invariantes del aprendizaje de las reglas de transformación:

• SAE Invariante: Un par codificador-decodificador entrenado para mapear todas las activaciones dentro de una órbita (por ejemplo, todas las rotaciones de una imagen) a una única representación canónica. Esto aprende las características invariantes.
• Módulo de Transformación Lineal: Un conjunto separado de matrices aprendibles ($M, R$) que modelan cómo la reconstrucción canónica se transforma de vuelta al espacio de entrada original. Para un grupo generado por $g, h$, el modelo aprende matrices tales que:
  $$M_p R_q D(E(\psi(g^p h^q x))) = \psi(g^p h^q x)$$
  Esto permite que el SAE capture la estructura de simetría sin forzar al diccionario a ser equivariante, lo que requeriría un número impráctico de parámetros.

3. Estrategia de Entrenamiento

El entrenamiento ocurre en dos pasos independientes (que pueden paralelizarse):

1. Entrenamiento del SAE Invariante: El codificador y el decodificador se optimizan para minimizar la pérdida de reconstrucción en entradas canónicas mientras se impone la invarianza. Los autores experimentan con tres objetivos de pérdida:
   • $L_{Canónica}$: Reconstruir una muestra canónica de cualquier miembro de la órbita.
   • $L_{Latente}$: Imponer la invarianza directamente en los códigos latentes ($E(\psi(x)) \approx E(\psi(gx))$).
   • $L_{Salida}$: Imponer la invarianza en las salidas del SAE.
2. Aprendizaje de la Transformación: Las matrices $M$ y $R$ se optimizan para minimizar el error entre las activaciones transformadas predichas y las activaciones transformadas reales del modelo base.

Contribuciones Clave

1. Extensión Teórica: El artículo extiende la Hipótesis de Representación Lineal a datos simétricos, definiendo formalmente cómo se comportan los coeficientes de las características y los vectores bajo acciones de grupo.
2. Innovación Algorítmica: El diseño de SAEs Equivariantes que separan el aprendizaje de características invariantes del modelado de la transformación, evitando la explosión de parámetros de los diccionarios equivariantes de grupo.
3. Validación Empírica: Demostración a través de conjuntos de datos sintéticos y del mundo real (formas, imágenes de galaxias, células sanguíneas) de que los SAEs Equivariantes superan a los baselines de vanilla, Archetypal y Group Cross-coder en recuperación de características y tareas de sondeo (probing) descendentes.
4. Perspectiva de Métricas: El descubrimiento de que la calidad de la reconstrucción está inversamente correlacionada con la utilidad de la característica en entornos simétricos. Los SAEs Equivariantes a menudo logran una mejor calidad de reconstrucción (menor MSE) que los vanilla SAEs, pero generan características significativamente más interpretables y útiles.

Resultados Experimentales

Los autores evaluaron su método en:

• Modelos de Juguete (Toy Models): Datos sintéticos con características equivariantes de verdad conocida.
• Dataset de Formas: Imágenes sintéticas con simetrías $C_4$ (rotaciones) y $D_4$ (rotaciones + reflejos), utilizando modelos base MLP, CNN y ViT.
• GalaxyMNIST: Imágenes de galaxias reales con simetrías $C_4/D_4$.
• MLL23: Imágenes de células sanguíneas periféricas utilizando un modelo fundacional DinoBloom.

Hallazgos Clave:

• Recuperación de Características: En los modelos de juguete, los SAEs Equivariantes lograron una recuperación de características superior (Similitud de Coseno Media) y detección (F1 Score) en comparación con los baselines, especialmente en regímenes de sobrecompletitud donde la interferencia es alta.
• Sondeo Descendente (Downstream Probing): En todos los datasets, los sondas entrenados en los latentes de los SAEs Equivariantes alcanzaron una mayor precisión en la clasificación de propiedades de entrada (por ejemplo, presencia de forma, orientación, tipo de célula) que los sondos entrenados en los latentes de los vanilla SAEs. En algunos casos (por ejemplo, GalaxyMNIST), igualaron el rendimiento de los sondos entrenados directamente en las activaciones del modelo base.
• Reconstrucción vs. Utilidad: Surgió una tendencia constante donde los SAEs Equivariantes tenían un error de reconstrucción más alto (peor MSE) que los vanilla SAEs. Sin embargo, esta "peor" reconstrucción correspondió a características que estaban más alineadas con la estructura semántica real y eran más útiles para tareas descendentes.
• Modelado de Transformación: Las matrices lineales aprendidas ($M, R$) explicaron con éxito más del 97% de la varianza en las transformaciones de activación para la mayoría de los modelos, validando la hipótesis de que las simetrías de entrada inducen transformaciones lineales en el espacio de activación.

Significado y Reivindicaciones

El artículo afirma que alinear los priors de los métodos de interpretabilidad con la estructura conocida de los datos (específicamente las simetrías) es crucial para obtener explicaciones fieles. Los autores destacan dos implicaciones principales:

1. La Reconstrucción no es Interpretabilidad: Depender únicamente de métricas de reconstrucción (como el MSE) para evaluar los SAEs es peligroso, particularmente para datos simétricos. Un modelo que reconstruye mal las entradas puede, aun así, descubrir las características más útiles e interpretables.
2. Interpretabilidad de Dominio Específico: Para aplicaciones científicas (por ejemplo, biología, astronomía) donde las simetrías de los datos son fundamentales, los SAEs estándar son insuficientes. Los SAEs Equivariantes proporcionan un camino para descubrir características que son robustas a estas simetrías, lo que potencialmente conduce a nuevos descubrimientos científicos y a sistemas de IA más fiables.

El trabajo sugiere que la investigación futura en interpretabilidad debe ir más allá del codificado disperso genérico e incorporar priors estructurales de dominio específico para evitar los problemas de la unidentificabilidad.