Accelerated Learning of High Dimensional Functions with a Tensor-Featured Training Network

The CPD is a higher-order extension of the singular value decomposition and represents a function of $d$ variables (expressed as an order-$d$ tensor) as a sum of outer products of single variables, i.e. given $\mathcal{T}\in\mathbb{R}^{I_{1},I_{2},\dots,I_{d}}$

$\displaystyle t_{a,b,\dots,n}\overset{\mathrm{CPD}}{=}\sum_{i=1}^{R_{\mathrm{CP}}}\lambda_{i}{\bf a}_{i}\circ{\bf b}_{i}\circ\dots\circ{\bf n}_{i}$

(1)

where $\circ$ defines the vector outer product, ${\bf a}_{i}\in\mathbb{R}^{I_{1}},{\bf b}_{i}\in\mathbb{R}^{I_{1}},\dots,{\bf n}_{i}\in\mathbb{R}^{I_{d}}$, $\lambda_{i}$ is a scaling constant that allows the vectors $[a_{i},\dots n_{i}]$ be unit normalized, and $R_{\mathrm{CP}}$ is the CP rank, i.e. the smallest set of vectors such that Eq. 1 satisfies the equality. All the vectors which span the same mode of the tensor can be group into a factor matrix, i.e. ${\bf A}=[{\bf a}_{1},\dots,{\bf a}_{R_{\mathrm{CP}}}]\in\mathbb{R}^{I_{a}\times R_{\mathrm{CP}}}$. Therefore, the CPD reduces the storage complexity of a $d$-dimensional function from $N^{d}$ where $N\approx I_{1}\approx I_{2},...$ to $NDR_{\mathrm{CP}}$. Although finding the exact CP rank is formally an NP-hard problem, it is straightforward and easy to discover an approximate CPD of $\mathcal{T}$ given a desired rank $R$. By construction, the CPD defines a set of 1-dimensional basis functions such that a function of $d$ inputs as a separable set of products of functions of single inputs.

The most standard way to optimize the rank-$R$ CPD is via the alternating least squares (ALS) optimization.[13, 2] The ALS method splits the optimization of the CPD into $d$-linear least square problems. Unfortunately, the CPD-ALS optimization is non-convex. However, based on our experience and results, we find that the non-convex nature of the CPD-ALS has minimal effect on discovering accurate decomposition’s of tensor for most cases in physics and chemistry.[16] As we will demonstrate in Section IV, this is also true for this work. A more serious issue with the CPD-ALS is that the optimization suffers from the curse of dimensionality, i.e. the computational cost of the method scales exponentially with the dimension of the function being decomposed.

In general, the task of the CPD-ALS is to iteratively solve a collection of extremely overdetermined least squares problems. Because these problems are over-determined, a direct result of the matricization of higher order tensors along a single mode of the tensor, the CPD-ALS a perfect candidate for randomized linear algebra. There are two main ways to apply randomization to the CPD-ALS, via sampled least squares[1, 14] or via sketched and sampled least squares.[26, 21] In this work, we will only consider the sampled least squares method because of its low-cost and structure preserving properties. In the literature, it has been demonstrated that the number of samples required for the CPD randomized ALS (CPD-RALS) is proportional to the rank of the CPD and has no direct dependence on the dimensionality of the function being decomposed.[1, 4] Therefore, in this work, we will assume that the CP rank of every tensor is practically small compared to the number of interpolation points along a single dimension and does not grow (significantly) with increasing tensor dimension. The RALS procedure therefore reduces the computational complexity of the CPD from exponential in function dimension to polynomial and independent of function dimension. Using the CPD-RALS, the most time intensive component of the algorithm is the cost of sampling the target function which grows as $dSN$ where $d$ is the dimension of the function, $S$ is the number of samples per mode and $N$ is the number of interpolation points along a single dimension. To minimize the overall cost of this step, we will assume that single set of samples gathered from a uniform distribution is sufficient for optimizing the CPD-RALS.

