在人工智能领域,Transformer模型因其卓越的性能而成为各种基础模型的主导架构。然而,随着模型规模的不断扩大,其计算成本和训练难度也急剧增加。为了解决这一问题,来自北京大学、谷歌和马普所的研究人员提出了一种名为TokenFormer的新型架构,它通过将模型参数也视为Token,实现了Transformer的灵活扩展。
Transformer模型在处理各种任务时表现出色,但当需要扩展模型规模时,问题就出现了。传统上,扩展Transformer模型规模需要改变其核心架构组件,如通道维度,这通常意味着需要从头开始重新训练整个模型。随着模型规模的增长,这种策略变得越来越不可持续,因为计算成本和训练时间都会急剧增加。
TokenFormer的创新之处在于它利用了注意力机制,不仅用于输入Token之间的计算,还用于Token与模型参数之间的交互。通过将模型参数也视为Token,TokenFormer用一种称为Token-Parameter注意力(Pattention)的新型层取代了Transformer中的所有线性投影层。
在TokenFormer中,输入Token作为查询,模型参数作为键和值。这种重新表述允许模型参数以一种渐进且高效的方式进行扩展,而无需从头开始重新训练。通过不断添加新的键值参数对,TokenFormer能够从124M参数扩展到1.4B参数,同时保持与从头开始训练的Transformer相当的性能,但训练成本大大降低。
TokenFormer的架构设计基于注意力机制,包括一个全新的Token-Parameter注意力(Pattention)层。该层使用一组可学习的Token来表示模型参数,并允许输入Token对它们进行注意力计算。当模型规模扩大时,TokenFormer通过添加新的可学习Token来扩展现有的键值参数集,同时保持特征维度不变,从而最小化对其他计算的影响。
TokenFormer的架构设计还考虑了训练的稳定性和效率。例如,它使用一种修改后的softmax操作来计算Pattention分数,以确保梯度的稳定性。此外,它还使用非参数化的层归一化来确保模型参数的可扩展性。
为了评估TokenFormer的性能,研究人员进行了广泛的实验,包括语言模型和视觉模型的训练和评估。在语言模型方面,TokenFormer在各种零样本下游任务上表现出与标准Transformer相当的性能,同时训练成本大大降低。在视觉模型方面,TokenFormer在ImageNet-1K图像分类任务上也表现出与标准Vision Transformer相当的性能。
此外,研究人员还比较了TokenFormer与其他模型扩展方法的性能,如Net2Net。实验结果表明,TokenFormer在模型扩展方面具有更好的性能和效率,特别是在处理长文本序列时。
TokenFormer的提出为Transformer模型的扩展提供了一种创新的解决方案。通过将模型参数也视为Token,并利用注意力机制进行计算,TokenFormer实现了Transformer的灵活扩展,同时大大降低了训练成本。实验结果表明,TokenFormer在各种任务上表现出与标准Transformer相当的性能,同时具有更好的扩展性和效率。
然而,TokenFormer也存在一些挑战和限制。例如,它的设计和实现相对复杂,可能需要更多的计算资源和时间来训练和优化。此外,TokenFormer的扩展性主要体现在参数维度上,而对其他维度的扩展可能需要进一步的研究和探索。