Supervised Learning with Quantum-Inspired Tensor Networks

E. Miles Stoudenmire, David J. Schwab

I Introduction

The connection between machine learning and statistical physics has long been appreciated Hopfield (1982); Amit et al. (1985); Derrida et al. (1987); Amit et al. (1987); Seung et al. (1992); Engel and Van den Broeck (2001); Malzahn and Opper (2005); Hinton et al. (2006); Mezard and Montanari (2009), but deeper relationships continue to be uncovered. For example, techniques used to pre-train neural networks Hinton et al. (2006) have more recently been interpreted in terms of the renormalization group Mehta and Schwab (2014). In the other direction there has been a sharp increase in applications of machine learning to chemistry, material science, and condensed matter physics Fischer et al. (2006); Hautier et al. (2010); Rupp et al. (2012); Saad et al. (2012); Snyder et al. (2012); Pilania et al. (2013); Arsenault et al. (2014); Amin et al. (2016); Carrasquilla and Melko (2017), which are sources of highly-structured data and could be a good testing ground for machine learning techniques.

A recent trend in both physics and machine learning is an appreciation for the power of tensor methods. In machine learning, tensor decompositions can be used to solve non-convex optimization tasks Anandkumar et al. (2014a, b) and make progress on many other important problems Phan and Cichocki (2010); Bengua et al. (2015); Novikov et al. (2015), while in physics, great strides have been made in manipulating large vectors arising in quantum mechanics by decomposing them as tensor networks Bridgeman and Chubb (2016); Schollwöck (2011); Evenbly and Vidal (2011). The most successful types of tensor networks avoid the curse of dimensionality by incorporating only low-order tensors, yet accurately reproduce very high-order tensors through a particular geometry of tensor contractions Evenbly and Vidal (2011).

Another context where very large vectors arise is in non-linear kernel learning, where input vectors x\mathbf{x} are mapped into a higher dimensional space via a feature map Φ(x)\Phi(\mathbf{x}) before being classified by a decision function

The feature vector Φ(x)\Phi(\mathbf{x}) and weight vector WW can be exponentially large or even infinite. One approach to deal with such large vectors is the well-known kernel trick, which only requires working with scalar products of feature vectors, allowing these vectors to be defined only implicitly Muller et al. (2001).

In what follows we propose a rather different approach. For certain learning tasks and a specific class of feature map Φ\Phi, we find the optimal weight vector WW can be approximated as a tensor network, that is, as a contracted sequence of low-order tensors. Representing W as a tensor network and optimizing it directly (without passing to the dual representation) has many interesting consequences. Training the model scales linearly in the training set size; the cost for evaluating an input is independent of training set size. Tensor networks are also adaptive: dimensions of tensor indices internal to the network grow and shrink during training to concentrate resources on the particular correlations within the data most useful for learning. The tensor network form of W presents opportunities to extract information hidden within the trained model and to accelerate training by using techniques such as optimizing different internal tensors in parallel Stoudenmire and White (2013). Finally, the tensor network form is an additional type of regularization beyond the choice of feature map, and could have interesting consequences for generalization.

One of the best understood types of tensor networks is the matrix product state Östlund and Rommer (1995); Schollwöck (2011), also known as the tensor train decomposition Oseledets (2011). Matrix product states (MPS) have been very useful for studying quantum systems, and have recently been proposed for machine learning applications such as learning features of images Bengua et al. (2015) and compressing the weight layers of neural networks Novikov et al. (2015). Though MPS are best suited for describing one-dimensional systems, they are powerful enough to be applied to higher-dimensional systems as well.

There has been intense research into generalizations of MPS better suited for higher dimensions and critical systems Verstraete and Cirac (2004); Vidal (2007); Evenbly and Vidal (2009). Though our proposed approach could generalize to these other types of tensor networks, as a proof of principle we will only consider the MPS decomposition in what follows. The MPS decomposition approximates an order-N tensor by a contracted chain of N lower-order tensors shown in Fig. 1. (Throughout we will use tensor diagram notation; for a brief review see Appendix A.)

Representing the weights WW of Eq. (1) as an MPS allows us to efficiently optimize these weights and adaptively change their number by varying WW locally a few tensors at a time, in close analogy to the density matrix renormalization group algorithm used in physics White (1992); Schollwöck (2011). Similar alternating least squares methods for tensor trains have also been explored in applied mathematics Holtz et al. (2012).

