Yongil Kim

Seoul National University

[ICLR2024] LOFTQ: LORA-FINE-TUNING-AWARE QUANTIZATION FOR LARGE LANGUAGE MODELS

01 Apr 2024 » LLM, PLM

[pdf][github]

Yixiao Li1∗, Yifan Yu1∗, Chen Liang1, Pengcheng He2, Nikos Karampatziakis2, Weizhu Chen2, Tuo Zhao1
Equal contribution, 1 Li, Yu, Liang, and Zhao are affiliated with Georgia Institute of Technology. Correspondence to yixiaoli@gatech.edu, yyu429@gatech.edu, and tourzhao@gatech.edu., 2 He, Karampatziakis, and Chen are affiliated with Microsoft Azure.  

image

Abstract

  • (Quantization and LoRA) LLM 을 Finetuning 하기 위해 필수불가결한 요소가 Quantization 이고, 최근 LoRA fine-tuning 기법을 통한 quantization 연구도 활발하다. 기존의 연구들은 quantization 과 LoRA 를 같이 적용하였을 때, full fine-tuning 과 비교하여 consistent gap 이 있음을 한계점으로 지적한다.
  • ( LoftQ ) 이에 저자들은 LLM 을 quantize 하면서 동시에, LoRA fine-tuning 을 위한 proper low-rank initialization 을 찾는 LoRA-Fine-Tuning-aware Quatization, LoftQ 를 제안한다. 이 방법론은 full-precision model 과 quantized model 사이의 discrepancy 를 경감시켜 downstream task 에서의 generalization 성능을 향상시킨다.
  • (Experiment) NLU, QA, Summarization, NLG task 등에 적용하였을 때, 기존의 quantization method 보다 우수한 성능을 보이고, 특히 어려운 2-bit 이나 2/4-bit mixed precision regime 에서 강력한 성능을 보임을 확인한다.

1. Introduction

▶ LLM and their costs
Large Language Model (LLM) 이 자연어 이해 (NLU) 와 자연어 생성 (NLG) 에서, 다른 모델들과 (LLM 이 아닌 모델들)과 비교가 불가능할 정도로 압도적인 성능을 보인다. 그러나 그들은 extensive computational and memory cost 를 요구한다. 특히 Training 을 어렵게 할 뿐 아니라, deploying 이나 테라포밍 단계에서 매우 많은 resource 를 요구한다.

▶ Quantization and LoRA
이 extensive requirement 를 해결하기 위해, qunatization 이 pivotal compression technique 으로 많은 연구가 되고 있다. Quantization 기법은 high-preicison numerical value 를 discrete value set 으로 변환시키는 것이다. 보통 model 들이 16-bit float format 으로 저장되어있는 것을 4-bit integer format 으로 quantization 시키면 storage overhead 가 75% 나 줄어드는 것이다.

Low-Rank Adaptation (LoRA) 는 quantized pre-trained model 을 downstream task 에 효과적으로 adaptation 을 시킬 수 있는 매우 중요한 방법이다. 이 방법은 fully fine-tuned weight 과 pre-trained weight 의 차이는 low-rank property 를 보인다는 점을 가정한다. 이 가정으로 그 차이점을 low-rank matrix 를 활용해 표현한다. 그 결과, pre-trianed weight 은 고정한 채, low-rank matrix 만을 solely train 하여 효과적인 task adaptation 이 가능하게 한다.

기존에는 보통 pre-trained model 을 quantizing 할 때, 추후의 LoRA fine-tuning 의 중요성은 무시 한채, quantization 기술에만 집중하였다. 예를 들어, QLoRA 의 경우 LoRA 에서 사용되는 fixup initialization 을 상속받아(inherit), quantized pre-trained model 에 zero initialized low-rank adapter 를 붙인다. 이렇게 될 경우, 2-bit regime 같은 극단적인 low-bit situation 에서 qunatization 학습을 위한 bix approximation 방법이 LoRA finetuning 의 initialization 에 영향 을 미칠 수 있다. 아래의 그림 왼쪽(a) 처럼, QLoRA 의 quantized pre-trained model 은 3-bit level 이하에서 심각한 degradation 이 있다. 이 initialization 에서의 일탈(deviation)은 fine-tuning performance 에 큰 나쁜 영향을 미친다. 오른쪽 (b) 그림 처럼, QLoRA 를 적용하면 quantization bit 이 작아질 수록 fine-tuning performance 가 크게 감소한다. QLoRa 가 3-bit level 이하에선 실패하는 것을 보이는 것은 noteworthy 하다.

