Scalable Efficient Training of Large Language Models with Low-dimensional Projected Attention (2024)

Abstract: Improving the effectiveness and efficiency of large language models (LLMs) simultaneously is a critical yet challenging research goal. In this paper, we find that low-rank pre-training, normally considered as efficient methods that will compromise performance, can be scalably effective when reduced parameters are precisely targeted. Specifically, by applying low-dimensional module only to the attention layer — resolves this issue and enhances both effectiveness and efficiency. We refer to this structure as *Low-dimensional Projected Attention (LPA)* and provide an explanatory analysis. Through extensive experimentation at parameter scales of 130M, 370M, and scaling up to 3B, we have validated the effectiveness and scalability of LPA. Our results show that LPA model can save up to 12.4% in time while achieving an approximate 5% improvement in test perplexity (ppl) and on downstream tasks compared with vanilla Transformer.

## introduction improving large language models' (llms) @xcite @xcite @xcite effectiveness and efficiency simultaneously presents challenges due to inherent tradeoffs, which remains a critical research goal in the research field. among series methods proposed to alleviate this issue, parameter-efficient fine-tuning @xcite @xcite offer valuable insights. notably, low-rank or lowdimension techniques such as lora @xcite demonstrate on-par or even enhanced performance over traditional full-parameter fine-tuning with reduced computational resources. intuitively, besides the fine-tuning phase, adapting lora's principles to the pre-training phase * corresponding authors through low-rank decomposition is both viable and promising, which can yield substantial benefits if effectiveness is maintained. however, existing studies have found that the direct low-rank pre-training often compromises the effectiveness. to reduce such effects, strategies such as iteratively accumulating low-rank updates @xcite or integrating low-rank decomposition directly into the gradient @xcite have been suggested. whether it's the original lora or these improved methods, they all involve performing low-rank decomposition and updates on "amounts of change" (weights or gradients), and do not reduce the number of parameters in the model itself, which face obstacles in maintaining efficiency during subsequent inference and fine-tuning stages. therefore, an ideal scenario would be permanently reducing the number of parameters (computational load) through efficient methods, without compromising or even enhancing the performance of pre-trained models. to achieve this goal, is it feasible to directly perform low-rank decomposition on the matrices in the model itself, rather than on the changes? current limited research suggests that existing lowrank pre-training methods experience performance losses and uncertainties @xcite , with even fewer studies exploring more direct approaches. however, in this paper, we demonstrate that such direct low-rank pre-training is feasible, provided that the parameters to be reduced are more precisely targeted. specifically, we describe the reduction of parameters as replacing the original matrices with low-dimensional modules. we find that using low-dimensional modules in the feed-forward neural (ffn) layers or across all layers negatively impacts the model's effectiveness. however, we observe that employing them in the attention layers consistently allows the model to outperform the original transformer. we refer to this structure as low-dimensional projected at-tention (lpa), provide an explanation, and experimentally demonstrate its ability to reliably enhance both the efficiency and effectiveness of the model. we validate the effectiveness of the lpa model on two transformer model configurations, assessing both pre-training and downstream task performance. with a particular focus on the scalability of lpa model, we observe that it remains effective even when the model parameters scale up to 3b. furthermore, our study explores the effects of the hyperparameter on lpa, the necessity of integrating the low-dimensional module into every sublayer of the attention layer, and how to distribute any extra parameters effectively. the code of this work will be publicly available at https://github.com/tsinghuac3i/lpa . ## related work low-rank parameter-efficient fine-tuning. parameter-efficient fine-tuning optimize only a tiny portion of parameters while keeping the majority of the neural network frozen @xcite @xcite @xcite , saving significant time and computational costs and achieving performance comparable to full parameter fine-tuning on many tasks @xcite . low-rank adaptation (lora) is one of the most effective and influential parameter-efficient fine-tuning methods, having found widespread application @xcite . the lora method involves freezing the weights w 0 of the pre-trained model while training two low-rank decomposition matrices w u and w d , resulting in the output of the lora module being represented as z ← w 0 x + w u w d x. we drew inspiration from lora and its improvement works, adapting them to the pre-training process to enhance effectiveness and efficiency of the model. ## low-dimensional projected attention we use a low-dimensional module for replacing the original weight matrix, and observe the varying effects of incorporating the low-dimensional structure in different modules. we provide an explanatory analysis of these findings and propose the low-dimensional projected attention (lpa). additionally, we examine the efficiency of this approach. ## experiments extensive experiments are conducted to validate the effectiveness of lpa across models of various scales, particularly emphasizing its efficacy with the 3.23b models. furthermore, we investigate the impact of hyperparameter r on lpa, whether applying the low-dimensional module to all sublayers in the attention layer is necessary, and the allocation of surplus parameters. 2016 ), consisting of 0.1b tokens, and set r of lpa to 256. for model setting 2, we pretrain the models using 2.6b tokens from the pile dataset @xcite for the 130m parameter model and 6.8b tokens for the 370m parameter model, with the lpa architecture r set to 128 or 256. detailed model configurations and training hyperparameters are provided in ## conclusion this paper demonstrates that low-rank pre-training can enhance both the effectiveness and efficiency of llms when reduced parameters are precisely targeted. by incorporating low-dimensional modules specifically in the attention layers, we develop the low-dimensional projected attention (lpa), which outperforms transformers without the efficiency compromises. our empirical analysis and experiments show that lpa maintains its effective-ness even as model parameters scale up to 3b. additionally, we explore the impact of hyperparameters and the optimal reallocation of surplus parameters, providing a robust framework for future enhancements in llm pre-training.