Infinitely Wide Tensor Networks as Gaussian Process

As is well known, Gaussian Process can be understood as the generalization of the multi-variate Gaussian distribution on the vector space to infinite dimensional function space which is defined by the mean function $\mu(x)$ and the covariance function $C(x,x^{\prime})$. Here $\mu(x)$ represents the mean of the random variable $X$ with index $x$, and the covariance $C(a,x^{\prime})$ represents the covariance between two random variable $X$ and $X^{\prime}$ with index $x$ and $x^{\prime}$. In the application case with finite size training data set, the finite dimensional distribution (f.d.d.s.) of the Gaussian Process is used to model the data set, and then the mean function degenerate to the finite dimensional vector $\mathbf{\mu}$ and the covariance function degenerate to the covariance matrix $C_{ij}$ which defines the f.d.d.s. of the GP, namely the multi-variates Gaussian distribution $\mathcal{N}(\mathbf{\mu},C_{ij})$.

The nice point of the GP is that the Bayesian updating of the GP prior is analytically tractable and simple. For a GP prior as follows,

We can get the joint distribution of the training data $\mathbf{x}$ and the $\mathbf{x}^{*}$ easily as

The predictive posterior distribution $p(f^{*}|\mathbf{x}^{*},\mathbf{x},f)$ can be obtained analytically as

It is easy to extend above analysis to the noisy label model, namely introducing noise to the label $\mathbf{y}=f(\mathbf{x})+\mathcal{\epsilon}$.

For a non-parametric model, it can adapt itself to simulate the data after the data comes but we can still use the Maximum Likelihood Estimator (M.L.E.) or the Maximum A Posterior (M.A.P.) to optimize the hyper-parameters. We can get the nice marginal likelihood as follows,

where $C_{\mathbf{y}}=C_{f}+\sigma_{n}^{2}\mathbb{I}$.

The training data only get into first term of above formula which evaluates how good our model fits the given training data, however the second term only depends on the covariance matrix which is a model complexity penalty term. We can undertand the competition between the data fitting and the model capacity by considering the length-scale. The model capacity will increase as the length-scale decrease, namely the model complexity penalty will increase although the data-fitting will get better. From the computation perspective, we know to get the optimal point of the marginal likelihood, we need to compute the inverse of the covariance matrix $C_{\mathbf{y}}^{-1}$ which is of order $O(n^{3})$, where $n$ is the training sample size. For application in the big data set, the inverse matrix computation will be highly time-consuming so we can consider using matrix approximation method such as outer-product approximation. See  (Williams and Rasmussen, 2006) for a good review on GP application in Machine Learning.

For an infinite-width neural network with response function $f(\mathbf{x})$ as

where $\mathbf{w}^{[i]}$ and $b^{[i]}$ represent the weights and the bias in the $i$’th layer and $a_{i}(\mathbf{w};\mathbf{x})$ represents the activation function. Since all parameters in the neural network are independent and identically distributed (i.i.d.), by the Central Limit Theorem (C.L.T.), we know as $i\to\infty$,

where $X$ belongs to the Normal distribution. Since $b^{[1]}$ is normal and independent with $\sum_{i}w_{i}^{[1]}a(w_{i}^{[0]};x)$, we know $f(\mathbf{x})$ is normal random variable. We will get a sequence of random variables $\{f(\mathbf{x}^{(i)}),i\in\{1,\cdots,m\}\}$ by evaluating the infinite neural network on the data set. For the neural network with several special non-linear activation functions, analytical expressions of the Gaussian Process are obtained  (Williams, 1997).

Above discussion is an intuition explaining of the equivalence between GP and infinite-width neural networks, here we will propose our first theorem on the equivalence between GP and neural networks and provide a strict proof in the Appendix A.

For a $n$ nodes Tensor Network, actually Matrix Product States (M.P.S.)  (Biamonte and Bergholm, 2017; Orús, 2014; Cichocki, 2014), the scalar response $f(\mathbf{x})$ is as follows,

where

Here $\phi^{s_{i}}(x_{i})$ represents the function which maps each component of the input data $x_{i}$ into a higher dimensional feature space. $A^{s_{i}}_{\alpha_{i}\alpha_{i+1}}$ represents the rank $3$ tensor node, where $s_{i}$ is the physical index which contracts with the $s_{i}$ index in the kernel $\phi^{s_{i}}$. $\alpha_{i}$ is the bond index by which we can tune the representation power of the tensor network.

