通过 Token 扩展实现 Transformer 通用高效的训练
机器学习
2024-04-02 v1 人工智能
计算与语言
计算机视觉与模式识别
摘要
Vision Transformers (ViTs) 的卓越性能通常需要极大的训练成本。现有方法试图加速 ViTs 的训练，但通常在准确率下降的同时忽视了方法的通用性。同时，它们打破了原始 Transformer 的训练一致性，包括超参数、架构和策略的一致性，这使其难以广泛应用于不同的 Transformer 网络。在本文中，我们提出了一种新颖的 token 增长方案——Token Expansion (简称 ToE)，以实现 ViTs 的一致性训练加速。我们引入了“初始化-扩展-合并”流程，以保持原始 Transformer 中间特征分布的完整性，防止在训练过程中丢失关键的可学习信息。ToE 不仅能无缝集成到 Transformer（如 DeiT 和 LV-ViT）的训练和微调过程中，而且对高效训练框架（如 EfficientTrain）同样有效，且无需扭曲原始训练超参数、架构或引入额外的训练策略。大量实验表明，ToE 以无损方式实现了 ViTs 训练约 1.3 倍的加速，甚至比全 token 训练基线获得了性能提升。代码可在 https://github.com/Osilly/TokenExpansion 获取。
引用
@article{arxiv.2404.00672,
title = {A General and Efficient Training for Transformer via Token Expansion},
author = {Wenxuan Huang and Yunhang Shen and Jiao Xie and Baochang Zhang and Gaoqi He and Ke Li and Xing Sun and Shaohui Lin},
journal= {arXiv preprint arXiv:2404.00672},
year = {2024}
}
备注
Accepted to CVPR 2024. Code is available at https://github.com/Osilly/TokenExpansion