This paper is organized as follows: we propose our general approach then describe an algorithm for optimizing the weight vector WW in MPS form. We test our approach, both on the MNIST handwritten digit set and on two-dimensional toy data to better understand the role of the local feature-space dimension dd. Finally, we discuss the class of functions realized by our proposed models as well as a possible generative interpretation.

Those wishing to reproduce our results can find sample codes based on the ITensor library ITe at: https://github.com/emstoudenmire/TNML

II Encoding Input Data

The most successful use of tensor networks in physics so far has been in quantum mechanics, where combining NN independent systems corresponds to taking the tensor product of their individual state vectors. With the goal of applying similar tensor networks to machine learning, we choose a feature map of the form

The tensor Φs1s2⋯sN\Phi^{s_{1}s_{2}\cdots s_{N}} is the tensor product of the same local feature map ϕsj(xj)\phi^{s_{j}}(x_{j}) applied to each input xjx_{j}, where the indices sjs_{j} run from 11 to dd; the value dd is known as the local dimension. Thus each xjx_{j} is mapped to a dd-dimensional vector, which we require to have unit norm; this implies each Φ(x)\Phi(\mathbf{x}) also has unit norm.

The full feature map Φ(x)\Phi(\mathbf{x}) can be viewed as a vector in a dNd^{N}-dimensional space or as an order-NN tensor. The tensor diagram for Φ(x)\Phi(\mathbf{x}) is shown in Fig. 2. This type of tensor is said be rank-1 since it is manifestly the product of NN order-1 tensors. In physics terms, Φ(x)\Phi(\mathbf{x}) has the same structure as a product state or unentangled wavefunction.

For a concrete example of this type of feature map, consider inputs which are grayscale images with NN pixels, where each pixel value ranges from 0.0 for white to 1.0 for black. If the grayscale pixel value of the jthj^{\text{th}} pixel is xj∈x_{j}\in, a simple choice for the local feature map ϕsj(xj)\phi^{s_{j}}(x_{j}) is

and is illustrated in Fig. 3. The full image is represented as a tensor product of these local vectors. From a physics perspective, ϕsj\phi^{s_{j}} is the normalized wavefunction of a single qubit where the “up” state corresponds to a white pixel, the “down” state to a black pixel, and a superposition corresponds to a gray pixel.

While our choice of feature map Φ(x)\Phi(\mathbf{x}) was originally motivated from a physics perspective, in machine learning terms, the feature map Eq. (2) defines a kernel which is the product of NN local kernels, one for each component xjx_{j} of the input data. Kernels of this type have been discussed previously (Vapnik, 2000, p. 193) and have been argued to be useful for data where no relationship is assumed between different components of the input vector prior to learning Waegeman et al. (2012).

Though we will use only the local feature map Eq. (3) in our MNIST experiment below, it would be interesting to try other local maps and to understand better the role they play in the performance of the model and the cost of optimizing the model.

III Multiple Label Classification

IV MPS Approximation

A tensor network approximates the exponentially large set of components of a high-order tensor in terms of a much smaller set of parameters whose number grows only polynomially in the size of the input space. Various tensor network approximations impose different assumptions, or implicit priors, about the pattern of correlation of the local indices when viewing the original tensor as a distribution. For example, a MERA network can explicitly model power-law decaying correlations while a matrix product state (MPS) has exponentially decaying correlations Evenbly and Vidal (2011). Yet an MPS can still approximate power-law decays over quite long distances.

In typical physics applications the MPS bond dimension mm can range from 10 to 10,000 or even more; for the most challenging physics systems one wants to allow as large a bond dimension as possible since a larger dimension means more accuracy. However, when using MPS in a machine learning setting, the bond dimension controls the number of parameters of the model. So in contrast to physics, taking too large a bond dimension might not be desirable as it could lead to overfitting.

V Sweeping Algorithm for Optimizing Weights

Inspired by the very successful density matrix renormalization group (DMRG) algorithm developed in physics, here we propose a similar algorithm which “sweeps” back and forth along an MPS, iteratively minimizing the cost function defining the classification task.

For concreteness, let us choose to optimize the quadratic cost

