SambaNova SN40L: 利用Dataflow和專家組合(COE)來克服AI記憶牆的大模型

摘要

GPT-4等整體式大型語言模型(LLM)為現代生成AI應用鋪路。然而，大規模訓練、服務及維護整體式LLM仍然極其昂貴和充滿挑戰。現代AI加速器計算能力與記憶體比例的不成比例增長已經造成了記憶體壁障，需要新的方法來部署AI。最近的研究顯示，許多小型專家模型的組合，每個模型參數量級更小，可以匹配或超越整體式LLM的能力。專家組合(CoE)是一種模組化方法，可降低訓練和服務的成本和複雜性。然而，在使用傳統硬體時，此方法呈現兩個關鍵挑戰：(1) 沒有融合操作，較小的模型具有較低的運算強度，使高利用率更具挑戰性；以及(2) 託管大量模型在動態切換它們時，可能非常昂貴或速度慢。

本文描述如何結合專家組合、流式dataflow和三層記憶體系統來克服AI記憶體壁障。我們介紹Samba-CoE，一個擁有150個專家和總計一兆參數的專家組合系統。我們將Samba-CoE部署在SambaNova SN40L可重構dataflow單元(RDU)上——一種為企業推論和訓練應用程式而共同設計的商用dataflow加速器架構。該晶片推出一個新的三層記憶體系統，具有晶片內分散式SRAM、封裝內HBM和封裝外DDR DRAM。一個專用的RDU間網路能夠跨多個插槽進行向上和向外擴展。我們展示了在8個RDU插槽上執行各種基準測試相較於非融合基線的2倍到13倍加速。我們表明，對於CoE推理部署，8插槽RDU節點相較於DGX H100和DGX A100，將機器規模縮小高達19倍、加速模型切換時間15倍到31倍，並實現3.7倍和6.6倍的整體加速。

參考論文: https://arxiv.org/pdf/2405.07518

I. 導言

近期大型語言模型(LLM)的訓練和推理進展令世人驚艷。ChatGPT和Gemini等先進的生成AI/ML應用程式建立在整體式LLM之上，可能擁有數十億或兆級參數。它們使用從網路抓取的數兆級token所組成的精心策劃資料集進行訓練。然而，訓練和服務一個先進整體式LLM既非常昂貴，也是一項複雜的系統工程挑戰。訓練需要建造和操作一台由數千台主機、專用網路、電力和冷卻基礎設施以及數千個加速器（通常是GPU或TPU）組成的超級電腦。訓練和服務數千億參數所需的高昂成本和專業知識使得許多學術研究人員和小型組織難以獲得先進AI能力，尤其是在需要內部部署時。例如，估計訓練OpenAI的GPT-4的計算成本為7800萬美元，而谷歌的Gemini Ultra為1.91億美元。為超大規模公司或任何需要持續在其資料上訓練和更新強大AI模型的組織建造和部署大型整體模型可能是不可持續的。最後，適應整體模型的系統已經將計算TFLOP的增長速度遠大於記憶體頻寬和容量，產生了記憶體壁障，記憶體系統無法再有效地供給計算。

ML研究社群已經做出回應，提供了一個由更小、更模組化的模型組成的生態系統，這些模型同樣強大，但訓練和服務成本更低、更容易。像7B參數的Llama 3, Llama 2和Mistral 7B等小型模型通常已足夠。它們可能無法在一般任務套件上匹配大型模型的性能，但在針對它們進行微調的一組較窄專業任務上，小型模型可以提供優越的精確度。例如，Flan-T5-XL只有3B參數，但在MMLU分數上超過175B參數的GPT-3近10%。像這樣的案例加強了社群在透過將基礎模型專業化到一個領域、透過將基礎模型微調到特定任務或一組任務，以及透過將大型模型蒸餾或壓縮到小型模型中來構建和訓練小型模型的活動。此外，這些較小模型的組合已被證明展現出匹配大型整體模型的新興行為。它們讓更廣泛的社群能夠使用AI。

我們相信，未來成功的AI系統將有效地託管和執行許多小型模型。這反映在學術界成功追求的方向以及業界正在採用的新產品中。類似CoE的複合AI系統在推進AI前沿方面發揮了關鍵作用。在本文中，我們將這些由專業化小型模型組成的模組化系統稱為專家組合(CoE)。

CoE由幾個小型專家模型組成，協同完成一項任務。一個專家的輸出決定下一步要執行哪個專家。運行一個專家涉及將模型參數權重載入到加速器的主記憶體中，然後執行模型。因此，執行CoE涉及一系列模型切換和模型執行。如圖1所示，目前先進的AI加速器無法有效地處理這一系列操作。

有效加速CoE歸結為在最小化模型切換成本的同時有效執行專家模型。我們將其細分為三個關鍵要求：

1) 積極的運算子融合和流水線平行化，以有效執行專家模型。較小的模型具有較低的運算強度和運算子之間複雜的存取模式。傳統的運算子融合技術在任意存取模式下取得的成功有限。

2) 高頻寬記憶體，以在生成式推理期間利用權重和中間結果中的時間和空間局部性，以及

3) 高容量記憶體，以最小化切換成本並儲存許多專家模型的參數

在本文中，我們描述了一種硬體/軟體解決方案，透過解決上述挑戰來克服記憶體壁障。

首先，我們描述Samba-CoE，一個擁有150個7B專家的兆級參數CoE系統，以及如何有效運行它需要硬體支援積極的運算子融合和新穎的記憶體系統。我們介紹SambaNova SN40L可重構dataflow單元(RDU)，這是一種結合流式dataflow平行化與新穎的三層記憶體系統的商用dataflow加速器，包含大型晶片內SRAM、HBM和直接連接到加速器的DDR DRAM。

RDU的流式dataflow架構允許我們將數百個具有任意存取模式的複雜運算子融合到單一核心呼叫中——無需程式設計師手動編寫該核心。這透過利用晶片內硬體支援流水線、資料和張量平行化的混合來提供巨大的加速。我們的積極融合技術遠超過傳統架構使用的先進技術的能力。

使用TSMC 5nm技術製造的SN40L RDU是一種採用2.5D晶片上晶圓封裝(CoWoS)的晶片級設計，包含兩個SN40L可重構dataflow晶片(RDD)和HBM。每個SN40L RDU插槽使用1040個分散式模式計算單元(PCU)提供638 BF16 TFLOP的峰值計算性能。這些由1040個分散式模式記憶體單元(PMU)補充，合計提供數百TBps的晶片內記憶體頻寬以及PMU內和PMU之間的高銀行級平行度。靈活的晶片內位址產生邏輯為任意張量記憶體存取模式提供高頻寬。SN40L中的三層記憶體是：520 MiB晶片內PMU SRAM、64 GiB封裝內HBM和高達1.5 TiB的DDR DRAM（使用可插拔DIMM）。在單個SN40L節點中，模型以超過1 TB/s的速度從DDR載入到HBM。

我們在幾個真實世界基準測試上量化並討論流式dataflow平行化的影響，顯示相較於優化的基線加速2倍到13倍。我們在包含八個SN40L RDU插槽和一個主機的單個SN40L節點上部署Samba-CoE。我們討論了SN40L節點上Samba-CoE的性能，並與DGX A100和DGX H100進行比較。我們表明，對於CoE推理部署，SN40L將機器規模減少高達19倍，將模型切換時間加速15倍到31倍，相較於DGX H100和DGX A100實現3.7倍到6.6倍的整體加速。

本文的組織如下：第二節描述我們的兆級參數CoE Samba-CoE。第三節描述流式dataflow及其轉化為關鍵硬體需求的挑戰。第四節詳細描述SN40L硬體架構，並列出與先前RDU相比的主要變化。第五節描述管理DDR和HBM的軟體支援。第六節量化流式dataflow的好處以及Samba-CoE的性能。第七節討論從硬體/軟體協同設計過程中學到的關鍵經驗。第八節涵蓋相關工作。我們在第九節作結。

II. 背景：專家組合 

在本節中，我們描述在SN40L上構建和部署的一個CoE實例，名為Samba-CoE。圖2顯示了Samba-CoE從提示到回應的流程。

Samba-CoE由多個專家模型和一個路由器模型組成。每個專家都在特定領域中進行了微調。我們利用了開源社群中在編程、數學和語言翻譯等領域微調的幾個優秀專家模型。路由器是另一個專業模型，可動態地將每個輸入提示分配給最相關的專家。例如，與數學相關的查詢將被路由到數學專家，而編程問題將去到編程專家。

Samba-CoE的靈感來自混合專家(MoE)架構，但有一些關鍵區別。雖然MoE和CoE都比傳統的密集型整體模型更稀疏地啟動，但MoE比CoE更不靈活。MoE需要像整體模型一樣作為單一模型進行訓練/微調，而CoE是由獨立且異構的專家模型組成，這些模型彼此獨立地訓練/微調。CoE也更強大：先前的工作已經表明CoE可以優於MoE以及GPT-3.5和GPT-4等大型整體模型, 。我們注意到CoE和MoE是正交的技術，可以輕鬆結合：CoE可以利用內部實現為MoE的專家模型。

為簡單起見，在本文中，路由器模型和專家模型都來自Llama2-7B。請注意，路由器和專家模型不需要是同質的——它們可以是具有不同參數數量的不同架構。選擇Llama2-7B作為本工作的基礎是因為其方便的大小、令人印象深刻的能力以及強大的社群支援。CoE概念和Samba-CoE系統不限於Llama2。

III. CoE的硬體需求

CoE執行時間細分為模型執行時間和模型切換時間，如圖1所示。最小化CoE執行時間可用於在每個用戶基礎上減少機器佔用空間或在給定佔用空間下增加支持的用戶數量。為了減少模型執行時間，我們展示了流式dataflow相對於傳統運算子融合的優勢。為了最小化模型切換時間，我們說明了同時需要高容量的加速器本地DDR介面和HBM的必要性。

A. 流式Dataflow

傳統運算子融合不足：運算子融合是一種常見的優化技術，用於增加運算強度並提高硬體利用率。融合還減少了執行模型所需的核心呼叫數量，並攤銷了核心啟動開銷。然而，專家模型通常包含低運算強度的運算子，加上涉及洗牌和轉置的複雜存取模式。複雜的存取模式嚴重限制了GPU上融合的功效。像PyTorch2和TensorRT等框架對融合明確不支持的模式有文檔記錄的限制。因此，許多複雜的融合核心仍然是為GPU手動編寫的。

圖3描述了Monarch FFT分解的簡化版本，張量形狀標註在邊上。表I顯示了融合對運算強度的影響。較高的運算強度允許應用程式針對給定的目標加速器實現屋頂線性能。例如，A100 GPU的TFLOPS/TBps比率約為300/2=150，這意味著運算強度小於150 FLOP/byte的核心在A100上受記憶體限制。在表I中，前兩行在A100上受記憶體限制，最後一行受計算限制。

然而，由於以下原因，GPU無法融合圖3的所有部分：

1) 僵化的記憶體層次結構和程式設計模型造成資料移動瓶頸：GPU核心使用線程塊網格啟動。網格結構在核心的持續時間內是固定的。融合Gemm0和Mul將是微不足道的。然而，Transpose迫使線程存取其他SM中線程的值，觸發通過共享快取和HBM在SM之間交換資料。由於沒有其他方式在SM之間傳輸資料，這種缺乏靈活性在共享快取和HBM上造成瓶頸。

2) 晶片內SRAM容量不足迫使Transpose的輸出具體化到HBM，防止融合機會。 

3) 運算子之間沒有利用流水線平行化：更高階的Monarch FFT分解（在第VI節中研究）創建許多小的矩陣乘法，其大小為32×32×32或更小，這無法有效利用所有SM。然而，在所有矩陣乘法和元素級運算子之間存在大量的流水線級平行度。GPU SIMT程式設計模型沒有提供直接的方式將圖3中的運算子作為流水線執行。 

流式Dataflow實現流水線和任意存取模式的自動融合：

與傳統融合不同，流式dataflow將運算子作為粗粒度流水線執行。張量被切塊並通過此流水線流式傳輸。切塊之間可以具有任意的讀寫存取模式。

圖4描繪了空間融合的實現。藍色框表示晶片內緩衝區單元，灰色框表示晶片內計算單元。運算子Gemm0、Mul和Gemm1作為粗粒度流水線中的階段執行。中間的藍色記憶體單元用作解耦的階段緩衝區，保存中間結果。由於Gemm0和Gemm1佔總操作的較大部分，因此分配更多計算單元給它們。通過使用適當數量的記憶體單元來匹配各自階段所需的輸入和輸出頻寬。例如，邏輯階段緩衝區I0被分割為兩個記憶體單元I00和I01，以匹配Gemm0所需的輸入頻寬。出於容量原因，緩衝區S0-S3被分割為四個記憶體單元。轉置操作被融合到緩衝區T0*和T1*中作為存取模式。

我們從上面的觀察中提煉出以下晶片內架構特性列表，以實現流式dataflow：

1) 可組合的記憶體單元：單個記憶體單元提供固定的容量和頻寬。由於晶片內張量的容量和頻寬需求各不相同，硬體應支援跨記憶體單元對邏輯位址進行可程式設計的交錯處理。

2) 位址產生頻寬和靈活性：高資料頻寬需要高位址產生頻寬。此外，每個記憶體單元應支援非阻塞的並發讀寫，以有效實現階段緩衝區。簡而言之，位址產生硬體應允許以高輸送量為任意複雜的位址表達式產生多個並發位址。

3) 全身和流式計算：ML加速器架構通常實現全身陣列以提高GEMM等運算的計算密度。然而，在許多ML模型中，GEMM經常伴隨著需要高輸送量流式計算能力的元素級運算子和約簡。

4) 一對多、多對一和資料重排：生產者和消費者單元數量的差異會在RDN上創建一生產者對多消費者和多生產者對一消費者的流量流。對於一對多，需要硬體支援在互連中從來源創建分支路徑到軟體決定的一組目的地。對於多對一流量，來自不同路徑的資料可能以亂序方式到達目的地。這個亂序序列通常必須在目的地重新排序，以匹配程式的預期。換句話說，硬體必須提供一個協議來重新排序資料流。最後，需要軟體控制的頻寬管理（例如節流）和路由來滿足具有不同頻寬需求的流。所有這些硬體支援需要能夠被編譯器內的自動佈局和布線演算法使用。

B. 模型託管和切換成本

HBM的有限容量限制了當託管在GPU或TPU上時CoE中的專家數量。僅使用HBM，運行大型CoE需要(a)使用更多機器來提供HBM容量，這增加了成本、複雜了部署並引入了負載平衡挑戰，或者(b)使用主機的記憶體，如圖1所示，這增加了切換延遲。直接連接到加速器的更高容量DDR記憶體可以同時減少模型託管和模型切換成本。此外，專家參數在CoE中表現出時間局部性，因為它們被多次使用（例如，在自迴歸解碼期間）。HBM通過充當DDR和SRAM之間的軟體管理快取層，在利用這種時間資料局部性方面發揮關鍵作用。

因此，我們得出結論，執行較小模型組合的系統需要兩種類型的晶片外記憶體：(1)高頻寬記憶體，以利用專家參數的時間局部性，以及(2) 高容量記憶體，以在小佔用空間中儲存專家參數。

在下一節中，我們將描述基於上述原則構建的SN40L可重構dataflow單元。

IV. SN40L硬體架構

SN40L dataflow加速器使用TSMC的5FF工藝製造，並使用晶片上晶圓封裝(CoWoS)多晶片封裝技術封裝為雙晶片插槽。圖5顯示了SN40L的顯著組件，描述如下。

RDU Tile：粗粒度可重構dataflow核心陣列。包括模式計算單元(PCU)、模式記憶體單元(PMU)以及連接在稱為可重構dataflow網路(RDN)的二維網格互連中的位址產生和合併單元(AGCU)。

記憶體介面：SN40L與兩層晶片外記憶體介面——HBM和DDR。這兩個記憶體空間都由軟體管理。DDR層可以有高達1.5 TiB的峰值記憶體容量，峰值頻寬超過200 GB/s。HBM層每個插槽有64 GiB的容量，峰值頻寬約為2 TB/s。

晶片間(D2D)介面：SN40L tile組件可以直接在兩個晶片之間流式傳輸資料，而無需通過晶片外記憶體。

主機介面：SN40L通過PCIe連結與主機x86 CPU介面。此介面支援主機和設備晶片外記憶體之間的DMA以及主機和tile之間的直接通信。

對等(P2P)介面：將SN40L連接到其他SN40L RDU。第IV-D節中描述的對等協議提供了實現集體通信原語的基本體。

頂級網路(TLN)：此網路將SN40L tile連接到主機、記憶體和對等介面。

圖6展示了一個具有關鍵dataflow組件的SN40L tile：PCU、PMU、RDN交換機和AGCU。以下小節將更詳細地描述它們。

A. 模式計算單元(PCU)

SN40L中的PCU提供SN40L中的全身和流式計算能力。PCU的資料路徑由頭部、體部和尾部組成。頭部消耗輸入dataflow並驅動體部。PCU的體部可配置為輸出固定的全身陣列或具有多個向量計算級的流水線SIMD核心。尾部執行特殊的元素級函數並驅動輸出FIFO。這實現了GEMM等運算、元素級運算或約簡的高效執行。

