LeCun力挺，馬毅教授五年集大成之作：完全數學可解釋的白盒Transformer，性能不輸ViT

過去十多年，AI的快速發展主要是工程實踐上的進步，AI理論並沒有起到指導演算法開發的作用，經驗設計的神經網路依然是一個黑盒子。

而隨著ChatGPT的爆火，AI的能力也被不斷誇大、炒作，甚至到了威脅、綁架社會的地步，讓Transformer架構設計變得透明已刻不容緩！

最近，馬毅教授團隊發布了最新研究成果，設計了一個完全可用數學解釋的白盒Transformer模型CRATE， 並在真實世界資料集ImageNet-1K上取得了接近ViT的效能。

程式碼連結：https://github.com/Ma-Lab-Berkeley/CRATE

論文連結：https://arxiv.org/abs/2306.01129

在這篇論文中，研究人員認為，表示學習的目標是壓縮和轉換資料（例如token集合）的分佈，以支援在不相干子空間（incoherent subspace）上的低維高斯分佈混合，最終表徵的品質可以透過稀疏率降低（sparse rate reduction）的統一目標函數來度量。

從這個角度來看，流行的深度網路模型，如Transformer等可以自然地被認為是實現迭代方案（realizing iterative schemes）以逐步優化該目標。

特別是，研究結果表明標準Transformer區塊可以從對該目標的互補部分的交替優化中派生出：多頭自註意力運算符可以被視為透過最小化有損編碼率來壓縮token集合的梯度下降步驟，而隨後的多層感知器可以被視為嘗試稀疏化token的表示。

這項發現也促進設計了一系列在數學上完全可解釋的白盒Transformer類別深度網路架構，儘管設計上很簡單，但實驗結果表明，這些網路確實學會了優化設計目標：壓縮和稀疏化了大規模真實世界視覺資料集（如ImageNet）的表示，並實現了接近高度工程化Transformer模型（ViT）的性能。

圖靈獎得主Yann LeCun對馬毅教授的工作也表示贊同，認為Transformer使用LISTA（Learned Iterative Shrinkage and Thresholding Algorithm）類似的方法增量地優化稀疏壓縮。

馬毅教授於1995年獲得清華大學自動化與應用數學雙學士學位，並於1997年獲加州大學柏克萊分校EECS碩士學位，2000年獲數學碩士學位與EECS博士學位。

2018年馬毅教授加入加州大學柏克萊分校電子工程與電腦科學系，今年1月加入香港大學出任資料科學研究院院長，最近又接任香港大學計算系主任。

主要研究方向為3D電腦視覺、高維度資料的低維度模型、可擴展性最佳化和機器學習，最近的研究主題包括大規模3D幾何重構和互動以及低維模型與深度網路的關係。

讓Transformer變成白盒子

這篇論文的主要目的在於用一個更統一的框架以設計類似Transformer的網路結構，從而實現數學上的可解釋性和良好的實際性能。

為此，研究人員提出學習一個增量映射（incremental mappings）序列，以獲得輸入資料（token集合）的最小壓縮和最稀疏的表徵，優化一個統一的目標函數，即稀疏率降低。

這個框架統一了「Transformer模型和自註意力」、「擴散模型和降噪」、「結構化查找和率降低」（Structure-seeking models and rate reduction）三種看似不同的方法，並表明類似Transformer的深層網絡層可以自然地從展開迭代優化（unrolling iterative optimization）方案中導出， 以增量地優化稀疏率降低目標。

#已對應的目標

Self-Attention via Denoising Tokens Towards Multiple Subspaces

研究人員使用一個理想化的token分佈模型表明，如果朝著低維子空間系列迭代去噪，相關的評分函數就會呈現出類似Transformer中的自註意力操作符的顯式形式。

Self-Attention via Compressing Token Sets through Optimizing Rate Reduction

研究人員將多頭自註意力層推導為一個展開的梯度下降步，以最小化速率降低的有損編碼率部分，從而展現了將自註意力層解釋為壓縮token表徵的另一種解釋方法。

