Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Initialization and training of matrix product state probabilistic models

Xun Tang, Yuehaw Khoo, Lexing Ying

TL;DR

The authors address trainability issues of matrix product state (MPS) based probabilistic models, including Born machines, when modeling distributions with nonlocal correlations. They identify a causality trap under standard gradient descent and show local minima in MPS tomography, then propose two remedies: natural gradient descent (NGD) to perform updates in tensor space via a projected gradient flow, and TTNS-Sketch warm-start initialization to begin optimization near the ground truth. NGD yields rapid convergence to the global minimum and avoids degeneracies observed with GD and 2-site DMRG, while TTNS-Sketch initialization prevents the causality trap and enables fast convergence from a favorable starting point. Together, these strategies enhance reliability and efficiency of quantum-inspired generative modeling and quantum-state tomography for systems with strong boundary or nonlocal interactions, with $p_{\theta}$ and $\ ext{NLL}$ guiding optimization and $I_{X_1,X_n}$ informing model sufficiency under nonlocal correlations.

Abstract

Modeling probability distributions via the wave function of a quantum state is central to quantum-inspired generative modeling and quantum state tomography (QST). We investigate a common failure mode in training randomly initialized matrix product states (MPS) using gradient descent. The results show that the trained MPS models do not accurately predict the strong interactions between boundary sites in periodic spin chain models. In the case of the Born machine algorithm, we further identify a causality trap, where the trained MPS models resemble causal models that ignore the non-local correlations in the true distribution. We propose two complementary strategies to overcome the training failure -- one through optimization and one through initialization. First, we develop a natural gradient descent (NGD) method, which approximately simulates the gradient flow on tensor manifolds and significantly enhances training efficiency. Numerical experiments show that NGD avoids local minima in both Born machines and in general MPS tomography. Remarkably, we show that NGD with line search can converge to the global minimum in only a few iterations. Second, for the BM algorithm, we introduce a warm-start initialization based on the TTNS-Sketch algorithm. We show that gradient descent under a warm initialization does not encounter the causality trap and admits rapid convergence to the ground truth.

Initialization and training of matrix product state probabilistic models

TL;DR

The authors address trainability issues of matrix product state (MPS) based probabilistic models, including Born machines, when modeling distributions with nonlocal correlations. They identify a causality trap under standard gradient descent and show local minima in MPS tomography, then propose two remedies: natural gradient descent (NGD) to perform updates in tensor space via a projected gradient flow, and TTNS-Sketch warm-start initialization to begin optimization near the ground truth. NGD yields rapid convergence to the global minimum and avoids degeneracies observed with GD and 2-site DMRG, while TTNS-Sketch initialization prevents the causality trap and enables fast convergence from a favorable starting point. Together, these strategies enhance reliability and efficiency of quantum-inspired generative modeling and quantum-state tomography for systems with strong boundary or nonlocal interactions, with and guiding optimization and informing model sufficiency under nonlocal correlations.

Abstract

Modeling probability distributions via the wave function of a quantum state is central to quantum-inspired generative modeling and quantum state tomography (QST). We investigate a common failure mode in training randomly initialized matrix product states (MPS) using gradient descent. The results show that the trained MPS models do not accurately predict the strong interactions between boundary sites in periodic spin chain models. In the case of the Born machine algorithm, we further identify a causality trap, where the trained MPS models resemble causal models that ignore the non-local correlations in the true distribution. We propose two complementary strategies to overcome the training failure -- one through optimization and one through initialization. First, we develop a natural gradient descent (NGD) method, which approximately simulates the gradient flow on tensor manifolds and significantly enhances training efficiency. Numerical experiments show that NGD avoids local minima in both Born machines and in general MPS tomography. Remarkably, we show that NGD with line search can converge to the global minimum in only a few iterations. Second, for the BM algorithm, we introduce a warm-start initialization based on the TTNS-Sketch algorithm. We show that gradient descent under a warm initialization does not encounter the causality trap and admits rapid convergence to the ground truth.
Paper Structure (20 sections, 3 theorems, 34 equations, 9 figures, 3 algorithms)

This paper contains 20 sections, 3 theorems, 34 equations, 9 figures, 3 algorithms.

Table of Contents

  1. Introduction
  2. Causality trap in the Born machine algorithm
  3. Background in Born machine
  4. Causality trap under a periodic Ising model
  5. Local minima in MPS state tomography
  6. Background in MPS tomography
  7. Local minima under a periodic TFIM model
  8. Natural gradient descent algorithm for MPS optimization
  9. Algorithm
  10. Gradient flow perspective
  11. Numerical experiment with NGD
  12. Avoiding causality trap with TTNS-Sketch initialization
  13. Warm-start initialization of BM with TTNS-Sketch
  14. Numerical experiment with TTNS-Sketch initialization
  15. Conclusion and outlook
  16. ...and 5 more sections

Key Result

Proposition 1

For any site $i \in \{ 1, \ldots, n\}$, we let $S_{i} = \{(\delta G_{k})_{k = 1}^{n} \mid \delta G_{k} = 0 \, \forall \, k \not = i\}$ and we consider In other words, we let $\delta \theta_t$ be the solution to the minimization task in eqn: NGD for TN from only changing the tensor component at $G_i$. Then, one has where $\Pi_{i}$ denotes the projection onto the tangent space of varying $q_{\thet

Figures (9)

  • Figure 1: Illustration of NGD and warm initialization for the Born machine algorithm. One sees that the gradient descent method and the 2-site DMRG method do not converge to the optimal log-likelihood level. The ground truth model is a periodic ferromagnetic Ising model, and the experiment details are in \ref{['sec: analysis']}.
  • Figure 2: Graphical representations of the underlying model $p^{\star}$ (Fig. \ref{['fig:circle']}) and of the mis-specified model $p^{\mathrm{causal}}$ (Fig. \ref{['fig:line_16']}).
  • Figure 3: Performance of Born machine algorithm for the periodic spin system model in \ref{['eqn: periodic model']}. The models are initialized randomly and trained under gradient descent. The NLL gap is $0.33$, which coincides with the mutual information level of $(X_1, X_n)$ in $p^{\star}$. \ref{['appendix: KL analysis']} shows that the NLL gap and the mutual information level are approximately equal under the causality trap.
  • Figure 4: Plot of total variation (TV) distance of the trained Born machine model with respect to the true model $p^{\star}$ and the causal model $p^{\mathrm{causal}}$. The setting is the same as in \ref{['fig:result of GD']}. One can see that the trained BM model is much closer to the causal model than to the true model. The TV distance is defined by $\lVert p -p' \rVert_{\mathrm{TV}} = \frac{1}{2}\lVert p - p' \rVert_{1}$.
  • Figure 5: Performance of MPS tomography algorithm for the ground state of the periodic TFIM model in \ref{['eqn: TFIM']}. The models are initialized randomly and trained under gradient descent.
  • ...and 4 more figures

Theorems & Definitions (3)

  • Proposition 1
  • Proposition 2
  • Lemma 3