Balancing Cost and Benefit with Tied-Multi Transformers

Raj Dabre, Raphael Rubino, Atsushi Fujita

Keywords: tied models, encoder-decoder, multi-layer softmaxing, depth prediction, model compression

TL;DR: Training multiple transformers with tied parameters, depth selection, and further compression

Abstract: This paper proposes a novel procedure for training multiple Transformers with tied parameters which compresses multiple models into one enabling the dynamic choice of the number of encoder and decoder layers during decoding. In sequence-to-sequence modeling, typically, the output of the last layer of the N-layer encoder is fed to the M-layer decoder, and the output of the last decoder layer is used to compute loss. Instead, our method computes a single loss consisting of NxM losses, where each loss is computed from the output of one of the M decoder layers connected to one of the N encoder layers. A single model trained by our method subsumes multiple models with different number of encoder and decoder layers, and can be used for decoding with fewer than the maximum number of encoder and decoder layers. We then propose a mechanism to choose a priori the number of encoder and decoder layers for faster decoding, and also explore recurrent stacking of layers and knowledge distillation to enable further parameter reduction. In a case study of neural machine translation, we present a cost-benefit analysis of the proposed approaches and empirically show that they greatly reduce decoding costs while preserving translation quality.

Community Implementations: [ 3 code implementations](https://www.catalyzex.com/paper/arxiv:2002.08614/code)