圖7展示了PCU作為2D全身陣列和SIMD核心的示意圖。全身陣列加速矩陣乘法，如圖3中的Gemm0和Gemm1。輸入到全身陣列通過稱為廣播緩衝區的結構從左到右和從上到下流式傳輸。累積結果通過尾部單元從左到右排出到輸出FIFO。矩陣乘法可以在多個PCU上進一步並行化，類似於圖4中的描述。作為SIMD核心，PCU以完全流水線方式執行並行的多維張量運算，如圖4中的Mul。每個SIMD階段支持FP32、BF16和INT32格式的常見算術、邏輯和位元運算。PCU可以配置為實現可選的跨通道約簡網路，如圖中的藍色三角形所示。還支援純SIMD方式的通道內約簡。跟蹤迴圈反覆運算的計數器在達到程式設計的最大值時產生控制事件，表明迴圈已完成執行。

尾部支援超越函數、隨機數產生、隨機取整和格式轉換。尾部操作可以與體部的計算融合和流水線化。

一個操作可以以資料並行、張量並行或流水線並行的方式在多個PCU上並行化。資料並行通過分割輸入和輸出來創建由不同PCU處理的多個獨立資料流來實現。張量並行通過分叉到資料並行流，然後將它們連接來實現。流水線並行通過將多個PCU鏈接在一起以融合運算並增加運算強度來實現。

B. 模式記憶體單元(PMU)

SN40L中的PMU提供晶片內記憶體容量、頻寬和運算子融合所需的位址靈活性。圖8顯示了PMU的高級塊圖。PMU用於儲存晶片內張量，如輸入、參數、元資料和中間結果。例如，圖4中的所有藍色塊對應於PMU。PMU的關鍵組件如下所述。

Scratchpad banks：每個PMU都包含一個程式設計員管理的scratchpad記憶體，它被組織為SRAM實例陣列。SRAM陣列整體支援併發寫入和讀取。

標量ALU管道：PMU包含幾個整數ALU級，可配置為同時產生讀取和寫入scratchpad的位址。PMU ALU實現了一組經常用於位址計算的特殊複雜指令，如位域提取和移位設置。這使得能夠以更少的ALU級更有效地生成複雜的位址，因此延遲更少。ALU管道還有一條路徑，可以從標量RDN攝取標量作為操作數，並將計算值作為標量輸出到標量RDN中。此路徑有助於位址可組合性：如果需要，可以將複雜的整數計算分解並映射到多個PMU上。

如第III節所述並在圖4中所示（緩衝區T00-T03和T10-T13），空間融合核心中的階段緩衝區需要併發讀寫，它們可能具有不同的存取模式。雖然不是普遍正確的，但我們已經偶然觀察到張量的寫入和讀取存取模式經常相互影響對方的複雜性；複雜的寫入存取模式通常實現更簡單的讀取存取模式，反之亦然。標量ALU管道允許軟體利用這一見解。它可以分割為獨立的讀取和寫入位址產生管道，具有軟體配置的分配給每個管道的級數。

位址謂詞和Banking：圖4顯示，由於容量需求（S0-S3）、頻寬需求（W00-W01和I00-I01）或兩者兼而有之（T00-T03和T10-T13），單個邏輯張量可以跨越多個PMU。PMU透過提供鉤子以程式設計方式控制跨PMU的張量位址交錯來實現這一點。具體而言，每個PMU可以程式設計為該PMU的一系列有效位址，或為每個產生的位址提供一個可程式設計的謂詞位元。如果在程式設計的範圍內或有效謂詞，則處理位址，否則丟棄。此外，使用軟體可程式設計的銀行位元位置將位址映射到scratchpad銀行。

資料對齊單元：資料對齊單元支援常見的張量轉換運算，如轉置、跨通道向量置換、向量非對齊存取、查找表(LUT)、資料格式和資料佈局轉換。要轉置的張量以特殊的對角線條紋格式跨scratchpad bank寫入，這使得能夠以完整頻寬讀取相同張量的常規和轉置格式。此功能實現了透過將轉置運算子融合到圖4中緩衝區T00-T03和T10-T13之間的讀寫存取模式優化中來實現轉置運算子。

C. 可重構Dataflow網路（RDN）

RDN是晶片內可程式設計互連，可促進PCU、PMU和AGCU之間的通信。RDN由三個物理網狀架構組成——向量、標量和控制。向量和標量網狀是以封包為基礎的。控制網狀是電路交換的，由一束可以單獨路由的單位元線組成。向量網狀是張量資料的主要管道。標量網狀主要用於傳輸元資料，如位址，但在某些情況下也可以用於攜帶資料或控制。控制網狀用於攜帶分佈式粗粒度流控制的token以及集體編排圖的執行。控制token通常對應於計數器的「完成」事件，表示迴圈的結束，如第IV-A節中簡要討論的。

RDN使用非阻塞交換機網格實現。如圖7和8所示，從交換機到單元的輸入標量和向量封包落在輸入FIFO中，透過輸出FIFO退出。向量和標量網狀上的傳輸在每一跳都受到基於信用的流控制。封包流在RDN上的通信單元之間也透過軟體token、細粒度硬體信用和硬體中的前向進度保證的組合進行端到端流控制。所有三個RDN網狀的路由表都由軟體使用編譯器中的佈局和布線（PnR）層配置。

現在我們描述支援第III節中確定的關鍵通信模式的機制。

• 多播和可程式設計路由：標量和向量網狀上的封包路由是透過2-D維度順序路由或軟體控制的靜態流路由動態完成的。在靜態流路由中，軟體為封包流分配流ID域，該域隨封包一起攜帶。在每個交換機埠對流ID域進行解碼，並在將封包轉發到下一個目的地之前重新分配。靜態流路由機制透過交換機支援封包多播。
• 多對一和資料重排：向量封包包含一個稱為序列ID的中繼資料域，這是在晶片內支援任意多對一流的主要機制。單元的向量輸出埠配備了可程式設計邏輯，用於為每個輸出向量產生序列ID。這樣，軟體就可以程式設計序列ID，以表示跨多個來源的給定運算的邏輯向量順序。序列ID域用作PMU中的輸入運算元，以計算寫入位址來重新排序封包。

D. 位址產生和合併單元（AGCU）

AGCU是RDU tile到本地設備記憶體（HBM/DDR）、主機記憶體、遠程RDU設備記憶體和透過TLN到遠程RDU tile的可重構dataflow橋。在tile側，它像dataflow核心一樣暴露RDN向量、標量和控制埠。在TLN側，它產生讀寫請求並合併回應。它配備了標量位址產生管道和計數器，與PMU邏輯有一些相似之處（不包括SRAM）。位址產生管道由計數器組成，並提供記憶體管理的位址轉換層。

對等通信：AGCU支援對等（P2P）通信協議，以直接在不同插槽上的RDU tile之間流式傳輸資料，而無需涉及DDR或HBM。P2P協議提供了在RDU之間構建集體通信原語所需的原語，如AllReduce。

核心啟動編排：AGCU實現了一個核心啟動機制，由三個命令的序列組成：程式載入、參數載入和核心執行。運行模型涉及執行核心啟動的時間表，可以是軟體編排的或硬體編排的。軟體編排允許更靈活地調度核心，並提供更多主機軟體對模型執行的可見性。然而，軟體編排會產生影響性能的開銷。硬體編排將靜態核心時間表卸載到AGCU中的專用硬體，這顯著降低了開銷。在第VI節中，我們量化並討論了軟體與硬體編排執行對各種基準的影響。

E. SN40L中的新功能

先前的工作涵蓋了從學術原型到第一代工業產品的經驗教訓。本文首次深入探討了生產RDU的內部工作原理。

SN40L對於更廣泛的工作負載相較於先前的RDU是一個顯著的改進，它還改善了軟體可用性。SN10主要專注於訓練。我們討論一些關鍵的增強功能。

HBM：SN40L是第一個包含HBM的RDU。這是由需要繼續將大型模型放入緊湊的佔用空間（使用DDR）但也需要改善更多受記憶體限制的推理工作負載的加速（使用HBM）所驅動的。添加HBM記憶體層對CoE的可行性至關重要，但它也能夠在整體模型上獲得更好的性能。

GEMM和小張量：PCU的全身陣列得到改進，可以更有效地處理較小的矩陣乘法。PCU和PMU增加了對具有整合偏置（可以節省硬體資源）的GEMM以及對有效的晶片內pad產生而無需具體化pad值的支援。

動態dataflow：SN40L包括幾項改進，以更好地支援動態dataflow。原始SN10 RDU要求所有封包目的地都按每個dataflow核心進行靜態配置。這在執行將封包在多個目的地之間解多工的序列化一對多運算（如張量「拆分」運算）時被證明是有局限性的。SN40L在PCU中增加了對每個封包的動態目的地的支援。另外，SN40L PMU可以透過增強資料對齊單元，使其具有高輸送量可程式設計通道改組和遮罩，更好地支援FFT和排序等運算子。

晶片內頻寬利用率：在SN10上，必須全域分配有限的RDN流量流ID池。這意味著，如果具有流ID X和Y的兩個流經過相同的交換機，則這些ID在這兩個流使用的任何交換機中都不可用。這為編譯器創建了一個複雜的全域最佳化問題，並浪費了晶片內頻寬。在SN40L中，流ID分配被改為類似於MPLS[15]的靈活方案。每個流在編譯時分配一個流ID序列，該序列被程式設計到交換機流表中。當封包流過交換機時，封包中的流ID域使用交換機流表重寫。這種靈活性導致更好的晶片內頻寬利用率和可用性。

增強的電源管理：SN10需要軟體幫助緩解電壓下垂（L di/dt），但實現過於保守，導致性能損失高達25%。SN40L透過先進的時脈技術和更複雜的硬體電源管理更好地緩解了電壓下垂。

V. 軟體支援

在本節中，我們描述如何在SN40L上部署Samba-CoE。

Samba-CoE由150個Llama-7B專家組成，總共超過1T參數。它部署在單個具有8個RDU插槽的SN40L節點上。圖9顯示了Samba-CoE如何利用DDR和HBM，以及單個提示的簡化事件流。所有150個專家的權重保存在高容量DDR中，而路由器的權重保存在HBM中。單個Samba-CoE推理有三個高級步驟：(1)運行路由器以確定專家，(2)將專家從DDR複製到HBM，以及(3)運行專家。

在模型切換之後，專家模型的權重在自迴歸解碼迴圈中被多次讀取以產生多個token。Samba-CoE中這種固有的模型層次時間局部性透過將權重移動到HBM來利用。

然而，這確實給軟體堆疊帶來了管理多個非一致記憶體空間的額外負擔。在本節中，我們討論在應用層和SN40L的低層次執行時期驅動程式之間構建的軟體基礎設施，以實現這一願景。

A. 記憶體分配

我們這個系統設計的目標之一是使設備記憶體管理對應用開發人員透明。為此，我們在SN40L編譯器中加入了自動異構設備記憶體管理。基本假設是只要一切都適合，預設就使用HBM。因此，我們將DDR用於兩個主要目的：

1) 當給定模型的常駐記憶體太大而無法容納在HBM中時，將資料從HBM溢出到DDR，以及 2) 保存CoE的所有其他當前非活動的模型。請注意，雖然在HBM和DDR之間移動資料會產生顯著的性能成本，但它仍然比將所有內容一直溢出到主機的記憶體要便宜得多，如圖1所量化的那樣。在本節中，我們將重點關注第一種使用情況，並在第V-B節中討論第二種情況。

我們發現需要積極的垃圾回收才能將表II中的Llama 7B等模型放入每個插槽64 GiB的HBM容量中。然而，典型的動態垃圾回收方案對於這種使用情況而言開銷太大，因為它會要求我們經常在應用程式中間將控制權返回給CPU以重新組織SN40L的記憶體。相反，我們利用了SN40L的程式設計模型既沒有動態記憶體分配也沒有指標別名的事實，因此我們可以靜態執行符號生命週期分析，並透過在其生命週期不重疊時將多個邏輯符號分配給相同的設備虛擬位址來實現垃圾回收。

最後的問題是當記憶體仍然不適合時要將哪些符號溢出到DDR。因為這兩種記憶體之間存在一個數量級的頻寬差異，所以我們分析每個符號的時間局部性及其傳輸佔用空間，以估計該符號在整個應用程式中的總頻寬需求。然後按總傳輸大小對符號進行排序，以便我們首先溢出具有最小頻寬需求的符號。

在實踐中，我們發現對於表II中的LLM模型，權重獲得最高優先級以保留在HBM中，而激活符號和其他中間結果可以在必要時溢出。我們正在研究對該演算法的進一步改進。

B. CoE執行時期

我們為創建CoE模型架構的主要設計目標之一是每個單獨的模型專家的生命週期應該是獨立的：開發、編譯、訓練、微調、量化、維護、服務和共享。這為軟體模組化提供了顯著的飛躍，優於整體模型。模型精確度和性能的進展非常快，CoE使得逐步利用最新創新變得容易。

為了實現這一目標，我們需要在執行時期將任意數量的獨立編譯的模型鏈接在一起，並在新請求從應用層進入時在它們之間動態切換。該方法類似於動態鏈結器/載入器在傳統軟體應用程式中的工作方式。

我們在先前描述的靜態管理器之上添加了一個輕量級動態記憶體管理器。它可以使用主機執行時期有效實現，因為它只需要在整個專家模型的邊界運行。每個編譯的模型二進位檔案提前告訴我們該模型需要多少HBM和DDR空間。然後，CoE執行時期與低層次設備驅動程式介面，以動態分配DDR中每個模型的記憶體塊。此分配包括編譯器打算用於HBM的記憶體段；這些記憶體段最初也儲存在DDR中。當請求特定模型時，CoE執行時期透過將打算用於HBM的記憶體段「激活」，然後將控制權轉移到編譯的應用程式二進位檔案以正常運行。一旦應用程式退出，控制權返回到CoE執行時期，等待下一個請求。

CoE運行時嘗試在HBM中同時保持盡可能多的模型處於活動狀態，並在達到容量限制時使用LRU驅逐策略。如果下一個請求是針對同一模型的，則可以立即恢復，無需額外開銷。如果請求需要一個新模型，而新模型無法適應剩餘的HBM空間，則從HBM中驅逐最舊模型的HBM記憶體，並換入新模型的HBM記憶體。為避免不必要地複製回大型唯讀權重，編譯器在適當的地方將符號標註為唯讀，然後運行時跳過將任何此類符號複製回DDR。

VI. 案例研究

在本節中，我們使用幾個真實世界的訓練和推理基準來量化並討論運算子融合的影響。然後，我們在SN40L節點上評估Samba-CoE，並與DGX A100和DGX H100進行比較。

A. 運算子融合的影響

1) 基準和設置：表II描述了用於量化運算子融合影響的一組語言模型基準。基準包括各種大小的幾個訓練和推理工作負載。在推理工作負載中，我們將「預填充」階段——生成第一個token——與「解碼」階段——透過自迴歸解碼生成第二個和後續token分開。自迴歸解碼步驟利用KV緩存，與構建KV緩存的預填充階段相比，計算量和運算強度要低得多。FlashFFTConv的計算圖結構是第III節中描述的激勵示例的複雜版本。

除了FlashFFTConv之外的所有實驗都在包含8個SN40L插槽和一個主機的系統上進行評估，FlashFFTConv是一個較小的核心，我們在單個SN40L上進行評估。我們在三種配置中測量並比較每個基準的性能：

• 非融合：模型中的每個PyTorch運算子在SN40L上作為一個或多個核心執行，中間結果具體化到DDR或HBM。核心由軟體安排執行。每個核心仍然被並行化以在SN40L上有效運行。
• 融合+軟體編排（SO）：使用自動編譯器優化和程式設計員提示的組合將運算子融合到更少的核心中。核心調度從主機上運行的軟體執行。
• 融合+硬體編排（HO）：與上面相同的融合核心，但核心調度使用第IV-D節中描述的功能卸載到硬體。

2) 基準結果：圖10顯示了運算子融合對所有基準的影響。藍色條形圖量化了運算子融合的影響。橙色條形圖量化了硬體編排核心啟動的影響。

運算子融合加速 如果沒有融合，FlashFFTConv的運算強度非常低，性能受到影響。在SN40L上進行融合時，整個FlashFFTConv基準使用單個核心啟動執行，類似於圖4中所示的簡化圖。增加的運算強度提供了13倍於非融合基線的加速。

推理預填充、訓練和sparseGPT基準實現了1.5倍到3倍的加速範圍。使用空間融合，在具有足夠粗粒度流水線平行性的情況下實現更高的運算強度，以保持晶片內單元被佔用。對於這些基準，批次大小和序列大小提供了所需級別的流水線平行性，以利用融合的好處。

儘管自迴歸解碼推理基準的運算強度較低，但它們透過消除核心啟動和不必要的HBM流量的開銷而從融合中獲益。我們觀察到從1倍到13倍的加速，其中Mistral實現了最高的加速。

圖11比較了以融合和非融合配置執行各種基準所涉及的核心呼叫啟動數量。對於llama7B-4k-inf-prefill，11倍的比率意味著使用運算子融合，llama7B-4k-inf-prefill執行的核心啟動數量是其非融合基線的11倍。此比率量化了在給定基準中執行的融合級別。對於某些基準，較高的數字意味著更積極的融合，如FlashFFTConv和sparseGPT的情況。對於其他基準如llama70B，較高的數字也是由於模型的較大規模。