Our strategy for reducing this cost function will be to vary only two neighboring MPS tensors at a time within the approximation Eq. (5). We could conceivably just vary one at a time but as will become clear, varying two tensors leads to a straightforward method for adaptively changing the MPS bond dimension.

Finally, when proceeding to the next bond, it would be inefficient to fully project each training input over again into the configuration in Fig. 6(b). Instead it is only necessary to advance the projection by one site using the MPS tensor set from a unitary matrix after the SVD as shown in Fig. 7(c). This allows the cost of each local step of the algorithm to remain independent of the size of the input space, making the total algorithm scale only linearly with input space size.

The scaling of the above algorithm is d3 m3 N NL NTd^{3}\,m^{3}\,N\,N_{L}\,N_{T}, where recall mm is the MPS bond dimension; NN the number of input components; NLN_{L} the number of labels; and NTN_{T} the number of training inputs. In practice, the cost is dominated by the large number of training inputs NTN_{T}, so it would be very desirable to reduce this cost. One solution could be to use stochastic gradient descent, but while our experiments at blending this approach with the MPS sweeping algorithm often reached single-digit classification errors, we could not match the accuracy of the full gradient. Mixing stochastic gradient with MPS sweeping thus appears to be non-trivial but we believe it is a promising direction for further research.

Finally, we note that a related optimization algorithm was proposed for hidden Markov models in Ref. Rolfe and Cook, 2010. However, in place of our SVD above, Ref. Rolfe and Cook, 2010 uses a non-negative matrix factorization. In a setting where negative weights are allowed, the SVD is the optimal choice because it minimizes the distance between the original tensor and product of factorized tensors. Furthermore, our framework for sharing weights across multiple labels and our use of local feature maps has interesting implications for training performance and for generalization.

VI MNIST Handwritten Digit Test

To test the tensor network approach on a realistic task, we used the MNIST data set, which consists of grayscale images of the digits zero through nine Yann LeCun . The calculations were implemented using the ITensor library ITe . Each image was originally 28×2828\times 28 pixels, which we scaled down to 14×1414\times 14 by averaging clusters of four pixels; otherwise we performed no further modifications to the training or test sets. Working with smaller images reduced the time needed for training, with the tradeoff being that less information was available for learning.

To approximate the classifier tensors as MPS, one must choose a one-dimensional ordering of the local indices s1,s2,…,sNs_{1},s_{2},\ldots,s_{N}. We chose a “zig-zag” ordering shown in Fig. 8, which on average keeps spatially neighboring pixels as close to each other as possible along the one-dimensional MPS path. We then mapped each grayscale image x\mathbf{x} to a tensor Φ(x)\Phi(\mathbf{x}) using the local map Eq. (3).

Using the sweeping algorithm in Section V to train the weights, we found the algorithm quickly converged in the number of passes, or sweeps over the MPS. Typically only two or three sweeps were needed to see good convergence, with test error rates changing only hundreths of a percent thereafter.

Test error rates also decreased rapidly with the maximum MPS bond dimension mm. For m=10m=10 we found both a training and test error of about 5%; for m=20m=20 the error dropped to only 2%. The largest bond dimension we tried was m=120m=120, where after three sweeps we obtained a test error of 0.97% (97 misclassified images out of the test set of 10,000 images); the training set error was 0.05% or 32 misclassified images.

VII Two-Dimensional Toy Model

To better understand the modeling power and regularization properties of the class of models presented in Sections II and III, consider a family of toy models where the input space is two-dimensional (N=2N=2). The hidden distribution we want to learn consists of two distributions, PA(x1,x2)P_{A}(x_{1},x_{2}) and PB(x1,x2)P_{B}(x_{1},x_{2}), from which we generate training data points labeled AA or BB respectively. For simplicity we only consider the square region x1∈x_{1}\in and x2∈x_{2}\in.

To train the model, each training point (x1,x2)(x_{1},x_{2}) is mapped to a tensor

When selecting a model, our main control parameter is the dimension dd of the local indices s1s_{1} and s2s_{2}. For the case d=2d=2, the local feature map is chosen as in Eq. 3. For d>2d>2 we generalize ϕsj(xj)\phi^{s_{j}}(x_{j}) to be a normalized dd-component vector as described in Appendix B.

