DNArch: Learning Convolutional Neural Architectures by Backpropagation

David W. Romero* 1 Neil Zeghidour 2

Abstract
We present Differentiable Neural Architectures(DNArch), a method that learns the weights andthe architecture of CNNs jointly by backpropagation. DNArch enables learning (i) the size of
convolutional kernels, (ii) the width of all layers,(iii) the position and value of downsampling layers, and (iv) the depth of the network. DNArch
treats neural architectures as continuous entities
and uses learnable differentiable masks to control
their size. Unlike existing methods, DNArch is
not limited to a (small) predefined set of possible
components, but instead it is able to discover CNN
architectures across all feasible combinations ofkernel sizes, widths, depths and downsampling.
Empirically, DNArch finds effective architectures
for classification and dense prediction tasks onsequential and image data. By adding a loss termthat controls the network complexity, DNArch
constrains its search to architectures that respect apredefined computational budget during training.

Figure 1. DNArch views neural architectures as if they were defined in a continuous multidimensional space, and uses learnablemasks to learn their length by backpropagation. This example
shows how DNArch learns the width of a layer by applying a differentiable mask m(⋅ ; θ) to the channel dimension of the input.
Different values of θ lead to a different number of channels.

nents via backpropagation, it requires (i) defining a (small)set of possible components beforehand, (ii) computing andkeeping their responses in memory during training, and (iii)
retraining the found architecture from scratch to remove theeffect of other components in the output.
In this paper, we introduce Differentiable Neural Architectures (DNArch), a method that jointly learns the weights
and the entire architecture of a CNN by backpropagation.
Specifically, DNArch learns the weights as well as (i) thesize of convolutional kernels at each layer, (ii) the numberof channels at each layer, (iii) the position and resolutionof downsampling layers, and (iv) the depth of the network.
To this end, DNArch treats neural architectures as entities
defined in a multidimensional continuous space with dimensions corresponding to network attributes, e.g., depth, width,etc., and uses learnable differentiable masks along each dimension to control their length (Fig. 1). Unlike DARTS
methods (Liu et al., 2018; Shen et al., 2022), DNArch does
not require a predefined set of components to choose from,but instead is able to explore among all feasible values, e.g.,
all kernel sizes between 1×1 and N×N for a N×N image.
This is a result of the truly continuous nature of DNArch,
which, unlike DARTS, does not require multiple instantiations of the same layer for different parameter values (Fig. 2).
Instead, DNArch explores the parameter space by modifying
the learnable parameters of the differentiable masks (Fig. 3),
making it a much more scalable NAS method. Moreover,since both the architecture and the weights are optimized ina single run, no retraining is needed after training.

1. Introduction
Tailoring Convolutional Neural Networks (CNNs) (LeCunet al., 1998) to novel tasks and datasets requires substantial
human intervention and cross-validation to find a good architecture, e.g. appropriate kernel sizes, width, depth, etc.
This has motivated exploring the space of architectures in an
automatic fashion, by developing architecture search algorithms. Although these methods can find good architectures,
they must solve an expensive discrete optimization problem that involves training and evaluating candidate architectures in each iteration. Differentiable Architecture Search(DARTS) (Liu et al., 2018) addresses this issue by allowing
the network to consider a set of predefined possible components in parallel, e.g., convolutions with kernels of size 3×3,
5×5, 7×7, and adjusting their contribution using learnableweights (Fig. 2). Although DARTS is able to select compo*

Work done during an internship at Google Research. 1 Vrije
Universiteit Amsterdam. 2 Google Research, Paris, France. Correspondence to: David W. Romero <d.w.romeroguzman@vu.nl>.

2. Differentiable masking
Consider a function f ∶ [a, b] → R, which we want to be
non-zero only in a subset [c, d] ⊆ [a, b]. To this end, we
can multiply f with a mask m whose values are non-zero

Published at the Differentiable Almost Everything Workshop of the
40 th International Conference on Machine Learning, Honolulu,
Hawaii, USA. July 2023. Copyright 2023 by the author(s).

1DNArch: Learning Convolutional Neural Architectures by Backpropagation

Figure 2. DARTS learns the size of convolutional kernels using
backpropagation to select among predefined options.
Figure 4. Gaussian and sigmoid masks.

5 7 5 7 5 3

Figure 3. DNArch learns the size of convolutional kernels by modifying the parameters of the differentiable mask m(⋅ ; θ). Different
θ values lead to different sizes.

only on [c, d], e.g., a rectangular mask Π[c,d] (x)=1[c,d] .
However, as the gradient of Π[c,d] (x) is either zero or nondefined, it is not possible to learn the interval in which it isnon-zero by backpropagation. To overcome this limitation,
we can use a parametric differentiable mask m(⋅ ; θ) whose
interval of non-zero values is defined by its parameters θ. As
the mask is differentiable w.r.t. its parameters θ, we can now
use backprop to learn the interval on which it is non-zero.

Figure 5. The Continuous CNN architecture (Romero et al., 2022).

3.1. Learning the size of convolutional kernels
Introduced in Romero et al. (2021a), differentiable masking
can be combined with CKConvs (Romero et al., 2021b) tolearn the size of convolutional kernels by backpropagation
by parameterizing convolutional kernels ψ as the product ofa small neural network MLPψ and a learnable differentiablemask m(⋅ ; θ), i.e., ψ(ci )=MLPψ (ci ) ⋅ m(ci ; θ) (Fig. 6).
Note that, we can construct the kernel only for non-zero values of the mask by following the method given in Appx. A.3.

Here, we consider two types of masks: a Gaussian mask
mgauss parameterized by its mean and variance θ={µ, σ 2 },and a Sigmoid mask msigm parameterized by its offset andits temperature θ={µ, τ } (Eqs. 3, 4). Tm is a predefined
threshold below which the mask is zero (Fig. 4). Additional
details can be found in Appx. A.

3.2. Learning downsampling layers
We can use differentiable masking to learn downsampling byapplying a differentiable mask on the Fourier domain. The
Fourier transform F represents a function f ∶ RD → R in
terms of its spectrum f˜ ∶ RD → C, which map frequencies ω
to the amount of that frequency in the input f˜(ω). A useful
identity is that cropping high frequencies in the Fourier
domain equals downsampling in the spatial domain.

Materializing parameters only for non-zero mask values.
Parts of the differentiable masks will map to zero basedon the value of the parameters θ. Hence, it would be awaste of compute and memory to materialize the mask andthe corresponding network parameters, e.g., channels ch ∈[10, N] in Fig. 1, to zero them out next. Fortunately, we cantake advantage of the invertible form of the Gaussian andSigmoid masks to materialize parameters only for values forwhich the mask is non-zero (Appx. A.3). Thus significantly
reducing the compute and memory costs of DNArch.

To learn downsampling, we use a learnable sigmoid mask
msigm to perform a learnable low-pass filtering on the input
by multiplying the spectrum of the input with the mask
msigm . By doing so, all frequencies above the mask’s cutoff
frequency ωmax =Tm becomes zero (Fig. 7). An importantconsequence of low-pass filtering is that as the spectrum of
the signal becomes zero above ωmax , the low-passed signal
can be represented at a lower resolution determined by ωmax .
With F, F−1 be the Fourier and inverse Fourier transform,
crop>ωmax be an operator that crops values above ωmax , andfdown represent the downsampled signal f , we have that:

3. Learning neural architectures by backprop
Most DNArch components, e.g., the learning of the network’s width and depth, are not limited to convolutionalarchitectures. However, here we aim to show how DNArch
can be used to learn as many components of a neural architecture as possible. To that end, we take a general-purpose
convolutional architecture: the CCNN (Knigge et al., 2023),
and make all its architectural components learnable. CCNNs
are an ideal base network for DNArch due to their ability tomodel global context on inputs of any resolution, length anddimensionality. This (i) prevents the formation of poor architectures due to insufficient receptive fields, and (ii) allows
DNArch to be used on tasks on data of arbitrary length anddimensionality without changing the base network (neededin existing methods. Tu et al. (2022) shows many examples).

fdown =F−1 [crop>ωmax (F[f ] ⋅ msigm ( ⋅; θ))] .

(1)