In this work we will not explicitly derive the sampled CPD-RALS method. We utilize a modified version of the canonical version developed by Battaglino et. al. where samples are determined via the leverage score distribution of each factor matrix. Our method deviates from the canonical method in the following way. We found that the magnitude of elements in the sampled Khatri-Rao product (KRP) associated CPD-RALS decay’s rapidly with dimensionality of the tensor. To accommodate this issue, we introduce an intermediate normalization where rows of the KRP are normalized to 1 such that

	$\displaystyle Ax$	$\displaystyle=B$		(2)
	$\displaystyle(l\hat{A})(xl^{-1})$	$\displaystyle=B$
	$\displaystyle\hat{A}\bar{x}$	$\displaystyle=B.$

The intermediate normalization tensor $l$ can, without loss of generality, be factored into the column-wise normalization of $\bar{x}$ and, therefore, absorbed into $\lambda$ in Eq. 1. We have found that this process, in general, stabilizes the LS solver in the CPD-RALS procedure.

Recently, it has been shown that there is a connection between popular tensor decomposition strategies, kernel machines[22] and Gaussian process(GP).[24] The GP is a mechanism for modeling functions and allows the user to incorporate prior knowledge of the function (i.e. random inputs on the domain and associated noisy observations). GP are fully described by two components, a mean $\mu\in\mathbb{R}$ which is usually chosen to be zero and a covariance or kernel matrix $k\in\mathbb{R}^{\Omega_{i}\times\Omega_{i}}$ where $\Omega_{i}$ is the extent of the domain in the $i$th dimension. Given a set of points and their associated observations, one can determine random functions on the domain by sampling the multivariate gaussian associated with this mean and kernel matrix. The kernel matrix is commonly approximated using a finite number of basis functions[17, 19]

$\displaystyle k(x,x^{\prime})=\varphi(x)^{T}\Lambda\varphi(x^{\prime})$

(3)

where $\varphi(x)\in\mathbb{R}^{\Omega_{i}\times M}$ are basis functions and $\Lambda\in\mathbb{R}^{M\times M}$ are basis function weights. One can recognize that each CPD factor matrix are explicitly constructed as a collection of $R$, 1-dimensional basis functions for a on set of observed points on a predefined collocation grid.

The CDP-ALS iteratively constructs each 1-dimensional basis matrix (CPD factor matrix) in the following way. For convenience we limit this explanation to a three dimensional function $\mathcal{T}\in\mathbb{R}^{I_{a}\times I_{b}\times I_{c}}$ but this idea can be extended to arbitrary dimension without loss of generality. First, the loss function is defined for the $I_{a}$th mode

$\displaystyle f(A)=\|T_{a}-A^{*}[B\otimes C]^{T}\|^{2}_{2}$

(4)

where $T_{a}$ is the so called matriciation of the third order tensor $\mathcal{T}$ that permutes and reshapes the tensor to be $T_{a}\in\mathbb{R}^{I_{a}\times I_{b}I_{c}}$, a short and long matrix. Also $\otimes$ defines the Khatri-Rao tensor product (KRP) which is the column-wise Kronecker product of two tensors. The KRP defines the following product $W_{r}=b_{r}\circ c_{r}$ where is a matrix $W_{r}\in\mathbb{R}^{I_{b}\times I_{c}}$ for each value of $r\in\mathbb{R}^{R}$. By definition $B$ and $C$ are also collections of 1D basis which were either chosen at random or drawn randomly a different GP. Next, one solves for $A^{*}$ by taking the gradient of the loss function and solving the following least squares problem

$\displaystyle T_{a}(B\otimes C)=A^{*}(B\otimes C)^{T}(B\otimes C).$

(5)

This equation consists of two parts: $T_{a}(B\otimes C)\in\mathbb{R}^{I_{a}\times R}$ which down-folds the high dimensional observations of the function into a single dimensional set of mean values for each basis basis function by taking the inner product of the higher order tensor and all other degrees of freedom. The second part is the construction of the squared KRP, $(B\otimes C)^{T}(B\otimes C)$. Using Tensor Algebra, this expression can be rewritten as $(B^{T}B)*(C^{T}C)$ where $*$ defines the Hadamard or element-wise product of two matrices. Written in this way one can see that $B^{T}B\in\mathbb{R}^{R\times R}$ and $C^{T}C\in\mathbb{R}^{R\times R}$ are both approximations to 1-dimensional kernel matrices. Therefore the squared KRP can be viewed as a product kernel[19] of many 1-dimensional kernel matrices. As a result, the above solution to the CPD-ALS can be viewed as a method that leverages GP to iteratively draw and refine 1-dimensional basis functions for high dimensional functions. By leveraging the basis function approximation and kernel product structure, the CPD-ALS is able to draw GP basis functions with information across the entire domain of the function with a relatively low cost. Unfortunately, one can see that the construction of the left hand side of Eq. 5 suffers from the curse of dimensionality.

