[2410.23168] TokenFormer: Rethinking Transformer Scaling with Tokenized Model Parameters

“TokenFormer: Rethinking Transformer Scaling with Tokenized Model Parameters”由Haiyang Wang等人撰写，提出了TokenFormer架构，通过将模型参数视为令牌，利用注意力机制增强架构灵活性，实现高效的增量缩放，减少训练成本。

  1. 研究背景

    • Transformer的优势与局限：Transformer在各领域表现出色，但缩放成本高，因其依赖固定参数的线性投影，调整架构需从头训练，资源消耗大。
    • 相关研究的不足：大规模训练可提升模型能力，但增加模型尺寸成本高；先前模型重用方法会破坏预训练知识分布，影响收敛。

  2. TokenFormer架构

    • 核心创新：引入Token-Parameter Attention（Pattention）层，将模型参数视为令牌，通过交叉注意力管理输入令牌与参数令牌的交互，增强模型扩展性。

    • 架构设计

      • Pattention层：引入两组可学习参数令牌，通过修改的softmax操作计算注意力分数，实现输入与参数的动态交互，支持模型参数的无缝扩展。
      • 整体架构：遵循预归一化Transformer设计，在多头自注意力和前馈层中用Pattention层替代线性投影，将输入数据和模型参数统一为令牌表示，实现基于注意力的统一计算。
      • 架构配置：参考标准Transformer配置，如GPT - 2，确保与预训练生态系统的兼容性。

  3. 渐进式模型缩放

    • 缩放方法：通过向预训练模型添加新的键值参数对来扩展模型，新参数初始化为零，保留预训练知识，加速收敛。此方法可在不改变输入输出维度的情况下集成任意数量参数，提高训练效率。

    • 实验验证

      • 数据集与训练设置：使用OpenWebText Corpus训练，对比从头训练的Transformer，TokenFormer在减少训练预算的同时达到相近性能。例如，从124M参数模型逐步扩展到1.4B参数，TokenFormer仅需额外30B令牌，而Transformer需300B令牌，且TokenFormer在1.4B参数时测试困惑度为11.77，与Transformer的11.63相近。
      • 实验结果分析：TokenFormer在语言和视觉任务基准测试中表现出竞争力，如在零样本下游任务和ImageNet - 1K图像分类任务中，其性能与标准Transformer相当，证明了架构的有效性和表达能力。

  4. 对比标准Transformer

    • 与Net2Net对比：Net2Net通过复制神经元扩展网络宽度，TokenFormer在模型缩放时能更好地保持输出分布，收敛更快，损失更低，在增量缩放性能上表现更优。
    • 计算成本优势：传统Transformer缩放时增加通道维度会提高计算成本，尤其在处理长文本时，token - token交互成本增加。TokenFormer通过解耦token - token交互成本与模型缩放，在参数增加时保持计算成本可控，处理长序列时优势明显。

  5. 消融实验

    • 优化Pattention层的softmax函数：用GeLU函数替代指数非线性变换，并将L1归一化改为L2归一化，提高了梯度稳定性，使模型在ImageNet分类基准上性能提升。
    • 非参数层归一化：将Transformer的层归一化改为非参数变体，去除可训练权重和偏差，实验表明模型性能不受影响。

  6. 未来工作方向

    • 扩展混合专家范式：将TokenFormer视为混合专家框架的极端实例，有望降低token - parameter交互的计算成本，开发更高效的基础模型。
    • 推进参数高效调整：TokenFormer的缩放方法可用于快速适应新任务或数据集，通过添加新参数令牌实现参数高效调整。
    • 集成视觉和语言模型：利用TokenFormer的参数高效调整能力，统一视觉和语言模型的参数，进行视觉 - 语言对齐和指令调整。
    • 设备 - 云协作：TokenFormer可作为设备 - 云协作中云端知识库，实现设备实时处理与云端密集任务的协同。
    • 增强模型可解释性：基于注意力机制，TokenFormer在token - parameter交互中的注意力特性有助于提高模型可解释性。