Transformer稱霸的原因找到了？OpenAI前核心員工揭開注意力頭協(xié)同工作機(jī)理

這家由 OpenAI 前核心員工組成的 AI 創(chuàng)業(yè)公司，正在努力打開 Transformer 的黑箱。

在過去的兩年里，人工智能在自然語言處理領(lǐng)域的突破達(dá)到了驚人的水平。很多重要的突破都是基于谷歌在 2017 年提出的 Transformer 架構(gòu)。
但在模型之下，我們?nèi)匀徊恢?Transformer 對(duì)其處理的單詞做了什么。從普遍的理解來看，它們能夠以某種方式關(guān)注多個(gè)單詞，從而可以立即進(jìn)行「全局」分析。但這究竟是如何做到的，或者這種說法是否是理解 Transformer 的正確方式，都還不清楚。我們知道「成分」，但卻不了解「配方」。
現(xiàn)在，Anthropic 公司的研究人員進(jìn)行的兩項(xiàng)研究試圖從根本上探究 Transformer 在處理和生成文本時(shí)會(huì)做什么。Anthropic 公司成立于 2021 年 5 月，團(tuán)隊(duì)主要由一批前 OpenAI 核心員工組成，致力于提高 AI 的安全性及可解釋性。
第一篇論文發(fā)表于 2021 年的 12 月，作者著眼于架構(gòu)的簡(jiǎn)化版本，并充分解釋了它們的功能。以色列理工學(xué)院的 Yonatan Belinkov 表示：「這些簡(jiǎn)化架構(gòu)很好地描述了 Transformer 在非常簡(jiǎn)單的情況下是如何工作的。我對(duì)這項(xiàng)工作非常感興趣。它很有趣，同時(shí)很有前途，也比較獨(dú)特和新穎。」

論文鏈接：https://transformer-circuits.pub/2021/framework/index.html
作者還表明，簡(jiǎn)單的 Transformer 具備從學(xué)習(xí)基本語言模式到語言處理的一般能力?！改銜?huì)看到能力的飛躍」，哈佛大學(xué)的 Martin Wattenberg 說道。研究團(tuán)隊(duì)已經(jīng)開始破解 Transformer 的「配方」了。
在 2022 年發(fā)表的第二篇論文中，研究者表明，負(fù)責(zé)這種能力的相同組件在最復(fù)雜的 Transformer 中也能發(fā)揮作用。雖然這些模型背后的數(shù)學(xué)原理在很大程度上仍然難以理解，但該研究的結(jié)果提供了一種理解途徑?！杆麄?cè)诤?jiǎn)化版模型中發(fā)現(xiàn)的東西可以遷移到更大的模型中」，Conjecture 公司和研究小組 EleutherAI 的 Connor Leahy 說道。

論文鏈接：https://transformer-circuits.pub/2022/in-context-learning-and-induction-heads/index.html
理解 Transformer 的難點(diǎn)在于它們的抽象。傳統(tǒng)程序遵循可理解的過程，例如每當(dāng)看到「green」一詞時(shí)，就輸出「grass」一詞，而 Transformer 將「green」一詞轉(zhuǎn)換為數(shù)字，然后將它們與某些值相乘。這些值（也稱為參數(shù)）決定下一個(gè)單詞是什么。Transformer 在訓(xùn)練的過程中得到微調(diào)，模型能夠?qū)W會(huì)產(chǎn)生最佳輸出，但目前尚不清楚模型學(xué)習(xí)的是什么。
大多數(shù)機(jī)器學(xué)習(xí)程序?qū)?shù)學(xué)打包為模塊化的神經(jīng)元。Transformer 則加入了一種額外的成分，稱為注意力頭（attention head），這些頭分層排列（就像神經(jīng)元一樣）。但是頭的運(yùn)算方式與神經(jīng)元不同。頭的作用通常被理解為「允許程序記住輸入的多個(gè)單詞」，但這種解釋并不確切。
「注意力機(jī)制顯然是有效的。它取得了很好的效果，」Wattenberg 說。「但問題是它在做什么？我的猜測(cè)是它正在做很多我們不知道的事情?！?br />為了更好地理解 Transformer 的工作原理，Anthropic 的研究人員簡(jiǎn)化了架構(gòu)，去掉了所有的神經(jīng)元層和除了一層或兩層注意力頭之外的所有層。這讓他們發(fā)現(xiàn)了 Transformer 和簡(jiǎn)化模型之間的聯(lián)系。
理解 Transformer
先假設(shè)一種最簡(jiǎn)單的語言模型，稱為「二元模型」，它再現(xiàn)了基本的語言模式。例如在大型文本上訓(xùn)練時(shí)，二元模型會(huì)記下單詞「green」之后最常出現(xiàn)的單詞（例如「grass」）并記住它。然后，在生成文本時(shí)，它會(huì)重現(xiàn)相同的模式。通過為每個(gè)輸入單詞記住一個(gè)相關(guān)的后續(xù)單詞，它可以獲得非?；镜恼Z言知識(shí)。
研究人員表明，具有一層注意力頭的 Transformer 模型做了類似的事情：它再現(xiàn)了所記憶的內(nèi)容。假設(shè)你給它一個(gè)特定的輸入，比如「Doctor Smith went to the store because Doctor …」這個(gè)輸入被稱為 prompt 或上下文。對(duì)我們來說，下一個(gè)詞是顯而易見的 ——Smith。
在訓(xùn)練好的單層模型中，注意力頭可以分兩步進(jìn)行此預(yù)測(cè)。首先，它查看上下文中的最后一個(gè)詞（Doctor），并在它（訓(xùn)練期間）已經(jīng)學(xué)習(xí)的上下文中搜索一個(gè)特定的詞以與最后一個(gè)詞相關(guān)聯(lián)。然后，對(duì)于找出的單詞，查找另一個(gè)它已知相關(guān)聯(lián)的單詞，就像二元模型一樣。最后將這個(gè)關(guān)聯(lián)的詞移動(dòng)到模型的輸出中。
對(duì)于這個(gè)例子，研究人員表示，根據(jù)最后一個(gè)詞「Doctor」，注意力頭從訓(xùn)練中知道：要搜索一個(gè)通用名稱的詞。在句子前面找到名字「Smith」時(shí)，注意力頭會(huì)查看它學(xué)到的與「Smith」相關(guān)聯(lián)的內(nèi)容，并將該詞移動(dòng)到輸出。（在這種情況下，模型已經(jīng)學(xué)會(huì)將同一個(gè)詞「Smith」與找到的詞「Smith」相關(guān)聯(lián)。）整個(gè)過程的凈效果是模型將「Smith」這個(gè)詞從上下文復(fù)制到輸出。

