過去十多年，AI的飛速發(fā)展主要是工程實踐上的進步，AI理論并沒有起到指導(dǎo)算法開發(fā)的作用，經(jīng)驗設(shè)計的神經(jīng)網(wǎng)絡(luò)依然是一個黑盒。

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

最近，馬毅教授團隊發(fā)布了最新研究成果，設(shè)計了一個完全可用數(shù)學(xué)解釋的白盒Transformer模型CRATE，并在真實世界數(shù)據(jù)集ImageNet-1K上取得了接近ViT的性能。

代碼鏈接：https://github.com/Ma-Lab-Berkeley/CRATE

論文鏈接：https://arxiv.org/abs/2306.01129

在這篇論文中，研究人員認為，表示學(xué)習(xí)的目標(biāo)是壓縮和轉(zhuǎn)換數(shù)據(jù)（例如token集合）的分布，以支持在不相干子空間（incoherent subspace）上的低維高斯分布混合，最終表征的質(zhì)量可以通過稀疏率降低（sparse rate reduction）的統(tǒng)一目標(biāo)函數(shù)來度量。

從這個角度來看，流行的深度網(wǎng)絡(luò)模型，如Transformer等可以很自然地被認為是實現(xiàn)迭代方案（realizing iterative schemes）以逐步優(yōu)化該目標(biāo)。

特別是，研究結(jié)果表明標(biāo)準(zhǔn)Transformer塊可以從對該目標(biāo)的互補部分的交替優(yōu)化中派生出：多頭自注意力運算符可以被視為通過最小化有損編碼率來壓縮token集合的梯度下降步驟，而隨后的多層感知器可以被視為嘗試稀疏化token的表示。

這一發(fā)現(xiàn)也促進設(shè)計了一系列在數(shù)學(xué)上完全可解釋的白盒Transformer類深度網(wǎng)絡(luò)架構(gòu)，盡管設(shè)計上很簡單，但實驗結(jié)果表明，這些網(wǎng)絡(luò)確實學(xué)會了優(yōu)化設(shè)計目標(biāo)：壓縮和稀疏化了大規(guī)模真實世界視覺數(shù)據(jù)集（如ImageNet）的表示，并實現(xiàn)了接近高度工程化Transformer模型（ViT）的性能。

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

馬毅教授于1995年獲得清華大學(xué)自動化與應(yīng)用數(shù)學(xué)雙學(xué)士學(xué)位，并于1997年獲加州大學(xué)伯克利分校EECS碩士學(xué)位，2000年獲數(shù)學(xué)碩士學(xué)位與EECS博士學(xué)位。

2018年馬毅教授加入加州大學(xué)伯克利分校電子工程與計算機科學(xué)系，今年1月加入香港大學(xué)出任數(shù)據(jù)科學(xué)研究院院長，最近又接任香港大學(xué)計算系主任。

主要研究方向為3D計算機視覺、高維數(shù)據(jù)的低維模型、可擴展性優(yōu)化和機器學(xué)習(xí)，最近的研究主題包括大規(guī)模3D幾何重構(gòu)和交互以及低維模型與深度網(wǎng)絡(luò)的關(guān)系。

讓Transformer變白盒

這篇論文的主要目的在于用一個更統(tǒng)一的框架以設(shè)計類似Transformer的網(wǎng)絡(luò)結(jié)構(gòu)，從而實現(xiàn)數(shù)學(xué)上的可解釋性和良好的實際性能。

為此，研究人員提出學(xué)習(xí)一個增量映射（incremental mappings）序列，以獲得輸入數(shù)據(jù)（token集合）的最小壓縮和最稀疏的表征，優(yōu)化一個統(tǒng)一的目標(biāo)函數(shù)，即稀疏率降低。

這個框架統(tǒng)一了「Transformer模型和自注意力」、「擴散模型和降噪」、「結(jié)構(gòu)化查找和率降低」（Structure-seeking models and rate reduction）三種看似不同的方法，并表明類似Transformer的深層網(wǎng)絡(luò)層可以自然地從展開迭代優(yōu)化（unrolling iterative optimization）方案中導(dǎo)出， 以增量地優(yōu)化稀疏率降低目標(biāo)。

映射的目標(biāo)

Self-Attention via Denoising Tokens Towards Multiple Subspaces

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

Self-Attention via Compressing Token Sets through Optimizing Rate Reduction

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

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

研究人員展示了在Transformer塊中緊隨多頭自注意力層后面的多層感知機可以被解釋為（并且可以被替換為）一個層，該層通過構(gòu)建token表征稀疏編碼來逐步優(yōu)化稀疏率降低目標(biāo)剩余部分。

