3行程式碼就能視覺化Transformer的奧義 | 開源

博雯 發自 凹非寺量子位 報道 | 公眾號 QbitAI

Transformer背後的計算模型是什麼？

就是最資深的煉丹er，也很難直觀地解釋這一結構中的資訊如何流動，資料又是如何計算的。

但現在，只要寫上幾行程式碼，就能將Transformer的流程視覺化。

△RASP程式碼被編譯成一個兩層三頭的Transformer架構。

這就是RASP，以色列科學家最新研發的一種簡單的序列處理語言。

他們希望通過這一語言去計算和解決更高層次的抽象問題。

熱心網友立馬點贊：

我愛這篇論文，因為我在閱讀機器學習論文時老是有「視覺化Transformer」的想法，現在終於可以拿這個去玩兒了。而且用的都是一些簡單的高階函數，估計移植到Python或者放在Jupyter上跑也不難。

那麼，這個叫做RASP的語言到底如何「像Transformer一樣思考」的呢？

這門新方法是怎麼形成的

讓我們先追溯到上個世紀，RNN被抽象成為有限狀態自動機（Finite State Automata）這一計算模型的例子。

這一例子說明了RNN與自動機之間存在著非常直觀的相似性。

經過這種抽象之後，圍繞RNN架構的變體或訓練模型就展開了更為直觀的討論和思考。

而Transformer就不夠直觀，因為它在處理序列時擁有獨特的資訊流約束。

到了2020年，出現了基於Transformer提出的可以識別Dyck-k語言的構架。

這時的Transformer網路在一階段的邏輯公式中作為輸入時，能夠學習並展現出一種多步驟的邏輯推理。

這就對研究團隊產生了啟發：能不能編寫一種程式，對輸入表示式也進行類似的邏輯推理呢？

RASP（Access Sequence Processing Language）就這樣誕生了。

將Transformer編碼器的基本元件：注意力（attention）和前饋計算層（feed-forward computation）對映到簡單的原語（primitives）中，然後圍繞它們形成一種新的編碼語言。

這就是RASP，全稱限制性訪問序列處理語言。

當然，你也可以將RASP認為是一種Transformer結構的計算方法：

將Transformer網路的技術細節抽象而出，使其支援符號化程式，然後「編譯」到Transformer硬體上，再定義一系列的注意力和多層感知器操作。

這就是像Transformer一樣思考。

怎樣一個計算模型

好，現在來看看這個RASP到底是由什麼構成的：

內建的序列操作符（built-in s-ops）在RASP中程式設計的目標就是將這些序列操作符組合成起來，去計算最終的目標任務。

元素性操作（Elementwise Operations）反映了Transformer的前饋子層。

其他操作所有除元素性操作之外的操作，比如寬度選擇（selector_width operation）。

而Transformer的核心是什麼？注意力機制。

RASP允許每個輸入定義選擇一種注意力模式，再通過加權平均，將輸入值混合成一個新的輸出。

那麼它是如何進行編譯呢？

當你輸入這樣3行程式碼：

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

可以看到，第一層的頭並不均勻地關注整個序列，而是在序列的最後位置上展示出了偏向性。

而Transformer第二層的注意頭與程式中的反轉選擇器（flip selector）的行為完全對應：

表現怎麼樣？

RASP能夠有效地預測一個Transformer解決某任務所需的最大層數和頭數嗎？

研究者在每個任務上訓練4個規定大小的Transformer，然後測試它們的準確性：結果是大多數Transformer都達到了99.5%以上的準確率。

這讓RASP可以作為一種用來推斷Transformer擅長哪些任務，或其架構如何影響這些任務的工具。

而在根據RASP預測尺寸訓練的Transformer減少它的頭數和層數時，準確率會發生大幅下降**。

△用L、H表示編譯的RASP程式所預測的層數和頭數。

這也表示RASP擁有非常嚴格的約束性，在邊界的緊密性（Tightness of the bound）上達到了極高的精度。

團隊介紹

論文一作Gail Weiss是以色列理工學院的博士生，目前在學院內的計算機科學系中工作。

其中二作Yoav Goldberg來自以色列的巴伊蘭大學，目前就職於艾倫人工智慧研究院（Allen Institute for Artificial Intelligence）。

論文地址：https://arxiv.org/abs/2106.06981

下載：https://github.com/tech-srl/RASP

參考連結：https://news.ycombinator.com/item?id=27528004