Unlike regular downsampling, e.g., max-pooling, spectral
downsampling (Rippel et al., 2015) considers the spectral
content of the input during downsampling, and thus prevents
aliasing (Fig. 7, middle down), which has negative effects
for learning (Zhang, 2019; Vasconcelos et al., 2021). Moreinformation w.r.t. the positioning of learnable downsampling and its use for dense predictions is given in Appx. B.
2

DNArch: Learning Convolutional Neural Architectures by Backpropagation

Figure 6. Learning kernel sizes with differentiable masking.

Figure 8. (a) Positioning of masks for the learning of the width. (b)Learning the network’s depth with differential masking.

Lobj and the complexity loss Lcomp weighted by a factor λ:
L = Lobj + λ Lcomp .

(2)

By minimizing this loss, DNArch is encouraged to findarchitectures that meet the desired computational budget
while still achieving good performance on the end task.

4. Experiments

Figure 7. Learning downsampling with differentiable masking.

A detailed description of the experimental setup and thenomenclature used can be found in Appx. E.1. We evaluate
DNArch on sequential and image datasets for classificationand dense prediction tasks. On 1D, we use the Long Range
Arena (LRA) benchmark (Tay et al., 2020). On 2D, weperform image classification on the CIFAR10/100 datasets(Krizhevsky et al., 2009) and report results on two dense prediction tasks from the NAS-Bench-360 benchmark (Tu et al.,2022): DarcyFlow and Cosmic. A detailed description of
the datasets used can be found in Appx. D.

3.3. Learning the width of the network
To learn the width of a layer we apply a differentiable maskm(⋅ ; θ) along the channel dimension of its feature representations (Fig. 1). To learn the width of all layers in the
network, we apply differentiable masks with independent
learnable parameters along the channel dimensions of all the
network components that change the network’s width, i.e.,
all Conv and PWLinear layers. That is, we learn an independent mask for the input (Nin ), the middle (Nmid ) and theoutput (Nout ) channels of each residual block in the network(Fig. 8). Layers that do not change the network’s width, e.g.,
GELU, have their width determined by the previous mask.

DNArch without computational constraints. First, we use
DNArch to improve the expressiveness and computational
efficiency of a CCNN4,140 . We start using DNArch to learn
the size of all convolutional kernels, and then we learn boththe kernel sizes and downsampling layers to simultaneously
improve the expressiveness and the computational efficiencyof the CCNN4,140 . Note that the learning is solely driven by
the objective loss Lobj , i.e., Lcomp is not used.

3.4. Learning the depth of the network
To learn the network’s depth we view the number of residual
blocks as a continuous axis with values [1, 2, .., D] corresponding to the index of each block, and use a differentiablemask m(⋅, θ) along this axis to dynamically mask out blocks
based on the value of the mask parameters θ (Fig. 8b). To
ensure that information flows from the input to the output of
the network regardless of the value of the mask parameters,
we only apply the mask on the residual branch

Results. We observe that using DNArch to learn kernel
sizes consistently improves the accuracy of the base architecture (DNArchK models in Tabs. 1-3). Interestingly, foundDNArch architectures perform on par, and even surpass, architectures specifically designed for each tasks, e.g., S4 (Guet al., 2021) for sequences and NFOs (Li et al., 2020) for
2D PDEs, with a remarkably lower number of trainable parameters. In contrast to DARTS methods, e.g., DASH (Shenet al., 2022), DNArch can be applied across all tasks without
the need to manually change the base architecture. When additionally learning downsampling, we observe that DNArchfinds high-performant architectures with improved computational efficiency (DNArchK,R models). Interestingly, foundmodels often exhibit slight accuracy improvements.

3.5. Putting it all together
By using the techniques described in Sections 3.1-3.4,
DNArch uses backpropagation to learn the weights, thesize of convolutional kernels at each layer, the number ofchannels at each layer, the position and resolution of downsampling layers, and the depth of a convolutional network.
Learning architectures under computational constraints.
As outlined in Appx. C, we can ensure that the architectures
searched by DNArch respect a predefined computational
complexity by including a regularization term Lcomp thatreflects the complexity of the current candidate architecture
based on its mask parameters. To this end, we define theoptimization loss L as the sum of the task objective loss

DNArch with computational constraints. Next, we utilize DNArch to learn entire neural architectures that respect a predefined computational budget. We start with base3

DNArch: Learning Convolutional Neural Architectures by Backpropagation

Table 1. Performance on the LRA benchmark. × denotes random guessing. Highest per-section scores are in bold and the overall best scores
are underlined. For DNArch, values in parenthesis indicate the computational cost of the architecture relative to the target complexity.
M ODELL IST O PS

T EXT

R ETRIEVALI MAGEPATHFINDER

PATH -X

AVG .

Transformer (Vaswani et al., 2017)
Reformer (Kitaev et al., 2020)
Performer (Choromanski et al., 2020)
BigBird (Zaheer et al., 2020)

36.37
37.27
18.01
36.05

64.27
56.10
65.40
64.02

57.46
53.40
53.82
59.29

42.44
38.07
42.77
40.83

71.40
68.50
77.05
74.87

×
×
×
×

53.66
50.56
51.18
54.17

Mega (O(L2 )) (Ma et al., 2022)
Mega-chunk (O(L)) (Ma et al., 2022)

63.14
58.76

90.43
90.19

91.25
90.97

90.44
85.80

96.01
94.41

97.98
93.81

88.21
85.66

S4D (Gu et al., 2022)
S4 (Gu et al., 2021)
S5 (Smith et al., 2022)

60.47
59.60
61.50

86.18
86.82
89.31

89.46
90.90
91.40

88.19
88.65
88.00

93.06
94.20
95.33

91.95
96.35
98.58

84.89
86.09
87.35

FNet (Lee-Thorp et al., 2021)
Luna-256 (Ma et al., 2021)
CCNN4,140 (Romero et al., 2022)

35.33
37.25
44.85

65.11
64.57
83.59

59.61
79.29
×

38.67
47.38
87.62

77.80
77.72
91.36

×
×
×

54.42
59.37
76.86

CCNN4,140 (Global Kernels)
DNArchK (CCNN4,140 )DNArchK,R (CCNN4,140 )DNArchK,R,W,D (CCNN4,140 )

55.65
59.90
60.15(0.80×)60.55(1.01×)

87.80
88.28
88.50(0.75×)89.03(1.00×)

90.55
90.66
91.08(0.78×)91.22(1.02×)

85.51
86.07
86.55(0.82×)87.20(1.02×)

94.26
93.46
94.05(0.89×)
94.95(1.00×)

91.15
89.93
91.15(0.82×)
91.71(1.01×)

84.15
84.72
85.25
85.78

Table 2. Performance on Dense Tasks of NAS-Bench-360.
M ODELDARCY F LOW
rel. l 2 loss

Table 3. Performance on Image Classification Datasets.

C OSMIC
1 - AUROC

Expert*

0.008

0.13

WRN (Zagoruyko & Komodakis, 2016)
DenseNAS (Fang et al., 2020)
DARTS (Liu et al., 2018)
Auto-DL (Liu et al., 2019)DASH (Shen et al., 2022)

0.073
0.100
0.026
0.049
0.060

0.24
0.38
0.229
0.495
0.190

CCNN4,140 (Global Kernels)
DNArchK (CCNN4,140 )DNArchK,R (CCNN4,140 )DNArchK,R,W,D (CCNN4,140 )0.002989
0.002970
0.002929(0.79×)
0.002285(1.01×)

0.059
0.058
0.056(0.82×)
0.055(1.01×)

M ODELCIFAR10CIFAR100

WRN (Zagoruyko & Komodakis, 2016)
DenseNAS (Fang et al., 2020)
DARTS (Liu et al., 2018)
DASH (Shen et al., 2022)

-

76.65
74.51
75.98
75.63

CCNN4,140 (Global Kernels)
DNArchK (CCNN4,140 )DNArchK,R (CCNN4,140 )DNArchK,R,W,D (CCNN4,140 )

90.52
92.51
92.77(0.82×)
93.47(1.01×)

64.72
69.01
68.96(0.85×)
72.98(1.05×)

CCNN6,380 (Global Kernels)
DNArchK,R,W,D (CCNN6,380 )94.18
95.03(1.00×)

72.29
76.37(1.02×)

