AI 技术AI 推理 · 15/18
#TokenFormer#Parameter Token#Pattention#渐进扩容#模型架构#增量训练

AI 推理系列(十五):TokenFormer——如果模型参数也是 Token,会发生什么?

围绕大模型每次扩容都要重新训练的成本,解释 TokenFormer 如何用 Token-Parameter Attention 将部分固定矩阵计算改写为可查询的 Parameter Tokens,并分析渐进扩容、参数复用、训练稳定性、推理开销与工程生态边界。

一个已经训练到 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 LayerToken-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低秩增量参数冻结基础模型额外低秩分支低成本任务适配
TokenFormerParameter Token 数量保留并追加参数 TokenPattention 与参数访问渐进扩容和参数组织重构

这些方案也并非必然互斥。理论上,参数 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 后,累计训练成本、最终质量和线上运行成本能否同时形成可验证的优势。

资料来源

  1. TokenFormer: Rethinking Transformer Scaling with Tokenized Model Parameters
  2. TokenFormer ICLR 2025 Conference Paper
  3. TokenFormer Official Implementation