PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Younes Javanmard, Tanmoy Pandit, Masoud Mardani

arXiv:2603.28534v1cs.CLphysics.data-an

TL;DR

Transformer language models have quadratic parameter scaling with hidden dimension, making deployment on resource-constrained hardware expensive. This paper applies MPO decomposition to PicoGPT, replacing linear layers with MPO layers and training their cores using PyTorch autograd. At χ = 16, the model retains 97.7% of dense-baseline token accuracy with 18.8% of its parameters, though the current implementation reconstructs the dense matrix on each forward pass.

  • Problem

    Transformer language models’ parameter counts scale quadratically with hidden dimension, making deployment on resource-constrained hardware expensive.

    Source: §1 · p. 2

  • Method

    The paper replaces PicoGPT’s linear layers with MPO layers, using balanced factorisation schemes and standard PyTorch autograd to train the cores.

    Source: §1 · p. 3

  • Results

    At χ = 16, the MPO model retains 97.7% of the dense baseline’s token accuracy with only 18.8% of its parameters.

    Source: §7 · p. 10

  • Takeaways & Limitations

    MPO parameterisation achieves 5–13× compression per transformer block, and its PyTorch implementation uses standard parameters without custom backward code.

    Source: §7 · p. 10

    The current MPOLinear forward pass reconstructs the full dense matrix, demonstrating parameter compression rather than inference-time memory or FLOP reduction.

    Source: §6 · p. 10

Abstract

from arXiv · show

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.

1 Introduction

MPOs represent weight matrices as chains of low-rank cores, with bond dimension controlling the accuracy–compression trade-off. This paper applies the representation to PicoGPT and introduces balanced factorisations, TT-SVD initialization, and autograd-based training without a custom backward pass.

1 Introduction

  • MPOs factorise weight matrices into chains of low-rank cores, offering an alternative compression framework to pruning and low-rank factorisation.The bond dimension χ controls the accuracy–compression trade-off: increasing it systematically recovers the dense weight.Source: §1 · p. 2
  • The paper applies MPO compression to PicoGPT and introduces an autograd-compatible MPOLinear layer, balanced schemes for five weight shapes, and a benchmark across χ ∈ {4, 8, 16, 32}.The contributions also include analysis of reconstruction error, training dynamics, and the accuracy–compression Pareto frontier.Source: §1 · p. 3

2.1 Tensor Train and Matrix Product Operators

  • An MPO represents a weight matrix after reshaping its input and output dimensions into local site dimensions and expressing the result through cores.For balanced factorizations with bounded local dimensions and fixed bond dimension, parameter count grows linearly with the number of sites, unlike the dense count.Source: §2.1 · p. 3

2.2 TT-SVD Algorithm

  • TT-SVD converts a dense weight matrix into an MPO through a sequential SVD sweep, and its approximation error decreases monotonically as admissible bond dimensions increase.The approximation satisfies a standard Frobenius-norm bound based on singular values discarded at each unfolding step.Source: §2.2 · p. 4

2.3 Gradient Flow Through MPO Layers

  • PyTorch autograd computes MPO-core gradients through the contraction chain, using left and right environments around each local core without a custom backward pass.This gradient structure is familiar from single-site tensor-network optimization methods such as DMRG.Source: §2.3 · p. 4

3 Architecture and Implementation

The section describes PicoGPT’s GPT-2-style architecture and how its linear weights are factorised and implemented as MPO chains. It specifies balanced factorisation choices and an MPOLinear module that reconstructs weights through contractions.

3.1 PicoGPT Architecture

  • PicoGPT uses pre-norm residual blocks, sinusoidal positional encoding, causal multi-head self-attention, and a two-layer ReLU FFN, with six linear layers per block and a language-model head.The block layers are query, key, value, and output projections plus FFN up- and down-projections.Source: §3.1 · p. 5

3.2 MPO Factorisation Schemes

  • For each PicoGPT linear-layer shape, the factorisation choices aim for balanced physical dimensions and exact integer factorizations.The stated example uses [8,8,8] output and [4,4,8] input dimensions for W1 = (512,128), with L = 3.Source: §3.2 · pp. 5–6, Table 2 · p. 5
  • Reported parameter counts include all trainable parameters, including biases and uncompressed components, and are given for factorised matrices and the full network.The section distinguishes matrix-level counts in Table 2 from network-level counts in Tables 3 and 4.Source: §3.2 · pp. 5–6

3.3 The MPOLinear Module

  • MPOLinear stores L MPO cores as trainable tensors and contracts them to reconstruct the weight matrix before applying the linear transformation.Its forward pass uses sequential torch.tensordot contractions and F.linear; initialization targets reconstructed-weight variance on the same order as the dense baseline.Source: §3.3 · p. 6, Figure 1 · p. 6

4 Experimental Setup

The experiments evaluate dense and MPO models on character-level Tiny Shakespeare prediction under shared training and evaluation procedures. They test both random-initialized MPO training and TT-SVD-based compression followed by fine-tuning, while focusing reported results on train-from-scratch models.

4.1 Dataset and Tokenisation

  • Tiny Shakespeare provides approximately 1.1×10^6 characters, tokenized into 65 symbols and split 90/10 into training and validation sets.Source: §4.1 · p. 6

4.2 Training Protocol

  • Models train for 2,000 steps with AdamW, batch size 32, sequence length 256, gradient clipping, and a warmup-plus-cosine learning-rate schedule.Validation loss and token accuracy are computed every 100 steps and averaged over 20 random validation batches of size 8; reported curves and statistics are single-seed unless noted.Source: §4.2 · pp. 6–7

