Neural Parameter Allocation Search

1 INTRODUCTION . Training neural networks requires ever more computational resources , with GPU memory often being a significant limitation ( Rajbhandari et al. , 2021 ) . Methods such as checkpointing ( e.g. , Chen et al. , 2016 ; Gomez et al. , 2017 ; Jain et al. , 2020 ) and out-of-core algorithms ( e.g. , Ren et al. , 2021 ) have been developed to reduce memory from activations and improve training efficiency . Yet even with such techniques , Rajbhandari et al . ( 2021 ) find that to train efficiently , model parameters require significantly greater memory bandwidth than activations , indicating parameters are a key limit on future growth . One solution to this problem is cross-layer parameter sharing , which reduces the memory needed to store parameters . This can also reduce the cost of communicating model updates in distributed training ( Lan et al. , 2020 ; Jaegle et al. , 2021 ) and federated learning ( Konečný et al. , 2016 ; McMahan et al. , 2017 ) , as the model is smaller , and can help avoid overfitting ( Jaegle et al. , 2021 ) . However , prior work in cross-layer parameter sharing ( e.g. , Dehghani et al. , 2019 ; Savarese & Maire , 2019 ; Lan et al. , 2020 ; Jaegle et al. , 2021 ) has two significant limitations . First , they rely on suboptimal hand-crafted techniques for deciding where and how to share parameters . Second , they rely on models having many identical layers . This limits the network architectures they apply to ( e.g. , DenseNets ( Huang et al. , 2017 ) have few such layers ) and their parameter savings is only proportional to the number of identical layers . To move beyond these limits , we introduce Neural Parameter Allocation Search ( NPAS ) , a novel task which generalizes existing parameter sharing approaches . In NPAS , the goal is to identify where and how to distribute parameters in a neural network to produce a high-performing model using an arbitrary , fixed parameter budget and no architectural assumptions . Searching for good sharing strategies is challenging in many neural networks due to different layers requiring different numbers of parameters or weight dimensionalities , multiple layer types ( e.g. , convolutional , fully-connected , recurrent ) , and/or multiple modalities ( e.g. , text and images ) . Hand-crafted sharing approaches , as in prior work , can be seen as one implementation of NPAS , but they can be complicated to create for complex networks and have no guarantee that the sharing strategy is good . Trying all possible permutations of sharing across layers is computationally infeasible even for small networks . To our knowledge , we are the first to consider automatically searching for good parameter sharing strategies . By supporting arbitrary parameter budgets , NPAS explores two novel regimes . First , while prior work considered using sharing to reduce the number of parameters ( which we refer to as low-budget NPAS , LB-NPAS ) , we can also increase the number of parameters beyond what an architecture typically uses ( high-budget NPAS , HB-NPAS ) . HB-NPAS can be thought of as adding capacity to the network in order to improve its performance without changing its architecture ( e.g. , without increasing the number of channels that would also increase computational time ) . Second , we consider cases where there are fewer parameters available to a layer than needed to implement the layer ’ s operations . For such low-budget cases , we investigate parameter upsampling methods to generate the layer ’ s weights . A vast array of other techniques , including pruning ( Hoefler et al. , 2021 ) , quantization ( Gholami et al. , 2021 ) , knowledge distillation ( Gou et al. , 2021 ) , and low-rank approximations ( e.g. , Wu , 2019 ; Phan et al. , 2020 ) are used to reduce memory and/or FLOP requirements for a model . However , such methods typically only apply at test/inference time , and actually are more expensive to train due to requiring a fully-trained large network , in contrast to NPAS . Nevertheless , these are also orthogonal to NPAS and can be applied jointly . Indeed , we show that NPAS can be combined with pruning or distillation to produce improved networks . Figure 1 compares NPAS to closely related tasks . To implement NPAS , we propose Shapeshifter Networks ( SSNs ) , which can morph a given parameter budget to fit any architecture by learning where and how to share parameters . SSNs begin by learning which layers can effectively share parameters using a short preliminary training step . In this , all layers are implemented from a single shared set of parameters , and a representation is learned for each layer based on how it uses the parameters . Layers that use parameters in a similar way are then good candidates for sharing during the main training step . When training , SSNs generate weights for each layer by down- or upsampling the associated parameters as needed . We demonstrate SSN ’ s effectiveness by using them to implement high- and low-budget NPAS on a variety of networks , including vision , text , and multi-modal vision/language tasks . E.g. , a LB-NPAS SSN implements a WRN-50-2 ( Zagoruyko & Komodakis , 2016 ) using 19M parameters ( 69M in the original ) and achieves an Error @ 5 on ImageNet classification ( Deng et al. , 2009 ) 3 % lower than a WRN with the same budget . Similarity , we achieve a 1 % boost to SQuAD v2.0 ( Rajpurkar et al. , 2016 ) with 18M parameters ( 334M in the original ) over ALBERT ( Lan et al. , 2020 ) , prior work for parameter sharing in Transformers ( Vaswani et al. , 2017 ) . For HB-NPAS , we achieve a 1–1.5 % improvement in Error @ 1 on CIFAR ( Krizhevsky , 2009 ) by adding capacity compared to a traditional network . In summary , our key contributions are : • We introduce Neural Parameter Allocation Search ( NPAS ) , a novel task in which the goal is to implement a given network architecture using any parameter budget . • To solve NPAS , we propose Shapeshifter Networks ( SSNs ) , which automate parameter sharing . To our knowledge , SSNs are the first approach to automatically learn where and how to share parameters and to share parameters between layers of different sizes or types . • We benchmark SSNs for LB- and HB-NPAS and show they create high-performing networks when either using few parameters or adding network capacity . • We also show that SSNs can be combined with knowledge distillation and parameter pruning to boost performance over such methods alone . 2 NEURAL PARAMETER ALLOCATION SEARCH ( NPAS ) . In NPAS , the goal is to implement a neural network given a fixed parameter budget . More formally : Neural Parameter Architecture Search ( NPAS ) : Given a neural network architecture with layers ` 1 , . . . , ` L , which each require weights w1 , . . . , wL , and a fixed parameter budget θ , train a high-performing neural network using the given architecture and parameter budget . Any general solution to NPAS ( i.e. , that works for arbitrary θ or network ) must solve two subtasks : 1 . Parameter mapping : Assign to each layer ` i a subset of the available parameters . 2 . Weight generation : Generate ` i ’ s weights wi from its assigned parameters , which may be any size . Prior work , such as Savarese & Maire ( 2019 ) and Ha et al . ( 2016 ) , are examples of weight generation methods , but in limited cases , e.g. , Savarese & Maire ( 2019 ) does not support there being fewer parameters than weights . To our knowledge , no prior work has automated parameter mapping , instead relying on hand-crafted heuristics that do not generalize to many architectures . Note weight generation must be differentiable so gradients can be backpropagated to the underlying parameters . NPAS naturally decomposes into two different regimes based on the parameter budget relative to what would be required by a traditional neural network ( i.e. , ∑L i |wi| versus |θ| ) : • Low-budget ( LB-NPAS ) , with fewer parameters than standard networks ( ∑L i |wi| < |θ| ) . This regime has traditionally been the goal of cross-layer parameter sharing , and reduces memory at training and test time , and consequentially reduces communication for distributed training . • High-budget ( HB-NPAS ) , with more parameters than standard networks ( ∑L i |wi| > |θ| ) . This is , to our knowledge , a novel regime , and can be thought of as adding capacity to a network without changing the underlying architecture by allowing a layer to access more parameters . Note , in both cases , the FLOPs required of the network do not significantly increase . Thus , HB-NPAS can significantly reduce FLOP overhead compared to larger networks . The closest work to ours are Shared WideResNets ( SWRN ) ( Savarese & Maire , 2019 ) , Hypernetworks ( HN ) ( Ha et al. , 2016 ) , and Lookup-based Convolutional Networks ( LCNN ) ( Bagherinezhad et al. , 2017 ) . Each method demonstrated improved low-budget performance , with LCNN and SWRN focused on improving sharing across layers and HN learning to directly generate parameters . However , all require adaptation for new networks and make architectural assumptions . E.g. , LCNN was designed specifically for convolutional networks , while HN and SWRN ’ s benefits are proportional to the number of identical layers ( see Figure 3 ) . Thus , each method supports limited architectures and parameter budgets , making them unsuited for NPAS . LCNN and HN also both come with significant computational overhead . E.g. , the CNN used by Ha et al . requires 26.7M FLOPs for a forward pass on a 32×32 image , but weight generation with HN requires an additional 108.5M FLOPs ( 135.2M total ) . In contrast , our SSNs require 0.8M extra FLOPs ( 27.5M total , 5× fewer than HN ) . Across networks we consider , SSN overhead for a single image is typically 0.5–2 % of total FLOPs . Note both methods generate weights once per forward pass , amortizing overhead across a batch ( e.g. , SSN overhead is reduced to 0.008–0.03 % for batch size 64 ) . HB-NPAS is also reminiscent of mixture-of-experts ( e.g. , Shazeer et al. , 2017 ) ; both increase capacity without significantly increasing FLOPs , but NPAS allows this overparameterization to be learned without architectural changes required by prior work . NPAS can be thought of as searching for efficient and effective underlying representations for a neural network . Methods have been developed for other tasks that focus on directly searching for more effective architectures ( as opposed to their underlying representations ) . These include neural architecture search ( e.g. , Bashivan et al. , 2019 ; Dong & Yang , 2019 ; Tan et al. , 2019 ; Xiong et al. , 2019 ; Zoph & Le , 2017 ) and modular/self-assembling networks ( e.g. , Alet et al. , 2019 ; Ferran Alet , 2018 ; Devin et al. , 2017 ) . While these tasks create computationally efficient architectures , they do not reduce the number of parameters in a network during training like NPAS ( i.e. , they can not be used to train very large networks or for federated or distributed learning applications ) , and indeed are computationally expensive . NPAS methods can also provide additional flexibility to architecture search by enabling them to train larger and/or deeper architectures while keeping within a fixed parameter budget . In addition , the performance of any architectures these methods create could be improved by leveraging the added capacity from excess parameters when addressing HB-NPAS .

Summary: Parameter sharing can reduce memory footprint of neural networks and memory bandwidth requirements, but existing methods require manually tuning the sharing strategy. This paper uses a small phase of training to cluster the learned layer representations by groups. This allows networks to be scaled from small to large parameter counts (to be clear, number of trainable parameters) without changing the model architecture. Importantly, this procedure does not change the number of FLOPs in the model.   Experiments across a wide set of tasks and networks compare this approach with either (1) SWRN from Savarasese & Maire, 2019), or (2) existing hand-tuned parameter scaling from families of networks such as EfficientNet, DenseNet, or ALBERT.