本文是一篇專(zhuān)注于減輕大型Transformer語(yǔ)言模型訓(xùn)練代價(jià)的工作。作者提出了一種基于低秩更新的ReLoRA方法。過(guò)去十年中深度學(xué)習(xí)發(fā)展階段中的一個(gè)核心原則就是不斷的“堆疊更多層（stack more layers），因此作者希望探索能否同樣以堆疊的方式來(lái)提升低秩適應(yīng)的訓(xùn)練效率，實(shí)驗(yàn)結(jié)果表明，ReLoRA在改進(jìn)大型網(wǎng)絡(luò)的訓(xùn)練方面更加有效。

論文鏈接： https://arxiv.org/abs/2307.05695 代碼倉(cāng)庫(kù)： https://github.com/guitaricet/peft_pretraining

一段時(shí)間以來(lái)，大模型（LLMs）社區(qū)的研究人員開(kāi)始關(guān)注于如何降低訓(xùn)練、微調(diào)和推理LLMs所需要的龐大算力，這對(duì)于繼續(xù)推動(dòng)LLMs在更多的垂直領(lǐng)域中發(fā)展和落地具有非常重要的意義。目前這一方向也有很多先驅(qū)工作，例如從模型結(jié)構(gòu)上創(chuàng)新的RWKV，直接替換計(jì)算量較大的Transformer架構(gòu)，改用基于RNN范式的新架構(gòu)。還有一些方法從模型微調(diào)階段入手，例如在原有LLMs中加入?yún)?shù)量較小的Adapter模塊來(lái)進(jìn)行微調(diào)。還有微軟提出的低秩自適應(yīng)（Low-Rank Adaptation，LoRA）方法，LoRA假設(shè)模型在任務(wù)適配過(guò)程中對(duì)模型權(quán)重的更新量可以使用低秩矩陣進(jìn)行估計(jì)，因而可以用來(lái)間接優(yōu)化新加入的輕量級(jí)適應(yīng)模塊，同時(shí)保持原有的預(yù)訓(xùn)練權(quán)重不變。目前LoRA已經(jīng)成為大模型工程師必備的一項(xiàng)微調(diào)技能，但本文作者仍然不滿(mǎn)足于目前LoRA所能達(dá)到的微調(diào)效果，并進(jìn)一步提出了一種可疊加的低秩微調(diào)方法，稱(chēng)為ReLoRA。

本文來(lái)自馬薩諸塞大學(xué)洛厄爾分校的研究團(tuán)隊(duì)，作者團(tuán)隊(duì)將ReLoRA應(yīng)用在具有高達(dá)350M參數(shù)的Transformer上時(shí)，展現(xiàn)出了與常規(guī)神經(jīng)網(wǎng)絡(luò)訓(xùn)練相當(dāng)?shù)男阅?。此外，本文作者還觀察到ReLoRA的微調(diào)效率會(huì)隨著模型參數(shù)規(guī)模的增加而不斷提高，這使得其未來(lái)有可能成為訓(xùn)練超大規(guī)模（通常超過(guò)1B參數(shù)）LLMs的新型手段。

一、引言雖然目前學(xué)術(shù)界和工業(yè)界都在不斷推出自家的各種基座模型，但不可否認(rèn)的是，完全預(yù)訓(xùn)練一個(gè)具有初等推理能力的LLMs仍然需要非常龐大的算力，例如大家熟知的LLaMA-6B模型[1]就需要數(shù)百個(gè)GPU才能完成訓(xùn)練，這種規(guī)模的算力已經(jīng)讓絕大多數(shù)學(xué)術(shù)研究小組望而卻步了。在這種背景下，參數(shù)高效微調(diào)（PEFT）已經(jīng)成為了一個(gè)非常具有前景的LLMs研究方向。具體來(lái)說(shuō)，PEFT方法可以在消費(fèi)級(jí)GPU（例如RTX 3090或4090）上對(duì)十億級(jí)語(yǔ)言或擴(kuò)散模型進(jìn)行微調(diào)。因此本文重點(diǎn)關(guān)注PEFT中的低秩訓(xùn)練技術(shù)，尤其是LoRA方法。作者思考到，過(guò)去十年中深度學(xué)習(xí)發(fā)展階段中的一個(gè)核心原則就是不斷的“堆疊更多層（stack more layers）”，例如ResNet的提出可以使我們將卷積神經(jīng)網(wǎng)絡(luò)的深度提升到100層以上，并且也獲得了非常好的效果。因此本文探索能否同樣以堆疊的方式來(lái)提升低秩適應(yīng)的訓(xùn)練效率呢？

