Dynamic Layer Tying for Parameter-Efficient Transformers
Tamir David Hay, Lior Wolf
TL;DR
This work tackles Transformer overparameterization by introducing a reinforcement-learning driven dynamic layer tying mechanism that runs during a single training pass. An RL agent selects, for each layer, whether to train independently or to copy weights from earlier layers, enabling extensive weight sharing and substantial memory savings while preserving or improving perplexity on language modeling tasks. The approach yields replication of a large majority of layers (roughly 75%+), reduces trainable parameters from billions to the hundreds of millions or less, and achieves up to an order of magnitude reduction in training memory, with strong results on GPT-2 and BERT across multiple datasets. Overall, the method demonstrates a highly parameter-efficient form of neural architecture search that adapts topology on-the-fly, with potential extensions to LoRA and downstream tasks.
Abstract
In the pursuit of reducing the number of trainable parameters in deep transformer networks, we employ Reinforcement Learning to dynamically select layers during training and tie them together. Every few iterations, the RL agent is asked whether to train each layer $i$ independently or to copy the weights of a previous layer $j<i$. This facilitates weight sharing, reduces the number of trainable parameters, and also serves as an effective regularization technique. Experimental evaluations validate that our model modestly outperforms the baseline transformer model with regard to perplexity and drastically reduces the number of trainable parameters. In particular, the memory consumption during training is up to one order of magnitude less than the conventional training method.