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: here we specify the shape to which to tensorize the weights, and use rank=0.5, which means automatically determine the rank so that the factorization uses approximately half the number of parameters.
fact_linear = tltorch.FactorizedLinear.from_linear(linear, auto_tensorize=False, in_tensorized_features=(4, 4), out_tensorized_features=(2, 5), rank=0.5)
The tensorized weights will have the following shape:
torch.Size([4, 4, 2, 5])
Note that you can also let TensorLy-Torch automatically determine the tensorization shape. In this case we just instruct it to
out_tensorized_features to have length 2:
fact_linear = tltorch.FactorizedLinear.from_linear(linear, auto_tensorize=True, n_tensorized_modes=2, rank=0.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)
Finally, during the forward pass, you can reconstruct the full weights (
implementation='reconstructed') and perform a regular linear layer forward pass.
ALternatively, you can let TensorLy-Torch automatically direction contract the input tensor with the factors of the decomposition (
which can be faster, particularly if you have a very small rank, e.g. very small factorization factors.