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.
翻译:近期,多种架构通过改变自注意力块设计以实现线性成本推理(LCI),从而提升Transformer语言模型的效率。该领域一种显著方法是状态空间机(SSMs)架构,其在语言建模任务上表现出与自注意力Transformer相当的性能。然而,此类架构变更要求完全从头预训练权重,给希望使用新架构的研究者和从业者带来了巨大成本。在更传统的线性注意力工作中,研究者提出了通过"交换-微调"框架用线性注意力近似全注意力的方法。受此启发,我们提出跨架构迁移学习(XATL),其中LCI与基于自注意力的Transformer之间的共享组件(如层归一化、MLP、输入/输出嵌入)的权重,可直接从已预训练的模型参数迁移至新架构。我们通过不同规模和替代注意力架构的实验验证了该方法的效果,结果表明在相同计算预算下,\methodabbr可将训练时间显著缩短至原来的2.5倍,并在语言模型基准测试中收敛至更优最小值,模型性能提升高达2.6%。