Mechanistic Interpretability and Causal Feature Steering of Neural Quantum States via Sparse Autoencoders

Resumen Técnico: Interpretabilidad Mecanicista y Control de Características Causales de Estados Cuánticos Neuronales mediante Autoencoders Dispersos

Planteamiento del Problema
Los Estados Cuánticos Neuronales (NQS) han surgido como potentes ansätze variacionales para representar funciones de onda de muchos cuerpos cuánticos, logrando éxito tanto en la representación del estado fundamental como en la dinámica en tiempo real en sistemas de espín y fermiónicos. Sin embargo, a pesar de su eficacia empírica, los NQS siguen siendo en gran medida "cajas negras". Existe una brecha fundamental en la comprensión de sus mecanismos internos: al estar entrenados únicamente mediante objetivos variacionales (como la minimización de la energía) sin una optimización explícita para observables físicos específicos, no está claro cómo los NQS capturan con precisión cantidades como los parámetros de orden o las funciones de correlación. Esta opacidad dificulta la confianza y la robustez del modelo, particularmente en regímenes donde los puntos de referencia numéricos son difíciles de obtener. Los enfoques existentes se han centrado mayoritariamente en restricciones arquitectónicas externas o análisis de parámetros, dejando una falta de análisis post hoc sobre cómo las redes organizan internamente la información física.

Metodología
Los autores proponen un marco sistemático utilizando Autoencoders Dispersos (SAEs), una herramienta de la interpretabilidad mecanicista, para descomponer las activaciones internas de los NQS en características dispersas e interpretables.

1. Arquitectura NQS: El estudio utiliza NQS basados en transformadores autorregresivos. La amplitud de la función de onda $\psi(\sigma)$ se factoriza en una probabilidad de Born $p(\sigma)$ y una fase $\phi(\sigma)$. El transformador actúa como un modelo de secuencia causal, prediciendo configuraciones de espín sitio por sitio.
2. Aplicación de SAE: El SAE se aplica al flujo residual (residual stream) de la última capa del transformador. Este aprende un diccionario sobrecompleto de direcciones de características $\{f_k\}$ y coeficientes dispersos $z_k$ tales que el vector de activación $h$ se aproxima como una combinación lineal dispersa: $h \approx \sum z_k f_k$.
3. Entrenamiento No Supervisado: Crucialmente, el SAE se entrena completamente de forma no supervisada sobre las activaciones del flujo residual. No recibe etiquetas físicas (por ejemplo, valores de magnetización) y se optimiza únicamente para reconstruir las activaciones con una penalización de dispersión (norma $L_1$).
4. Correlación Característica-Observable: Tras el entrenamiento, los autores realizan un "agrupamiento medio" (mean-pooling) de los coeficientes dispersos sobre la cadena de espín para obtener magnitudes escalares de las características. Estos se correlacionan con observables físicos (por ejemplo, magnetización, magnetización escalonada, correladores de media cadena) computados desde las mismas configuraciones muestreadas.
5. Intervención Causal (Control/Steering): Para establecer la causalidad, los autores realizan intervenciones dirigidas. Reescalan la fuerza de activación de una característica específica ($z_{k^*} \to \alpha z_{k^*}$) mientras mantienen las demás fijas, reconstruyen el flujo residual modificado y observan el cambio resultante en el observable físico. Esto se realiza mientras se monitorea la energía variacional para asegurar que la intervención permanezca localizada.

Contribuciones Clave y Resultados

• Descubrimiento de Características Interpretables: En el Modelo de Ising de Campo Transverso (TFIM) en 1D, el SAE descubrió autónomamente características dispersas que se correlacionan casi perfectamente ($|r| > 0.97$) con observables físicos, incluyendo la magnetización ($M_z$), la magnetización absoluta ($|M_z|$) y los correladores de media cadena ($C_{zz}$). Este descubrimiento ocurrió a través de la fase paramagnética, la fase ferromagnética y el punto crítico, a pesar de que el SAE no tenía conocimiento previo de estas cantidades.
• Control Causal: Los autores demostraron que estas características desempeñan un papel causal en las predicciones de los NQS. El reescalado de la fuerza de una característica ($\alpha$) dirigió de forma suave y monótona el observable correspondiente (por ejemplo, la magnetización) mientras la energía variacional permaneció casi inalterada (desviación relativa $< 0.02\%$). Esto confirma que las características no son meros artefactos de correlación, sino coordenadas internas que controlan las predicciones físicas.
• Dinámica en Tiempo Real: El marco se extendió a la dinámica en tiempo real bajo el TFIM. Se encontró que una única dirección de característica estática rastrea la magnetización dependiente del tiempo $M_z(t)$ durante toda la trayectoria de evolución. Además, el escalado de la fuerza de activación de esta característica a través del tiempo alteró sistemáticamente la dinámica de la magnetización, probando que los NQS codifican información dinámica en un sistema de coordenadas internas estáticas.
• Generalización a Sistemas 2D: El enfoque fue validado en el modelo de Heisenberg Antiferromagnético 2D (un sistema no integrable y 2D). El SAE identificó una característica fuertemente correlacionada con la magnetización escalonada (un parámetro de orden espacialmente estructurado que distingue subredes). El control de esta característica logró controlar con éxito la magnetización escalonada, demostrando la aplicabilidad del método más allá de los sistemas de 1D y exactamente resolubles.

Significancia y Reivindicaciones
El artículo sostiene que los NQS no son meros aproximadores funcionales de las amplitudes de la función de onda, sino que codifican representaciones de información física ricas e interpretables.

• Herramienta de Diagnóstico: El enfoque proporciona un método para evaluar si un NQS entrenado ha aprendido una representación físicamente sensata. La presencia de características que controlan causalmente los parámetros de orden sirve como evidencia de un aprendizaje exitoso, mientras que su ausencia podría diagnosticar un fallo del modelo.
• Herramienta de Intervención: El trabajo establece una ruta para la manipulación controlada a nivel de característica de los NQS. Al ajustar un único coeficiente escalar en el flujo residual, es posible sintonizar observables físicos específicos sin necesidad de reentrenar o realizar un ajuste fino (fine-tuning) de toda la red.
• Intercambio Bidireccional: Los autores postulan una relación sinérgica entre la física de muchos cuerpos cuánticos y la interpretabilidad mecanicista. Los NQS proporcionan un entorno con observables físicos de verdad fundamental para probar las afirmaciones de interpretabilidad, mientras que las herramientas de interpretabilidad ofrecen métodos para abrir la "caja negra" de las funciones de onda variacionales, conduciendo potencialmente a simulaciones cuánticas más fiables y transparentes.

Los autores mantienen la modestia respecto al alcance, señalando que este trabajo se centra exclusivamente en arquitecturas de transformadores y flujos residuales de la última capa, dejando abiertas preguntas sobre otras arquitecturas (por ejemplo, RBMs, RNNs) y mecanismos de autoatención resueltos por capas.