Learning to Grow Pretrained Models for
Efficient Transformer Training

Peihao Wang1	Rameswar Panda2	Lucas Torroba Hennigen4	Philip Greengard3

Leonid Karlinsky2	Rogerio Feris2	David Cox2	Atlas Wang1	Yoon Kim4

1 University of Texas at Austin	2 MIT-IBM Watson AI Lab	3 Columbia University	4 MIT

ICLR 2023

Abstract

Scaling transformers has led to significant breakthroughs in many domains, leading to a paradigm in which larger versions of existing models are trained and released on a periodic basis. New instances of such models are typically trained completely from scratch, despite the fact that they are often just scaled-up versions of their smaller counterparts. How can we use the implicit knowledge in the parameters of smaller, extant models to enable faster training of newer, larger models? This paper describes an approach for accelerating transformer training by {learning to grow} pretrained transformers, where we learn to linearly map the parameters of the smaller model to initialize the larger model. For tractable learning, we factorize the linear transformation as a composition of (linear) width- and depth-growth operators, and further employ a Kronecker factorization of these growth operators to encode architectural knowledge. Extensive experiments across both language and vision transformers demonstrate that our learned Linear Growth Operator (LiGO) can save up to 50% computational cost of training from scratch, while also consistently outperforming strong baselines that also reuse smaller pretrained models to initialize larger models.

---

Qualitative Results

Qualitative examples showing that LiGO can accelerate BERT training time by ~40% with ~45% FLOPs saving.

---

Paper & Code

		Peihao Wang, Rameswar Panda, Lucas Torroba Hennigen, Philip Greengard, Leonid Karlinsky, Rogerio Feris, David Cox, Atlas Wang, Yoon Kim Learning to Grow Pretrained Models for Efficient Transformer Training International Conference on Learning Representations (ICLR), 2023 [PDF] [Code]

---