To understand how the flexibility of the model grows with increasing dd, consider the case where PAP_{A} and PBP_{B} are overlapping distributions. Specifically, we take each to be a multivariate Gaussian centered respectively in the lower-right and upper-left of the unit square, and to have different covariance matrices. In Fig. 9 we show the theoretically optimal decision boundary that best separates AA points (crosses, red region) from BB points (squares, blue region), defined by the condition PA(x1,x2)=PB(x1,x2)P_{A}(x_{1},x_{2})=P_{B}(x_{1},x_{2}). To make a training set, we sample 100 points from each of the two distributions.

Next, we optimize the toy model for our overlapping training set for various dd. The decision boundary learned by the d=2d=2 model in Fig. 10(a) shows good agreement with the optimal one in Fig. 9. Because the two sets are non-separable and this model is apparently well regularized, some of the training points are necessarily misclassified—these points are colored white in the figure.

The d=3d=3 decision boundary shown in Fig. 10 begins to show evidence of overfitting. The boundary is more complicated than for d=2d=2 and further from the optimal boundary. Finally, for a much larger local dimension d=6d=6 there is extreme overfitting. The decision boundary is highly irregular and is more reflective of the specific sampled points than the underlying distribution. Some of the overfitting behavior reveals the structure of the model; at the bottom and top of Fig. 10(c) there are lobes of one color protruding into the other. These likely indicate that the finite local dimension still somewhat regularizes the model; otherwise it would be able to overfit even more drastically by just surrounding each point with a small patch of its correct color.

VII.2 Non-Linear Decision Boundary

To test the ability of our proposed class of models to learn highly non-linear decision boundaries, consider the spiral shaped boundary in Fig. 11(a). Here we take PAP_{A} and PBP_{B} to be non-overlapping with PAP_{A} uniform on the red region and PBP_{B} uniform on the blue region.

In Fig. 11(b) we show the result of training a model with local dimension d=10d=10 on 500 sampled points, 250 for each region (crosses for region AA, squares for region BB). The learned model is able to classify every training point correctly, though with some overfitting apparent near regions with too many or too few sampled points.

VIII Interpreting Tensor Network Models

If the functions {ϕs(x)}\{\phi^{s}(x)\}, s=1,2,…,ds=1,2,\ldots,d form a basis for a Hilbert space of functions over x∈x\in, then the tensor product basis

forms a basis for a Hilbert space of functions over x∈×N\mathbf{x}\in^{\times N}. Moreover, if the basis {ϕs(x)}\{\phi^{s}(x)\} is complete, then the tensor product basis is also complete and f(x)f(\mathbf{x}) can be any square integrable function.

Next, consider the effect of restricting the local dimension to d=2d=2 as in the local feature map of Eq. (3) which was used to classify grayscale images in our MNIST benchmark in Section VI. Recall that for this choice of ϕ(x)\phi(x),

Thus if x^\hat{\mathbf{x}} is a black and white image with pixel values of only x^j={0,1}\hat{x}_{j}=\{0,1\}, then f(x^)f(\hat{\mathbf{x}}) is equal to a single component Ws1s2…sNW_{s_{1}s_{2}\ldots s_{N}} of the weight tensor. Because each of these components is an independent parameter (assuming no further approximation of WW), f(x^)f(\hat{\mathbf{x}}) is a highly non-linear, in fact arbitrary, function when restricted to these black and white images.

Returning to the case of grayscale images x\mathbf{x} with pixels xj∈x_{j}\in, f(x)f(\mathbf{x}) cannot be an arbitrary function over this larger space of images for finite dd. For example, if one considers the d=2d=2 feature map Eq. (3), then when considering the dependence of f(x)f(\mathbf{x}) on only a single pixel xjx_{j} (all other pixels being held fixed), it has the functional form acos⁡(π/2 xj)+bsin⁡(π/2 xj)a\cos(\pi/2\,x_{j})+b\sin(\pi/2\,x_{j}) where aa and bb are constants determined by the (fixed) values of the other pixels.

VIII.2 Implicit Feature and Kernel Selection

