Tensor Train Low-rank Approximation (TT-LoRA): Democratizing AI with Accelerated LLMs

Let $P_{\theta}(y,x)$ be a pre-trained language model, parameterized by $\theta$. For example, $P_{\theta}(y,x)$ could represent a versatile multi-task learning model, such as DistilBERT [22], DeBERTa [23], or LLaMA-3 [1], all of which are developed from the Transformer architecture [24] initially proposed by Vaswani et al. in 2017. We consider adapting the pre-trained model to downstream tasks, such as text classification, summarizing, question answering, and sentiment analysis. Each downstream task is represented by a training dataset of input-output pairs: $\mathcal{D}:{(x_{i},y_{j})}_{i=1}^{N},_{j=1}^{M}$. Here, $x_{i}$ represents a sequence of input tokens. The output $y_{j}$ can vary depending on the task, for instance, it may be a sequence of tokens for text generation tasks, categorical labels for classification tasks, or continuous values for regression tasks. For example, for sentiment analysis, $x_{i}$ is a social media post, and $y_{j}$ is the categorical label; for the question-answering task, $x_{i}$ is the question, and $y_{j}$ is the answer; for the summarizing task, $x_{i}$ is the article and $y_{j}$ is the summary of the corresponding article. Consider $P_{\theta}(y,x)$ to be a text generation task-based pre-trained model. During the full fine-tuning, the model is initialized with its pre-trained weights $\theta_{0}$ and subsequently updated to $\theta_{0}+\Delta\theta$. Here, $\Delta\theta$ represents the modifications made to the pre-trained weights, specifically tailored to enhance performance on the downstream task. These adjustments are derived by optimizing the following objective function: [19]:

Here, the objective function aims to maximize the cumulative log probability of correctly predicting each output token $y_{t}$, given the input sequence $x$ and all preceding output tokens.

A major drawback of full fine-tuning is that all the model parameters need to be trained for each downstream task, which makes the dimension of $|\Delta\theta|$ equal to $|\theta_{0}|$. In this paper, we introduce a parameter-efficient fine-tuning strategy that significantly decreases the number of trainable parameters from $\theta_{0}$ to $\phi$, where the dimension of $|\Delta\phi|\ll|\Delta\theta|$.

Transformer-based embedding models, such as BERT, depicted in Figure 3(a), comprise multiple dense layers, including numerous encoder blocks and a fully connected feed-forward neural network. Each encoder block, detailed in Figure 3(b), contains a complex sub-layer arrangement, specifically multi-head attention layer and a fully connected feed-forward neural network. The weight matrices associated with these dense layers possess full rank, ensuring that each input feature uniquely contributes to the output and maximizes the learning capacity of the layer. However, when adapting to a specific downstream task, Aghajanyan et al. [25] show that the pre-trained LLMs have a low intrinsic dimension and Hu et al. [19] hypothesize that the updates to the weights of the pre-trained models also have a low intrinsic rank, highlighting a potential area for model compression during adaptation to downstream tasks. Therefore, in this paper, for a pre-trained weight matrix $W_{0}\in\mathbb{R}^{m\times n}$ and its corresponding update during adaptation $\Delta W\in\mathbb{R}^{m\times n}$, we constrain the update through the proposed TT-LoRA approach by presenting $\Delta W$ with low-rank representations, as shown in Figure 3(c).

