Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Xingtai Lv, Ning Ding, Kaiyan Zhang, Ermo Hua, Ganqu Cui, Bowen Zhou

TL;DR

This paper finds that low-rank pre-training, normally considered as efficient methods that will compromise performance, can be scalably effective when reduced parameters are precisely targeted and resolved by applying low-dimensional module only to the attention layer.

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, applying the 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 the vanilla Transformer.

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

TL;DR

This paper finds that low-rank pre-training, normally considered as efficient methods that will compromise performance, can be scalably effective when reduced parameters are precisely targeted and resolved by applying low-dimensional module only to the attention layer.

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, applying the 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 the vanilla Transformer.

Paper Structure

This paper contains 21 sections, 2 theorems, 6 equations, 6 figures, 10 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Low-rank Parameter-efficient Fine-tuning.
  4. Low-rank Pre-training for Neural Network.
  5. Low-dimensional Projected Attention
  6. Low-dimensional Module
  7. Position Optimization of Low-dimensional Module
  8. Explanation for Position Optimization
  9. Methodological Efficiency
  10. Experiments
  11. Effectiveness of LPA
  12. Experimental Settings.
  13. Results and analysis.
  14. Scaling up to 3.23B
  15. Downstream Tasks Performance
  16. ...and 6 more sections

Key Result

Lemma 1

In the attention layer, for the input vector $\mathbf{x}_i \in \mathbb{R}^{1 \times d_\text{in}}$ of the $i$-th input token, the corresponding output $\mathbf{z}_i \in \mathbb{R}^{1 \times d_\text{out}}$ satisfies indicating that $\mathbf{z}_i$ is dependent on all the vectors in the input $\mathbf{x}$, especially for the computation in the Key, Value layers.

Figures (6)

  • Figure 1: An illustration of the Low-dimensional Projected Attention (LPA). The calculations in $\texttt{softmax}$ function measure the relationships between input tokens.
  • Figure 2: Training loss for the 2.43B LPA model, the 3.23B Same-Dim Transformer, and the 2.49B Transformer with nearly the same parameter count as the LPA model.
  • Figure 3: Training loss for Transformer and LPA models with different $r$. The darker curves correspond to larger values of $r$ in LPA.
  • Figure :
  • Figure :
  • ...and 1 more figures

Theorems & Definitions (2)

  • Lemma 1
  • Lemma 2