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

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

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

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

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

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

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

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

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

圖靈獎(jiǎng)得主Yann LeCun對(duì)馬毅教授的工作也表示贊同，認(rèn)為T(mén)ransformer使用LISTA（Learned Iterative Shrinkage and Thresholding Algorithm）類(lèi)似的方法增量地優(yōu)化稀疏壓縮。

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

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

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

讓Transformer變白盒

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

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

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

映射的目標(biāo)

Self-Attention via Denoising Tokens Towards Multiple Subspaces

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

Self-Attention via Compressing Token Sets through Optimizing Rate Reduction

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

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

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

CRATE

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

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

實(shí)驗(yàn)部分

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

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

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

模型架構(gòu)

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

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

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

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

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

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

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

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

性能對(duì)比

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

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

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

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

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