CRATE

結(jié)合上述理解，研究人員創(chuàng)建了一個全新的的白盒Transformer架構(gòu)CRATE（Coding RAte reduction TransformEr），學(xué)習(xí)目標(biāo)函數(shù)、深度學(xué)習(xí)架構(gòu)和最終學(xué)習(xí)到的表征都完全可以用數(shù)學(xué)解釋，其中每一層執(zhí)行交替最小化算法（alternating minimization algorithm）的一個步驟，以優(yōu)化稀疏率降低目標(biāo)。

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

實驗部分

研究人員的實驗?zāi)繕?biāo)不僅僅是在使用基本設(shè)計的情況下與其他精心設(shè)計的Transformer競爭，還包括：

1、與通常僅在端到端性能上評估的經(jīng)驗設(shè)計的黑盒網(wǎng)絡(luò)不同，白盒設(shè)計的網(wǎng)絡(luò)可以查看深層架構(gòu)的內(nèi)部，并驗證學(xué)習(xí)網(wǎng)絡(luò)的層是否確實執(zhí)行其設(shè)計目標(biāo)，即對目標(biāo)進行增量優(yōu)化。

2、盡管CRATE架構(gòu)很簡單，但實驗結(jié)果應(yīng)當(dāng)驗證該架構(gòu)的巨大潛力，即可以在大規(guī)模真實世界的數(shù)據(jù)集和任務(wù)上取得與高度工程化Transformer模型相匹配的性能。

模型架構(gòu)

通過變化token維度、頭數(shù)和層數(shù)，研究人員創(chuàng)建了四個不同規(guī)模的CRATE模型，表示為CRATE-Tiny，CRATE-Small，CRATE-Base和CRATE-Large

數(shù)據(jù)集和優(yōu)化

文中主要考慮ImageNet-1K作為測試平臺，使用Lion優(yōu)化器來訓(xùn)練具有不同模型規(guī)模的CRATE模型。

同時還評估了CRATE的遷移學(xué)習(xí)性能：在ImageNet-1K上訓(xùn)練的模型作為預(yù)訓(xùn)練模型，然后在幾個常用的下游數(shù)據(jù)集（CIFAR10/100、Oxford Flowers、Oxford-IIT-Pets）上對CRATE進行微調(diào)。

CRATE的層實現(xiàn)設(shè)計目標(biāo)了嗎？

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

結(jié)果表明，CRATE與原始的設(shè)計目標(biāo)非常契合：一旦學(xué)習(xí)完畢，基本上通過其層逐漸學(xué)習(xí)對表示進行壓縮和稀疏化。

在其他規(guī)模的CRATE模型以及中間模型檢查點上測量壓縮和稀疏化項后可以發(fā)現(xiàn)，實驗結(jié)果依然非常一致，具有更多層的模型往往能更有效地優(yōu)化目標(biāo)，驗證了之前對每個層角色的理解。

性能對比

通過測量ImageNet-1K上的最高準(zhǔn)確率以及在幾個廣泛使用的下游數(shù)據(jù)集上的遷移學(xué)習(xí)性能來研究所提出的網(wǎng)絡(luò)的經(jīng)驗性能。

由于設(shè)計的架構(gòu)在注意力塊（MSSA）和MLP塊（ISTA）中都利用了參數(shù)共享，所以CRATE-Base模型（2208萬）與ViT-Small（2205萬）的參數(shù)數(shù)量相似。

可以看到，在模型參數(shù)數(shù)量相似的情況下，文中提出的網(wǎng)絡(luò)實現(xiàn)了與ViT相似的ImageNet-1K和遷移學(xué)習(xí)性能，但CRATE的設(shè)計更簡單，可解釋性強。

此外，在相同的訓(xùn)練超參數(shù)下，CRATE還可以繼續(xù)擴展，即通過擴大模型的規(guī)模不斷提高性能，而在ImageNet-1K上直接擴大ViT的規(guī)模并不總是能帶來一致的性能改善。

也就是說，CRATE網(wǎng)絡(luò)盡管簡單，但已經(jīng)可以在大規(guī)模的真實世界數(shù)據(jù)集上學(xué)習(xí)所需的壓縮和稀疏表示，并在各種任務(wù)（如分類和遷移學(xué)習(xí)）上取得與更工程化Transformer網(wǎng)絡(luò)（如ViT）相當(dāng)?shù)男阅堋?/span>