賈佳亞韓松團隊新作:兩行代碼讓大模型上下文窗口倍增 | GitHub熱榜
只要兩行代碼+11個小時微調,就能把大模型4k的窗口長度提高到32k。
規模上,最長可以擴展到10萬token,一口氣就能讀完長篇小說的多個章節或中短篇小說。
賈佳亞韓松聯合團隊提出的這個基于LoRA的全新大模型微調方法,登上了GitHub熱榜。
這種方式叫做LongLoRA,由來自香港中文大學和MIT的全華人團隊聯合出品。
在一臺8個A100組成的單機上,增大窗口長度的速度比全量微調快數倍。
網友看了之后不禁表示,這個效率實在是令人印象深刻:
那么,用LongLoRA微調之后,模型會有什么樣的變化呢?
一口氣讀完一部小說
研究團隊的實驗當中使用的模型是Llama 2。
經過LongLoRA方法微調之后,Llama 2-7B的窗口長度最高可提升到10萬token。
實測發現,微調后的模型可以一口氣讀完一部小說,然后回答各種問題。
比如總結一下大劉在《三體》第三部中體現的中心思想,比總結內容還高出了一個層次。
模型給出的答案是與外星文明首次接觸的危險性、星際旅行之困難與人類文明之脆弱,以及團結協作的重要性等內容。
的確每條在原著中都有所體現,而且也比較全面了。
除了對整部作品進行概括提煉,局部內容當然也可以詢問。
小說中的角色也能對答如流,比如《西游記》中孫悟空是怎么開花成長的。
模型告訴我們,孫悟空很有智慧,但又有一顆頑皮的心,在伴隨唐僧取經的過程中走向了成熟。
這次的總結依舊是很到位。
而且不僅是單個角色,不同人物之間復雜的關系也能了如指掌。
提問的方式可以簡單粗暴些,直接要求描述這本書(《哈利波特》)中的人物關系。
模型以哈利·波特為中心,介紹了他的朋友韋斯萊、赫敏,敵人馬爾福,以及鄧布利多教授等人物。
除了看小說,LongLoRA微調后的Llama還可以讀論文,生產力一下子就提高了(喜)。
無論是整體概括還是局部詢問,微調后的模型都能準確地給出答案:
△中文部分為谷歌機翻
為了從宏觀上把握模型的表現,研究團隊用了如下數據集進行了測試:
- PG19:來自書籍的長篇文檔數據集,用來測試語言建模效果。
- Proof-pile:來自arXiv的數學論文數據集,用來測試語言建模效果。
- LongQA:作者自行構建的長序列問答數據集,用于有監督的微調。
- LongChat:第三方構建的長對話理解數據集,用來測試長序列敘述理解效果。
結果顯示,LongLoRA在PG19和Proof-pile上的困惑度與全量微調接近。
在問答數據集上,LongLoRA微調出的模型表現也很優異,長文本理解方面更是達到了SOTA水平。
當然,LongLoRA的意義不僅在于提高了窗口長度,關鍵在于用更少的消耗提高了窗口長度。
以7B參數量的Llama-2為例,如果使用全量微調,從4k提升到32k,在一臺8個A100的單機上需要五天。
而改用LongLoRA方式,則只用11.3小時就能完成,連半天都不到,效率提升近十倍。
如果提升到65k,全量微調所需時間將超過1000小時,LongLoRA卻只用52.4小時。
那么LongLoRA又是怎么做到的呢?
“大而化小”降低計算量
LongLoRA建立在LoRA的基礎之上,引入了一種稱為“移位短注意力”(shift short attention)的機制。
這種機制只需要兩行代碼就能實現:
Transformer架構的核心是自注意力(Self-attention)計算。
短注意力就是將訓練文本劃分為多個組,使自注意力計算在每個組內分別進行,從而達到降低運算量的目的。
而在這一過程中注意力頭也被進行了分組,通過注意力頭的位移,就實現了組間的信息交互。
劃分出的每個組之間有重疊部分,確保了數據可以在全文中流通。
這樣一來,每次計算都只需要對組內的token進行操作,運算量大大降低。
除了對輸入進行分割之外,LongLoRA相比于Lora還可以微調embedding層和normalization層。
這兩項內容占的參數量很小,以Llama 2-7B為例,embedding層只占1.94%,normalization層更是不到十萬分之四。
消融實驗結果表明,除了核心的Attention層,這兩個占比很小的部分也起到了重要作用。
除了核心的短注意力機制,研究團隊引入了DeepSpeed和FlashAttention方式,進一步降低了訓練消耗。
目前,LongLoRA微調過后不同參數量和窗口長度的Llama 2已經開源,感興趣的話可以到GitHub頁面中查看。
論文地址:https://arxiv.org/abs/2309.12307
GitHub項目頁:https://github.com/dvlab-research/LongLoRA