來源：DeepTech深科技

眾所周知，目前的大模型大多基于 Transformer 架構(gòu)。Transformer 的核心結(jié)構(gòu)是多頭自注意力模型（multi-head self-attention model）。

大模型的一個重要能力是所謂的“上下文學(xué)習(xí)”。具體來說，當(dāng)大模型的參數(shù)訓(xùn)練好之后，用戶和大模型的交互方式，是通過提供上文來獲得大模型的下文，這時大模型的參數(shù)是固定的。

當(dāng)所提供的上文包含一些關(guān)于同一主題的輸入輸出例子時，大模型可以根據(jù)給的這些例子，學(xué)到這些例子背后的主題，從而可以在給到一個新輸入時，回答正確的輸出。

比如，上文可以是：

（來源：資料圖）

這時，Claude 3 大模型的回答是：

（來源：資料圖）

由圖可知，Claude 根據(jù)這些例子意識到“+”其實(shí)代表著減法，故能針對“10+5=？”這一新問題給出正確回答。

上下文學(xué)習(xí)，是大模型的一個基礎(chǔ)能力。使用大模型時的其他更復(fù)雜方式比如 Chain-of-thought reasoning，都是以此為基礎(chǔ)。

但是，從原理來看上下文學(xué)習(xí)的機(jī)制并不是很清楚。很大原因在于大模型作為一個系統(tǒng)，它不僅非常復(fù)雜，而且模型參數(shù)非常多，訓(xùn)練數(shù)據(jù)也非常大。

為了更好地理解上下文學(xué)習(xí)，美國斯坦福大學(xué)團(tuán)隊(duì)曾在 GPT-2 架構(gòu)之下，針對大模型如何使用上下文，學(xué)習(xí)解決簡單的回歸預(yù)測進(jìn)行了研究。

其發(fā)現(xiàn)當(dāng)使用簡單函數(shù)的數(shù)據(jù)來訓(xùn)練大模型時，訓(xùn)練好的大模型，可以通過上下文學(xué)到這些簡單函數(shù)。

一個特別的例子便是線性函數(shù)。這時的訓(xùn)練數(shù)據(jù)是一些線性數(shù)據(jù) x_1，w x_1，…，x_n，w x_n，其中 w 是高斯隨機(jī)向量。

換句話說，每個“句子”里都有 n 個線性函數(shù)的例子，而這個線性函數(shù)是隨機(jī)的。

（來源：https://arxiv.org/pdf/2208.01066.pdf）

以此為啟發(fā)，美國耶魯大學(xué)助理教授楊卓然和團(tuán)隊(duì)，希望可以從理論上研究這種訓(xùn)練過程是否收斂、以及收斂到哪里，也希望厘清多頭自注意力結(jié)構(gòu)到底是如何實(shí)現(xiàn)上下文學(xué)習(xí)的。

隨后，他和所在團(tuán)隊(duì)考慮了一個最簡單的模型：一層多頭自注意力模型。

（來源：arXiv）

具體來說在本次課題之中，他們研究了訓(xùn)練多頭自注意力模型（multi-head self-attention model）的優(yōu)化問題。

尤其是，他們回答了這樣一個問題：在使用一層多頭自注意力模型（one-layer multi-head self-attention model）進(jìn)行上下文學(xué)習(xí)時：

首先，梯度優(yōu)化算法是否能夠收斂？

其次，梯度優(yōu)化算法收斂到的解統(tǒng)計(jì)性質(zhì)如何？

再次，從網(wǎng)絡(luò)結(jié)構(gòu)的角度看，多頭自注意力模型是如何進(jìn)行上下文學(xué)習(xí)的？

期間，他們所使用的訓(xùn)練數(shù)據(jù)是多任務(wù)線性模型（multi-task linear regression）。

特別地，每個線性模型的參數(shù) G 在一個固定的正交基下，有一個分塊對角的分解。

