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

Resumo Técnico: Interpretabilidade Mecanística e Direcionamento de Características Causais de Estados Quânticos Neurais via Autoencoders Esparsos

Declaração do Problema
Os Estados Quânticos Neurais (NQS) emergiram como poderosos ansatzes variacionais para representar funções de onda de muitos corpos quânticos, alcançando sucesso tanto na representação do estado fundamental quanto na dinâmica em tempo real em sistemas de spins e fermiônicos. No entanto, apesar de sua eficácia empírica, os NQS permanecem amplamente "caixas pretas". Existe uma lacuna fundamental na compreensão de seus mecanismos internos: treinados puramente sob objetivos variacionais (como a minimização da energia) sem otimização explícita para observáveis físicos específicos, não está claro como os NQS capturam com precisão quantidades como parâmetros de ordem ou funções de correlação. Essa opacidade dificulta a confiança e a robustez do modelo, particularmente em regimes onde benchmarks numéricos são difíceis de obter. As abordagens existentes focaram majoritariamente em restrições arquiteturais externas ou análise de parâmetros, deixando uma carência de análise post hoc sobre como as redes organizam internamente a informação física.

Metodologia
Os autores propõem um framework sistemático utilizando Autoencoders Esparsos (SAEs), uma ferramenta da interpretabilidade mecanística, para decompor as ativações internas dos NQS em características esparsas e interpretáveis.

1. Arquitetura NQS: O estudo utiliza NQS baseados em transformers autorregressivos. A amplitude da função de onda $\psi(\sigma)$ é fatorada em uma probabilidade de Born $p(\sigma)$ e uma fase $\phi(\sigma)$. O transformer atua como um modelo de sequência causal, prevendo configurações de spin sítio a sítio.
2. Aplicação de SAE: O SAE é aplicado ao fluxo residual (residual stream) da camada final do transformer. Ele aprende um dicionário sobrecompleto de direções de características $\{f_k\}$ e coeficientes esparsos $z_k$ de tal forma que o vetor de ativação $h$ é aproximado como uma combinação linear esparsa: $h \approx \sum z_k f_k$.
3. Treinamento Não Supervisionado: Crucialmente, o SAE é treinado inteiramente de forma não supervisionada nas ativações do fluxo residual. Ele não recebe rótulos físicos (ex: valores de magnetização) e é otimizado apenas para reconstruir as ativações com uma penalidade de esparsidade (norma $L_1$).
4. Correlação Característica-Observável: Após o treinamento, os autores realizam o "mean-pooling" dos coeficientes esparsos sobre a cadeia de spins para obter intensidades de características escalares. Estas são correlacionadas com observáveis físicos (ex: magnetização, magnetização estagiada, correladores de meia-cadeia) computados a partir das mesmas configurações amostradas.
5. Intervenção Causal (Direcionamento): Para estabelecer causalidade, os autores realizam intervenções direcionadas. Eles reescalam a intensidade da ativação de uma característica específica ($z_{k^*} \to \alpha z_{k^*}$) enquanto mantêm as outras fixas, reconstroem o fluxo residual modificado e observam a mudança resultante no observável físico. Isso é feito enquanto monitoram a energia variacional para garantir que a intervenção permaneça localizada.

Principais Contribuições e Resultados

• Descoberta de Características Interpretáveis: No Modelo de Ising de Campo Transverso (TFIM) 1D, o SAE descobriu autonomamente características esparsas que se correlacionam quase perfeitamente ($|r| > 0,97$) com observáveis físicos, incluindo magnetização ($M_z$), magnetização absoluta ($|M_z|$) e correladores de meia-cadeia ($C_{zz}$). Esta descoberta ocorreu através das fases paramagnética, ferromagnética e o ponto crítico, apesar de o SAE não possuir conhecimento prévio dessas quantidades.
• Controle Causal: Os autores demonstraram que essas características desempenham um papel causal nas previsões do NQS. Reescalar a intensidade de uma única característica ($\alpha$) direcionou de forma suave e monotônica o observável correspondente (ex: magnetização) enquanto deixou a energia variacional quase inalterada (desvio relativo $< 0,02\%$). Isso confirma que as características não são meros artefatos correlacionados, mas sim coordenadas internas que controlam as previsões físicas.
• Dinâmica em Tempo Real: O framework foi estendido para a dinâmica em tempo real sob o TFIM. Uma única direção de característica estática foi capaz de rastrear a magnetização dependente do tempo $M_z(t)$ ao longo de toda a trajetória de evolução. Além disso, escalar a intensidade de ativação desta característica ao longo do tempo alterou sistematicamente a dinâmica da magnetização, provando que os NQS codificam informações dinâmicas em um sistema de coordenadas internas estáticas.
• Generalização para Sistemas 2D: A abordagem foi validada no modelo de Heisenberg Antiferromagnético 2D (um sistema não integrável de 2D). O SAE identificou uma característica fortemente correlacionada com a magnetização estagiada (um parâmetro de ordem espacialmente estruturado que distingue sublatices). O direcionamento desta característica controlou com sucesso a magnetização estagiada, demonstrando a aplicabilidade do método além de sistemas 1D e exatamente solúveis.

Significância e Alegações
O artigo alega que os NQS não são meros aproximadores funcionais de amplitudes de função de onda, mas codificam representações internas ricas e interpretáveis de informação física.

• Ferramenta Diagnóstica: A abordagem fornece um método para avaliar se um NQS treinado aprendeu uma representação fisicamente sensível. A presença de características que controlam causalmente os parâmetros de ordem serve como evidência de um aprendizado bem-sucedido, enquanto sua ausência poderia diagnosticar falhas do modelo.
• Ferramenta de Intervenção: O trabalho estabelece uma rota para a manipulação controlada de características ao nível de feature. Ao ajustar um único coeficiente escalar no fluxo residual, é possível ajustar observáveis físicos específicos sem necessidade de retreinar ou realizar o ajuste fino (fine-tuning) de toda a rede.
• Troca de Via Dupla: Os autores postulam uma relação sinérgica entre a física de muitos corpos quânticos e a interpretabilidade mecanística. Os NQS fornecem um cenário com observáveis físicos de verdade fundamental para testar alegações de interpretabilidade, enquanto ferramentas de interpretabilidade oferecem métodos para abrir a "caixa preta" das funções de onda variacionais, potencialmente levando a simulações quânticas mais confiáveis e transparentes.

Os autores mantêm-se modestos quanto ao escopo, observando que este trabalho foca exclusivamente em arquiteturas transformer e fluxos residuais de camadas finais, deixando em aberto questões relativas a outras arquiteturas (ex: RBMs, RNNs) e mecanismos de auto-atenção resolvidos por camada.