TensorGPT: Efficient Tensor-Train Compression for LLMs

• TensorGPT is a tensor compression method that replaces high-dimensional parameter matrices with a series of low-dimensional tensor cores to reduce memory usage.
• The approach employs TT-SVD and matrix product state techniques to achieve up to 38.4× compression, significantly easing deployment on resource-constrained hardware.
• It integrates end-to-end training and optimized forward/backward propagation strategies, maintaining robust language modeling performance despite dramatic parameter reductions.

TensorGPT refers to a class of compression approaches for LLMs, particularly those within the GPT family, based on tensor-train decomposition (TTD) and matrix product state (MPS) representations. TensorGPT targets drastic reductions in model memory and storage requirements—making transformer architectures feasible for deployment on edge or resource-constrained devices—by exploiting the inherent low-rank structure of both token embeddings and weight matrices. It also includes schemes for end-to-end training and inference with tensorized layers. The fundamental innovation lies in replacing high-dimensional parameter matrices with a sequence of low-dimensional "cores," yielding significant parameter compression while preserving language modeling capacity (Xu et al., 2023, Chekalina et al., 2023).

1. Challenges in LLMs and Token Embeddings

State-of-the-art LLMs such as GPT-2 allocate a substantial portion of total parameters to the token embedding matrix. For instance, GPT-2's embedding table $W \in \mathbb{R}^{V \times d}$ (with $V = 50\,257$, $d = 768$) constitutes approximately 31% of its overall 1.5 B parameters. The resulting tens of millions of parameters in $W$ not only increase storage burdens but also preclude on-device deployment, such as on a Raspberry Pi, which is restricted by limited RAM/flash capacity. The situation is aggravated if the vocabulary is dynamic, as naïvely recomputing a full SVD for every new token insertion is computationally infeasible (Xu et al., 2023).

These bottlenecks also affect the fully connected layers in the Transformer’s feed-forward blocks, where each linear projection contains millions of parameters—complicating training, storage, and inference, especially for non-datacenter applications (Chekalina et al., 2023).

2. Tensor-Train Decomposition and Matrix Product States

Tensor-Train Decomposition (TTD), introduced by Oseledets (2011), expresses high-dimensional tensors $$\mathcal{X} \in \mathbb{R}^{I_1 \times \cdots \times I_N}$$ as products of $N$ cores $\{G^{(n)}\}$, each of size $R_{n-1} \times I_n \times R_n$, with $R_0 = R_N = 1$. The representation is:

$$\mathcal{X}_{i_1,\dots,i_N} \approx \sum_{r_0,\dots,r_N} G^{(1)}_{r_0, i_1, r_1} G^{(2)}_{r_1, i_2, r_2} \cdots G^{(N)}_{r_{N-1}, i_N, r_N}$$

In physics, this format is known as the Matrix Product State (MPS). TT-SVD, a sequential SVD procedure, facilitates automated extraction of these cores (see Alg. 1 in (Xu et al., 2023)). By choosing the mode sizes $V = 50\,257$0 and number of cores $V = 50\,257$1 to approximate the original dimension and by truncating SVDs to lower TT-ranks, substantial parameter compression is achieved.

For weight matrices $V = 50\,257$2, a closely related Tensor Train Matrix (TTM) factorization views $V = 50\,257$3 as a $V = 50\,257$4-way tensor and factorizes it into $V = 50\,257$5 matrix-valued TT cores, reducing storage from $V = 50\,257$6 to a sum scaling as $V = 50\,257$7, where $V = 50\,257$8 is the uniform TT-rank and $V = 50\,257$9, $d = 768$0 are mode sizes for input/output (Chekalina et al., 2023).

3. Methodology and Implementation in TensorGPT

Token Embedding Compression (MPS Approach)

Given the embedding matrix $d = 768$1, TensorGPT compresses each token embedding $d = 768$2 as follows:

1. Padding & Tensorization: Pad $d = 768$3 to a product of small mode sizes, ideally a power of two. For GPT-2, $d = 768$4 is padded to $d = 768$5, giving $d = 768$6, $d = 768$7.
2. Per-token MPS/TT-SVD: Reshape $d = 768$8 to an order-$d = 768$9 tensor $W$0 and perform TT-SVD, producing cores $W$1 for $W$2.
3. Storage: Retain the set of TT cores across all tokens, vastly reducing memory versus storing the dense $W$3 (Xu et al., 2023).

This per-token compression is training-free, requires no extra data, and leverages a scan over TT-rank choices $W$4 (e.g., $W$5) to trade off reconstruction accuracy and storage.

Tensor Train Matrix for Projections

For fully-connected layers, the Tensor Train Matrix (TTM) representation proceeds as follows:

• Reshape: $W$6 is viewed as a $W$7-way tensor, dimensions $W$8; $W$9.
• TTM Factorization: $$\mathcal{X} \in \mathbb{R}^{I_1 \times \cdots \times I_N}$$0, where each $$\mathcal{X} \in \mathbb{R}^{I_1 \times \cdots \times I_N}$$1 is a rank-$$\mathcal{X} \in \mathbb{R}^{I_1 \times \cdots \times I_N}$$2 four-way core.
• Forward and Backward Propagation: Efficient einsum-based forward contraction and custom memory-aware backward passes are implemented. The optimal (Einsum forward + Full-Matrix backward) reduces per-layer gradient memory from $$\mathcal{X} \in \mathbb{R}^{I_1 \times \cdots \times I_N}$$31 GB to $$\mathcal{X} \in \mathbb{R}^{I_1 \times \cdots \times I_N}$$4192 MB (Chekalina et al., 2023).

4. Parameter Reduction and Empirical Performance

Embedding Layer Compression

Theoretical compression ratio is:

$$\mathcal{X} \in \mathbb{R}^{I_1 \times \cdots \times I_N}$$5

For typical choices ($$\mathcal{X} \in \mathbb{R}^{I_1 \times \cdots \times I_N}$$6, $$\mathcal{X} \in \mathbb{R}^{I_1 \times \cdots \times I_N}$$7 for $$\mathcal{X} \in \mathbb{R}^{I_1 \times \cdots \times I_N}$$8), this yields $$\mathcal{X} \in \mathbb{R}^{I_1 \times \cdots \times I_N}$$9, i.e., a 3.31-fold reduction. Aggressive choices ($N$0) can achieve up to $N$1 compression (Xu et al., 2023). For GPT-2, this method theoretically results in $N$2 fewer total model parameters, with a compression ratio $N$3 for the embedding layers.

Method	TT-ranks	η (×)	Generation Loss
Original	–	1.00	13.71
Direct TT on $N$4 (2-core)	[1,32,1]	23.64	10.14
MPS on each $N$5	[1,2,4,\ldots,4,2,1]	3.31	9.01
MPS on each $N$6 (max η)	[1,1,...,1]	38.40	20.77

A compression of $N$7 yields a lower generation loss (9.01) than the original (13.71), indicating slight regularization from the low-rank approximation. Maximal compression reduces performance but maintains core semantic fidelity (Xu et al., 2023).

TTM for Feed-Forward Layers

By replacing GPT-2’s dense MLP matrices with TTM layers, the following parameter reductions and effects on perplexity are observed (Chekalina et al., 2023):

Model	#Params (M)	% of Baseline	Perplexity
GPT-2 Small	124.4	100%	17.55
TTM-16	68.1	54%	21.33
TTM-64	83.6	67%	18.08

A 40% size reduction leads to a perplexity increase of only ≈3.2, with downstream GLUE and summarization performance matching or slightly underperforming standard GPT-2 and surpassing SVD-compressed and distilled baselines.

5. Computational Complexity and Hardware Implications

The offline TT-SVD compression for embeddings has $N$8 per token cost. Online reconstruction requires $N$9 per token, and storage is reduced to $\{G^{(n)}\}$0, which is significantly less than the $\{G^{(n)}\}$1 for the original matrices. For hardware-limited environments, such as Raspberry Pi, this enables models to fit entirely in memory or flash (Xu et al., 2023).

In the TTM approach for projections, TTM layers' forward passes have complexity $\{G^{(n)}\}$2; custom backpropagation routines ensure memory usage is comparable to dense layers (Chekalina et al., 2023).

6. Insights, Limitations, and Broader Applicability

A key empirical finding is that GPT embeddings admit a strong low-rank structure. Certain tensor modes exhibit lower mean absolute errors after reconstruction, suggesting interpretable subspace significance. While aggressive compression ($\{G^{(n)}\}$3) degrades generation quality, moderate compression ($\{G^{(n)}\}$4) can preserve or even improve downstream task performance—potentially via regularization effects (Xu et al., 2023).

Practical implementation guidelines recommend padding embedding dimensions to powers of two, initializing with low TT-ranks, and compressing new tokens in dynamic vocabularies without affecting existing TT cores.

Current limitations include the need to integrate position/mask embeddings into the MPS format and to validate full end-to-end inference speedup in hardware. Ongoing extensions involve generalizing TT-based compression to other transformer blocks, including BERT, GPT-3, and OPT, by replacing dense projections with TTM layers, and designing TT-based transformer blocks for entire model compression and acceleration (Xu et al., 2023, Chekalina et al., 2023).

7. Conclusion and Future Directions

TensorGPT demonstrates that LLMs, especially those from the GPT family, have substantial low-rank redundancy in both embeddings and dense projections. By leveraging tensor-train decomposition and matrix product state representations, up to 38.4× parameter compression is achievable without additional fine-tuning, and with little or no loss in performance for moderate compression settings. The approach is broadly applicable across transformer-based architectures and is particularly advantageous for on-device applications constrained by memory and storage. A plausible implication is that further efficiency gains are attainable by integrating positional/mask encodings into the tensorized formats and by implementing TT-based end-to-end transformer blocks, paving the way for highly compressed and efficient LLM deployments (Xu et al., 2023, Chekalina et al., 2023).