In an effort to minimize the cost of the CPD-ALS, the CPD-RALS modifies the above process. The gradient of the CPD-RALS loss function for the $I_{a}$th mode can be written as

$\displaystyle(T_{a}S_{a})[S^{T}_{a}(B\otimes C)]=A^{*}[S^{T}_{a}(B\otimes C)]^{T}[S^{T}_{a}(B\otimes C)]$

(6)

where $S_{a}\in\mathbb{R}^{I_{b}I_{c}\times I_{s}}$ is a sparse, column-wise selection matrix chosen specifically for the $I_{a}$th subproblem. By constructing the randomized solver in this way we may significantly reduce the cost of inner-product down-folding process from $I_{a}I_{b}I_{c}R$ to $I_{a}I_{s}R$, provided that $I_{s}\propto R$ and $R\ll I_{b}I_{c}$.[1, 4] However, this sampling process changes the form of our kernel matrix The sampling matrix first down-folds the high-dimensional space orthogonal to the $I_{a}$ into a relatively small set of observations. In this work, we choose these observations by first constructing a leverage score distribution of the kernel basis function. By taking advantage of the KRP structure of the kernel basis functions, we may randomly sample the leverage score distribution of each basis matrix independently. The row-wise leverage score distribution of each factor matrix measures how far away observations on each quadrature point are from all the other observations and thus the leverage scores are correlated with the degree of misprediction of a least squares model. Therefore, by importance sampling the KRP using the leverage score distribution, we are attempting to determine a compact basis of observations which contribute most strongly to the basis-approximated kernel matrix. After choosing this compact set of basis points, we now have a single set of kernel basis functions. In this process you can recognize that our normalization process in Eq. 2 is effectively a modification of our basis function weights, $\Lambda$.

First we investigate the effect of using a pretrained DNN and as a tensor feature. For this section we attempt to learn a function which is the solution to the the nonlinear elliptic equation inside a bounded domain,

$\displaystyle u(x)=\sin(\frac{\pi}{2}(1-|x|_{2})^{2.5}),$

(7)

using the DNN structure described previously. We consider the problem in three dimensions $d=5,10,$ and $15$. In this experiment, we first use canonical optimization strategies to train a DNN over a 3000 training epoch. Throughout this work we will referred to this canonical training as Conventional Training. After this training we initialize two new DNN each with 1 additional feature added to their input layers. For one DNN, this new feature is the output of the first DNN evaluated at the input training coordinates. We call this NN Featured training. In the other DNN, we provide the approximated output of the first DNN via CPD interpolated at the input training coordinates. As explained early, for off-grid points, we use a nearest neighbor interpolation. We denote this second DNN as a Tensor-Featured trained model.

In Fig. 1 we compare accuracy of the three training strategies (conventional, NN featured and Tensor-Featured) on the set of validation points over 3000 gradient updates. For reference, we also show the error of the CPD interpolation of the conventionally trained NN on the same set of validation points. One can see, across the board, decomposing the conventionally optimized DNN into a TN can significantly improve the accuracy of learned function. As pointed out previously, we believe this is because the process of optimizing the tensor decomposition using non-local information which effectively smooths the DNN function over the domain.

With respect to Tensor-Featured training, in Fig. 1 one can also see that using the DNN directly as a feature does not practically provide any improvement over the conventional training. On the other hand with the CPD interpolation as a feature, we find an improvement in accuracy over the conventional training. Unfortunately, although there is an improvement in accuracy with the Tensor-Featured training in the validation set, the improvement is not dramatic and, in all cases, the CPD approximation alone outperforms the Tensor-Featured DNN. We believe this phenomena is not a failure in the Tensor-Featured learning but, instead, associated with the accuracy of the conventionally trained model relative to the ground truth.

To clarify this statement, consider the heatmaps in Fig. 2 which illustrates true and predicted values of Eq. 7 on a 2000 point interpolation grids with all but 2 modes of the tensor fixed to the 500th interpolation point.

