「可解釋性」一直是深度學習難以繞過的難題，用戶無法理解模型的運行機制，也就無法放心地把模型應用到實際場景中。

最近來自蘇黎世聯(lián)邦理工學院和DeepMind的研究人員提出了一種全新的模型構造工具Tracr，直接由人來根據「已知的機制」針對不同任務編寫代碼，然后由Tracr編譯成模型的權重，讓模型的解釋變得更輕松！

論文鏈接：?https://arxiv.org/pdf/2301.05062.pdf?

代碼鏈接：?https://github.com/deepmind/tracr?

Tracr的輸入為采用領域特定語言RASP編寫的代碼，輸出為標準的、僅包含解碼器的、類似GPT的Transformer結構的權重。

在實驗部分，研究人員使用Tracr創(chuàng)建了一系列ground truth Transformers，實現了包括計算token頻率、排序和Dyck-n括號檢查等程序。

如何解釋模型？

可解釋性是理解機器學習模型的重要手段，但由于模型的實際運行方式仍然不清晰，所以目前的研究成果大多難以評估。

其中一個工作機理可解釋性（Mechanistic interpretability）嘗試對神經網絡進行逆向工程（reverse engineering）來對模型實現的算法給出一個機理上的解釋，并在一系列任務中取得進展，包括用于圖像分類的卷積神經網絡，Transformer語言模型等。

但該方法仍然存在相關工具少、對模型機制解釋不深、仍然需要研究人員進行創(chuàng)造性的解釋等問題。

評估機理解釋的標準方法結合了許多臨時實驗的證據，然而，由于這樣做的成本很高，許多方法只能在玩具模型中進行評估，或在真實模型中的少數不重要的circuits上進行評估。

Tracr的解決方案則是，通過將人類可讀的代碼「編譯」成神經網絡的權重，直接解決缺乏基礎機理解釋的問題。

也就是說，Tracr實際上起到的就是類似編譯器的作用。

 Tracr中主要涉及三個組件：

1. RASP代碼

RASP即Restricted Access Sequence Processing Language，是2021年提出的一個用于表達Transformer計算的語言，可以作為計算模型用來描述Transformers，并配有相應的解釋器運行RASP代碼。

可以把RASP程序視為一個計算圖，當根據給定的輸入token序列時，圖上的每個節(jié)點都會取一個特定值。

RASP語言中包括兩個基本的節(jié)點類型；序列操作（Sequence Operations， s-op），返回輸入值的token序列及indices序列；元素操作（Elementwise operations）；選擇-聚合操作等。

在大多數情況下，可以將RASP操作直接映射到Transformer模型的組件上，包括embedding, MLP層和Attention層。

2. 對RASP語言的修改

雖然RASP的操作可以直接映射到Transformers中，但仍然需要對RASP語言進行修改以對模型權重進行翻譯。

3. craft，Transformer的匯編語言

如果說RASP算是編譯的高級語言，那么craft就是匯編語言，它提供了比在純權重矩陣上操作更多的抽象性。

craft可以表示具有基維度（basis dimensions）的向量空間，及其相應的操作，可以用基方向的標簽來定義投影或其他線性操作。重要的是，craft抽象化了跟蹤權重矩陣中的填充的需要。

Tracr: Transformer編譯器

Tracr的代碼使用Python進行編寫，并將RASP實現嵌入到Python中，從而可以在Python中直接編寫RASP程序，比較方便地對變量編碼（variable encoding）進行標注。

在Tracr中，RASP程序是一個數據結構，通過向每個操作傳遞依賴關系來逐步構建，同時對RASP程序做了一些基本的簡化。

Tracr將RASP程序翻譯成Transformer權重的過程主要包括六個步驟：

1. 構建計算圖

追蹤整個程序以創(chuàng)建一個代表計算過程的有向圖。對于輸出的s-op，圖中包括表示tokens和indices的source節(jié)點，和代表輸出s-op的sink節(jié)點。

2. 推理s-op值

對于每個s-op，需要決定如何將其嵌入到剩余流中；為了使用類別編碼，需要知道一個s-op可以取哪些值。