MLP via Iterative Shrinkage-Thresholding Algorithms (ISTA) for Sparse Coding

研究人員展示了在Transformer區塊中緊接著多頭自註意力層後面的多層感知機可以被解釋為（並且可以被替換為）一個層，該層通過構建token表徵稀疏編碼來逐步優化稀疏率降低目標剩餘部分。

CRATE

結合上述理解，研究人員創建了一個全新的的白盒Transformer架構CRATE（Coding RAte reduction TransformEr），學習目標函數、深度學習架構和最終學習到的表徵都完全可以用數學解釋，其中每一層執行交替最小化演算法（alternating minimization algorithm）的一個步驟，以優化稀疏率降低目標。

可以注意到，CRATE在建置的每個階段都選擇了盡可能最簡單的建置方式，只要新建置的部分保持相同的概念角色，就可以直接替換，並獲得一個新的白盒架構。

實驗部分

研究人員的實驗目標不僅是在使用基本設計的情況下與其他精心設計的Transformer競爭，還包括：

1、與通常僅在端到端效能上評估的經驗設計的黑盒網路不同，白盒設計的網路可以查看深層架構的內部，並驗證學習網路的層是否確實執行其設計目標，即對目標進行增量最佳化。

2、儘管CRATE架構很簡單，但實驗結果應當驗證該架構的巨大潛力，即可以在大規模真實世界的資料集和任務上取得與高度工程化Transformer模型相匹配的效能。

模型架構

#透過變化token維度、頭數和層數，研究人員創建了四個不同規模的CRATE模型，表示為CRATE-Tiny，CRATE-Small，CRATE-Base和CRATE-Large

##資料集和最佳化

文中主要考慮ImageNet-1K作為測試平台，使用Lion優化器來訓練具有不同模型規模的CRATE模型。

同時也評估了CRATE的遷移學習效能：在ImageNet-1K上訓練的模型作為預訓練模型，然後在幾個常用的下游資料集（CIFAR10/100、Oxford Flowers、Oxford-IIT-Pets）上對CRATE進行微調。

CRATE的圖層實現設計目標了嗎？

隨著層索引的增加，可以看到CRATE-Small模型在大多數情況下的壓縮和稀疏化項目都得到了提升，最後一層稀疏度量的增加是由於用於分類的額外線性層。

結果表明，CRATE與原始的設計目標非常契合：一旦學習完畢，基本上透過其層逐漸學習對錶示進行壓縮和稀疏化。

在其他規模的CRATE模型以及中間模型檢查點上測量壓縮和稀疏化項後可以發現，實驗結果依然非常一致，具有更多層的模型往往能更有效地優化目標，驗證了先前對每個層角色的理解。

性能對比

#透過測量ImageNet-1K上的最高準確率以及在幾個廣泛使用的下游資料集上的遷移學習性能來研究所提出的網路的經驗性能。

由於設計的架構在註意力區塊（MSSA）和MLP區塊（ISTA）中都利用了參數共享，所以CRATE-Base模型（2208萬）與ViT-Small（2205萬）的參數數量相似。

可以看到，在模型參數數量相似的情況下，文中提出的網路實現了與ViT相似的ImageNet-1K和遷移學習效能，但CRATE的設計更簡單，可解釋性強。

此外，在相同的訓練超參數下，CRATE還可以繼續擴展，即透過擴大模型的規模不斷提高效能，而在ImageNet-1K上直接擴大ViT的規模並不總是能帶來一致的性能改善。

也就是說，CRATE網路儘管簡單，但已經可以在大規模的真實世界資料集上學習所需的壓縮和稀疏表示，並在各種任務（如分類和遷移學習）上取得與更工程化Transformer網路（如ViT）相當的效能。

以上是LeCun力挺，馬毅教授五年集大成之作：完全數學可解釋的白盒Transformer，性能不輸ViT的詳細內容。更多資訊請關注PHP中文網其他相關文章！