One can see that, though Fig. 1 implies that the CPD approximation is a better representation of the true function, none of these models look significantly more accurate in this representation. With this example, we demonstrate that because the reference function for the CPD (the conventionally trained network) is not representation of the ground truth, its addition to the feature list in the Tensor-Featured training doesn’t significantly improve the DNN training. Therefore, we hypothesize that, if the conventional training is more similar of the ground truth, then the Tensor-Featured training procedure will be significantly accelerated.

In Fig. 3, we improve the accuracy of the conventionally trained DNN by increasing the training sample batch size from 1,000 to 10,000. With a qualitatively more accuracy reference, the CPD approximation generates a function that captures more character of the true function. Finally, equipped with this CPD, the Tensor-Featured training is greatly improved and the mean squared error in the validation set falls dramatically. In general, these results demonstrate that introduced features provide context for DNN training. This 2-step optimization process of training DNN with a fixed feature and decomposing partially trained DNN into CPD can be repeated iteratively in an attempt refine the Tensor-Featured DNN.

For example, we train the 15 dimensional example from above in a 3 step optimization procedure. First a DNN is trained conventionally using 10,000 samples and 3000 training epoch. After the training session, the DNN is discretized and decomposed into a CPD. This CPD is passed as a feature into a new DNN which trained for 3000 iterations using the same 10,000 training points with the value of the CPD interpolation at each training point added to the feature list. After this second training session, the resulting DNN is again decomposed into a CPD and this CPD is passed to the training of a third DNN. In Fig. 4 we show the same 2D heatmaps for the CPD of the Tensor Featured trained DNN (i.e. Fig. 3(d)) and the DNN constructed from the second round of Tensor-Featured training. While it is hard to tell if this network is an improvement on the first iteration, the data becomes clear when considering the relative L2 error and absolute mean squared error on the 2D slices in Table 1. From these results one can see that this multi-step training process, iteratively reduces the error in the trained DNNs.

Because the tensor features provide context to the DNN optimization, we require that the initial training be relatively close (in some measure) to the global extrema. If the original network is not within epsilon of the solution, the Tensor-Featured training doesn’t help dramatically, if at all. In the example above, we were able to improve the canonical training solution by increasing the number of batched training point samples. In general, it may not be possible to improve the accuracy of conventional training so easily, for example there may be a limitation on the total number of available samples. We are therefore interested in developing a strategy that introduces tensor features into the DNN training process from the start. So far, we have only considered decomposing partially trained DNN however practically, we are able to introduce any one or many features into the Tensor-Featured training process. In the next section we will explore the idea of adding tensor features beyond TN decomposed DNN.

One of the biggest challenges of the Tensor-Featured training procedure as presented in the previous section is need for features to provide meaningful context to a DNN’s training. If a target function is unknown and the convergence of conventional DNN training is slow then the Tensor-Featured training as presented thus far will not, in general, improve DNN optimization. Therefore, to improve optimization, we consider introducing alternative and additional contextual features.

In this section we define two types of features. The first type of features are elementary functions which can be evaluated quickly for points in the domain of the unknown function. Practically these features could be discretized and decomposed using TN methods but, because the functions can be evaluated quickly, we choose to simply access elements of the functions in a matrix-free manner. The first class of features we will call Rank-1 features. The second type of feature are more complex functions which may be either slow to evaluate for points in the functions domain or otherwise limit the Tensor-Featured optimization. For this second class of features, we will leverage tensor decomposition strategies to address these issues. An example of this type of feature are the partially trained DNN from the previous section. We call this second class of features TN features.

In this section we will attempt to learn a more difficult function, the solution to the initial boundary value problem of the hyperbolic (wave) equation

$\displaystyle u(t,x)=(\exp(t^{2})-1)\sin(\frac{\pi}{2}(1-|x|)^{2.5})$

(8)

with $d=15$. In this study, we only consider the following rank-1 features

• •

  $y=\sum_{i}x_{i}$

• •

  $y=\|x\|_{2}$

• •

  $y=\sin(\|x\|_{2})$

• •

  $y=\cos(\|x\|_{2})$

• •

  $y=\sum_{i}\sin(x_{i})$

• •

  $y=\sum_{i}\cos(x_{i}))$

• •

  $y=\exp(\|x\|_{2})$