image

▶ LoftQ
이에 저자들은 LoRA-Fine-Tuning-aware Quantization (LoftQ) 방법론을 제안한다. 이 것은 pre-trained model 중에서 quantization 과 LoRA fine-tuning 을 모두 필요로 하는 모델을 타겟으로 한다. 이 framework 은 low-rank approximation 과 quantization 을 active 하게 통합한다. 이 시너지(synergy)는 아래 그림처럼 original pre-trained model 과 quantized model 사이의 discrepancy 를 크게 줄여준다. 결과적으로, 추후의 LoRA fine-tuning 을 위한 효과적인 initialization point 를 제공하여 downstream task 의 improvement 를 이끌어낸다.

image

▶ Experiments
저자들은 LoftQ framework 을 NLU, QA, Summarization, NLG 태스크들에 광범위하게 적용해본다. 그 결과, 4-bit quantization 에서 XSum 에서 1.1, CNN/DailyMadil 에서 0.8 gain 을 얻었다. LoftQ 는 특히 low-bit scenario 에서 효과적인데, 2-bit Normal float 과 2-bit uniform quantization 환경에서, MNLI 에서 8%, SQuAD1.1 에서 10% gain 을 얻었다.

2. Background

2.1. Transformer Models

Multi-head Attention (MHA) + Feed Forward Network (FFN)

※ 논문참고

2.2. Qunatization

Quantization
N-bit quantization : 32-bit floating point number 같은 high-preicision number $X^{HP} \in \mathbb{R}$가 주어졌을 때, N-bit integer $X^{INT} \in \mathbb{R}$ 로 변환하는 것이다.

image

$F(\cdot):\mathbb{R} -> [0,1]$ 은 normalization function 이다. Uniform Quntization 은 $F(X) = (X-X_{min})/(X_{max}-X_{min})$ 이다. QLoRA 에서는 4-bit NormlaFLoat Quantization (NF4) 방법을 제안한다. 이 것은 $X ~ N(0, \sigma^2)$ 을 가정하여 $F(X) = \Phi(X/\sigma)$, where $\Phi(\cdot)$ is cumulative distribution 가 된다.

Dequantization
아래의 Lookup Table $T$ 를 활용하여,

image

$X^{INT}$ 를 high preicision counterpart $X^D \in \mathbb{R}$ 로 변환한다. 따라서, dequantization 은 아래와 같이 표현된다.

image

Simulated Quantization for Matric
Matrix Multiplication 을 quantized representation 으로 direct 하게 적용하는 방법도 가능하다. 이를 simulated quantization for matrices 라고 하고, quantized weight matrix 들이 encoded integer 로 저장이 되고, high-precision matrix 을 simulate 하기 위해 dequantized 되어 활용된다. Simulated quantization 을 위해서는 high-precision matrix 부터 simulated high-precision amtrix 로의 mapping 만 필요하다.

image

2.3. Low-rank Adaptation

Low-Rank Adaptation (LoRA) 는 small weight matrix $A$ 와 $B$ 를 frozen pre-trained weight matrix $W$ 에 붙인다. 따라서 linear trasnformation $Y=XW$ 가 아래의 식으로 reformulate 된다.

image

$A$ 와 $B$ 의 init 은 pre-trained weight 과의 align 을 위해서이고, fine-tuning 때는 $W$는 fixed 된 채, $A$ 와 $B$ 만 SGD type 의 optimization method 를 통해 update 된다.

