Efficient Generative Modeling with Unitary Matrix Product States Using Riemannian Optimization

Haotong Duan; Zhongming Chen; Ngai Wong

cs.LGMar 12, 2026

Efficient Generative Modeling with Unitary Matrix Product States Using Riemannian Optimization

Authors: Haotong Duan, Zhongming Chen, Ngai Wong

Organizations: Department of Mathematics, School of Sciences, Hangzhou Dianzi University, Hangzhou 310018, China. · Department of Electrical and Electronic Engineering, The University of Hong Kong, Hong Kong.

Abstract

Tensor networks, which are originally developed for characterizing complex quantum many-body systems, have recently emerged as a powerful framework for capturing high-dimensional probability distributions with strong physical interpretability. This paper systematically studies matrix product states (MPS) for generative modeling and shows that unitary MPS, which is a tensor-network architecture that is both simple and expressive, offers clear benefits for unsupervised learning by reducing ambiguity in parameter updates and improving efficiency. To overcome the inefficiency of standard gradient-based MPS training, we develop a Riemannian optimization approach that casts probabilistic modeling as an optimization problem with manifold constraints, and further derive an efficient space-decoupling algorithm. Experiments on Bars-and-Stripes and EMNIST datasets demonstrate fast adaptation to data structure, stable updates, and strong performance while maintaining the efficiency and expressive power of MPS.

Figures & tables

Panel of Figure 1 — Fig. 1: Optimization flowchart of the space-decoupling method, where dashed lines indicate the optimization process.

Operation	Time complexity
Two SVDs (lines 5 and 10)	$O(2nr_{\max}^{2}+r_{\max}^{3})$
Compute $\nabla f(X)$ via ( 22 )	$O(\|\mathcal{T}\|dn^{2}r_{\max}^{2})$
Compute $(\bar{K},\bar{V}_{p})$ via ( 26 ) (line 8)	$O(n^{2}r_{\max}^{3})$
Update $X$ (line 9)	$O(n^{2}r_{\max}^{3}))$
Object	Space complexity
MPS cores $A^{(1)},\dots,A^{(d)}$	$dnr_{\max}^{2}$

TABLE I: Summary of time and space complexity for the main operations and objects of Algorithm 1 .

Panel of Figure 3 — Fig. 2: Schematic of sequential sample generation from a UMPS. All tensors except $A^{(d)}$ are gauged to be left-canonical, Each tensor $A^{(i)}$ generates one bit $v_{i}$ conditionally on the bits to its right.

Panel of Figure 4 — Fig. 3: (a) Illustration of column-major flattening of a $16\times 16$ image into a 256-dimensional vector $v=(p_{1},p_{2},\dots,p_{16},p_{17},\dots,p_{256})^{\top}$ . (b) A subset of the Bars-and-Stripes dataset, with images of size $16\times 16$ .

Panel of Figure 5 — Fig. 4: (a) The blue curve represents the average NLL as a function of the number of loops, while the orange bars indicate the cumulative computation time required to reach the corresponding loops. (b) Images generated by the model at loops from 1 to 5 as indicated in (a). In experiments above, we set $r_{\rm max}=500$ , $|\mathcal{T}|=400$ , and the learning rate to $0.007$ .

$loops$	1	2	3	4	5
$r_{\rm mean}$	7.8945	31.1133	122.4883	470.7852	470.7852
$r_{\rm max}$	8	32	128	500	500

TABLE II: Under the experimental conditions shown in Fig. 4 , as the number of loops increases, the table shows the average bond dimension

r_{\rm mean}

and the largest bond dimension

r_{\rm max}

in the UMPS updated up to this step.

Panel of Figure 7 — Fig. 5: (a) Average NLL as a function of the dataset size under different $r_{\rm max}$ . For each curve, the circular markers from left to right correspond to the NLL values obtained at $|\mathcal{T}|=100,200,300,400,$ and $500$ , respectively. (b) The average computation time per experiment shown in (a).

Panel of Figure 8 — Fig. 6: Panels (a) and (b) correspond to the negative log-likelihood (NLL) convergence of the optimization algorithm of the MPS model and UMPS-SD algorithm for maximum bond dimensions $r_{\rm max}=200$ and $r_{\rm max}=400$ , respectively, with a training set of $|\mathcal{T}|=100$ samples. The blue line represents the optimization algorithm of the MPS model, and the red line represents the UMPS-SD algorithm.

$loops$	MPS		UMPS-SD
	NLL	Time/s	NLL	Time/s
1	154.74	1.30	167.70	1.68
2	94.48	3.14	80.69	5.61
3	62.25	7.86	13.01	30.25
4	46.85	26.08	-	-
5	40.14	53.61	-	-

TABLE III: Under the same experimental conditions as Fig. 6a , NLL and training time (s) for MPS and UMPS-SD models with learning rate

1\times 10^{-3}