也就是說如果能找到這組基，這個線性模型就可以分解成 H 個獨(dú)立的線性模型。

對于每個參數(shù) G，能夠生成 L 個（x，y）對，并且可以讓 Transformer 推測一個隨機(jī)的 q 所對應(yīng)的 y 是什么。

（來源：arXiv）

在這種多任務(wù)線性數(shù)據(jù)上，課題組使用梯度流來訓(xùn)練 Transformer，進(jìn)而研究這一算法的收斂問題。

通過此，他們發(fā)現(xiàn)：梯度流算法的確是收斂的。并且收斂有三個階段——（a）預(yù)熱階段、（b）任務(wù)分配階段 、以及（c）最終收斂階段。

在（a）這一預(yù)熱階段，損失函數(shù)緩慢下降。

在（b）這一“任務(wù)分配階段”，損失函數(shù)迅速下降。并且，softmax 函數(shù)使得每一個自注意力頭只關(guān)注多任務(wù)線性模型的一個任務(wù)，該團(tuán)隊(duì)把這一現(xiàn)象稱為“任務(wù)分配”。

在（c）最終收斂階段，每個自注意力頭繼續(xù)對它被分配的任務(wù)求解，最終達(dá)到收斂。

（來源：arXiv）

此外，他們還描述了梯度流學(xué)習(xí)的極限模型的上下文學(xué)習(xí)預(yù)測誤差。

當(dāng) （d/L） 趨于零時，誤差衰減到零，其中 d 是線性模型的維數(shù)，L 是上下文學(xué)習(xí)中（x，y）例子的數(shù)量。

并且，該團(tuán)隊(duì)還證明多頭自注意力模型，顯著好于單頭自注意力模型。所預(yù)測的誤差相差 H 倍之多，其中 H 是注意力頭的個數(shù)。

換句話說，注意力頭的個數(shù)越多，性能差距越大。

據(jù)介紹，該團(tuán)隊(duì)的分析主要基于對自注意力權(quán)重的分解。

自注意力機(jī)制里主要有兩類權(quán)重：QK 權(quán)重（query-key）和 OV 權(quán)重（output-value）。其中，QK 權(quán)重反應(yīng)著 query 和 key 的關(guān)系。

簡單來說，就是給定了 query q（新的輸入）和過去的例子（x，y）的關(guān)系。

而 QK 權(quán)重反應(yīng)著 attention 對每一個過去的例子的重視程度。

OV 權(quán)重反應(yīng)著輸出和每一個輸入例子（x，y）的關(guān)系，即 attention 如何通過組合上下文學(xué)習(xí)中的例子從而得到輸出。

需要注意的是在回歸問題里面，q 是一個輸入，和 x 有一樣的維度，輸出和 y 有一樣的維度。

課題組發(fā)現(xiàn)，QK 權(quán)重和 OV 權(quán)重都是分塊的，并且 QK 權(quán)重的 X-X 分塊和 OV 權(quán)重的 Y 分塊最為重要。

也就是說，在回歸問題里只需使用 q 和例子里的 x 比較得到注意力值（attention score）。

在輸出時，只需要根據(jù)注意力值（attention score）來合并例子里的那些 y。

而通過利用數(shù)據(jù)的線性結(jié)構(gòu)，他們發(fā)現(xiàn) QK 權(quán)重和 OV 權(quán)重的分塊結(jié)構(gòu)，可以被梯度流算法保持。

更特別的是，因?yàn)槎嗳蝿?wù)線性模型的參數(shù) G 可以在某個基下分解，憑借此他們證明 QK 權(quán)重和 OV 權(quán)重也是可以被分解的。

這樣一來，就可以把參數(shù)的梯度流化，簡為奇異值的梯度流，這時就只需要分析奇異值的變化。

其中，總共有 H*（d_x + d_y）個奇異值，每個頭的 QK 權(quán)重有 d_x 個奇異值，OV 權(quán)重有 d_y 個奇異值。