which map $\mathbb{R}^{n}\rightarrow\mathbb{R}$. Our decision to use these rank-1 features is relatively arbitrary. Any function is acceptable as long as it map $\mathbb{R}^{k}\rightarrow\mathbb{R}^{m}$ where $k\leq n$ and $\mathbb{R}^{m}$ is the dimensionality of function being learned. In follow-up studies we will investigate methods to automatically choose a good set of Rank-1 features. Furthermore, it should be noted that we will only consider the CPD of a previously trained DNN as a TN feature though in future work other TN features will be investigated.

In Fig. 5, we consider the effect of using different features in the Tensor-Featured training. In Fig. 5(a) we add just a single to the training and find that the feature does influence the training efficiency. One can see that adding a feature that does not align with the problem does not significantly improve the DNN optimization performance. But adding a feature that does align with the problem improves the training.

In Fig. 5(b), we consider the effect of adding multiple features in parallel to the Tensor-Featured training. In this figure, we incrementally introduce the features from the rank-1 feature list above in the order of the list. One can see that introducing multiple features seems to systematically improve the accuracy over the conventional training and the best single feature training in Fig. 5(a). In our tests we found that most of the functions in the rank-1 feature list do not significantly improve the DNN training but together they outperform any single feature. However, we note that this process can fail when a feature is added that puts the DNN close to a local extrema. For example in Fig. 5(c), we add the conventionally trained DNN into feature list and find that the optimization quickly converge back to the Conventional Training solution. Though one can see that replacing the conventionally trained DNN with its CPD approximation can resolve this issue.

Using Tensor-Featured training with rank-1 features, we can now improve the efficiency of training a DNN with no prior knowledge of the function. Furthermore, the rank-1 features do not depend on a pretrained DNN, like the CPD approximation did in the previous section. In Fig. 6, we compare the accuracy of the conventional training to the Tensor-Featured training with only rank-1 features by plotting the true and predicted function values of points in the validation set. Note that the points are ordered via the value of the true function from most negative to most positive. One can see that the Tensor-Featured training using only rank-1 features more accurately predicts nearly all points in the validation set, compared to conventional training. Furthermore, the Tensor-Featured DNN training converges more quickly when the training batch size is increased.

After training the DNN with rank-1 features, we can discretize and decompose the trained DNN as a TN. We compare the accuracy of the Tensor-Featured trained DNN to those interpolated by the CPD of this DNN in Fig. 7(a). From this figure, one can see that the decomposition process has little impact on the accuracy of predicting values in the validation set. Furthermore, by decomposing the DNN we, effectively, cannibalize all of the features from the preliminary DNN optimization into a single function. We demonstrate this in Fig. 7(b) where we attempt to improve on the first Tensor-Featured training. In the curve marked ‘NN Feature‘ we use exactly the DNN trained via the Rank-1 Tensor Featured Training into the optimization of a new DNN. To do this optimization, we must evaluate the original DNN on the batch of training and validation points. This evaluation may be relatively slow, especially if there were many input features in this DNN or the DNN is very deep. The curve marked ‘TN Feature‘ uses only CPD interpolation of the Rank-1 Featured Training. Similar to the previous results, the CPD approximation of the DNN is more accurate than the DNN. For the curve marked ‘Rank-1 and TN Features‘ we combine the CPD with the original Rank-1 features and find that the new Tensor-Featured Training can become more accurate than the original Tensor-Featured Training. Because the original features were well captured by the CPD approximation, we are also free to modify the feature list in the DNN retraining. For example for the curve marked ‘Modified Rank-1 and TN Features‘ we remove features 1, 2, 5 and 6 from the input and find a result that outperforms the CPD approximation alone. In these optimization curves we often see oscillation in the errors and will investigate this phenomenon in future studies. This two step process of training DNN and updating features can be repeated iteratively until some convergence is reached.

Finally in Fig. 8, we the accuracy of training a DNN to learn Eq. 8 in 40 dimensions with either conventional training or Tensor-Featured training using only rank-1 features. One can see that the Tensor-Featured training more accurately learns the objective function in just 3000 epoch. Though comparing this figure to Fig. 6(b), one can see that increasing the dimensionality of the problem slows the convergence of both training procedures. In Table 2, we show the cost comparing conventional training to Tensor-Featured with only rank-1 features and see that the Tensor-Featured training is only marginally more expensive in both optimization time and storage complexity.