Dynamic Layer Tying for Parameter-Efficient Transformers

TL;DR

This paper introduces a reinforcement learning-based method for dynamically tying transformer layers to reduce parameters and memory usage, while maintaining or improving performance.

Contribution

It presents a novel dynamic layer tying approach using RL to select and share layers during training, enhancing efficiency and regularization.

Findings

Model outperforms baseline in perplexity

Reduces trainable parameters significantly

Memory consumption during training is greatly decreased

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.