N EURAL P ARAMETER A LLOCATION S EARCH
A BSTRACT
Training neural networks requires increasing amounts of memory.
Parameter sharing can reduce memory and communication costs, but existing methods assume networks have many identical layers and utilize hand-crafted sharing strategies that fail to generalize.
We introduce Neural Parameter Allocation Search (NPAS), a novel task where the goal is to train a neural network given an arbitrary, ﬁxed parameter budget.
NPAS covers both low-budget regimes, which produce compact networks, as well as a novel high-budget regime, where additional capacity can be added to boost performance without increasing inference FLOPs.
To address
NPAS, we introduce Shapeshifter Networks (SSNs), which automatically learn where and how to share parameters in a network to support any parameter budget without requiring any changes to the architecture or loss function.
NPAS and
SSNs provide a complete framework for addressing generalized parameter sharing, and can also be combined with prior work for additional performance gains.
We demonstrate the effectiveness of our approach using nine network architectures across four diverse tasks, including ImageNet classiﬁcation and transformers.
I NTRODUCTION
Training neural networks requires ever more computational resources, with GPU memory being a signiﬁcant limitation (Rajbhandari et al., 2021).
Method 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 efﬁciency.
Yet even with such techniques, Rajbhandari et al.
ﬁnd that model parameters require signiﬁcantly greater memory bandwidth than activations during training, indicating parameters are a key limit on future growth.
One solution is cross-layer parameter sharing, which reduces the memory needed to store parameters, which can also reduce the cost of communicating model updates in distributed training (Lan et al., 2020; Jaegle et al., 2021) and federated learning (Kone ˇ cn ´ y et al., 2016; McMahan et al., 2017), as the model is smaller, and can help avoid overﬁtting (Jaegle et al., 2021).
However, prior work in parameter sharing ( e.g ., Dehghani et al., 2019; Savarese & Maire, 2019; Lan et al., 2020; Jaegle et al., 2021) has two signiﬁcant limitations.
First, they rely on suboptimal hand-crafted techniques for deciding where and how sharing occurs.
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, ﬁxed 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 ﬁrst to consider automatically searching for good parameter sharing strategies.
[Figure]
Pruning/Quantizing Networks makes training more efficient þ þ ý ý supports any parameter budget (low and high ) þ ý ý ý generalizes to any architecture þ ý þ þ focuses on improving task performance þ ý þ ý
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 (Hoeﬂer 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 ﬁt any architecture by learning where and how to share parameters.
SSNs begin by learning which layers can effectively share parameters using a short pretraining step, where all layers are generated from a single shared set of 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 in high- and low-budget NPAS on a variety of networks, including vision, text, and 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 (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 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 ﬁrst 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.
N EURAL P ARAMETER A LLOCATION S EARCH (NPAS)
In NPAS, the goal is to implement a neural network given a ﬁxed parameter budget.
More formally:
Neural Parameter Allocation Search
NPAS): Given a neural network architecture with layers (cid:96) 1 , . . .
, (cid:96) L , which each require weights w 1 , . . . , w L , and a ﬁxed 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:
Parameter mapping: Assign to each layer (cid:96)
i
a
subset
of the available
parameters.
Weight generation: Generate (cid:96)
i ’s weights w i 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 ., (cid:80) Li | w i | versus | θ | ):
• Low-budget (LB-NPAS), with fewer parameters than standard networks ( (cid:80) Li | w i | < | θ | ). 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 ( (cid:80) Li | w i | > | θ | ).
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 signiﬁcantly increase.
Thus, HB-NPAS can signiﬁcantly 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 speciﬁcally for convolutional networks, while HN and SWRN’s beneﬁts 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 signiﬁcant 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 signiﬁcantly 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 efﬁcient 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 efﬁcient architectures, they do not reduce the number of parameters in a network during training like NPAS ( i.e ., they cannot 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 ﬂexibility to architecture search by enabling them to train larger and/or deeper architectures while keeping within a ﬁxed 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. [Figure]
S HAPESHIFTER N ETWORKS FOR NPAS
We now present Shapeshifter Networks (SSNs), a framework for addressing NPAS using generalized parameter sharing to implement a neural network with an arbitrary, ﬁxed parameter budget.
Figureprovides an overview and example of SSNs, and we detail each aspect below.
An SSN consists of a provided network architecture with layers (cid:96) 1 ,...,L , and a ﬁxed budget of parameters θ , which are partitioned into P parameter groups (both hyperparameters) containing parameters θ 1 ,...,P .
Each layer is associated with a single parameter group, which will provide the parameters used to implement it.
This mapping is learned in a preliminary training step by training a specialized SSN and clustering its layer representations (Section 3.2).
To implement each layer, an SSN morphs the parameters in its associated group to generate the necessary weights; this uses downsampling (Section 3.1.1) when the group has more parameters than needed, or upsampling (Section 3.1.2) when the group has fewer parameters than needed.
SSNs allow any number of parameters to “shapeshift” into a network without necessitating changes to the model’s loss, architecture, or hyperparameters, and the process can be applied automatically.
Finally, we note that SSNs are simply one approach to NPAS.
Appendices B-D contain ablation studies and discussion of variants we found to be less successful.
3. W EIGHT G ENERATION
Weight generation implements a layer (cid:96) i , which requires weights w i , using the ﬁxed set of parameters in its associated parameter group θ j .
(We assume the mapping between layers and parameter groups has already been established; see Section 3.2.)
There are three cases to handle:
| w i | =
| θ j | (exactly enough parameters): The parameters are used as-is.
| w i | < | θ j | (excess parameters): We perform parameter downsampling (Section 3.1.1).
| w i | > | θ j | (insufﬁcient parameters):
We perform parameter upsampling (Section 3.1.2).
We emphasize that, depending on how layers are mapped to parameter groups, both down- and upsampling may be required in an LB- or HB-NPAS model.
3.1. P ARAMETER D OWNSAMPLING
When a parameter group θ j provides more parameters than needed to implement a layer (cid:96) i , we perform template-based downsampling to generate w i .
To do this, we ﬁrst split θ j into up to K (a hyperparameter) templates T 1 ,...,K i , where each template T ki is the same dimension as w i .
If θ j does not evenly divide into templates, we ignore excess parameters.
To avoid uneven sharing of parameters between layers, the templates for each layer are constructed from θ j in a round-robin fashion.
These templates are then combined to produce w i ; if only one template can be produced we instead use it directly.
We present two different methods of learning to combine templates.
To simplify presentation, we will assume there are exactly K templates used.
WAvg (Savarese & Maire, 2019)
This learns a vector
α i ∈
R K which is used to produce a weighted average of the templates: w i = (cid:80) K k
=
α k i T k i .
The α i are initialized orthogonally to the α s of all other layers in the same parameter group.
While efﬁcient, this only implicitly learns similarities between layers.
Empirically, we ﬁnd that different layers often converge to similar α s, limiting sharing.
Emb To address this, we can instead more directly learn a representation of the layer using a layer embedding .
We use a learnable vector φ i ∈ R E , where E is the size of the layer representation; we use E = 24 throughout, as we found it to work well.
A linear layer, which is shared between all layers in the parameter group and parameterized by W j ∈ R K × E and b j ∈ R K , is then used to construct an α i for the layer, which is used as in WAvg.
That is, α i =
W j φ i + b j and w i =
(cid:80)
Kk
=
α ki T ki .
We considered more complex methods ( e.g ., MLPs, nonlinearities), but they did not improve performance.
While both methods require additional parameters, this is quite small in practice.
WAvg requires
K additional parameters per layer.
Emb requires E = 24 additional parameters per layer and
KE + K = 24 K + K parameters per parameter group.
3.1. P ARAMETER U PSAMPLING
If instead a parameter group θ j provides fewer parameters than needed to implement a layer (cid:96) i , we upsample θ j to be the same size as w i .
As a layer will use all of the parameters in θ j , we do not use templates.
We consider two methods for upsampling below.
Inter As a na ¨ ıve baseline, we use bilinear interpolation to directly upsample θ j .
However, this could alter the patterns captured by parameters, as it effectively stretches the receptive ﬁeld.
In practice, we found fully-connected and recurrent layers could compensate for this warping, but it degraded convolutional layers compared to simpler approaches such as tiling θ j .
Mask To address this, and avoid redundancies created by directly repeating parameters, we propose instead to use a learned mask to modify repeated parameters.
For this, we ﬁrst use n = (cid:100)| w i | / | θ j |(cid:101) tiles of θ j to be the same size as w i (discarding excess in the last tile).
We then apply a separate learned mask to each tile after the ﬁrst ( i.e ., there are n − 1 masks).
All masks are a ﬁxed “window” size, which we take to be 9 by default (to match the size of commonly-used 3 × 3 kernels in CNNs), and are shared within each parameter group.
To apply, masks are multiplied element-wise over sequential windows of their respective tile.
While the number of additional parameters depends on the amount of upsampling required, as the masks are small, this is negligible.
3. M APPING L AYERS TO P ARAMETER G ROUPS
We now discuss how SSNs can automatically learn to assign layers to parameter groups in such a way that parameters can be efﬁciently shared.
This is in contrast to prior work on parameter sharing ( e.g ., Ha et al., 2016; Savarese & Maire, 2019; Jaegle et al., 2021), which required layers to be manually assigned to parameter groups.
Finding an optimal mapping of layers to parameter groups is challenging, and a brute-force approach is computationally infeasible.
We rely instead on SSNs learning a representation for each layer as part of the template-based parameter downsampling process, and then use this representation to identify similar layers which can effectively share parameters.
To do this, we perform a short preliminary training step in which we train a small ( i.e ., low parameter budget)
SSN version of the model using a single parameter group and a modiﬁed means of generating templates for parameter downsampling.
Speciﬁcally, for a layer (cid:96) i , we split θ into K (cid:48) evenly-sized templates T 1 ,...,K (cid:48) i .
Since we wish to use downsampling-based weight generation, each T k (cid:48)
i is then resized with bilinear interpolation to be the same size as w i .
Next, we train the SSN as usual, using WAvg or Emb downsampling with the modiﬁed templates for weight generation (there is no upsampling).
By using a small parameter budget and template-based weight generation where each template comes from the same underlying parameters, we encourage signiﬁcant sharing between layers so we can measure the effectiveness of sharing.
We found that using a budget equal to the number of weights of the largest single layer in the network to work well.
Further, this preliminary training step is short, and requires only 10–15% of the typical network training time.
Finally, we construct the parameter groups by clustering the learned layer representations into P groups.
As the layer representation, we take the α i or φ i learned for each layer by WAvg or Emb downsampling (resp.).
We then use k -means clustering to group these representations into P groups, which become the parameter groups used by the full SSN.
E XPERIMENTS
Our experiments include a wide variety of tasks and networks in order to demonstrate the broad applicability of NPAS and SSNs.
We adapt code and data splits made available by the authors and report the average of ﬁve runs for all comparisons except ImageNet and ALBERT, which average three runs.
A more detailed discussion on SSN hyperparameter settings can be found in Appendices B-D.
In our paper we primarily evaluate methods based on task performance, but we demonstrate that
SSNs improve reduce training time and memory in distributed learning settings in Appendix G.
Compared Tasks.
We brieﬂy describe each task, datasets, and evaluation metrics.
For each model, we use the authors’ implementation and hyperparameters, unless noted (more details in Appendix A).
Image Classiﬁcation.
For image classiﬁcation the goal is to recognize if an object is present in an image.
This is evaluated using Error@ k , i.e ., the portion of times that the correct category does not appear in the top k most likely objects.
We evaluate SSNs on CIFAR-10 and CIFAR100 (Krizhevsky, 2009), which are composed of 60K images of 10 and 100 categories, respectively, and ImageNet (Deng et al., 2009), which is composed of 1.2M images containing 1,000 categories.
We report Error@1 on CIFAR and Error@5 for ImageNet.
Image-Sentence Retrieval.
In image-sentence retrieval the goal is match across modalities (sentences and images).
This task is evaluated using Recall@K= { 1, 5, 10 } for both cross-modal directions (six numbers), which we average for simplicity.
We benchmark on Flickr30K (Young et al., 2014) which contains 30K/1K/1K images for training/testing/validation, and COCO (Lin et al., 2014), which contains 123K/1K/1K images for training/testing/validation.
For both datasets each image has about ﬁve descriptive captions.
We evaluate SSNs using EmbNet (Wang et al., 2016) and ADAPTT2I (Wehrmann et al., 2020).
Note that ADAPT-T2I has identical parallel layers ( i.e ., they need different outputs despite having the same input), which makes sharing parameters challenging.
Phrase Grounding.
Given a phrase the task is to ﬁnd its associated image region.
Performance is measured by how often the predicted box for a phrase has at least 0.5 intersection over union with its ground truth box.
We evaluate on Flickr30K Entities (Plummer et al., 2017) which augments
Flickr30K with 276K bounding boxes for phrases in image captions, and ReferIt (Kazemzadeh et al., 2014), which contains 20K images that are evenly split between training/testing and 120K region descriptions.
We evaluate our SSNs with SimNet (Wang et al., 2018) using the implementation from
Plummer et al.
(2020) that reports state-of-the-art results on this task.
Question Answering.
For this task the goal is to answer a question about a textual passage.
We use SQuAD v1.1 (Rajpurkar et al., 2016), which has 100K+ question/answer pairs on 500+ articles, and SQuAD v2.0 (Rajpurkar et al., 2018), which adds 50K unanswerable questions.
We report F and EM scores on the development split.
We compare our SSNs with ALBERT (Lan et al., 2020), a recent transformer architecture that incorporates extensive, manually-designed parameter sharing.
4. R ESULTS
We begin our evaluation in low-budget (LB-NPAS) settings.
Figure 3 reports results on image classiﬁcation, including WRNs (Zagoruyko & Komodakis, 2016), DenseNets (Huang et al., 2017), and EfﬁcientNets (Tan & Le, 2019); Table 1 contains results on image-sentence retrieval and phrase grounding.
For each task and architecture we compare SSNs to same parameter-sized networks without sharing.
In image classiﬁcation, we also report results for SWRN (Savarese & Maire, 2019) sharing; but note it cannot train a WRN-28-10 or WRN-50-2 with fewer than 12M or 40M parameters, resp.
We show that SSNs can create high-performing models with fewer parameters than SWRN is capable of, and actually outperform it using 25% and 60% fewer parameters on C-100 and ImageNet, resp.
Table 1 demonstrates that these beneﬁts generalize to vision-language tasks.
In Table 2 we also compare SSNs with ALBERT (Lan et al., 2020), which applies manually-designed parameter sharing to BERT (Devlin et al., 2019), and ﬁnd that SSN’s learned parameter sharing outperforms ALBERT.
This demonstrates that SSNs can implement large networks with lower memory requirements than is possible with current methods by effectively sharing parameters.
We discuss the runtime and memory performance implications of SSNs extensively in Appendix G.
In short, by reducing parameter counts, SSNs reduce communication costs and memory.
For example,
[Figure]
CIFAR
[Figure]
[Table]
[Figure]
[Table]
our SSN-ALBERT-Large trains about 1 . 4 × faster using 128 GPUs than BERT-Large (in line with results for ALBERT)
[Table]
and reduces memory requirements by about 5 GB ( 1 / 3 of total).
As mentioned before
[Table]
knowledge distillation and parameter pruning can help create more efﬁcient models at test time
[Table]
although they cannot reduce memory requirements during training like SSNs.
Tables 3 and 4 show our approach can be used to accomplish a similar goal as these tasks.
Comparing our LB-NPAS results in Table 4 and the lowest parameter setting of HRank, we report a 1.5% gain over pruning methods even when using less than half the parameters.
We note that one can think of our SSNs in the high budget setting (HP-NPAS) as mixing together a set of random initializations of a network by learning to combine the different templates.
This setting’s beneﬁt is illustrated in Table 3 and Table
[Table]
where our HB-NPAS models report a 1-1.5% gain over training a traditional network.
As a reminder, in this setting we precompute the weights of each layer once training is complete so they require no additional overhead at test time.
That said, best performance on both tasks comes from combining our SSNs with prior work on both tasks.
4. A NALYSIS OF S HAPESHIFTER N ETWORKS
In this section we present ablations of the primary components of our approach.
A complete ablation study, including, but not limited to comparing the number of parameter groups (Section 3.2) and number of templates ( K in Section 3.1.1) can be found in Appendices B-D.
Table 5 compares the strategies generating weights from templates described in Section 3.1.1 when using a single parameter group.
In these experiments we set the number of parameters as the amount required to implement the largest layer in a network.
For example, the ADAPT-T2I model requires 14M parameters, but its bidirectional GRU accounts for 10M of those parameters, so all SSN variants in this experiment allocate 10M parameters.
Comparing to the baseline, which involves modifying the original model’s number and/or size of ﬁlters so they have the same number of parameters as our [Table] [Table] [Table]
[Table] SSNs, we see that the variants of our SSNs perform better, especially on ImageNet where we reduce
Error@5 by 3%.
Generally we see a slight boost to performance using Emb over WAvg.
Also note that prior work in parameter sharing, e.g ., SWRN (Savarese & Maire, 2019), can not be applied to the settings in Table 5, since they require parameter sharing between layers of different types and different operations, or have too few parameters.
E.g ., the WRN-28-10 results use just under 4M parameters, but, as shown in Figure 3(a), SWRN requires a minimum of 12M parameters.
In Table 6 we investigate one of the new challenges in this work: how to upsample parameters so a large layer operation can be implemented with relatively few parameters (Section 3.1.2).
For example, our SSN-WRN-28-10 results use about 0.5M parameters, but the largest layer deﬁned in the network requires just under 4M weights.
We ﬁnd using our simple learned Mask upsampling method performs well in most settings, especially when using convolutional networks.
For example, on CIFARit improves Error@1 by 2.5% over the baseline, and 1.5% over using bilinear interpolation (inter).
While more complex methods of upsampling may seem like they would improve performance ( e.g ., using an MLP, or learning combinations of basis ﬁlters), we found such approaches had two signiﬁcant drawbacks.
First, they can slow down training time signiﬁcantly due to their complexity, so only a limited number of settings are viable.
Second, in our experiments we found many had numerical stability issues for some datasets/tasks.
We believe this may be due to trying to learn the local parameter patterns and the weights to combine them concurrently.
Related work suggests this can be resolved by leveraging prior knowledge about what these local parameter patterns should represent (Denil et al., 2013), i.e ., you deﬁne and freeze what they represent.
However, prior knowledge is not available in the general case, and data-driven methods of training these local ﬁlter patterns often rely on pretraining steps of the fully-parameterized network ( e.g ., Denil et al., 2013).
Thus, they are not suited for NPAS since they can not training large networks with low memory requirements, but addressing this issue would be a good direction for future work.
Table 7 compares approaches for mapping layers to parameter groups using the same number of parameters as the original model.
We see a small, but largely consistent improvement over using a traditional (baseline) network using SSNs.
Notably, our automatically learned mappings (auto) perform on par with manual groups.
This demonstrates that our automated approach can be used without loss in performance, while being applicable to any architecture, making them more ﬂexible than hand-crafted methods.
This ﬂexibility does come with a computational cost, as our preliminary step that learns to map layers to parameter groups resulted in a 10-15% longer training time for equivalent epochs.
That said, parameter sharing methods have demonstrated an ability to converge [Figure]
[Figure]
[Figure]
v 3 x 3 x
C o n v 3 x 3 x C o n v 3 x 3 x
C o n v 3 x 3 x C o n v 3 x 3 x C o n v 3 x 3 x C o n v 3 x 3 x
C o n v 3 x 3 x C o n v 3 x 3 x C o n v 3 x 3 x
C o n v 3 x 3 x C o n v 3 x 3 x
F C 640 x
Conv 1x1x Conv 1x1x
Conv 1x1x
O u t pu t
[Figure]
[Table]
[Figure]
[Table]
[Figure]
Layer
No Sharing
SSN (1 Param Group) Layer Layer Parameter groups for a WRN-28𝑃 = 14 parameter groups (maximum sharing possible with the manual assignments of Savarese & Maire, 2019)
[Table]
our automatically learned mappings (auto) pro vide more ﬂexibility without compared SSNs (ours) Dataset [Table]
[Table]
faster (Bagherinezhad et al., 2017; Lan et al., 2020).
Thus, exploring more efﬁcient training strategies using NPAS methods like SSNs will be a good direction for future work.
Figure 4(a) compares the 3 × 3 kernel ﬁlters at the early, middle, and ﬁnal convolutional layers of a WRN-16-2 of a traditional neural network (no parameter sharing) and our SSNs where all layers belong to the same parameter group.
We observe a correspondence between ﬁlters in the early layers, but this diverges in deeper layers.
This suggests that sharing becomes more difﬁcult in the ﬁnal layers, which is consistent with two observations we made about Figure 4(b), which visualizes parameter groups used for SSN-WRN-28-10 to create 14 parameter group mappings.
First, we found the learned parameter group mappings tended to share parameters between layers early in the network, opting for later layers to share no parameters.
Second, the early layers tended to group layers into 3–4 parameter stores across different runs, with the remaining 10–11 parameter stores each containing a single layer.
Note that these observations were consistent across different random initializations.
C ONCLUSION
We propose NPAS, a novel task in which the goal is to implement a given, arbitrary network architecture given a ﬁxed parameter budget.
This involves identifying how to assign parameters to layers and implementing layers with their assigned parameters (which may be of any size).
To address NPAS, we introduce SSNs, which automatically learn how to share parameters.
SSNs beneﬁt from parameter sharing in the low-budget regime—reducing memory and communication requirements when training—and enable a novel high-budget regime that can improve model performance.
We show that SSNs boost results on ImageNet by 3% improvement in Error@5 over a same-sized network without parameter sharing.
Surprisingly, we also ﬁnd that parameters can be shared among very different layers.
Further, we show that SSNs can be combined with knowledge distillation and parameter pruning to achieve state-of-the-art results that also reduce FLOPs at test time.
One could think of SSNs as spreading the same number of parameters across more layers, increasing effective depth, which beneﬁts generalization (Telgarsky, 2016), although this requires further exploration. Acknowledgements.
This work is funded in part by grants from the National Science Foundation and DARPA.
This project received funding from the European Research Council (ERC) under the
European Union’s Horizon 2020 program (grant agreement MAELSTROM, No. 955513). N.D. is supported by the ETH Postdoctoral Fellowship. We thank the Livermore Computing facility for the use of their GPUs for some experiments. E THICS S TATEMENT
Neural Parameter Allocation Search and Shapeshifter Networks are broad and general-purpose, and therefore applicable to many downstream tasks.
It is thus challenging to identify speciﬁc cases of beneﬁt or harm, but we note that reducing memory requirements can have broad implications.
In particular, this could allow potentially harmful applications to be cheaply and widely deployed ( e g ., facial recognition, surveillance) where it would otherwise be technically or economically infeasible.
R EPRODUCIBILITY S TATEMENT
NPAS is a task which can be implemented in many different ways; we deﬁne it formally in Section 2.
SSNs, our proposed solution to NPAS, are presented in detail in Section and Figure 2 provides an illustration and example for weight generation methods.
AppendixA also provides a thorough discussion of the implementation details.
To further aid reproducibility, we publicly release our SSN code at https://github.com/BryanPlummer/SSN .
A D ESCRIPTION OF C OMPARED T ASKS A. I MAGE -S ENTENCE R ETRIEVAL
In bidirectional image-sentence retrieval when a model is provided with an image the goal is to retrieve a relevant sentence and vice versa.
This task is evaluated using Recall@K= { 1, 5, 10 } for both directions (resulting in 6 numbers), which we average for simplicity.
We benchmark methods on two common datasets: Flickr30K (Young et al., 2014) which contains 30K/1K/1K images for training/testing/validation, each with ﬁve descriptive captions, and MSCOCO (Lin et al., 2014), which contains 123K/1K/1K images for training/testing/validation, each image having roughly ﬁve descriptive captions.
EmbNet (Wang et al., 2016).
This network learns to embed visual features for each image computed using a 152-layer Deep Residual Network (ResNet) (He et al., 2016) that has been trained on
ImageNet (Deng et al., 2009) and the average of MT GrOVLE
(Burns et al., 2019) language features representing each word into a shared semantic space using a triplet loss.
The network consists of two branches, one for each modality, and each branch contains two fully connected layers (for a total of four layers).
We adapted the implementation of Burns et al .
(Burns et al., 2019) 1 , and left all hyperparameters at the default settings.
Speciﬁcally, we train using a batch size of 500 with an initial learning rate of 1e-4 which we decay exponentially with a gamma of 0.794 and use weight decay of 0.001.
The model is trained until it has not improved performance on the validation set over the last 10 epochs.
This architectures provides a simple baseline for parameter sharing with our Shapeshifter
Networks (SSNs), where layers operate on two different modalities.
ADAPT-T2I (Wehrmann et al., 2020).
In this approach word embeddings are aggregated using a bidirectional GRU (Cho et al., 2014) and its hidden state at each timestep is averaged to obtain a fullsentence representation.
Images are represented using 36 bottom-up image region features (Anderson et al., 2018) that are passed through a fully connected layer.
Then, each sentence calculates scaling and shifting parameters for the image regions using a pair of fully connected layers that both take the full-sentence representation as input.
The image regions are then averaged, and a similarity score is computed between the sentence-adapted image features and the fully sentence representation.
Thus, there are four layers total (3 fully connected, 1 GRU) that can share parameters, including the two parallel fully connected layers ( i.e ., they both take the full sentence features as input, but are expected to have different outputs).
We adapted the author’s implementation and kept the default hyperparameters 2 .
Speciﬁcally, we use a latent dimension of 1024 for our features and train with a batch size of 105 using a learning rate of 0.001.
This method was selected since it achieves high performance and also included fully connected and recurrent layers, as well as having a set of parallel layers that make effectively performing cross-layer parameter sharing more challenging.
A. P HRASE G ROUNDING
Given a phrase the goal of a phrase grounding model is to identify the image region described by the phrase.
Success is achieved if the predicted box has at least 0.5 intersection over union with the ground truth box.
Performance is measured using the percent of the time a phrase is accurately localized.
We evaluate on two datasets: Flickr30K Entities (Plummer et al., 2017) which augments the Flickr30K dataset with 276K bounding boxes associated with phrases in the descriptive captions, and ReferIt (Kazemzadeh et al., 2014) which contains 20K images that are evenly split between training/testing and 120K region descriptions.
SimNet (Wang et al., 2018).
This network contains three branches that each operate on different types of features.
One branch passes image regions features computed with a 101-layer ResNet that have been ﬁne-tuned for phrase grounding using two fully connected layers.
A second branch passes MT GrOVLE features through two fully connected layers.
Then, a joint representation is computed for all region-phrase pairs using an elementwise product.
Finally, the third branch passes these joint features through three fully connected layers (7 total), where the ﬁnal layer acts as a classiﬁer indicating the likelihood that phrase is present in the image region.
We adapt the code from Plummer et al.
and keep all hyperameters at their default settings.
Speciﬁcally, we use a pretrained Faster R-CNN model (Ren et al., 2015) ﬁne-tuned for phrase grounding by Plummer et al.
(2020) on each dataset to extract region features.
Then we encode each phrase by averaging
MT GrOVLE features (Burns et al., 2019) and provide the image and phrase features as input to our model.
We train our model using a learning rate of 5e-5 and a ﬁnal embedding dimension of 256 until it no longer improves on the validation set for 5 epochs (typically resulting in training times of 15epochs).
Performing experiments on this model enables us to test how well our SSNs generalize to another task and how well it can adapt to sharing parameters with layers operating on three types of features (just vision, just language, and a joint representation).
A. I MAGE C LASSIFICATION
For image classiﬁcation the goal is to be able to recognize if an object is present in an image.
Typically this task is evaluated using Error@K, or the portion of times that the correct category doesn’t appear in the top k most likely objects.
We evaluate our Shapeshifter Networks on three datasets: CIFARand CIFAR-100 (Krizhevsky, 2009), which are comprised of 60K images of 10 and 100 categories, respectively, and ImageNet (Deng et al., 2009), which is comprised of 1.2M images containing 1,000 categories.
We report Error@1 for both CIFAR datasets and Error@5 for ImageNet.
In these appendices, we also report Error@1 for ImageNet.
Wide Residual Network (WRN) (Zagoruyko & Komodakis, 2016).
WRN modiﬁed the traditional
ResNets by increasing the width k of each layer while also decreasing the depth d , which they found improved performance.
Different variants are identiﬁed using WRN- d - k .
Following Savarese et al .
(Savarese & Maire, 2019), we evaluate our Shapeshifter Networks using WRN-28-10 for CIFAR and WRN-50-2 for ImageNet.
We adapt the implementation of Savarese et al . 4 and use cutout (DeVries & Taylor, 2017) for data augmentation.
Speciﬁcally, on CIFAR we train our model using a batch size of 128 for 200 epochs with weight decay set at 5e-4 and an initial learning rate of 0.
which we decay using a gamma of 0.2 at 60, 120, and 160 epochs.
Unlike the vision-language models discussed earlier, these architecture include convolutional layers in addition to a fully connected layer used to implement a classiﬁer, and also have many more layers than the shallow vision-language models.
DenseNet (Huang et al., 2017).
Unlike traditional neural networks where each layer in the network is computed in sequence, every layer in a DenseNet using feature maps from every layer which came before it.
We adapt PyTorch’s ofﬁcial implementation 5 using the hyperparameters as set in Huang et al .
(Huang et al., 2017).
Speciﬁcally, on CIFAR we train our model using a batch size of 96 for 300 epochs with weight decay set at 1e-4 and an initial learning rate of 0.1 which we decay using a gamma of 0.1 at 150 and 225.
These networks provide insight into the effect depth has on learning
SSNs, as we use a 190-layer DenseNet-BC conﬁguration for CIFAR.
However, due to their high computational cost we provide limited results testing only some settings.
EfﬁcientNet (Tan & Le, 2019).
EfﬁcientNets are a class of model designed to balance depth, width, and input resolution in order to produce very parameter-efﬁcient models.
For ImageNet, we adapt an existing PyTorch implementation and its hyperparameters 6 , which are derived from the ofﬁcial
TensorFlow version.
We use the EfﬁcientNet-B0 architecture to illustrate the impact of SSNs on very parameter-efﬁcient, state-of-the-art models.
On CIFAR-100 we use an EfﬁcientNet with Network
Deconvolution (ND)
(Ye et al., 2020), which results in improved results with similar numbers of epochs for training.
We use the author’s implementation 7 , and train each model for 100 epochs (their best performing setting).
Note that our best error running different conﬁgurations of their model (35.88) is better than those in their paper (37.63), so despite the relatively low performance it is comparable to results from their paper.
A. Q UESTION A NSWERING
In question answering, a model is given a question and an associated textual passage which may contain the answer, and the goal is to predict the span of text in the passage that contains the answer.
We use two versions of the Stanford Question Answering Dataset (SQuAD), SQuAD v1.1 (Rajpurkar et al., 2016), which contains 100K+ question/answer pairs on 500+ Wikipedia particles, and SQuAD v2.0, which augments SQuAD v1.1 with 50K unanswerable questions designed adversarially to be similar to standard SQuAD questions.
For both datasets, we report both the F1 score, which captures the precision and recall of the chosen text span, and the Exact Match (EM) score.
ALBERT (Lan et al., 2020)
ALBERT is a version of the BERT (Devlin et al., 2019) transformer architecture that applies cross-layer parameter sharing.
Speciﬁcally, the parameters for all components of a transformer layer are shared among all the transformer layers in the network.
ALBERT also includes a factorized embedding to further reduce parameters.
We follow the methodology of BERT and ALBERT for reporting results on SQuAD, and our baseline ALBERT scores closely match those reported in the original work.
This illustrates the ability of NPAS and SSNs to develop better parameter sharing methods than manually-designed systems for extremely large models.
B E XTENDED R ESULTS WITH A DDITIONAL B ASELINES
Below we provide additional results with more baseline methods for the three components of our
SSNs: weight generator (Section B.1), parameter upsampling (Section B.4), and mapping layers to parameter groups (Section B.3).
We provide ablations on the number of parameter groups and templates used by our SSNs in Section C and Section D, respectively.
B. A DDITIONAL M ETHODS THAT G ENERATE L AYER W EIGHTS FROM T EMPLATES
Parameter downsampling uses the selected templates T ki for a layer (cid:96) i to produce its weights w i .
In
Section 3.1.1 of the paper we discuss two methods of learning a combination of the T ki to generate w i .
Below in Section B.2 we provide two simple baseline methods that directly use the candidates.
Table 8 compares the baselines to the methods in the main paper that learn weighted combinations of templates, where the learned methods typically perform better than the baselines.
B. D IRECT T EMPLATE C OMBINATION
Here we describe the strategies we employ that require no parameters to be learned by weight generator, i.e ., they operate directly on the templates T ki .
Round Robin (RR) reuses parameters of each template set as few times as possible.
The scheme simply returns the weights at index k mod K in the (ordered) template set T i at the k th query of a parameter group.
Candidate averaging (Avg) averages all candidates in T i to provide a naive baseline for using multiple candidates.
A signiﬁcant drawback of this approach is that, if K is large, this can result in reusing parameters (across combiners) many times with no way to adapt to a speciﬁc layer, especially when the size of the parameter group is small.
B. A DDITIONAL P ARAMETER M APPING R ESULTS
Table 9 compares approaches that map layers to parameter groups using the same number of parameters as the original model.
We see a small, but largely consistent improvement over using a traditional (baseline) network.
Notably, our learned grouping methods (WAvg, Emb) perform on par, and sometimes better than using manual mappings.
However, our approach can be applied to any architecture to create a selected number of parameter groups, making them more ﬂexible than hand-crafted methods.
For example, in Table 10, we see using two groups often helps to improve performance when using very few parameters, but it is not clear how to effectively create two groups by hand for many networks.
[Table]
[Table]
of the paper)
Each method uses the same number of parameters and the baseline represents no parameter sharing.
Notably, these results demonstrate how our automatically learned groups can provide greater ﬂexibility when sharing parameters while performing on par or better than manually created groups (which are illustrated in
Figure 5).
[Table]
[Table]
[Table]
[Table]
COCO 81.
[Table]
[Table]
[Table]
[Table]
COCO
[Table]
[Table]
( ↑ better)
SimNet
F30K Entities 71.
ReferIt 61.
[Table]
( ↓ better)
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table] [Table] [Table] [Table]
[Table] [Table] [Table]
E@
[Table] [Table] [Table] [Table]
[Table]
P ARAMETER U PSAMPLING
In Table 10 we provide extended results comparing the parameter upsamping methods.
We additionally compare with a further na¨ıve baseline of simply repeating parameters until they are the appropriate size.
We ﬁnd that Mask upsampling is always competitive, and typically moreso when two parameter groups are used.
B. C OMPARISON W ITH H YPERNETWORKS
In Table 11 we compare our SSNs on Wide ResNets (Ha et al., 2016) to the same networks implemented using Hypernetworks (Ha et al., 2016) for CIFAR-10, using the results reported in their paper.
We can see that, for the same parameter budget, SSNs outperform Hypernetworks.
[Figure]
when upsampling.
“Reduced
Baseline” uses no parameter sharing, but adjusts the number and/or size of ﬁlters so they have the same number of parameters as our SSNs.
The original model’s performance is reported in the ﬁrst column of results in Table 9.
Orig Num % of Params
[Table]
Dataset
[Table]
C E FFECT OF THE N UMBER OF P ARAMETER G ROUPS P
A signiﬁcant advantage of using learned mappings of layers to parameter groups, described in
Section 3.2, is that our approach can support any number of parameter groups, unlike prior work that required manual grouping and/or heuristics to determine which layers shared parameters ( e.g ., Lan et al., 2020; Savarese & Maire, 2019). In this section we explore how the number of parameter groups
[Table]
[Table]
[Table]
effects performance on the image classiﬁcation task.
We do not benchmark bidirectional retrieval and phrase grounding since networks addressing these tasks have few layers, so parameter groups are less important (as shown in Table 7).
Table 12 reports the performance of our SSNs when using different numbers P parameter groups.
We ﬁnd that when training with few parameters (ﬁrst line) low numbers of parameter groups work best, while when more parameters are available larger numbers of groups work better (second line).
In fact, there is a signiﬁcant drop in performance going from 4 to 8 groups when training with few parameters as seen in the ﬁrst line of Table 12.
This is due to the fact that starting at 8 groups some parameter groups had too few weights to implement their layers, resulting in extensive parameter upsampling.
This suggests that we may be able to further improve performance when there are few parameters by developing better methods of implementing layers when too few parameters are available.
D E FFECT OF THE N UMBER OF T EMPLATES K
Table 13 reports the results using different numbers of templates.
We ﬁnd that varying the number of templates only has a minor impact on performance most of the time.
We note that more templstes tends to lead to reduced variability between runs, making results more stable.
As a reminder, however, the number of templates does not guarantee that each layer will have enough parameters to construct them.
Thus, parameter groups only use this hyperparameter when many weights are available to it ( i.e ., it can form multiple templates for the layers it implements).
This occurs for the phrase grounding and bidirectional retrieval results at the higher maximum numbers of templates.
E S CALING SSN S TO LARGER NETWORKS
Table 14 demonstrates the ability of our SSNs to signiﬁcantly reduce the parameters required, and thus the memory required to implement large Wide ResNets so they fall within speciﬁc bounds.
For example, Table 14(b) shows larger and deeper conﬁgurations continue to improve performance even when the number of parameters remains largely constant.
Comparing the ﬁrst line of Table 14(a) and the last line of Table 14(c)
we see that SSN-WRN-76-12 outperforms the fully-parameterized WRN28-10 network by 0.6% on CIFAR-100 while only using just over half the parameters, and comes within 0.5% of WRN-76-12 while only using 13.0% of its parameters.
We do note that using a SSN does not reduce the number of ﬂoating point operations, so although our SSN-WRN-76-12 model uses fewer parameters than the WRN-28-10, it is still slower at both test and train time.
However, our results help demonstrate that SSNs can be used to implement very large networks with lower memory [Table] [Table] requirements by effectively sharing parameters.
This enables us to train larger, better-performing networks than is possible with traditional neural networks on comparable computational resources [Table]
F I MAGE C LASSIFICATION N UMBERS
We provide raw numbers for the results in Figure 3 in Table 15 (CIFAR-100) and Table 16 (ImageNet).
G P ERFORMANCE I MPLICATIONS OF NPAS AND SSN S
Our SSNs can offer several performance beneﬁts by reducing parameter counts; notably, they can reduce memory requirements storing a model and can reduce communication costs for distributed training.
We emphasize that LB-NPAS does not reduce FLOPs, as the same layer operations are implemented using fewer parameters.
Should fewer FLOPs also be desired, SSNs can be combined [Table] with other techniques [Table]
such as pruning.
Additionally
[Table] we note that our implementation has not been extensively optimized [Table] and further performance improvements could likely be achieved with additional engineering.
G. C OMMUNICATION C OSTS FOR D ISTRIBUTED T RAINING
Communication for distributed data-parallel training is typically bandwidth-bound [Table] and thus employs bandwidth-optimal allreduces, which are linear in message length (Chan et al., 2007).
Thus, we expect communication time to be reduced by a factor proportional to the parameter savings achieved by
NPAS, all else being equal.
However, frameworks will typically execute allreduces layer-wise as soon as gradient buffers are ready to promote communication/computation overlap in backpropagation; reducing communication that is already fully overlapped is of little beneﬁt.
Performance beneﬁts are thus sensitive to the model, implementation details [Table] and the system being used for training.
For CNNs, we indeed observe minor performance improvements, as the number of parameters is typically small.
When using 64 V100 GPUs for training WRN-50-2 on ImageNet, we see a 1.04 × performance improvement in runtime per epoch when using SSNs with 10.5M parameters (15% of the original model).
This is limited because most communication is overlapped.
We also observe small performance improvements in some cases because we launch fewer allreduces, resulting in less demand for SMs and memory bandwidth on the GPU.
These performance results are in line with prior work on communication compression for CNNs ( e.g ., Renggli et al., 2019).
For large transformers, however, we observe more signiﬁcant performance improvements.
The
SSN-ALBERT-Large is about 1.4 × faster using 128 GPUs than the corresponding BERT-Large model.
This is in line with the original ALBERT work (Lan et al., 2020), which reported that training
ALBERT-Large was 1.7 × faster than BERT-Large when using 128 TPUs.
Note that due to the differences in the systems for these results, they are not directly comparable.
We would also reiterate that for some applications where communication is more costly, say, for federated learning applications ( e.g . McMahan et al.
; Kone ˇ cn ´ y et al.
, our approach would be even more beneﬁcial due to the decreased message length.
G. M EMORY S AVINGS
LB-NPAS and SSNs reduce the number of parameters, which consequentially reduces the size of the gradients and optimizer state ( e.g ., momentum) by the same amount.
It does not reduce the storage requirements for activations, but note there is much work on recomputation to address this ( e.g ., Chen et al., 2016; Jain et al., 2020).
Thus, the memory savings from SSNs is independent of batch size.
For SSN-ALBERT-Large, we use 18M parameters (5% of BERT-Large, which contains 334M parameters).
Assuming FP32 is used to store data, we save about 5 GB of memory in this case (about 1 / 3 of the memory used)