Tensorized Linear Layers

Linear layers are parametrized by matrices. However, it is possible to tensorize them, i.e. reshape them into higher-order tensors in order to compress them.

You can do this easily in TensorLy-Torch:

import tltorch
import torch

Let’s create a batch of 4 data points of size 16 each:

data = torch.randn((4, 16), dtype=torch.float32)

Now, imagine you already have a linear layer:

linear = torch.nn.Linear(in_features=16, 10)

You can easily compress it into a tensorized linear layer:

fact_linear = tltorch.FactorizedLinear.from_linear(linear, (4, 4), (2, 5), rank=0.5)
torch.Size([4, 4, 2, 5])

You can also create tensorized layers from scratch:

fact_linear = tltorch.FactorizedLinear(in_tensorized_features=(4, 4),
                                       out_tensorized_features=(2, 5),
                                       factorization='tucker', rank=0.5)