硬體編排核心啟動加速：現在我們討論由於硬體編排核心啟動獲得的改進。在這裡，我們看到與前面討論相反的趨勢：自迴歸解碼推理基準實現了顯著的1.4倍到8倍的加速。在這些基準中，核心具有非常短的執行時間，其中大部分時間都由載入權重和其他輸入所主導。因此，核心啟動開銷開始佔總時間的更大部分。將核心調度卸載到SN40L可以消除軟體調度的開銷，從而提供加速。

另一方面，推理預填充和訓練基準僅看到最多1.1倍的改進。在這些基準中，每個核心執行的時間要長得多，因此核心啟動開銷被攤銷了。融合的FlashFFTConv基準只有一個核心呼叫，在兩種核心調度方法下花費相同的持續時間。

B. 專家組合(Composition of Experts, COE)

我們量化在SN40L節點上部署專家數量不斷增加的Samba-CoE的延遲和系統佔用空間，並與DGX A100和DGX H100進行比較。表III總結了結果。我們研究了兩種專家數量不斷增加的情況：單個節點上的延遲影響，以及系統佔用空間影響以維持相同的TP8延遲。

延遲影響：我們模擬兩個用例：聊天機器人對話用例，每個輸入提示產生20個輸出token，以及翻譯用例，產生200個輸出token。模型和路由器在所有平台上以張量並行的8插槽（TP8）方式映射。路由器和KV緩存總是在HBM中。SN40L節點數是在真實硬體上測量的。由於Samba-CoE未部署在DGX上，我們使用已發佈的模型延遲數和基於DGX規格的樂觀模型切換估計值來估計延遲。總延遲包括運行路由器、複製所需的專家權重以及運行專家的時間。研究了兩個專家情況：生成20個輸出token和生成200個輸出token。我們樂觀地假設在DGX上，HBM和主機記憶體的整個容量都可用於保存權重和KV緩存（僅HBM）。

我們報告批次大小（BS）= 1和BS=8情況下的延遲。請注意，「批次」適用於整個Samba-CoE模型，而不是單個專家。BS=8意味著Samba-CoE模型收到8個批次提示。首先使用BS=8運行路由器以獲得每個提示的專家。樣本中的每個提示可能需要不同的專家，因為批次中的樣本彼此無關。然後將所需的專家複製到HBM。然後按順序運行每個（提示，專家）對。

圖12比較了三個平台的延遲。我們從兩個廣泛類別分析這些結果：

50個專家以下：所有專家都適合HBM，因此性能受專家執行時間的限制。SN40L節點生成20個token的速度比DGX A100快2倍，比DGX H100快1.5倍。對於200個token，加速數字分別為3.2倍和2.3倍。請注意，生成式推理受記憶體限制，SN40L節點具有與A100相當的HBM頻寬（低於H100）。加速效果清楚地展示了流式dataflow的好處：整個解碼器層被融合到單個核心呼叫中，使用近90%的PCU和PMU，並飽和接近85%的HBM頻寬。此外，由於模型主要包含多個相同的解碼器層，SN40L幾乎看不到核心啟動開銷。

50個專家以上：當專家溢出到主機DRAM時，DGX上的延遲會激增（大約50個7B專家）。對於BS=8，SN40L節點分別實現了6.6倍和3.7倍於DGX A100和DGX H100的加速。對於BS=1，SN40L節點分別實現了4.8倍和2.8倍的加速。BS=8需要複製更多的專家，因此佔總時間的更大部分。圖1a和1b顯示了專家切換與模型執行的時間細分。在SN40L節點上，借助超過1 TB/s的DDR到HBM總頻寬，複製時間比DGX A100快31倍（提供32 GB/s的主機到GPU頻寬），比H100快16倍（提供64 GB/s的主機到GPU頻寬）。DGX在150個專家時耗盡記憶體。

系統佔用空間影響：接下來，我們量化隨著專家數量增加以維持每個平台上相同TP8延遲的系統佔用空間。實現這一點需要消除GPU上的專家複製。因此，所有專家都應位於GPU HBM中。SN40L DDR到HBM的切換成本計入SN40L的延遲。

單個SN40L節點可以以TP8延遲保存和服務多達850個專家的CoE。使用DGX實現這一點需要19個DGX節點才能將所有專家保存在HBM中。

VII. 經驗教訓

SN40L是一個複雜的系統，是多年來許多軟硬體工程師集體智慧和辛勤工作的結晶。在本節中，我們討論了在此過程中學到的三個經驗教訓。

軟體管理頻寬：軟體必須管理來自各種實體的頻寬：Tile級單元通信、HBM、DDR、晶片間、對等和主機頻寬。頻寬轉化為TLN和RDN中的一個或多個並發資料流。為了從HBM等單元利用更多頻寬，需要創建更多的載入和儲存資料流。相反，需要更少頻寬的單元應分配更少的資源，以避免過度提供和浪費。在編譯器中構建靜態頻寬模型對於啟用適當的頻寬分配和流量管理至關重要。對良好的靜態頻寬模型的投資在其他方面得到了回報：應用程式可以靜態地分析和調整性能到第一個數量級。

性能偵錯：積極的運算子融合和流水線化會在晶片內創建大量並發流量流，從而產生頻寬瓶頸。我們注意到，頻寬問題通常歸結為兩件事之一：網路擁塞或記憶體銀行衝突。SN40L交換機和PMU中的性能計數器統計停頓次數，有助於識別SN40L tile中的熱點。在RDN擁塞方面，

我們觀察到，如果不加管理，突發流量很容易減慢整個核心的速度。硬體中可程式設計的封包節流功能使軟體能夠減少突發行為並緩解許多RDN擁塞問題。為了處理PMU銀行衝突，我們觀察到PMU經常被程式設計為任意張量形狀的雙緩衝區，如果這些緩衝區靜態映射到不同的銀行，則可以避免銀行衝突。第IV-B節中描述的可程式設計銀行位元幫助根據這一見解採取行動，消除了這種多緩衝區配置的銀行衝突。

流水線計算和集體通信：對於我們的編譯器的更高層而言，跨插槽將資料/張量/流水線並行dataflow的映射問題類似於將它們映射到插槽內的問題。借助第IV-D節中描述的硬體對等協議，集體通信運算子可以與其他計算融合並流水線化到同一個核心中，就像任何其他運算子組一樣。此外，SN40L之間的流式對等協議在許多情況下避免了到HBM的躍點，這有助於節省一些通信資源。

VIII. 相關工作

1) 商用AI加速器：在本節中，我們廣泛討論其他商用AI加速系統。由於這是一個競爭激烈的市場領域，關於當代AI硬體的許多細節通常是未知的或模糊的。

NVIDIA A100、H100、AMD MI300X、Google TPUv4、和Intel Gaudi等架構都是帶有HBM的AI加速器。雖然這些架構在其程式設計模型、記憶體系統和擴展能力方面差異很大，但它們沒有有效執行大型CoE和巨型模型所需的三層記憶體系統。因此，如第VI-B節所述，在它們上部署CoE會產生低效率。此外，SN40L中的流式dataflow提供了相較於上述加速器的獨特優勢，如第VI節所量化的。先前的研究量化並利用了TPU和GPU上的運算子融合，但沒有執行本文中描述的積極融合級別。最後，SN40L提供了每個插槽約2.5倍於最近宣佈的NVIDIA GH200的總記憶體容量，這使得在SN40L上支援更大的CoE模型成為可能。

Graphcore、Cerebras、Groq和SambaNova的早期RDU提供了替代的AI加速器。然而，它們都缺乏有效執行CoE所需的三層記憶體系統。據我們所知，SN40L是唯一已經成功演示在單個節點部署萬億參數CoE和其他巨型模型並達到所報告性能的系統。

2) 研究dataflow架構：關於dataflow架構的主題有幾篇先前的出版物，涵蓋了計算、記憶體、互連和程式設計模型的各個方面，如綜述論文所述。據我們所知，SN40L是第一個結合流式dataflow與三層記憶體系統的dataflow架構，並量化其對真實世界基準的影響。

3) 運算子融合：傳統的運算子融合是一個廣泛研究的主題。然而，映射在PCU和PMU上的流式dataflow流水線通常包含20多個運算子（見圖11），這些運算子是由編譯器從Python框架層自動生成的。相比之下，傳統的運算子融合針對1-5個運算子，這些運算子通常是手寫的，並且具有存取模式限制。

4) 參數高效微調技術（PEFT）：像LoRA這樣的技術通常用於將專家權重縮小到應用於基礎模型的小型低秩適配器。然而，在幾種情況下，PEFT技術無法達到與監督微調（SFT）相同的品質。因此，較小的專家模型通常是使用額外訓練或SFT專業化的整個模型（在撰寫本文時，HuggingFace上有超過9000個Llama 2變體）。

IX. 結論

在本文中，我們將專家組合（CoE）描述為大型整體式LLM的一種模組化和成本效益高的替代方案。我們描述了具有150個專家的Samba-CoE，並說明了CoE的硬體需求。然後，我們介紹了SN40L dataflow加速器和SN40L節點，它們旨在使用流式dataflow和新穎的三層記憶體系統來解決記憶體牆問題。SN40L的記憶體系統包括晶片內分散式SRAM、晶片外HBM和高容量DDR DRAM。我們討論了跨DDR和HBM管理地址空間的軟體影響，以及在SN40L上部署Samba-CoE的執行時期複雜性。我們證明了流式dataflow相較於非融合基線提供了2倍到13倍的好處。我們表明，在SN40L節點上部署Samba-CoE將機器佔用空間減少高達19倍，將專家複製時間加速15倍到31倍，並相較於DGX H100和DGX A100分別實現3.7倍到6.6倍的整體加速。

我的觀點

這篇論文展示了SambaNova在解決大型語言模型(LLM)訓練和推理面臨的挑戰方面取得的重大進展。作者提出了一種創新的硬體/軟體協同設計方法,通過結合專家組合(CoE)、流式dataflow和新穎的三層記憶體系統來克服AI記憶體壁障。

我認為CoE是一個非常有前景的研究方向。相較於訓練和服務單一的大型整體式模型,CoE提供了更高的模組化、靈活性和成本效益。它允許針對特定任務或領域專業化和優化較小的專家模型,同時仍能透過路由器實現整體系統的強大功能。這種方法不僅使學術界和較小組織更容易獲得尖端AI能力,而且還為持續改進和更新复合AI系統開闢了道路。

SN40L硬體架構在支援CoE方面給我留下了深刻的印象。流式dataflow設計實現了積極的運算子融合和自動流水線化,大大提高了加速器的效率。三層記憶體系統在提供所需的容量和頻寬來儲存和切換專家參數方面發揮了關鍵作用。晶片內SRAM、HBM和大容量的直連DDR DRAM的獨特組合使SN40L能夠在單個節點上高效地執行具有數兆參數的CoE。

此外,我欣賞論文強調了硬體/軟體協同設計在實現卓越系統性能方面的重要性。記憶體管理、頻寬分配、運算子融合和集體通信等方面的軟體創新與硬體功能緊密結合,以充分利用底層架構。這凸顯出AI系統設計中跨學科協作的價值。

SN40L與DGX A100和DGX H100的性能比較結果令人印象深刻。在CoE推理部署中,SN40L在延遲、機器規模和模型切換時間方面的顯著優勢說明了其架構創新的有效性。隨著AI工作負載變得越來越複雜和多樣化,這種專為CoE等新興範例量身定制的加速器將變得越來越重要。

我認為這項工作為探索複合AI系統和專用硬體架構開闢了令人興奮的道路。它不僅推進了LLM訓練和推理的技術發展,而且還強調了系統層面和跨層協同設計對於驅動AI創新的重要性。我期待看到未來在CoE、記憶體系統和dataflow加速器方面的更多研究。

從MHA、MQA、GQA到MLA的演變歷程

前言

本文將與大家一起梳理從MHA、MQA、GQA到MLA的演變歷程，並重點介紹MLA的設計思路。

MHA

MHA（多頭注意力）是開創性論文《Attention is all you need》中提出的一種注意力形式，可以說是當前主流大型語言模型（LLM）的基礎工作。數學上，多頭注意力是多個獨立的單頭注意力的拼接。假設輸入的向量序列為 \(x_1, x_2, \ldots, x_l\)，其中 \(x_i \in \mathbb{R}^d\)，那麼MHA可以記為：

$$
\begin{aligned}
o_t &= [o_t^{(1)}, o_t^{(2)}, \ldots, o_t^{(h)}] \\
o_t^{(s)} &= \text{Attention}(q_t^{(s)}, k_{\leq t}^{(s)}, v_{\leq t}^{(s)}) \\
&\equiv \sum_{i \leq t} \frac{\exp(q_t^{(s)} k_i^{(s)\top})}{\sum_{i \leq t} \exp(q_t^{(s)} k_i^{(s)\top})} v_i^{(s)} \\
q_t^{(s)} &= x_t W_q^{(s)} \in \mathbb{R}^{d_k}, \quad W_q^{(s)} \in \mathbb{R}^{d \times d_k} \\
k_i^{(s)} &= x_i W_k^{(s)} \in \mathbb{R}^{d_k}, \quad W_k^{(s)} \in \mathbb{R}^{d \times d_k} \\
v_i^{(s)} &= x_i W_v^{(s)} \in \mathbb{R}^{d_v}, \quad W_v^{(s)} \in \mathbb{R}^{d \times d_v}
\end{aligned}
$$

為了簡化，這裡省略了注意力矩陣的縮放因數。在實踐中，常見的設置是 \(d_k = d_v = \frac{d}{h}\)。例如，對於LLAMA2-7b來說，有 \(d = 4096, h = 32, d_k = d_v = 128\)；對於LLAMA2-70b則是 \(d = 8192, h = 64, d_k = d_v= 128\)。這裡只考慮了主流自回歸大型語言模型所使用的Causal Attention，這意味著在逐個token生成時，新預測出來的第 \( t+1 \) 個token並不會影響已經計算好的 \( k_{\leq t}^{(s)}, v_{\leq t}^{(s)} \)，因此這部分結果我們可以快取下來供後續生成調用，避免不必要的重複計算，這就是所謂的KV快取（KV Cache）。

瓶頸

一個自然的問題是：為什麼降低KV快取的大小如此重要？

眾所周知，一般情況下大型語言模型的推理都是在GPU上進行，單張GPU的顯存是有限的，一部分我們要用來存放模型的參數和前向計算的啟動值，這部分依賴於模型的體量，選定模型後它就是個常數；另外一部分我們要用來存放模型的KV快取，這部分不僅依賴於模型的體量，還依賴於模型的輸入長度，也就是在推理過程中是動態增長的，當Context長度足夠長時，它的大小就會占主導地位，可能超出一張卡甚至一台機（8張卡）的總顯存量。

在GPU上部署模型的原則是：能一張卡部署的，就不要跨多張卡；能一台機部署的，就不要跨多台機。這是因為「卡內通信帶寬 > 卡間通信帶寬 > 機間通信帶寬」，由於「木桶效應」，模型部署時跨的設備越多，受設備間通信帶寬的拖累就越大，事實上即便是單卡H100內SRAM與HBM的帶寬已經達到了3TB/s，但對於短Context來說這個速度依然還是推理的瓶頸，更不用說更慢的卡間、機間通信了。

所以，減少KV快取的目的就是要實現在更少的設備上推理更長的Context，或者在相同的Context長度下讓推理的batch size更大，從而實現更快的推理速度或者更大的吞吐總量。當然，最終目的都是為了實現更低的推理成本。

MQA

MQA（多查詢注意力），是減少KV快取的一次非常樸素的嘗試，首次提出自《Fast Transformer Decoding: One Write-Head is All You Need》這篇2019年的論文，這也意味著早在LLM成為熱門話題之前，減少KV快取就已經是研究人員非常關注的一個議題了。

MQA的思路很簡單，直接讓所有注意力頭（Attention Head）共用同一個K（Key）、V（Value），用公式來說，就是取消了MHA中所有的 \( k \)、\( v \) 的上標 \( (s) \)：

$$
\begin{aligned}
    o_t &= [o_t^{(1)}, o_t^{(2)}, \ldots, o_t^{(h)}] \\
    o_t^{(s)} &= \text{Attention}(q_t^{(s)}, k_{\leq t}, v_{\leq t}) \\
    &\equiv \sum_{i \leq t} \frac{\exp(q_t^{(s)} k_i^\top)}{\sum_{i \leq t} \exp(q_t^{(s)} k_i^\top)} v_i \\
    q_t^{(s)} &= x_t W_q^{(s)} \in \mathbb{R}^{d_k}, \quad W_q^{(s)} \in \mathbb{R}^{d \times d_k} \\
    k_i &= x_i W_k \in \mathbb{R}^{d_k}, \quad W_k \in \mathbb{R}^{d \times d_k} \\
    v_i &= x_i W_v \in \mathbb{R}^{d_v}, \quad W_v \in \mathbb{R}^{d \times d_v}
\end{aligned}
$$

使用MQA的模型包括PaLM、StarCoder、Gemini等。很明顯，MQA直接將KV快取減少到了原來的 \( \frac{1}{h} \)，這是非常顯著的，單從節省顯存角度看已經是極限了。

