作者 | 汪海洋
单位 | 北京大学博士生
研究方向 | 通用模型的架构设计
新一代通用灵活的网络结构 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 + 的浏览量)。
目前代码、模型和项目主页均已放出：
论文标题：
TokenFormer: Rethinking Transformer Scaling with Tokenized Model Parameters
论文链接：
https://arxiv.org/pdf/2410.23168
代码链接：
Backward Lens: Projecting Language Model Gradients into the Vocabulary Space
开源模型：
https://arxiv.org/abs/2402.12865
背景介绍
得益于其处理各种数据的灵活性，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 进行交互，提供了增量模型扩展所需的灵活性。
应用：天生的增量式 Model Scaling
实验结果
未来研究方向
扫描二维码添加小助手微信
