3行代碼就能可視化Transformer的奧義
本文經(jīng)AI新媒體量子位(公眾號ID:QbitAI)授權(quán)轉(zhuǎn)載,轉(zhuǎn)載請聯(lián)系出處。
Transformer背后的計算模型是什么?
就是最資深的煉丹er,也很難直觀地解釋這一結(jié)構(gòu)中的信息如何流動,數(shù)據(jù)又是如何計算的。
但現(xiàn)在,只要寫上幾行代碼,就能將Transformer的流程可視化。

△RASP代碼被編譯成一個兩層三頭的Transformer架構(gòu)。
這就是RASP,以色列科學(xué)家最新研發(fā)的一種簡單的序列處理語言。
他們希望通過這一語言去計算和解決更高層次的抽象問題。
熱心網(wǎng)友立馬點贊:
我愛這篇論文,因為我在閱讀機器學(xué)習論文時老是有“可視化Transformer”的想法,現(xiàn)在終于可以拿這個去玩兒了。而且用的都是一些簡單的高階函數(shù),估計移植到Python或者放在Jupyter上跑也不難。

那么,這個叫做RASP的語言到底如何“像Transformer一樣思考”的呢?

這門新方法是怎么形成的
讓我們先追溯到上個世紀,RNN被抽象成為有限狀態(tài)自動機(Finite State Automata)這一計算模型的例子。
這一例子說明了RNN與自動機之間存在著非常直觀的相似性。
經(jīng)過這種抽象之后,圍繞RNN架構(gòu)的變體或訓(xùn)練模型就展開了更為直觀的討論和思考。
而Transformer就不夠直觀,因為它在處理序列時擁有獨特的信息流約束。

到了2020年,出現(xiàn)了基于Transformer提出的可以識別Dyck-k語言的構(gòu)架。
這時的Transformer網(wǎng)絡(luò)在一階段的邏輯公式中作為輸入時,能夠?qū)W習并展現(xiàn)出一種多步驟的邏輯推理。
這就對研究團隊產(chǎn)生了啟發(fā):能不能編寫一種程序,對輸入表達式也進行類似的邏輯推理呢?
RASP(Access Sequence Processing Language)就這樣誕生了。

將Transformer編碼器的基本組件:注意力(attention)和前饋計算層(feed-forward computation)映射到簡單的原語(primitives)中,然后圍繞它們形成一種新的編碼語言。
這就是RASP,全稱限制性訪問序列處理語言。
當然,你也可以將RASP認為是一種Transformer結(jié)構(gòu)的計算方法:
將Transformer網(wǎng)絡(luò)的技術(shù)細節(jié)抽象而出,使其支持符號化程序,然后“編譯”到Transformer硬件上,再定義一系列的注意力和多層感知器操作。
這就是像Transformer一樣思考。
怎樣一個計算模型
好,現(xiàn)在來看看這個RASP到底是由什么構(gòu)成的:
- 內(nèi)置的序列操作符(built-in s-ops)
在RASP中編程的目標就是將這些序列操作符組合成起來,去計算最終的目標任務(wù)。
- 元素性操作(Elementwise Operations)
反映了Transformer的前饋子層。
- 其他操作
所有除元素性操作之外的操作,比如寬度選擇(selector_width operation)。
而Transformer的核心是什么?注意力機制。
RASP允許每個輸入定義選擇一種注意力模式,再通過加權(quán)平均,將輸入值混合成一個新的輸出。
那么它是如何進行編譯呢?
當你輸入這樣3行代碼:

RASP就能將序列操作符和輸入對的編譯流程可視化:

可以看到,第一層的頭并不均勻地關(guān)注整個序列,而是在序列的最后位置上展示出了偏向性。
而Transformer第二層的注意頭與程序中的反轉(zhuǎn)選擇器(flip selector)的行為完全對應(yīng):

表現(xiàn)怎么樣?
RASP能夠有效地預(yù)測一個Transformer解決某任務(wù)所需的最大層數(shù)和頭數(shù)嗎?
研究者在每個任務(wù)上訓(xùn)練4個規(guī)定大小的Transformer,然后測試它們的準確性:結(jié)果是大多數(shù)Transformer都達到了99.5%以上的準確率。

這讓RASP可以作為一種用來推斷Transformer擅長哪些任務(wù),或其架構(gòu)如何影響這些任務(wù)的工具。
而在根據(jù)RASP預(yù)測尺寸訓(xùn)練的Transformer減少它的頭數(shù)和層數(shù)時,準確率會發(fā)生大幅下降**。

△用L、H表示編譯的RASP程序所預(yù)測的層數(shù)和頭數(shù)。
這也表示RASP擁有非常嚴格的約束性,在邊界的緊密性(Tightness of the bound)上達到了極高的精度。
團隊介紹
論文一作Gail Weiss是以色列理工學(xué)院的博士生,目前在學(xué)院內(nèi)的計算機科學(xué)系中工作。
其中二作Yoav Goldberg來自以色列的巴伊蘭大學(xué),目前就職于艾倫人工智能研究院(Allen Institute for Artificial Intelligence)。
論文地址:
https://arxiv.org/abs/2106.06981
下載:
https://github.com/tech-srl/RASP
參考鏈接:
https://news.ycombinator.com/item?id=27528004