plexity Ctarget during the whole training. This indicates
that: (i) DNArch only searches among architectures close to
the target computational complexity, and that (ii) the computational overhead of DNArch is negligible. As a result,the cost of using DNArch on top of a CCNN is comparable
to the cost of training the base CCNN network.

CCNN6,380 (Global Kernels)
0.004521
0.059
DNArchK,R,W,D (CCNN6,380 )0.001763(1.00×) 0.048(1.00×)∗
FNO (Li et al., 2020) and deepCR (Zhang & Bloom, 2020).

CCNN4,140 and CCNN6,380 networks, and allow DNArch
to learn their width, depth, kernel sizes and downsampling.
We define the target complexity Lcomp as the complexity ofthe base CCNN networks. In other words, we use DNArch
to find better convolutional architectures of computational
complexity roughly equal to that of the base networks.

Architectures found by DNArch. We observe that foundarchitectures are very diverse, even within each architecture (Tabs. 6-8). For instance, some residual blocks havea bottleneck structure, some an expanded structure, andothers have monotonically decreasing or increasing widths.
Interestingly, the resolution of found architectures for classification, e.g., Text, often follow the style of U-Nets, andnot the monotonically decreasing pattern commonly used in
handcrafted networks. On dense tasks, found architectures
resemble U-Nets, and even concatenated U-Nets, e.g., the1.5× U-Net architectures found for the Cosmic task.

Results. Our results (DNArchK,R,W,D models) show that
DNArch finds neural architectures with higher accuracy
than the base CCNN networks but with the same complexity. In addition, we observe that learning more architecturalcomponents consistently leads to better results, thus supporting the claim that gradient-steered architectures can be
more beneficial than handcrafted ones. Furthermore, using base architectures with larger complexity and capacity( CCNN6,380 vs. CCNN4,140 ) consistently leads to better
results. This result is encouraging for the use of DNArch tolarge architectures, e.g., LLMs (Brown et al., 2020).

5. Conclusion
We presented Differentiable Neural Architectures (DNArch)a method that jointly learns the weights and the architecture of CNNs by backpropagation. DNArch finds effective
architectures across a broad set of tasks, and can include
additional loss terms to encourage the discovery of neuralarchitectures with desirable properties. Appx. F, G discuss
the limitations and promising future directions for DNArch.

Computational complexity of DNArch. To assess the applicability of DNArch, we also analyze its computational
overhead. To this end, we plot the evolution of the relative complexity (Ccurr /Ctarget ) during training (Fig. 9).
Interestingly, we observe that the theoretical complexity ofcandidate architectures Ccurr stays close to the target com4

DNArch: Learning Convolutional Neural Architectures by Backpropagation

Referencestask-specific to a general purpose CNN. In The EleventhInternational Conference on Learning Representations,2023. URL https://openreview.net/forum?
id=ZW5aK4yCRqU.

Bengio, Y., Léonard, N., and Courville, A. Estimating or
propagating gradients through stochastic neurons for conditional computation. arXiv preprint arXiv:1308.3432,
2013.

Krizhevsky, A., Hinton, G., et al. Learning multiple layersof features from tiny images. 2009.

Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D.,
Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G.,
Askell, A., et al. Language models are few-shot learners.
Advances in neural information processing systems, 33:
1877–1901, 2020.

LeCun, Y., Bottou, L., Bengio, Y., and Haffner, P. Gradientbased learning applied to document recognition. Proceedings of the IEEE, 86(11):2278–2324, 1998.
Lee-Thorp, J., Ainslie, J., Eckstein, I., and Ontanon, S. Fnet:
Mixing tokens with fourier transforms. arXiv preprint
arXiv:2105.03824, 2021.

Choromanski, K., Likhosherstov, V., Dohan, D., Song, X.,
Gane, A., Sarlos, T., Hawkins, P., Davis, J., Mohiuddin,A., Kaiser, L., et al. Rethinking attention with performers.
arXiv preprint arXiv:2009.14794, 2020.

Li, Z., Kovachki, N., Azizzadenesheli, K., Liu, B., Bhattacharya, K., Stuart, A., and Anandkumar, A. Fourier
neural operator for parametric partial differential equations. arXiv preprint arXiv:2010.08895, 2020.

Chowdhery, A., Narang, S., Devlin, J., Bosma, M., Mishra,G., Roberts, A., Barham, P., Chung, H. W., Sutton, C.,
Gehrmann, S., et al. Palm: Scaling language modelingwith pathways. arXiv preprint arXiv:2204.02311, 2022.

Linsley, D., Kim, J., Veerabadran, V., Windolf, C., and Serre,
T. Learning long-range spatial dependencies with horizontal gated recurrent units. Advances in neural information
processing systems, 31, 2018.

Fang, J., Sun, Y., Zhang, Q., Li, Y., Liu, W., and Wang, X.
Densely connected search space for more flexible neuralarchitecture search. In Proceedings of the IEEE/CVF
conference on computer vision and pattern recognition,
pp. 10628–10637, 2020.

Liu, C., Chen, L.-C., Schroff, F., Adam, H., Hua, W., Yuille,
A. L., and Fei-Fei, L. Auto-deeplab: Hierarchical neuralarchitecture search for semantic image segmentation. In
Proceedings of the IEEE/CVF conference on computer
vision and pattern recognition, pp. 82–92, 2019.

Ghodrati, A., Bejnordi, B. E., and Habibian, A. Frameexit:
Conditional early exiting for efficient video recognition.
In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 15608–15618,
2021.

Liu, H., Simonyan, K., and Yang, Y. Darts: Differentiablearchitecture search. arXiv preprint arXiv:1806.09055,
2018.

Gu, A., Goel, K., and Ré, C. Efficiently modeling longsequences with structured state spaces. arXiv preprint
arXiv:2111.00396, 2021.

Loshchilov, I. and Hutter, F. Sgdr: Stochastic gradient descent with warm restarts. arXiv preprint
arXiv:1608.03983, 2016.

Gu, A., Gupta, A., Goel, K., and Ré, C. On the parameterization and initialization of diagonal state space models.
arXiv preprint arXiv:2206.11893, 2022.

Loshchilov, I. and Hutter, F. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05101, 2017.

Hendrycks, D. and Gimpel, K. Gaussian error linear units(gelus). arXiv preprint arXiv:1606.08415, 2016.

Ma, X., Kong, X., Wang, S., Zhou, C., May, J., Ma, H., andZettlemoyer, L. Luna: Linear unified nested attention.
Advances in Neural Information Processing Systems, 34:
2441–2453, 2021.

Ioffe, S. and Szegedy, C. Batch normalization: Accelerating
deep network training by reducing internal covariate shift.
In International conference on machine learning, pp. 448–
456. pmlr, 2015.

Ma, X., Zhou, C., Kong, X., He, J., Gui, L., Neubig, G., May,J., and Zettlemoyer, L. Mega: moving average equipped
gated attention. arXiv preprint arXiv:2209.10655, 2022.

Kitaev, N., Kaiser, Ł., and Levskaya, A. Reformer: The
efficient transformer. arXiv preprint arXiv:2001.04451,
2020.

Maas, A., Daly, R. E., Pham, P. T., Huang, D., Ng, A. Y., andPotts, C. Learning word vectors for sentiment analysis.
In Proceedings of the 49th annual meeting of the association for computational linguistics: Human language
technologies, pp. 142–150, 2011.

Knigge, D. M., Romero, D. W., Gu, A., Gavves, E., Bekkers,
E. J., Tomczak, J. M., Hoogendoorn, M., and jakob Sonke,
J. Modelling long range dependencies in $n$d: From5

DNArch: Learning Convolutional Neural Architectures by Backpropagation

Nangia, N. and Bowman, S. R. Listops: A diagnostic dataset
for latent tree learning. arXiv preprint arXiv:1804.06028,
2018.

Tancik, M., Srinivasan, P., Mildenhall, B., Fridovich-Keil,S., Raghavan, N., Singhal, U., Ramamoorthi, R., Barron, J., and Ng, R. Fourier features let networks learnhigh frequency functions in low dimensional domains.
Advances in Neural Information Processing Systems, 33:
7537–7547, 2020.

