Speeding up neural networks using TensorNetwork in Keras

Posted by Marina Munkhoeva, PhD student at Skolkovo Institute of Science and Technology and AI Resident at Alphabet's X, Chase Roberts, Research Engineer at Alphabet's X, and Stefan Leichenauer, Research Scientist at Alphabet's X
, and how it can be used to supercharge a feed-forward neural network in TensorFlow. TensorNetwork is an open-source library ), and here we’re going to focus on answering the second. The basic idea is called “tensorizing” a neural network and has its roots in a .

TN Layers

The basic idea of a TN-enhanced neural network is to replace one or more of the layers of the network with a TN layer. You can think of the TN layer as a compressed version of the original layer. The best results can be expected when you start out with a dense, highly-connected layer that has a lot of potential for compression. For example, consider the following:
This network has a two-dimensional input, one-dimensional output, and contains two hidden layers with 1024 neurons each. The total number of parameters (including weights and biases) is (2+1)*1024 + (1024+1)*1024 + (1024 +1)*1 = 1,053,697. The bulk of these parameters, (1024+1)*1024 = 1,049,600 of them, are associated with the second hidden layer. We’re going to replace those with a TN layer that is less than 1% of the original size.

The 1,049,600 parameters of the hidden layer consist of 1024*1024 = 1,048,576 weights arranged in a 1024-by-1024 matrix, plus 1024 biases. Our savings are going to come from replacing the 1024-by-1024 matrix of weights by a tensor network. To preserve compatibility with the rest of the network, we don’t change the fact that there are 1024 inputs and 1024 outputs. That way we can treat the TN layer as a drop-in replacement for the original layer.

Ordinarily, we think of the 1024 inputs to the layer as a simple array X of shape (1024,). These inputs are multiplied by the weight matrix W to produce a new vector Y. The next step is to apply a nonlinearity to Y, but we won’t focus on that because the TN modifications only affect the linear part of the layer. In tensor network notation, we would write the computation of Y as follows:

Ordinarily, the weights W of a neural network layer act on the inputs X as a matrix multiplication, producing the output Y = WX. Here we illustrate that matrix multiplication diagrammatically.

At this point, we can discuss tensorizing the layer. First, we reshape the input array into something like (32,32) instead of (1024,). Pictorially, the reshape looks like

Reshaping the input X from a vector of shape (1024,) to an array of shape (32,32) is the first step of the tensorization process. In principle any reshape is allowed, we just chose this particularly simple one for our example.

Now, instead of applying the (1024,1024) weight matrix W, we apply a tensor network operation consisting of two cores:

Diagrammatic representation of the tensorized weight multiplication. The output Y can be reshaped into a vector if needed.

We need to make a choice about the dimensionality of the leg connecting the two cores (often called the bond dimension). For simplicity, we’ll take it to be two-dimensional. Then each core has two thirty two-dimensional legs along with the single two-dimensional leg, meaning that the shape is (32,32,2). The bond dimension controls the number of parameters in the model, and by appropriately choosing the bond dimension we achieve a good parameter reduction rate with little or no performance drop in many cases.

The result of using a TN layer is that we’ve replaced the 1,048,576 weights of the fully-connected weight matrix with the 2*(32*32*2) = 4,096 parameters of the tensor network. That’s a tremendous reduction! Even after accounting for the other layers, the total model size is down to 9,217 parameters, compared to the original 1,053,697.

For this simple example, it would not be too hard to translate the two-core tensor network into an einsum expression. However, for a tensor network with more cores and more complex connectivity, debugging or extending the einsum can be quite difficult. Thus, we instead use the TN library for a more object oriented way of constructing the network. In the code examples below we’ll stick to the simple two-core example to make things easy to follow, but at the end we’ll come back to discuss other possibilities.

Code for a TN Layer

Here is some sample code (also available for tensor networks in embedding layers, for convolutional layers.

Tensorization of neural networks is still in its infancy. There are still many unanswered questions, and many experiments to try. All of the tensor networks considered in this post are of the (see also and was designed to handle any possible tensor network, and specifically to do so in the context of machine learning as we have presented here. We look forward to seeing all of the new and exciting results that you can cook up using it.