而 d_x 是 x 的維度，d_y 是 y 的維度，也就是多任務(wù)線性模型的任務(wù)數(shù)量。

并且，每個自注意力頭的 OV 權(quán)重的奇異值，反應(yīng)著自注意力頭對于對應(yīng)任務(wù)的重視程度。

隨后，該團(tuán)隊(duì)開始分析這些奇異值的變化。他們發(fā)現(xiàn)自注意力頭的任務(wù)分配基于“OV 權(quán)重–每個任務(wù)”的原則，來將最大的自注意力頭分給對應(yīng)的奇異值。

比如，第一個任務(wù)被分配給了第一個奇異值最大的自注意力頭。

最終，在梯度流達(dá)到收斂之后，根據(jù)任務(wù)分配機(jī)制，每個自注意力頭的 OV 權(quán)重只有唯一一個非零的奇異值。

（來源：arXiv）

陳思宇和王天浩分別是第一作者和第三作者，楊卓然擔(dān)任通訊作者。其中，王天浩將于 2025 年秋入職美國加州大學(xué)圣地亞哥分校。

圖 | 楊卓然（來源：楊卓然）

不過，課題組仍然覺得自己對于 transformer 和上下文學(xué)習(xí)的理解還非常粗淺。

目前，他們只研究了一層自注意力模型。后續(xù)，他們希望能夠研究多層的自注意力模型。

與此同時，目前他們只研究了線性模型。因此，他們也非常希望研究非線性的上下文學(xué)習(xí)問題。

此外，目前課題組給到 transformer 的輸入，是獨(dú)立同分布的（x，y）輸入輸出數(shù)據(jù)對，這里輸入并沒有任何復(fù)雜的前后依賴結(jié)構(gòu)。

但是，實(shí)際用來訓(xùn)練 transformer 的數(shù)據(jù)都是文本數(shù)據(jù)，里面有復(fù)雜的依賴結(jié)構(gòu)，針對此他們也將繼續(xù)加以探索。

參考資料：

1.https://arxiv.org/pdf/2403.00993

排版：羅以

相關(guān)知識

耶魯大學(xué)研究揭示，身體虛弱與抑郁風(fēng)險之間的聯(lián)系
心理學(xué)部教授團(tuán)隊(duì)在Appetite發(fā)文揭示超重肥胖人群食物相關(guān)抑制控制的神經(jīng)機(jī)制
肥胖人的福音？中山六院研究團(tuán)隊(duì)連續(xù)發(fā)文揭示肥胖機(jī)制
何清華教授團(tuán)隊(duì)發(fā)文揭示超重/肥胖者的食物決策沖動性及認(rèn)知神經(jīng)機(jī)制
心理學(xué)部團(tuán)隊(duì)在《Nature Communications》發(fā)文揭示睡眠健康的多維度特點(diǎn)及其和腦網(wǎng)絡(luò)的協(xié)同變化規(guī)律
齊魯理工學(xué)院：培養(yǎng)齊魯文化孕育下的高素質(zhì)人才
科研團(tuán)隊(duì)揭示遺傳因素在母嬰健康中的重要作用—新聞—科學(xué)網(wǎng)
減肥神藥會讓心臟變??？科學(xué)家建議關(guān)注人心臟結(jié)構(gòu)和功能的變化；40多篇論文齊發(fā)！“人類細(xì)胞圖譜”發(fā)布重大進(jìn)展
宮婷婷團(tuán)隊(duì)系統(tǒng)揭示間歇性禁食對健康的影響
科學(xué)網(wǎng)—健康記錄顯示，流行的減肥藥與自殺念頭之間沒有聯(lián)系

網(wǎng)址: ?耶魯團(tuán)隊(duì)揭示多頭自注意力結(jié)構(gòu)的上下文學(xué)習(xí)機(jī)制，證明梯度流算法的收斂性 http://www.gysdgmq.cn/newsview484921.html