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

新一代通用靈活的網絡結構 TokenFormer: Rethinking Transformer Scaling with Tokenized Model Parameters 來啦！

TokenFormer 不僅像原始 Transformer 一樣 Token 化了 input data，并且 Token 化了網絡參數，將 attention 機制拓展到 Token 和 parameters 的交互中，最大化了 Transformer 的靈活性，真正得到了一個 Fully attention-based 的網絡結構。

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

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

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

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

背景介紹

得益于其處理各種數據的靈活性，Transformer 網絡結構在各個 AI 領域都取得了巨大的成功。

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

Attention 促進了 Token-Token 之間的交互，使現代通用基礎模型能夠將多模態數據編碼成統一的 Token 序列，并有效捕捉它們之間的復雜依賴關系。

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

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

為實現這一目標，研究團隊引入了 TokenFormer。統一 Token-Token 和 Token-Parameters Interaction 的計算。其 Token-Parameter attention 具有靈活性，并能夠處理可變數量的參數，從而本質上最大化了 Transformer 的靈活性，增強了模型的可擴展性。

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

該團隊希望 TokenFormer 作為一種通用的網絡結構，不僅在 incremental model scaling 上有貢獻，還在 Sparse Inference, Parameter-Efficient Tuning, Vision and Language Models, Device-Cloud Collaboration 和 Model Interpretability 等領域有更多的貢獻。

方法

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

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

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

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

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

研究團隊使用 Pattention Layer 替換掉標準 Transformer 中的所有的 linear projection，最大化 Transformer 的靈活性。

應用：天生的增量式 Model Scaling

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

假設已經訓練好了一個 TokenFormer，其 key parameters 和 value parameters 計為 和。

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

然后使用 pattention layer，讓 input data 與 Parameter tokens 進行交互。

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

實驗結果

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

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

在相同規模、相同模型尺寸下， TokenFormer 在大大增加靈活性的前提下達到了比 Transformer 更好的 zero-shot 性能。這里研究團隊 follow 了 pythia 標準的訓練代碼以及數據集：Pile (300B)。上述結果展現了 TokenFormer 在語言模型建模上的能力。

Visual Modeling: 為了進一步驗證 TokenFormer 的表達能力，研究團隊還和標準的 vision transformer 進行了對比。

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

未來研究方向

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

研究團隊認為 Tokenformer 是專家混合（MoE）框架的極致實例化，其中每一組鍵 - 值參數對都充當一個獨立的專家。這種創新的類 MoE 架構有可能顯著減少與 Token-Parameter 交互相關的計算成本。

新的參數高效微調范式

Tokenformer 的擴展方法通過集成額外的 key-value parameter pairs，展現了一種參數高效的微調策略。當面對新任務或數據集時，該模型可以通過加入新的 Token Parameters 來擴展其預訓練參數，從而快速適應特定任務需求。

整合視覺和語言模型

利用 Tokenformer 的參數高效微調能力，可以實現視覺和語言模態的無縫集成。具體方法是將預訓練的 Visual Tokenformer 和 Language Tokenformer 的 key-value parameter Tokens 統一為一個參數集，然后引入新的 Trainable Tokens 來執行視覺 - 語言對齊和指令微調。

端云協同

Tokenformer 可以在設備 - 云協作中充當云端知識庫，為設備端的大語言模型（LLM）提供支持，其中每組 key-value parameter tokens 代表一個可學習模式，通過設備進行實時處理，并利用云端執行密集任務。

增強模型的可解釋性

由于 Tokenformer 完全基于注意力機制，它自然受益于在 Token-Parameter 交互中與注意力相關的可解釋性特性。這一特點增強了模型的可解釋性，為 AI 社區開發更透明、易理解的模型貢獻力量。