Radev, D. R., Muthukrishnan, P., Qazvinian, V., and AbuJbara, A. The acl anthology network corpus. Language
Resources and Evaluation, 47(4):919–944, 2013.

Tay, Y., Dehghani, M., Abnar, S., Shen, Y., Bahri, D., Pham,P., Rao, J., Yang, L., Ruder, S., and Metzler, D. Long
range arena: A benchmark for efficient transformers.
arXiv preprint arXiv:2011.04006, 2020.

Riad, R., Teboul, O., Grangier, D., and Zeghidour, N.
Learning strides in convolutional neural networks. arXiv
preprint arXiv:2202.01653, 2022.
Rippel, O., Snoek, J., and Adams, R. P. Spectral representations for convolutional neural networks. Advances inneural information processing systems, 28, 2015.

Teerapittayanon, S., McDanel, B., and Kung, H.-T.
Branchynet: Fast inference via early exiting from deepneural networks. In 2016 23rd International Conference
on Pattern Recognition (ICPR), pp. 2464–2469. IEEE,
2016.

Romero, D. W., Bruintjes, R.-J., Tomczak, J. M., Bekkers,
E. J., Hoogendoorn, M., and van Gemert, J. C. Flexconv:
Continuous kernel convolutions with differentiable kernel
sizes. arXiv preprint arXiv:2110.08059, 2021a.

Tu, R., Roberts, N., Khodak, M., Shen, J., Sala, F., andTalwalkar, A. Nas-bench-360: Benchmarking neural architecture search on diverse tasks. In Thirty-sixth Conference on Neural Information Processing Systems Datasets
and Benchmarks Track, 2022.

Romero, D. W., Kuzina, A., Bekkers, E. J., Tomczak,
J. M., and Hoogendoorn, M. Ckconv: Continuous
kernel convolution for sequential data. arXiv preprint
arXiv:2102.02611, 2021b.

Vasconcelos, C., Larochelle, H., Dumoulin, V., Romijnders,R., Le Roux, N., and Goroshin, R. Impact of aliasing
on generalization in deep convolutional networks. In
Proceedings of the IEEE/CVF International Conference
on Computer Vision, pp. 10529–10538, 2021.

Romero, D. W., Knigge, D. M., Gu, A., Bekkers, E. J.,
Gavves, E., Tomczak, J. M., and Hoogendoorn, M. Towards a general purpose cnn for long range dependencies
in nd. arXiv preprint arXiv:2206.03398, 2022.
Ronneberger, O., Fischer, P., and Brox, T. U-net: Convolutional networks for biomedical image segmentation. In International Conference on Medical image computing andcomputer-assisted intervention, pp. 234–241. Springer,
2015.

Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones,
L., Gomez, A. N., Kaiser, Ł., and Polosukhin, I. Attention is all you need. Advances in neural information
processing systems, 30, 2017.
Zagoruyko, S. and Komodakis, N. Wide residual networks.
arXiv preprint arXiv:1605.07146, 2016.

Schuster, T., Fisch, A., Gupta, J., Dehghani, M., Bahri, D.,
Tran, V. Q., Tay, Y., and Metzler, D. Confident adaptive
language modeling. arXiv preprint arXiv:2207.07061,
2022.

Zaheer, M., Guruganesh, G., Dubey, K. A., Ainslie, J., Alberti, C., Ontanon, S., Pham, P., Ravula, A., Wang, Q.,
Yang, L., et al. Big bird: Transformers for longer sequences. Advances in Neural Information Processing
Systems, 33:17283–17297, 2020.

Shen, J., Khodak, M., and Talwalkar, A. Efficient architecture search for diverse tasks. arXiv preprint
arXiv:2204.07554, 2022.

Zhang, K. and Bloom, J. S. deepcr: cosmic ray rejection
with deep learning. The Astrophysical Journal, 889(1):24, 2020.

Sitzmann, V., Martel, J., Bergman, A., Lindell, D., and Wetzstein, G. Implicit neural representations with periodic
activation functions. Advances in Neural Information
Processing Systems, 33:7462–7473, 2020.

Zhang, R. Making convolutional networks shift-invariant
again. In International conference on machine learning,pp. 7324–7334. PMLR, 2019.

Smith, J. T., Warrington, A., and Linderman, S. W. Simplified state space layers for sequence modeling. arXiv
preprint arXiv:2208.04933, 2022.
Srivastava, N., Hinton, G., Krizhevsky, A., Sutskever, I.,
and Salakhutdinov, R. Dropout: a simple way to preventneural networks from overfitting. The journal of machinelearning research, 15(1):1929–1958, 2014.
6

DNArch: Learning Convolutional Neural Architectures by Backpropagation

Appendix

B. Learning Downsampling Layers

A. Differentiable Masking

B.1. Combining learnable downsampling andconvolution

A.1. Definition

The previous method requires mapping inputs to the Fourier
domain and back to learn downsampling. Fortunately, CCNNs as well as most methods that rely on global convolutions, e.g., CKConv (Romero et al., 2021b), S4 (Guet al., 2021), rely on the Fourier convolution theorem:(f ∗ ψ)=F−1 [F[f ] ⋅ F[ψ]] to compute convolutions withlarge kernels efficiently. This means that CCNNs alreadyuse a Fourier and inverse Fourier transforms in each residual block to compute convolutions. Hence, we can avoid
recomputing these steps by placing the learnable downsampling operation within the Fourier convolution. Specifically,
we can simultaneously compute downsampling and convolution by applying the differentiable mask msigm and thecropping operations crop>ωmax before returning from theFourier domain back to the spatial domain. That is:1

Gaussian mask:⎧(x−µ)2⎪⎪m̄ = exp (− 21 σ2 )mgauss (x ; µ, σ 2 ) = ⎨⎪⎪⎩0
Sigmoid mask:⎧1⎪⎪m̄ = 1− 1+exp(−τ (x−µ))msigm (x ; µ, τ ) = ⎨⎪⎪⎩0m̄ ≥ Tmelsem̄ ≥ Tm
else

(3)

(4)

A.2. Multidimensional masks
N-dimensional masks can be constructed by combining N
1D masks, each with their own parameters. For example,
the Gaussian mask used to learn the size of convolutional
kernels in Fig. 3 is constructed as:(f ∗ψ)down= F−1 [crop>ωmax (msigm (⋅ ; θ) ⋅ F[f ] ⋅ F[ψ])] . (7)2mgauss (x, y; {{µX , µY }, {σX, σY2 }})