as illustrated in Fig. 12(a). To say the UU and VV tensors are left or right orthogonal means when viewed as matrices Uαj−1sj αjU_{\alpha_{j-1}s_{j}}\,^{\alpha_{j}} and Vαj−1 sjαjV^{\alpha_{j-1}}\,_{s_{j}\alpha_{j}} these tensors have the property U†U=IU^{\dagger}U=I and VV†=IVV^{\dagger}=I where II is the identity; these orthogonality conditions can be understood more clearly in terms of the diagrams in Fig. 12(b). Any MPS can be brought into the form Eq. (17) through an efficient sequence of tensor contractions and SVD operations similar to the steps in Fig. 7(b).

The above interpretation implies that training an MPS model uncovers a relatively small set of important features and simulatenously learns a decision function based only on these reduced features. This picture is similar to popular interpretations of the hidden and output layers of shallow neural network models Nielsen (2015). A similar interpretation of an MPS as learning features was first proposed in Ref. Bengua et al., 2015, though with quite a different scheme for representing data than what is used here. It is also interesting to note that an interpretation of the UU and VV tensors as combining and projecting features into only the mm most important combinations can be applied at any bond of the MPS. For example, the tensor Uαjsjαj+1U^{\alpha_{j+1}}_{\alpha_{j}s_{j}} tensor at site jj can be viewed as defining a vector of mm features labeled by αj+1\alpha_{j+1} by forming linear combinations of products of the features ϕsj(xj)\phi^{s_{j}}(x_{j}) and the features αj\alpha_{j} defined by the previous UU tensor, similar to the contraction in Fig. 7(c).

VIII.3 Generative Interpretation

This condition is automatically satisfied for tensor-product feature maps Φ(x)\Phi(\mathbf{x}) of the form Eq. (2) if the constituent local maps ϕs(x)\phi^{s}(x) have the property

that is, if the components of ϕs\phi^{s} are orthonormal functions with respect to the measure dμ(x)d\mu(x). Furthermore, if one wants to demand, after mapping to feature space, that any input x\mathbf{x} itself defines a normalized distribution, then we also require the local vectors to be normalized as

Unfortunately neither the local feature map Eq. (3) nor its generalizations in Appendix B meet the first criterion Eq. (20). A different choice that satisfies both the orthogonality condition Eq. (20) and normalization condition Eq. (21) could be

However, this map is not suitable for inputs like grayscale pixels since it is anti-periodic over the interval x∈x\in and would lead to a periodic probability distribution. An example of an orthogonal, normalized map which is not periodic on x∈x\in is

This local feature map meets the criteria Eqs. (20) and (21) if the integration measure chosen to be dμ(x)=2dxd\mu(x)=2dx.

As a basic consistency check of the above generative interpretation, we performed an experiment on our toy model of Section VII, using the local feature map Eq. (23). Recall that our toy data can have two possible labels AA and BB. To test the generative interpretation, we first generated a single, random “hidden” weight tensor WW. From this weight tensor we sampled NsN_{s} data points in a two step process:

where recall LnL_{n} is the known correct label for training point nn.

We repeated this procedure multiple times for various sample sizes NsN_{s}, each time computing the Kullback-Liebler divergence of the learned versus exact distribution

IX Discussion

We have introduced a framework for applying quantum-inspired tensor networks to multi-class supervised learning tasks. While using an MPS ansatz for the model parameters worked well even for the two-dimensional data in our MNIST experiment, other tensor networks such as PEPS, which are explicitly designed for two-dimensional systems, may be more suitable and offer superior performance. Much work remains to determine the best tensor network for a given domain.

Representing the parameters of our model by a tensor network has many useful and interesting implications. It allows one to work with a family of non-linear kernel learning models with a cost that is linear in the training set size for optimization, and independent of training set size for evaluation, despite using a very expressive feature map (recall in our setup, the dimension of feature space is exponential in the size of the input space). There is much room to improve the optimization algorithm we described, adopting it to incorporate standard tricks such as mini-batches, momentum, or adaptive learning rates. It would be especially interesting to investigate unsupervised techniques for initializing the tensor network.

Additionally, while the tensor network parameterization of a model clearly regularizes it in the sense of reducing the number of parameters, it would be helpful to understand the consquences of this regularization for specific learning tasks. It could also be fruitful to include standard regularizations of the parameters of the tensor network, such as weight decay or L1L_{1} penalties. We were surprised to find good generalization without using explicit parameter regularization. For issues of interpretability, the fact that tensor networks are composed only of linear operations could be extremely useful. For example, it is straightforward to determine directions in feature space which are orthogonal to (or projected to zero by) the weight tensor WW.

