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

技术摘要：通过稀疏自编码器实现神经量子态的机械解释性与因果特征操控

问题陈述
神经量子态（NQS）已成为表示量子多体波函数的强大变分拟设，在自旋和费米子系统的基态表示及实时动力学方面取得了成功。然而，尽管具有经验上的有效性，NQS 在很大程度上仍是“黑盒”。由于其仅针对变分目标（如能量最小化）进行训练，而未针对特定物理可观测量进行显式优化，因此存在一个根本性的空白：目前尚不清楚 NQS 如何准确捕捉诸如序参数或相关函数之类的物理量。这种不透明性阻碍了模型的信任度和鲁棒性，尤其是在数值基准难以获得的机制中。现有方法主要集中在外部架构约束或参数分析上，缺乏关于网络如何内部组织物理信息的后验（post hoc）分析。

方法论
作者提出了一个使用**稀疏自编码器（SAE）**的系统性框架，该工具源自机械解释性领域，用于将 NQS 的内部激活分解为稀疏且可解释的特征。

1. NQS 架构： 本研究采用了基于自回归 Transformer 的 NQS。波函数振幅 $\psi(\sigma)$ 被分解为 Born 概率 $p(\sigma)$ 和相位 $\phi(\sigma)$。Transformer 作为一个因果序列模型，逐位预测自旋构型。
2. SAE 应用： SAE 被应用于最终 Transformer 层的残差流（residual stream）。它学习了一个过完备的特征方向字典 $\{f_k\}$ 和稀疏系数 $z_k$，使得激活向量 $h$ 可以近似为稀疏线性组合：$h \approx \sum z_k f_k$。
3. 无监督训练： 至关重要的是，SAE 完全在残差流激活上进行无监督训练。它不接收任何物理标签（例如磁化强度值），而是仅通过优化以重建激活并带有稀疏惩罚（$L_1$ 范数）。
4. 特征-可观测量相关性： 训练完成后，作者通过对自旋链进行“均值池化”来获得标量特征强度。这些强度与从相同采样构型中计算出的物理可观测量（如磁化强度、交错磁化强度、半链相关函数）进行相关性分析。
5. 因果干预（操控）： 为了建立因果关系，作者进行了针对性的干预。他们通过缩放特定特征的激活强度（$z_{k^*} \to \alpha z_{k^*}$）同时保持其他特征固定，从而重建修改后的残差流，并观察物理可观量的变化。在此过程中，他们同时监测变分能量，以确保干预是局部化的。

核心贡献与结果

• 发现可解释特征： 在一维横场伊辛模型（TFIM）中，SAE 自主发现了与物理可观测量（包括磁化强度 $M_z$、绝对磁化强度 $|M_z|$ 和双自旋相关函数 $C_{zz}$）近乎完美相关（$|r| > 0.97$）的稀疏特征。尽管 SAE 预先并不知道这些物理量，但这种发现跨越了顺磁相、铁磁相以及临界点。
• 因果控制： 作者证明了这些特征在 NQS 预测中起到了因果作用。缩放单个特征强度（$\alpha$）可以平滑且单调地操控相应的可观测量（如磁化强度），同时使变分能量的变化极小（相对偏差 $< 0.02\%$）。这证实了这些特征不仅是相关的伪影，更是控制物理预测的内部坐标。
• 实时动力学： 该框架被扩展到 TFIM 下的实时动力学。研究发现，单个静态特征方向可以在整个演化轨迹中追踪随时间变化的磁化强度 $M_z(t)$。此外，跨时间缩放该特征的激活强度可以系统地改变磁化强度动力学，证明了 NQS 将动力学信息编码在静态内部坐标系中。
• 推广至二维系统： 该方法在二维海森堡反铁磁模型（一个非可积的二维系统）上得到了验证。SAE 识别出一个与交错磁化强度（一种区分子格的空间结构化序参数）强相关的特征。操控该特征成功控制了交错磁化强度，证明了该方法在超越一维及可解系统方面的适用性。

意义与主张
本文声称，NQS 不仅仅是波函数振幅的函数逼近器，还编码了丰富的、具有可解释性的物理信息内部表示。

• 诊断工具： 该方法提供了一种评估训练好的 NQS 是否学习到了具有物理意义的表示的方法。存在能够因果控制序参数的特征，可以作为成功学习的证据；而特征的缺失则可以诊断模型失效。
• 干预工具： 这项工作为 NQS 的特征级受控操作建立了路径。通过调整残差流中的单个标量系数，可以在无需重新训练或微调整个网络的情况下，调节特定的物理可观测量。
• 双向交换： 作者认为量子多体物理学与机械解释性之间存在协同关系。NQS 提供了一个具有地面真值（ground-truth）物理可观量的环境来测试解释性主张，而解释性工具则提供了打开变分波函数“黑盒”的方法，有望实现更可靠、更透明的量子模拟。

作者对研究范围保持谦逊，指出目前的工作仅限于 Transformer 架构和最终层残差流，对于其他架构（如 RBM、RNN）以及层级解析的自注意力机制仍留有待解决的问题。