중요한 것은 만약 $A$ 와 $B$ 가 quantized backbone $Q=q_N (W)$ 에 붙여진다면, 위의 initialization 을 통한 $Q+AB^T$는 더 이상 pre-trained weight $W$ 와 같지 않아 discrepancy 가 생긴다.

3. METHOD : LoftQ (LoRA-Fine-Tuning-aware Quantization)

3.1. LoRA-Aware Quantization

$N$-bit quantized weight $Q \in \mathbb{R}_N^{d_1 \times d_2}$ 와 low-rank approximation $A \in \mathbb{R}^{d_1 \times r}$, $A \in \mathbb{R}^{d_2 \times r}$ 을 활용하여, original high-precision pre-trained weight $W \in \mathbb{R}^{d_1 \times d_2}$ 를 LoRA fine-tuning 의 initialization 으로 approximate 한다. 즉, Fine-tuning 전에 아래의 objective 를 최소화하게 network 를 initialze 한다.

image

여기서 $|| \cdot ||_F$ 는 Frobenious norm 이다. 이 objective 는 low-rank adapter $A$, $B$ 에 더불어, qunatized backbone $Q$ 의 init value 를 동시에 optimize 하여, 추후 LoRA fine-tuning 을 고려한 설계이다. 기존의 방법에서는 추후 LoRA fine-tuning 을 무시한채 $W$ 를 $Q$ 로 바꾸는 것만 신경썼고, 이러한 것은 notable degradation 을 불러온다.

3.2. Alternating optimization

저자들은 위의 objective 를 quantization 과 Singular value decomposition (SVD)를 번갈아가며(alternating) 최소화 문제를 푼다.

Quantization
$t$ 번 째 step 의 quantization 은 아래와 같다.

image

$q_N ( \cdot )$은 여러 quantization function 이 가능한데, QLoRA 와 같이 NF4 를 적용하였다.

SVD
$t$ 번째 quantization step 이후, SVD 를 적용한다. Quantization Residual $R_t = W - Q_t$ 에 대해,

image image

로 SVD 를 적용한다. 이후, $A$ 와 $B$의 rank-$r$ approximation 을 $R_t$로 부터 얻는다.

image

지금까지의 과정을 아래의 알고리즘으로 정리할 수 있다.

image

$T=1$ 일 때는 QLoRA 의 $Q_1$과 정확히 일치한다. $T=1$ 만으로도 quantization discrepancy 를 줄이는데 효과적이지만, (즉 QLoRA 도 효과적이지만), alternating optimization 방법이 pre-trained weight $W$ 와 더 가까운 initialization 을 제공하여 성능 향상이 있음을 추후에 보인다.

3.3. Applying to LORA Fine-tuning

LoRA fine-tuning 때는 integer weight 은 고정하고 low-rank adapter 만 AdamW 로 학습한다. Forward pass 에서, interger weight 은 lookup table 을 통해 dequantization 이 된다. Backward pass 에서, gradient 와 optimizer 는 low-rank adapter $A$, $B$ 에만 적용된다.

4. Experiments

Quantization Methods

  • Uniform quantization
  • 4-bit NF4 (Gaussian quantization)
  • 2-bit NF2 (Gaussian quantization)

Baselines

  • Full fine-tuning
  • Full precision LoRA (LoRA)
  • QLoRA

4.1. Encoder-only Model : DEBERTa-v3

Models and Datasets

  • Model : DeBERTaV3-base
  • Benchmark : GLUE (w/o WNLI), SQuADv1.1, ANLI

Implementation Details

  • Learning Rates : {1e-5, 5e-5, 1e-4, 5e-4}
  • Quantize Entire Backbone and quantize the embedding layer for higher compression efficiency

Main Results
image

