Cross-layer parameter sharing in the Transformer

🚀 Feature Request

Implement cross-layer parameter sharing as in Universal Transformer [1] and ALBERT [2]

Motivation

Transformers implementing cross-layer parameter sharing have been shown to perform well in different tasks in the literature while being more sample and memory-efficient.

Pitch

The Transformer model in transformer.py should have arguments "--encoder-cross-layer-weight-sharing" and "--decoder-cross-layer-weight-sharing". If they are set, then instead of instantiating N (encoder/decoder) layers, only one is instantiated, and the input is recurrently fed to the same layer N times. Step embeddings and transition functions would need to be implemented, as in [1]. In addition, the conditional computation mechanism of [1] could be implemented.

Once it was first introduced, it was a specific feature of the Universal Transformer. However, now it is a general trick that can be used in a variety of tasks. The Fairseq community would benefit from having this feature, which seems a generic building block instead of a very specific method that users should implement on their own.

Alternatives

Let each user of Fairseq implement it, or let each user use another library instead of Fairseq if they need cross-layer weight-sharing.

Additional context

ALBERT is already implemented in Huggingface (https://github.com/huggingface/transformers) and there are other reference implementations in PyTorch.

[1] Universal Transformer: https://arxiv.org/abs/1807.03819 [2] ALBERT: https://arxiv.org/abs/1909.11942