效果方面，目前看來大部分任務的損失都比較有限，且MQA的支持者相信這部分損失可以通過進一步訓練來彌補回來。此外，注意到MQA由於共用了K、V，將會導致注意力的參數量減少了將近一半，而為了模型總參數量的不變，通常會相應地增大FFN（Feed Forward Network）/GLU（Gated Linear Unit）的規模，這也能彌補一部分效果損失。

GQA

然而，也有人擔心MQA對KV快取的壓縮太嚴重，以至於會影響模型的學習效率以及最終效果。為此，一個MHA與MQA之間的過渡版本GQA（分組查詢注意力）應運而生，出自論文《GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints》，是去年的工作。

事後看來，GQA的思想也很樸素，它就是將所有Head分為 \( g \) 個組（ \( g \) 可以整除 \( h \)），每組共用同一對K、V，用數學公式表示為：

$$
\begin{aligned}
    o_t &= [o_t^{(1)}, o_t^{(2)}, \ldots, o_t^{(h)}] \\
    o_t^{(s)} &= \text{Attention}(q_t^{(s)}, k_{\leq t}^{(\lceil s\frac{g}{h} \rceil)}, v_{\leq t}^{(\lceil s\frac{g}{h} \rceil)}) \\
    &\equiv \sum_{i \leq t} \frac{\exp(q_t^{(s)} k_{(\lceil s\frac{g}{h} \rceil)i}^\top)}{\sum_{i \leq t} \exp(q_t^{(s)} k_{(\lceil s\frac{g}{h} \rceil)i}^\top)} v_{(\lceil s\frac{g}{h} \rceil)i} \\
    q_t^{(s)} &= x_t W_q^{(s)} \in \mathbb{R}^{d_k}, \quad W_q^{(s)} \in \mathbb{R}^{d \times d_k} \\
    k_i^{(\lceil s\frac{g}{h} \rceil)} &= x_i W_k^{(\lceil s\frac{g}{h} \rceil)} \in \mathbb{R}^{d_k}, \quad W_k^{(\lceil s\frac{g}{h} \rceil)} \in \mathbb{R}^{d \times d_k} \\
    v_i^{(\lceil s\frac{g}{h} \rceil)} &= x_i W_v^{(\lceil s\frac{g}{h} \rceil)} \in \mathbb{R}^{d_v}, \quad W_v^{(\lceil s\frac{g}{h} \rceil)} \in \mathbb{R}^{d \times d_v}
\end{aligned}
$$

這裡的 \( \lceil \cdot \rceil \) 是向上取整符號。GQA提供了從MHA到MQA的自然過渡，當 \( g = h \) 時就是MHA，當 \( g = 1 \) 時就是MQA。當 \( 1 < g < h \) 時，它只將KV快取壓縮到 \( \frac{g}{h} \)，壓縮率不如MQA，但同時也提供了更大的靈活性，效果上更有保障。

GQA最知名的使用者，大概是Meta開源的LLAMA2-70B，以及LLAMA3全系列，此外使用GQA的模型還有TigerBot、DeepSeek-V1、StarCoder2、Yi、ChatGLM2、ChatGLM3等，相比使用MQA的模型更多（ChatGLM雖然在它的介紹中說自己是MQA，但實際是 \( g=2 \) 的GQA）。
在llama2/3-70B中，GQA的 \( g=8 \)，其他用了GQA的同體量模型基本上也保持了這個設置，這並非偶然，而是同樣出於推理效率的考慮。我們知道，70B這個體量的模型，如果不進行極端的量化，那麼不可能部署到單卡（A100/H100 80G）上。單卡不行，那麼就能單機了，一般情況下一台機可以裝8張卡，剛才我們說了，Attention的每個Head實際上是獨立運算然後拼接起來的，當 \( g=8 \) 時，正好可以每張卡負責計算一組K、V對應的Attention Head，這樣可以在盡可能保證K、V多樣性的同時最大程度上減少卡間通信。

MLA

有了MHA、MQA、GQA的鋪墊，我們理解MLA（多頭潛在注意力）就相對容易一些了。DeepSeek-V2的技術報告裡是從低秩投影的角度引入MLA的，以至於有部分讀者提出「為什麼LoRA提出這麼久了，直到MLA才提出對KV快取低秩分解的做法」之類的疑問。

然而，筆者認為低秩投影這個角度並不貼近本質，因為要說低秩投影的話，事實上只要我們將GQA的所有K、V疊在一起，就會發現GQA也相當於在做低秩投影：

$$
\begin{aligned}
    [k_i^{(1)}, \dots, k_i^{(g)}, v_i^{(1)}, \dots, v_i^{(g)}] &\equiv c_i \in \mathbb{R}^{g(d_k+d_v)} \\
    c_i &= x_i[W_k^{(1)}, \dots, W_k^{(g)}, W_v^{(1)}, \dots, W_v^{(g)}] \in \mathbb{R}^{d \times g(d_k+d_v)}
\end{aligned}
$$

這裡我們將所有\( k_i^{(s)} \)和\( v_i^{(s)} \)拼在一起記為\( c_i \)，相應的投影矩陣也拼在一起記為\( W_c \)。注意到一般都有\( d_c = g(d_k + d_v) < d \)，所以\( x_i \)到\( c_i \)的轉換就是一個低秩投影。因此，MLA的本質改進不是低秩投影，而是在低秩投影之後的處理。

Part 1

GQA在投影之後做了什麼呢？首先它將向量對半分為兩份分別作為K、V，然後每一份又均分為\( g \)份，每一份複製\( h/g \)次，以此來「湊」夠\( h \)個Attention Head所需要的K、V。我們知道分割、複製都是簡單的線性變換，所以MLA的第一個想法是將這些簡單的線性變換換成一般的線性變換，以增強模型的能力：

$$
\begin{aligned}
    o_t &= [o_t^{(1)}, o_t^{(2)}, \ldots, o_t^{(h)}] \\
    o_t^{(s)} &= \text{Attention}(q_t^{(s)}, k_t^{(s)}, v_t^{(s)}) \\
    &\equiv \sum_{i \leq t} \frac{\exp(q_t^{(s)} k_i^{(s)\top})}{\sum_{i \leq t} \exp(q_t^{(s)} k_i^{(s)\top})} v_i^{(s)} \\
    k_i^{(s)} &= c_i W_k^{(s)} \in \mathbb{R}^{d_k}, \quad W_k^{(s)} \in \mathbb{R}^{d_c \times d_k} \\
    v_i^{(s)} &= c_i W_v^{(s)} \in \mathbb{R}^{d_v}, \quad W_v^{(s)} \in \mathbb{R}^{d_c \times d_v} \\
    q_t^{(s)} &= x_t W_q^{(s)} \in \mathbb{R}^{d_k}, \quad W_q^{(s)} \in \mathbb{R}^{d \times d_k}
\end{aligned}
$$

理論上這樣能增加模型能力，但別忘了GQA的主要目的是減少KV快取。MLA的這個做法，通過不同的投影矩陣讓所有的K、V Head都變得各不相同，那麼KV快取的大小就恢復成跟MHA一樣大了，違背了GQA的初衷。

對此，MLA發現，我們可以結合Dot-Attention的具體形式，通過一個簡單但不失巧妙的恒等變換來規避這個問題。首先，在訓練階段還是照常進行，此時優化空間不大；然後，在推理階段，我們利用

$$
q_t^{(s)}k_i^{(s)\top} = (x_t W_q^{(s)})(c_i W_k^{(s)})^\top = x_t (W_q^{(s)} W_k^{(s)\top}) c_i^\top
$$

這意味著推理階段，我們可以將\( W_q^{(s)} W_k^{(s)\top} \)合併起來作為Q的投影矩陣，那麼\( c_i \)則取代了原本的\( k_i^{(s)} \)，同理，在\( o_t \)後面我們還有一個投影矩陣，於是\( v_i^{(s)} \)的\( W_v^{(s)} \)也可以吸收到後面的投影矩陣中去，於是等效地\( v_i^{(s)} \)也可以用\( c_i \)代替，也就是說此時KV快取只需要存下所有的\( c_i \)就行，而不必存下所有的\( k_i^{(s)} \)、\( v_i^{(s)} \)。注意到\( c_i \)跟\( s \)無關，也就是說是所有頭共用的，即MLA在推理階段它可以恒等變換為一個MQA。

再次強調，本文的主題一直都是減少KV快取，那到目前為止，MLA做到了什麼呢？答案是通過不同的投影矩陣來增強了GQA的能力，並且推理時可以保持同樣大小的KV快取。那麼反過來，如果我們只需要跟GQA相近的能力，那麼是不是就可以再次減少KV快取了？換言之，\( d_c \)沒必要取\( g(d_k + d_v) \)，而是取更小的值（DeepSeek-V2取了512），從而進一步壓縮KV快取，這就是MLA的核心思想。

（注：這裡有一個細節，就是\( W_q^{(s)} W_k^{(s)\top} \)合併成一個矩陣的恒等變換，理論上只有在無限精度下才成立，實際上如果我們使用單精度尤其是BF16的話，經過變換後的精度損失往往還是挺明顯的，經過多層累積後可能放大到比較可觀的程度，這裡可能要根據實際誤差看要不要做一些後處理。）

Part 2

一切似乎都很完美，看上去一個又好又省的理想設計就要出爐了。不過別急，當我們再深入思考一下就會發現，到目前為止的MLA有一個難以繞開的缺陷——不相容RoPE（旋轉位置編碼）。

剛才我們提到，MLA之所以能保持與GQA相同大小的KV緩存，其關鍵一步是“將\(W_q^{(s)} W_k^{(s)^\top}\)”合併成一個與位置無關的矩陣作為Q的投影矩陣。但是，如果加入RoPE（旋轉位置編碼），這一步就無法實現了。RoPE是一個與位置相關的\(d_k \times d_k\)的分塊對角矩陣\(R_m\)，滿足\(R_m R_n^\top = R_{m-n}\)。在MLA中加入RoPE後，會使得\(W_q^{(s)} W_k^{(s)^\top}\)之間多插入了一個項\(R_{t-i}\)：

$$
q_i^{(s)} = x_i W_q^{(s)} R_i, \quad k_i^{(s)} = c_i W_k^{(s)} R_i \\
q_t^{(s)} k_i^{(s)^\top} = (x_t W_q^{(s)} R_t)(c_i W_k^{(s)} R_i)^\top = x_t (W_q^{(s)} R_{t-i} W_k^{(s)^\top}) c_i^\top
$$

這裡的\(W_q^{(s)} R_{t-i} W_k^{(s)^\top}\)就無法合併為一個固定的投影矩陣了（與位置差\(t-i\)相關），從而MLA的想法無法與RoPE結合實現。

前段時間，我也很榮幸地與DeepSeek團隊討論過這個問題，但這個問題可以說非常本質，所以當時我實際上也沒能提出什麼有效的建議。最簡單的方式是放棄RoPE，換用其他基於Attention Bias的位置編碼，如ALIBI，但DeepSeek的實驗顯示它明顯不如RoPE（注意，MLA不是不能加RoPE，而是加了RoPE之後無法用恒等變換技巧來減少KV緩存）。我也提議過換Sandwich，它不像ALIBI單調衰減到負無窮，估計效果會好些，但感覺是治標不治本。還有一個折中的辦法是將\(q_i^{(s)}\)的輸入也改為\(c_i\)，然後RoPE加在\(c_i\)之後，即：

$$
q_i^{(s)} = c_i R_i W_q^{(s)}, \quad k_i^{(s)} = c_i R_i W_k^{(s)}
$$
這樣\(R_i\)就可以吸收到\(c_i\)中去，但這樣就沒有\(R_m R_n^\top = R_{m-n}\)的運算了，此時的RoPE不再是通過絕對位置實現相對位置，而單純是加在Q、K上的絕對位置資訊，讓模型自己想辦法提煉相對位置資訊。
最後發布的MLA，採取了一種混合的方法——每個注意力頭的Q、K新增\( d_r \)個維度用來添加RoPE，其中K新增的維度每個頭共用：

$$
o_t^{(s)} = \text{Attention}(q_t^{(s)}, k_{\leq t}^{(s)}, v_{\leq t}^{(s)}) \\
q_t^{(s)} = [x_t W_{qc}^{(s)}, x_t W_{qr}^{(s)} R_t] \in \mathbb{R}^{d_k + d_r}, \quad W_{qc}^{(s)} \in \mathbb{R}^{d \times d_k}, \quad W_{qr}^{(s)} \in \mathbb{R}^{d \times d_r} \\
k_i^{(s)} = [c_i W_{kc}^{(s)}, x_i W_{kr}^{(s)} R_i] \in \mathbb{R}^{d_k + d_r}, \quad W_{kc}^{(s)} \in \mathbb{R}^{d_c \times d_k}, \quad W_{kr}^{(s)} \in \mathbb{R}^{d \times d_r} \\
v_i^{(s)} = c_i W_{vc}^{(s)} \in \mathbb{R}^{d_v}, \quad W_{vc}^{(s)} \in \mathbb{R}^{d_c \times d_v}
$$

這樣一來，沒有RoPE的維度就可以重複「Part 1」的操作，在推理時KV快取只需要存下所有的\( c_i \)，新增的帶RoPE的維度就可以用來補充位置資訊，並且由於所有頭共用，所以也就只有在K快取這裡增加了\( d_r \)個維度，原論文取了\( d_r = \frac{d_k}{2} = 64 \)，相比原本的\( d_c = 512 \)，增加的幅度不大。

Part 3

最後有一個細節，就是MLA的最終版本，還將Q的輸入也改為了低秩投影形式，這與減少KV快取無關，主要是為了減少訓練期間參數量和相應的梯度（原論文說的是啟動值，個人表示不大理解）所占的顯存：

