LoRA:語言模型微調的計算資源最佳化策略

發表於2024-02-11

編者按:隨著資料量和計算能力的增加,大模型的引數量也在不斷增加,同時進行大模型微調的成本也變得越來越高。全引數微調需要大量的計算資源和時間,且在進行切換下游任務時代價高昂。

本文作者介紹了一種新方法 LoRA,可以在保持模型效能的同時大幅減少微調的引數量和所需資源。

LoRA透過引入兩個低秩適配矩陣,用矩陣乘法的方法替換大部分引數。實驗證明,LoRA 在多項 NLP 任務上的表現與許多微調方法(如Adapter 和 PreLayer 等)相當或更好。與全引數微調相比,LoRA降低了可訓練引數數量 10,000 倍,GPU 記憶體需求減少 3 倍,儲存需求減少 10,000 倍,訓練速度提高 25 %。

LoRA 為大語言模型的高效多工微調提供了一種有效途徑。作者認為 LoRA 可以推廣到更多模型結構,有望加深我們對模型微調機制的理解。

作者 | Arxiv Dives

編譯 | 嶽揚

🚢🚢🚢歡迎小夥伴們加入AI技術軟體及技術交流群,追蹤前沿熱點,共探技術難題~

一、  背景知識

Paper: https://arxiv.org/abs/2106.09685

Published: October 16th, 2021, by Microsoft and CMU

這篇文章涉及的數學內容較多,但幸運的是,涉及的線性代數內容僅涉及基礎的加法和乘法運算,我相信我們都能夠理解。

首先,我將簡要概述相關的數學原理及其原因,接下來,我們將深入探討論文的細節,以及它們如何應用於 GPT-2 和 GPT-3 等 transformers 模型。

最重要的一點是:LoRA 減少了可訓練引數(trainable parameters)的數量,從而減少了訓練時間和 GPU 記憶體的使用量,同時保持了輸出的質量。

LLM(顧名思義)的規模非常大。用於微調的資料集(fine-tuning datasets)通常比模型的預訓練資料集小得多。當資料集比較小的時候,LoRA 只需更新較少的權重,這即是 LoRA 的優勢所在。

二、LoRA 的工作原理

如果你熟悉矩陣乘法,那麼應該知道 AxM 矩陣和 MxB 矩陣相乘得到的結果是一個 AxB 矩陣。

https://www.youtube.com/watch?app=desktop&v=2spTnAiQg4M

假設在神經網路中有一個 MxM 的預訓練密集層(pre-trained dense layer)(權重矩陣)W。

例如,這個 Keras 模型有 3 個 size 為 512x512 的密集層(dense layers):

然後再初始化兩個密集層 A 和 B,它們的 shapes 分別為 M x R 和 R x M。

R(秩)遠遠小於 M。研究表明,R 的取值在1和4之間效果較好。

所以,舉個例子,假設密集層擁有 512x512= 262,144 個引數。

因此,可以有一個 size 為 512x4 和一個 size 為 4x512 的矩陣,每個矩陣只有2048個引數,總共4096個引數。

密集層的原始方程式為:

Y = Wx + b

LoRA 將其修改為:

Y = Wx + b + BAx

其中,x 是一個 512x1 的向量,是神經網路的輸入,b 是一個 512x1 的偏置向量。

矩陣乘法的數學公式如下:

  Dimensions of each variable:
  W = 512x512
  x = 512x1
  b = 1x512
  B = 512x4 (New params)
  A = 4x512 (New params)
  
  Dimensions fully laid out:
  Y = (512x512) * (512x1) + (1x512) + (512x4) * (4x512) * (512x1)

但在這種情況下,我們只訓練 A 和 B 兩個矩陣,每個矩陣只有2048個引數。因此,透過使用LoRA方法,可將可訓練引數的數量從 262,144 減少到 4,096 個。

三、可以最佳化神經網路的哪些部分?

在訓練/執行神經網路時,我們需要考慮哪些部分可以進行最佳化?

1.  總體模型大小(Total model size)

● 模型的磁碟佔用空間,使用 serverless 時,透過網路傳輸模型所需的模型大小,需要佔用的 RAM 大小,需要佔用的 GPU 大小,需要佔用的 CPU 大小

2.  推理時的batch size(Inference batch size)

● batch size,序列長度(sequence length),data size

3.  訓練所需的記憶體

● 所有模型引數 + 可訓練引數的梯度

如果你還記得反向傳播演算法(backpropagation)的工作原理,你需要計算每個偏導數並將它們儲存在記憶體中,以便進行反向傳播。這意味著對於傳統的全引數微調,所需的記憶體使用量將增加一倍。

https://soumya997.github.io/2022-03-20-pytorch-params/

LoRA 只訓練秩分解矩陣(rank decomposition matrices)(A和B),從而減少了總體訓練所需的記憶體。

這些少量的介面卡(adapter)權重可以合併到實際模型本身中,因此它們不會影響推理過程或總體模型的大小。

五、為什麼沒有額外的推理時間?

LoRA 的原始方程為:

Y = Wx + b + BAx

由於加法的傳遞性質,我們可以將其重新表達為:

Y = Wx + BAx + b

或者將 x 因數合併為

Y = (W + BA)x + b

也就是說,我們可以簡單地將 (W + BA) 替換為新的 W1,然後恢復原來的線性方程。

W1 = (W + BA)

我們用一組新的權重恢復到原來的方程:

Y = W1*x + b

這就意味著,如果我們將原始模型和介面卡(adapter)的權重合並,我們所進行的計算基本上與原始模型相同!

五、深入探究這篇論文

目前自然語言處理的正規化是先在大量通用資料上進行預訓練,然後再對一項或多項特定任務進行微調。對於大模型來說,全面微調所有引數的成本變得過高。

以 GPT-3 為例,其擁有 175B 個引數,這意味著現在需要加倍儲存所有梯度來進行訓練,更不用說如果要儲存多個微調後的模型,還需要將每個模型的全套引數儲存下來。

LoRA 可以將可訓練引數的數量減少 10,000 倍,GPU 記憶體的需求減少 3 倍。

在實際應用中,記憶體使用量的減少程度取決於模型的大小。

https://keras.io/examples/nlp/parameter_efficient_finetuning_...

雖然 LoRA 具有更高的訓練吞吐量,而且沒有額外的推理延遲,但其效能與微調(fine-tuning)相當,甚至更好。

5.1 Introduction

自然語言處理中的許多應用依賴於將一個大引數量的通用模型適應於多個下游應用。

例如,可能有一個通用模型,可以用最常見的 next words 完成大量的英語句子。人類語言的一個問題在於對於同一個句子可能有多種有效的延續方式。

想想人們對不同的話題有多少不同的看法。其中很多觀點都基於他們的過去經驗,人們在討論和交流觀點時經常會產生分歧和辯論。

例如,你希望下游模型能夠用你的聲音總結文字,或者能夠將自然語言翻譯成 SQL 查詢語句,或者讓微調後的模型比基礎模型更有趣,這些都可以透過微調來實現。

對整個模型進行端到端(end to end)微調的一個缺點是,新模型包含的引數和舊模型一樣多。如果要進行 N 次微調,這就意味著每個新模型的儲存佔用空間和記憶體都要線性增加。

一些人透過為新任務學習外部模組或在模型中新增新層來解決了這個問題,但這會增加推理延遲。

https://pyimagesearch.com/2019/06/03/fine-tuning-with-keras-a...

5.2 LoRA 中的“Rank”指的是什麼?

他們強調了一個事實,即過引數化(over-parameterized)模型(引數量遠超訓練資料量的模型)實際上具有相對簡單的內在空間維度,並假設在模型適應或調整過程中,權重的變化具有“低內在秩(low instrinsic rank)” (譯者注:即其權重調整過程可能並不需要在所有這些維度上進行。實際上,有效的權重調整可能僅僅發生在一個相對較小的子空間內)。

矩陣的“秩(rank)”是指其中線性無關的列或行的數量。

我們可以將神經網路中的線性無關性(linear independence)理解為“每組權重對決策的影響有多大”。

秩為零的矩陣將是一個全零矩陣。

如果你有一個看起來像這樣的矩陣:

1 2 3 4
2 4 6 8
5 3 9 7

我們可以看到這個矩陣前兩行其實是彼此的倍數,所以它們會繼續指向相同的方向。但是,第三行則帶我們走向完全不同的方向。

在下圖中,秩為2的矩陣將形成一個二維平面,因為所有向量都在同一平面上,而秩為3的矩陣則更像一個立方體,因為每個向量指向不同的方向。

https://peterbloem.nl/blog/pca-4?ref=blog.oxen.ai

神經網路的維度通常非常高,而我們的大腦很難想象或理解這麼高維度的空間。

即使原始資料的維度非常高(高達12,228),使用低秩(甚至是1或2)也能夠有效地表示資料。

這種技術的優勢包括:

● 可以共享同一個預訓練模型,也可以為不同的任務構建許多更小的 LoRA 模型。

● LoRA 提高了訓練效率,降低了硬體門檻。

● 簡單的線性設計允許權重可合併,不會帶來推理延遲。

● LoRA 可應用於許多模型架構和先前的方法,因為它是一個簡單的密集層(dense layer)。

在本例中將 LoRA 應用於 Transformer 架構。因此,下一節一起來了解有哪些變數,分別代表著什麼。

5.3 現有解決方案還不夠好嗎?

論文承認,這絕不是一個新問題。遷移學習(Transfer learning)有多種方法提高模型自適應(Model Adaptation)的效率,包括引數和計算的效率。。

論文作者特別關注了在模型自適應(Model Adaptation)過程中使用 adapter layers 以及最佳化輸入層(input layers)或 prompts 的方法。

adapter layers 雖然體積小,但必須按順序處理,而不是並行處理,因此會增加額外的延遲。

他們對 A 從高斯分佈中隨機選擇初始值來初始化矩陣中的元素,並將矩陣 B 初始化為零,這樣在訓練開始時,矩陣 B 乘以矩陣 A 的結果為零。

