Compressing Transformer Language Models via Matrix Product Operator Decomposition: A Case Study on PicoGPT

TL;DR

This paper explores MPO decomposition as a method to compress transformer language models, achieving significant size reduction while maintaining high accuracy, demonstrated on a small GPT-2 style model.

Contribution

It introduces MPO-based linear layers for transformers, showing effective compression and competitive performance on a character-level language model.

Findings

Up to 13x compression per transformer block at bond dimension 4.

Model with 16 bond dimension retains 97.7% accuracy with 191,872 parameters.

MPO compression outperforms low-rank and pruning methods in efficiency and effectiveness.

Abstract

Transformer-based language models achieve strong performance across NLP tasks, but their quadratic parameter scaling with hidden dimension makes deployment on resource-constrained hardware expensive. We study Matrix Product Operator (MPO) decomposition as a principled compression method for transformers. MPO factorises weight matrices into chains of low-rank cores, with approximation quality controlled by the bond dimension chi. We replace every nn.Linear layer in PicoGPT, a GPT-2-style character-level language model with about 1M parameters, with an MPOLinear module parameterised as an MPO chain. Cores are initialised either by TT-SVD from pretrained dense weights or from random initialisation, and trained using standard PyTorch autograd without a custom backward pass. We derive balanced factorisation schemes for the five distinct weight shapes in PicoGPT and evaluate bond dimensions chi in {4, 8, 16, 32} on Tiny Shakespeare. MPO compression achieves up to 13x compression per transformer block at chi = 4. At chi = 16, the model uses 191,872 parameters instead of 1,020,224 while retaining 97.7% of baseline token accuracy (51.6% vs 52.8%). Reconstruction error follows the expected trend and is lower for three-site than two-site factorisations at the same bond dimension. The chi = 8 model gives the best accuracy per parameter, exceeding the dense baseline by 2.7x on this metric. These results show that MPO parameterisation is a practical and theoretically grounded alternative to low-rank methods and unstructured pruning for transformer compression.