$$
\begin{aligned}
    o_t^{(s)} &= \text{Attention}(q_t^{(s)}, k_{\leq t}^{(s)}, v_{\leq t}^{(s)}) \\
    q_t^{(s)} &= [c'_i W_{qc}^{(s)}, c'_i W_{qr}^{(s)} R_i] \in \mathbb{R}^{d_k + d_r}, \quad W_{qc}^{(s)} \in \mathbb{R}^{d'_c \times d_k}, \quad W_{qr}^{(s)} \in \mathbb{R}^{d'_c \times d_r} \\
    k_i^{(s)} &= [c_i W_{kc}^{(s)}, x_i W_{kr}^{(s)} R_i] \in \mathbb{R}^{d_k + d_r}, \quad W_{kc}^{(s)} \in \mathbb{R}^{d_c \times d_k}, \quad W_{kr}^{(s)} \in \mathbb{R}^{d \times d_r} \\
    v_i^{(s)} &= c_i W_{vc}^{(s)} \in \mathbb{R}^{d_v}, \quad W_{vc}^{(s)} \in \mathbb{R}^{d_c \times d_v} \\
    c'_i &= x_i W'_{c} \in \mathbb{R}^{d'_c}, \quad W'_{c} \in \mathbb{R}^{d \times d'_c}
\end{aligned}
$$

注意\( k_i^{(s)} \)中的第二項，帶RoPE的部分，其輸入還是\( x_i \)而不是\( c_i \)，這裡保持了原論文的設定，不是筆誤，\( d'_c \)原論文的取值是1536，跟\( d_c = 512 \)不同。同時，我們把帶RoPE的MHA放在下面，方便大家對比：

$$
\begin{aligned}
    o_t^{(s)} &= \text{Attention}(q_t^{(s)}, k_{\leq t}^{(s)}, v_{\leq t}^{(s)}) \\
    q_t^{(s)} &= x_i W_q^{(s)} R_i \in \mathbb{R}^{d_k}, \quad W_q^{(s)} \in \mathbb{R}^{d \times d_k} \\
    k_i^{(s)} &= x_i W_k^{(s)} R_i \in \mathbb{R}^{d_k}, \quad W_k^{(s)} \in \mathbb{R}^{d \times d_k} \\
    v_i^{(s)} &= x_i W_v^{(s)} \in \mathbb{R}^{d_v}, \quad W_v^{(s)} \in \mathbb{R}^{d \times d_v}
\end{aligned}
$$

可以發現，其實在訓練階段，除了多了一步低秩投影以及只在部分維度加RoPE外，MLA與Q、K的Head Size由\( d_k \)換成\( d_k + d_r \)的MHA基本無異。

小結

本文簡單概述了多頭注意力的演變歷程，特別是從MHA向MQA、GQA，最終到MLA的變化理念，最後詳細展開了對MLA的介紹。在本文中，MLA被視為GQA的一般化，它用投影矩陣的方式替代了GQA的分割、重複，並引入了一個恆等變換技巧來進一步壓縮KV快取，同時採用了一種混合方法來相容RoPE。總的來說，MLA稱得上是一種非常實用的注意力變體，其創新在於如何在保持推理效率的同時減少存儲和計算資源的需求。

這種設計不僅顯示了深入的技術見解，也反映出當前大型語言模型推理優化的趨勢——即在最大限度減少資源消耗的同時，盡可能保持或提升模型的性能。未來的研究可能會繼續在這條路上進行探索，尋找更有效的方法來解決推理時的瓶頸問題，特別是在處理更大規模的數據和模型時。這些技術的進步將對AI的實用性和可達性產生深遠的影響。

Microsoft: 用於TTS的語言模型VALL-E

微軟引入了一種用於文字轉語音合成 (TTS) 的語言模型方法。具體而言,微軟使用從現成的神經音訊編解碼器模型中得到的離散編碼訓練了一個神經編解碼器語言模型 (稱為 VALL-E),並將 TTS 視為條件語言建模任務,而不是像之前的工作那樣進行連續訊號迴歸。在預訓練階段,微軟將 TTS 訓練資料擴大到 60K 小時的英語語音,比現有系統大數百倍。VALL-E 出現了上下文學習能力,只需一個 3 秒鐘未見說話者的錄音作為聲學提示,就可以用於合成高品質的個人化語音。實驗結果表明,在語音自然度和說話者相似度方面,VALL-E 顯著優於最先進的零樣本 TTS 系統。此外,微軟發現 VALL-E 可以在合成中保留聲學提示中的說話者情緒和聲學環境。範例請見 https://aka.ms/valle。

與以前的 pipeline (例如音素→Mel 頻譜圖→波形) 不同,VALL-E 的 pipeline 是音素→離散編碼→波形。VALL-E 根據音素和聲學編碼提示產生離散音訊編解碼器編碼,分別對應於目標內容和說話者的聲音。VALL-E 直接支援各種語音合成應用,例如零樣本 TTS、語音編輯以及與 GPT-3 等其他生成式 AI 模型結合的內容創作。

1 引言

神經網路和端到端建模的發展在過去十年中給語音合成帶來了巨大突破。目前,級聯文字轉語音 (TTS) 系統通常利用一個聲學模型和一個聲碼器的 pipeline,使用 Mel 頻譜圖作為中間表示。儘管先進的 TTS 系統能夠從單個或多個說話者合成高品質的語音,但它仍然需要來自錄音室的高品質乾淨資料。從網際網路爬取的大規模資料無法滿足要求,而且總是導致性能下降。由於訓練資料相對較少,目前的 TTS 系統在泛化方面仍然存在不足。在零樣本場景中,說話者相似度和語音自然度會大幅下降。

為了解決零樣本 TTS 問題,現有工作利用說話者自適應和說話者編碼方法,需要額外的微調、複雜的預設計特徵或繁重的結構工程。

與其為這個問題設計一個複雜而特定的網路,不如盡可能地用大量和多樣化的資料訓練模型,這受到文字合成領域成功的啟發。近年來,隨著文字語言模型中資料增加,人們見證了顯著的性能提升,從 16GB 的未壓縮文字,到 160GB,再到 570GB,最後是大約 1TB。將這一成功轉移到語音合成領域,微軟引入了 VALL-E,這是第一個利用大規模、多樣化和多說話者語音資料的基於語言模型的 TTS 框架。如圖 1 所示,為了合成個人化語音 (例如零樣本 TTS),VALL-E 以 3 秒鐘註冊錄音的聲學編碼和音素提示為條件產生相應的聲學編碼,分別約束說話者和內容資訊。最後,產生的聲學編碼用於合成最終的波形。來自音訊編解碼器模型的離散聲學編碼使我們能夠將 TTS 視為條件編解碼器語言建模,並且可以利用基於提示的大型模型技術 (如 GPT)用於 TTS 任務。聲學編碼還允許我們在推理期間使用不同的採樣策略在 TTS 中產生不同的合成結果。

微軟使用 LibriLight 訓練 VALL-E,這是一個包含 60K 小時英語語音的語料庫,有 7000 多個獨特的說話者。原始資料是純音訊的,因此微軟採用語音識別模型產生轉錄。與之前的 TTS 訓練資料集 (如 LibriTTS)相比,微軟的資料包含更多噪音和不準確的轉錄,但提供了不同的說話者和韻律。微軟相信,通過利用大數據,所提出的方法對噪音具有魯棒性,並且泛化性能良好。值得注意的是,現有的 TTS 系統總是使用幾十小時的單說話者資料或幾百小時的多說話者資料進行訓練,比 VALL-E 小幾百倍。表 1 總結了 VALL-E 的創新,即一種用於 TTS 的語言模型方法,使用音訊編解碼器編碼作為中間表示,利用大規模和多樣化的資料,從而具有強大的上下文學習能力。

微軟在 LibriSpeech 和 VCTK 資料集上評估 VALL-E,其中所有測試說話者在訓練語料庫中都是未見過的。在語音自然度和說話者相似度方面,VALL-E 顯著優於最先進的零樣本 TTS 系統,在 LibriSpeech 上的比較平均意見得分 (CMOS) 提高了 +0.12,相似度平均意見得分 (SMOS) 提高了 +0.93。VALL-E 在 VCTK 上也以 +0.11 SMOS 和 +0.23 CMOS 的改進擊敗了基準。它甚至在 VCTK 上獲得了 +0.04 的 CMOS 分數,表明未見說話者的合成語音與人類錄音一樣自然。此外,定性分析表明,VALL-E 能夠用相同的文字和目標說話者合成不同的輸出,這可能有利於為語音識別任務創建偽資料。我們還發現,VALL-E 可以保持聲學提示中的聲學環境 (例如混響) 和情緒 (例如憤怒)。

總之,微軟做出了以下貢獻:

- 微軟提出了 VALL-E,這是第一個像 GPT-3 那樣具有強大上下文學習能力的 TTS 框架,它將 TTS 視為語言模型任務,使用音訊編解碼器編碼作為中間表示,取代傳統的 Mel 頻譜圖。它具有上下文學習能力,並支援基於提示的方法進行零樣本 TTS,不需要像之前的工作那樣進行額外的結構工程、預設計的聲學特徵和微調。

- 微軟通過利用大量半監督資料在說話者維度上構建了一個廣義 TTS 系統,表明簡單地擴大半監督資料的規模在 TTS 中被低估了。 

- VALL-E 能夠用相同的輸入文字提供多樣化的輸出,並保持聲學提示中的聲學環境和說話者情緒。

- 微軟驗證了 VALL-E 通過提示在零樣本場景中合成具有高說話者相似度的自然語音。評估結果表明,VALL-E 在 LibriSpeech 和 VCTK 上顯著優於最先進的零樣本 TTS 系統。

微軟鼓勵讀者在範例頁面 https://aka.ms/valle 上聽樣本。

3 背景:語音量化 

由於音訊通常儲存為一系列 16 位整數值,生成模型需要在每個時間步輸出 $2^{16}=65,536$ 個機率才能合成原始音訊。此外,超過萬赫茲的音訊採樣率導致了極長的序列長度,使得原始音訊合成更加難以處理。為此,需要進行語音量化以壓縮整數值和序列長度。μ-law 變換可以將每個時間步量化為 256 個值並重建高品質的原始音訊。它在語音生成模型 (如 WaveNet) 中被廣泛使用,但由於序列長度沒有減少,推理速度仍然很慢。最近,向量量化在自監督語音模型中被廣泛應用於特徵提取,如 vq-wav2vec 和 HuBERT。後續工作表明自監督模型中的編碼也可以重建內容,並且推理速度比 WaveNet 更快。然而,說話者身份已經被丟棄,重建品質也很低。AudioLM 在自監督模型的 k-means 權杖和神經編解碼器模型的聲學權杖上訓練語音到語音語言模型,從而實現高品質的語音到語音生成。

在本文中,微軟沿用 AudioLM 的做法,利用神經編解碼器模型將語音表示為離散權杖。為了壓縮音訊以進行網路傳輸,編解碼器模型能夠將波形編碼為離散聲學編碼並重建高品質波形,即使說話者在訓練中是未見過的。與傳統的音訊編解碼器方法相比,基於神經網路的編解碼器在低位元速率下明顯更好,微軟相信量化後的權杖包含了關於說話者和錄音條件的充分資訊。與其他量化方法相比,音訊編解碼器具有以下優勢:1)它包含豐富的說話者資訊和聲學資訊,與 HuBERT 編碼相比,可以在重建中維持說話者身份。2)有一個現成的編解碼器解碼器將離散權杖轉換為波形,不需要像在頻譜上操作的基於 VQ 的方法那樣額外訓練聲碼器。3)它可以減少時間步長以提高效率,解決 μ-law 變換中的問題。 

微軟採用預先訓練的神經音訊編解碼器模型 EnCodec 作為標記器。EnCodec 是一個卷積編碼器-解碼器模型,其輸入和輸出都是可變位元速率的 24 kHz 音訊。編碼器為 24 kHz 的輸入波形產生 75 Hz 的嵌入,這是採樣率的 320 倍減少。每個嵌入由殘差向量量化 (RVQ) 建模,其中選擇了 8 個具有 1024 個條目的階層量化器,如圖 2 所示。此配置對應於用於 24 kHz 音訊重建的 6K 位元速率下的 EnCodec。在此設定下,給定一個 10 秒的波形,離散表示是一個 \(750 \times 8\) 的矩陣,其中 \(750 = \lfloor \frac{24,000 \times 10}{320} \rfloor\) 是下採樣的時間步,8是量化器的數量。選擇其他位元速率設定也是可以的。更大的位元速率對應更多的量化器和更好的重建品質。例如,如果選擇 12K 位元速率下的 EnCodec,則需要 16 個量化器,10 秒波形對應一個 \(750 \times 16\) 的矩陣。使用所有量化器的離散編碼,EnCodec 的卷積解碼器生成實值嵌入並以 24 kHz 重建波形。

4 VALL-E

4.1 問題公式化:將 TTS 視為條件編解碼器語言建模 

給定一個資料集 \(\mathcal{D}={x^i,y^i\\}\),其中 \(y\) 是一個音訊樣本, \(x={x_0,x_1,\ldots,x_L\\}\) 是其對應的音素轉錄,微軟使用預先訓練的神經編解碼器模型對每個音訊樣本進行編碼為離散聲學編碼,表示為 \(\text{Encodec}(y)=C^{T\times 8}\),其中 \(C\) 表示二維聲學編碼矩陣, \(T\) 是下採樣的話語長度。每個聲學編碼矩陣的行向量 \(c_{t,:}\) 表示第 \(t\) 幀的 8 個編碼,每個聲學編碼矩陣的列向量 \(c_{:,j}\) 表示來自第 \(j\) 個碼本的編碼序列,其中 \(j\in{1,\ldots,8\\}\)。量化後,神經編解碼器解碼器能夠重建波形,表示為 \(\text{Decodec}(C)\approx\hat{y}\)。

零樣本 TTS 要求模型為未見過的說話者合成高品質的語音。