論文指出,LoRA 可以進行全引數微調,因為一旦將 LoRA 的秩設為預訓練權重矩陣的秩,LoRA會逐漸收斂到與原始模型相似的狀態,從而實現對原始模型的訓練。

當模型部署到生產環境中時,可以進行 W = W + BA 的計算並將結果儲存下來,然後像往常一樣進行推理。當需要將微調後的模型切換到另一個任務時,可以透過簡單的數學操作來恢復原始的模型權重W,而且不會佔用太多額外的記憶體空間。

5.4 將 LoRA 應用於 Transformers

論文只將 LoRA 應用於注意力機制(attention mechanism)中的查詢(Query)、鍵(Key)和值(Value)權重矩陣,而不將其應用於其他密集層(dense layers)。

本文還將 LoRA 應用於 Transformer 的其他部分留作 “未來的工作” 。

5.5 實際優勢 Practical Benefits

LoRA帶來的最顯著好處是記憶體和儲存空間的減少。

對於 GPT-3 175B,該技術將 VRAM 消耗從 1.2TB 減少到 350GB。這是相當驚人的,訓練 GPT-3 需要如此大量的 VRAM。

由於 r=4 且僅調整 Q 和 V 矩陣,微調(fine-tuning)過程中生成的 checkpoint 大小約為 35MB,而不是 350GB......因此減少了 10,000 倍。

請注意,在模型的部署過程中仍然需要 350GB 儲存空間,但儲存 100 個微調後的模型僅需要 350GB + 35MB  100 ≈ 354GB,而不是 100  350GB ≈ 35TB。

這樣就可以建立許多定製模型,並且可以在儲存預訓練權重(pre-trained weights)的機器上隨時切換模型。

論文作者還觀察到,與全引數微調相比,訓練速度提高了 25%,因為不需要為絕大多數引數計算梯度。

5.6 根據經驗和實驗證明

論文評估了 LoRA 在 RoBERTa、DeBERTa 和 GPT-2 上的下游任務效能,然後將其擴充套件到了 GPT-3。

論文在多個基準測試上進行了評估:

● GLUE(General Language Understanding Evaluation),這個基準測試包含許多子任務。

● WikiSQL

● SAMsum(conversation summarization)

● Baselines

我們不會對所有 adapters 進行比較,但我們可以看到,在 RoBERTa 上使用 LoRA 時,LoRA 甚至比許多 adapters 更有競爭力。

即使可訓練引數的數量要小得多,但是當應用於 GPT-2 時,效能超過了許多 Adapter 和 PreLayer 方法。

在上述任務中,GPT-3 也是如此。GPT-3 的執行成本要高得多,所以作者會更加謹慎地選擇對 GPT-3 進行評估的基準測試,以避免過高的執行成本。

5.7 什麼是 prompt engineering?

他們承認 prompt engineering 可以用來最大限度地提高通用模型在特定任務中的效能,並指出微調 GPT-3 與 prompt engineering 進行比較並不十分科學,因此在 prompting 和微調之間還沒有進行過很多比較。

此外,David在談話中提到,prompt engineering 比全引數微調更不穩健,更容易受到提示語注入攻擊(prompt injection hacks)的影響。

5.8 使用多少秩(Rank)以及應用於哪些權重(weights)?

在評估過程中,較低秩的模型表現優於較高秩的模型,這可能令人感到驚訝。

5.9 Subspace Similarity

(譯者注:"Subspace Similarity"指的是線上性代數中,兩個向量空間之間的相似性。)他們使用奇異值分解來研究模型中不同子空間之間的相似性。他們發現,在觀察模型中的不同子空間時,直到維度為1時,這些子空間之間的相似性仍然很高。這可能對於理解為什麼較低秩的模型表現更好具有重要意義。

下面是對列向量之間子空間相似性的另一個視覺化圖示。可以看到,很多列向量的值都接近於零,這意味著它們非常相似,只有那些 top ranks 的列向量才會顯示出差異。(譯者注:"top ranks" 可能指的是奇異值分解中最大的奇異值對應的向量。)

這些關於權重的研究提出了一個問題:如果有如此多的引數是線性相關的,那麼一般來說,大語言模型到底需要多少引數呢?

六、 結論和未來展望

對大語言模型進行微調的成本過高,尤其是在需要切換不同任務時。

LoRA有助於降低訓練成本,並實現快速任務切換。

由於 LoRA 是一種與模型架構無關的技術,因此可以與許多其他方法和模型結合使用。

微調或 LoRA 背後的機制尚不清楚,他們認為透過研究矩陣的秩,可以更容易地理解 LoRA 的工作機制,相較於全引數微調的方法。

他們認為模型的許多其他部分也可以應用 LoRA,而且他們的許多設定都是基於啟發式方法(heuristics)選擇的。

作者在微調 Llama 時使用 LoRA 不會增加硬體要求,同時我們也看到很多人將其應用於 stable diffusion(用於影像生成),我認為很多雲服務可能也都在使用 Lora 來提高不同任務的準確性。

本文經原作者授權,由Baihai IDP編譯。如需轉載譯文,請聯絡獲取授權。

原文連結:

https://blog.oxen.ai/arxiv-dives-how-lora-fine-tuning-works/