Memory-Efficient Differentiable Transformer Architecture Search

Findings of the Association for Computational Linguistics: ACL-IJCNLP 2021, pages 4254–4264 August 1–6, 2021. ©2021 Association for Computational Linguistics 4254 Memory-Efficient Differentiable Transformer Architecture Search Yuekai Zhao†, Li Dong‡, Yelong Shen‡, Zhihua Zhang§, Furu Wei‡, Weizhu Chen‡ †Academy for Advanced Interdisciplinary Studies, Peking University §School of Mathematical Sciences, Peking University ‡Microsoft Corporation {yuekaizhao@,zhzhang@math.}pku.edu.cn {lidong1,yeshe,fuwei,wzchen}@microsoft.com Abstract Differentiable architecture search (DARTS) is successfully applied in many vision tasks. However, directly using DARTS for Trans- formers is memory-intensive, which renders the search process infeasible. To this end, we propose a multi-split reversible network and combine it with DARTS. Specifically, we de- vise a backpropagation-with-reconstruction al- gorithm so that we only need to store the last layer’s outputs. By relieving the memory bur- den for DARTS, it allows us to search with larger hidden size and more candidate opera- tions. We evaluate the searched architecture on three sequence-to-sequence datasets, i.e., WMT’14 English-German, WMT’14 English- French, and WMT’14 English-Czech. Exper- imental results show that our network con- sistently outperforms standard Transformers across the tasks. Moreover, our method com- pares favorably with big-size Evolved Trans- formers, reducing search computation by an order of magnitude. 1 Introduction Current neural architecture search (NAS) studies have produced models that surpass the performance of those designed by humans (Real et al., 2019; Lu et al., 2020). For sequence tasks, efforts are made in reinforcement learning-based (Pham et al., 2018) and evolution-based (So et al., 2019; Wang et al., 2020) methods, which suffer from the huge compu- tational cost. Instead, gradient-based methods (Liu et al., 2018; Jiang et al., 2019; Yang et al., 2020) are less demanding in computing resources and easy to implement, attracting many attentions recently. The idea of gradient-based NAS is to train a super network covering all candidate operations. Different sub-graphs of the super network form the search space. To find a well-performing sub- graph, Liu et al. (2018) (DARTS) introduced search parameters jointly optimized with the network 300 400 500 600 700 800 900 Hidden size 6000 8000 10000 12000 14000 16000 Memory consumption in MB DARTS in Transformers Out of memory Ours Figure 1: Memory comparison of using our reversible networks and Transformers as the backbone model of DARTS. Experiments are run on a single step of forward-backward pass on a batch of 3584 tokens with a NVIDIA P100 GPU. Limited by GPU memory, DARTS in Transformers has to search in small sizes while evaluating in large sizes, which will cause perfor- mance gaps (Chen et al., 2019). weights. Operations corresponding to the largest search parameters are kept for each intermediate node after searching. A limitation of DARTS is its memory inefficiency because it needs to store the intermediate outputs from all its candidate op- erations. This is much more pronounced when we apply Transformers (Vaswani et al., 2017) as the backbone of DARTS (the operation set is detailed in Section 2.5). As shown in Figure 1, memory consumption grows extremely fast as we increase the hidden size d, quickly running out of memory as …

Content truncated to 500 words