原文標題：《像 Transformer 一樣思考！DeepMind 發(fā)布全新模型設計工具 Tracr：從可解釋邏輯反向搭建模型》代碼直接編譯成 Transformer 模型，做實驗從未如此輕！「可解釋性」一直是深學習難以繞過的難題，用無法理解模型的運行機制也就無法放心地把模型應到實際場景中。最近來自黎世聯邦理工學院和 DeepMind 的研究人員提出了一種全新的模型構工具 Tracr，直接由人來根據「已知的機制」對不同任務編寫代碼，然由 Tracr 編譯成模型的權重，讓模型的解釋得更輕松！論文鏈接：https://arxiv.org/ pdf / 2301.05062.pdf代碼鏈接：https://github.com/ deepmind / tracrTracr 的輸入為采用領域特定語 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 可以取哪些值。因為計算確定的，基于有限的輸入匯量和上下文大小，所有點都有一組有限的輸出值所以第二步的主要操作就對圖進行遍歷，并為每個點標注出其可能的輸出；注使用簡單的啟發(fā)式方法確保找到的是 s-op 值集合的超集。3. 獨立地翻譯 s-ops獨立考慮計算圖中的每個節(jié)點，將其轉化為一個 craft 組件；元素操作翻譯為 MLP 塊，選擇-聚合操作翻譯為注意力塊。使人工設計的 MLP 和注意力模塊庫來近似數字和類輸入和輸出的任意函數將具有分類輸入和輸出的 MLPs 作為查找表使；帶有數字輸入和輸出的 MLP 使用基于通用函數近似定理的明重結構。對于意層，把選擇器翻譯成??_{????} 運算符，把相應的聚合運算翻譯成??_{????} 運算符。目前只支持對分類輸入的注意駁。4. 把組件分配到層上為了構建一個 Transformer 模型，需要將計算圖中的所有 craft 組件分配給各層。理想情下，目標是找到最小的模來進行所需的計算，但一可以將其表述為一個具有個約束條件的組合優(yōu)化問：Transformer 結構具有交替的注意力和 MLP 層，而且所有相互依賴的計算都需要有正的順序。出于范圍上的考，可以用啟發(fā)式方法解決個問題。首先，計算出從入到一個給定節(jié)點的最長徑，其路徑長度是可以分給該節(jié)點的層數的一個上；然后應用額外的啟發(fā)式法，將層與可以并行計算塊結合起來。這種方法可返回一個正確但有時是次的層分配。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 模型新方法的試驗臺。參考料：https://arxiv.org/pdf/2301.05062.pdf本文來自微信公眾號：新智元 （ID：AI_era）