An interpretable closed form for entanglement entropy from bitstrings, guided by a graph neural network

Technical Summary: An Interpretable Closed Form for Entanglement Entropy from Bitstrings, Guided by a Graph Neural Network

Problem Statement
In Rydberg-atom array quantum simulations, the empirical bitstring distribution (measurement outcomes in the occupation basis) is the most accessible observable. However, the bipartite von Neumann entanglement entropy ($S_{vN}$), a critical diagnostic for quantum many-body matter, is far more costly to obtain. Full state tomography scales exponentially with system size $N$, and even tomography-free protocols require significant overhead in copies, gates, or post-processing. While the bitstring distribution does not uniquely determine $S_{vN}$ (as different states can share a diagonal distribution but differ in entanglement), it tightly constrains it. Existing estimators face a trade-off: analytic bounds based on classical mutual information ($I$) are interpretable but coarse (underestimating $S_{vN}$ by $\sim 0.10$ nats), while accurate Graph Neural Network (GNN) surrogates are black boxes that require retraining for every new regime (phase, geometry, or Hamiltonian). The goal is to find an estimator that combines the accuracy of learned models with the interpretability and portability of analytic forms.

Methodology
The authors employ a "GNN-guided" approach to derive a closed-form estimator:

1. Mechanistic Probing of the GNN: Instead of fitting an analytic form to hand-picked observables, the authors probe a pre-trained GNN (from [5]) to identify which inputs drive its $S_{vN}$ predictions. Through causal interventions (attention routing, activation patching, and edge removal), they establish that the GNN's prediction is routed specifically through the bipartition-boundary edges ($E_b = A \times B$). The network re-routes attention based on whether an edge crosses the cut, and removing boundary edges significantly degrades performance.
2. Feature Restriction and Search: Guided by the GNN's focus on boundary correlators, the authors restrict an exhaustive search to a curated shortlist of 25 physics scalars derived from the bitstring distribution, specifically those computed over the boundary set $E_b$.
3. Exhaustive Symbolic Search: An exhaustive search over all $\binom{25}{6} = 177,100$ subsets of six features is performed using Ordinary Least Squares (OLS). The winner is a six-term linear closed form (Eq. 1).
4. Validation and Scaling: The form is validated against ground truth (Exact Diagonalization for $N \le 20$ and Density Matrix Renormalization Group (DMRG) for $N$ up to 100). The authors analyze the scaling behavior of the coefficients and test transferability across different Hamiltonian families and geometric shifts.

Key Contributions

1. A Six-Term Closed Form (Eq. 1): The authors present a human-readable, hyperparameter-free linear estimator:
   $$\hat{S}_{vN} = b_1 \log(I+1) + b_2 \sum_{E_b} \sqrt{|c|} + b_3 \sum_{E_b} \tanh |c| + b_4 (\max_{E_b} |c|)^2 + b_5 \min_{E_b} c + b_6 \min(H_A, H_B) I$$
   where $c$ represents two-point correlators, $I$ is mutual information, and $H$ is single-site entropy.
2. GNN-Guided Featurization: The work demonstrates that causal probing of a neural network can effectively identify the relevant physical subspace (boundary correlators) for a specific prediction task, bypassing the need for manual feature engineering.
3. Scaling Laws and Calibration Rules: The paper establishes that while the coefficients of the closed form are not scale-invariant (frozen coefficients fail at large $N$), the form itself holds up to $N=100$. Two coefficients follow clean $1/N$ laws, and the anchor coefficient grows with downward curvature. This allows for a "label-budget rule": a few dozen labels at a target size can recalibrate the form to match fine-tuned ensembles.

Results

• Accuracy: The closed form achieves a Mean Absolute Error (MAE) of 0.0244 nats on the in-distribution test set. This is 6.4 times the error of the base GNN (0.0038 nats) but significantly better than the classical mutual information baseline (0.1002 nats).
• Portability (Out-of-Distribution): When applied unchanged (fit once) to six out-of-distribution (OOD) pools (varying system sizes, structural shifts, and random 2D positions), the closed form outperforms the base GNN on five of the six pools and ties the sixth. This is attributed to the form's geometry-nonparametric nature; unlike the GNN, it does not ingest explicit atomic coordinates.
• Cross-Hamiltonian Transfer: When refitted per family, the six-feature form outperforms analytic baselines ($I$ and filtered $I$) on all tested Hamiltonian families (TFIM, XXZ, J1-J2, Rydberg, and random-sign). While XGBoost on the same features is more accurate, it lacks interpretability.
• Scaling to $N=100$: Using independent DMRG data, the authors show that a global fit with frozen coefficients fails at large scales (MAE grows to $>1$ nat). However, refitting the coefficients per size maintains an MAE of 25–50 mnat. The fitted scaling laws for the coefficients allow for label-free deployment at roughly 40–80 mnat.
• Robustness: The form is robust to shot noise, outperforming mutual information baselines with as few as ~160 measurement shots.

Significance and Claims
The paper claims that Eq. 1 occupies a unique niche in the accuracy-interpretability trade-off. It is not a "distillation" of the GNN in the sense of reconstructing its internal weights; rather, it is an independent, low-dimensional ground-truth fit on the same physical quantities the GNN utilizes. The GNN encodes these quantities in a redundant, high-dimensional code, whereas Eq. 1 provides a readable projection.

The significance lies in providing a transparent, transferable, and tunable estimator. It allows experimentalists to inspect the physics driving the entropy estimate, apply the model across different geometries and phases without retraining, and calibrate it with a minimal number of labels. The authors conclude that while nonlinear models (like XGBoost or the GNN) offer higher raw accuracy given large labeled datasets, the closed form offers superior transparency and characterized scaling behavior, fixing a "label-budget rule" for large-scale quantum simulations.