在這項工作中,微軟將零樣本 TTS 視為條件編解碼器語言建模任務。微軟訓練了一個神經語言模型,以在給定音素序列 \(x\) 和聲學提示矩陣 \(\tilde{C}^{T'\times 8}\) 的條件下產生聲學編碼矩陣 \(C\),優化目標是 \(\max p(C|x,\tilde{C})\)。這裡, \(\tilde{C}\) 是通過相同的神經編解碼器從註冊錄音中得到的。微軟期望神經語言模型分別從音素序列和聲學提示中學習提取內容和說話者資訊。在推理期間,給定一個音素序列和一個未見過的說話者的 3 秒鐘註冊錄音,聲學編碼矩陣首先由訓練好的語言模型估計,然後神經編解碼器解碼器合成高品質的語音。

 4.2 訓練:條件編解碼器語言建模

神經語音編解碼器模型允許我們對離散音訊表示進行操作。由於神經編解碼器模型中的殘差量化,權杖具有分層結構:來自先前量化器的權杖恢復聲學屬性 (如說話者身份),而連續的量化器學習精細的聲學細節。每個量化器都經過訓練以對來自先前量化器的殘差進行建模。受此啟發,微軟以分層方式設計了兩個條件語言模型。 

對於來自第一個量化器 \(c_{:,1}\) 的離散權杖,微軟訓練了一個自迴歸 (AR) 解碼器語言模型。它以音素序列 \(x\) 和聲學提示 \(\tilde{C}_{:,1}\) 為條件,公式為:

$$p(c_{:,1}|x,\tilde{C}_{:,1};\theta_{AR})=\prod_{t=0}^T p(c_{t,1}|c_{<t,1},\tilde{c}_{:,1},x;\theta_{AR})$$

由於 VALL-E 是一個僅解碼器 LM,所以 \(\tilde{c}_{:,1}\) 和 \(c_{:,1}\) 的連接是一個完整的序列,在訓練中微軟不區分它們或插入特定的權杖。在推理中只預測 \(c_{:,1}\) 而前綴 \(\tilde{c}_{:,1}\) 是給定的。

對於第二個到最後一個量化器的離散權杖 \(c_{:,j\in[2,8]}\),微軟訓練了一個非自迴歸 (NAR) 語言模型。由於權杖在 NAR 方式下無法相互訪問,為了約束說話者身份,聲學提示矩陣 \(\tilde{C}\) 被用作聲學提示。因此,該模型以音素序列 \(x\)、聲學提示 \(\tilde{C}\) 和屬於先前碼本的預測聲學權杖 \(C_{:,<j}\) 為條件:

$$p(C_{:,2:8}|x,\tilde{C};\theta_{NAR})=\prod_{j=2}^8 p(c_{:,j}|C_{:,<j},x,\tilde{C};\theta_{NAR})$$

AR 模型和 NAR 模型的組合在語音品質和推理速度之間提供了很好的權衡。一方面,產生的語音速率應該與註冊錄音一致,並且由於不同說話者的說話速度可能非常不同,因此很難為不同的說話者訓練長度預測器。在這種情況下,AR 模型在聲學序列長度預測方面具有更大的靈活性,是一個更自然的選擇。另一方面,對於連續階段,由於輸出槽的數量遵循第一階段的序列長度,NAR 可以將時間複雜度從 \(O(T)\) 減少到 \(O(1)\)。總的來說, \(C\) 的預測可以建模為:

$$p(C|x,\tilde{C};\theta)=p(c_{:,1}|\tilde{C}_{:,1},X;\theta_{AR})\prod_{j=2}^8 p(c_{:,j}|c_{:,<j},x,\tilde{C};\theta_{NAR})$$

4.2.1 自迴歸編解碼器語言建模 

自迴歸語言模型產生來自第一個量化器的權杖。它包括音素嵌入 \(W_x\)、聲學嵌入 \(W_a\)、transformer 解碼器和預測層。為了產生具有特定內容的語音,微軟使用音素序列作為語言模型的音素提示。因此,模型輸入是 \(x\) 和 \(c_{:,1}\) 的連接,在每個序列之後附加兩個特殊的 <EOS> 權杖。微軟分別為提示和輸入權杖計算正弦位置嵌入。對於因果 transformer 模型,每個權杖 \(c_{t,1}\) 可以參與 \((x,c_{\leq t,1})\),如圖 3 的左側所示。該模型被優化以最大化第一個碼本中下一個權杖的機率。微軟共享輸出投影層的參數與聲學嵌入 \(W_a\) 的參數。

在 AR 模型中,微軟沒有在訓練中明確提取音訊片段作為提示。訓練過程是純因果語言模型訓練。通過這種方式,任何前綴序列 \(c_{<t,1}\) 都被視為序列 \(c_{\geq t,1}\) 後面部分的提示。在推理過程中,給定註冊錄音,應該將註冊錄音的音素序列和要合成的音素序列連接在一起。同時,註冊錄音的聲學權杖序列被用作 AR 解碼中的前綴,如公式 1 所示。微軟將在實驗中研究這種設置的優越性。  

4.2.2 非自迴歸編解碼器語言建模

當微軟通過 AR 模型獲得第一個量化器編碼時,使用非自迴歸 (NAR) 模型產生其他七個量化器的編碼。NAR 模型的架構與 AR 模型類似,只是它包含八個獨立的聲學嵌入層。在每個訓練步驟中,隨機採樣一個訓練階段 \(i\in[2,8]\)。模型被訓練以最大化來自第 \(i\) 個量化器碼本的聲學權杖。來自第 1 階段到第 \(i-1\) 階段的聲學權杖被嵌入並相加作為模型輸入:

$$e_{c_{t,j}}=W_a^j c_{t,j}$$ 

$$e_{c_t}=\sum_{j=1}^{i-1} e_{c_{t,j}}$$

其中表示索引選擇。音素序列也被視為語言模型的提示。此外,為了克隆給定說話者的獨特聲音,微軟還使用來自註冊語音的聲學權杖作為聲學提示。具體來說,首先用神經編解碼器模型將註冊語音標記為 \(\tilde{C}^{T\times 8}\)。來自所有八個碼本的嵌入表示相加為聲學提示 \(e_{\tilde{c}_t}=\sum_{j=1}^8 e_{\tilde{c}_{t,j}}\)。為了預測來自第 $i$ 個碼本的聲學權杖,transformer 輸入是 \((e_x,e_{\tilde{c}},e_{c_{:,<i}})\) 的串聯。位置嵌入也是分別為提示和聲學序列計算的。目前階段 \(i\) 通過自適應層歸一化 (AdaLN) 算子注入到網路中,即 \(\text{AdaLN}(h,i)=a_i\text{LayerNorm}(h)+b_i\),其中 \(h\) 是中間激活, \(a_i\) 和 \(b_i\) 是從階段嵌入的線性投影獲得的。與 AR 不同的是,NAR 模型允許每個權杖在自注意力層中參與所有輸入權杖。微軟還共享聲學嵌入層和輸出預測層的參數,這意味著第 \(j\) 個預測層的權重與第 \(j+1\) 個聲學嵌入層相同。 

4.3 推理:通過提示進行上下文學習

上下文學習是基於文字的語言模型的一種令人驚訝的能力,它能夠在不進行額外參數更新的情況下預測未見輸入的標籤。對於 TTS 而言,如果模型能夠在沒有微調的情況下為未見說話者合成高品質的語音,則認為該模型具有上下文學習能力。然而,現有 TTS 系統的上下文學習能力並不強,因為它們要麼需要額外的微調,要麼在未見說話者上顯著下降。

對於語言模型,提示是在零樣本場景中實現上下文學習所必需的。微軟設計提示和推理如下。首先將文字轉換為音素序列,並將註冊錄音編碼為聲學矩陣,形成音素提示和聲學提示。AR 模型和 NAR 模型都使用這兩個提示。對於 AR 模型,微軟在給定提示的條件下使用基於採樣的解碼,因為觀察到 beam search 可能會導致 LM 陷入無限循環。此外,基於採樣的方法可以顯著增加輸出的多樣性。對於 NAR 模型,微軟使用貪婪解碼來選擇機率最高的權杖。最後,使用神經編解碼器解碼器在給定八個編碼序列的條件下產生波形。聲學提示可能與要合成的語音在語義上相關或不相關,從而產生兩種情況:

VALL-E:微軟的主要興趣是為未見過的說話者產生給定的內容。給模型一個文字句子、一段註冊語音及其相應的轉錄。將註冊語音的轉錄音素附加到給定句子的音素序列前作為音素提示,並使用註冊語音的第一層聲學權杖 \(\tilde{c}_{:,1}\) 作為聲學前綴。使用音素提示和聲學前綴,VALL-E 產生給定文字的聲學權杖,克隆該說話者的聲音。

VALL-E-continual:在此設置中,使用整個轉錄和話語的前 3 秒作為音素和聲學提示,並要求模型產生後續內容。推理過程與 VALL-E 設置相同,只是註冊語音和產生的語音在語義上是連續的。

5 實驗

5.1 實驗設置

資料集:微軟使用 LibriLight 作為訓練資料,其中包含 60K 小時來自英語有聲讀物的未標記語音。LibriLight 中不同說話者的數量約為 7000。微軟在 960 小時帶標籤的 LibriSpeech 上訓練了一個混合 DNN-HMM ASR 模型。一旦訓練了混合模型,就對未標記的語音資料進行解碼並轉換為最佳的音素級別對齊路徑,其中幀移為 30ms。EnCodec 模型用於為 60K 小時的資料產生聲學編碼矩陣。 

模型:AR 模型和 NAR 模型具有相同的 transformer 架構,包括12層、16個注意力頭、1024維嵌入維度、4096維前饋層和 0.1 的 dropout。LibriLight 中波形的平均長度為 60 秒。在訓練期間,隨機將波形裁剪為 10 秒到 20 秒之間的隨機長度。其相應的音素對齊用作音素提示。微軟刪除了偏對齊音素序列中的連續重複。對於 NAR 聲學提示權杖,從同一話語中選擇 3 秒的隨機片段波形。

使用 16 個 NVIDIA TESLA V100 32GB GPU 對模型進行了 800k 步的訓練,每個 GPU 的批次大小為 6k 個聲學權杖。微軟使用 AdamW 優化器優化模型,在前 32k 次更新中將學習率預熱到 $5\times 10^{-4}$ 的峰值,然後線性衰減。 

基準:微軟選擇 SOTA 零樣本 TTS 模型 YourTTS 作為基準,該模型在 VCTK、LibriTTS 和 TTS-Portuguese 的組合資料集上進行訓練。微軟使用其發布的檢查點。

自動指標:微軟使用 SOTA 說話者驗證模型 WavLM-TDNN 來評估提示 (解壓縮的註冊語音) 和合成語音之間的說話者相似度。WavLM-TDNN 在 VoxSRC Challenge 2021 和 2022 排行榜上獲得第一名。它在 Vox1-O、Vox1-E 和 Vox1-H 上分別達到了 0.383、0.480 和 0.986 的平均等錯誤率 (EER)。WavLM-TDNN 預測的相似度分數在 \([-1,1]\) 的範圍內,其中較大的值表示輸入樣本的相似度較高。 

微軟還評估了模型的合成穩健性。神經 TTS 系統存在穩健性問題,由於注意力對齊錯誤,有時會出現刪除、插入和替換錯誤。微軟對產生的音訊執行 ASR,並計算相對於原始轉錄的字錯誤率 (WER)。在此實驗中,微軟使用在 LibriSpeech 960h 上微調的 HuBERT-Large 模型作為 ASR 模型,這是一個沒有語言模型融合的基於 CTC 的模型。

人工評估:微軟通過眾包計算比較平均意見得分 (CMOS) 和相似度平均意見得分 (SMOS),其中 12 名和 6 名母語者分別被邀請為 CMOS 和 SMOS 貢獻者。SMOS 的量表從 1 到 5,增量為 0.5。CMOS 的範圍從 -3 (新系統比基準差得多) 到 3 (新系統比基準好得多),間隔為 1。CMOS 是語音自然度的指標,SMOS 衡量語音是否與原始說話者的聲音相似。

5.2 LibriSpeech 評估

微軟首先使用 LibriSpeech 進行零樣本 TTS 評估,因為 LibriLight 訓練資料和 LibriSpeech test-clean 資料之間沒有說話者重疊。微軟使用 LibriSpeech test-clean 中長度在 4 到 10 秒之間的樣本,得到 2.2 小時的子集。對於每個樣本合成,VALL-E 隨機選擇同一說話者的另一個話語,並裁剪 3 秒的語音片段作為註冊語音。每個實驗運行三次,報告平均分數。VALL-E-continual 使用真實語音的前 3 秒作為註冊語音。

表 2 顯示了客觀評估結果。微軟首先計算真實語音的 WER 分數和說話者相似度分數作為上限。為了比較說話者相似度,微軟使用測試集中來自同一說話者的語音對。與 YourTTS 基準相比,微軟的模型在穩健性和說話者相似度方面都有顯著優勢,表明產生的語音高度忠實於給定文字和給定註冊語音。此外,在 VALL-E-continual 設置中,字錯誤率可以進一步降低,因為前 3 秒的聲學權杖是從真實語音中提取的。微軟還將穩健性與其他基於語音到語音 LM 的生成模型 (如 GSLM 和 AudioLM) 進行了比較,這些模型使用音訊潛在編碼作為輸入。GSLM 使用 HuBERT 編碼作為輸入,並使用 Tacotron2 模型和 WaveGlow 聲碼器重建波形。微軟運行其開源代碼使用發布的模型並評估結果。由於 HuBERT 編碼丟棄了說話者身份,因此它在說話者得分方面表現很差。對於 AudioLM,微軟列出了他們論文中報告的 WER 分數,該分數是由 Conformer Transducer 模型獲得的。實驗結果表明,在穩健性方面,VALL-E 優於其他基於語音到語音 LM 的生成系統。一個主要原因是 VALL-E 使用偽音素進行訓練,而不是 HuBERT/w2v-BERT 編碼,這與輸入文字相比具有更好的對齊品質。

微軟隨機採樣 LibriSpeech test-clean 中每個說話者的一個話語進行人工評估,得到 40 個測試用例。表 3 顯示了人工評估結果。在 SMOS 方面,VALL-E 非常接近真實語音,表明合成語音與測試中給定的未見說話者相似。它以 +0.93 SMOS 顯著優於基準,證明了 VALL-E 在零樣本場景中的有效性。在自然度方面,VALL-E 以 +0.12 CMOS 擊敗基準,表明所提出的方法可以合成比基準更自然和逼真的語音。

消融研究:在本節中,微軟進行了詳細的消融實驗。首先研究 NAR 模型。微軟訓練了具有不同數量提示的三個 NAR 模型。設置 NAR-no prompt 在沒有任何提示的情況下進行訓練。設置 NAR-phn prompt 僅使用音素序列作為提示進行訓練,設置 NAR-2 prompts 同時使用音素提示和聲學權杖提示作為條件。在評估中,使用真實的第一級聲學權杖作為模型輸入,並計算 WER 和說話者相似度分數。結果列於表 4 中。結果表明,即使聲學輸入權杖是真實的,沒有任何提示的模型在 ASR 和說話者相似度評估方面的表現都很差。當添加音素提示時,WER 從 19.6 大幅降低到 3.0。這表明音素提示主要有助於產生的內容。在 NAR-2 prompts 中,模型可以從聲學權杖提示中學習說話者資訊,從而提高說話者評估品質。

微軟進一步對 AR 模型進行了消融實驗。在這些實驗中,始終使用 NAR-2 prompts 設置作為 NAR 模型。在表 5 中,可以看到,當移除聲學提示 (w/o 聲學提示) 時,它只能獲得 0.236 的說話者相似度得分,表明提示對於說話者身份極其重要。即使 NAR 模型可以看到提示,AR 模型的提示也對說話者相似度有很大貢獻。

5.3 VCTK 評估 

微軟在包含 108 個說話者的 VCTK 上評估模型,其中沒有一個說話者在訓練期間被觀察到。由於 YourTTS 在 VCTK 中將 97 個說話者視為訓練,微軟分別評估了 YourTTS 在全部 107 個說話者和 11 個未見說話者上的性能。對於每個說話者,隨機選擇三個 3s/5s/10s 的話語作為提示,另一個話語的文字作為文字提示。

微軟首先使用前面描述的說話者驗證指標評估兩個模型。從表 6 可以看出,即使基準在訓練中看到了 97 個說話者,VALL-E 也優於基準,表明微軟的模型能夠合成具有更高說話者相似度的語音。當在公平設置下 (11 個說話者) 與基準進行比較時,性能差距變得更大,尤其是當只有 3s 提示可用時。通過比較不同長度的提示,可以看出微軟的模型能夠在提示變長時產生更相似的語音,這與直覺一致。 

微軟採樣 60 個說話者進行人工評估,每個說話者一個話語,其中 11 個是未見說話者,49 個說話者在 YourTTS 中被視為已見。VALL-E 沒有看到任何 60 個說話者。在模型合成期間,每個說話者有一個 3 秒的註冊錄音。表 7 顯示了微軟的方法與基準和真實語音的比較。SMOS 的比較表明,VALL-E 的說話者相似度優於基準,即使基準在訓練中看到了一些說話者。並排 CMOS 評估表明,VALL-E 比 YourTTS 高 +0.23,在語音自然度方面表現明顯更好。此外,VALL-E 在真實語音上獲得 +0.04 CMOS,表明在該資料集上與人類錄音沒有統計學上的顯著差異。與 LibriSpeech 上的評估結果相比,VALL-E 在與真實語音的比較中表現出更好的 CMOS 分數,主要原因是 VCTK 中平均句子長度更短,並且一些真實話語也有嘈雜的環境。在說話者相似度方面,VCTK 更具挑戰性,因為它包含具有各種口音的說話者,而訓練資料和 LibriSpeech 測試資料不包含各種口音的說話者。

5.4 定性分析

多樣性:以前的 TTS 系統在輸入文字和輸出波形之間具有強的一對一映射,因為 Mel 頻譜產生基於每一步的重建,沒有隨機性。由於 VALL-E 使用基於採樣的方法產生離散權杖,對於相同的輸入文字,其輸出由於推理中的隨機性而多樣化。給定一個句子和一個註冊錄音,微軟運行推理過程兩次並在圖 4 中可視化其波形。在圖 4(a) 中,可以觀察到兩個樣本具有不同的長度和短語持續時間,其中第一個具有更快的語音速率。在圖 4(b) 中,可以觀察到兩個樣本的重音不同。第二個輸出用更大的振幅強調單詞 "must",而第一個輸出沒有。微軟在演示頁面上提供了更多樣本。

多樣性對某些下游場景很重要。例如,語音識別總是從具有不同說話者和聲學環境的多樣化輸入中受益,而以前的 TTS 系統無法滿足這一點。考慮到 VALL-E 的多樣性特點,它是產生語音識別偽資料的理想候選。

聲學環境維持:另一個有趣的發現是聲學提示和產生之間的聲學環境一致性。當聲學提示有混響時,VALL-E 也可以合成有混響的語音,而基準輸出乾淨的語音。微軟的解釋是,與基準使用的資料相比,VALL-E 在包含更多聲學條件的大規模資料集上進行訓練,因此 VALL-E 可以在訓練期間學習聲學一致性,而不僅僅是乾淨的環境。微軟在演示頁面上展示了一致性。

說話者情緒維持:情緒 TTS 是語音合成的一個經典子主題,它合成具有所需情緒的語音。傳統方法總是在有監督的情緒 TTS 資料集上訓練模型,其中語音對應於轉錄和情緒標籤。微軟發現 VALL-E 可以在零樣本設置下保留提示中的情緒。微軟從 EmoV-DB 中選擇聲學提示,這是一個包含五種情緒語音的資料集,VALL-E 能夠在語音合成中保持與提示相同的情緒,即使模型沒有在情緒 TTS 資料集上進行微調。微軟在演示頁面上放置了音訊樣本。

6 結論、局限性和未來工作

微軟引入了 VALL-E,這是一種使用音訊編解碼器編碼作為中間表示的 TTS 語言模型方法。微軟使用 60K 小時的語音資料對 VALL-E 進行了預訓練,並展示了在零樣本場景中的上下文學習能力。微軟在 LibriSpeech 和 VCTK 上實現了新的最先進的零樣本 TTS 結果。此外,VALL-E 可以在合成中保持聲學環境和說話者情緒,並在不同的基於採樣的解碼過程中提供多樣化的輸出。

儘管取得了顯著進展,VALL-E 仍然存在幾個問題。

合成穩健性:微軟觀察到,在語音合成中,某些詞可能不清楚、遺漏或重複。這主要是因為音素到聲學語言部分是一個自迴歸模型,其中存在無序的注意力對齊,並且沒有解決該問題的約束。在基於普通 Transformer 的 TTS 中也觀察到了這種現象,通過應用非自迴歸模型或修改建模中的注意力機制來解決。未來,微軟希望利用這些技術來解決這個問題。

資料覆蓋:即使微軟使用 60K 小時的資料進行訓練,它仍然無法覆蓋每個人的聲音,尤其是口音說話者。在 VCTK 上的結果比在 LibriSpeech 上的結果差,也暗示了對口音說話者的覆蓋不足。此外,說話風格的多樣性還不夠,因為 LibriLight 是一個有聲讀物數據集,其中大多數話語都是朗讀風格。未來,微軟將進一步擴大訓練數據,以提高模型在韻律、說話風格和說話者相似度方面的性能。微軟相信,通過我們的方法,通過模型和數據的擴展,零樣本 TTS 任務幾乎可以解決。

模型結構:現在,微軟使用兩個模型來預測不同量化器的編碼。一個有前途的方向是使用一個大型通用模型來預測它們。另一個有趣的方向是在框架中使用完全 NAR 模型來加速模型推理。 

更廣泛的影響:由於 VALL-E 可以合成保持說話者身份的語音,它可能存在濫用模型的潛在風險,例如欺騙聲音識別或模仿特定說話者。為了降低此類風險,可以構建檢測模型來區分音訊片段是否由 VALL-E 合成。在進一步開發模型時,微軟也將付諸實踐微軟 AI 原則。 

以下是我的一些思考:

從技術倫理的角度,我們要警惕 VALL-E 被濫用的風險,如用於製造虛假語音、侵犯他人隱私等。這需要研發人員保持警惕,並制定相應的防範機制。

語音合成只是人機交互的一個方面,如何將其與語音識別、自然語言理解等技術更好地結合,來創造無縫的交互體驗,還有待進一步探索。

從產業化的角度,VALL-E 在客製化語音、有聲內容創作等領域有廣闊的應用前景。但是,如何平衡模型性能和部署成本,如何適配不同的硬體環境,也是工程師們需要考慮的問題。

從學術研究的角度,VALL-E 提出的一些思路,如語言模型範式、提示學習等,可能啟發其他領域的研究,如計算機視覺、機器翻譯等。跨學科的創新,往往能產生意想不到的火花。

展望未來,隨著模型和數據規模的進一步增長,以及範式的不斷創新,人工智慧生成語音與真人的差距會越來越小。這對傳統的配音行業、有聲內容生產模式,都會帶來深遠的影響。我們要積極擁抱變革,也要未雨綢繆,做好應對準備。

Transformer 語言模型內部運作機制的技術

這篇論文是一篇全面性的綜述,詳細介紹了目前用來解釋 Transformer 語言模型內部運作機制的技術。本文聚焦在生成式的 decoder-only 架構。作者在結論部分概述了這些模型的已知內部機制,揭示了當前流行方法和這個領域的積極研究方向之間的聯繫。 

論文出處: https://arxiv.org/pdf/2405.00208

數學符號說明

在本文中，我們採用以下數學符號：

- \( n \)：序列長度

- \( V \)：詞彙表

- \( t = \langle t_1, t_2, \ldots, t_n \rangle \)：輸入的 token 序列

- \( x = \langle x_1, x_2, \ldots, x_n \rangle \)：輸入的 token 嵌入序列

- \( d \)：模型維度

- \( d_h \)：注意力頭維度

- \( d_{FFN} \)：前饋網路維度

- \( H \)：注意力頭數量

- \( L \)：層數

- \( x_{l,i} \in \mathbb{R}^d \)：在位置 \( i \)，第 \( l \) 層的殘差串流狀態

- \( x_{\text{mid},l,i} \in \mathbb{R}^d \)：在位置 \( i \)，第 \( l \) 層，經過注意力區塊後的殘差串流狀態

- \( f_c(x) \in \mathbb{R}^d \)：在最後一個位置，組件 \( c \) 的輸出表示

- \( f_l(x) = x_{l,n} \in \mathbb{R}^d \)：在最後一個位置，第 \( l \) 層的殘差串流狀態

- \( A_{l,h} \in \mathbb{R}^{n \times n} \)：在第 \( l \) 層第 \( h \) 個注意力頭的注意力矩陣

- \( W_{l,h}^Q, W_{l,h}^K, W_{l,h}^V \in \mathbb{R}^{d \times d_h} \)：在第 \( l \) 層第 \( h \) 個注意力頭的 queries、keys 和 values 權重矩陣

- \( W_{l,h}^O \in \mathbb{R}^{d_h \times d} \)：在第 \( l \) 層第 \( h \) 個注意力頭的輸出權重矩陣

- \( W_{l}^{\text{in}} \in \mathbb{R}^{d \times d_{FFN}}, W_{l}^{\text{out}} \in \mathbb{R}^{d_{FFN} \times d} \)：在第 \( l \) 層前饋網路的輸入和輸出權重矩陣

- \( W_E \in \mathbb{R}^{d \times |V|} \) 和 \( W_U \in \mathbb{R}^{|V| \times d} \)：嵌入（Embedding）和去嵌入（Unembedding）矩陣

一個僅解碼器的模型 \( f \) 有 \( L \) 層，對一個序列的嵌入 \( x = \langle x_1, x_2, \ldots, x_n \rangle \) 進行操作，這些嵌入表示 tokens \( t = \langle t_1, t_2, \ldots, t_n \rangle \)。每個嵌入 \( x \in \mathbb{R}^d \) 是嵌入矩陣 \( W_E \in \mathbb{R}^{|V| \times d} \) 的一個列向量，其中 \( V \) 是模型詞彙表。中間層的表示，例如在位置 \( i \) 和層 \( l \)，表示為 \( x_{l,i} \)。\( X \in \mathbb{R}^{n \times d} \) 代表將序列 \( x \) 表示為一個矩陣，嵌入堆疊為列。同樣地，對於中間表示，\( X_{l, \leq i} \) 是層 \( l \) 的表示矩陣，最多到位置 \( i \)。

遵循最近關於Transformers可解釋性的文獻，本文採用殘差串流的觀點來呈現架構。在這個觀點下，每個輸入嵌入都會透過注意力區塊和前饋網路區塊的向量相加來更新，產生殘差串流狀態（或中間表示）。最後一層的殘差串流狀態會透過去嵌入矩陣 \( W_U \in \mathbb{R}^{d \times |V|} \) 投影到詞彙空間，並透過softmax函數正規化以獲得詞彙上的機率分布，從中取樣出新的token。

在Transformer層中，LayerNorm是一種常見的操作，用於穩定深度神經網路的訓練過程。給定一個表示 \( z \)，LayerNorm的計算為 \((z-\mu(z))/\sigma(z) \odot \gamma + \beta\)，其中 \( \mu \) 和 \( \sigma \) 分別計算平均值和標準差，\( \gamma \in \mathbb{R}^d \) 和 \( \beta \in \mathbb{R}^d \) 是學習到的逐元素轉換和偏差。LayerNorm可以從幾何角度解釋，將減去平均值的操作視為將輸入表示投影到由法向量 \([1, 1, \ldots, 1] \in \mathbb{R}^d\) 定義的超平面上，然後將結果表示映射到 \(\sqrt{d}\) 範數的超球面。

注意力區塊由多個注意力頭組成。在解碼步驟 \( i \)，每個注意力頭從先前位置（≤ \( i \)）的殘差串流讀取，決定要關注哪些位置，從那些位置收集資訊，最後將其寫入當前的殘差串流。

注意力頭計算如下：

\[ \text{Attn}_{l,h}(X_{l-1, \leq i}) = \sum_{j \leq i} a_{l,h,i,j} x_{l-1,j} W_{l,h}^V W_{l,h}^O = \sum_{j \leq i} a_{l,h,i,j} x_{l-1,j} W_{l,h}^{OV} \]

可學習的權重矩陣 \( W_{l,h}^V \in \mathbb{R}^{d \times d_h} \) 和 \( W_{l,h}^O \in \mathbb{R}^{d_h \times d} \) 合併為 OV 矩陣 \( W_{l,h}^V W_{l,h}^O = W_{l,h}^{OV} \in \mathbb{R}^{d \times d} \)，也稱為 OV (output-value) 電路。對於每個查詢（query）\( i \)，給定當前的鍵（key）（\( \leq i \)），注意力權重計算為：

\[ a_{l,h,i} = \text{softmax}\left( \frac{x_{l-1,i} W_{l,h}^Q (X_{l-1, \leq i} W_{l,h}^K)^T}{\sqrt{d_k}} \right) \]

其中 \( W_{l,h}^Q \in \mathbb{R}^{d \times d_h} \) 和 \( W_{h}^Q (W_{h}^K)^T = W_{h}^{QK} \in \mathbb{R}^{d \times d} \) 組合為 QK (query-key) 電路 \[ W_{h}^Q (W_{h}^K)^T = W_{h}^{QK} \in \mathbb{R}^{d \times d} \]

QK 和 OV 電路可以視為負責分別從殘差串流讀取和寫入的單元。注意力區塊的輸出是各個注意力頭的總和，隨後加回到殘差串流：

\[ \text{Attn}_l(X_{l-1, \leq i}) = \sum_{h=1}^H \text{Attn}_{l,h}(X_{l-1, \leq i}) \]

\[ x_{\text{mid},l,i} = x_{l-1,i} + \text{Attn}_l(X_{l-1, \leq i}) \]

前饋網路區塊由兩個可學習的權重矩陣組成：\( W_{l}^{\text{in}} \in \mathbb{R}^{d \times d_{FFN}} \) 和 \( W_{l}^{\text{out}} \in \mathbb{R}^{d_{FFN} \times d} \)。\( W_{l}^{\text{in}} \) 從殘差串流狀態 \( x_{\text{mid},l,i} \) 讀取，其結果通過一個逐元素的非線性激活函數 \( g \)，產生神經元激活。這些激活再由 \( W_{l}^{\text{out}} \) 轉換以產生輸出 \( \text{FFN}(x_{\text{mid},i}) \)，然後加回到殘差串流：

\[ \text{FFN}_l(x_{\text{mid},l,i}) = g(x_{\text{mid},l,i} W_{l}^{\text{in}}) W_{l}^{\text{out}} \]

\[ x_{l,i} = x_{\text{mid},l,i} + \text{FFN}_l(x_{\text{mid},l,i}) \]

前饋網路的計算可以等同於鍵值記憶體檢索，\( W_{l}^{\text{in}} \) 中的列向量充當輸入序列上的模式檢測器（鍵），每個神經元激活加權了 \( W_{l}^{\text{out}} \) 的列向量（值）。元素式非線性在前饋網路內部創建了一個特權基底，鼓勵特徵與基底方向對齊。

預測層包含一個去嵌入矩陣 \( W_U \in \mathbb{R}^{d \times |V|} \)，有時還有一個偏差。最後一個殘差串流狀態透過這個線性映射轉換，將表示轉換為下一個 token 的 logits 分布，再通過 softmax 函數轉換為概率分布。由於模型組件透過相加與殘差串流交互作用，未正規化的分數(logits)是透過組件輸出的線性投影獲得的。基於線性變換的性質，可以重新排列傳統的前向傳遞公式，使每個模型組件直接貢獻於預測的 token 的 logits:

\[ f(x) = x_L^n W_U = \left( \sum_{l=1}^L \sum_{h=1}^H \text{Attn}_{l,h}(X_{l-1}^{\leq n}) + \sum_{l=1}^L \text{FFN}_l(x_{\text{mid},l}^n) + x^n \right) W_U \]

\[ = \sum_{l=1}^L \sum_{h=1}^H \text{Attn}_{l,h}(X_{l-1}^{\leq n}) W_U + \sum_{l=1}^L \text{FFN}_l(x_{\text{mid},l}^n) W_U + x^n W_U \]

注意力頭 logits 更新 前饋網路 logits 更新

這個分解在定位負責預測的組件時扮演重要角色，因為它允許測量每個組件對預測 token 的 logits 的直接貢獻。

殘差網路就像淺層網路的集成，每個子網路定義了計算圖中的一條路徑。將前向傳遞分解為：

\[ f(x) = xW_U + xW_1^{OV} W_U + xW_1^{OV} W_2^{OV} W_U + xW_2^{OV} W_U \]

直接路徑 完整 OV 電路

虛擬注意力頭 (V-composition)

連結輸入嵌入與去嵌入矩陣的路徑稱為直接路徑。穿過單個 OV 矩陣的路徑稱為完整 OV 電路。涉及兩個注意力頭的路徑稱為虛擬注意力頭，執行 V-composition，因為兩個頭的順序寫入和讀取被視為 OV 矩陣的組合。

理解語言模型的內部運作機制意味著定位前向傳遞中負責特定預測的元素（輸入元素、表示和模型組件）。本文介紹了兩種不同類型的方法，允許定位模型行為：輸入歸因和模型組件歸因。

輸入歸因方法通常用於通過估計輸入元素（在語言模型的情況下是 tokens）對定義模型預測的貢獻來定位模型行為。對於像語言模型這樣的神經網路模型，梯度資訊經常被用作歸因目的的自然度量。基於梯度的歸因在這個上下文中涉及 Transformer 在點 \( x \) 處的一階泰勒展開，表示為 \( \nabla f(x) \cdot x + b \)。結果梯度 \( \nabla f_w(x) \in \mathbb{R}^{n \times d} \) 直觀地捕捉了模型對輸入中的每個元素在預測 token \( w \) 時的敏感度。雖然歸因分數是針對輸入 token 嵌入的每個維度計算的，但它們通常在 token 層面匯總，以獲得對個別 token 影響的更直觀概述。這通常是透過取梯度向量相對於第 \( i \) 個輸入嵌入的 \( L_p \) 範數來完成的：

\[ A_{\text{Grad},f_w}^{(x \rightarrow t_i)} = \| \nabla_{x_i} f_w(x) \|_p \]

透過梯度向量與輸入嵌入 \( \nabla_{x_i} f_w(x) \cdot x_i \) 的點積，稱為梯度 × 輸入方法，可以將這種敏感度轉化為重要性估計。然而，這些方法已知會出現梯度飽和和碎裂問題。這一事實促使引入了積分梯度和 SmoothGrad 等方法來過濾嘈雜的梯度資訊。例如，積分梯度近似基線輸入 \( \tilde{x} \) 和輸入 \( x \) 之間直線路徑上的梯度積分：

\[ \int_0^1 \nabla_{x_i} f_w(\tilde{x} + \alpha(x - \tilde{x})) d\alpha \]

隨後提出了適應文本輸入離散性的改編。最後，基於層次相關傳播 (Layer-wise Relevance Propagation, LRP) 的方法已廣泛應用於研究基於 Transformer 的語言模型。這些方法對梯度傳播使用自定義規則，以分解每一層的組件貢獻，確保它們的總和在整個網路中保持不變。

另一個流行的方法族是透過添加噪聲或消融輸入元素並測量對模型預測的結果影響來估計輸入重要性。例如，可以移除位置 \( i \) 處的輸入 token，結果機率差異 \( f_w(x) - f_w(x_{-x_i}) \) 可以用作其重要性的估計。如果給予 \( w \) 的 logit 或機率沒有改變，我們可以得出第 \( i \) 個 token 沒有影響的結論。

雖然原始模型內部資料（如注意力權重）通常被認為提供了不忠實的模型行為解釋，但最近的方法提出了注意力權重的替代方案來測量中間 token 級別的歸因。其中一些替代方案包括使用值加權向量和輸出值加權向量的範數，或使用向量距離來估計貢獻。这些方法的共同策略涉及使用注意力展開等技術聚合反映上下文混合模式的中間每層歸因，得到輸入歸因分數。

一個重要的限制是，歸因的輸出 token 屬於一個大的詞彙空間，在下一個詞預測中經常有語義上等價的 token 競爭機率質量。在這種情況下，歸因分數很可能誤代了驅動模型預測的幾個重疊因素，如語法正確性和語義適當性。最近的工作通過提出這些方法的對比公式來解決這個問題，為模型預測 token \( w \) 而不是替代 token \( o \) 提供反事實解釋。

另一個輸入歸因的維度涉及識別影響推理時特定模型預測的有影響力的訓練樣本。這些方法通常被稱為訓練數據歸因 (TDA) 或實例歸因方法，並被用來識別數據中的人工因素和語言模型預測中偏差的來源。最近的方法提出通過訓練運行模擬來執行 TDA。雖然已建立的 TDA 方法的適用性受到質疑，特別是由於其低效率，但這個領域最近的工作產生了更有效的方法，可以大規模應用於大型生成模型。

早期關於 Transformer 語言模型組件重要性的研究強調了模型能力的高度稀疏性。例如，即使刪除模型中相當一部分的注意力頭，也可能不會使其下游性能惡化。這些結果激發了一條新的研究路線，研究語言模型中的各種組件如何貢獻於其廣泛的能力。

讓我們稱組件 \( c \)（注意力頭或前饋網路）在特定層上對最後一個 token 位置的輸出表示為 \( f_c(x) \)。等式 (10) 中提出的分解允許我們測量每個模型組件對輸出 token \( w \in V \) 的直接 logit 歸因 (DLA):

\[ A_{DLA,f_w}^{(x \rightarrow c)} = f_c(x) W_U[:,w] \]

其中 \( W_U[:,w] \) 是 \( W_U \) 的第 \( w \) 列，即 token \( w \) 的去嵌入向量。實際上，組件 \( c \) 的 DLA 表示 \( c \) 對預測 token 的 logit 的貢獻，使用第 2.2 節中描述的模型組件的線性特性。

我們可以將模型的計算視為一個因果模型，並使用因果工具來闡明每個模型組件 \( c \in C \) 在不同位置對預測的貢獻。因果模型可以看作是一個有向無環圖 (DAG)，其中節點是模型計算，邊是激活。

我們可以通過改變前向傳遞中由模型組件計算的某些節點值 \( f_c(x) \) 到另一個值 \( \tilde{h} \) 來干預模型，這被稱為激活修補。我們可以使用 do 算子表示這個干預：\( f(x \mid \text{do}(f_c(x) = \tilde{h})) \)。然後我們測量修補後預測的變化：

\[ A_{\text{Patch}, f}^{(x \rightarrow c)} = \text{diff}(f(x), f(x \mid \text{do}(f_c(x) = \tilde{h}))) \]

其中 \(\text{diff}(\cdot, \cdot)\) 函數的常見選擇包括 KL 散度和 logit/機率差異。修補後的激活 (\(\tilde{h}\)) 可以來自各種來源。一種常見的方法是創建一個具有分布 \(P_{\text{patch}}\) 的反事實數據集，其中一些關於任務的輸入信號被反轉。這種方法會導致兩種不同類型的消融：

- 重採樣干預，其中修補後的激活是從 \(P_{\text{patch}}\) 的單個樣本中獲得的，即 \(\tilde{h} = f_c(\tilde{x}), \tilde{x} \sim P_{\text{patch}}\)。

- 平均干預，其中多個 \(P_{\text{patch}}\) 樣本的平均激活用於修補，即 \(\tilde{h} = \mathbb{E}_{\tilde{x} \sim P_{\text{patch}}}[f_c(\tilde{x})]\)。

另外，修補激活的其他來源包括：

- 零干預，其中激活被替換為空向量，即 \(\tilde{h} = 0\)。

- 噪聲干預，其中新的激活是通過在受擾動的輸入上運行模型獲得的，例如 \(\tilde{h} = f_c(x + \epsilon), \epsilon \sim N(0, \sigma^2)\)。

在設計因果干預實驗時需要考慮的一個重要因素是設置的生態效度，因為零消融和噪聲消融可能使模型偏離自然激活分布，最終破壞組件分析的有效性。

機制可解釋性 (Mechanistic Interpretability) 子領域專注於將神經網路逆向工程為人類可理解的算法。MI 的最新研究旨在揭示電路的存在，電路是一組共同交互作用以解決任務的模型組件（子圖）。激活修補、logit 歸因和注意力模式分析是電路發現的常用技術。

邊緣修補和路徑修補利用每個模型組件輸入是其殘差串流中先前組件輸出之和這一事實，並考慮直接連接成對模型組件節點的邊。路徑修補將邊緣修補方法推廣到多條邊，允許更細粒度的分析。例如，使用等式中描述的淺層網路的前向傳遞分解，我們可以將圖中的單層 Transformer 視為由以下組件組成：

\[ f(x) = \text{Attn}(X^{\leq n}) W_u + \text{FFN}(\text{Attn}(X^{\leq n}) + x_n)W_u + x_n W_u \]

這裡的各個路徑包括：

- 直接從Attn到logits。

- 通過FFN再到logits的間接路徑。

其中每個發送節點Attn\( _L(X_{L-1}^{\leq n}) \)的副本都與單個路徑相關。在這個例子中，分別修補每個發送節點副本允許我們估計Attn\( _L(X_{L-1}^{\leq n}) \)對輸出logits\( f(x) \)的直接和間接效應。一般來說，我們可以將路徑修補應用於網路中的任何路徑，並測量頭部之間的組合、FFN或這些組件對logits的影響。

基於因果干預的電路分析存在幾個缺點：

1. 它需要為要評估的任務設計輸入模板，以及反事實數據集（即定義\( P_{\text{patch}} \)）方面的大量努力。

2. 在獲得組件重要性估計後，需要人工檢查和領域知識來分離重要的子圖。

3. 已經表明，干預會在下游組件的行為中產生二階效應，在某些情況下甚至會引發類似於自我修復的補償行為。這種現象會使得難以得出關於每個組件作用的結論。

為了克服這些限制，Conmy等人提出了一種自動電路發現（ACDC）算法，通過迭代移除計算圖中的邊來自動識別電路的過程。然而，這個過程需要大量的前向傳遞（每個修補元素一次），在研究大型模型時變得不切實際。

修補的一個有效替代方法是基於梯度的方法，這些方法已經擴展到超越輸入歸因，以計算中間模型組件的重要性。例如，給定token預測\( w \)，為了計算中間層\( l \)的歸因，表示為\( f_l(x) \)，計算梯度\( \nabla f_w(f_l(x)) \)。Sarti等人將等式中的對比梯度歸因公式擴展到使用單次前向和後向傳遞來定位對正確延續的預測貢獻大於錯誤延續的組件。Nanda等人提出了邊緣歸因修補（EAP），包括修補前後預測差異的線性近似，以估計計算圖中每個邊的重要性。這種方法的主要優點是它只需要兩次前向傳遞和一次後向傳遞就可以獲得圖中每個邊的歸因分數。

另一個研究方向是在較低層次的神經網路中尋找可解釋的高層次因果抽象概念。這些方法涉及大量的計算搜索，並假設高層次變數與單元或神經元組對齊。為了克服這些限制，Geiger等人提出了分佈式對齊搜索（Distributed Alignment Search, DAS），它在通過梯度下降找到的低層次表示空間的非基底對齊子空間上執行分佈式交換干預（Distributed Intervention Interchange, DII）。DAS干預在使用語法評估尋找具有因果影響的特徵方面被證明是有效的，同時在分離實體的個別屬性的因果效應方面也有顯著效果。

最後，在Transformer語言模型的內部行為方面，本文總結了以下主要發現：

注意力區塊：

• 位置頭：有些頭主要關注相對於正在處理的token的特定位置，如token本身、前一個token或下一個位置。
• 子詞連接頭： 專門關注屬於與當前處理的token相同詞的前一個子詞token。
• 語法頭： 一些頭關注具有與被處理token相關語法角色的token，明顯多於隨機基準。
• 重複token：** 關注上下文中同一個token的先前出現。
• 複製頭： OV矩陣表現出複製行為。
• 歸納頭： 由兩個在不同層組合的頭組成，讓語言模型完成模式。一個早期的前一個token頭將第一個token A的資訊複製到B的殘差流，然後一個下游的歸納頭關注token B，增加B的可能性。
• 複製抑制頭： 如果它們出現在上下文中並且當前殘差流正在自信地預測它，則減少它們關注的token的logit分數。

前饋網路區塊：

• 神經元的輸入行為：有些神經元僅在特定位置範圍上激發；技能神經元，其激活與輸入提示的任務相關；概念特定神經元，其反應可用於預測上下文中概念的存在。
• 神經元的輸出行為：有些神經元促進與特定語義和句法概念相關的token的預測；一小部分後層神經元負責做出在語言上可接受的預測；抑制不可能延續的神經元。
• 多語義神經元： 早期層的大多數神經元專門用於n-gram集合，充當n-gram偵測器，大多數神經元在大量n-gram上激發。

殘差流：

• 殘差流可以被視為Transformer中的主要通訊通道。直接路徑主要對應於bi-gram統計，而網路中的最新偏差根據詞頻轉移預測，促進高頻token。
• 一些組件執行記憶體管理，以移除存儲在殘差流中的資訊。例如，有負特徵值的OV矩陣關注當前位置的注意力頭，以及輸入和輸出權重具有較大負餘弦相似性的前饋網路神經元。
• 在殘差流中發現了離群值維度。這些維度展現出相對於其他維度的大幅度，與各向異性表示的生成相關聯。消融離群值維度已被證明會顯著降低下游性能，表明它們編碼任務特定的知識。

多組件行為：

• 歸納機制： 是兩個組件（注意力頭）組合在一起以完成模式的一個明顯例子。最近的證據表明，多個注意力頭協同工作，在給定上下文樣本時創建描述任務的「函數」或「任務」向量。
• Variengien 和 Winsor 研究了涉及回答可以在上下文中找到答案的請求的上下文檢索任務
• 作者確定了一種在子任務和模型之間通用的高層次機制。具體而言，中間層處理請求，然後由後層的注意力頭執行從上下文中檢索實體的步驟。
• 在 GPT-2 Small 中發現了用於間接賓語識別 (IOI) 任務的電路，主要包括：
• 重複信號：重複token頭和涉及前一個token頭的歸納機制表明S（John）的重複性。這個資訊被最後位置的S-抑制頭讀取，它們在殘差流中寫入一個token信號，表明S被重複，以及S1 token的位置信號。
• 名稱複製： 後層的名稱移動頭將它們在上下文中關注的名稱的資訊複製到最後的殘差流。然而，先前層S-抑制頭的信號修改了名稱移動頭的query，使得重複的名稱（在S1和S2中）受到較少關注，有利於複製間接賓語（IO），從而推動其預測。

這篇論文全面而深入地介紹了目前用於解釋 Transformer 語言模型內部運作的技術,並總結了通過這些方法得到的關於模型內部機制的見解。作者強調,雖然在可解釋性研究方面取得了顯著進展,但將這些見解應用於調試和改進未來模型的安全性和可靠性,為開發人員和用戶提供更好的工具來與之交互並理解影響其預測的因素,仍然是一個巨大的挑戰。未來可解釋性研究的發展將面臨從在模型組件空間運作的方法和分析轉向人類可解釋空間(即從模型組件到特徵和自然語言解釋)的挑戰性任務,同時仍然忠實地反映模型行為。

此外,跨學科研究將在擴大可解釋性分析的範圍方面發揮關鍵作用,以考慮從人的角度看模型解釋的感知和交互維度。最終,作者認為,確保對先進語言模型的內部機制的開放和便利訪問,將仍然是這一領域未來進展的基本先決條件。

這篇綜述論文對於理解 Transformer 語言模型的內部運作機制,以及目前在這個領域的最新研究進展,提供了全面而詳盡的概覽。透過系統性地介紹各種可解釋性技術,並深入探討它們揭示的模型內部行為,本文為相關研究人員提供了寶貴的參考。

不過,正如作者所指出,將這些見解應用到實際中仍面臨諸多挑戰。未來的可解釋性研究需要在忠實反映模型行為的同時,努力向更加貼近人類直觀理解的方向發展。這需要不同學科領域的通力合作。此外,開放模型內部機制的訪問,或許是這一領域能夠取得突破性進展的關鍵。

證據順序(order of evidence)

 我今天在ICLR學到一個很有趣的術語,叫做證據順序(order of evidence)。

在檢索增強生成(Retrieval-Augmented Generation, RAG)的過程中,技術團隊會將檢索到的前幾名文件作為證據,並提示(prompt)給大型語言模型(Large Language Models, LLMs)。通常提示的順序是基於餘弦相似度(cosine similarity)的:匹配度最高的文件總是最先被提示。

然而,不同的大型語言模型可能對證據的順序有不同的「偏好」:

1. ChatGPT偏好排名最前面的證據,這很好。

2. GPT-4對證據的順序沒有偏好,這意味著相似度分數並沒有被考慮,你只需要決定前K個證據,所有證據都會被平等對待。

3. 令人驚訝的是,Llama2和PaLM偏好最後一個證據,所以你需要先反轉排名列表,然後再提示給大型語言模型:)

該研究提出一個系統性的框架,用於引出大型語言模型的參數記憶(parametric memory),並構建相應的反記憶(counter-memory)。技術團隊設計了一系列檢查,如從參數記憶到答案的蘊涵(entailment),以確保引出的參數記憶確實是大型語言模型的內部信念。對於反記憶,該研究沒有採用啟發式地編輯參數記憶,而是直接指示大型語言模型生成一個與參數記憶在事實上相矛盾的連貫段落。

在獲得大量參數記憶和反記憶對之後,技術團隊接著檢查大型語言模型在不同的知識衝突情況下的行為,包括1)當只有反記憶作為外部證據時,以及2)當參數記憶和反記憶都存在時。該研究的調查揭示了大型語言模型看似矛盾的行為。一方面,與先前的觀點不同,該研究發現,如果外部證據是連貫且令人信服的,大型語言模型可以高度接受外部證據,即使那與它們的參數記憶相矛盾。另一方面,當存在一些與其參數記憶一致的信息時,儘管同時被提供相矛盾的證據,大型語言模型也表現出強烈的確認偏差(confirmation bias)。這些結果對進一步開發和部署工具增強型和檢索增強型大型語言模型提出了值得仔細考慮的重要啟示。可在以下網址找到相關資源:

https://github.com/OSU-NLP-Group/LLM-Knowledge-Conflict

$$MR=\frac{f_m}{f_m+f_c}$$, 

(1)

其中\(f_m\)是參數記憶答案的頻率,\(f_c\)是反記憶答案的頻率。記憶率(memorization ratio)越高,表示大型語言模型越依賴其參數記憶,而記憶率越低則表示更頻繁地採用反記憶。

該研究進行了一系列實驗來探究大型語言模型在知識衝突下的行為。在單一證據的情況下,當反記憶是唯一呈現的證據時,技術團隊發現大型語言模型實際上對外部證據的接受度很高,即使它與其參數記憶相矛盾。這與先前的結論相矛盾,該研究認為這是因為通過其框架構建的反記憶更加連貫和令人信服。然而,這也表明大型語言模型可能很容易被誤導,例如來自惡意第三方工具的錯誤信息。

在多個證據的情況下,當同時存在支持其參數記憶和與之矛盾的證據時,大型語言模型表現出強烈的確認偏差,傾向於堅持其參數記憶。這揭示了大型語言模型在公正地協調多個相互矛盾的證據方面可能面臨挑戰,這是生成式搜索引擎常見的情況。此外,該研究表明,大型語言模型對證據的偏好受到證據的流行度(popularity)、順序(order)和數量(quantity)的影響,這些都可能不是工具增強型大型語言模型所希望的特性。

