馬毅團隊造出白盒Transformer 可解釋的大模型要來了嗎?
AGI 到底離我們還有多遠?在ChatGPT 引發的新一輪AI 爆發之後,一部分研究者指出,大語言模型具備透過觀察進行因果歸納的能力,但缺乏自己主動推理新的因果場景的能力。相較於樂觀預測的觀點,這意味著AGI 仍然是一個複雜而遙遠的目標。一直以來,AI 社群內有一種觀點:神經網路的學習過程可能只是對資料集的壓縮。
近日,柏克萊和香港大學的馬毅教授領導的一個研究團隊給出了自己的最新研究結果:包括GPT-4 在內的當前AI 系統所做的正是壓縮。
透過新提出的深度網路架構CRATE,他們透過數學方式驗證了這一點。
而更值得注意的是,CRATE 是一種白盒Transformer,其不僅能在幾乎所有任務上與黑盒Transformer 相媲美,而且還具備非常出色的可解釋性。
基於此,馬毅教授也在Twitter 上分享了一個有趣的見解:既然目前的AI 只是在壓縮數據,那麼就只能學習到數據中的相關性/ 分佈,所以就並不真正具備因果或邏輯推理或抽象思考能力。因此,當今的AI 還算不是AGI,即便近年來在處理和建模大量高維度和多模態資料方面,深度學習在實驗中取得了巨大的成功。
但很大程度上,這種成功可以歸功於深度網路能有效學習資料分佈中可壓縮的低維結構,並將該分佈轉換為簡約的(即緊湊且結構化的)表徵。這樣的表徵可用來幫助許多下游任務,例如視覺、分類、辨識和分割、生成。
表徵學習是透過壓縮式編碼和解碼實現的
為了更形式化地表述這些實踐背後的共同問題,我們可以將給定資料集的樣本看作是高維空間?^D 中的隨機向量x。
通常來說,x 的分佈具有比所在空間低得多的內在維度。一般來說,學習某個表徵通常是指學習一個連續的映射關係,如f (?),其可將x 轉換成另一個空間?^d(通常是低維空間)中的所謂特徵向量z。人們希望透過這樣一種映射:
能以一種緊湊且結構化的方式找到x 的低維內在結構並使用z 來表示它,然後藉此幫助分類或生成等後續任務。特徵z 可被視為原始資料x 的(學習到的)緊湊編碼,因此映射f 也稱為編碼器。
這樣一來,表徵學習的基礎問題(也也就是這份研究關注的核心問題)便成了:
為了衡量表徵的優劣,有什麼有數學原理保證且有效的衡量方法?
從概念上講,表徵z 的品質取決於它為後續任務找到x 的最相關和充分資訊的程度以及它表示該資訊的效率。
長期以來,人們都相信:所學到的特徵的「充分性」和「優良度」應根據具體任務而定義。舉個例子,在分類問題中,z 只需足以用於預測類別標籤y 即可。
為了理解深度學習或深度網路在這種類型的表徵學習中的作用,Tishby and Zaslavsky (2015) 在論文《Deep learning and the information bottleneck principle》中提出了資訊瓶頸框架,其提出:衡量特徵優良度的一種方法是最大化z 和y 之間的互訊息,同時最小化z 和x 之間的互資訊。
然而,近年來普遍通行的做法是先預先訓練一個大型深度神經網路(有些時候也被稱為基礎模型)來學習與任務無關的表徵。之後再針對多個具體任務對學習到的表徵進行微調。研究顯示這種方法能有效且有效率地處理許多不同資料模態的實踐任務。
請注意,這裡的表徵學習與針對特定任務的表徵學習非常不同。對於針對特定任務的表徵學習,z 只需能預測出特定的y 就足夠了。在與任務無關的情況下,所學到的表徵z 需要編碼幾乎所有與資料x 的分佈有關的關鍵資訊。也就是說,所學習到的表徵z 不僅是x 的內在結構的更緊湊和結構化表徵,還能以一定的可信度恢復出x。
因此,在與任務無關的情況下,人們自然會問:對於學習到的(特徵)表徵,一個衡量其優良度的有原理保證的度量應該是什麼?
研究者認為,一種有效方法(可能是唯一方法)是:為了驗證表徵z 是否已經編碼了有關x 的足夠信息,可以看透過如下(逆)映射(也被稱為解碼器或生成器)能從z 多好地恢復出x:
由於編碼器f 通常是有損壓縮,因此我們不應期望其逆映射能精確地恢復出x,而是恢復出一個近似
我們通常會尋找最優的編碼和解碼映射,使得解碼得到的
與x 最接近- 無論是樣本方面(例如,透過最小化預期均方誤差)或是在寬鬆的分佈意義上。
研究者將上述這個過程稱為壓縮式編碼和解碼或壓縮式自動編碼。這個想法與自動編碼器的原始目標高度相容,而自動編碼器則可被看作是經典的主成分分析泛化用於x 有線性的低維結構的情況。
過去十一年來,大量實驗已經清楚地表明:深度網路能夠非常有效地建模非線性編碼和解碼映射。
深度學習的幾乎所有應用都依賴於實現這樣的編碼或解碼方案,其方式是部分或完全地學習f 或g,當然它們可以分開或一起學習。
儘管從概念上講,解碼器g 應該是編碼器f 的「逆」映射,但在實踐中,我們一直不清楚編碼器和解碼器的架構有何關聯。在許多案例中,解碼器的架構設計與編碼器的關聯不大,通常是透過實驗測試和消融實驗選取的。
可以想見,一個優秀的表徵學習理論架構應能清楚揭示編碼器和解碼器架構之間的關係。而這正是這項研究希望達成的目標。
研究者總結了先前提出的相關方法,並將其分成了以下幾種情況:
透過壓縮打開現代深度網路的黑盒。
Transformer 模型和壓縮。
去雜訊擴散模型和壓縮。
促進低維度的測量:稀疏性和率下降。
展開最佳化:一個用於網路解釋和設計的統一範式。
詳情參看原論文。
這項研究的目標和貢獻
他們搭建了理論和實踐之間的橋樑。為此,這項研究提出了一個更完整和統一的框架。
一方面,這個新框架能對基於深度網路的許多看似不同的方法提供統一的理解,包括壓縮式編碼/ 解碼(或自動編碼)、率下降和去噪擴散。
另一方面,該框架可以指導研究者推導或設計深度網路架構,而這些架構不僅在數學上是完全可解釋的,而且在大規模現實世界圖像或文字資料集上的幾乎所有學習任務上都能獲得頗具競爭力的性能。
基於上述觀察,他們提出了一個白盒深度網路理論。更具體而言,他們為學習緊湊和結構化的表徵提出了一個統一的目標,也就是一個有原理保證的優良度量。對於學習到的表徵,該目標旨在既優化其在編碼率下降方面的內在複雜性,也優化其在稀疏性方面的外在複雜性。他們將該目標稱為稀疏率下降(sparse rate reduction)。圖3 給出了這個目標背後的直觀思想。
為了優化這個目標,他們提出可以學習一個增量映射序列,其能模擬展開目標函數的某些類似梯度下降的迭代最佳化方案。這自然地會得到一個類似Transformer 的深度網路架構,並且它完全是一個「白盒」—— 其最佳化目標、網路算子和學習到的表徵在數學上是完全可解釋的。
他們將這個白盒深度架構命名為CRATE 或CRATE-Transformer,這是Coding-RATE transformer 的縮寫。他們也透過數學方式證明這些增量映射在分佈的意義上是可逆的,並且它們的逆映射本質上由同一類數學算子構成。
因此,可以將幾乎完全一樣的CRATE 架構用於編碼器、解碼器或自動編碼器。如圖4 給出了一個自動編碼過程,其中每個編碼層f^?? 和解碼層g^{L-??} 是(部分)可逆的。
下圖給出了CRATE 白盒深度網路設計的「主循環」。
在將輸入資料X 預處理為一系列token Z^1 後,CRATE 會建立一個深度網絡,可將資料轉換為低維子空間的規範配置,其做法是針對分佈的局部模型進行連續壓縮生成Z ^{?+1/2},以及針對一個全域字典執行稀疏化,得到Z^{?+1}。透過重複堆疊這些模組並使用反向傳播訓練模型參數,可以獲得強大且可解釋的資料表徵。
下面則給出了CRATE 編碼器架構的一層。其完整架構就是將這些層串連起來,再加上一些初始token 化器、預處理頭和最後的針對特定任務的頭。
下圖比較了編碼器層和解碼器層,可以看到兩者是部分可逆的。
更多理論和數學描述請參閱原論文。
實驗評估
為了證明這個框架確實能將理論和實踐串連起來,他們在圖像和文字資料上執行了廣泛的實驗,在傳統Transformer 擅長的多種學習任務和設定上評估了CRATE 模型的實際表現。
下表給出了不同大小的CRATE 在不同資料集上的Top-1 準確度。
表2 給出了CRATE-Base 模型與MAE-Base 模型在訓練和驗證集上的平均重建損失。
令人驚訝的是,儘管其概念和結構很簡單,但CRATE 在所有任務和設定上都足以與黑盒版的對應方法媲美,這些任務包括透過監督學習進行影像分類、影像和語言的無監督掩碼補全、影像資料的自監督特徵學習、透過下一個詞預測的語言建模。
此外,CRATE 模型在實踐上還有其它優勢,每一層和網絡算子都有統計和幾何意義、學習到的模型的可解釋性顯著優於黑盒模型、其特徵具有語義含義(即它們可輕鬆用於將物件從背景中分割出來以及將其分成共用部件)。
下圖便給出了在每層? 的逐層token Z^? 表徵的可視化。
下圖展示了來自監督式CRATE 的自註意力圖。
注意由於資源限制,他們在實驗中沒有刻意追求當前最佳,因為那需要大量工程開發或微調。
儘管如此,他們表示這些實驗已經令人信服地驗證了新提出的白盒深度網路CRATE 模型是普遍有效的,並為進一步的工程開發和改進奠定了堅實的基礎。