ElastFormer：通过自蒸馏在 Transformer 中实现冗余降低
机器学习
2024-11-26 v1 人工智能
摘要
我们提出 ElastiFormer，这是一种后训练技术，可将预训练的 Transformer 模型适配为具备可变推理计算能力的弹性模型。ElastFormer 引入小型路由模块（仅增加约 0.00006% 可训练参数），动态选择每个网络层要处理的子集参数和输入标记。路由模块使用自蒸馏损失进行训练，以最小化预训练模型输出与其弹性版本之间的差异。由于 ElastiFormer 不对预训练 Transformer 模型的模态作出假设，可轻松应用于包括因果语言建模、图像建模及视觉-语言建模在内的所有模态。我们表明，对于 Transformer 层的不同组件，可实现 20% 至 50% 的计算节省，这可以通过添加极低秩 LoRA 权重（秩为 1）通过相同的蒸馏目标进一步降低。最后，通过比较在不同 ImageNet 子集上训练的路由，我们表明 ElastiFormer 能抵抗训练领域的变化。
关键词
引用
@article{arxiv.2411.15281,
title = {ElastiFormer: Learned Redundancy Reduction in Transformer via Self-Distillation},
author = {Junzhang Liu and Tingkai Liu and Yueyuan Sui and Stephen Xia},
journal= {arXiv preprint arXiv:2411.15281},
year = {2024}
}