因為計算是確定的，基于有限的輸入詞匯量和上下文大小，所有節(jié)點都有一組有限的輸出值。

所以第二步的主要操作就是對圖進行遍歷，并為每個節(jié)點標注出其可能的輸出；標注使用簡單的啟發(fā)式方法，確保找到的是s-op值集合的超集。

3. 獨立地翻譯s-ops

獨立考慮計算圖中的每個節(jié)點，并將其轉化為一個craft組件；元素操作翻譯為MLP塊，選擇-聚合操作翻譯為注意力塊。

使用人工設計的MLP和注意力模塊庫來近似數字和分類輸入和輸出的任意函數；將具有分類輸入和輸出的MLPs作為查找表使；帶有數字輸入和輸出的MLP使用基于通用函數近似定理的明確結構。

對于注意層，把選擇器翻譯成??_{????}運算符，把相應的聚合運算翻譯成??_{????}運算符。

目前只支持對分類輸入的注意力。

4. 把組件分配到層上

為了構建一個Transformer模型，需要將計算圖中的所有craft組件分配給各層。

理想情況下，目標是找到最小的模型來進行所需的計算，但一般可以將其表述為一個具有幾個約束條件的組合優(yōu)化問題：Transformer結構具有交替的注意力和MLP層，而且所有相互依賴的計算都需要有正確的順序。

出于范圍上的考慮，可以用啟發(fā)式方法解決這個問題。

首先，計算出從輸入到一個給定節(jié)點的最長路徑，其路徑長度是可以分配給該節(jié)點的層數的一個上限；然后應用額外的啟發(fā)式方法，將層與可以并行計算的塊結合起來。

這種方法可以返回一個正確但有時是次優(yōu)的層分配。

5. 構造一個craft模型

直接對模型組件的輸入和輸出空間進行求和作為構建的殘差流空間（residual stream space ）。

換句話說，將每個s-op嵌入到它自己的正交子空間中，這個子空間在整個網絡中只保留給它使用。

然后按照層分配決定的順序遍歷計算圖，并將組件堆疊起來，得到一個用craft表示的完整Transformer。

6. 組裝Transformer權重

最后，將模型的craft表示轉換為具體的模型權重。

首先把并行的MLP層合并為一個層，再把并行的注意力頭合并為一個層。其中在注意層中，把??_{????}和??_{????}矩陣分成????、????、????、????重量矩陣。

然后調整所有權重的形狀，并將其連接到Transformer架構中就可以推斷出模型配置（深度、層寬、殘差流大小等），以適應所創(chuàng)建的元素。

只需重新實現步驟6，就可以直接將Tracr擴展到支持任何其他Transformer的實現。

在可解釋性研究中的應用

Tracr可以加速受控實驗的驗證過程，以測試關于Transformer的計算結構的特定假設；通過這種方式，它也相當于成了可解釋性研究的實驗平臺。

研究人員對token計數、排序等例子編寫了RASP程序。

可解釋性工具的測試案例

編譯的模型可以很自然地作為測試「解釋忠實性」的基礎，并提供了一種偽造可解釋性技術給出的解釋的方法。

發(fā)展到最后，這些模型可以被用來建立可解釋性工具的測試案例庫，反過來又可以實現定量的評價指標。

替換模型組件

評估對模型如何工作的理解的另一種方法是用手工編碼的組件替換模型的部分內容。

例如有研究人員通過用他們自己理想化的實現來替換模型的組件以測試他們對Transformer如何實現模塊化加法的理解，結果發(fā)現該思路可以提高下游任務的性能，相當于有力地證明了所提出的解釋是正確的。

雖然Tracr是將一個算法編譯成一個完整的Transformer模型，但也可以通過修改代碼調整為只編譯訓練模型中的一部分，從而使得更容易評估對大型模型的理解。

理解模型現象和開發(fā)新技術

除了評估之外，編譯的模型可以作為研究circuits級別現象和開發(fā)解釋Transformer模型新方法的試驗臺。