Cross-Architecture Transfer Learning for Linear-Cost Inference Transformers
This reduces computational costs for researchers and practitioners adopting efficient transformer architectures, though it is incremental as it builds on existing transfer learning ideas.
The paper tackles the high cost of pretraining new linear-cost inference transformer architectures from scratch by proposing Cross-Architecture Transfer Learning (XATL), which transfers weights from pre-trained models to reduce training time by up to 2.5x and improve model performance by up to 2.6% on language modeling benchmarks.
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.