本文提出了一種基于低秩更新的ReLoRA方法，來(lái)訓(xùn)練和微調(diào)高秩網(wǎng)絡(luò)，其性能優(yōu)于具有相同可訓(xùn)練參數(shù)數(shù)量的網(wǎng)絡(luò)，甚至能夠達(dá)到與訓(xùn)練100M+規(guī)模的完整網(wǎng)絡(luò)類(lèi)似的性能，對(duì)比效果如上圖所示。具體來(lái)說(shuō)，ReLoRA方法包含（1）初始化全秩訓(xùn)練、（2）LoRA 訓(xùn)練、（3）參數(shù)重新啟動(dòng)、（4）鋸齒狀學(xué)習(xí)率調(diào)度（jagged learning rate schedule）和（5）優(yōu)化器參數(shù)部分重置。作者選擇目前非?；馃岬淖曰貧w語(yǔ)言模型進(jìn)行實(shí)驗(yàn)，并且保證每個(gè)實(shí)驗(yàn)所使用的GPU計(jì)算時(shí)間不超過(guò)8天。二、本文方法作者首先從兩個(gè)矩陣之和的秩入手，通常來(lái)說(shuō)，矩陣相加的后秩的上界會(huì)比較緊湊，對(duì)于矩陣，，然后存在矩陣，，使得矩陣之和的秩高于或。作者希望利用這一特性來(lái)制定靈活的參數(shù)高效訓(xùn)練方法，然后從LoRA算法開(kāi)始入手，LoRA可以將模型權(quán)重的更新量 分解為一組低秩矩陣乘積 ，如下式所示，其中 是固定縮放因子。

在具體操作時(shí)，LoRA通常是加入新的可學(xué)習(xí)參數(shù) 和 來(lái)實(shí)現(xiàn)，這些參數(shù)可以在訓(xùn)練后合并回原始參數(shù)中。因此，盡管上述方程允許我們?cè)谟?xùn)練階段中實(shí)現(xiàn)總參數(shù)更新量高于任意單個(gè)矩陣的秩，但其仍然受到的限制。因此作者想到通過(guò)不斷疊加這一過(guò)程來(lái)突破限制達(dá)到更好的訓(xùn)練效果。這首先需要對(duì)LoRA過(guò)程進(jìn)行重新啟動(dòng)，就可以在訓(xùn)練階段不斷合并每次得到的 和 來(lái)得到累加的權(quán)重更新量，計(jì)算公式如下：但是，想要對(duì)已經(jīng)完成的LoRA過(guò)程重新啟動(dòng)并不容易，這需要對(duì)優(yōu)化器進(jìn)行精細(xì)的調(diào)整，如果調(diào)整不到位，會(huì)導(dǎo)致模型在重啟后立即與之前的優(yōu)化方向出現(xiàn)分歧。例如Adam優(yōu)化器在更新時(shí)主要由先前步驟中所累積梯度的一階矩和二階矩引導(dǎo)。實(shí)際上，梯度矩平滑參數(shù) 和 通常非常高，因而在重新啟動(dòng)時(shí)的秩上界為 ，相應(yīng)的梯度矩 和 都是滿(mǎn)秩的，在合并參數(shù)后就會(huì)使用先前的舊梯度來(lái)優(yōu)化 朝向與 相同的子空間方向。 為了解決這個(gè)問(wèn)題，作者提出了ReLoRA方法，ReLoRA在合并和重新啟動(dòng)期間可以對(duì)優(yōu)化器進(jìn)行部分重置，并在隨后的預(yù)熱中過(guò)程中將學(xué)習(xí)率設(shè)置為0。具體來(lái)說(shuō)，作者提出了一種鋸齒狀學(xué)習(xí)率調(diào)度算法，如下圖所示，在每次對(duì)ReLoRA參數(shù)進(jìn)行重置時(shí)，都會(huì)將學(xué)習(xí)率設(shè)置為零，并執(zhí)行快速（50-100 步）學(xué)習(xí)率預(yù)熱使其回到與重置前相同的水平范圍內(nèi)。