. The time metric indicates the cumulative time required to reach this NLL value.

$loops$	UMPS-SD(1)		UMPS-SD(2)		UMPS-SD(3)
	NLL	Time/s	NLL	Time/s	NLL	Time/s
1	171.99	1.62	164.41	1.59	165.76	1.62
2	79.90	5.59	81.14	5.48	79.97	5.55
3	7.35	30.73	6.568	30.23	12.30	30.11

TABLE IV: Under the same experimental conditions as Fig. 6a , from (1) to (3), the learning rates are set to

2.5\times 10^{-4}

5\times 10^{-4}

, and

2\times 10^{-3}

, respectively.

Panel of Figure 11 — Fig. 7: (a) Generated from the MPS model with a maximum bond dimension $r_{\max}=400$ , trained on a subset of $|\mathcal{T}|=300$ samples randomly selected from the EMNIST dataset for $l_{\max}=25$ . (b) Generated by the proposed UMPS model with $l_{\max}=4$ , with all other settings kept identical to (a).

Panel of Figure 12 — Fig. 8: Image reconstruction from half of the images, which are selected from the training set. For both experiments, the maximum bond dimension is set to $r_{\rm max}=200$ and $|\mathcal{T}|=100$ . The update details are summarized in Table III . The black areas in the image are the given pixels, and the blue areas are the pixels completed by the model. (a) The image reconstructed by the MPS model; (b) The image reconstructed by the UMPS model.

Fig. 9: The curves illustrate how the NLL varies with the bond dimension $r_{\rm max}$ , where the training set $\mathcal{T}$ is selected from the EMNIST-Letters dataset with a fixed size of $|\mathcal{T}|=150$ . All marked points correspond to the setting of $l_{\max}=4$ . In the images, black pixels denote the given observations, while the blue pixels are reconstructed by the UMPS model. The figure presents the reconstruction results obtained for $r_{\rm max}=120$ , $150$ , and $175$ , respectively.

Tensor Networks

Tensor networks represent high-dimensional data and quantum states as interconnected low-rank tensors to enable efficient computation, learning, and simulation.

See all 49 papers

Quantum Tensor Network Learning with DMRG
Aug 19, 2026Gustav J L Jäger, Martin B Plenio, Hans-Martin RieserTensor Networks Quantum Machine Learning
Deutsches Zentrum für Luft- und Raumfahrt e.V. – Institut für KI-Sicherheit Wilhelm-Runge-Straße 10, 89081 Ulm, Germany · Universität Ulm – Institut für Theoretische Physik and IQST, Albert-Einstein-Allee 11, 89081 Ulm, Germany
Quantum-inspired tensor networks in machine learning models
Apr 15, 2026Guillermo Valverde, Igor García-Olaizola, Giannicola Scarpa +1Tensor Networks Quantum-Inspired ML
Vicomtech Foundation, Basque Research and Technology Alliance (BRTA), Donostia-San Sebastián 20009, Spain. · University of Deusto, Avda. de las Universidades, 24, Bilbao 48007, Spain. · Escuela Técnica Superior de Ingeniería de Sistemas Informáticos, Universidad Politécnica de Madrid, Madrid 28031, Spain. +1
Stochastic Optimization of Tree Tensor Networks
Sep 1, 2026Marius Willner, Maximilian Scharf, André Uschmajew +2Tensor Networks Riemannian Optimization
Institute of Mathematics, University of Augsburg, 86159 Augsburg, Germany · Centre for Advanced Analytics and Predictive Sciences, University of Augsburg, 86159 Augsburg, Germany · Tensor AI Solutions GmbH, 89284 Pfaffenhofen an der Roth, Germany +1

Aug 19, 2026J/K move · Enter open · S save

Quantum Tensor Network Learning with DMRG

Gustav J L Jäger, Martin B Plenio, Hans-Martin Rieser

Deutsches Zentrum für Luft- und Raumfahrt e.V. – Institut für KI-Sicherheit Wilhelm-Runge-Straße 10, 89081 Ulm, Germany · Universität Ulm – Institut für Theoretische Physik and IQST, Albert-Einstein-Allee 11, 89081 Ulm, Germany

Tensor Networks Quantum Machine Learning TT Decomposition Neural Network Optimization

Open paper AskPDF arXiv

Tensor Networks are a relatively new machine learning approach. The architectures proposed initially are inspired by approaches from quantum many-body physics simulations. One common layout is the matrix product state (MPS) also known as a tensor train optimized with gradient descent techniques. We introduce a global normalization condition, so that the MPS represents a quantum state. We investigate two optimization methods that find the locally optimal tensors and compare them regarding their effectiveness. One is based on gradient descent and the other on an adaptation of DMRG.

Abstract

Tensor Networks

Explore similar work

Quantum Tensor Network Learning with DMRG

Quantum-inspired tensor networks in machine learning models

Stochastic Optimization of Tree Tensor Networks