Source-linked AI summary
Efficient GPT Model Pre-training using Tensor Train Matrix Representation
Viktoriia Chekalina, Georgii Novikov, Julia Gusak, Ivan Oseledets, Alexander Panchenko
TL;DR
Large transformer models are costly to train and deploy because their parameter counts demand memory, time, and electricity. This paper replaces GPT-2 fully connected layers with TTM layers and customizes their forward and backward operations. The resulting model has up to 40% fewer parameters and performs close to the original across language modeling and downstream tasks.
Problem
Training and employing large models requires memory, time, and electricity proportional to their parameter count.
Source: §1 · p. 1Method
The paper replaces GPT-2 fully connected layers with TTM-represented weights and customizes forward and backward operations through the layer.
Source: §4 · p. 2, §4.2 · p. 3Results
The GPT-based model has up to 40% fewer parameters while performing close to the original on language modeling, GLUE, and text summarization.
Source: §1 · p. 1Takeaways & Limitations
The authors conclude that GPT-2 TTM can replace GPT-2 in memory-restricted environments and that TTM layers can reduce the effective size of Transformer architectures.
Source: §7 · p. 6Custom-layer models must be trained from scratch, requiring several industrial GPUs for weeks, and the model was not validated on few-shot tasks.
Source: §8 · p. 6
Abstract
from arXiv · showhide
Large-scale transformer models have shown remarkable performance in language modelling tasks. However, such models feature billions of parameters, leading to difficulties in their deployment and prohibitive training costs from scratch. To reduce the number of the parameters in the GPT-2 architecture, we replace the matrices of fully-connected layers with the corresponding Tensor Train Matrix~(TTM) structure. Finally, we customize forward and backward operations through the TTM-based layer for simplicity and the stableness of further training. % The resulting GPT-2-based model stores up to 40% fewer parameters, showing the perplexity comparable to the original model. On the downstream tasks, including language understanding and text summarization, the model performs similarly to the original GPT-2 model. The proposed tensorized layers could be used to efficiently pre-training other Transformer models.
1 Introduction
Large language models deliver strong NLP results, but their parameter scale makes training and deployment costly. The paper replaces GPT-2 fully connected layers with TTM layers and reports a model with up to 40% fewer parameters and performance close to the original.
1 Introduction
- Training and employing models with vast parameter counts requires memory, time, and electricity proportional to model size.Source: §1 · p. 1
- The authors replace fully connected layers with sequential TTM containers and train the custom-layer architecture from scratch because the weight matrix is generally full-rank.They evaluate the pretrained model on in-domain and out-of-domain language modeling and downstream tasks.Source: §1 · p. 1
- The GPT-based model has up to 40% fewer parameters while performing close to the original on language modeling, GLUE, and text summarization.The paper presents these results as the outcome of its custom TT layer.Source: §1 · p. 1
2 Related work
Related work covers model compression through distillation, smaller GPT-like models, and Tensor Train representations. Prior TT work addressed embeddings and computer-vision linear layers, but had not tested TT representations in generative Transformers.
2 Related work
- Distillation was applied to BERT and GPT-2, while OPT provides a smaller model that imitates GPT-3 using more efficient training and particular datasets.Source: §2 · p. 1
- Tensor Train methods have reduced embedding-layer size, and TT-format linear layers have compressed computer-vision models.Source: §2 · p. 1
- The authors identify generative Transformers as an untested setting for TT representations.Source: §2 · p. 1
3 Math Background
The background introduces tensor contraction and Tensor Train Matrix representation, which factorizes a tensor into core tensors. It then describes applying sequential contractions to propagate a signal through a TTM weight matrix.
3 Math Background
- Tensor contraction joins two tensors along shared axes and returns a tensor over their uncontracted axes.Source: §3 · p. 2
- In TTM format, each tensor element is computed from core tensors, with boundary ranks R0 and RM equal to 1.The passage defines the rank as the sequence of intermediate ranks.Source: §3 · p. 2
- A TTM weight matrix is represented by M cores, and forward propagation contracts the input sequentially with those cores.Source: §3 · p. 2
- Sequential contractions produce an output of shape (B, J1, ..., JM), equivalent to (B, Dout).Source: §3 · p. 2
4 Efficient TTM Layer
The section replaces dense linear-layer weights with TTM cores and develops forward and backward computations intended to manage memory and time costs. It selects a forward and backward strategy based on memory and time, while noting that memory use can vary with rank.
4 Efficient TTM Layer
- TTM layers represent a linear-layer weight matrix as a sequence of cores, replacing the dense weight representation.The TTM layer's weight matrix has shape Din × Dout and is represented with M cores.Source: §4 · p. 2
- Memory consumption increased significantly at rank 16 in the measured GPT-2 TTM training iteration.The paper also compares memory footprints for fully connected and TTM layers under custom and PyTorch signal-propagation strategies.Source: §4 · p. 2
4.1 Forward pass
- Forward propagation contracts the input sequentially with TTM cores, with the contraction schedule affecting computation and saved-activation memory.The Einsum approach optimizes the schedule and can share intermediate results; a fixed schedule skips that optimization.Source: §4.1 · pp. 2–3
4.2 Backward pass
- Automatic differentiation stores many intermediate activations for TTM layers, motivating alternative backward-pass methods designed to use less memory.The paper describes Full Einsum and Full Matrix approaches alongside autodiff.Source: §4.2 · p. 3
- The backward contraction schedule is optimized along the batch axis, improving complexity when batch size multiplied by sequence length is large, as in Transformers.The passage ties this condition to the proposed backward computation.Source: §4.2 · p. 3
- The selected implementation pairs Einsum Forward with Full Matrix Backward, chosen as the most optimal pair according to memory and time.The backward method computes the weight gradient and then gradients for the TTM cores.Source: §4.2 · pp. 3–4
5 Singular Value decomposition (SVD) Layer
The SVD layer replaces a fully connected layer with two sequential linear layers built from truncated factors of its weight matrix. This approximates the original weight matrix while reducing the layer's parameter count.
5 Singular Value decomposition (SVD) Layer
- A weight matrix W is factorized into U, Σ, and V, then truncated factors define two sequential linear layers that replace the original layer.The truncated factors are Ur, Σr, and Vr.Source: §5 · p. 4
- Truncated SVD approximates the original weight matrix as W ≈ W2W1 and the layer output using the two replacement linear layers.The output approximation is described with the original input and bias.Source: §5 · p. 4
- For a layer with dimensions nin × nout, parameter count changes from nin × nout to r × (nin + nout) after truncated SVD.Source: §5 · p. 4
6 Experiments: end-to-end training of GPT-2 with custom layers
The experiments compare GPT-2 with TTM layers against the original model and parameter-matched SVD, evaluating language modeling and downstream tasks. TTM-72 retains comparable performance while reducing parameters, though results vary by setting.
6 Experiments: end-to-end training of GPT-2 with custom layers
- Models were trained from scratch against GPT-2 and parameter-matched SVD; TTM’s full-rank representation offers a broader weight space than truncated SVD.The paper’s rationale is that TTM can represent full-rank matrices, whereas the compared SVD structure is truncated.Source: §6 · p. 4
6.4 Natural Languge Understanding - GLUE
- On GLUE, classical models and TTM-layer models achieved approximately equal results, periodically overtaking one another; GPT-2 TTM-72 performed worse on Acceptability, QNLI, and MNLI in the benchmark.The pre-trained TTM-72 model received a task-compatible head and one training epoch; reported results average five random-seed runs, with standard deviation no more than 0.0008.Source: §6.4 · p. 5
6.5 Text summarization
- With 3,000 CNN/Daily Mail examples, classical GPT-2 and TTM-72 show similar summarization behavior; SVD-50 is somewhat worse except on ROUGE-L.The reported ROUGE scores match those in the cited paper’s pipeline.Source: Table 7 · p. 5, §6.5 · p. 5
6.1 Hyperparameter selection
- Core factor dimensions were made approximately equal, while ranks targeted 50–90% reduction for GPT-small and 40% for GPT-medium.The GPT-2-small fully connected layers use dimensions [768, 3072], factored as 768 = 4 × 6 × 8 × 4 and 3072 = 8 × 8 × 6 × 8.Source: §6.1 · p. 5
6.2 In-domain language modelling task
- Over 30% parameter reduction corresponds to a half a percent decrease in perplexity; over 40% reduction leads to a 3% drop.This is the in-domain Wikitext-103 evaluation, with GPT-2-small TTM ranks 16, 32, 64, and 80.Source: §6.2 · p. 5
6.3 Out-domain language modelling task
- TTM-72 saves 40% versus full GPT-2 and reaches perplexity 31, versus 55 for similarly sized SVD-50.OPT saves 7% of full GPT-2; its perplexity is described as the best among compressed models.Source: §6.3 · p. 5
7 Conclusion
Custom Tensor Train Matrix layers make GPT-2 40% lighter while performing without a significant quality drop across in-domain, out-of-domain, and downstream tasks. The authors also report that lower-rank structures such as SVD are less expressive when trained from scratch under the same conditions.
7 Conclusion
- Custom TTM layers replace fully connected layers in GPT-2, producing a model 40% lighter without a significant quality drop across language modeling and downstream tasks.The conclusion identifies memory-restricted environments as a use case and says TTM layers can reduce the effective size of any Transformer architecture.Source: §7 · pp. 5–6
- When trained from scratch under the same conditions, lower-rank structures such as SVD are less expressive than TTM models.Source: §7 · p. 6
8 Limitations
The custom-layer model must be trained from scratch, requiring substantial GPU resources and expertise. The work also lacks few-shot validation, and results can vary with the training pipeline.
8 Limitations
- Training the custom-layer model from scratch requires several industrial GPUs for several weeks and equipment that may be limited in academia.The passage also identifies a load-bearing power supply as necessary equipment.Source: §8 · p. 6
- The model was not validated on few-shot tasks, which the authors identify as a measure of a pre-trained model’s generalization ability.Source: §8 · p. 6
- Training large models requires expertise, and performance between identical architectures can vary significantly with the training pipeline.Source: §8 · p. 6