研究作者：Chris Olah（左）、Catherine Olsson（中）和 Nelson Elhage（右）
「在這里，我們實(shí)際上可以理解注意力的作用，」研究的共同作者之一 Chris Olah 說。
但這種「記憶」只能幫模型到這里了。思考一下，當(dāng)史密斯的名字變成一個(gè)虛構(gòu)的名字時(shí)會(huì)發(fā)生什么，比如「Gigamuru」。對(duì)我們來說，句子如何完成仍然很簡(jiǎn)單 ——Gigamuru，但模型在訓(xùn)練期間不會(huì)看到虛構(gòu)的詞。因此，它不可能記住它與其他詞之間的任何關(guān)系，也不會(huì)生成它。
當(dāng)研究一個(gè)更復(fù)雜的模型 —— 一個(gè)有兩層注意力頭的模型時(shí)，Anthropic 團(tuán)隊(duì)找到了一種解決方案。該方案依賴于注意力頭獨(dú)有的能力：它們不僅可以將信息移動(dòng)到輸出上，還可以移動(dòng)到上下文中的其他地方。
借助這種能力，第一層的頭部學(xué)習(xí)利用前驅(qū)詞的信息來注釋上下文中的每個(gè)詞。然后，第二個(gè)頭可以搜索以單詞 “Doctor”（在本例中為 “Gigamuru”）之前的單詞，并像單層模型中的頭一樣，將其移動(dòng)到輸出上。研究人員將與前一層中頭協(xié)作的后一層中的注意力頭稱為感應(yīng)頭（induction head）。感性頭不僅有記憶的功能，而且可以做一些類似于抽象推理或?qū)崿F(xiàn)算法的事情。
感應(yīng)頭能讓兩層模型做到更多的事情，但它們與全尺寸 Transformer 的相關(guān)性尚不清楚，它們有數(shù)百個(gè)注意力頭協(xié)同工作。在第二篇論文中，研究人員發(fā)現(xiàn)這些結(jié)果得到了延續(xù)：感應(yīng)頭似乎對(duì)一些最復(fù)雜多層架構(gòu)的顯著行為做出了重大貢獻(xiàn)。
在這些行為中，算術(shù)的能力得到彰顯，因?yàn)槟Ｐ椭槐挥?xùn)練來完成文本。例如，如果給出重復(fù) prompt：「問：48 加 76 等于多少？答：124；問：48 加 76 等于多少？A：___」一個(gè)完整的模型會(huì)得到正確的答案，并且在獲得足夠多的非重復(fù) promp 后，它就能正確回答它從未見過的算術(shù)問題。這種從上下文中明顯學(xué)習(xí)新能力的現(xiàn)象稱為「上下文學(xué)習(xí)」。
這種現(xiàn)象令人費(fèi)解，因?yàn)閺纳舷挛闹袑W(xué)習(xí)是不可能的。這是因?yàn)闆Q定模型性能的參數(shù)僅在訓(xùn)練期間進(jìn)行調(diào)整，而沒有在模型處理輸入上下文時(shí)進(jìn)行調(diào)整。
感應(yīng)頭至少解決了部分難題。它們解釋了上下文學(xué)習(xí)的簡(jiǎn)單、重復(fù)形式的可能性，同時(shí)提供了所需要的 —— 復(fù)制模型未經(jīng)訓(xùn)練使用的新詞（如「Gigamuru」或「124」）的能力。
論文作者之一、Anthropic 團(tuán)隊(duì)成員 Catherine Olsson 表示：「感應(yīng)頭更可能適應(yīng)各種模式，即使它有些許奇怪或新奇?！?br />該研究進(jìn)一步在多層模型中應(yīng)用感應(yīng)頭，表明感性頭能夠參與更新穎的上下文學(xué)習(xí)形式，例如學(xué)會(huì)翻譯語言。
該研究為理解 transformer 提供了一個(gè)新的視角。它們不僅在獲取知識(shí)，而且還在學(xué)習(xí)處理根本沒有學(xué)過的東西。這或許可以解釋 transformer 現(xiàn)在的「霸主」地位。

原文鏈接：https://www.quantamagazine.org/researchers-glimpse-how-ai-gets-so-good-at-language-processing-20220414/