image

  • Table1 은 NF2 에 대한 결과, Table2 는 2-bit Uniform Quantization 에 대한 결과.
  • 모든 rank, qunatization method, dataset 에 대해 QLoRA 보다 좋은 성능을 보인다.
  • Table2 의 MNLI-m 에서 88.0% 정확도를 달성하여, QLoRA 를 8% 이긴다.
  • NF2 의 SST 와 SQuAD 에서 full fine-tuning 과 유사한 결과를 보인다.
  • 2-bit 에서 QLoRA 는 COLA 에서 실패하는데 비해, LoftQ는 60.5 로 높은 수치를 기록한다.

4.2. Encoder-Decdoer Model : BART

Models and Datasets

  • Model : BART-large
  • Benchmark : Summarization taskXSum, CNN/DailyMail
  • Metric : ROUGE 1/2/L

Main Results
image

  • QLoRA 를 NF4 와 Uniform 에 대해 rank-8, rank-16 에서 모두 앞선다.
  • 심지어 XSum 에서 Full precision 보다도 더 좋은 성능을 보인다. 이에 대한 분석은 뒤에서 진행한다.

image

  • NF2 quantization 에 대해, QLoRA 는 전혀 성능을 내지 못하지만, LoftQ 는 좋은 성능을 보인다.

4.3. Decoder-only Model : LLaMA-2

Models and Datasets

  • Model : LLaMA2-7b, LLaMA-2-13b
  • Benchmark : NLG taskGSM8K, WikiText-2
  • Metric : Accuracy for GSM8K, perplexity for WikiText-2

Main Results
image

  • WikiText-2 에서 모든 setting 에서 QLoRA 보다 좋은 성능을 보인다.
  • 역시 2-bit 에서 QLoRA 는 생성에 실패하지만, 7.85 ppl 을 달성한다.
  • GSM8K 에서도 QLoRA 는 생성에 실패하지만, 26.5% acc 를 달성한다.
  • Mixed-precision quantization scenario 에서 LoftQ 의 포테셜을 확인할 수 있다.

4.4. Analysis

Effectiveness of Alternating Optimization
Alternating optimization step $T$ 를 달리하며 실험 분석을 해본다. 앞서 말했듯, $T=1$ 일 때, QLoRA 와 동일하다. 모든 task 와 model 에 대하여, minimal alternating step 만으로 주효한 성능 향상이 있다. 이는 Quantized weight 과 original weight 사이의 discrepancy 를 rapid 하게 줄인다.

흥미롭게도, alternating step 이 너무 높으면 성능이 오히려 약간 낮아지는데 ($T=10$ 에서 MNLI, 그리고 XSUM 에서 $T$들) gap 이 작아질 수록 alternating step 이 gap 을 minimize 하는데 어려움을 겪기 때문이라고 분석한다.

image

5. Discussion

Start with quantization or SVD in the alternating optimization?
LoftQ 는 quantization -> SVD 순서로 alteranting optimization 이 구성되는데, SVD -> Quantization 으로 바꾸면,

image

아래와 같이 여전히 SVD 를 먼저해도 좋은 결과지만, 원래대로 Quantization 을 먼저하는 것이 조금 더 좋은 성능을 보인다.

LoftQ better than Full-precision LoRA?
Table3 와 Table5 에서, XSUM 에 대해 Full-precision LoRA 보다 LoftQ 가 더 좋았다. 저자들은 LoftQ 의 low-rank adapter 가 non-zero init 이고, Full-precision LoRA 는 zero-init 이기 때문에, 이러한 unexpected phenomenon 이 일어난다고 분석한다. 이 zero initialization 이 fine-tuning 을 unstable 하게 한다는 분석이다.

Conclusion

We propose LoftQ, a quantization framework for LLMs, which alternatively applies quantization and low-rank approximation to the original high-precision pre-trained weights, to obtain an initialization for the subsequent LoRA fine-tuning. Experiments on natural language understanding, question answering, summarization, and natural language generation show that our framework remarkably surpasses existing methods, e.g., QLoRA, for quantizing encoder-only, encoder-decoder, and decoder-only models. We have not observed our method exhibiting worse performance over QLoRA. Moreover, our quantization framework demonstrates effectiveness and robustness particularly in low-bit quantization regimes, e.g., the 2-bit level.