M.P.S. is introduced to approximate a quantum state $\Ket{\psi}$ efficiently in quantum physics  (Orús, 2019; Chabuda et al., 2020; Schröder et al., 2019; McMahon et al., 2020). To fully describe a $n$-party quantum state $\Ket{\psi}$, we need $2^{n}$ basis which grows exponentially with respecting to $n$. By expanding the $n$-party quantum state with a chain of product of matrices, we can approximate the original quantum state with less parameters. For the M.P.S., the number of the basis grows linearly as $n$ increase.

Here we write a general n-party state as a matrix product chain with $O(s\alpha^{2}n)$ parameters.

Tensor Networks as a powerful machine has been widely studied to construct statistical learning model in Machine Learning community recently and also great achievements have been obtained in different tasks such as supervised pattern recognition (classification and regression), natural language processing and density estimation (unsupervised generative model)  (Stoudenmire and Schwab, 2016; Han et al., 2018; Pestun et al., 2017; Novikov et al., 2018). Since the special structure of tensor network, the model easily blows out or decays to dead nodes during training process, efficient and robust initialization strategy was proposed in  (Guo and Draper, 2021). It is known that the correlation function of tensor networks decays exponentially  (Evenbly and Vidal, 2011b, a), so it is discussed to extend the tensor network to more complicated structures such as multi-scale tree structure to reduce the correlation decay problem in machine learning application  (Stoudenmire, 2018). More extension work of tensor network can be found in  (Cheng et al., 2020; Li and Zhang, 2018; Cheng et al., 2020; Liu et al., 2018; Glasser et al., 2018).

For the infinite-width pure MPS, no activation functions and no latent variables are introduced and then the response function $\psi(\mathbf{x})$ is linear respecting to the weights $A^{\alpha_{i}}_{\alpha_{i}\alpha_{i+1}}$ in each tensor node. The representation power of the pure tensor network is limited since the pure MPS can only represent the linear function. To increase the representation capacity of MPS, we can introduce another layer of neural network with activation functions. Here we propose two methods to combine the MPS with the neural network:

• •

  Neural network kernel MPS;

• •

  MPS hidden layer neural network.

For the neural network kernel MPS, we use one layer neural network as the kernel function and ’non-linearity’ is introduced by the hidden layer and the activation function in neural network. For the MPS hidden layer neural network, the response $\psi^{l}(\mathbf{w},\mathbf{A};\mathbf{x})$ of tensor network is fed into the neural network by which the linearity is injected into the model.

In the following sections, we will discuss these two hybrid tensor neural network and the theorems on their Gaussian Process equivalence properties.

There are three indices in each tensor node $\mathbf{A}^{s_{i}}_{\alpha_{i}\alpha_{i+1}}$, so at first step we can contract the kernel $\Phi^{s_{1},\cdots,s_{n}}(\mathbf{x})$ with the indices sequence $\{s_{1},\cdots s_{n}\}$ in tensor network, and then we get

where

It is easy to verify that all the components of the $Z_{\alpha_{i}\alpha_{i+1}}$ are i.i.d. since all components of one specific tensor node $\mathbf{A}^{s_{i}}_{\alpha_{i}\alpha_{i+1}}$ are i.i.d. We note here that in our assumption of Theorem 2, we just assume that all tensor nodes are independent, but not necessary to be i.i.d. For two specific implementations of the bond indices, namely two different sequences of value of the indices, $\{i_{1},\cdots,i_{n}\}$ and $\{j_{1},\cdots,j_{n}\}$, the product of each chain of the nodes are i.i.d.. We introduce a new notation for the production of the chain,

We know all the different possible evaluation of the indices sequence $\{\alpha_{1}\cdots\alpha_{n}\}$ of $K$ are i.i.d. As the width of the tensor network goes to infinity, the number of all the possible implementation of the indices sequence goes to infinity. By the C.L.T, $\psi$ converges to a GP.  

Based on the work of Theorem 2 proof, to prove Theorem 3, we just need to show that all the components of the neural kernel are independent. It is obvious that all component of the output of the neural kernel $a_{i}(\mathbf{w};\mathbf{x})$ are independent. By reshaping the output of the neural kernel $f^{l}$, we get the kernel $\Phi(\mathbf{x})$ and all the components are independent. Then by setting the width of the MPS to be infinite, using C.L.T and Theorem 3, the response of the neural kernel mps $\psi$ will converge to GP.  

By taking the infinite width limit of the hidden layer MPS, the response of the MPS $\psi^{l}$ are all i.i.d. which follow the normal distribution by Theorem 2, then it is easy to verify that $\sum_{l}w_{l}a(\psi^{l})$ follows gaussian distribution by C.L.T. and then $b+\sum_{l}w_{l}a(\psi^{l})$ is normal random variable. So $f$ converge to the GP as $l$ goes to infinity.