本論文第一作者是汪海洋，北京大學(xué)20級博士生，目前主要關(guān)注是通用模型的架構(gòu)設(shè)計和學(xué)習(xí)算法。指導(dǎo)教授主要包括王立威，北京大學(xué)智能學(xué)院教授；Bernt Schiele，德國馬普計算所教授；Federico Tombari 谷歌人工智能科學(xué)家等。

新一代通用靈活的網(wǎng)絡(luò)結(jié)構(gòu) TokenFormer: Rethinking Transformer Scaling with Tokenized Model Parameters 來啦！

TokenFormer 不僅像原始 Transformer 一樣 Token 化了 input data，并且 Token 化了網(wǎng)絡(luò)參數(shù)，將 attention 機(jī)制拓展到 Token 和 parameters 的交互中，最大化了 Transformer 的靈活性，真正得到了一個 Fully attention-based 的網(wǎng)絡(luò)結(jié)構(gòu)。

這種方式打破了原有人們區(qū)別看待 data 和 model 的觀念，即所有的計算都?xì)w納為不同類型的 Token（e.g., data, param token）通過靈活的 attention 來交互。得益于這一靈活的性質(zhì)，TokenFormer 允許 incremental scaling model size，基于訓(xùn)好的模型上增量的拓展新的更大的模型，大大節(jié)省了計算的開銷：

這項名為 TokenFormer 的新工作，由谷歌，馬普計算所和北大的研究者提出，在 Twitter，HackerNews, Reddit 上得到廣泛的討論和關(guān)注 (Twitter 上有 150K + 的瀏覽量)。

目前代碼、模型和項目主頁均已放出：

• 論文鏈接：https://arxiv.org/pdf/2410.23168
• 開源代碼：https://github.com/Haiyang-W/TokenFormer
• 開源模型：https://huggingface.co/Haiyang-W

背景介紹

得益于其處理各種數(shù)據(jù)的靈活性，Transformer 網(wǎng)絡(luò)結(jié)構(gòu)在各個 AI 領(lǐng)域都取得了巨大的成功。

Transformer 模型通常將處理單個 Token 所需的計算分為兩個部分：與其他 Token 的交互（Token-Token Interaction）和涉及模型參數(shù)的計算（Token-Parameter Interaction）。

Attention 促進(jìn)了 Token-Token 之間的交互，使現(xiàn)代通用基礎(chǔ)模型能夠?qū)⒍嗄B(tài)數(shù)據(jù)編碼成統(tǒng)一的 Token 序列，并有效捕捉它們之間的復(fù)雜依賴關(guān)系。

相反，Token-Parameter 計算主要依賴于固定的 linear projection，大大限制 model size 的 scaling。Scaling model 是通常改變模型結(jié)構(gòu)，往往需要從頭訓(xùn)練整個模型，帶來了過多的資源消耗，使其越來越不切實(shí)際。

在本文中，研究團(tuán)隊使用 token 這一概念建模所有的計算，即將 model parameters 也視為一種 token，網(wǎng)絡(luò)的計算統(tǒng)一為各種不同的 token ( e.g., data tokens and parameter tokens) 之間通過 attention 來進(jìn)行交互，大大增強(qiáng)了 Token-Parameter 交互的靈活性，從而能夠增量式的擴(kuò)展模型參數(shù)，有效地重用先前訓(xùn)練的模型，從而顯著降低了訓(xùn)練負(fù)擔(dān)。

為實(shí)現(xiàn)這一目標(biāo)，研究團(tuán)隊引入了 TokenFormer。統(tǒng)一 Token-Token 和 Token-Parameters Interaction 的計算。其 Token-Parameter attention 具有靈活性，并能夠處理可變數(shù)量的參數(shù)，從而本質(zhì)上最大化了 Transformer 的靈活性，增強(qiáng)了模型的可擴(kuò)展性。

TokenFormer 提供一種新的看待模型的視角，即網(wǎng)絡(luò)的計算就是一些 Tokens 相互任意交互。基于這些 Tokens （e.g., data token, parameter token, memory token）和 attention 機(jī)制可以靈活地構(gòu)造任意的網(wǎng)絡(luò)結(jié)構(gòu)。

該團(tuán)隊希望 TokenFormer 作為一種通用的網(wǎng)絡(luò)結(jié)構(gòu)，不僅在 incremental model scaling 上有貢獻(xiàn)，還在 Sparse Inference, Parameter-Efficient Tuning, Vision and Language Models, Device-Cloud Collaboration 和 Model Interpretability 等領(lǐng)域有更多的貢獻(xiàn)。

方法

Tokenformer 的核心創(chuàng)新是 Token-Parameter Attention（Pattention） Layer，它結(jié)合了一組 Trainable Tokens 作為 model parameters，并通過 cross-attention 來管理 Input Token 與這些 Parameter Tokens 之間的交互。

通過這種方式，Pattention 層引入了一個額外的維度 —Parameter Token 的數(shù)量，這一維度獨(dú)立于輸入和輸出維度。此解耦方式使得輸入數(shù)據(jù)可以與 variable number of parameters 進(jìn)行交互，提供了增量模型擴(kuò)展所需的靈活性。  