To achieve a low-rank approximation of $\Delta W$, we utilize tensor decomposition [26], more specifically, Tensor Train (TT)[27] decomposition. TT decomposition decomposes a tensor into a series of low-rank, small, three-dimensional tensors (cores). The product of these low-rank tensor cores provides an accurate approximation of the original tensor, significantly reducing its dimensionality while preserving essential structure and information. The fundamental property of TT decomposition is that each core tensor interacts only with its immediate predecessor and successor. This localized interaction allows operations, such as tensor core multiplication to approximate the original tensor, to be performed in a step-by-step manner, focusing on smaller, manageable pieces rather than the entire tensor at once. As a result, the complexity of tensor operations is significantly simplified, leading to substantial reductions in computational and memory requirements [13]. However, TT decomposition is particularly applicable to high-dimensional tensors. Therefore, in TT-LoRA, the matrix $\Delta W$, which is initialized using a random Gaussian distribution, is first represented as a $d$-dimensional tensor $\Delta\mathcal{W}\in\mathbb{R}^{k_{1}\times\dots\times k_{d}}$. Here, $\prod_{i=1}^{d}k_{i}=m\times n$. The $d$-dimensional tensor $\Delta\mathcal{W}$ is then decomposed into $d$ number of small tensor cores $\mathcal{C}_{1},\dots,\mathcal{C}_{d}$. The shape of each tensor core can be defined as $\mathcal{C}_{i}\in\mathbb{R}^{r_{i-1},k_{i},r_{i}}$, given the TT rank $[r_{0},\dots,r_{d}]$, where the first ($r_{0}$) and last ($r_{d}$) TT ranks are 1. Consequently, the total number of parameters in the tensor train decomposition of $\Delta\mathcal{W}$ can be represented as:

Here, $r_{i-1}\times k_{i}\times r_{i}$ is the size of the $i$-th tensor core. During fine-tuning, the pre-trained weight matrix $W_{0}$ remains frozen and does not receive gradient updates, while the $\Delta\mathcal{W}$ contains trainable parameters. The adapted weight matrix $W_{adapted}$ is given by:

Using TT decomposition:

Thus, this adapted layer applies a linear transformation to an input $x$ can be described as follows:

Here, $W_{0}$ represents the pre-trained weight matrix of shape $m\times n$, $\Delta W$ is the update to the weight matrix in full fine-tuning of shape $m\times n$, $\Delta\mathcal{W}$ is the d-dimensional tensorized matrix of shape $k_{1}\times\dots\times k_{d}$, $\mathcal{C}_{i}$ is the $i$-th tensor core of shape $r_{i-1}\times k_{i}\times r_{i}$ and $\alpha$ is a fixed scaling factor used to scale the update before adding to the pre-trained weight.

The compression ratio of TT-LoRA is closely related to the choice of TT ranks and the structure of d-dimensional tensor $\Delta\mathcal{W}$. For instance, consider a $W_{0}$ with dimensions $768\times 2304$. In full fine-tuning, $\Delta W$ will have the exact dimensions of $768\times 2304$, resulting in $1,769,472$ trainable parameters. However, with TT-LoRA, considering $\Delta\mathcal{W}$ as 7D tensor of shape $12\times 8\times 8\times 3\times 8\times 8\times 12$ with TT-rank of $5$ results in $1,135$ trainable parameters, achieving a remarkable $1560$x compression ratio. Reducing the TT-rank and/ or increasing the tensor dimension increases the compression ratio even further. To obtain the optimal TT-rank and tensor dimensions for our proposed method, we did a thorough hypermarameter search which is discussed in Section IV-D.

We initially conducted experiments on (RoBERTa) Robustly Optimized BERT Pretraining Approach [28], which is an optimized method for training BERT (Bidirectional Encoder Representations from Transformers) [34], a transformer-based LLM. Developed by researchers at Meta AI, RoBERTa revises BERT’s pretraining methodology to improve the model’s performance in several ways, such as dynamic masking, eliminating the next sentence prediction loss, and increasing the batch size while decreasing the learning rate. Apart from RoBERTa, We performed experiments on DeBERTa (Decoding-enhanced BERT with disentangled Attention) [23], a recent variant of BERT trained on a larger scale. Developed by Microsoft, DeBERTa improves the BERT architecture by introducing a novel disentangled attention mechanism that separately models the content and position, enhancing the model’s ability to understand contextual relationships in text. We utilized the pre-trained RoBERTa-base and DeBERTa-base from the HuggingFace Transformers library.

