Token化一切，甚至網路！北大&谷歌&馬普所提出TokenFormer，Transformer從來沒有這麼靈活過！

机器之心發表於2024-11-14
谷歌ORM
圖片
AIxiv專欄是機器之心釋出學術、技術內容的欄目。過去數年，機器之心AIxiv專欄接收報導了2000多篇內容，覆蓋全球各大高校與企業的頂級實驗室，有效促進了學術交流與傳播。如果您有優秀的工作想要分享，歡迎投稿或者聯絡報導。投稿郵箱：liyazhou@jiqizhixin.com；zhaoyunfeng@jiqizhixin.com

本論文第一作者是汪海洋，北京大學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 社群開發更透明、易理解的模型貢獻力量。

相關文章