最後,該研究框架的有效性也表明,大型語言模型可以生成令人信服的錯誤信息,這帶來了潛在的道德風險。技術團隊希望他們的工作能為理解、改進和部署工具增強型大型語言模型提供一個可靠的評估基準和有用的見解。

該研究強調了一個嚴重的問題:大型語言模型可以被指示編造連貫且令人信服的虛假信息。這凸顯了如果不加以控制,這些模型可能被濫用的風險。作為研究人員,解決這一迫切問題是他們的責任。大型語言模型被濫用的風險需要強有力的保障措施和預防措施,這需要更廣泛的研究界共同努力。為此,技術團隊承諾謹慎分發通過其研究生成的數據,確保其僅用於研究目的。他們的目標是在最大限度地發揮大型語言模型所提供的益處的同時,降低風險。

該研究的實驗利用了三個通過API訪問的封閉源大型語言模型,以及五個開源的大型語言模型。為了提高可重複性,技術團隊在附錄C中包含了實驗中使用的提示(prompt)。關於封閉源大型語言模型的版本,他們在所有測試中使用了ChatGPT-0301、GPT-4-0314和PaLM2的Chat-Bison-001。

該研究還討論了其發現的更廣泛影響以及潛在的解決方案。高度接受性是一把雙刃劍。一方面,它意味著可以有效地補救大型語言模型的過時或不正確的參數知識,這對檢索增強生成等方法是有益的。另一方面,隨著大型語言模型越來越多地與外部工具連接,例如ChatGPT插件和最近的語言代理如AutoGPT,高度接受外部輸入引起了擔憂——大型語言模型可能很容易被來自惡意第三方工具的誤導或操縱信息所欺騙。

確認偏差是一個非常不受歡迎的特性,特別是對於生成式搜索引擎或大型語言模型的類似應用(例如多文檔摘要),在這些應用中,以公正的方式協調多個可能相互矛盾的信息非常重要。

在潛在的解決方案方面,對於高度接受性帶來的風險,應該採用驗證和監控系統,以防止第三方工具提供不當信息給大型語言模型。對於確認偏差,根據部署情景,通過微調(fine-tuning)或人類反饋強化學習(RLHF)進行進一步的調整以減少偏差可能是一個有前景的方向。最後,從生成式搜索引擎的角度來看,引用答案的來源並讓用戶更加知情,並判斷最終答案,可能是一種更可靠的方式。

總的來說,該研究對工具增強型大型語言模型在知識衝突下的行為進行了全面而受控的調查,揭示了一些看似矛盾但值得關注的特性。這些發現為進一步開發和部署此類系統提供了重要啟示,同時也引發了一些值得深思的道德問題。技術團隊希望他們的工作能夠為相關研究提供一個紮實的基礎,推動工具增強型大型語言模型朝著更加可靠、穩健的方向發展。

在該研究的附錄中,技術團隊提供了更多關於實驗設置的詳細信息。他們在每個步驟的數據集規模如表B.3所示。此外,他們還在表B.5中報告了不一致類型的分佈情況。表B.6展示了一些大型語言模型在答案不一致方面的示例。在表B.7中,技術團隊展示了最終數據集中的更多樣本。

在表B.8中,該研究展示了一些即使在僅給出反記憶證據的情況下,大型語言模型仍然固執地給出參數記憶答案的例子。通過人工仔細審查50個隨機選擇的樣本,技術團隊發現其中34個例子是由於反記憶中的歧義、導致無法接受反記憶的常識問題或高度暗示性的問題。這意味著只有一小部分大型語言模型在單一來源設置下對參數記憶表現出固執,再次證實了大型語言模型在此設置下保持開放。

技術團隊還探討了更複雜的知識衝突情景。他們感興趣的問題是:如果向大型語言模型提供不相關的證據會怎樣?當提供不相關的證據時,大型語言模型應該 1)如果沒有證據明確支持任何答案,則應避免回答;2)忽略不相關的證據,並根據相關證據回答問題。為了設置實驗,技術團隊將句子BERT嵌入檢索到的與問題相似度最高的不相關段落視為不相關證據(即與問題中提到的實體無關的句子)。表B.7顯示了在POPQA上的實驗結果。

總之,該研究通過系統地構建高質量的參數記憶和反記憶,並設計一系列受控實驗,深入調查了工具增強型大型語言模型在知識衝突下的行為。這些發現揭示了當前大型語言模型在處理外部證據時存在的一些問題,包括易受誤導、確認偏差以及對證據順序和數量的敏感性等。這些結果為進一步改進和部署工具增強型大型語言模型提供了重要啟示,同時也凸顯了潛在的道德風險。技術團隊希望他們的工作能夠為相關研究提供一個紮實的基礎,推動工具增強型大型語言模型朝著更加可靠、穩健的方向發展。他們也呼籲研究界和業界共同努力,制定相應的應對措施,最大限度發揮大型語言模型的潛力,同時降低其被濫用的風險。

提出一些個人的見解和想法。

首先,該研究揭示了工具增強型大型語言模型在處理外部證據時存在的一些問題,這對於開發和部署此類系統具有重要啟示意義。技術團隊提出的框架和實驗設置為相關研究提供了一個紮實的基礎,有助於研究者更好地理解和改進這些模型。未來的研究可以在此基礎上,進一步探索如何減少大型語言模型的確認偏差,提高其處理不一致證據的能力,以及如何設計更加穩健的知識融合機制等。

其次,該研究也凸顯了大型語言模型潛在的道德風險。大型語言模型強大的生成能力如果被濫用,可能會產生連貫且令人信服的錯誤信息,對社會造成負面影響。因此,在開發和部署這些模型時,研究界和業界都需要高度重視這一問題,並採取相應的應對措施。這可能包括制定嚴格的數據使用和分發規範、設計有效的驗證和監控機制、加強模型的可解釋性和可控性等。

最後,工具增強型大型語言模型作為一個新興的研究方向,其潛力和局限性都有待進一步探索。該研究為這一領域提供了寶貴的見解和啟發,同時也揭示了一些亟待解決的問題。未來的研究可以在這些發現的基礎上,進一步拓展工具增強型大型語言模型的應用場景,如在問答、對話、知識圖譜構建等任務中的應用,並不斷優化其性能和可靠性。同時,也需要密切關注這一技術的最新進展,評估其對社會的潛在影響,並做好相應的準備和應對。

該研究對於理解和改進工具增強型大型語言模型具有重要意義。它為相關研究提供了新的視角和方法,同時也提出了一些值得深入探討的問題。技術團隊的工作有助於推動這一領域的發展,讓我們在充分發揮大型語言模型潛力的同時,也能更好地應對其帶來的挑戰與風險。