Parameter Prediction for Unseen Deep Architectures

1 Introduction . Consider the problem of training deep neural networks on large annotated datasets , such as ImageNet [ 1 ] . This problem can be formalized as finding optimal parameters for a given neural network a , parameterized by w , w.r.t . a loss function L on the dataset D = { ( xi , yi ) } Ni=1 of inputs xi and targets yi : arg min w ∑N i=1 L ( f ( xi ; a , w ) , yi ) , ( 1 ) where f ( xi ; a , w ) represents a forward pass . Equation 1 is usually minimized by iterative optimization algorithms – e.g . SGD [ 2 ] and Adam [ 3 ] – that converge to performant parameters wp of the architecture a . Despite the progress in improving the training speed and convergence [ 4–7 ] , obtaining wp remains a bottleneck in large-scale machine learning pipelines . For example , training a ResNet-50 [ 8 ] on ImageNet can take many GPU hours [ 9 ] . With the ever growing size of networks [ 10 ] and necessity of training the networks repeatedly ( e.g . for hyperparameter or architecture search ) , the classical process of obtaining wp is becoming computationally unsustainable [ 11–13 ] . A new parameter prediction task . When optimizing the parameters for a new architecture a , typical optimizers disregard past experience gained by optimizing other nets . However , leveraging this past experience can be the key to reduce the reliance on iterative optimization and , hence the high computational demands . To progress in that direction , we propose a new task where iterative optimization is replaced with a single forward pass of a hypernetwork [ 14 ] HD . To tackle the task , 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . ∗Part of the work was done while interning at Facebook AI Research . HD is expected to leverage the knowledge of how to optimize other networks F . Formally , the task is to predict the parameters of an unseen architecture a /∈ F using HD , parameterized by θp : ŵp = HD ( a ; θp ) . The task is constrained to a dataset D , so ŵp are the predicted parameters for which the test set performance of f ( x ; a , ŵp ) is similar to the one of f ( x ; a , wp ) . For example , we consider CIFAR-10 [ 15 ] and ImageNet image classification datasets D , where the test set performance is classification accuracy on test images . Approaching our task . A straightforward approach to expose HD to the knowledge of how to optimize other networks is to train it on a large training set of { ( ai , wp , i ) } pairs , however , that is prohibitive2 . Instead , we follow the bi-level optimization paradigm common in meta-learning [ 16–18 ] , but rather than iterating over M tasks , we iterate over M training architectures F = { ai } Mi=1 : arg min θ ∑N j=1 ∑M i=1 L ( f ( xj ; ai , HD ( ai ; θ ) ) , yj ) . ( 2 ) By optimizing Equation 2 , the hypernetwork HD gradually gains knowledge of how to predict performant parameters for training architectures . It can then leverage this knowledge at test time – when predicting parameters for unseen architectures . To approach the problem in Equation 2 , we need to design the network space F and HD . For F , we rely on the previous design spaces for neural architectures [ 19 ] that we extend in two ways : the ability to sample distinct architectures and an expanded design space that includes diverse architectures , such as ResNets and Visual Transformers [ 20 ] . Such architectures can be fully described in the form of computational graphs ( Fig . 1 ) . So , to design the hypernetwork HD , we rely on recent advances in machine learning on graph-structured data [ 21–24 ] . In particular , we build on the Graph HyperNetworks method ( GHNs ) [ 24 ] that also optimizes Equation 2 . However , GHNs do not aim to predict large-scale performant parameters as we do in this work , which motivates us to improve on their approach . By designing our diverse space F and improving on GHNs , we boost the accuracy achieved by the predicted parameters on unseen architectures to 77 % ( top-1 ) and 48 % ( top-5 ) on CIFAR-10 [ 15 ] and ImageNet [ 1 ] , respectively . Surprisingly , our GHN shows good out-of-distribution generalization and predicts good parameters for architectures that are much larger and deeper compared to the ones seen in training . For example , we can predict all 24 million parameters of ResNet-50 in less than a second either on a GPU or CPU achieving ∼60 % on CIFAR-10 without any gradient updates ( Fig 1 , ( b ) ) . Overall , our framework and results pave the road toward a new and significantly more efficient paradigm for training networks . Our contributions are as follows : ( a ) we introduce the novel task of predicting performant parameters for diverse feedforward neural networks with a single hypernetwork forward pass ; ( b ) we introduce DEEPNETS-1M – a standardized benchmark with in-distribution and out-of-distribution architectures to track progress on the task ( § 3 ) ; ( c ) we define several baselines and propose a GHN model ( § 4 ) that performs surprisingly well on CIFAR-10 and ImageNet ( § 5.1 ) ; ( d ) we show that our model learns a strong representation of neural network architectures ( § 5.2 ) , and our model is useful for initializing neural networks ( § 5.3 ) . Our DEEPNETS-1M dataset , trained GHNs and code is available at https : //github.com/facebookresearch/ppuda . D eco d er 2Training a single network ai can take several GPU days and thousands of trained networks may be required . 2 2 Background . We start by providing a brief background about the network design spaces leveraged in the creation of our DEEPNETS-1M dataset of neural architectures described in § 3 . We then cover elements of graph hypernetworks that we leverage when designing our specific GHN HD in § 4 . 2.1 Network Design Space of DARTS . DARTS [ 19 ] is a differentiable NAS framework . For image classification tasks such as those considered in this work , its networks are defined by four types of building blocks : stems , normal cells , reduction cells , and classification heads . Stems are fixed blocks of convolutional operations that process input images . The normal and reduction cells are the main blocks of architectures and are composed of : 3×3 and 5×5 separable convolutions , 3×3 and 5×5 dilated separable convolutions , 3×3 max pooling , 3×3 average pooling , identity and zero ( to indicate the absence of connectivity between two operations ) . Finally , the classification head defines the network output and is built with a global pooling followed by a single fully connected layer . Typically , DARTS networks have one stem block , 14-20 cells , and one classification head , altogether forming a deep computational graph . The reduction cells , placed only at 1/3 and 2/3 of the total depth , decrease the spatial resolution and increase the channel dimensionality by a factor of 2 . Summation and concatenation are used to aggregate outputs from multiple operations within each cell . To make the channel dimensionalities match , 1×1 convolutions are used as needed . All convolutional operations use the ReLU-Conv-Batch Norm ( BN ) [ 7 ] order . Overall , DARTS enables defining strong architectures that combine many principles of manual [ 25 , 8 , 26 , 27 ] and automatic [ 24 , 28–33 ] design of neural architectures . While DARTS learns the optimal task-specific cells , the framework can be modified to permit sampling randomly-structured cells . We leverage this possibility for the DEEPNETS-1M construction in § 3 . Please see § A.1 for further details on DARTS . 2.2 Graph HyperNetwork : GHN-1 . Representation of architectures . GHNs [ 24 ] directly operate on the computational graph of a neural architecture a . Specifically , a is a directed acyclic graph ( DAG ) , where nodes V = { vi } |V |i=1 are operations ( e.g . convolutions , fully-connected layers , summations , etc . ) and their connectivity is described by a binary adjacency matrix A ∈ { 0 , 1 } |V |×|V | . Nodes are further characterized by a matrix of initial node features H0 = [ h01 , h 0 2 , ... , h 0 |V | ] , where each h 0 v is a one-hot vector representing the operation performed by the node . We also use such a one-hot representation for H0 , but in addition encode the shape of parameters associated with nodes as described in detail in § B.1 . Design of the graph hypernetwork . In [ 24 ] , the graph hypernetwork HD consists of three key modules . The first module takes the input node features H0 and transforms them into d-dimensional node features H1 ∈ R|V |×d through an embedding layer . The second module takes H1 together with A and feeds them into a specific variant of the gated graph neural network ( GatedGNN ) [ 34 ] . In particular , their GatedGNN mimics the canonical order π of node execution in the forward ( fw ) and backward ( bw ) passes through a computational graph . To do so , it sequentially traverses the graph and performs iterative message passing operations and node feature updates as follows : ∀t ∈ [ 1 , ... , T ] : [ ∀π ∈ [ fw , bw ] : ( ∀v ∈ π : mtv = ∑ u∈Nπv MLP ( htu ) , h t v = GRU ( h t v , m t v ) ) ] , ( 3 ) where T denotes the total number of forward-backward passes ; htv corresponds to the features of node v in the t-th graph traversal ; MLP ( · ) is a multi-layer perceptron ; and GRU ( · ) is the update function of the Gated Recurrent Unit [ 35 ] . In the forward propagation ( π = fw ) , N πv corresponds to the incoming neighbors of the node defined by A , then in the backward propagation ( π = bw ) it similarly corresponds to the outgoing neighbors of the node . The last module uses the GatedGNN output hidden states hTv to condition a decoder that produces the parameters ŵ v p ( e.g . convolutional weights ) associated with each node . In practice , to handle different parameter dimensionalities per operation type , the output of the hypernetwork is reshaped and sliced according to the shape of parameters in each node . We refer to the model described above as GHN-1 ( Fig . 1 ) . Further subtleties of implementing this model in the context of our task are discussed in § B.1 . 3 DEEPNETS-1M . The network design space of DARTS is limited by the number of unique operations that compose cells , and the low variety of stems and classification heads . Thus , many architectures are not realizable within this design space , including : VGG [ 25 ] , ResNets [ 8 ] , MobileNet [ 33 ] or more recent ones such as Visual Transformer ( ViT ) [ 20 ] and Normalization-free networks [ 36 , 37 ] . Furthermore , DARTS does not define a procedure to sample random architectures . By addressing these two limitations we aim to expose our hypernetwork to diverse training architectures and permit its evaluation on common architectures , such as ResNet-50 . We hypothesize that increased training diversity can improve hypernetworks ’ generalization to unseen architectures making it more competitive to iterative optimizers . Extending the network design space . We extend the set of possible operations with non-separable 2D convolutions3 , Squeeze & Excite4 ( SE ) [ 40 ] and Transformer-based operations [ 41 , 20 ] : multihead self-attention ( MSA ) , positional encoding and layer norm ( LN ) [ 42 ] . Each node ( operation ) in our graphs has two attributes : primitive type ( e.g . convolution ) and shape ( e.g . 3×3×512×512 ) . Overall , our extended set consists of 15 primitive types ( Table 1 ) . We also extend the diversity of the generated architectures by introducing VGG-style classification heads and ViT stems . Finally , to further increase architectural diversity , we allow the operations to not include batch norm ( BN ) [ 7 ] and permit networks without channel width expansion ( e.g . as in [ 20 ] ) . Architecture generation process . We generate different subsets of architectures ( see the description of each subset in the next two paragraphs and in Table 1 ) . For each subset depending on its purpose , we predefine a range of possible model depths ( number of cells ) , widths and number of nodes per cell . Then , we sample a stem , a normal and reduction cell and a classification head . The internal structure of the normal and reduction cells is defined by uniformly sampling from all available operations . Due to a diverse design space it is extremely unlikely to sample the same architecture multiple times , but we ran a sanity check using the Hungarian algorithm [ 43 ] to confirm that ( see Figure 6 in § A.2 for details ) . In-distribution ( ID ) architectures . We generate a training set of |F| = 106 architectures and validation/test sets of 500/500 architectures that follow the same generation rules and are considered to be ID samples . However , training on large architectures can be prohibitive , e.g . in terms of GPU memory . Thus , in the training set we allow the number of channels and , hence the total number of parameters , to be stochastically defined given computational resources . For example , to train 3Non-separable convolutions have weights of e.g . shape 3×3×512×512 as in ResNet-50 . NAS works , such as DARTS and GHN , avoid such convolutions , since the separable ones [ 38 ] are more efficient . Non-separable convolutions are nevertheless common in practice and can often boost the downstream performance . 4The Squeeze & Excite operation is common in many efficient networks [ 39 , 12 ] . our models we upper bound the number of parameters in the training architectures to around 3M by sampling fewer channels if necessary . In the evaluation sets , the number of channels is fixed . Therefore , this pre-processing step prior to training results in some distribution shift between the training and the validation/test sets . However , the shift is not imposed by our dataset . Out-of-distribution ( OOD ) architectures . We generate five OOD test sets that follow different generation rules . In particular , we define WIDE and DEEP sets that are of interest due the stronger downstream performance of such nets in large-scale tasks [ 47 , 48 , 10 ] . These nets are often more challenging to train for fundamental [ 49 , 50 ] or computational [ 51 ] reasons , so predicting their parameters might ease their subsequent optimization . We also define the DENSE set , since networks with many operations per cell and complex connectivity are underexplored in the literature despite their potential [ 27 ] . Next , we define the BN-FREE set that is of interest due to BN ’ s potential negative side-effects [ 52 , 53 ] and the difficulty or unnecessity of using it in some cases [ 54–56 , 36 , 37 ] . We finally add the RESNET/VIT set with two predefined image classification architectures : commonlyused ResNet-50 [ 8 ] and a smaller 12-layer version of the Visual Transformer ( ViT ) [ 20 ] that has recently received a lot of attention in the vision community . Please see § A.1 and § A.2 for further details and statistics of our DEEPNETS-1M dataset .

Summary: This paper proposes and analyzes a method for generating trained weights already optimized for a given single task for a large class of diverse architectures using Graph hypernetworks. They introduce a new DEEPNETS-1M benchmark with in-distribution and out of distribution architectures. The work also proposed three main improvements to improve performance of their GHN: normalization of predicted parameters, enhanced long-range interactions in the GHN, and meta-batching of architectures.  They demonstrate that their method is able to produce weights that do surprisingly well on unseen and diverse networks in less than 1 second on CPU or GPU for CIFAR-10 and Imagenet. They demonstrate that even for large Resnet-50s which are out of the training distribution, they could achieve 60% accuracy on CIFAR10.