Skip to content

Abhijoy Sarkar

Notes on AI, agents, and building things that work.

I co-authored a tensor-network compression paper. Here is what the tensor networks actually do.

I am the third author on CompactifAI, which compresses large language models with quantum-inspired tensor networks. The paper reports LLaMA-7B at 30% of its parameters and 7% of its memory footprint for a 2–3% accuracy drop, and the number people quote back at me is the 93% memory figure.

What almost nobody asks is which part of the pipeline does the work. So I implemented the decomposition from scratch, ran it on GPT-2, and measured. The answer surprised me, and it is not the tensor networks.

The decomposition is about sixty lines

A linear layer is a matrix W of shape (out, in). Factor both dimensions (out = o₁·o₂, in = i₁·i₂), reshape into a four-index tensor, interleave so that each input axis sits beside its matching output axis, and split with a truncated SVD. What you keep is two small cores instead of one big matrix.

W(o1o2),(i1i2)rGo1i1r(1)Gro2i2(2)

Two cores replace one matrix; r is the bond rank, and the sum over it is the contraction.

The interleave is the entire trick. Grouping (o₁,i₁) against (o₂,i₂) factorises the operator. Grouping outputs against inputs instead would just be a low-rank approximation of W, which is a different and weaker method and, as it turns out below, a stronger one in practice.

oir(o1i1+o2i2)

Dense on the left, factorised on the right. The saving exists only while r stays small.

For a GPT-2 MLP matrix, 3072×768, that is 2,359,296 parameters dense against 204,800 at bond rank 64. An 11.5× saving. At full rank the round trip is exact to 9.6e−07, so the implementation is right.

Then you run it on a real model and it falls over.

Trained weights are barely more compressible than noise

The premise underneath every low-rank compression method is that training concentrates structure into a subspace much smaller than the matrix: that the weights are redundant in a way you can exploit. That is testable. Take GPT-2’s first MLP matrix, and a Gaussian matrix of identical shape and standard deviation, and compress both at matched parameter budgets.

0.0 0.2 0.4 0.6 0.8 1.0 random noise, MPO GPT-2, MPO GPT-2, plain SVD 10% 25% 50% 75% parameters kept, as a share of the dense matrix relative error
Relative reconstruction error at matched parameter budgets. Lower is better. The gap between the trained model and random noise is the redundancy the method is supposed to exploit.

At half the parameters, GPT-2 under MPO scores 0.607. Random noise under MPO scores 0.616. That is a 1.5% advantage for a matrix that took thousands of GPU-hours to train over one drawn from torch.randn.

And plain truncated SVD beats the tensor-network split on the trained weights at every budget: 0.532 against 0.607 at half, 0.360 against 0.452 at three-quarters. The fancy method loses to the obvious one.

The spectrum says why. You need rank 501 of a possible 768 to capture 90% of the spectral energy in that matrix. Random noise needs 585. Both are close to full rank, because a 3072×768 weight matrix in a well-trained transformer is not low-rank in any sense that saves you parameters, and the MPO bond rank that would capture 90% costs more storage than the dense matrix it replaces.

What that does to the model

Replacing every MLP matrix in all twelve GPT-2 blocks, at a rank that saves 11.4% of total parameters:

baseline            perplexity   35.45   124,439,808 params
after MPO rank 552  perplexity  768.83   110,210,304 params  (88.6% of original)

An 11% parameter saving costs a 22× increase in perplexity. The model is destroyed. Held-out WikiText-2, non-overlapping 512-token windows.

So the healing is the method

Every practical pipeline of this kind retrains after factorising. The paper calls it healing, and having now watched the decomposition on its own, I understand why it is not an optional polish step. The factorisation is not supposed to preserve the function. It is supposed to hand gradient descent a smaller parameter space in which to re-find it.

Three hundred steps, training the MPO cores only and freezing everything else:

after MPO rank 552  perplexity  768.83
after healing       perplexity  601.04   (88s, 42.4M trainable, 300 steps)
recovered 22.9% of the gap back to baseline

Twenty-three percent, and the model is still seventeen times worse than where it started. Meanwhile the healing loss fell to 0.37: the cores were busily memorising the thirty-two thousand tokens I gave them while held-out perplexity barely moved.

That is the real shape of the thing. Healing is not a patch you apply after compression. It is a training run, with a training run’s appetite for data and compute, and my laptop-sized version of it is nowhere near enough. The paper retrains on a real corpus, and that is the difference between their result and mine.

A methodological note, because I got this wrong first

My first healing run reported perplexity 1.86 against a baseline of 40, and “105% of the gap recovered.” I had healed and evaluated on the same tokens. It was not recovery, it was memorisation, and it looked like a triumph.

A compression result measured on the text you retrained on is worthless, and it fails in the flattering direction, which is the direction you are least likely to interrogate. The fix is one line and the script now asserts the slices are disjoint.

What I would tell someone reading the paper

Read the abstract carefully. It says “a combination of CompactifAI with quantization” achieves the 93% memory reduction. Not tensor networks alone. I have seen that number quoted as a tensor-network result, including by people who then ask me why their own decomposition did nothing, and having now run the experiment I think the honest decomposition of that headline is:

  • The tensor-network split chooses the shape of the smaller parameter space. On its own, on these matrices, it is worse than plain SVD.
  • Quantisation does a large share of the memory reduction, and is orthogonal to any of this.
  • Healing does the accuracy recovery, and it is a real training run.
  • Layer sensitivity profiling decides where you can afford to be aggressive, which the paper does and my flat sweep does not.

None of that makes the paper wrong. It makes the folklore wrong: the version where tensor networks are a magic compressor you apply to a checkpoint. What I would say now, having implemented it, is that the decomposition is the cheap part and the retraining is the method, and if you are evaluating a compression technique the first question to ask is how much retraining is hiding inside the number.

Run it

experiments/compactifai/
  tensor_train.py     the MPO decomposition, ~60 lines, torch only
  spectra.py          trained vs random at matched budgets
  compress_gpt2.py    swap in the cores, measure perplexity
  heal.py             retrain the cores, disjoint eval

Everything above runs on a laptop in a few minutes on GPT-2. All of it is in the blog repo, including the raw JSON, so the numbers can be reproduced or contradicted.

Caveats worth stating plainly: this is GPT-2 at 124M parameters and the paper is LLaMA-7B, my healing budget is three hundred steps on thirty-two thousand tokens against their real retraining, and I only touched MLP weights. A negative result at this scale does not refute a positive result at theirs. It does show what the decomposition contributes on its own, which is the question I actually wanted answered.

Discover more from Abhijoy Sarkar

Subscribe now to keep reading and get access to the full archive.

Continue reading