Cross-Architecture Transfer Learning for Linear-Cost Inference Transformers

📄 arXiv: 2404.02684v1 📥 PDF

作者: Sehyun Choi

分类: cs.CL, cs.AI, cs.LG

发布日期: 2024-04-03

备注: Preprint

💡 一句话要点

提出跨架构迁移学习以降低Transformer推理成本

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 跨架构迁移学习 线性成本推理 Transformer模型 状态空间机器 自注意力机制 模型训练效率 自然语言处理 机器学习

📋 核心要点

  1. 现有的线性注意力方法在架构变更时需要从头预训练，导致高昂的时间和资源成本。
  2. 本文提出跨架构迁移学习（XATL），通过直接转移共享组件的权重来减少预训练成本。
  3. 实验结果显示，XATL在训练时间上减少最多2.5倍，并在相同计算预算下提升模型性能2.6%。

📝 摘要（中文）

近年来，多个架构被提出以提高Transformer语言模型的效率，尤其是通过改变自注意力模块以实现线性成本推理（LCI）。其中，状态空间机器（SSMs）架构在语言建模任务中表现出与自注意力Transformer相当的性能。然而，这种架构变更需要从头开始进行全面预训练，给研究者和从业者带来了巨大的成本。为此，本文提出了跨架构迁移学习（XATL），直接将LCI和自注意力Transformer之间共享组件的权重从已预训练模型中转移到新架构。实验表明，该方法在不同规模和替代注意力架构上显著减少了训练时间，最多可达2.5倍，并在相同计算预算下在语言模型基准上收敛到更好的最小值，提升幅度可达2.6%。

🔬 方法详解

问题定义：本文旨在解决在Transformer架构变更时需要从头进行全面预训练的问题，这一过程耗时且资源消耗巨大。

核心思路：提出跨架构迁移学习（XATL），通过转移共享组件的权重（如层归一化、MLP、输入/输出嵌入）来降低预训练成本，从而加速新架构的训练。

技术框架：XATL方法包括几个主要模块：首先，识别LCI和自注意力Transformer之间的共享组件；其次，从已预训练模型中提取这些组件的权重；最后，将这些权重应用于新架构，以实现快速收敛。

关键创新：XATL的核心创新在于通过权重转移来避免从头预训练的需求，这与传统方法的全量预训练形成鲜明对比，显著降低了训练时间和资源消耗。

关键设计：在实现过程中，关键设计包括选择合适的共享组件进行权重转移，以及在训练过程中调整学习率和损失函数，以确保新架构的有效收敛。具体的参数设置和网络结构设计在实验中进行了详细验证。

🖼️ 关键图片

fig_0
img_1

📊 实验亮点

实验结果表明，XATL方法在不同规模和替代注意力架构上显著提高了训练效率，训练时间最多减少2.5倍，并在相同计算预算下，模型性能提升达2.6%。这些结果展示了XATL在实际应用中的巨大潜力。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理、机器翻译和对话系统等，能够帮助研究者和开发者在资源有限的情况下快速部署高效的Transformer模型。未来，XATL方法可能会推动更多高效架构的开发与应用，降低AI模型的训练成本。

📄 摘要（原文）

Recently, multiple architectures has been proposed to improve the efficiency of the Transformer Language Models through changing the design of the self-attention block to have a linear-cost inference (LCI). A notable approach in this realm is the State-Space Machines (SSMs) architecture, which showed on-par performance on language modeling tasks with the self-attention transformers. However, such an architectural change requires a full pretraining of the weights from scratch, which incurs a huge cost to researchers and practitioners who want to use the new architectures. In the more traditional linear attention works, it has been proposed to approximate full attention with linear attention by swap-and-finetune framework. Motivated by this approach, we propose Cross-Architecture Transfer Learning (XATL), in which the weights of the shared components between LCI and self-attention-based transformers, such as layernorms, MLPs, input/output embeddings, are directly transferred to the new architecture from already pre-trained model parameters. We experimented the efficacy of the method on varying sizes and alternative attention architectures and show that \methodabbr significantly reduces the training time up to 2.5x times and converges to a better minimum with up to 2.6% stronger model on the LM benchmarks within the same compute budget.