中文

Transformer 作为学习基因的线性扩展

人工智能 2023-12-21 v2 机器学习

摘要

我们提出扩展共享的 Transformer 模块以生成并初始化不同深度的 Transformer,从而适应多样的资源约束。类比基因的可扩展性,我们将此类模块称为学习基因。为了确定扩展机制,我们深入研究了层位置与其对应权重值之间的关系,发现线性函数能很好地近似这一关系。基于此发现,我们提出了作为学习基因线性扩展的 Transformer,这是一种灵活生成并初始化不同深度 Transformer 的新方法。具体而言,为了学习学习基因,我们首先构建一个由学习基因线性扩展而成的辅助 Transformer,随后采用软蒸馏对其进行训练。接着,我们可以通过线性扩展训练好的学习基因来生成并初始化不同深度的 Transformer,从而支持多样的下游场景。在 ImageNet-1K 上的大量实验表明,与许多从零开始训练的独立模型相比,TLEG 取得了相当或更好的性能,同时将训练成本降低了约 2 倍。在迁移到多个下游分类数据集时,TLEG 大幅超越了现有的初始化方法(例如,在 iNat 2019 上提升 +6.87%,在 CIFAR-100 上提升 +7.66%)。在需要生成不同深度模型以适应不同资源约束的情况下,与预训练加微调方法相比,TLEG 取得了相当的结果,同时将用于初始化这些模型的存储参数量减少了约 19 倍,预训练成本减少了约 5 倍。在迁移固定参数集以初始化不同模型时,与预训练方法相比,TLEG 展现出更好的灵活性和有竞争力的性能,同时将用于初始化的存储参数量减少了约 2.9 倍。

关键词

引用

@article{arxiv.2312.05614,
  title  = {Transformer as Linear Expansion of Learngene},
  author = {Shiyu Xia and Miaosen Zhang and Xu Yang and Ruiming Chen and Haokun Chen and Xin Geng},
  journal= {arXiv preprint arXiv:2312.05614},
  year   = {2023}
}