# Tensorization of neural network layers

## Introduction

First proposed by Novikov et al. 2015, tensorization of neural network aims to represent large weight matrices in neural networks with product of smaller, high dimensional tensors, which largely reduces the computational cost. For more details and applications please refer to my talk. In this post, we see how easy it has become to implement such tensorized neural networks.

## 1. Einstein summation notation.

Einstein summation notation can express tensor contraction in an extremely elegant fashion. There are only two conventions to follow:

- The mode (“dimensionality”) of a tensor is denoted using indices, which is in contrast to the usual matrix notation, where one denotes an entry in a tensor with the indices.
- The contraction is carried out over all modes whose indices appear more than once. Note: the order in which the contraction of multiple modes take place is not denoted explicitly.

One could express a large amount of tensor operations using Einstein summation notation in numpy, torch, tensorflow and other libraries.

Here are some basic usage:

```
import numpy as np
import torch
np.random.seed(123)
torch.manual_seed(123)
```

#### 1.1. Summation over all elements

```
a = torch.randn(5)
torch.einsum('n->', a)
# equivalent to a.sum()
```

#### 1.2. Summation over rows or columns

```
A = torch.randn(5, 3)
torch.einsum('nm->m', A)
# equivalent to A.sum(axis=0)
```

#### 1.3. Inner product of vectors

```
a = torch.randn(5)
b = torch.randn(5)
torch.einsum('n,n->', a, b)
# equivalent to torch.inner(a, b)
```

#### 1.4. Outer product of vectors

```
a = torch.randn(5)
b = torch.randn(3)
torch.einsum('n,m->nm', a, b)
# equivalent to torch.outer(a, b)
```

#### 1.5. Matrix vector product

```
A = torch.randn(3, 5)
b = torch.randn(5)
torch.einsum('mn,n->m', A, b)
# equivalent to torch.matmul(A, b)
```

#### 1.6. Matrix matrix product

```
A = torch.randn(3, 5)
B = torch.randn(5, 4)
torch.einsum('mn,nk->mk', A, B)
# equivalent to torch.matmul(A, B)
```

#### 1.7. Tensor contraction

```
A = torch.randn(5, 10, 7)
B = torch.randn(5, 7, 2)
torch.einsum('mnk,mkr->nr', A, B)
```

which is equivalent to

```
C = torch.zeros(10, 2)
for m in range(5):
for n in range(10):
for k in range(7):
for r in range(2):
C[n, r] += A[m, n, k] * B[m, k, r]
```

or

```
C = torch.zeros(10, 2)
for n in range(10):
for r in range(2):
C[n, r] = (A[:, n, :] * B[:, :, r]).sum()
```

## 2. Tensorizing neural network layers

We create a simple regression task as an isolated layer.

```
import numpy as np
import torch
from sklearn.datasets import make_regression
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_squared_error, r2_score
np.random.seed(123)
torch.manual_seed(123)
```

### 2.1. Differentiability

Obviously, tensor contraction are differentiable linear operations. The `einsum`

API is thus also differentiable.

```
W = torch.randn(3, 5, requires_grad=True)
x = torch.randn(5)
y = torch.randn(3)
y_hat = torch.einsum('NM,M->N', W, x)
loss = ((y - y_hat)**2).sum()
loss.backward()
W.grad
# tensor([[ 2.3307, 2.4076, 4.2191, -3.5718, -4.3023],
# [ 3.2671, 3.3749, 5.9143, -5.0069, -6.0309],
# [-4.2287, -4.3682, -7.6550, 6.4805, 7.8060]])
```

As we know, the 1st order derivative of such a layer is $-2 \cdot (y - W x )x^T$. One could check with

```
-torch.outer(y-(torch.matmul(W, x)), x)*2
# tensor([[ 2.3307, 2.4076, 4.2191, -3.5718, -4.3023],
# [ 3.2671, 3.3749, 5.9143, -5.0069, -6.0309],
# [-4.2287, -4.3682, -7.6550, 6.4805, 7.8060]], grad_fn=<MulBackward0>)
```