Pattention Layer: 具體來說，就是讓 input data 作為 query, 研究團(tuán)隊引入了兩組具有 n 個可學(xué)習(xí)的 Tokens：代表 key， 表示 value。輸出如下：

其中 Θ 是改進(jìn)的 softmax，為了防止梯度 exponential 帶來的梯度問題，

這里 f () 是任意非線性函數(shù)，默認(rèn)使用 gelu。。

研究團(tuán)隊使用 Pattention Layer 替換掉標(biāo)準(zhǔn) Transformer 中的所有的 linear projection，最大化 Transformer 的靈活性。

應(yīng)用：天生的增量式 Model Scaling

有了 TokenFormer 這一靈活的性質(zhì)，可以延伸出很多應(yīng)用。這里以增量式 model scaling 為例。

假設(shè)已經(jīng)訓(xùn)練好了一個 TokenFormer，其 key parameters 和 value parameters 計為 和。

如上圖所示，加入新的重新初始化的 key-value parameter pairs，計為 和，進(jìn)而組合成新的 key-value set,

然后使用 pattention layer，讓 input data 與 Parameter tokens 進(jìn)行交互。

這里直觀的理解就是每個 Key-Value 代表一種學(xué)好的 pattern，其組成一個巨大的知識庫。文中的 incremental scaling 就是在原有的知識庫上進(jìn)一步拓展訓(xùn)練。

實(shí)驗結(jié)果

增量式 model scaling：如下右圖所示，模型在已經(jīng)訓(xùn)好的 124M 的模型的基礎(chǔ)上，采用增量式訓(xùn)練，只用十分之一的數(shù)據(jù)就可以達(dá)到從頭訓(xùn)練策略相近的性能，讓模型可以不斷迭代，真正地活起來了。

Language Modeling：如下表所示，研究團(tuán)隊比較了 Transformer-based 的模型和 TokenFormer 在語言建模上的能力。

在相同規(guī)模、相同模型尺寸下， TokenFormer 在大大增加靈活性的前提下達(dá)到了比 Transformer 更好的 zero-shot 性能。這里研究團(tuán)隊 follow 了 pythia 標(biāo)準(zhǔn)的訓(xùn)練代碼以及數(shù)據(jù)集：Pile (300B)。上述結(jié)果展現(xiàn)了 TokenFormer 在語言模型建模上的能力。

Visual Modeling: 為了進(jìn)一步驗證 TokenFormer 的表達(dá)能力，研究團(tuán)隊還和標(biāo)準(zhǔn)的 vision transformer 進(jìn)行了對比。

在 ImageNet-1K 的監(jiān)督訓(xùn)練的 setting 上，使用相同的訓(xùn)練策略， TokenFormer 的性能超過了 vision-transformer，驗證了其在 visual modeling 上的能力。

未來研究方向

極致的專家混合（Mixture-of-Experts）范式

研究團(tuán)隊認(rèn)為 Tokenformer 是專家混合（MoE）框架的極致實(shí)例化，其中每一組鍵 - 值參數(shù)對都充當(dāng)一個獨(dú)立的專家。這種創(chuàng)新的類 MoE 架構(gòu)有可能顯著減少與 Token-Parameter 交互相關(guān)的計算成本。

新的參數(shù)高效微調(diào)范式

Tokenformer 的擴(kuò)展方法通過集成額外的 key-value parameter pairs，展現(xiàn)了一種參數(shù)高效的微調(diào)策略。當(dāng)面對新任務(wù)或數(shù)據(jù)集時，該模型可以通過加入新的 Token Parameters 來擴(kuò)展其預(yù)訓(xùn)練參數(shù)，從而快速適應(yīng)特定任務(wù)需求。

整合視覺和語言模型

利用 Tokenformer 的參數(shù)高效微調(diào)能力，可以實(shí)現(xiàn)視覺和語言模態(tài)的無縫集成。具體方法是將預(yù)訓(xùn)練的 Visual Tokenformer 和 Language Tokenformer 的 key-value parameter Tokens 統(tǒng)一為一個參數(shù)集，然后引入新的 Trainable Tokens 來執(zhí)行視覺 - 語言對齊和指令微調(diào)。

端云協(xié)同

Tokenformer 可以在設(shè)備 - 云協(xié)作中充當(dāng)云端知識庫，為設(shè)備端的大語言模型（LLM）提供支持，其中每組 key-value parameter tokens 代表一個可學(xué)習(xí)模式，通過設(shè)備進(jìn)行實(shí)時處理，并利用云端執(zhí)行密集任務(wù)。

增強(qiáng)模型的可解釋性

由于 Tokenformer 完全基于注意力機(jī)制，它自然受益于在 Token-Parameter 交互中與注意力相關(guān)的可解釋性特性。這一特點(diǎn)增強(qiáng)了模型的可解釋性，為 AI 社區(qū)開發(fā)更透明、易理解的模型貢獻(xiàn)力量。