There exist tensor network coarse-graining approaches for purely classical systems Efrati et al. (2014); Evenbly and Vidal (2015), which could possibly be used instead of our approach. However, mapping the data into an extremely high-dimensional Hilbert space is likely advantageous for producing models sensitive to high-order correlations among features. We believe there is great promise in investigating the power of quantum-inspired tensor networks for many other machine learning tasks.

Note: while preparing our final manuscript, Novikov et al. Novikov et al. (2016) published a related framework for parameterizing supervised learning models with MPS (tensor trains).

We would like to acknowledge helpful discussions with Juan Carrasquilla, Josh Combes, Glen Evenbly, Bohdan Kulchytskyy, Li Li, Roger Melko, Pankaj Mehta, U. N. Niranjan, Giacomo Torlai, and Steven R. White. This research was supported in part by the Perimeter Institute for Theoretical Physics. Research at Perimeter Institute is supported by the Government of Canada through Industry Canada and by the Province of Ontario through the Ministry of Economic Development & Innovation. This research was also supported in part by the Simons Foundation Many-Electron Collaboration.

Appendix A Graphical Notation for Tensor Networks

Though matrix product states (MPS) have a relatively simple structure, more powerful tensor networks, such as PEPS and MERA, have such complex structure that traditional tensor notation becomes unwieldy. For these networks, and even for MPS, it is helpful to use a graphical notation. For some more complete reviews of this notation and its uses in various tensor networks see Ref. Bridgeman and Chubb, 2016; Cichocki, 2014.

The basic graphical notation for a tensor is to represent it as a closed shape. Typically this shape is a circle, though other shapes can be used to distinguish types of tensors (there is no standard convention for the choice of shapes). Each index of the tensor is represented by a line emanating from it; an order-N tensor has N such lines. Figure 14 shows examples of diagrams for tensors of order one, two, and three.

To indicate that a certain pair of tensor indices are contracted, they are joined together by a line. For example, Fig. 15(a) shows the contraction of an an order-1 tensor with the an order-2 tensor; this is the usual matrix-vector multiplication. Figure 15(b) shows a more general contraction of an order-4 tensor with an order-3 tensor.

Graphical tensor notation offers many advantages over traditional notation. In graphical form, indices do not usually require names or labels since they can be distinguished by their location in the diagram. Operations such as the outer product, tensor trace, and tensor contraction can be expressed without additional notation; for example, the outer product is just the placement of one tensor next to another. For a network of contracted tensors, the order of the final resulting tensor can be read off by simply counting the number of unpaired lines left over. For example, a complicated set of tensor contractions can be recognized as giving a scalar result if no index lines remain unpaired.

Finally, we note that a related notation for sparse or structured matrices in a direct-sum formalism can be used, and appears extensively in Ref. Fishman and White, 2015.

Appendix B Higher-Dimensional Local Feature Map

As discussed in Section II, our strategy for using tensor networks to classify input data begins by mapping each component xjx_{j} of the input data vector x\mathbf{x} to a dd-component vector ϕsj(xj)\phi^{s_{j}}(x_{j}), sj=1,2,…,ds_{j}=1,2,\ldots,d. We always choose ϕsj(xj)\phi^{s_{j}}(x_{j}) to be a unit vector in order to apply physics techniques which typically assume normalized wavefunctions.

For the case of d=2d=2 we have used the mapping

A straightforward way to generalize this mapping to larger dd is as follows. Define θj=π2xj\theta_{j}=\frac{\pi}{2}x_{j}. Because (cos⁡2(θj)+sin⁡2(θj))=1(\cos^{2}(\theta_{j})+\sin^{2}(\theta_{j}))=1, then also

Expand the above identity using the binomial coeffiecients (nk)=n!/(k!(n−k)!)\binom{n}{k}=n!/(k!(n-k)!).

This motivates defining ϕsj(xj)\phi^{s_{j}}(x_{j}) to be

where recall that sjs_{j} runs from 11 to dd. The above definition reduces to the d=2d=2 case Eq. (26) and guarantees that ∑sj∣ϕsj∣2=1\sum_{s_{j}}|\phi^{s_{j}}|^{2}=1 for larger dd. (These functions are actually a special case of what are known as spin coherent states.)

References