ReLoRA通過(guò)序列疊加的方式僅訓(xùn)練一小組參數(shù)就可以實(shí)現(xiàn)與全秩訓(xùn)練相當(dāng)?shù)男阅?，并且遵循LoRA方法的基礎(chǔ)原則，即保持原始網(wǎng)絡(luò)的凍結(jié)權(quán)重并添加新的可訓(xùn)練參數(shù)。乍一看，這種方式可能顯得計(jì)算效率低下，但我們需要清楚的是，這種方法可以通過(guò)減小梯度和優(yōu)化器狀態(tài)的大小，來(lái)顯著提高顯存效率。例如Adam優(yōu)化器狀態(tài)消耗的顯存通常是模型權(quán)重占用的兩倍。通過(guò)大幅減少可訓(xùn)練參數(shù)的數(shù)量，ReLoRA可以在相同的顯存條件下使用更大的batchsize大小，從而最大限度地提高硬件效率，ReLoRA的整體操作細(xì)節(jié)如下圖所示。

三、實(shí)驗(yàn)效果為了清晰的評(píng)估ReLoRA方法的性能，作者將其應(yīng)用在各種規(guī)模大?。?0M、130M、250M 和 350M）的Transformer模型上，并且都在C4數(shù)據(jù)集上進(jìn)行訓(xùn)練和測(cè)試。為了展現(xiàn)ReLoRA方法的普適性，作者重點(diǎn)考察NLP領(lǐng)域的基礎(chǔ)語(yǔ)言建模任務(wù)。模型架構(gòu)和訓(xùn)練超參數(shù)設(shè)置基本與LLaMA模型保持一致。與LLaMA不同的是，作者在實(shí)驗(yàn)中將原始的注意力機(jī)制（使用float32進(jìn)行 softmax計(jì)算）替換為了Flash注意力[2]，并且使用bfloat16精度進(jìn)行計(jì)算，這樣操作可以將訓(xùn)練吞吐量提高50-100%，且沒(méi)有任何訓(xùn)練穩(wěn)定性問(wèn)題。此外，使用ReLoRA方法訓(xùn)練的模型參數(shù)規(guī)模相比LLaMA要小得多，最大的模型參數(shù)才僅有350M，使用8個(gè)RTX4090上訓(xùn)練了一天時(shí)間就可以完成。 下圖展示了本文方法與其他方法的性能對(duì)比效果，可以看到ReLoRA顯著優(yōu)于低秩LoRA方法，證明了我們提出的修改的有效性。此外，ReLoRA還實(shí)現(xiàn)了與滿(mǎn)秩訓(xùn)練（Full training）相當(dāng)?shù)男阅?，并且我們可以觀察到，隨著網(wǎng)絡(luò)規(guī)模的增加，性能差距逐漸縮小。有趣的是，ReLoRA 唯一無(wú)法超過(guò)的基線(xiàn)模型是僅具有60M參數(shù)的最小模型。這一觀察結(jié)果表明，ReLoRA在改進(jìn)大型網(wǎng)絡(luò)的訓(xùn)練方面更加有效，這與作者最開(kāi)始研究探索一種改進(jìn)大型網(wǎng)絡(luò)訓(xùn)練方法的目標(biāo)是一致的。

此外，為了進(jìn)一步判斷ReLoRA是否能夠通過(guò)迭代低秩更新來(lái)實(shí)現(xiàn)相比LoRA更高的秩更新訓(xùn)練，作者繪制了ReLoRA、LoRA和全秩訓(xùn)練的熱啟動(dòng)權(quán)重與最終權(quán)重之間差異的奇異值譜。如下圖所示，下圖說(shuō)明了LoRA和ReLoRA之間對(duì)于 、、 和 奇異值的顯著差異，可以看到ReLoRA在所有四個(gè)矩陣參數(shù)上均得到了最小的奇異值。

四、總結(jié)本文是一篇專(zhuān)注于減輕大型Transformer語(yǔ)言模型訓(xùn)練代價(jià)的工作，作者選取了一條非常具有前景的方向，即低秩訓(xùn)練技術(shù)，并且從最樸素的低秩矩陣分解 (LoRA) 方法出發(fā)，利用多個(gè)疊加的低秩更新矩陣來(lái)訓(xùn)練高秩網(wǎng)絡(luò)，為了實(shí)現(xiàn)這一點(diǎn)，作者精心設(shè)計(jì)了包含參數(shù)重新啟動(dòng)、鋸齒狀學(xué)習(xí)率調(diào)度算法和優(yōu)化器參數(shù)重置等一系列操作，這些操作共同提高了ReLoRA算法的訓(xùn)練效率，在某些情況下甚至能夠達(dá)到與全秩訓(xùn)練相當(dāng)?shù)男阅?，尤其?shí)在超大規(guī)模的Transformer網(wǎng)絡(luò)中。作者通過(guò)大量的實(shí)驗(yàn)證明了ReLoRA的算法可行性和操作有效性，不知ReLoRA是否也會(huì)成為大模型工程師一項(xiàng)必備的算法技能呢？