Tree tensor networks for generative modeling

Matrix product states (MPSs), a tensor network designed for one-dimensional quantum systems, were recently proposed for generative modeling of natural data (such as images) in terms of the “Born machine.” However, the exponential decay of correlation in MPSs restricts its representation power heavily for modeling complex data such as natural images. In this work, we push forward the effort of applying tensor networks to machine learning by employing the tree tensor network (TTN), which exhibits balanced performance in expressibility and efficient training and sampling. We design the tree tensor network to utilize the two-dimensional prior of the natural images and develop sweeping learning and sampling algorithms which can be efficiently implemented utilizing graphical processing units. We apply our model to random binary patterns and the binary MNIST data sets of handwritten digits. We show that the TTN is superior to MPSs for generative modeling in keeping the correlation of pixels in natural images, as well as giving better log-likelihood scores in standard data sets of handwritten digits. We also compare its performance with state-of-the-art generative models such as variational autoencoders, restricted Boltzmann machines, and PixelCNN. Finally, we discuss the future development of tensor network states in machine learning problems.

Figure 1

(a) Tree-structure factor graph, where each block denotes a random variable with a value of $-1$ or 1, in which the blue (purple) block represents a hidden (visible) variable. The edge between two blocks introduces a factor function $f^{\left(k\right)}$ of those two variables. By adjusting those factor functions, the model could obtain the appropriate joint probability. (b) Tree tensor network, where $x_i$ denote the value of the $i\mathrm{th}$ pixel of the data set. Each yellow circle denotes a two- or three-order tensor. The edge between two tensors denotes a share index of tensors, which is also called a virtue index in the literature and will be contracted later. The exposed edge denotes the so-called physical index of tensors; those indices would ultimately be determined by the data set. For one of the given configurations of the physical indices, the probability of the configuration is proportional to the final scale value of the TTN after contracting all the virtue indices.

Figure 2

(a) The TTN with 2D structure. Changing the 1D order of data with the 2D order is equivalent to using the TTN with 2D structure replacing Fig. 1. (b) The 1D order of data. (c) The 2D order of data.

Figure 3

(a) Training the NLL of the TTN Born machine as a function of the data size $\left|\mathcal{T}\right|$; the system size is $N=16$. (b) Training the NLL of the TTN Born machine as a function of the system size $N$; the data size $\left|\mathcal{T}\right|=50$.

Figure 4

Comparison between the TTN and MPS Born machines trained on ten random patterns with different system sizes. As the system size become larger, MPS can no longer reach the theoretical limit of the NLL when $D_{\max }$ equals the number of samples, while the TTN is almost unaffected by the system size. This is because the structure of the TTN can capture the long-range dependences better.

Figure 5

Training the NLL and sampling images for a $\left|\mathcal{T}\right|=100$ binarized MNIST data set. $ln\left(\right|\mathcal{T}\left|\right)$ is the theoretical minimum of the NLL. The TTN exactly remembers all the information of the images when $D_{\max }=\left|\mathcal{T}\right|$.

Figure 6

Two-site correlation of pixels extracted from the original data (first row), the MPS (second row), and the TTN model (third row). We randomly choose three pixels at the tenth row of the images. The $D_{\max }$ of the TTN is 50; the $D_{\max }$ of the MPS is 100, which corresponds to the models with the smallest test NLL.

Figure 7

Images generated by the MPS and TTN that are trained on $\left|\mathcal{T}\right|=50000$ training images. (a) Part of the training images. (b) MPS with $D_{\max }=100$, test NLL = 101.45, (c) 1D TTN with $D_{\max }=50$, test NLL = 96.88, and (d) 2D TTN with $D_{\max }=50$, test NLL = 94.25.