4.3 Compression Modes

  • The study tests MPO models trained from random initialization and dense models compressed by TT-SVD before 500-step fine-tuning.Although both workflows were implemented, the main results focus on train-from-scratch models; broader post-training compression and fine-tuning comparisons are left to future work.Source: §4.3 · p. 7

5 Results

The results show an accuracy–compression trade-off for MPO models: χ = 16 retains nearly all dense-model accuracy with substantially fewer parameters, while reconstruction quality and training behavior vary with bond dimension and factorisation.

5.1 Compression Statistics

  • Table 3 reports parameter counts and compression ratios against the dense baseline; reconstruction error is relative Frobenius error averaged across linear layers.Source: §5.1 · p. 7, Table 3 · p. 7

5.2 Per-Layer Reconstruction Error

  • Across χ, three-site MPOs achieve lower error per parameter than two-site ones; W1 has the lowest reconstruction error at every χ.The authors attribute W1’s advantage to distributing correlations across more local cores.Source: Figure 2 · p. 8, §5.2 · p. 8

5.3 Training Dynamics

  • Higher bond dimensions converge faster and to lower training and validation cross-entropy losses in the train-from-scratch setting.Source: Figure 3 · p. 9
  • 51.6% validation accuracy for MPO χ = 16 retains 97.7% of the dense model’s 52.8% while using 191872 versus 1020224 parameters (5.3× compression).Source: §5.3 · p. 8, Figure 4 · p. 9

5.4 Summary of Results

  • MPO χ = 8 scores highest on the informal per-parameter accuracy proxy; the paper treats the accuracy–compression Pareto frontier as primary.Source: §5.4 · p. 9, Table 4 · p. 9

6 Discussion

The discussion identifies constraints at low and high bond dimensions, the small scale of the test model, and an implementation that demonstrates parameter compression rather than inference-time savings. It also contrasts MPO structure with low-rank methods and notes an alternative optimization approach.

Accuracy gap at low χ

  • At χ = 4, per-layer reconstruction error exceeds 40%; training from scratch only partly compensates for insufficient representational capacity.The discussion suggests that a better factorization could reduce the gap at fixed χ.Source: §6 · p. 10

Diminishing returns at high χ

  • 0.8 percentage points separate χ = 16 and χ = 32 in accuracy, while parameter count more than doubles from 192 K to 409 K.For this model size, the discussion identifies χ = 16 as a practical sweet spot.Source: §6 · p. 10

Model scale

  • PicoGPT’s small, pedagogical scale makes it a controlled testbed, but these results do not establish performance on modern large-scale language models.Source: §6 · p. 10

Relation to low-rank methods

  • Unlike LoRA, MPO represents the entire weight in a multi-factor structure; its advantage over rank-χ factorization grows with L.Source: §6 · p. 10

Connection to DMRG

  • ALS offers a learning-rate-free alternative for compress-then-finetune, but non-convexity may still leave local minima.Source: §6 · p. 10

Limitations

  • The current forward pass materializes the dense weight matrix, so the implementation demonstrates parameter compression, not inference-time memory or FLOP reduction.Source: §6 · p. 10

7 Conclusion

MPO compression of all linear layers in a GPT-2-style transformer achieves 5–13× parameter compression per block. At χ = 16, the model retains 97.7% of the dense baseline’s token accuracy with 18.8% of its parameters.

7 Conclusion

  • 5–13× parameter compression per transformer block; at χ = 16, the MPO model retains 97.7% of dense-baseline token accuracy with 18.8% of its parameters.The conclusion describes compression of all linear layers in a GPT-2-style transformer.Source: §7 · p. 10
  • Future directions include inference-time FLOP reduction, ALS training at low χ, dynamic bond adaptation, larger models, and deeper tensor-network architectures.The paper suggests larger models may offer improved compression ratios.Source: §7 · pp. 10–11
  • The results support tensor-network structure as an interpretable inductive bias linking approximation quality to a controllable structural parameter.Source: §7 · p. 11

Code Availability

The full PyTorch implementation of MPO-PicoGPT is publicly available, with core code, benchmark scripts, and reproduction instructions.

Code Availability

Plain Language Summary

The paper asks whether MPOs can compact language-model weights without substantial accuracy loss. It reports 5–13× compression while retaining over 97% character-prediction accuracy, with compact tensors trained through standard gradient descent.

Plain Language Summary

  • The paper asks whether language-model weights can be stored more compactly without losing much accuracy.Source: Plain Language Summary · p. 11
  • An MPO represents a large weight matrix as a chain of smaller tensors connected by a bond, with bond dimension χ controlling compression.Source: Plain Language Summary · p. 11
  • The approach compresses each weight matrix by a factor of 5 to 13 while retaining over 97% of character-prediction accuracy.The compact tensors are learned automatically using standard gradient descent in PyTorch.Source: Plain Language Summary · p. 11

Funding

The research received no specific grant from funding agencies in the listed sectors.

Funding

  • The research received no specific grant.Source: Funding · p. 11
  • No specific grant was received from a funding agency.Source: Funding · p. 11
  • The statement covers funding agencies in the public, commercial, and not-for-profit sectors.Source: Funding · p. 11

Declarations

The authors report no known competing financial interests or personal relationships that could have influenced the work. The work uses public-domain data and required no ethical approval.

  • The authors declare no known competing financial interests or personal relationships that could have influenced the reported work.Source: Declarations · p. 11
  • The work involved no human participants, animal subjects, or sensitive personal data.Source: Declarations · p. 11
  • Tiny Shakespeare is publicly available and in the public domain, and no ethical approval was required.Source: Declarations · p. 11

A Heuristic justification of the MPO initialization scale

The initialization-scale rationale is a heuristic intended to explain its order-of-magnitude dependence on fan-in and bond dimension. It uses a crude variance estimate and is not presented as uniquely optimal.

Related papers

Loading 2603.28534v1…