一个已经训练到 70B 参数的模型需要扩展容量时,工程团队通常不能像给数据库加磁盘那样,直接把一批新参数挂到模型后面。增加隐藏维度会改变大量权重矩阵的形状;增加层数会改变优化路径和激活分布;新增模块还要决定怎样初始化、怎样与旧能力对齐。即使旧模型已经消耗了大量训练资源,扩容项目仍可能接近一次新的预训练。
TokenFormer 试图改变的正是这种“模型结构与参数规模一次性绑定”的关系。它把一部分原本组织在固定矩阵中的可学习容量,改写为一组可以被输入表示查询的 Parameter Tokens,并通过 Token-Parameter Attention,也就是 Pattention,完成交互。这样,扩展模型容量可以表现为追加新的参数 Token,而不必先把整个网络宽度重新定义。
固定矩阵为什么让扩容成为结构性改造
标准 Transformer 的主要可学习容量分布在线性投影、Attention 投影和前馈网络等矩阵中。输入向量进入某一层后,与固定形状的权重矩阵相乘,得到新的隐藏表示。矩阵中的元素可以在训练中更新,但矩阵的输入维度、输出维度和它在计算图中的位置通常预先确定。
这意味着“参数更多”并不是简单增加文件大小。假设某层从 4096 维扩展到 8192 维,与它相连的投影矩阵、归一化参数、残差路径和并行切分方案都可能变化。旧检查点中的矩阵无法直接完整填充新形状,新增区域如何初始化也会影响训练稳定性。即使采用 Net2Net、层复制或知识迁移等办法,工程上仍需处理新旧模型之间的结构映射。
扩大深度也存在类似问题。新增 Transformer Block 可以复制邻近层参数,但复制后的层是否需要缩放、哪些参数冻结、学习率如何设置以及旧能力是否退化,都需要额外验证。传统参数表示因此具有很强的结构耦合:参数数量和网络拓扑通常要一起设计。
TokenFormer 的出发点不是否定矩阵乘法,而是提出一个问题:如果一部分模型容量不再固定编码在矩阵形状里,而是表现为数量可变的参数对象,扩容是否可以从“修改网络结构”变成“扩展参数集合”?
Parameter Token 不是自然语言 Token,也不是单个权重
“把参数变成 Token”很容易被误解成把每个标量权重都转换成一个词元。TokenFormer 中的 Parameter Token 仍然是可学习的连续向量,不对应“猫”“数据库”或某个词义,也不等于矩阵中的一个独立数字。它更接近一组被统一组织、可以作为 Key/Value 参与计算的参数向量。
输入侧的数据 Token 经过模型得到查询表示;参数侧维护可学习的 Key/Value Parameter Tokens。Pattention 根据当前输入与参数 Key 的匹配关系,对参数 Value 做加权聚合,再产生输出表示。输入不同,访问参数集合时形成的权重也会不同,因此参数不再只通过一次固定线性映射作用于所有输入。
flowchart LR
A[输入 Token] --> B[构造查询表示]
C[Parameter Key Tokens] --> D[计算匹配分数]
B --> D
D --> E[归一化权重]
F[Parameter Value Tokens] --> G[加权聚合]
E --> G
G --> H[输出隐藏表示]
H --> I[进入下一模块]
这条路径与普通 Token-Token Attention 在形式上相似,但参与匹配的对象不同。普通 Self-Attention 让输入 Token 查询同一序列中的其他 Token;Pattention 让输入表示查询持久化的可学习参数集合。Parameter Tokens 属于模型检查点,是跨请求共享的模型参数,不是每个请求临时构造的上下文缓存。
Parameter Token 也不自动获得可解释语义。某个 Token 被频繁访问,并不表示它对应一个人类可命名的知识模块。多个 Parameter Tokens 可能共同表示某种变换,同一个 Token 也可能在不同输入下承担不同作用。将它们称为“参数记忆单元”可以帮助理解接口,但不能据此推断其内部已经形成离散知识槽位。
Pattention 与 Linear Layer 的计算边界
线性层对输入执行固定矩阵变换。对于给定权重,输入到输出的映射由矩阵元素共同决定,GPU 可以把它组织成高度成熟的大规模 GEMM。编译器、Tensor Core、量化工具和张量并行框架都长期围绕这种规则形状优化。
Pattention 则增加了输入相关的参数选择。每个输入表示需要与 Parameter Key Tokens 计算匹配分数,再根据权重聚合 Parameter Value Tokens。它提供了更灵活的参数访问形式,但也引入了 Attention Score、归一化、中间张量和数据布局问题。即使参数 Token 在所有请求间共享,当前输入对应的匹配权重通常仍要动态计算,不能像静态权重那样简单缓存最终输出。
| 维度 | 固定 Linear Layer | Token-Parameter Attention |
|---|---|---|
| 参数组织 | 固定形状权重矩阵 | 数量可扩展的 Key/Value 参数 Token |
| 输入相关性 | 所有输入经过同一矩阵结构 | 不同输入可形成不同参数聚合权重 |
| 扩容方式 | 常需改变宽度、深度或矩阵形状 | 可通过追加 Parameter Tokens 增加容量 |
| GPU 生态 | GEMM、量化与并行实现成熟 | 需要高效 Attention 与参数 Token 布局支持 |
| 主要风险 | 扩容时结构迁移成本高 | 参数访问成本、利用率和训练稳定性仍需验证 |
因此,TokenFormer 的价值不能只用“Attention 比 Linear 更灵活”概括。它用动态参数访问换取可扩展的参数组织,而这种交换是否划算,取决于训练复用节省的成本能否覆盖 Pattention 带来的计算与系统复杂度。
渐进扩容如何从修改形状变成追加 Token
TokenFormer 最具辨识度的应用是 Progressive Scaling。假设第一阶段模型拥有一组 Parameter Tokens,后续需要增加模型容量时,可以在同一参数集合中追加新的 Key/Value Tokens。已有 Token 和已有网络路径得到保留,新容量主要由新增 Token 提供。
论文给出了从 124M、354M、757M 到 1.4B 参数的渐进扩展实验。这些结果说明,在论文配置下,参数 Token 化可以支持多阶段增加容量,并复用前一阶段已经学习到的参数,而不必每次重新初始化一个完整的大模型。这里的重点是训练过程复用,不是简单把四个独立模型拼接起来。
一次扩容大致包含以下状态变化:
已有 Parameter Tokens
+
新增 Key/Value Parameter Tokens
↓
更新模型检查点与分布式切分
↓
继续训练新增容量及允许更新的参数
↓
验证旧能力、扩容收益和参数利用率
“旧参数保持不变”也需要结合具体训练策略理解。架构允许保留已有 Parameter Tokens,但工程团队仍要决定旧 Token 是否完全冻结、部分解冻还是与新增 Token 一起训练。完全冻结有利于控制旧能力漂移,却可能限制新旧参数协同;全部更新可能提高适应性,但会重新引入灾难性遗忘和优化不稳定风险。
追加参数 Token 还会改变 Pattention 的竞争关系。新增 Token 初始时可能几乎得不到权重,也可能因为初始化尺度异常抢占大量注意力。扩容成功不仅要求训练损失继续下降,还要确认新增容量真正被使用,而不是长期处于“存在于检查点但几乎不参与输出”的状态。
它与 MoE、LoRA 和增宽模型不是同一种扩展
TokenFormer 经常与 MoE 或参数高效微调放在一起讨论,但它们修改的是不同层面。
MoE 仍然把参数组织在多个专家网络中,Router 决定当前 Token 激活哪些专家。它的核心是稀疏使用大量参数,而不是改变参数的基本表示方式。新增专家可以扩展容量,但专家内部仍通常由固定矩阵构成,并且要处理负载均衡、专家并行和 All-to-All 通信。
LoRA 在既有权重旁添加低秩增量,目标通常是用较小的可训练参数适配任务。它保留基础矩阵,把更新限制在低秩分支中。LoRA 可以降低微调成本,却不等同于构建一个可以长期追加参数 Token 的通用扩容机制。
传统增宽或加深直接改变网络拓扑。它可以充分利用成熟的 Dense Transformer Kernel,但每次扩展都要解决参数映射和结构兼容。TokenFormer 则试图让容量变化发生在参数 Token 数量上,从而减少拓扑变化。
| 方法 | 扩展对象 | 旧参数复用方式 | 运行时主要代价 | 典型目标 |
|---|---|---|---|---|
| 增宽或加深 | 层数、隐藏维度、矩阵形状 | 复制、插值或迁移 | 更大的 Dense 计算 | 构建新的更大模型 |
| MoE | 专家数量 | 保留旧专家并增加新专家 | 路由、分桶和跨卡通信 | 稀疏扩大总参数容量 |
| LoRA | 低秩增量参数 | 冻结基础模型 | 额外低秩分支 | 低成本任务适配 |
| TokenFormer | Parameter Token 数量 | 保留并追加参数 Token | Pattention 与参数访问 | 渐进扩容和参数组织重构 |
这些方案也并非必然互斥。理论上,参数 Token 化可以与稀疏路由、低秩更新或混合架构结合,但每叠加一层动态机制,训练稳定性、Kernel 数量、并行切分和检查点管理都会更复杂。是否组合应由清晰的成本模型驱动,而不是因为概念上能够统一就全部加入。
训练复用收益不等于推理速度收益
TokenFormer 的主要动机是降低模型逐级扩容的训练浪费。论文实验表明,同等参数规模下,渐进扩展模型可以获得有竞争力的效果,并减少反复从头训练的需求。这个结论不能直接推导为“TokenFormer 推理一定比 Transformer 快”。
标准 Transformer 的 Dense Linear 与 MLP 已经拥有高度融合的 Kernel。Pattention 虽然结构规则,但需要读取 Parameter Tokens、计算匹配分数和聚合输出。Parameter Token 数量增加后,单个输入 Token 需要访问的参数集合也可能增长。如果没有稀疏化、分块或专用 Kernel,扩容带来的参数数量增长会转化为更高的计算和显存带宽压力。
在线服务还要考虑 Batch。Dense 模型可以让许多请求在同一层执行相同形状的矩阵乘法;Pattention 同样可以批量化,但参数 Token 数、序列长度和并行切分会共同决定张量形状。若模型阶段性扩容后导致 Kernel 形状频繁变化,推理引擎需要重新选择实现、重新编译图或调整并行计划。
模型权重在请求之间共享,并不代表 Parameter Tokens 能像 Prefix Cache 一样跳过计算。Prefix Cache 复用的是某段相同输入已经产生的 KV 状态;Parameter Tokens 是所有请求都要查询的模型参数。两者一个减少重复输入计算,一个改变模型内部参数访问,不能用同一套缓存命中率衡量。
扩容训练最容易失败在哪里
第一类失败是新增 Token 利用率过低。模型已经依赖旧参数形成稳定路径后,新 Token 可能难以获得足够梯度。训练损失仍可能缓慢下降,但新增容量没有承担有效功能。此时应观察 Parameter Attention 的权重分布、不同 Token 的累计访问量、梯度范数和被选中程度,而不只看总参数量。
第二类失败是新增 Token 破坏旧能力。如果新参数初始尺度过大,Pattention 权重可能突然偏向新增部分,使旧模型输出分布明显漂移。渐进扩容应在新阶段开始时进行学习率、归一化和初始化敏感性测试,并使用旧阶段验证集监控回归。
第三类失败是容量增加但数据不够。参数组织更灵活并不会创造训练信号。若新增阶段的数据规模、质量或任务覆盖不足,更大的 Parameter Token 集合可能只增加过拟合空间。扩容计划仍需要遵守数据、计算与参数规模之间的基本约束。
第四类失败来自分布式训练。Parameter Tokens 需要在设备间切分或复制,Pattention 的查询和聚合可能产生新的通信模式。新增 Token 后,旧的张量并行分区不一定仍然均衡;检查点格式、优化器状态和恢复流程也必须支持不同阶段的参数集合大小。
第五类失败是优化目标只奖励最终精度,却忽略总训练成本。渐进扩容是否有价值,应比较从小模型到目标规模的累计 Token、GPU 时间、通信量和失败重训成本,而不是只比较最后一个阶段的单次训练曲线。
生产验证需要同时观察训练与推理
评估 TokenFormer 应拆成两条路径。训练路径关注渐进扩容是否真的复用了旧能力;推理路径关注新参数组织是否能在目标硬件上高效执行。只验证其中一条,容易得到片面的结论。
训练侧至少应记录:每个扩容阶段的累计训练计算、达到目标损失所需 Token、旧任务回归、新任务提升、Parameter Token 访问分布、梯度分布和新增参数有效利用率。还要保留同规模从头训练基线,否则无法判断渐进扩容节省了多少成本,又付出了多少质量代价。
推理侧应观察:Prefill 与 Decode 延迟、tokens/s、批量扩展曲线、HBM 读写、Parameter Token 访问耗时、Kernel occupancy、显存峰值以及不同 Parameter Token 数量下的性能变化。如果模型扩容后质量提高但单请求延迟急剧上升,就需要判断业务目标更重视训练复用还是在线成本。
检查点兼容性也是独立指标。一个可持续扩容的模型体系必须能够识别阶段版本、Parameter Token 数量、优化器状态和分布式分片格式。若每次增加 Token 都需要手工迁移检查点或重写 Serving 引擎,架构层面的扩容灵活性会被运维成本抵消。
参数、数据与记忆能否真的统一成 Token
TokenFormer 提供了一个更广的研究视角:数据 Token、Parameter Token 和 Memory Token 都可以表示为向量集合,并通过某种 Attention 接口交互。这个视角有助于把模型容量、上下文和外部记忆放在相似的计算抽象中讨论。
但“统一接口”不等于“统一生命周期”。数据 Token 属于单次输入,随请求产生和消失;Memory Token 可能来自会话状态或外部存储;Parameter Token 属于训练得到的长期模型权重,需要版本控制、分布式保存和优化器更新。它们即使都参与 Attention,也具有不同的一致性、安全性和更新规则。
统一为 Token 还不能消除规模问题。数据 Token 太多会增加上下文计算,Parameter Token 太多会增加参数访问成本,Memory Token 太多则需要检索与淘汰。Attention 提供了交互形式,但系统仍要决定哪些对象参与本轮计算、如何分片、何时更新以及何时删除。
因此,TokenFormer 当前更适合被理解为一种参数表示和扩容机制的探索,而不是已经完成的“Token 世界”通用架构。它最有价值的贡献,是把模型扩容问题从“怎样重新设计更大的矩阵”转化为“怎样扩展并训练一个参数 Token 集合”。
当模型训练越来越昂贵时,这个问题具有直接工程意义。若渐进扩容能够在更大规模、更多任务和成熟推理引擎中继续成立,模型生命周期可能从一次性预训练转向持续增加容量;若 Pattention 的计算成本和生态改造过高,固定矩阵仍会凭借硬件效率占据主流。真正的判断标准不是概念是否统一,而是新增一批 Parameter Tokens 后,累计训练成本、最终质量和线上运行成本能否同时形成可验证的优势。