1. 導(dǎo)言

Transformer架構(gòu)已經(jīng)成為當(dāng)今大模型的基石，不管是NLP還是CV領(lǐng)域，目前的SOTA模型基本都是基于Transformer架構(gòu)的，比如NLP中目前的各種知名大模型，或者CV中的Vit等模型

本次介紹的論文標(biāo)題為：Tokenformer: Rethinking Transformer Scaling with Tokenized Model Parameters,” 顧名思義，本文提出了Tokenformer架構(gòu)，其優(yōu)勢(shì)在于增量學(xué)習(xí)能力：在增加模型尺寸時(shí)，無(wú)需從頭開(kāi)始重新訓(xùn)練模型，大大降低了成本。 本文代碼已開(kāi)源。

2. Transformer vs Tokenformer - 結(jié)構(gòu)比較

首先我們從頂層設(shè)計(jì)的角度，對(duì)于傳統(tǒng) Transformer 架構(gòu)和 本文提出的 Tokenformer 架構(gòu)進(jìn)行比較，如下圖所示：

2.1 Transformer 架構(gòu)

自注意力機(jī)制是Transformer的核心，主要包括以下幾個(gè)步驟：

如上圖所示，一個(gè)Transformer層主要由兩個(gè)部分組成：

1. 多頭自注意力機(jī)制（Multi-Head Self-Attention） ：輸入首先經(jīng)過(guò)一個(gè)線性投影模塊，以計(jì)算注意力模塊的輸入，即矩陣 Q、K 和 V。然后利用子注意力機(jī)制計(jì)算出Token之間的權(quán)重
2. 前饋神經(jīng)網(wǎng)絡(luò)（Feed-Forward Network, FFN） ：對(duì)于注意力層的輸出進(jìn)行投影，計(jì)算出下一層的輸入

2.2 Transformer 架構(gòu)的缺陷

傳統(tǒng)Transformer在處理token與參數(shù)的交互時(shí)，依賴于固定數(shù)量的線性投影，這限制了模型的擴(kuò)展性，這句話本身較難理解，因此接下來(lái)詳細(xì)論述架構(gòu)的缺陷。

2.2.1 模型的拓展性是什么

模型的拓展性（Scalability）指的是模型在需要更強(qiáng)大性能時(shí)，能夠有效地增加其規(guī)模（如參數(shù)數(shù)量、計(jì)算能力等）而不導(dǎo)致性能下降或計(jì)算成本過(guò)高的能力。

簡(jiǎn)而言之，拓展性好的模型可以在保持或提升性能的同時(shí)，靈活且高效地?cái)U(kuò)大其規(guī)模。

2.2.2 為什么說(shuō)傳統(tǒng)Transformer的固定線性投影限制了模型的擴(kuò)展性

3. TokenFormer的解決方案

為了解決模型維度固定導(dǎo)致的模型缺乏拓展性的問(wèn)題，TokenFormer提出了一種創(chuàng)新的方法，通過(guò)將模型參數(shù)視為tokens，并利用注意力機(jī)制來(lái)處理token與參數(shù)之間的交互，從而實(shí)現(xiàn)更高效、更靈活的模型擴(kuò)展。

3.1 模型參數(shù)Token化

參數(shù)Tokens：原本transformer模型的Q、K、V投影層不再是固定的矩陣，而是轉(zhuǎn)化為一組向量（tokens），例如：

3.2. Token-Parameter Attention（Pattention）層

Pattention層是TokenFormer的核心創(chuàng)新，它通過(guò)注意力機(jī)制來(lái)處理token與參數(shù)之間的交互。從而替代原本的Q,K,V，具體過(guò)程如下：

4. 總體結(jié)構(gòu)

為方便閱讀再把圖扔到這：

與傳統(tǒng)transformer結(jié)構(gòu)相同，其總體上也包括兩層：多頭自注意力層和前饋網(wǎng)絡(luò)層。

4.1 多頭自注意力（Single-Head Variant：

4.2 前饋網(wǎng)絡(luò)（Feed-Forward Network, FFN）

這里也可以看到，相對(duì)于Transformer，Tokenformer就是將所有的投影層從固定的全連接網(wǎng)絡(luò)也變成了Pattention層。

4.3 與transformer的比較

下方公式左側(cè)代表傳統(tǒng)Transformer的自注意力機(jī)制，右側(cè)代表tokenformer的自注意力機(jī)制：

從上邊的圖中可以清楚看到，相對(duì)于transformer，本論文只是將投影層與連接層替換成了新的層。

5. 可擴(kuò)展性

之前說(shuō)過(guò)，相對(duì)于transformer，tokenformer主要是解決可拓展性的問(wèn)題，那么假設(shè)我們要增加參數(shù)數(shù)量，或者要增加輸入維度，tokenformer如何進(jìn)行增量學(xué)習(xí)？

這樣，模型的參數(shù)量可以按需擴(kuò)展。

初始化策略：新增的參數(shù)tokens初始化為零，類(lèi)似于LoRA技術(shù)（Low-Rank Adaptation），確保模型能夠在保持原有知識(shí)的基礎(chǔ)上，快速適應(yīng)新的參數(shù)擴(kuò)展。

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

與從零重訓(xùn)練的 Transformer 相比，如上圖所示，Y 軸代表模型性能，X 軸代表訓(xùn)練成本。藍(lán)線代表使用 3000 億個(gè) token 從頭開(kāi)始訓(xùn)練的 Transformer 模型，不同的圓圈大小代表不同的模型大小。

其他線條代表 Tokenformer 模型，不同顏色代表不同的Token數(shù)量。例如，紅線從 1.24 億個(gè)參數(shù)開(kāi)始，擴(kuò)展到 14 億個(gè)參數(shù)，其訓(xùn)練集為從300B token中抽樣出的30B Token。最終版本模型的性能與相同規(guī)模的 Transformer 相當(dāng)，但訓(xùn)練成本卻大大降低。

黃線顯示，使用 60B個(gè)Token來(lái)訓(xùn)練的增量版本在更低的訓(xùn)練成本下，性能已經(jīng)比 Transformer 更優(yōu)。