Rethinking Neural Operations for Diverse Tasks

1 Introduction . Automated machine learning ( AutoML ) and neural architecture search ( NAS ) are often motivated by a vision of democratizing ML by reducing the need for expert design on a variety of tasks . While NAS has grown rapidly with developments such as weight-sharing [ 36 ] and “ NAS-benches ” [ 47 , 49 ] , most efforts focus on search spaces that glue together established primitives for well-studied tasks like vision and text [ 32 , 26 , 45 , 25 ] or on issues such as latency [ 8 , 13 ] . In this work , we revisit the broader vision of NAS and propose to move towards much more general search spaces while still exploiting successful network topologies . To do so we focus on expanding the set of operations , which is usually fairly small ; for example , that of the well-studied DARTS space has eight elements : a few types of convolution and pooling layers [ 32 ] . The baseline approach for expanding this set—adding operations one-by-one—scales poorly and will not result in new operations when faced with new types of data . Our core contribution is a re-imagining of NAS operation spaces that drastically expands this set in a principled fashion to include both standard operations as well as a wide range of new ones . To do so we exploit the fact that most standard operations used in modern NAS return linear transforms diagonalized by the discrete Fourier transform ( DFT ) . Replacing the DFT matrices in the diagonal decomposition by a more expressive family of efficient linear transforms known as Kaleidoscope or ∗ denotes equal contribution . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . K-matrices [ 10 ] yields the set of Expressive Diagonalization ( XD ) Operations , which comprise a large search space containing various types of grid-based convolutions and pooling , permutations , transposed convolutions , certain kinds of graph convolutions , the Fourier Neural Operator ( FNO ) [ 30 ] , and infinitely many more . This broad expressivity reflects the key insight of our work : that many of the most important neural operations in ML consist of multiple channels that apply weights w to inputs x by computing K diag ( Lw ) Mx ( 1 ) where the matrices K , L , and M are efficient ( to represent and apply ) and shared across channels . We leverage XD-operations to take critical steps towards a broader NAS that enables the discovery of good design patterns with limited human specification from data in under-explored domains . To do so we develop a simple procedure which transforms any backbone convolutional neural network ( CNN ) into an architecture search space by replacing its operations with XD-operations . This space is then searched using a simple weight-sharing algorithm that needs only a small amount of tuning to find effective operations . As a simple first demonstration , we show that XD-operations yield models that are 15 % more accurate than standard discrete search spaces on permuted CIFAR-10 , highlighting the fragility of standard NAS operation spaces on new datasets , and thus the need for XD-operations . As our main evaluation , we demonstrate the effectiveness of XD-operations in a series of applications showing that , starting from vanilla CNNs , they consistently outperform custom-designed operations . • Learning to solve partial differential equations ( PDEs ) : when substituted into a simple CNN backbone , XD-operations outperform convolutions and the dense prediction NAS method AutoDeepLab [ 31 ] , and even achieve lower error than custom-designed , state-of-the-art operations ( FNOs [ 30 ] ) across three problems with different dimensionalities ( Burgers ’ equation , Darcy Flow , and Navier-Stokes ) . Our method also maintains consistent performance across different resolutions , a major stated advantage of FNOs over previous methods . • Protein folding : on the task of predicting residue distances in a polypeptide chain—a key component of the protein folding problem—we substitute XD-operations into vanilla ResNets and achieve lower error than cyclically-dilated ResNets adapted specifically for this setting [ 1 ] . Furthermore , our ResNet-34 XD outperforms the reported error of the much deeper Dilated ResNet-258 . • Music modeling : on two next-note prediction tasks , we show that substituting XD-operations into an undilated CNN outperforms temporal convolutional networks ( TCNs ) —exponentially-dilated 1d CNNs that themselves outperform standard convolutional and recurrent networks [ 5 ] . Code to reproduce these results is available here : https : //github.com/nick11roberts/XD . Software to apply XD-operations can be found here : https : //github.com/mkhodak/relax . Related Work AutoML is a well-studied area , with most work focusing on fairly small hyperparameter spaces [ 7 , 24 ] or on NAS [ 12 ] . Most NAS operation spaces only contain a few operations such as convolutions [ 32 , 33 , 49 , 11 ] , which may not be useful for domains where CNNs are ineffective . Applications of NAS outside vision largely follow the same pattern of combining human-designed operations [ 35 , 43 ] . On the other extreme , AutoML-Zero [ 37 ] demonstrates the possibility of evolving all aspects of ML from scratch . We seek to establish a middle ground with large and domain-agnostic search spaces that still allow the use of well-tested methods , e.g . stochastic gradient descent ( SGD ) . Several papers have generalized the DFT to replace layers in deep nets [ 9 , 3 , 2 , 10 ] in order to speed up or add structure to models while reducing expressivity . In contrast , we can replace convolutions and other layers while increasing expressivity by extending their diagonalization via K-matrices . As discussed in Section 2 , using K-matrices for this directly is inefficient for input dimension > 1 . 2 The Expressive Diagonalization Relaxation . In this section we overview our main contribution : a large , general search space of neural operations . Formally , we view an architecture as a parameterizable object—a mapping from model weights to functions—described by a labeled directed acyclic graph ( DAG ) G ( V , E ) . Each edge in E has the form ( u , v , Op ) , where u , v ∈ V are nodes and Op is an operation that can be parameterized to define some transformation of the representation at node u ; node v aggregates the outputs of its incoming edges into a new representation . For example , the popular ResNet architecture [ 15 ] has many nodes with two incoming edges , one labeled by the convolution operation Conv and one by the identity ( skip-connect ) Id , whose outputs it sums and passes to outgoing edges with the same labels . Each architecture has a source node taking in input data and an output node returning a prediction . Neural architecture search is the problem of automatically selecting an operation for each edge of G to optimize an objective.1 For each edge e ∈ E a NAS algorithm must pick one element of a search space S = { Opa |a ∈ A } of operations specified by architecture parameters a ∈ A to assign to e ; in past work , A usually indexes a small set of operations . As an example , we will refer to a variant2 Sdiscrete of the DARTS search space with parameters Adiscrete = { 1 , . . . , 8 } where each operation is one of Zero , Id , MaxPool3×3 , AvgPool3×3 , Conv3×3 or 5×5 , or DilatedConv3×3,2 or 5×5,2 [ 32 ] . Our main contribution is a novel family of operations that comprise a search space containing almost all these operations , in addition to many others that have been found useful on different types of data . The starting point of our construction of these XD-operations is the simple observation that all the operations Op ∈ Sdiscrete listed above except MaxPool3×3 are linear , i.e . for any model weights w there exists a matrix Aw such that for all inputs x we have Op ( w ) ( x ) = Awx . More specifically , all seven of them return convolutions : to see this note that Zero , Id , and AvgPool3×3 each apply a convolution with filter 01×1 , 11×1 , and 13×3/9 , respectively . This means that most of the operations in the DARTS search space—which is representative of NAS operation spaces in computer vision—share the convolution ’ s diagonalization by the discrete Fourier transform ( DFT ) . Formally , if Aw ∈ Rn 2×n2 is the matrix representing a 2d convolution with filter w ∈ Rk of kernel size k ∈ [ n ] 2 , then for any 2d input x ∈ Rn2 we have Conv ( w ) ( x ) = Awx = F −1 diag ( Fw ) Fx ( 2 ) Here [ n ] = { 1 , . . . , n } , diag ( z ) denotes the diagonal matrix with entries z , w ∈ Rn2 is an appropriate zero-padding of w ∈ Rk , and F ∈ Cn2×n2 is the 2d DFT ( a Kronecker product of two 1d DFTs ) . This diagonalization explicates both the computational and representational efficiency of the DARTS operations , as the DFT and its inverse can be applied in time O ( n log n ) and stored with O ( n log n ) bits . It also suggests a natural way to dramatically expand the operation space while preserving these efficiencies : just replace matrices F and F−1 in ( 2 ) by any one of a general family of efficient matrices . Doing so yields the single-channel version of our expressive diagonalization ( XD ) operations : XD1α ( w ) ( x ) = Real ( Kdiag ( Lw ) Mx ) ( 3 ) Here architecture parameter α = ( K , L , M ) sets the matrices replacing F and F−1 in Equation 2 . 1It is often defined as selecting both operations and a graph topology [ 50 ] , but if the set of operations contains the zero-operation Zero then the former subsumes the latter . 2For memory-efficiency , all convolutions in the original DARTS search space are separable [ 32 ] . The main remaining question is the family of efficient matrices to use , i.e . the domain of the architecture parameters K , L , and M. For this we turn to the Kaleidoscope matrices , or K-matrices [ 10 ] , which generalize F and F−1 to include all computationally efficient linear transforms with short description length , including important examples such as sparse matrices and permutations . To obtain this general family , K-matrices allow the DFT ’ s butterfly factors—matrices whose products yield its efficient implementation—to take on different values . While a detailed construction of K-matrices can be found in the original paper , we need only the following useful properties : they are as ( asymptotically ) efficient to apply as DFTs , are differentiable and can thus be updated using gradient-based methods , and can be composed ( made “ deeper ” ) to make more expressive K-matrices . Specifying that K , L , and M in Equation 3 are K-matrices largely completes our core contribution : a new search space SXD of XD-operations with K-matrix architecture parameters . We give a full multi-channel formalization in N dimensions , as well as an overview of its expressivity , in Section 3 . First , we note some key aspects of this new search space : • Complexity : XD1α ( w ) requires three K-matrices and O ( 1 ) filter weights to represent , i.e . description length O ( n log n ) ; this is larger than a regular convolution ( which has no architecture parameters ) but is not quadratic in the input size like a linear layer . Applying XD1α requires multiplication by three K-matrices , yielding a theoretical per-channel time complexity of O ( n log n ) , matching the efficiency of convolutions . However , as XD-operations strictly generalize convolutions they are more expensive to apply in-practice ; we detail these costs both in the application sections and as appendix table , and we view improving upon them as an important future direction . • Initialization : a crucial advantage of XD-operations is that we can initialize or warm-start search using operations with known constructions . In particular , since we can recover convolutions ( 2 ) by setting architecture parameters K = F−1 , L = F , and M = F in Equation 3 , we can always start search with any CNN backbone . We use this extensively in experiments . • K-matrices : as they contain all efficient linear transforms , K-matrices can represent all functions returned by XD-operations , including convolutions . However , for input dimension and filter size > 1 the only known way is to apply K-matrices directly to flattened inputs x ∈ RnN , yielding much worse description length O ( nN log n ) . In contrast , as detailed in Section 3 , our diagonalization approach uses Kronecker products to apply DFTs to each dimension separately , yielding description length O ( n log n ) . It is thus the first ( and in some sense , “ right ” ) method to use such matrices to replace convolutions . Furthermore , diagonalization allows us to separate model weights w from architecture parameters α , letting the former vary across channels while fixing the latter . Finally , we address the fact that the architecture parameters of SXD are continuous , not discrete , contrasting with much of the NAS literature . This can be viewed as a natural extension of the weight-sharing paradigm [ 36 ] , in which continuous relaxation enables updating architecture parameters with gradient methods . For example , many algorithms traverse the relaxed DARTS search space S̃discrete = { ∑8 i=1 λiOpi |λi ≥ 0 , ∑8 i=1 λi = 1 } , defined via DARTS operations Opi ∈ Sdiscrete and architecture parameters λi in the 8-simplex ; most search spaces then require discretizing after search via a rounding procedure that maps from the simplex to Adiscrete . Note that the fully continuous nature of XD-operations means that we will only evaluate the final network returned by search . In particular , while some weight-sharing papers also report the correlation between true architecture performance and that indicated by the shared weights [ 46 ] , there is no obvious way to define a ranking or sampling distribution over XD-operations in order to do so . This also means that our final architecture will not be more efficient than the supernet , unlike other weight-sharing methods that do discretize .

Summary: The paper starts by introducing Kaleidoscope-matrices (K-matrices). The core idea here is:  1. Most of the interesting operations in NAS spaces are linear and can be expressed as a matrix multiplication $A_W x$ 2. The operations can therefore be diagonalized by a DFT 3. We can factorize the DFT matrices into “butterfly matrices” (such that, for a factorization of factor $N$, we can write $F =$ the product of the $N$ butterfly matrices and a permutation matrix).  4. By allowing this factor $N$ to vary, the matrix multiplication $A_W x$ can be generalized to express a wider family of operations   The core contribution of this paper appears to be using K-matrices in a NAS setting (which involves a reasonable amount of work): - The architecture parameters are continuous  - K-matrices can’t express parameter-free ops - Finding/evaluating new operations in a NAS setting  Their NAS process then looks like: 1. Take in a backbone and generalize it to a supernet  2. Simultaneously train the architecture parameters and the weights  3. Extract a net + weights for inference  The paper then explores a few unusual tasks and shows that this generalization process is worthwhile.