In RoBERTa and DeBERTa architecture, each encoder layer includes four weight matrices within the self-attention module ($W_{q}$, $W_{k}$, $W_{v}$, $W_{o}$) and a feed-forward neural network (FFNN). While TT-LoRA can be applied to any of the weight matrices in a neural network to reduce the number of trainable parameters, our experiments specifically target $W_{q}$ and $W_{v}$. This focus aligns with findings from the LoRA paper, which indicates that models achieve optimal performance when these particular weights are fine-tuned, while the remaining weights are kept frozen [19]. We performed hyperparameter optimization by fine-tuning the model with different TT-LoRA parameter initializations, the details of which will be presented later. Using the HyperBand optimizer [35], we identified the most efficient parameters. For each run, the model was fine-tuned for up to 20 epochs with an early stopping criterion of 5 epochs. Specifically, the training was halted if the validation loss did not improve for 5 consecutive epochs. The best model was selected based on the lowest observed validation loss from these runs.For reporting performance on the GLUE benchmark tasks, we use the following metrics: matched accuracy for MNLI, Matthews correlation coefficient for CoLA, Spearman correlation coefficient for STS-B, F1 score for both MRPC and QQP, and accuracy for all other tasks. Table I summarizes the downstream task performance comparison between TT-LoRA and other baseline PEFT methods.

As shown in Table I, TT-LoRA consistently achieves superior or comparable performance to other PEFT methods when FT DeBERTa on GLUE tasks, with no more than 0.2M trainable parameters. Consequently, TT-LoRA stands out for its efficiency by outperforming 9 out of 10 FT approaches in both model compression and accuracy.Prompt Tuning, which utilizes just 0.01M trainable parameters compared to TT-LoRA’s 0.02M, surpasses TT-LoRA in model compression but significantly suffers in performance. Nevertheless, TT-LoRA requires merely twice the trainable parameters of Prompt Tuning while achieving, on average, about 15% higher model accuracy. TT-LoRA also achieves achieves substantial model compression when FT RoBERTa, significantly reducing the number of trainable parameters. Specifically, TT-LoRA has reduced trainable parameters by approximately factors of $6,200\times$ for full FT, $31.5\times$ for LoRA, $5\times$ for BitFit, and $5\times$ for LoRETTA_adp. Despite this reduction, TT-LoRA outperforms other PEFT methods, except for full FT, in average model accuracy across various GLUE task sets.

Encouraged by the outcomes observed with DeBERTa and RoBERTa models, we extended our experiments to include the larger-scale Large Language Model at Meta AI (LLaMA) [36] models. LLaMA is a series of larger-scale language models designed for various natural language understanding and generation tasks. In our experiments, we utilized LLaMA-2-7B and LLaMA-3-8B models to assess the effectiveness of our proposed PEFT method. We chose these LLaMA models due to their extensive parameter sets, ranging into the billions, which provide a rigorous test environment to evaluate how well our PEFT approach can compress and optimize these substantial models while maintaining comparable performance levels. For our experiments, we utilized the pre-trained LLaMA2-7b and LLaMA3-8b models available in HuggingFace Transformers library. We conducted a comparative analysis of TT-LoRA against other baseline PEFT methods using the SuperGLUE benchmark tasks and the results are summarized in Table II. Similar to BERT family experiments, TT-LoRA has been applied to $W_{q}$ and $W_{v}$ weight matrices of the self-attention module of LLaMA models and has been trained over multiple epochs, stopping the process if the validation loss did not improve for 5 consecutive epochs. The best model is then chosen from these runs based on the lowest observed validation loss. For reporting performance on the SuperGLUE benchmark tasks, we used F1 score for both CB and WSC tasks, and accuracy for both BoolQ and COPA tasks.