2}) ⋅ mgauss (y ; {µY , σY2 })= mgauss (x ; {µX , σX

A.3. Materializing parameters only for non-zero mask
valuesB.2. Materializing functions only on the output
resolution

Parts of differentiable masks will map to zero based on
the value of the parameters θ. Therefore, it would be awaste of compute and memory to materialize the mask –and the corresponding network parameters, e.g., channels
ch ∈ [10, N] in Fig. 1– to zero them out next. Luckily, we
can take advantage of the invertible form of the Gaussian andSigmoid masks to materialize parameters only for values
for which the mask is non-zero. To this end, we find thevalue xTm for which the mask is equal to the threshold Tm ,i.e., xTm =x such that m(x; θ)=Tm , and only materialize themask and the corresponding network parameters for valuesof x for which the value of the mask is greater than Tm . By
inverting the mask equations (Eqs. 3, 4), we obtain xTm as:√(5)±xTm = µ ± −2σ 2 log(Tm ), andNote that Eq. 7 computes the convolution on the resolution of the input and downsamples next. This incurs in
an unnecessary overhead as the output of the convolution
will be downsampled directly after. A more efficient approach comes from inverting the order of these operations
to compute the convolution at the downsampled resolution.
Luckily, this can be achieved by using the method outlined
in Sec. A.3. Since the cutoff frequency of the mask corresponds to the coordinate at which the mask equals thethreshold, i.e., ωmax =xTm , it can be calculated using Eq. 6.
Next, since the cutoff frequency defines the minimum resolution required to faithfully represent the input, we cansimply downsample the input and convolutional kernel tothat resolution before the convolution to compute it on the
output resolution.

1− 1) ,
xTm = µ − τ1 log ( 1−Tm(6)

B.3. Learning subsampling for dense tasks

for Gaussian and Sigmoid masks, respectively. Consequently, we can make sure that all rendered values will
be used by only materializing the mask and related network
parameters for values of x within the range [−xTm , xTm ] for
Gaussian masks and [xmin , xTm ] for Sigmoid masks, where
xmin depicts the lowest coordinate indexing the mask.

Riad et al. (2022) apply learned downsampling on both theidentity and the residual branches of a residual block tolimit resolution of all representations after a specific residual block. In the context of DNArch, this is undesirable fortwo reasons: First, as we learn the whole network architecture during training, it is not known a priori what resolution
1

We note that the Fourier transform is not strictly necessary tolearn downsampling, e.g., for CNNs with local kernels. Leveraging
the Fourier convolution theorem, equivalent downsampling can beachieved by convolving the input with the inverse Fourier transformof the mask in the spatial domain (see Appx. B.4 for details).

7

DNArch: Learning Convolutional Neural Architectures by Backpropagation

C. Learning convolutional architectures undercomputational constraints

mappings will require at each layer. However, forcing theidentity branch to have the same resolution as the corresponding residual branch restricts all subsequent mappings
to be of maximum that resolution. Secondly, dense prediction tasks, e.g., segmentation, require the learned architecture to produce outputs that share the same resolution as the
input. However, if the identity branch is also downsampled,the output of the network would be of lower resolution evenfor a single level of downsampling in the network. This in
turn, would result in over-smoothed predictions.

We can ensure that the architectures searched by DNArch
respect a predefined computational complexity by including
an additional regularization term Lcomp that reflects thecomplexity of the current candidate architecture based on
its mask parameters. To this end, we define the optimization
loss L as the sum of the task objective loss Lobj and thecomplexity loss Lcomp weighted by a factor λ:
L = Lobj + λ Lcomp .

Based on these observations, we use downsampling only
on the residual branch and upsample features at the end ofeach residual block back to the resolution of the input. This
allows us to (i) have features at same resolution as the input
in the last layer, and (ii) learn U-Net (Ronneberger et al.,2015) like architectures.

By minimizing this loss, DNArch is encouraged to findarchitectures that meet the desired computational budget
while still achieving good performance on the end task.
C.1. Defining the complexity loss Lcomp
The purpose of Lcomp is to use the size of the learned masks
to estimate the total computation needed for a forward passof the network. Its construction is outlined below.

B.4. Learning downsampling in the spatial domain
Differentiable Masking learns downsampling by multiplying the spectrum f˜=F[f ] of a signal f with a differentiablemask m(⋅ ; θ), and cropping the output above the cutoff frequency of the mask ωmax next. However, it is not necessary
to perform this operation in the Fourier domain. Downsampling can also be learned directly in the spatial domain.

Layer-wise complexities. Let Clayer (L, Nin , Nout ) be thenumber of operations required in a given layer with an inputof length L and Nin and Nout input and output channels. To
estimate the number of computations required based on thecurrent size of the masks, we can substitute the lenght ofeach dimension with the size of the corresponding masks:
Clayer (size(mres ), size(mNin ), size(mNout )). As an example, consider a pointwise linear layer. A pointwise linear
layer lin ∶ RNin → RNout takes an input f of length L and
Nin channels and multiplies each element along the spatial dimensions of the input with a matrix of dimensions[Nin , Nout ] to produce an output of the same length, butwith Nout number of channels. The total operations requiredin this layer is given by Clin (f )=L ⋅ Nin ⋅ Nout .

The Fourier convolution theorem states that the spatial convolution is equivalent to a pointwise multiplication in the
Fourier domain. However, this equivalence works in both
directions. That is, we can equivalently say that the pointwise multiplication on the Fourier domain is equal to a
convolution on the spatial domain. Consequently, we canrepresent the pointwise multiplication of the spectrum of
the input F[f ] and the differentiable mask m(⋅ ; θ) as the
convolution of their inverse Fourier transforms. Formally:f˜ ⋅ m(⋅ ; θ) = F [F−1 [f˜] ∗ F−1 [m(⋅ ; θ)]]= F [F−1 [F[f ]] ∗ F−1 [m(⋅ ; θ)]]= F [f ∗ F−1 [m(⋅ ; θ)]](9)

Now, let us use three differentiable masks mres , mNin and
mNout to mask the resolution, input and output channels of
the linear layer. The total number of computations is now
given by:

(8)

Clin,masked = size(mres ) ⋅ size(mNin ) ⋅ size(mNout ).

In other words, we can perform the same operation in thespatial domain by convolution the original input signal fwith the inverse Fourier transform of the mask m(⋅ ; θ).

Since the size of the masks is now involved in the computation of the operations required, we can utilize it as anadditional source of feedback to update the masks by making the function size differentiable with regard to the mask
parameters. The same concept is used to calculate the costof other layers based on the size of the masks. A summaryof these costs can be found in Appx. C.2.

Defining the output resolution. Eq. 8 defines how lowpass filtering can be performed on the spatial domain, but it
does not provide information regarding the final resolutionof the operation. To derive the resolution of the output, we
can simply use Eqs. 10, 11 to analytically derive the sizeof the mask. Once the size of the mask is derived, we cansimply take the downsampled signal –after using Eq. 8–,
and downsample it to match the size of the mask.

Effect of the depth mask. To take into account the effectof the depth mask, we use it to determine the number ofresidual blocks in the network. If the number of operationsof a network with D residual blocks is denoted as Cnet,D ,the complexity of a network with masked depth is given
by Cnet,size(mdepth ) with size(mdepth ) the size of the depthmask.
8DNArch: Learning Convolutional Neural Architectures by Backpropagation

Computing the size of the masks. The size of a mask can
be calculated in a differentiable manner by determining thelength of the mask in continuous space and using that length
to estimate the change in size of the corresponding network
dimension. Specifically, the length at a time t is 2xtTm andxtTm −xmin , for Gaussian and Sigmoid masks, respectively(see Fig. 4). For some initial x0Tm and corresponding initiallength N, the size of a Gaussian and a Sigmoid mask at time
t is respectively:size(mgauss ) =2xtTm

size(msigm ) =xtTm − xmin

2x0Tm

N,

x0Tm − xmin

and

(10)

N.

(11)

Pointwise operations –GELU, DropOut, etc.–:
Cpointwise = L ⋅ Nin
Cpointwise,masked = size(mres ) ⋅ size(mNin )D. Dataset descriptions
D.1. The Long Range Arena benchmark
The Long Range Arena (Tay et al., 2020) consists of six
sequence modelling tasks with sequence lenghts ranging
from 1024 to over 16000. It encompasses modalities andobjectives that require similarity, structural, and visuospatial
reasoning. We follow the data preprocessing steps of Gu
et al. (2021), which we also include here for completeness.

Computational constraints as an additional loss. Let
Ccurr be the current complexity of the network and Ctarget
be the desired target complexity. We define the computational loss Lcomp as the relative l 2 difference between therelative complexity of the current network and the target:l2 (ListOps. An extended version of the dataset presented byNangia & Bowman (2018). The task involves computing
the integer result in the range zero to nine of a mathematical expression represented in prefix notation with brackets,e.g., [MAX29[MIN47]0] → 9. Characters are encoded as
one-hot vectors, with 17 unique values possible (opening
brackets and operators are grouped into a single token). The
sequences are of unequal length. Hence, the end of shorter
sequences is padded with a fixed indicator value to a maximum length of 2048. The task has 10 different classes
representing the possible integer results of the expression. It
consists of 96K training sequences, 2K validation sequences,and 2K test sequences. No data normalization is applied.

2

Ccurr
Ccurr
, 1) = ∥− 1.0∥ .
CtargetCtarget2(12)This form has two advantages over the alternative forml 2 (Ccurr , Ctarget ). It (i) prevents overflow that might occur
when comparing large values –Ccurr and Ctarget may easily
be of order 1e10 –, and (ii) allows for consistent tuning of λfor different tasks and complexities. In the alternative forml 2 (Ccurr , Ctarget ), λ might need to be tuned independently
for different complexity regimes.

Text. Based on the IMDB sentiment analysis dataset presented by Maas et al. (2011), the task is to classify movie
reviews as having a positive or negative sentiment. The
reviews are presented as a sequence of 129 unique integer
tokens padded to a maximum length of 4096. The dataset
contains 25K training sequences and 25K test sequences. No
validation set is provided. No data normalization is applied.

C.2. Computational complexity of masked network
components
In this section, we derive the computational complexity of
all layers used in the CCNN architecture with and without
the use of masks. The calculation of these complexities
follows the same reasoning as the pointwise linear layer
provided as example in the main text.

Retrieval. Based on the ACL Anthology network corpus
presented by Radev et al. (2013), the task is to classify
whether two given textual citations are equivalent. To accomplish this, each citation is separately passed throughan encoder, and passed to a final classifier layer. Denoting
X1 as the encoding for the first document and X2 as theencoding for the second document, four features are createdand concatenated together as:

With L, Nin and Nout the length, number of input channels and number of output channels of a given layer, andsize(mres ), size(mNin ), size(mNout ) the size of the masks
along the corresponding dimensions, the complexity of thelayers used in the CCNN architectures are given by:
Pointwise linear layer:

X = [X1 , X2 , X1 × X2 , X1 − X2 ],Clin (f ) = L ⋅ Nin ⋅ NoutClin,masked = size(mres ) ⋅ size(mNin ) ⋅ size(mNout )which are subsequently passed to a two layered MLP. The
goal of the task is to evaluate how well the network canrepresent the text by evaluating if the two citations are equivalent or not. Characters are encoded into a one-hot vector
with 97 unique values and sequences are padded to a maximum length of 4000. The dataset includes 147.086 training
pairs, 18.090 validation pairs, and 17.437 test pairs. No
normalization is applied.

Fourier convolution:

CFconv = L log (L)CFconv,masked = size(mres ) log (size(mres ))9DNArch: Learning Convolutional Neural Architectures by Backpropagation

Image. The Image task uses 32×32 images of the CIFAR10
dataset (Krizhevsky et al., 2009). It views the images as sequences of length 1024 that correspond to a one-dimensionalraster scan of the image. There are a total of 10 classes, 45K
training examples, 5K validation examples and 10K test examples. The RGB pixel values are converted to grayscaleintensities and then normalized to have zero mean and unit
variance across the entire dataset.

is to minimize the Mean Squared Error (MSE) between thepredicted fluid state and the ground truth.
Cosmic: Identifying Cosmic Ray Contamination. The
Cosmic task involves identifying and masking corruption
caused by charged particles collectively referred to as “cosmic rays” on images taken from space-based facilities. It
uses imaging data of local resolved galaxies collected fromthe Gubble Space Telescope. The input is an 128×128 image corresponding to the artifact of cosmic rays, and the output is a 2D grid of the same dimensions predicting whether
each pixel in the input is an artifact of cosmic rays. We
report the false-negative rate of identification results.

PathFinder. Based on the PathFinder challenge introduced
by Linsley et al. (2018), the task presents a 32×32 grayscale
image with an start and an end point depicted as small
circles. The task is to classify whether there is a dashed line(or path) joining the start and end points while presenting the
input as a one-dimension raster scan of the image, alike the
Image task. The dataset includes 160K training examples,
20Kvalidation examples and 20K test examples. The input
data is normalized to be in the range [-1, 1].

E. Experimental details
E.1. General remarks
Experimental setup. We use two CCNNs of different capacity as base networks: a CCNN4,140 –4 blocks, 140 channels, 200K parameters–, and a CCNN6,380 –6 blocks, 380
channels, 2M parameters–, and use DNArch to learn theirarchitectures. To understand the impact of learning each
network component, we also report results learning someand none of the neural architecture components.

Path-X. Path-X is an “extreme” version of the PathFinder
dataset, in which the input images are of size 128×128. As
a result, the input sequences are sixteen times longer with a
total length of 16384. Aside from this difference, the task is
identical to the PathFinder dataset.

Mask configurations. We initialize all the masks to matchthe architecture of the baseline CCNNs at the beginning oftraining. We use a Gaussian mask to learn kernel sizes as
in FlexConv (Romero et al., 2021a), and Sigmoid masks
to learn width, depth and downsampling. All masks use athreshold of Tm =0.1. All kernel masks are centered, i.e.,µ=0, and initialized to either be small or global, i.e., σ ∈[0.0325, 0.5]. Resolution masks are initialized to weightthe highest input frequency by 0.85, and the width and depthmasks are initialized to match the size of the base network’s
architecture. More information on hyperparameters, training
regimes, and experimental settings can be found in Appx. E.

D.2. Image classification datasets
CIFAR10 and CIFAR100.
The CIFAR10 dataset(Krizhevsky et al., 2009) consists of 60K real-world 32×32
RGB images uniformly drawn from 10 classes divided intotraining and test sets of 50K and 10K samples, respectively.
The CIFAR100 dataset (Krizhevsky et al., 2009) is similar to
the CIFAR10 dataset, with the difference that the images are
uniformly drawn from 100 different classes. For validation
purposes, we divide the training dataset of both CIFAR10
and CIFAR100 into training and validation sets of 45K and
5K samples, respectively. Both datsets are normalized tohave zero mean and unit variance across the entire dataset.

Notations. We use DNArch as an operator acting on a basenetwork and specify the learned components with indices
K, R, W, D representing kernel sizes, downsampling, widthand depth. DNArchK (CCNN4,140 ) indicates using DNArch
to learn only the kernel sizes of a CCNN4,140 .

D.3. NAS-Bench-360
NAS-Bench-360 (Tu et al., 2022) is a benchmark suite toevaluate Neural Architecture Search methods beyond image
classification. The benchmark is composed of ten tasks
spanning a diverse array of application domains, datset sizes,
problem dimensionalities, and learning objectives. In thiswork, we consider two tasks from the NAS-Bench-360 suite
which require dense predictions: The DarcyFlow (Li et al.,2020) and Cosmic (Zhang & Bloom, 2020) datasets.

Code. Our code is written in JAX and our experiments are
conducted on TPUs and GPUs. As outlined in the Limitations (Sec. F), JAX and TPU training prevent DNArch fromperforming operations that change the dimensions of arrays
during training. In addition to our JAX implementation, we
will also release a PyTorch implementation of DNArch that
supports these operations. We hope that this implementation
will make DNArch more flexible and accessible, especiallyin scenarios where it is crucial to keep candidate networks
close to the target complexity during the course of training.

DarcyFlow: Solving Partial Differential Equations. DarcyFlow aims to solve Partial Differential Equation (PDE)by using neural networks as a replacement for traditional
solvers. The input for this task is a 85×85 grid specifying
the initial conditions and coordinates of a fluid, and the output is a 2D grid of the same dimensions representing thefluid state at a later time. The ground truth for this task is
the result computed by a traditional solver, and the objective

The Continuous CNN and the CCNN residual block. The
CCNN architecture is shown in Fig. 5. It is composed by10

DNArch: Learning Convolutional Neural Architectures by Backpropagation

the updated parameter values lie outside that range. For
the Gaussian mask, we constraint the minimum admissible value of σ such that the length of the corresponding
dimension never collapses to a value of 1. This is to preventthe corresponding dimension to collapse such that it cangrow afterwards if required. The minimum value depends
on the resolution of the corresponding dimension, e.g., themaximum size of the convolutional kernel, and can be easily
calculated with Eq. 5.

an Encoder, a Decoder, and a number of CCNN residual blocks ResBlock (?). The Encoder is defined as a sequence of [PWLinear, BatchNorm (Ioffe & Szegedy, 2015),
GELU (Hendrycks & Gimpel, 2016)] layers. For tasks dealing with text, we additionally utilize an Embedding layer
mapping each token in the vocabulary to a vector representation of length equal to that used by Gu et al. (2021) (seeAppx. D.1). For dense prediction tasks, the Decoder is a
PWLinear layer, which is preceded by GlobalAvgPooling for global prediction tasks.

Note that the offset value of the Sigmoid mask µ could in
principle assume any value in R. However, if not controlled,µ could become too small and mask all values along aparticular dimension to zero. Similarly, if µ is too large,
the gradient of the mask at all positions would become
very small and it would difficult to update the mask. To
avoid these situations, we define minimum and maximum
values of µ such that for the lowest value, the mask at thelowest position is equal to 0.95, and for the largest value,
the mask at the highest position is equal 0.85. These values
are dependent on the value of the mask temperature τ , andcan be easily calculated with Eq. 6.

Batch Normalization in DNArch. As the architecture
is constantly changing during the search process, we use
batch-specific statistics for batch normalization instead of
the global moving average. This approach was adopted after early experiments showed that using the global moving
average leads to a significant discrepancy in the behavior of
the validation and training curves. Specifically, we observedthat while the training curves were converging to a good solution, the validation curves resembled random predictions.
This issue was resolved by deactivating the global moving
average in Batch Normalization layers.
Continuous convolutional kernels MLPψ . We parameterize our convolutional kernels as a 4-layer MLP with 128
hidden units and a Fourier Encoding (Tancik et al., 2020)of the form γ(x)=[cos(2πω0 Wx]), sin(2πω0 Wx)], withW ∈ RD×128 and ω0 a hyperparameter that acts as a prior on
the frequency content of the kernels (Romero et al., 2021b;
Sitzmann et al., 2020). In contrast to Romero et al. (2021b),
we utilize a single larger MLPψ to generate the kernels of theentire network. This allows the network MLPψ to administrate its capacity across all layers. Using different MLPψ ’s
for each layer as Romero et al. (2021b) is inadequate in
the learnable architectures setting as some layers can be
entirely erased. With our proposed solution, the capacityof the otherwise zeroed-out MLPψ is used to generate thekernels of the remaining layers.

Limiting the size of the mask to the maximum allowable
ranges. As outlined in the Limitations (Sec. F), we must
set a maximum allowable size for the width and depth of
the network on JAX. However, the maximum allowed value
for the parameters of the masks (see previous paragraph)allows both masks to grow beyond the point on which thetheoretical size of the masks is equal to the maximum allowable network size. For instance, for the maximum allowed
parameter values of a Sigmoid mask, the last channel, i.e.,the 280-th channel, would be weighted by a factor of 0.85.
Consequently, the theoretical size of the mask as calculated
by Eq. 11 will be well beyond 280. This value would lead toan unrealistic theoretical computational complexity that surpasses the real computational complexity the CCNN used.

Normalized relative positions. Following Romero et al.
(2021b;a), we normalize the coordinates going into MLPψ tolie in the space [−1, 1]D for D-dimensional kernels.

To overcome this issue, we limit the maximum size of themask calculated by Eq. 11 to be less or equal than the maximum allowable size, e.g., size(m)= min(size(m), 280). It
is important to note, however, that clipping the value of size
directly would stop the gradient flow for parameter values
leading to sizes larger 280. As a result, once the maximum
size is reached, the mask would not be able to contract
anymore. We avoid gradient flow stop by using clipping in
combination with a straight-through estimator (Bengio et al.,2013). As a result, we are able to propagate the gradientacross the clipping operation, and the resulting mask canstill be modified even in the cropping operation is used.

Parameters and hyperparameters of the differentiablemasks. We learn some of the parameters of the masks, andleave the others constant or treat them as a hyperparameter.
Specifically, for Gaussian masks, we only learn their width,
i.e., σ, and fix its mean to zero. For Sigmoid masks, welearn their offset µ and treat their temperature τ as a hyperparameter. For more information regarding the values of τ
used in our parameter tuning step, please refer to Appx. E.2.
Maximum and minimum allowable sizes for the learnable differentiable masks. We define some minimum andmaximum allowable sizes for the mask parameters, andreset them to these values after each training iteration if

E.2. Hyperparameters and training configurations
In this section, we include more information about the foundhyperparameters, the values that were considered during
11

DNArch: Learning Convolutional Neural Architectures by Backpropagation

Table 4. Hyperparameters used for the experiments with the target complexity of a CCNN4,140 .
E POCHS

BATCH
S IZE

L EARNING
R ATE

D ROP O UTW EIGHTD ECAYω0λτresolution

τchannelτdepthL IST O PS
T EXT
R ETRIEVALI MAGEPATH F INDER
PATH -X

50
100
50
210
210
80

50
50
50
50
50
32

0.005
0.02
0.001
0.02
0.005
0.001

0.0
0.2
0.1
0.1
0.0
0.0

0.01
0.01
0.01
0.001
0.001
0.0

27.5K
19.5K
21.5K
12.5K
21.5K
30K

5.0
0.1
0.1
0.1
0.1
0.1

50
50
50
25
50
100

25
25
25
25
25
25

8
8
8
8
8
8

CIFAR10CIFAR100

210
210

50
50

0.01
0.01

0.1
0.0

0.01
0.01

21.5K
6.5K

0.1
5.0

25
25

25
25

8
8

DARCY F LOW
C OSMIC

310
310

8
8

0.02
0.02

0.0
0.3

0.00010.0001

24.5K
5.5K

0.1
0.1

50
100

25
25

8
8

DATASET

Table 5. Hyperparameters used for the experiments with the target complexity of a CCNN6,380 .
DATASETE POCHS

BATCH
S IZE

L EARNING
R ATE

D ROP O UTW EIGHTD ECAYω0λτresolution

τchannelτdepthCIFAR10CIFAR100

210
210

50
50

0.005
0.01

0.0
0.0

0.01
0.01

21.5K
6.5K

0.1
0.1

50
25

50
25

16
8

DARCY F LOW
C OSMIC

310
310

12
4

0.01
0.01

0.2
0.2

0.0
0.01

7.5K
5.5K

0.1
0.1

50
50

25
50

8
8

consider two values τ ∈ [25, 50] which correspond to a minimum size of 10% and 5% of the corresponding dimension,but observe early during tuning that models prefer τ =25.
For the depth dimension, which is much more sparse thanthe channel and resolution dimensions, we consider twovalues τ ∈ [8, 16], which result on a minimum depth of 2
and 1 layers, respectively. We observe early during tuningthat models prefer τ =8.

hyperparameter tuning, and other training settings. The
final hyperparameters used are listed in Table 5.
Optimizer, learning rates and learning rate schedule.
All our models are optimized with AdamW (Loshchilov& Hutter, 2017) in combination with a cosine annealinglearning rate scheduler (Loshchilov & Hutter, 2016), anda linear learning rate warm-up stage of 10 epochs, except
for ListOps, Retrieval and Path-X for which we have awarm-up stage of 5 epochs.

Learning rate. The possible learning rate values exploredin this work are [0.0001, 0.0005, 0.001, 0.005, 0.01, 0.02].

Regularization. We utilize dropout (Srivastava et al., 2014)–as shown in Fig. 5– as well as weight decay during training.

Dropout. The possible dropout values explored in this work
are [0.0, 0.1, 0.2, 0.3].

E.2.1. H YPERPARAMETER TUNING

Weight decay. The possible weight decay values consideredin this work are [0.0, 0.0001, 0.001, 0.01, 0.05].

Frequency prior ω0 . The possible ω0 values explored inthis work are [1, 500, 1500, 2500, ...28500, 29500, 30000].

E.3. Computational cost of DNArch
Relative Computational Complexity

Tuning the value of λ. λ plays the role of controlling theweight of the computational loss Lcomp relative to the task
objective loss Lobj . In this work, we find two settings which
require different values of λ. One, given by the tasks thatconverge to a low prediction values relative to perfection,i.e., ListOps and CIFAR100, and for which the loss Lobj
remains relatively high at the end of training. The othergroup is given by all the other tasks, which converge to highprediction values –many even obtaining a perfect accuracy
on the train set–, and for which Lobj converges to values
close to zero. For the first group, we require a higher valueof λ such that the computational complexity loss Lcomp
remains relevant to the optimization objective. The final
values of λ used are 5.0 and 0.1, respectively.

1.05
1.04
1.03
1.02
1.01
1.00
0.99

0

20

40

Epoch

60

80

100

Figure 9. Relative complexity during the course of training on the
Text task. This behavior is consistent across all tasks.

F. Limitations

Tuning the temperature of the Sigmoid masks τ . For the
resolution mask, we consider three values of τ , [25, 50, 100]
which correspond to a minimum size of 10%, 5% and 2.5%of the corresponding dimension. For the channel mask, we

Training on TPU requires static shapes. We train ourmodels on TPUs, a type of accelerator that requires a staticcomputational graph derived for specific input and network
shapes via the XLA (Accelerated Linear Algebra) compiler.
12

DNArch: Learning Convolutional Neural Architectures by Backpropagation

at first, and progressively encourage it to converge to networks with the desired target complexity. Such a weighting
scheduling of λ could lead DNArch to find better architectures.

As a result, TPUs do not support operations that change
the shapes of arrays during training. This means that on
TPUs, DNArch can only perform masking modifications to
the network during training, i.e., setting certain channels tozero but still computing their outputs. At inference, however, the masks are fixed. Consequently, we can effectively
trim unused values to remove useless computations in a
way that is compatible with XLA. It is important to note that
this limitation is solely an implementation issue caused bynature of TPUs’ hardware and can be avoided by using libraries and hardware that support dynamic computationalgraphs, e.g., PyTorch and GPUs. While our results wereobtained using TPUs, we also provide a PyTorch implementation that avoids this issue, making it more flexible andaccessible, especially in scenarios where one needs to keep
candidate networks close to the target complexity Ctarget
during training.

Training DNArch with additional / multiple constraints.
Here, we only consider computational complexity as a constraint when training with DNArch. However, other properties such as memory efficiency, hardware-awareness androbustness are equally important. Designing regularization
terms that encourage other properties in DNArch as well
as exploring how different properties can be optimized inunison are important directions for further research.

H. Architectures Found by DNArch
Table 6. Architectures found by DNArch on LRA with the target
complexity of a CCNN4 , 140.

DNArch requires instantiating the largest possible architecture. While masking weights through a gradient updateis straightforward, increasing the number of active weights
requires those weights to be instantiated in memory. Thismeans that even with dynamic computational graphs, itis necessary to instantiate the largest possible architecturelearnable by DNArch in memory. To overcome this limitation, we set the maximum kernel size to the size of the input,
and limit the maximum network size along the depth andwidth dimensions to double the number of blocks and channels of the base network. While this trick allows DNArch
to easily shrink and grow representations within that range,
this restricts the potential sizes of optimal architectures andcan restrict the applicability of DNArch to very large models, e.g., LLMs (Brown et al., 2020; Chowdhery et al., 2022),
which can have billions of weights.

R ETRIEVAL

T EXTL IST O PSTASKG. Outlook and future work

PATH F INDERI MAGE

Input-dependent neural architectures. In this work, themask parameters are constant for all inputs within a task. An
alternative approach could use an additional neural network
MLPmask to predict the mask parameters based on context,
e.g., the current input, current task, etc. This would enable
the creation of context-dependent neural architectures suchas early-exit systems (Teerapittayanon et al., 2016; Ghodrati
et al., 2021; Schuster et al., 2022), but where the wholenetwork architecture is context-dependent. Consequently,
resulting architectures would providing finer control of persample / per-modality complexity than existing methods.

PATH -X

Dynamic weighting of Lcomp during training. DNArch
explores architectures with complexity similar to target complexity throughout training. This results from using a constant λ in Eq. 9. Alternatively, one could use a dynamic
value of λ during training to induce a different training behavior. For example, gradually increasing λ would allow
DNArch to explore architectures with larger complexity
13

D EPTHK ERNEL S IZER ESOLUTION

W IDTH[Nin , Nmid , Nout ]

8

266
569
1401
310
213
12
5
24

2048
632
1416
310
213
301
250
502

[150 189 145]
[150 168 168][176 186 162]
[166 175 153]
[154 159 163][168 128 162]
[170 158 153]
[153 171 165]

8

445
691
1420
415
1467
52
101
149

2284
2939
1420
1313
1467
594
932
1036

[180 217 205]
[208 176 153]
[152 152 120]
[120 120 147]
[147 118 135]
[134 173 153]
[150 156 183][180 92 192]

8

2
136
1013
1446
7
1
6
1

1913
2058
2363
2724
2604
2756
3545
3899

[29 33 172][184 174 183][205 171 161]
[188 164 115]
[29 29 163]
[29 35 154]
[71 110 147]
[71 88 137]

8

203
279
219
308
144
8
30
40

1024
1024
486
308
144
125
96
126

[118 155 147]
[146 172 164][173 166 196]
[199 197 196]
[207 197 92][106 29 75][78 28 110]
[104 51 104]

8

195
493
418
318
236
231
8
4

1024
770
507
318
236
231
251
253

[109 140 171][171 168 158]
[144 183 170][173 187 178]
[182 162 160]
[161 121 103]
[105 47 210]
[116 29 188]

5

2484
7204
3669
2323
513

15331
7204
3772
5496
4768

[280 174 157]
[177 280 159]
[167 280 98][123 164 164]
[136 128 195]

DNArch: Learning Convolutional Neural Architectures by Backpropagation

Table 8. Architectures found by DNArch on 2D datasets with thetarget complexity of a CCNN6,380 .
TASKD EPTHK ERNEL S IZE[Y X ]R ESOLUTION[Y X ]W IDTH[Nin , Nmid , Nout ]I MAGE C LASSIFICATION TASKSTASKD EPTHK ERNEL S IZE[Y X ]R ESOLUTION[Y X ]C IFAR 10

Table 7. Architectures found by DNArch on 2D datasets with thetarget complexity of a CCNN4,140 .
W IDTH[Nin , Nmid , Nout ]

12

[4 7]
[9 10][12 6][20 6][10 11][11 11][1 11][1 6][29 4][18 15]
[1 1]
[4 3]

[32 32][32 32][32 32][32 32][23 26][12 22][6 11][5 21][32 4][18 15][4 4][21 15][380 328 384]
[384 371 393]
[392 361 391]
[388 370 421]
[421 417 486]
[496 444 479]
[493 482 304]
[211 78 384]
[363 459 280]
[277 394 67][111 109 361]
[121 374 449]

12

[8 9]
[12 10][11 10][18 12][11 15][19 8][11 10]
[8 8]
[5 7]
[8 4]
[6 4]
[8 9]

[32 32][32 32][32 32][29 21][27 24][25 20][12 23]
[8 9]
[5 8]
[8 4]
[6 4]
[8 9]

[343 275 354]
[351 316 397]
[495 355 420]
[421 498 419]
[432 449 407]
[412 419 413]
[423 454 600]
[709 685 416]
[419 311 446]
[446 433 389]
[386 501 570]
[568 453 655]7

[54 49][43 47][65 69][67 66][85 85][85 85][68 68][54 49][70 72][85 85][85 85][85 85][85 85][85 85][435 428 289]
[499 393 284]
[496 434 281]
[323 412 275]
[319 369 271]
[306 379 258]
[521 435 271]12

[35 32][11 21][44 24][23 41][28 27][21 19][29 24][18 25][57 16][50 11][18 12][22 23][35 33][95 72][128 128][128 128][128 128][128 128][128 128][128 128][128 128][127 77][89 67][97 79][146 236 272]
[170 284 319]
[141 339 388]
[385 407 361]
[351 279 356]
[354 362 310]
[310 351 466]][396 292 183]
[179 210 580]
[273 250 63][347 400 77][171 241 79]

8

[9 7]
[12 8][25 7][9 10][1 13][1 10][5 1][6 5][32 32][32 32][32 20][9 15][5 15][6 11][15 4][11 7][142 139 145]
[145 160 157]
[158 186 182]
[186 208 168][169 177 150]
[151 139 156]
[154 115 110][108 41 166]8

[13 7][6 10][11 8][13 7][8 12]
[1 1]
[5 9]
[8 7]

[32 32][32 32][22 22][16 18][10 12][8 9][8 10][9 9][104 107 116]
[114 134 134]
[139 192 166][173 201 197]
[205 251 51][62 56 157]
[162 175 254]
[280 280 280]

C IFAR 100

C IFAR 100

C IFAR 10

I MAGE C LASSIFICATION TASKSD ENSE TASKS[80 72][22 22][85 85][156 280 107][180 204 78][280 280 50]

6[94 111][2 13][129 129][129 121][90 89][76 74][128 128][20 45][129 129][129 129][129 129][76 74][18 110 23][186 207 139]
[126 265 100]
[78 105 59][57 201 197]
202 145 216]

DARCY
F LOW

3

[43 38][22 22][76 76]C OSMICC OSMIC

DARCY
F LOW

D ENSE TASKS

14