Multi-split Reversible Transformers Can Enhance Neural Machine Translation (2021)

Abstract: Large-scale transformers have been shown the state-of-the-art on neural machine translation. However, training these increasingly wider and deeper models could be tremendously memory intensive. We reduce the memory burden by employing the idea of reversible networks that a layer’s input can be reconstructed from its output. We design three types of multi-split based reversible transformers. We also devise a corresponding backpropagation algorithm, which does not need to store activations for most layers. Furthermore, we present two fine-tuning techniques: splits shuffle and self ensemble, to boost translation accuracy. Specifically, our best models surpass the vanilla transformer by at least 1.4 BLEU points in three datasets. Our large-scale reversible models achieve 30.0 BLEU in WMT’14 En-De and 43.5 BLEU in WMT’14 En-Fr, beating several very strong baselines with less than half of the training memory.

## introduction transformers @xcite and their variants @xcite @xcite @xcite significantly enhance the performance of neural machine translation (nmt). but this often requires a large size of the hidden layer (e.g., @xcite used a dimension of 65k) or a deeper network by stacking more building blocks (e.g., @xcite used a 60-layer encoder). training large networks could be extremely memory intensive and might even require model parallelization across multiple gpus @xcite . as a result, reducing memory consumption is crucial to train wider and deeper networks efficiently. backpropagation (bp) is commonly used for training modern neural networks. bp needs to store layer activations to calculate the parameter gradients, which severely increases the memory burden. the idea of reversible networks can be a solution. during training, a reversible network layer's input can be reconstructed from its output. bp is run together with the reconstruction process, removing the need to store all layer activations except for the last layer. we extend the hidden dimension splitting approach by @xcite and design three types of reversible transformers, namely, simple reversible transformers (sim-rev), single dependent reversible transformers (sd-rev) and fully dependent reversible transformers (fd-rev). we also devise a corresponding bp algorithm for our reversible models, which is significantly memory efficient compared with the conventional bp. our reversible models rely on partitioning each layer's input and output into multiple equal splits. this multi-split feature inspires us to develop two fine-tuning techniques to further enhance translation accuracy. first, we randomly shuffle the output splits to encourage information sharing. second, we train distinct translation models based on different output splits in the final decoder layer and run model ensemble during inference. these two techniques are applied after model convergence. only a few epochs of fine-tuning are sufficient for boosting the translation performance. also, both techniques do not break the reversibility of our proposed models. we demonstrate that our reversible models can achieve similar or better performance than vanilla transformers do with less memory consumption. specifically, by employing reversible training and the fine-tuning techniques, our best models can surpass vanilla transformers by 1.5 bleu (iwslt'14 de-en), 2.0 bleu (wmt'19 en-lt) and 1.4 bleu (wmt'14 en-de). our large-scale models also beat several very strong nmt models with less than half the training memory on wmt'14 en-de (30.0 bleu) and wmt'14 en-fr (43.5 bleu). ## methodology we introduce reversible transformers in this section. the definition and benefits of layer reversibility are given in section 2.1. section 2.2 shows three types of reversible architectures based on partitioning the layer input along the hidden/embedding dimension. section 2.3 details the backpropagation algorithm we use. finally, in section 2.4, two techniques that can fit into the reversible training framework are introduced for further boosting model performance. ## analysis splits shuffle probability we study the impact of splits shuffle probability p. for all language pairs, we use the best models as mentioned in section 3.3 and section 3.4. the results are summarized in figure 5 . we find out that a medium p value (0.3 ≤ @xmath0.5) yields the largest bleu increase. a small p value is insufficient to enhance model generalization ability since the model has already been optimized for several epochs. meanwhile, a very large p value makes the model highly unstable and hard to converge. ## conclusion we have presented three types of multi-split based reversible transformers which outperform vanilla transformers. during backpropagation, activations for most layers need not be stored in memory because they can be reconstructed. furthermore, we have proposed two fine-tuning techniques, namely, splits shuffle and self ensemble. both techniques are easy to implement, and only a few fine-tuning epochs are sufficient for boosting translation performance. our approach has beaten several strong baselines in two large datasets with fewer model parameters and much less training memory. specifically, we have achieved 30.0 bleu points in wmt'14 en-de and 43.5 bleu points in wmt'14 en-fr. also, one can transform other network structures into their reversible versions by applying our methods. we would explore more computer vision or natural language processing tasks to widen reversible networks' applicability.