Table II highlights TT-LoRA’s performance on the LLaMA2-7B model, where it consistently outperforms all other PEFT methods on CB and WSC tasks and all but LoRETTA_rep on the BoolQ task. It achieves comparable performance on the COPA task and, on average, surpasses all competing PEFT methods. This higher model performance is achieved while attaining significant reductions in trainable parameters compared to our baselines by factors of approximately 67,384x (full fine-tuning, FT), 500x (Adapter), 41.9x (LoRA_r=8), 13.1x (Prefix), 5.1x (LoRETTA_rep), and 8.8x (LoRETTA_adp).

We extended our experiment by FT the LLaMA3-8B model on the SuperGLUE benchmark and compared its performance against LoRA. As shown in Table II, TT-LoRA either matched or exceeded the performance of LoRA across all tasks while achieving a remarkable reduction in trainable parameters, which is approximately $170.5\times$ fewer parameters compared to LoRA.

We evaluate the storage requirements of various PEFT methods applied to fine-tune the DeBERTa and LLaMA2-7B models. Specifically, we select the three top-performing PEFT methods from Table I and Table II. The storage calculations are based on the assumption that the model weights are stored with 16-bit precision. As demonstrated in Figure 4, TT-LoRA reduces the storage requirements for trainable parameters when used with DeBERTa, necessitating only $39$ KB. This represents a significant reduction in storage needs, achieving reductions by factors of approximately $2.5\times$, $5\times$, $15\times$, and $7000\times$ compared to LoRETTA_rep, LoRETTA_adp, LoRA_r=8, and full FT, respectively. A similar memory efficiency is achieved through TT-LoRA when used with LLaMA2-7B, requiring approximately $195$ KB of storage for trainable parameters. This storage efficiency significantly surpasses that of other PEFT methods, compared to our baselines reducing storage demands by factors of roughly $5\times$ (LoRETTA_rep), $8.8\times$ (LoRETTA_adp), $42\times$ (LoRA_r=8), and $65,925\times$ (full fine-tuning, FT). The reduced storage requirement of TT-LoRA establishes it as a highly efficient approach for FT LLMs, particularly on hardware with limited resources. This efficiency facilitates broader deployment options and ensures that larger-scale LLM capabilities are accessible even in constrained environments.

To optimize the performance of TT-LoRA, we conducted a comprehensive hyperparameter search using Ray Tune [37], a scalable hyperparameter optimization framework. The hyperparameters are selected to ensure a robust comparison of model performance across different configurations and tasks. We employed the HyperBand optimizer [35] to explore a discrete search space for hyperparameter combinations, aiming to minimize validation loss with fewer simulations. This approach allowed us to estimate the optimal parameters efficiently, performing around 200 searches. The search space included learning rates and specific model-related parameters such as tensor shapes, ranks, and alpha values, as highlighted in Table III. The best model configurations were determined based on the lowest validation loss observed during the tuning Ray Tune trials. The optimal hyperparameters for TT-LoRA, such as tensor shape, tensor rank, $\alpha$, and learning rate are summarized in Table IV. The hyperparameters used for other PEFT methods are demonstrated in [12].

In addition to identifying the optimal model configuration, the logs generated by Ray Tune provide a detailed analysis of the relationship between model performance and the number of trainable parameters. This insight is crucial for optimizing a PEFT method, especially in resource-constrained environments where efficiency is also important. Figure 5 illustrates how the number of trainable parameters, determined by specific tensor shapes and ranks, affects the accuracy of LLaMA2-7B model fine-tuned with TT-LoRA. The heatmap in Figure 5 depicts the $\alpha$ value, signifying the magnitude of updates to the pre-trained weights during fine-tuning. As shown in Figure 5, model accuracy tends to decrease as the number of trainable parameters is reduced. In addition to model parameters, the value of $\alpha$ also influences the model accuracy. Nonetheless, TT-LoRA maintains superior or comparable performance to other PEFT methods, even with fewer parameters, demonstrating its capability to effectively minimize trainable parameters within strict hardware limits.