### 2.2. A tensor-train layer

#### 2.2.1 We create a toy dataset and verify the existence of a solution:

```
n_samples = 256
n_informative = 128
n_input = 4096
n_output = 16
X, y, coef = make_regression(n_samples=n_samples, n_features=n_input, n_informative=n_informative, n_targets=n_output, coef=True)
X = (X - X.mean(0)) / X.std(0)
y = (y - y.mean(0)) / y.std(0)
lr = LinearRegression()
lr.fit(X, y)
lr_y_hat = lr.predict(X)
mean_squared_error(y.reshape(-1), lr_y_hat.reshape(-1)), r2_score(y.reshape(-1), lr_y_hat.reshape(-1))
```

#### 2.2.2. We decompose the input and output dimensions as

$$ 4096 = 8 \times 8 \times 8 \times 8, \ 16 = 2 \times 2 \times 2 \times 2 $$ and implement

```
m = 8
n = 2
X = X.reshape(n_samples, m, m, m, m)
y = y.reshape(n_samples, n, n, n, n)
X = torch.tensor(X, dtype=torch.float)
y = torch.tensor(y, dtype=torch.float)
```

#### 2.2.3. Define the rank of the decomposition:

```
r = 8
W1_init = np.random.randn(m, n, r)
W1_init /= (W1_init**2).sum()**0.5
W2_init = np.random.randn(m, n, r, r)
W2_init /= (W2_init**2).sum()**0.5
W3_init = np.random.randn(m, n, r, r)
W3_init /= (W3_init**2).sum()**0.5
W4_init = np.random.randn(m, n, r)
W4_init /= (W4_init**2).sum()**0.5
W1 = torch.tensor(W1_init, requires_grad=True, dtype=torch.float)
W2 = torch.tensor(W2_init, requires_grad=True, dtype=torch.float)
W3 = torch.tensor(W3_init, requires_grad=True, dtype=torch.float)
W4 = torch.tensor(W4_init, requires_grad=True, dtype=torch.float)
```

#### 2.2.4. Perform the gradient descent:

```
opt = torch.optim.Adam([W1, W2, W3, W4], lr=5e-2)
loss = torch.nn.MSELoss()
n_epochs = 10000
for i in range(n_epochs):
y_hat = torch.einsum('aei, bfij, cgjk, dhk, nabcd -> nefgh', W1, W2, W3, W4, X)
error = loss(y_hat.reshape(n_samples, -1), y.reshape(n_samples, -1))
error.backward()
opt.step()
opt.zero_grad()
```

Note that we only need a single line to define a tensor-train layer and the `backward()`

and `step()`

are boilerplate from torch.

Of course we could have implemented the update from scratch, such as

```
lr = 5e-2
for i in range(10000):
y_hat = torch.einsum('aei, bfij, cgjk, dhk, nabcd -> nefgh', W1, W2, W3, W4, X)
error = ((y-y_hat)**2).mean()
error.backward()
with torch.no_grad():
W1 -= lr*W1.grad
W2 -= lr*W2.grad
W3 -= lr*W3.grad
W4 -= lr*W4.grad
W1.grad.zero_()
W2.grad.zero_()
W3.grad.zero_()
W4.grad.zero_()
```

But at least for this example, `adam`

performs way better than plain SGD.

#### 2.2.5. Evaluation

```
y_pred = torch.einsum('aei, bfij, cgjk, dhk, nabcd -> nefgh', W1, W2, W3, W4, X).detach().numpy()
print(mean_squared_error(y.numpy().reshape(-1), y_pred.reshape(-1)))
print(r2_score(y.numpy().reshape(-1), y_pred.reshape(-1)))
# 0.055088233
# 0.9449117677629324
```

It cannot fit the data as perfectly as the full model did. But keep in mind that we only use 3.5% of the parameters:

```
tt_size = np.prod(W1_init.shape) + np.prod(W2_init.shape) + np.prod(W3_init.shape) + np.prod(W4_init.shape)
full_size = n_input * n_output
tt_size / full_size
# 0.03515625
```