Hyperparameter Optimization through Neural Network Partitioning
Authors: Bruno Kacper Mlodozeniec
Matthias Reisser
Christos Louizos
Venue: ICLR 2023 poster (2023)
Keywords: Hyperparameter optimization
invariances
data-augmentation
marginal likelihood
federated learning

Abstract: Well-tuned hyperparameters are crucial for obtaining good generalization behavior in neural networks. They can enforce appropriate inductive biases, regularize the model and improve performance --- especially in the presence of limited data. In this work, we propose a simple and efficient way for optimizing hyperparameters inspired by the marginal likelihood, an optimization objective that requires no validation data. Our method partitions the training data and a neural network model into $K$ data shards and parameter partitions, respectively. Each partition is associated with and optimized only on specific data shards. Combining these partitions into subnetworks allows us to define the "out-of-training-sample" loss of a subnetwork, i.e., the loss on data shards unseen by the subnetwork, as the objective for hyperparameter optimization. We demonstrate that we can apply this objective to optimize a variety of different hyperparameters in a single training run while being significantly computationally cheaper than alternative methods aiming to optimize the marginal likelihood for neural networks. Lastly, we also focus on optimizing hyperparameters in federated learning, where retraining and cross-validation are particularly challenging.

Published as a conference paper at ICLR 2023
HYPERPARAMETER OPTIMIZATION THROUGH NEURAL NETWORK PARTITIONING
Bruno Mlodozeniec'", Matthias Reisser Christos Louizos
TUniversity of Cambridge, *Qualcomm Al Research bkm280 cam.ac.m.... uk, (mreisser, cloui zos @qqti quw.eecc.cccccccoooommmmo com
ABSTRACT
Well-tuned hyperparameters are crucial for obtaining good generalization behavior in neural networks. They enforce appropriate inductive biases, regularize the model and improve performance - especially in the presence of limited data. In this work, propose a simple and efficient way for optimizing hyperparameters inspired by the marginal likelihood, an optimization objective that requires no validation data. Our method partitions the training data and a neural network model into K data shards and parameter partitions, respectively. Each partition is associated with and optimized only on specific data shards. Combining these partitions into subnetworks allows us to define the sett-of-training-aanne of subnetwork, í.e., the loss on data shards unseen by the subnetwork, as the objective for hyperparameter optimization. We demonstrate that we can apply this objective to optimize a variety of different hyperparameters in a single training run while being significantly computationally cheaper than alternative methods aiming to optimize the marginal likelihood for neural networks. Lastly, we also focus optimizing hyperparameters in federated learning, where retraining and cross-validation are particularly challenging.
loss
INTRODUCTION
Due to their remarkable generalization capabilities, deep neural networks have become the de-facto models for wide range of complex tasks. Combining large models, large-enough datasets, and sufficient computing capabilities enable researchers to train powerful models through gradient descent. Regardless of the data regime, however, the choice of hyperparameters - such neural architecture, data augmentation strategies, regularization, or which optimizer to choose - plays a crucial role in the final model's generalization capabilities. Hyperparameters allow encoding good inductive biases that effectively constrain the models' hypothesis space (e.g., convolutions for vision tasks), speed up learning, or prevent overfitting in the case of limited data. Whereas gradient descent enables the tuning of model parameters, accessing hyperparameter gradients is more complicated.
The traditional and general way optimize hyperparameters operates follows; 1) partition the dataset into training and validation data!, 2) pick a set of hyperparameters and optimize the model on the training data, 3) measure the performance of the model on the validation data and finally 4) use the validation metric as a way to score models or perform search over the space of hyperparameters. This approach inherently requires training multiple models and consequently requires spending resources on models that will be discarded. Furthermore, traditional tuning requires a validation set since optimizing the hyperparameters on the training set alone cannot identify the right inductive biases. A canonical example is data augmentations they are not expected to improve training set performance, but they greatly help with generalization. In the low data regime, defining a validation set that cannot be used for tuning model parameters is undesirable. Picking the right amount of validation data is hyperparameter in itself. The conventional rule of thumb to use ~ 10% of all data can result in significant overfitting, as pointed out by Lorraine et al. (2019). when has sufficiently large number of hyperparameters to tune. Furthermore, a validation set can challenging
Work done while at Qualcomm AI Research. Qualcomm AI Research is an initiative of Qualcomm Technologies, Inc. and/or its subsidiaries. a third partition, the holdout is used estimate the final model performance

---

Published as a conference paper at ICLR 2023
to obtain in many use cases. An example is Federated Learning (FL) (McMahan et al., 2017). which we specifically consider in our experimental section. In FL, each extra training run (for, e.g., a specific hyperparameter setting) comes with additional, non-trivial costs.
Different approaches have been proposed in order to address these challenges. Some schemes optimize hyperparameters during single training run by making the hyperparameters part of the model (e.g., learning dropout rates with concrete dropout (Gal et al., 2017), learning architectures with DARTs (Liu al., 2018) and learning data-augmentttions with schemes as in Benton et al. (2020); van der Wilk et al. (0. In cases where the model does not depend on the hyperparameters directly but only indirectly through their effect on the value of the final parameters (through optimization). schemes for differentiating through the training procedures have been proposed, such as Lorraine et al. (2019). Another way of optimizing hyperparameters without a validation set is through the canonical view on model selection (and hence hyperparameter optimization) through the Bayesian lens; the concept of optimizing the marginal likelihood. For deep neural networks, however, the marginal likelihood is difficult to compute. Prior works have therefore developed various approximations for its use in deep learning models and used those to optimize hyperparameters in deep learning, such as those of data augmentation (Schwöbel et al., 2021; Immer et al., 2022). Still, however, these come at a significant added computational expense and do not scale to larger deep learning problems.
This paper presents novel approach to hyperparameter optimization, inspired by the marginal likelihood, that only requires a single training run and no validation set. Our method is more scalable than previous works that rely marginal likelihood and Laplace approximations (which require computing or inverting Hessian (Immer et al., 2021)) and is broadly applicable to any hierarchical modelling setup.
MARGINAL LIKELIHOOD AND PRIOR WORK
In Bayesian inference, the rules of probability dictate how any unknown, such as parameters w or hyperparameterss d. should be determined given observed data D. Let p(w) be a prior over w and p(D|ut w. v)) be a likelihood for D with v being the hyperparameters. We are then interested in the posterior given the data p( u/ 'D, vb) p(D w. (t)pp(()/(()) is known as the marginal likelihood, as it measures the probability of observing the data given 4/, irrespective of the value of w: p(D())
The denominator term p(D)
p( w. d)
Marginal likelihood has many desirable properties that make it a good criterion for model selection and hyperparameter optimization. It intuitively implements the essence of Occam's Razor principle (MacKay, 2003, $ 28). In the PAC- AAC-Bayesian literature. it has been shown that higher marginal likelihood gives tighter frequentist upper bounds on the generalization performance class (McAllester, 1998; Germain et al., 2016). It also has close links to validavali (see section 2. and can be computed from the training data alone. However, computation of the marginal likelihood in deep learning models is usually prohibitively expensive and many recent works have proposed schemes to approximate the marginal likelihood for differentiable model selection (Lyle et al., 2020; Immer et al., 2021; 2022; Schwöbel et al., 2021).
given model
2.1 "LEARNING SPEED" PERSPECTIVE
Lyle et al. (2020); Fong and Holmes (2020) pointed out the correspondence between "learning speed" and marginal likelihood. Namely, the marginal likelihood of the data D conditioned on some hyperparameters z can be written
log = Σ log Ep(m)Dreeeee 180) [D w. +b 2
Σ Eppirlpis 1.0) (logp(Dp/uu ei]]
(1)
where (D,,..., De) is an arbitrary partitioning of the training dataset into C' shards or ? and DD + is the posterior over parameters of function fw X D, from the input domain X' to the target domain y after seeing data in shards through k. The right-hand side can be interpreted type of cross- validatio in which we fix ordering over the shards and measure the "validation" performance on each shard De using a model trained on the preceding shards D Disk-1-
We use the terms "chunk" and "shard" interchangeably.
2

---

For the sake of scalability and alignment with deep learning practice, we also focus on simple pointwise approximations g: (w) 8(w rZx) to the posteriors p(w Disk: v)). However, in contrast to prior work, we explicitly parametrize the learning procedure such that, at any given Published as a conference paper at ICLR 2023
Alternatively, it can be viewed as the learning speed of (probabilistic) model: i.e., measure of how quickly it learns to perform well on new shards of data after only having been fit to the previous shards (through exact Bayesian updating).
This perspective neatly illustrates why models with higher marginal likelihood can exhibit good inductive biases, e.g., encoded through di w and fw- Namely, such models can be expected to learn faster and generalize better after seeing fewer samples. For example, if the hypothesis space is constrained functions satisfying symmetries present in the data, we need fewer data to identify the correct function (Sokolic et al., 2017; Sannai et al., 2021). We argue that the "learning speed" aspect of marginal likelihood - i.e., measuring how well the model generalizes to new data the training set, having been trained only on the previous data points - is the key property making marginal likelihood a useful tool for selecting hyperparameters.
2.2 TRAINING SPEED FOR HYPERPARAMETEB OPTIMIZATION
Computing the "learning speed", requires samples from the posterior p( w (Dikk, v).. Unfortunately, in deep learning settings, such samples are impractical to obtain; thus, prior works have focused on more scalable alternatives. Lyle et al. (2020) propose to approximate the objective in Eq. by looking at the training speed during standard training a neural network by SGD. Specifically, they define the training speed as the reduction in the training loss after single SGD parameter update, summed over all updates in the first epoch. They argue that, during the first epoch of training, after the neural network parameters, w, have been updated with SGD steps using data from shards Di:k. they can be approximately used in place of the sample from the posterior [D1xk 1)) in Eq. 1. They extend the analogy to training past one epoch and use the training speed estimate for model selection (Ru et al., 2021). As pointed out by the authors, however, the analogy between learning speed and training speed somewhat breaks down after epoch of training. The network parameters have "seen" " every datapoint in the training set after 1 epoch, and hence the connection to measuring the model's generalization capability is weakened.
training iteration, we have access to a model that trained only on a subset of the data Dise. In doing so, we can approximate the objective in Eq. 1, and thus use it to optimize the hyperparameters during the entire training run.
PARTITIONED NEURAL NETWORKS
Our goal is to optimize the objective
CML (D.
Σ Eg 1(00) ((Dx [w, +)]]
kml
wrt. vh which is an approximation to the nower-bound presented in Eq. 1 above. In Appendix A, we show that the left-hand side is also lower-bound on the marginal likelihood under some unobtrusive conditions. As mentioned in Section 2.2, our goal is to propose architecture and a training scheme so that we can easily obtain models trained on only subsets of the data DisR for all k throughout training. We propose that each {gx (w)) lk=1 optimizes subset of the parameters of the neural network, in a manner that allows us to extract "subnetworks" from the main network that have been trained on specific chunks of data. We describe the partitioning scheme below.
Partitioning the parameters Denote the concatenations of the weights of a neural network w E R-N We can define partitioning ((w1...,, Wc), P) of the parameters into C partitions, such that w = P concat (w |,..,,, wc) for a permutation matrix P E {0, NxN. For ease of exposition, we drop the dependence on P. assuming that w is already arranged such that P is identity, P INNN.
Given the partitioning (Wi,... Wc) of the parameters, we then specify C' subnetworks with weights (k) Soucat(wi,-.. ...,, Wk, Wk41,.... wee), where w; are some default ..... (C) such that ws
wz (1)
or if the learning algorithm is heavily biased towards returning hypotheses that satisfy given invariance, e.g., through the use of prior.
3

---

Published as a conference paper at ICLR 2023
values not optimized during training". More specifically, the k-th subnetwork, w retains the first partitions from the weight partitioning and sets the remaining parameters to #k++10- Note that, if each we is only updated on chunks Disk the subnetwork w² is only comprised of weights that have been updated on Di:k- Thus, we can view the parameters of un! (R) as an approximation to q (un). Although, given that subset of the parameters in each w approximation to the true posterior over the weights given Dis it could be, intuitively, a reasonable approximation in function space".
(R) is fixed, this would likely be poor
Partitioned training Having partitioned the dataset into C chunks (Dı,..., Dz), we update each partition by optimising the negative log-likelihood on chunks Di:k using subnetwork w! by computing the following gradients:
Vwn C (PLik :
Σ
Vwl logp (v/=: w. .0).
(e..))DD:
We interleave stochastic gradient updates of each partition of the weights with updating the hyperpa- rameters using LML in Eq. 2:
deLML (2,
E Σ
V logp (u/se, w ,b))
This can be seen as the sum of the out-of-sample losses for each subnetwork wi The scheme is illustrated in Figure 1. For details of how the updates are scheduled in our experiments, see Appendix I. Note that, while we could incorporate the gradient of the first term from Eq. 1 corresponding to
Ego(w)) [logw)/oog/ p(D1 1 (w, })] in Eq.4, we chose to leave it out. Hence, the gradient of Eq. 4 is an estimate that can be viewed as an approximation
the conditional marginal likelihood log p (D2:c (D1, (,,,)) v)).
Conditional marginal likelihood has been shown to have many desirable properties for model selection and, in many cases, can be a better proxy for generalization (Lotfi et al., 2022).
Alternate:
Optimize parameters:
log (D:
(wi, th2, wa)
W.E.t.
Salmet.
log pl (Di2l (Wi, Wi2, 1.),
w.r.t 202
Subnet. 2
log Di:al (wi 222, 22)),
w.r.f. 2073
Subnet 3
Optimize hyper parameters 3/ on:
log pl (Dal (wi, W2 tdg). 1b) + log log (m11, wa, tha) v)
Weights: w (wi, W2, ws)
Subner
Subnet.
Figure Best viewed in colour. Illustration of the partitioning scheme for single hidden layer perceptron with C' = 3 chunks.
This procedure, inspired by the marginal likelihood, has several desirable properties compared to prior work. 1) Our objective is computationally efficient, with computational cost roughly corresponding to evaluating subnetworks on the training set. There is no need to compute nor invert a Hessian with
"e.g., th; could he the value of the weights at initialization, or r.: = O corresponding to pruning those parameters and obtaining proper subnetwork.
"Since a) the mapping from parameters to functions not bijective b) neural networks are highly overparameterised and can be heavily pruned while retaining performance (Frankle and Carbin, 2018), obtaining good fit to subset of the training data with subset of the model parameters should be possible. Furthermore. "scaling laws" indicate that the benefit of having more parameters becomes apparent mostly for larger dataset sizes (Kaplan al., 2020), thus it reasonable for subnetworks fit to more data have more learnable parameters. Optionally with an added negative log-prior regularization log (( (k)).
4

---

Published as a conference paper at ICLR 2023
respect to the weights, as in the Laplace approximation (Immer al., 2021; 2022). 2) Our objective is readily amenable to optimization by stochastic gradient descent; we do not have to iterate over the entire training set to compute single gradient update for the hyperparameters. 3) Compared to the training speed objective (Lyle et al., 2020), in our method, the training of the weights in each subnetwork progresses independently of the data in future chunks. Hence, it can be seen more truthfully measuring the generalization capability of a model using a given set of hyperparameters.
Partitioning Schemes There are several ways in which the neural network weights can be partitioned. In our experiments in Section 5, we partition the weights before beginning training by assigning a fixed proportion of weights in each layer to a given partition random. For each subnetwork, for the weight partitions corresponding to future chunks, we use the values of the weights at initialisation. For discussion of partitioning schemes, see Appendix C.
4 RELATED WORKS
Hyperparameter optimization in deep learning Many works have tackled the challenge of optimizing hyperparameters in deep learning. Works on implicit differentiation, such as the one by Lorraine et al. (2019), allow for optimizing training hyperparameters such as the learning rate, weight- decay, other hyperparameters that affect the final neural network weights only through the training routine. Other works have proposed ways to parameterize and optimize data-augmentattons (Cubuk et al., 2018; Li et al., 2020), search-spaces for neural network architectures, well as methods to optimize architectures using gradient-based optimization (Liu et al., 2018; Elsken et al., 2019). All of the above works have primarily relied optimizing hyperparameters are compatible with the objective defined in this work. Several works have also aimed to cast learning data augmentations as an invariance learning problem. They do so by parameterizing the model itself with data augmentations, and frame invariance learning as a model selection problem (van der Wilk et al., 2018; Benton et al., 2020; Schwöbel al., 2021; Nabarro al., 2022; Immer al., 2022). We compare against Benton et al. (2020) ("Augerino") and Immer et al. (2022) ( ("Differentiable Laplace") on this task in the experimental section.
separate validation set and
Hyperparameter optimization without validation set A limited number of works consider learning hyperparameters without a validation set in a deep learning context. Benton et al. (2020) propose simple method for learning invariances without validation set by regularising invariance hyperparameters to those resulting in higher invariance. They show that the invariances found tend to be insensitive to the regularisation strength, determined by another hyperparameter. However, the method relies on being able to priori define which hyperparameters lead to higher invariance through a suitable regularisation function. In more complex invariance learning settings, defining the regulariser can be challenging. For example, if data- auttenttgmtntiioon transformations were to be parameterized by a neural network (as proposed in Lorraine et al. (2019)), it is non- non-trivial to devise an adequate regulariser. We show that our method can be applied to such settings.
Other works focus on deriving tractable approximations to the marginal likelihood for deep neural networks. Schwöbel et al. (2021) propose only marginalising-out the parameters in the last layer of the neural network by switching it out for Gaussian Process. They treat the preceding layer effectively as hyperparameter, and optimize invariance parameters using the marginal likelihood. Although they show promising results on MNIST, they found they "were unable to learn invariances for CIFAR- (Schwöbel et al., 2021, $7) and highlighted the need marginalise lower layers as well. In contrast, our objective can be seen as being inspired by marginal likelihood where arbitrary network layers can be "marginalised"" and works datasets like CIFAR-10. Immer al. (2022) have adapted the Laplace approximation (Immer et al., 2021) to make it tractable for learning data augmentations. In contrast to Schwöbel et al. (2021), they approximately marginalize and performs favourably. Their approximation, however, requires approximations to a Hessian w.r.t. all network parameters; for that reason, their work reports results for architectures only up a ResNet- 14, whereas our method can easily scale to larger architectures.
all the network parameters,
Hyperparameter optimization in FL Improving byperparameter optimization is especially rele- vant to FL. Given the potential system level constraints (Wang et al., 2021), methods that optimize the hyperparameters and parameters in a single training run are preferred. On this note, Khodak et al. (2021) introduced FedEx and showed that it can successfully optimize the client optimizer
5

---

Published as a conference paper at ICLR 2023
hyperparameters. FedEx relies training/valication split the client level and uses REIN- FORCE type of gradient (Williams, 1992) estimator, which usually exhibits high variance and needs baselines to reduce it (Mohamed et al., 2020). This is in contrast partitioned networks, which use standard, low-variance backpropagation for the hyperparameters and no separate validation set per client. To optimize the other hyperparameters, Khodak al. (2021) wrapped FedEx with a traditional hyperparameter optimization strategy, the successive halving algorithm. This is orthogonal to our method and could be applied partitioned networks as well. In Zhou al. (2021), the authors perform hyperparameter search independently on each client with some off-the-shelf methods and then aggregate the results of the search the server once order identify the best hyperparameter setting. The main drawback of this method compared to partitioned networks is that when the local client datasets are small, client-specific validation set is not informative, and the aggregation happens only once. Finally, there is also the recent work from Seng et al. (2022) which performs hyperparameter optimization and neural architecture search in the federated setting. Similarly to prior works, it requires client-specific validation data in order to optimize the hyperparameters.
5
EXPERIMENTS
Partitioned Posthoc Diagonal Laplace
Input Mask Element
1.00
Likellnood E Estimate
-10
1.0
0.0
Averpp Log-likelhhod
Mask Probability
Log Ma Marginal
0.1
0.75
0.9
Gecrnay
-0.2
15
0.50
0.3
20
Train
0.8
-0.4
25
0.25
Test
0.5
44000
16000
20000
24000
28000
0.00
10 15 20 Num. inputs
25 30
25 30 Num. inputs
10 15
20
Iteration
(a)
(b)
Figure 2: (a) Demonstrating the ability of the marginal-likeliien likelihood inspired objective CML identify the correct model on a toy input selection task. We plot the hyperparameter objective, train maiin-aa and test--- set accuracy, and train and test set log- likeeiodd with the partitioned networks method (left), and the post-hoc diagonal Laplace method (Immer et al., 2021) (right). (b) Mask over input features learned by partitioned networks over time. The first 15 features are correctly identified.
Input Selection To demonstrate that LML is a good objective for model selection that captures the desirable properties of the marginal likelihood, we first deploy our method on the toy model selection task of Lyle et al. (2020): there the first 15 features are informative, and the remaining 15 are spurious
y Bern 2
y + E c + E15y E16. E30
E CI...ry E30
idd NV 1).
Informative
Sparions
We specify a fixed mask over the inputs prior to training, where the first inputs remain unmasked, and the remainder is masked. We expect that, given multiple models with different (fixed) masks over the inputs, the proposed objective will be able to identify the correct one - i.e., the one that keeps only the informative features. We train multiple fully connected neural networks (MLPs) training set of 1000 examples using our method and compare the final values of the LM. objective. The results shown in Figure 2a. LML. correctly identifies 15 input features the optimum, and correlates well with test accuracy and log-likelihoon Training loss and training accuracy, on the other hand, cannot alone disambiguate whether use 15 or more input features.
Differentiable input selection We further show that we can learn the correct mask over the inputs in differentiable manner using method during single training run. We parameterize learnable mask over the inputs with a concrete Bernoulli distribution (Maddison et al., 2016) and treat the parameters of the mask distribution proposed objective using our method. The evolution of the learned mask during training is shown in Figure 2b, where we see that we can correctly identify the first 15 informative features.
hyperparameter. We optimize them with respect to the

---

Published as a conference paper at ICLR 2023
Learning invariances through data-auggmentations Following previous literature on learning soft invariances through learning data augmentations (Nabarro et al., 2022; van der Wilk et al., 2015, 88; Benton et al., 2020; Schwöbel et al., 2021; Immer al., 2022), show that can learn useful affine image augmentations, resulting in gains in test accuracy. We specify affine data augmentations as part of a probabilistic model as done by van der Wilk et al. (2018), averaging over multiple data augmentation samples during training and inference. This allows us to treat the data-augmentation distribution as model hyperparameter rather than training hyperparameter. For datasets, we consider MNIST, CIFAR 10, TinyImagenet along with rotCIFAR 10 and rofTinylmagenet, variants where the datapoints are randomly rotated the beginning of training by angles sampled uniformly from [- n, x| (Immer et al., 2022). Experimental setup details are provided in Appendix I.
For the CIFARTO and roiCIFARIO datasets, we consider as baselines standard training with no augmentations, Augerino (Benton et al., 2020) and Differentiable Laplace (Immer et al., 2022). Following Immer et al. (2022), we use up ResNees (Zhang et al., 2019) for the architectures. The results can seen in Table I. There, we observe that partitioned networks outperform all baselines in the case of CIFARIO for both ResNet variants we consider. On RotCIFAR10, we observe that partitioned networks outperform the baseline and Augerino, but it is slightly outperformed by Differentiable Laplace, which optimizes additional prior hyperparameters. To demonstrate the scalability of partitioned networks, for the (rot)Tiny Imagenee experiments we consider a ResNet-50 architecture with GroupNormm22.. In Table we observe that in both cases, partitioned networks learn invariances successfully and improve upon the baseline. Relative to Augerino, we observe that partitioned networks either improve (Tinylmagenet) or are similar (rofTinylmagene).
Table 1: Test accuracy with learning affine augmentations on (rot)CIFAR10 and (rot)Tinylmagenet.
Method Diff. Laplace
Dataset
Architecture
Baseline
Augerino
Partitioned
RoICIFAR10
Resee
54.210.4
75.4
79.5
19.1
up
fix R
74.14u.5
79.0=1.v
84.240.8
86. 86.1+0.4
CIFARIO
up
fix up
ResNet11
79.510.3
83.0
88.110.2
89.1
Ro(Tinylmagene'
ResNet-50
31.540.6
44.5.0.2
00M
43.940.3
Tinylmagenet
ResNet-50
44.2
41.1
0OM
48.6
Imbuing a model with useful invariances is particularly useful in the low-data regime, due to better data efficiency. To show that, we perform experiments where we artificially reduce the size of the training dataset. The results be seen in Figure 3. We see that by learning augmentations with partitioned networks, we can drastically improve performance in the low-data regime upon baseline that does not learn augmentations, while performing favorably against prior works in most cases.
On MNIST, our method outperforms the last-layer marginal-likeliin Schwöbel al. (2021) in the large data regime but underperforms in the low-data regime. That is likely to be expected, as their work fits Gaussian Process (GP) at the last layer (Wilson et al., 2016), which is better tailored for the low-data regime and results into more flexible model (due to the GP corresponding to an additional, infinite width, layer). Since the MNIST-CNN is sufficiently small to fit multiple networks into memory, we also compare to a variant of our method where, instead of partitioning a single network, we train C' diffferent networks where network & is fit data D1sk. This serves as an upper bound on the performance of the partitioned networks. We see that by partitioning a single network, we can achieve almost equivalent accuracy. On CIFAR10, partitioned networks outperform all other works on all data sizes we considered. On 100, partitioned networks perform again favourably, but they are marginally outperformed by differentiable Laplace in the low-data regime. Compared to partitioned networks where we only optimize augmentations, differentiable Laplace also optimizes the precision of a Gaussian prior over the weights, which better combats overfitting in the low-data regime. On both the TinyImagenet and rot Imagenet experiments we observe that partitioned networks either outperform are similar to the baselines on all data sizes considered.
method (last-layer ML) by
OOt of memory error on 32GB Nvidia V100.
7

---

Published as a conference paper at ICLR 2023
0.75
0.50
0.25
Test Accuracy
0.99
Baselini
0.25
Last-layer
Augerino
0.98
Laplace
0.4
Partitioned (tns.)
0.50
0.2
5000
20000
60000
0.25
5
10
20
50
10
50
100
Dataset Size
Dataset Size (x1000)
Dataset Size (x1000)
(a) MNIST
(b) (ro))CIFARIO
(c) (rot)Tinylmagenet
Figure 3: Learning affine data augmentations on subsets of data. (b) uses ResNeet architecture
whereas (c) a ResNet-50 architecture. (b,c) Top: normal dataset, bottom: rotated dataset.
Comparisons to traditional training validation split We further perform comparisons between partitioned networks and the more traditional training/validationn split (denoted as validation set optimization) with additional finetuning to the task of learning data augmentations. This is realized as follows; we partition 20k CIFARIO examples into training and validation data of specific proportions. We then either train partitioned network (along with the hyperparameters on LMI.) these two chunks of data or train a standard network on the training set while using the validation set loss to obtain gradients for the data augmentation hyperparameters. For the validation set optimization baseline, once the hyperparameters are optimized, the resulting network finetuned on the whole dataset for 20 epochs. The results for varying chunk proportions are provided in Table 2.
Table Learning affine augmentations with ResNee on subset of CIFAR-10 (20k examples). NaN denotes that run crashed.
Chunk Proportions [0.7, 0.3)
Method
[0.3, 0.7]
[0.5, 0.5]
[0.8, 0.2]
[0.9, 0.1]
Partitioned
62.9%....5
03.0%.1..1
83.7% 83.7%4o.2
84.0%...6
84.6%
Validation optim.
NaN
78.9%
81.5%10.2
82.6%
83.4%.101
L +Finetune
NaN
81.3%
82.5%.1.2
83.5%
83.8%10.3
We can see that partitioned net- works (that do not employ ad- ditional finetuning) outperform validation set optimization with finetuning all settings we tried. The gap does get smaller when we move to the more tra- ditional 90/10 splits for train- ing/validation: a 10% proportion for validation data is enough to optimize a handful of hyper- parameters (just 6 scalars). To corroborate this claim, set up additional experiment; a Wide ResNet-20 on the full CIFARIO dataset, where the first two out of the three stages (13 convolution layers) are considered as hyperparameters. The results for this setting can be seen in Table 3. We see that 10% validation data are not enough, and the validation set optimization baseline performs poorly. This in contrast to partitioned networks, where with three chunks, we can learn all of these hyperparameters successfully. Note that, compared to Augerino, applying partitioned networks this setting is straightforward. To apply Augerino, one would have to come up with metric that can be used to regularize the feature extractor towards "higher invariance".
Table 3: Learning a feature extractor (first out of stages of Wide ResNet-20) as hyperparameter CIFAR10.
Method
Chunk Proportions
Test accuracy
Validation set optim.
00.9, 0.1]
59.6%
Partitioned
(0.1, 0.8,0.8,0.11 0.1]
87.3%+0.8
use
Partitioned networks for federated learning We consider federated learning (FL) (McMahan et al., 2017), setting where data is distributed across many clients. In this setting, there system properties that make hyperparameter optimization especially challenging (Wang et al., 2021). More specifically, obtaining validation set and performing multiple training with different
8

---

Published as a conference paper at ICLR 2023
hyperparameter settings might not be possible due to the additional communication and computation costs, and transient client availability (clients join and leave the training process at any time). Optimizing hyperparameters together with the model parameters in single run is therefore especially beneficial (Wang et al., 2021), and partitioned networks are a good fit for FI..
We extend centralized experimental setup to FL by splitting all N clients into C' non-overlapping chunks, such that each chunk is understood as the union of all clients' data shards that belong to that chunk. During federated training, client belonging to chunk & sequentially optimizes partitions WA:C: through sub-networks w (k:C) and computes gradient wrt. the hyperparameters ab. Note that partitions wisk remain unchanged and do not need to be communicated back to the server. This reduction in upload costs is a welcome property for FL, where upload costs can bottleneck system design. The server receives the (hyper-) parameter updates, averages them, and applies the result as a "gradient" to the server-side model in the traditional federated manner (Reddi et al., 2020). For partitioned networks, the hyperparameters that we optimize are the data augmentation parameters and, since we also include dropout in these architectures, the dropout rates (with the concrete relaxation from Maddison et al. (2016)). As a baseline, we consider the standard federated training without learning hyperparameters (denoted as FedAvg) as well as learning the augmentation parameters with Augerino Benton al. (2020). Please see Appendix J for a detailed explanation of our FL. setup.
Table 4 summarizes our results using different sub-sets and variations of MNIST and CIFARIO, where we also included rotMNIST Larochelle et al. (2007) as another dataset. We can see that partitioned networks allow training models that generalize better than both FedAvg and FedAvg with Augerino, at reduced communication costs. Especially when the true data -ata-genng process and underlying source of non-i..dd ness are explicitly accounted for - here in the form of rotation - the benefits of learning the augmentations with partitioned networks become apparent. For example, we observe that the rotated datasets, partitioned networks learn correctly increase the rotation angle.
Table 4: Validation accuracy averaged over the last 10 evaluations, each 10 rounds apart; standard- error is computed across 4 random seeds. All datasets are adapted to the federated setting and are synthetically split to be non-i.i.d. sampled as described in Appendix J.2.
Dataset & size Method
1MNIST 5k
RRoMNNIST 5k
JUpload [%)
1.25k
50k
1.25k
50k
FedAvg
95.4%
97.4%
80.5% 99.0%
90.4%
96.8%=o.1
100
FedAvg + Augerino
94.2%22..
96.4%.4.1
99.1%
79.5%
89.0%+2.0
95.3%=o.2
100
FedAvg + Partitioned
97.0%
98.3%20.0
99.2%...1
85.7%.o.9
93.5%
97.8%.o.1
77
1CIFARIO
RoCCIFARIO
AUpload 1%1
1.25k
5k
45k
1.25k
5k
45k
FedAvg FedAvg Augerino
50.2%
64.5%.o..
79.2%+a.7
35.6%=o.3
45.2%
53.9%=1. 1%.2%
100
49.9%
65.0%4o.2
79.9%+4.4
36.1%=oo
45.0%.4..
56. 56.4%=o.r
100
FedAvg Partitioned
50.8%11.0
64.8%.4.4
81.5%
37.1%...2
45.3%.o.3
60.6 60.6%11o2
91
DISCUSSION
We propose partitioned networks new method for hyperparameter optimization inspired by the marginal likelihood objective. It provides a general and scalable solution to finding hyperparameters in single training run without requiring access to a validation set while introducing less additional overhead to the training task than existing approaches. We showed that partitioned networks are applicable on a wide range of tasks; they can identify the correct model on illustrative toy examples, they can learn data augmentations in a way that improves data efficiency, they can optimize general feature extractors as hyperparameters and they can also optimize dropout rates. In the federated setting, partitioned networks allow us to overcome practical challenges, reduce the communication overhead and obtain better models. The notion of partitioned networks we propose in this work is novel to the literature and an orthogonal approach to many existing hyperparameter tuning algorithms. Like any other method, partitioned networks come with their own limitations, e.g., needing a partitioning strategy. We expand upon them in appendix H. We hope to see method successfully reducing the need to perform hyperparameter search through repeated training and thereby contribute to the community's effort to reduce its carbon footprint.
9

---

Published as a conference paper at ICLR 2023
REFERENCES
Gregory Benton, Mare Finzi, and Andrew G Wilson. (o0r 19d8572. https: https//gittuubccot on
Augerino, github, com-
learning-invariancen (on page 18)
Gregory Benton, Marc Finzi, Pavel Izmailov, and Andrew G Wilson. networks from training data, Advances in neural information processing systems, 33: 33:17605-17616. 2020. (on page 2, 5, 7, 9, 16, 18, 20, 24, 25)
Learning invariances in neural
Ekin D Cubuk, Barret Zoph, Dandelion Mane, Vijay Vasudevan, and Quoe V Le. Autoaugment: Learning augmentation policies from data. arXiv preprint arXiv: 8. (on page 5)
Kamal Dys. Cifarl0 resnet: 90+% % accuracy;less than 5 min. httpps:/www..kgggll./
code/ kmldas /e//aaflloooo wew.ettrreeeet-eee www.rreeeeeneeenueueeeeeee -eeerraa---eees-----------e-- than-thhhnnnnn Accessed: 2022-09- 17. (on page 26)
Thomas Elsken, Jan Hendrik Metzen, and Frank Hutter. Neural architecture search: A survey. The Journal of Machine Learning Research, 20(1):1997-2017, 2019. (on page 5)
Edwin Fong and Chris C Holmes. On the marginal likelihood and cross-validation. Biometrika, 107 (2):489-499) 2020. (on page 2)
Jonathan Frankle and Michael Carbin. The lottery ticket hypothesis: Finding sparse, trainable neural networks. arXiv preprint
2018. (on page 4)
Yarin Gal, Jiri Hron, and Alex Kendall. Concrete dropout. Advances in neural information processing systems, 30, 2017. (on page?> 2)
Pascal Germain, Francis Bach, Alexandre Lacoste, and Simon Lacoste-Julin Pac-bayesian theory meets bayesian inference. Advances in Neural Information Processing Systems, 29, 2016. (on page 2)
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance on imagenet classification. In Proceedings of the IEEE international conference on computer vision, pages 1026-10333 1034, 2015. (on page 23)
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Identity mappings in deep residual networks. In European conference on computer vision, pages 630-645, Springer, 2016. (on page 23)
Dan Hendrycks and Kevin Gimpel. 160.041666..8415. 2016. (on page 23)
Gaussian error linear units (gelus). arXiv preprint
Alexander Immer and Tycho F. A. van der Ouderaa. Learning invariances with laplace ap- https:
proximations (lila), github, lila. (on page 17)
Alexander Immer, Matthias Bauer, Vincent Fortuin, Gunnar Rätsch, and Khan Mohammad Emtiyaz. Scalable marginal likelihood estimation for model selection in deep learning. In International Conference Machine Learning, pages 4563-4573 PMLR, 2021. (on page 2, 6, 24)
Alexander Immer, Tycho A. der Ouderaa, Gunnar Rätsch, Vincent Fortuin, and Mark der Wilk. Invariance learning in deep neural networks with differentiable laplace approximations, 2022. URL httpt:t////c /atr///aaiii// http:/aaaivvv.....ia..///... v.org/ http:///aab//22.0553332 (on page 5, 7, 15, 16, 17, 18, 22, 23, 24, 25)
Sergey Toffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International conference on machine learning, pages 448-456. PMLR, 2015. (on page 23)
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.0m0 (08611, 2020. (on page
10

---

Published as a conference paper at ICLR 2023
Mikhail Khodak, Renbo Tu, Tian Li, Liam Li, Maria-Florina F Balcan, Virginia Smith, and Ameet Talwalkar. Federated byperparameter tuning: Challenges, baselines, and connections to weight- sharing. Advances in Neural Information Processing Systems, 34.19184-19197, 2021. (on page 5, 6)
Hugo Larochelle, Dumitru Erhan, Aaron Courville, James Bergstra, and Yoshua Bengio. An empirical evaluation of deep architectures problems with many factors of variation. In Proceedings of the 24th international conference on Machine learning, pages 783-480, 2007. (on page 9)
Ya Le and Xuan Yang. Tiny imagenet visual recognition challenge. CS 231N, 7(7):, 2015. (on page 24)
Yonggang Li, Guosheng Hu, Yongtao Wang, Timothy Hospedales, Neil M Robertson, and Yongxin
Yang. Dada: Differentiable automatic data augmentation. arXiv preprint arXiv:2003.03780, 2020.
(on page 5)
Hanxiao Liu, Karen Simonyan, and Yiming Yang. Darts: Differentiable architecture search. arXiv preprint 1806 09055, 2018. (on page 5)
Jonathan Lorraine, Paul Vicol, and David Duvenaud. Optimizing millions of hyperparameters by implicit differentiation. CoRR, Abs/1911.00599 2019. URL httt:://axxiv.org/abbl 1911.02590. (on page 1, 5)
Sanae Lotfi, Pavel Izmailov, Gregory Benton, Micah Goldblum, and Andrew Gordon Wilson Bayesian model selection, the marginal likelihood, and generalization. In Kamalika Chaud- huri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, editors, Proceedings of the 39th International Conference on Machine Learning, volume 162 of Pro- ceedings of Machine Learning Research, pages httpp://proceddaanoe http///oceeeen.g.l...c.c/ www.:ml.........mmmmmmmmmmmmmsee v2622112v// http://off22aahtttti. (on page 4)
PMLR, 17-23 Jul 2022. URL
Clare Lyle, Lisa Schut, Robin Ru, Yarin Gal, and Mark van der Wilk. A bayesian perspective on training speed and model selection. Advances Neural Information Processing Systems, 33: 10396-10408, 2020. (on page 2, 3, 5,5,0 6)
David JC MacKay. Information theory, inference and learning algorithms. Cambridge university press, 2003. (on page 2)
Chris J Maddison, Andriy Mnih, and Yee Whye Teh. The concrete distribution: A continuous relaxation of discrete random variables. arXiv preprint arXiv:1611.00712 2016. (on page 6, 9, 26, 27)
David A McAllester. Some pac-bayesian theorems. In Proceedings of the eleventh annual conference on Computational learning theory, pages 230-234, 1998. (on page 2)
Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas caaamunication-fffcccee learning of deep networks from decentralized data. In Artificial intelli- gence and statistics, pages 1273-1282. PMLR, 2017. (on page 2, 8)
Shakir Mohamed, Mihaela Rosca, Michael Figurnov, and Andriy Mnih. Monte carlo gradient estimation in machine learning. J. Mach. Learn. Res,, 21(132): 2020. (on page 6)
Seth Nabarro, Stoil Ganev, Adrià Garriga-Alonso, Vincent Fortuin, Mark van der Wilk, and Laurence Aitchison. Data augmentation in bayesian neural networks and the cold posterior effect. In Uncertainty in Artificial Intelligence, pages 1434-1444. PMLR, 2022. (on page 5,
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary Devito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performance deep learning library. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d' d'Alché-Bus E. Fox, and R. Gar- nett, editors, Advances in Neural Information Processing Systems 32, pages 8024-8035. Curran Associates, Inc., 2019. URL imperative st page 22)
htw.://oooo.oooeeaaaaaaaoaaaaaae (on
11

---

Published as a conference paper at ICLR 2023
Sashank Reddi, Zachary Charles, Manzil Zaheer. Zachary Garrett, Keith Rush, Jakub Konečný, Sanjiv Kumar, and Brendan McMahan. Adaptive federated optimization. arXiv preprint arXiv:2003.00295, 2020. (on page 9, 26, 27)
Robin Ru, Clare Lyle, Lisa Schut, Miroslav Fil, Mark der Wilk, and Yarin Gal. Speedy performance estimation for neural architecture search. Advances in Neural Information Processing Systems, 34:4079-4099 2021. (on page 3)
Akiyoshi Sannai, Masaaki Imaizumi, and Makoto Kawano. Improved generalization bounds of group invariant/equivariant deep networks via quotient feature spaces. In Uncerlainty in Artificial Intelligence, pages 771 1-1-7 PMLR, 2021. (on page 3)
Pola Schwöbel, Martin Jørgensen, Sebastian W. Ober, and Mark van der Wilk. Last layer marginal likelihood for invariance learning, 2021. URL https:/ http:////wwwwwwwwaa//a////esea/a 2tt.://////////10.0000011222 (on page 2, 5, 7, 15, 16, 23, 24, 26, 27)
Jonas Seng, Pooja Prasad, Devendra Singh Dhami, and Kristian Kersting, Hanf: Hyperparameter and neural architecture search in federated learning. arXiv preprint arXiv:2206. 12342, 2022. (on page 6)
Jure Sokolic, Raja Giryes, Guillermo Sapiro, and Miguel Rodrigues. Generalization enror of invariant classifiers. In Artificial Intelligence and Statistics, pages 1094-1103, PMLR, 2017. (on page 3)
Mark van der Wilk, Matthias Bauer, ST John, and James Hensman. Learning invariances using the marginal likelihood. Advances in Neural Information Processing Systems, 31, 2018. (on page 7, 16)
Bastiaan S Veeling, Jasper Linmans, Jim Winkens. Taco Cohen, and Max Welling, Rotation equiv- ariant cnns for digital pathology. International Conference on Medical image computing and computer-assisted intervention, pages 110-218. Springer, 201 8. (on page 15)
Jianyu Wang, Zachary Charles, Zheng Xu, Gauri Joshi, H Brendan McMahan, Maruan Al-Shedivat, Galen Andrew, Salman Avestimehr, Katharine Daly. Deepesh Data, et al. field guide to federated optimization. arXiv preprini arXiv: 107.06917, 2021. (on page 5, 8, 9)
Ronald J Williams. Simple statistical gradient-following algorithms for connectionist reinforcement learning. Machine learning, 8(3):229-22) 2566 1992. (on page 6)
Andrew Gordon Wilson, Zhiting Hu, Ruslan Salakhutdinov, and Eric P Xing, Deep kernel learning. In Artificial intelligence and statistics, pages 370-378. PMLR, 2016. (on page 7)
Yuxin Wu and Kaiming He. Group normalization. In Proceedings of the European conference on computer vision (ECCV), pages 3 391 2018. (on page 23, 26)
Sergey Zagoruyko and Nikos Komodakis. Wide residual networks. arXiv preprint arXiv: 160.00600 146, 2016. (on page 23)
Hongyi Zhang, Yann N Dauphin, and Tengyu Ma. Fixup initialization: Residual learning without normalization. arXiv preprint arXiv: /901. 00321, 2019. (on page 7, 18, 23)
Yi Zhou, Parikshit Ram, Theodoros Salonidis, Nathalie Baracaldo, Horst Samulowitz, and Heiko Ludwig. Flora: Single-shot hyper-parameter optimization for federated learning. arXiv preprint arXiv:21 2021. (on page 6)
12

---

Published as a conference paper at ICLR 2023
A LT. IS A LOWER-BOUND TO THE MARGINAL LIKELIHOOD
In this section, we show that the objective in equation is lower-bound on the marginal likelihood, under mild assumption on each approximate posterior qk (w). The aim approximate:
pogp(Dl>) log (P
(5)
k=1
Our partitioned approximation is given by:
Σ F (ww) (DDx/r, sp]
(6)
We can get the equation for the gap between quantities in and 6:
gap
(Dg 11.
E (( [log((Dulu,))
(7)
k=1
(ww)
[Dog () log logp(Dx/20, +)]
(D
(ww) (w)
log
(9)
p( & 2. 15)
p(w, (D))
D 1sks 2 po D v))((()))))) (Dp(w/Diik 1I:
1 (w)
log
(10)
p(w DD 1ike (DD ww, +b vppp(p u FoolDiiss
1, 1
k=1
p(w. Di
p(w 1I 15) po10) Diske vb)
9k-1 (w)
log
(11)
DKL ( Di:ks 2/1 - DKL [q=-1(w)] (1.k
(12)
We now make two assumptions
)²) This is motivated from the fact that qj (w) is trained on all data chunks Disk so it is expected to be a better approxima- tion to the posterior p(w) D DI.)), compared qu-((w) which only trained on D1
DKL [gk-1 (wn)
vb)) DEL [gx (ux))n) (Di:k
DKL (9c-1(w) [p(p)p((ps)) +)] 2 DKL ([p(w)))())) Since we are free choose the approximate posterior before seeing any data go (w)- we can set it to be equal to the prior p(w) which, together with the positivity of the KL divergence, trivially satisfies this assumption.
Therefore, by rearranging Eq. 12 and using our two assumptions we have that the gap is positive
gap = - DKL lgo (w))])((w)) + DKL I&c (w))://(/oollpp///P/ee..
C
DKL ((-1((w) [| Di. ] DKL [gx (w) || DIskY b)) 2 0,
(13)
and our approximation lower bound to the marginal likelihood, i.e.,
log logp(Dl)) >
Σ E (10) [logp(Dk(ee v/)
(14)
13

---

Published as a conference paper at ICLR 2023
B PARTITIONED NETWORKS AS A SPECIFIC APPROXIMATION TO THE MARGINAL LIKELIHOOD
In this section of the appendix, show that the partitioned neural networks presented in the paper are a particular instance of the approximation to the marginal likelihood shown in equation 2.
Consider a dataset D comprised of C' shards, i.e. D (DI 1....... ...... Dc), along with model, e.g., neural network, with parameters up & IR=x, a prior p(w) I (ww)0,, ( X) and likelihood p(D)w, w, )) with hyperparameters </. Assuming a sequence over the dataset chunks, we can write out the true marginal likelihood
log (gg((Dlp) =
Σ logp(Dr/DDit-1>) Σ &
log Eptw(pmee [p(Da 20, (0, +)]
(15)
2
Σ EpulDDus 1:41 ((ok ( w)]
(16)
Since the true posteriors p(w) Diy, 2)) for j .l....,C CJ are intractable, we can use variational inference approximate them with gø, (w) for j E (I,..., C). with 6j being the to-be-optimized parameters of the j 'tth variational approximation. Based on the result from Appendix A, when q6, 46, (w)
are optimized to match the respective posteriors pl p(w|D1.3. 15), we can use them approximate the marginal likelihood as
log (DI N
Σ E (20) [log
v)]
(17)
Partitioned networks correspond a specific choice for the sequence of approximating distribution families lsb (w. Specifically, we partition the parameter space into C chunks, i.e., E PR such that Σx Dwk Let The (wx) be base variational approximations over wx with parameters on. Each approximate distribution do. (w) is then defined terms of these base approximations, i.e.,
Dw. and we associate each parameter chunk wr with a data shard Dk-
K II mm-
(or(w) =
II ras (w;) ro, rol (wa)
ro(
(18)
where (-)- is some base distribution with no free parameters. In accordance with the assumptions in appendix A. then fit each q (w) by minimising the KL-divergence to p(w ((wwDD:k vb) the posterior after seeing k chunks:
DKL las. (w))lp( (i:kk v=) = Eas, (w) (w) (( 4 + DKL [go. (w))lp(w))
+ llog ((D1sk())))
(19)
(20)
Finding the optimum with respect to Φx:
arg min DKL. Ino. (w [|| (Disk +b]
(21)
dx
arg min E Ess (w) Dlg (D w. xb)] + DKL. (ao (w)llp(w)
(22)
d2
arg min
(wx) pog 2b)]
K II m==+11
K
+ DKL
II i=1
roy (w;) rou (wk)
ro(wm))
[I(wws)
(23)
arg min
E go, (w) (ww) (loog((Diiii w. +t)] DKT. Iro (wx) (wx))lp(w>)]]
(24)
We can now obtain partitioned networks by assuming that ro, (w)) = NW 10x, vI) for e (1...., C). ro(w) (u vI), with being the parameters initialization (i.e., before
14

---

to the layer to retrieve an equivalent computation to that with (h The gradients also need to be masked the backward pass adequately. No such simplification is possible with the random weight partitioning scheme; if we were to compute a backward pass for a single batch of examples using Published as a conference paper at ICLR 2023
update them data) and taking 0, i... in muchine-precision the weights deterministic. As noted in Section I.1, we scale the weight-decay regularizer for Φk (whenever used) differently for each partition k, such that it can be interpreted as regularization towards a prior. In the experiments where we do not regularize de according to ((wk) when we optimize them, this implicitly corresponds to + oo (i.e. the limiting behaviour when the variance of p(w) goes to infinity), which makes the contribution of the regularizer negligible.
C PARTITIONING SCHEMES
There are several ways in which we could aim to partition the weights of neural network. Throughout the experimental section 5, we partition the weights by assigning fixed proportion of weights in each layer to given partition at random. We call this approach random weight partitioning.
We also experimented with other partitioning schemes. For example, tried assigning fixed proportion of a layer's outputs (e.g., channels in a convolution layer) to each partition. All weights in given layer that specific output depends would then be assigned to that partition. We call this approach node partitioning. Both approaches are illustrated in Figure 4.
One benefit of the node partitioning scheme is that it makes it possible to update multiple partitions with a single batch; This is because we can make a forward pass at each linear or convolutional layer with the full network parameters w, and, instead, mask the appropriate inputs and outputs
different subnetworks for each example, the memory overhead would grow linearly with the number of subnetworks used.
In initial experiments, found both random weight partitioning and node partitioning performed similarly. In the experimental section 5, we focused on the former, as it's easier to reason about with relation e.g., dropout.
Throughout this work, partitioning happens prior to initiating training, and remains fixed throughout. might also be possible to partition the network parameters dynamically during training, which we leave for future work.
In node assignment
W111
W12
W13
W14
W15
W11
W0
W/1
W014
W15
W12
W22
W/23
W/24
W25
102
1022
1023
1024
1025
W11
W32
W33
W34
W33
Out
Wa1
1032
1033
1034
1/3
pode
W44
0742
W43
W44
W445
assignment
1041
1042
1043
1044
W44
W51
W752
W53
W54
W55
0551
0052
1053
0054
W55
W61
W62
W63
W61
W65
U61
1062
W/63
W061
W66
(a) Random weight partitioned
(b) Node partitioned
Figure 4: Figures showing how the weights within a single weight matrix W E g6x5 for a linear layer would be partitioned.
D SCALABILITY
In the paper, we claim that our method is scalable compared to Schwöbel et al. (2021) and Immer et al. (2022). What constraints the scalability of the mentioned prior works, however, is different.
For the Last Layer Marginal Likelihood, although the approach works on small datasets such as PCAM (Veeling et al., 2018) and MNIST, the authors report that they were unable to learn invariances
15

---

Published as a conference paper at ICLR 2023
on larger datasets such as CIFARIO. In (Schwöbel et al., 2021, section 7), they explore the issue of scalability in more detail, and showcase that last layer marginal likelihood is insufficient.
Differentiable Laplace performs well, scalability, however, is limited by the computational and memory complexity of their method, which go into in more detail in the section below.
complex datasets, such CIFARIO. Their
D.1 COMPLEXITY ANALYSIS
First, we consider the scalability of our algorithm in terms of computational and memory complexity. In particular, we show that our method scales much more favourably compared to Differentiable Laplace (Immer et al., 2022).
We present our analysis for feed-forward model of depth L, with layer widths D'. order to directly compare to Immer et al. (2022) and Benton et al. (2020), consider the complexities in the invariance learning setup (Benton al., 2020; van der Wilk et al., 2018) with augmentation samples. In other experiments, hyperparameter optimization setups, S can be taken to be 1. The notation is summarized Table 5.
N
Number of datapoints in dataset D Batch size
NB
S
Number of augmentation samples'
C
Output size (number of classes) Feedforward network layer widths
D
L
Feedforward network depth
Number of parameters (s.t. O(P) = OLLD² + DO))
Table 5: Notation for complexity analysis.
We consider the computational and memory costs of 1) obtaining a gradient with respect to the parameters 2) obtaining a gradient with respect to the hyperparameters, and 3) computing the value of the model/htperrarrmeeer selection objective for each method. All analysis assumes computation on a Monte-Carlo estimate of the objective on a single batch of data.
In Tables 6 and 7, assume that C' D, and hence, for the clarity of comparison, sometimes fold factor depending C' into a factor depending on D it's clearly smaller. This hiding of the factors was only done for Differentiable Laplace, which is the worst scaling method.
D.1.1 COMPUTATIONAL COMPLEXITY
Param.
Hyperparam. Backward
Hyperparam. Objective
Backward
Partitioned
O(NBPS) O(NPPS)
O(NPPS) O(NBPS) O(NPS + NCP +NCDLS + LD*)
Augerino
O(NPPS)
O(NBPS)
ONN
Diff. Laplace
O(NBPS)
+NCDLS + LD*)
Table 6: Computational Complexities. The two terms highlighted for Augerino can be computed in a single backward pass. For Differentiable Laplace, the terms in blue can be amortised over multiple hyperparameter backward passes. That is why, in their method, they propose updating the hyperparameters once every epoch on (possibly) multiple batches of data, rather than once on every batch as is done with Partitioned Networks and Augerino.
*This for the ease of comparison. Same upper hound complexities will hold network variable sizes Dr for E [L], where D = max; Dr OOlyy relevant for invariance learning.
16

---

Published as a conference paper at ICLR 2023
D.1.2 MEMORY COMPLEXITY
The memory complexities for Partitioned Networks, Augerino, and Differentiable Laplace are shown in Table 7. Crucially, the memory required to update the hyperparameters for Differentiable Laplace scales as O + P). with term depending on the square the network widths. This can become prohibitively expensive for larger models, and is likely the reason why their paper only considers experiments on architectures with widths up to a maximum of 256.
Param.
Hyperparam. Backward
Hyperparam. Objective
Backward
Partitioned Augerino
ONNSSLD P)
O(NBSLD P) O(NBSLD P) O(NBSLD + P)
O(NBSD P) O - P) O(NBSLD' +
ONNSLD + P)
Diff. Laplace
ONNSLD P)
Table 7: Memory Complexities Differences are highlighted in red.
D.2 PRACTICAL SCALABILITY
complexity analysis in big-O O notation as provided by us in the previous sections allows to understand scalability in the limit, but constant terms that manifest in practice are still of interest. In this section we aim present real timing measurements for our method in comparison to Augerino and Differential Laplace, and elaborate what overhead might be expected with respect standard neural network training.
The empirical timings measurements on an NVIDIA RIX 3080-10GB GPU are shown in Table We used a batch-size of 250, 200 for the MNIST and CIFAR 10 experiments respectively, and 20 augmentation samples, just like in our main experiments in Table I and Figure 3. As can be seen, the overhead from using partitioned network is fairly negligible compared to standard forward and backward pass. The one difference compared Augerino is, however, the fact that a separate forward-back backward pass needs to be made update the hyperparameters and regular parameters. This necessity is something that can be side-stepped with alternative partitioning schemes, as preliminarily mentioned in appendix C, and is interesting direction for future research.
MNIST
CIFARIO
Method
CNN
ExResNet-8 up
up 1ete
Augerino
x1
X1
x1
Param.
x1
x1
x1
Diff. Laplace
Hyperparam.
22015.6
x18.2
Param. x1.08
x1.17
x1.21
Partitioned
Hyperparam. x1.08
x1.08
x1.09
Table 8: Relative empirical time increase with respect to regular parameter update during standard training. The timing multipliers with respect to the baseline for 6 RexeeNet-8 are taken from the timings reported in (Immer et al., 2022, Appendix D.4). On the ResNet-14, we get an out-of-memory error during the hyperparam. update step with Differentiable Laplace on the NVIDIA RTX 3080- 10GB GPU when running with the official codebase (Immer and van der Ouderaa).
Memory Overhead Our proposed method's memory consumption scales in the same way as Augerino or vanilla neural network training. There is a minor constant memory overhead due to having to store the assignment of weights to partitions. In general, only log C' bits per parameter are necessary to store the partition assignments, where C is the number of chunks. In our implementation, only consider C < 28, and hence store the assignments in byte tensors. This means that the partitioned models require extra 25% memory for storing the parameters (when using 32bit floats to represent the parameters).
17

---

Published as a conference paper at ICLR 2023
If the "default" weight values (i.e. those denoted w; in Figure 1) are non-zero, there is additional overhead to storing those as well, which doubles the memory required to store the parameters. difference in performance when setting default weight values to 0 in architectures which normalisation layers are used (i.e. most modern architectures). As such, would in general recommend default values to the initialised values to be necessary for stability of training deep normalisation free architectures such as the fix architectures (Zhang et al., 2019) we used to compare with Differentiable up
We observed there
the default weight values to O. However, found setting
Laplace. As their method not compatible with BatchNorm, we used these architectures in our experiments, and hence used non-zero default values.
Lastly, if the default weight values are set to the (random) initialisation values, it possible write cleverer implementation in which only the random seeds are stored in memory, and the default values are re- e--genrrated every time they are need forward and backward pass. This would make the memory overhead from storing the default values negligible.
E NOTE ON AUGERINO
In replicating Augerino (Benton et al., 2020) within our code-base and experimenting with the implementation, we discovered pathological behaviour that is partly mirrored by the authors of Immer et al. (2022). In particular, note that the loss function (Benton et al., 2020, Equation (5)) proposed by the authors is problematic the sense that for any regularization strength À > the optimal loss value is negative infinity since the regularization term (negative L2-norm) is unbounded. In our experiments we observe that for a sufficiently- iterations, this behaviour indeed appears and training diverges. In practice, using Augerino therefore necessitates either careful tuning X, clipping the regularisation term method that introduces yet another hyperparameter), or other techniques such early stopping.
value of À and after sufficient number of
In the open-source repository for the submission (Benton et al.), it can seen that on many experiments the authors use "safe" variant of the objective, in which they clip the regulariser (without pass-through of the gradient) once the Lo -oorm of any of the hyperparameters becomes larger than arbitrary threshold. Without using this adjustment, found that the Augerino experiments on MNIST crashed every time with hyperparameters diverging to infinity.
F
SENSITIVITY TO PARTITIONING
F.1 SENSITIVITY IN TERMS FINAL PERFORMANCE
(a) 2 chunks
(b) 3 chunks
(c) chunks
[8.1
184.11
18.42.11
Ounnk hessncors
16.11
The
fresee
132211
Chunk hmsoetan
14.2211
daww
Tes
[8.3]
(3.2.A1
[111111
0.m9
[1.4]
0,500
0,900
[a.]
[3..
[2281
1841
(141
1441
(la)
18421
(A22))
(2141
Real
(8423)
(6222)
(12241 (LLAA
(2244)
Proprtions
4 am Prooriog
Aram Pipprtion
Figure 5: Learning affine augmentations on MNIST with CNN fit on all data. r- and y- ticks denote the ratios of parameters/datappints assigned to each partition/chunk respectively.
Partitioned networks allow for learning hyperparameters in single training run, however, they introduce an additional hyperparameter in doing so: the partitioning scheme. The practitioner needs to choose the number of chunks C, the relative proportions of data in each chunk, and the relative proportions of parameters assigned to each of the C' partitions wi. We investigate the sensitivity to the partitioning scheme here. We show that
results are fairly robust partitioning through grid-search over parameter partitions and chunk proportions on the affine augmentation learning task on MNIST with the CNN architecture we use throughout this work.
18

---

Published as a conference paper at ICLR 2023
(a) 2 chunks
(b) 3 chunks
(c) 4 chunks
D-800
[8.1]
[8.4.1]
0,000
Prepst ore
Ohuk regur
[4:11
13.2.11
14.32.1
11.31
[3.3.11
[1L11.3]
Chirk
[1.41
13.2.01
112334
0,007
(8.1
1841
(t,a)
[82.]
(3441
[12.51
(8321)
(12344)
(1448)
Panam, Proportons
From Pocorors
Aram Poppation
Figure 6: Learning affine augmentations RotMNIST with CNN fit all data. and y- ticks
denote the ratios of parametersvdatappinns assigned to each partition/chunk respectively.
Figure and Figure 6 show the test accuracy for choice of chunk and parameter proportions across two, three and four chunks. The proportions are to read as un-normalized distributions: for example, chunk proportions set to [1, 8] denotes that there are 8x many datapoints assigned to the second compared to the first. Each configuration was run with 2 random seeds, and we report the mean across those in the figure. The same architecture used the same for the main MNIST experiments in section 5 (see Appendix I.4 for details).
We observe that for various partition/datase--chuknngg configurations, all models achieve fairly similar final test accuracy. There is a trend for models with lot of parameters assigned to later chunks, but with few datapoints assigned later chunks, to perform worse. While these results show high level of robustness against the choice of additional hyperparameters introduced by our method, these results show an opportunity or necessity for choosing the right partitioning scheme order to achieve optimal performance.
F.2 SENSITIVITY IN TERMS OF HYPERPARAMETTER FOUND
To compare how the different partitioning schemes qualitatively impact the hyperparameters that the method identifies, also retrain vanilla models from scratch using the hyperparameter values found using partitioned networks. Namely, we take the final value of the hyperparameters learned with partitioned networks with given partitioning scheme, and plot the final test accuracy of vanilla neural network model trained from scratch with those byperparameters. The results are shown Figures 7 and 8.
(a) 2 chunks
(b) 3 chunks
(c) 4 chunks
0.90
Curr beworurr
16.11
Carnk Preertion
13.2.11
Chink bosortirs
14.22.11
0.9W
11.31
(3.3.11
[1111
[1.4]
[3.2.]]
(1234)
0.90
0.30
[L.]]
[L248]
0.41
(14)
(8441)
[8.2.1
Param Properioos (244))
(2.14)
(2.2.))
(82222)
(1224))
Propotionn
Bon,
Figure 7: Standard neural network trained on MNIST with CNN fit all data, with hyperparameters found using partitioned networks with chunk and parameter proportions corresponding to those in Figure 5. and y ticks denote the ratios of parameters/dataapinns assigned each partition/chunk respectively.
G HOW GOOD ARE THE HYPERPARAMETERS FOUND?
Here we show that the hyperparameters found by partitioned networks are also a good set of hyperparameters for vanilla neural networks retrained from scratch. This section expands the
19

---

Published as a conference paper at ICLR 2023
(a) 2 chunks
(b) 3 chunks
(c) 4 chunks
8-800
[8.1]
[8.4.1]
Ohuk regur
[4:11
0,000
Prepst one
13.2.11
14.32.13
Se Secn
11.33
(3.3.11
[1L11.3]
Chink
[1.41
13.2.01
112334]
0,507
(8.1
18.1
(t.a)
[82.1
(3441
Baal
(8321
(21111)
(12344)
(1240)
Panam Proportons
Poom Rocerars
Aram Poppation
Figure 8: Standard neural network trained RotMNIST with CNN fit all data, with hyper- parameters found using partitioned networks with chunk and parameter proportions corresponding to those in Figure 6. x- and y- ticks denote the ratios of parameters/dataponns assigned to each partition/chunk respectively.
experiment in section F.2. To validate this claim, conducted fairly extensive hyperparameter search on the affine augmentation learning task on RoMMNIST; we trained 200 models by first sampling of affine augmentation parameters uniformly at random from a predefined range" and then training a neural network model (that averages across augmentation samples at train and test time, as described in Benton et al. (2020)) with standard neural training with those hyperparameters fixed throughout.
In Figure we plot the final test-set performance of all the models trained with those hyperparameters sampled from a fixed range. Alongside, we show the hyperparameters and test-set performance of the partitioned networks as they progress throughout training. The partitioned networks consistently achieve final test-set performance as good as that of the best hyperparameter configurations iden- tified through extensive random sampling of the space. We also show the test-set performance of neural network models, trained through standard training, with hyperparameters fixed to the final hyperparameter values identified by the partitioned networks. The hyperparameters identified by partitioned networks appear also be good for regular neural networks; the standard neural networks with hyperparameters identified through partitioned training also outperform the extensive random sampling of the hyperparameter space. Furthermore, Figure 9 shows that partitioned networks do learn full rotation invariance on the RotMNIST task, i.e. when full rotation invariance is present in the data generating distribution.
0.99
0.88
0.97
0.36
3.0
0.2
0.4
0.5
0.2 Therelaikes
n2
0.0
0.2
344
0.5
0.2
2.4
0.5
0.0
0.2
0A
Tratulation
Rotation
Scale
Scale
Sheat
Randum Sampling
Pertilloned
Pertitioned
Rinel
Perlitioned Bune Fmnel- Betrained
Figure 9: The test-set performance plotted alongside (1D projections of) affine augmentation hyper- parameters on the RotMNIST task with MNIST-CNN. Final test-set accuracies are shown for the hyperparameters sampled randomly for neural network model trained through standard training with those hyperparameters fixed (+). For multiple partitioned networks runs, the plot shows the progres- sion of the identified hyperparameters and the test-set performance through the training run (--), as well as the final hyperparameters and test-set performance (0 o). Lastly, the plot also shows final test-set accuracies of models trained through standard training on the final hyperparameters identified through partitioned training
10 ranges were: Uniform(0, (0, T) for the maximum rotation, and Uniform(0, })) for all the remaining affine angmentation parameters (maximum shear, maximum and trannslatioon, and maximum x- and y- scale).
20

---

Published as a conference paper at ICLR 2023
H LIMITATIONS
As mentioned in the main text, our method improves upon existing work, but also comes with its limitations.
Complexity Inherent to our method - as presented in e.g. Figure is the necessity for an additional forward-backwand pass to update the hyperparameters. Consequently, hyperparameter optimization has additional costs which, however, are significantly less than the computational costs of existing work, as we discuss in more detail in Appendix D.1 and the experimental section. Furthermore, empirically, partitioned networks usually require more training iterations to converge.
Performance Assuming the optimal hyper-parameters are given, training the full, con-partitioned networks based on those optimal values can be expected to yield better performance compared to the final model found by partitioned training. Partitioning the network inherently constrains the network capacity, causing some loss of performance. Opportunities for alleviating this performance loss while still enjoying single-run hyperparameter optimization through partitioned training will be left to future work. These include for example adjusting training rounds or increasing network capacity in the first place.
Partitioning While partitioned networks allows for automatic optimization of, intuitively, hard to tune hyperparameters, such as augmentation parameters, they come with the additional limitation of requiring to partition both the data and the model. This introduces additional hyperparameter, namely, the partitioning strategy. While our default strategy of assigning more parameters and data to the first chunk works reasonably well on all of the experiments consider, if one targets obtaining the best possible performance on a given task, the partitioning strategy might need additional tuning. We provide some empirical results about the sensitivity partitioning in appendix F.1
I EXPERIMENTAL DETAILS
1.1 PARTITIONED TRAINING
Partitioned parameter update scheduling The gradient computation of Equation 3, as described in the main text, requires that the data-points for updating a given subnetwork ((()) come from the appropriate dataset chunks (ar, E Di.k for chunk k. Depending the partitioning scheme (Appendix C), evaluating different subnetworks for different chunks can or cannot be done a single mini-batch. More specifically, the random weight-partitioning requires a separate mini-batch per subnetwork (in order to keep the memory cost the same as for standard neural network training). An immediate question arising from chunked dataset and several partitions is to define the order and frequency of updates across subnetworks. In our experiments we define (non-uniform splits of the training dataset D across the chunks, which requires a tailored approach to sampling the data. More specifically, for a given (normalized) ratio of chunk-sizes (ui,...; Mc], each iteration of partitioned training proceeds as follows:
chose for our experiments
Sample a partition index k ~ Cat(ui...... uc)
2. Sample a mini-batch D of examples uniformly from Disk-
3. Evaluate log pDpp(Duuu ,)) using subnetwork (k) and
4. compute the (stochastic) gradient wrt. partition parameters uy (Eq. 3).
5. Update partition parameters wy using an optimizer, such as SGD or Adam.
This sampling scheme results in a data-point (æ. y) E DD from earlier chunks to be sampled more often. Concretely, the probability that example in chunk will be sampled is Eise u4 This done that each partition wi is updated with equal probability each of the examples in Di:k As result, we use with replacement sampling for the partitioned network training throughout the experimental section.
21

---

Published as a conference paper at ICLR 2023
Gradient optimization of partitioned parameters A consequence of per- per-partition updates with the random weight partitioning scheme (appendix C) is that, for a chosen partition wy to update, all other partitions do not receive gradient update. In other words, the gradient at each iteration is sparse. Consequently, many off-the-shelve momentum-based optimizers will not account correctly. Specifically, implement modifications to the PyTTrch Paszke al. (2019) provided optimizers that allow us to track per-partition momenta, number of steps, etc. Note that this creates disconnect between the number of iterations across all partitions and the number of iterations per-partitle per-partition. Doing so, however aligns the computational cost of training the partitioned network parameters with the cost of training regular neural network parameters. Regardless, we do not alter the way learning-rate schedulers behave in our experiments and anneal learning-rates according to the total number of iterations. Similarly, we report the total number of iterations when comparing against baselines that update all setwork-parameters per iteration.
While simple gradient-accummlatin scheme across mini-batches would result in a single gradient across all partitions, this approach inherently clashes with non-uniform partitioning [ui...., uc]. Instead, we chose to sequentially apply gradients computed on a single partition, as described in the previous paragraphs. A further advantage of this approach is that learning progress made by updating partition Wy immediately influences (and can improve) the prediction of subnetworks (k) (k+1)
w w.
w (C)
Gradient optimization of hyperparameters Our partitioned network scheme makes it easy to compute stochastic gradients of the hyperparameter objective LML in Eq. 4 using batch gradient descent optimization methods. After every update to randomly sampled network partition (see previous paragraph), we update hyperparamters & as follows:
sample a dataset chunk index k: Cat(" :.... ye D1.
Ratios are re-normalized to exclude
sample mini-batch D of examples uniformly from Dr (Note the choice of Dr. instead of D1.k).
Evaluate log ( (k-1) vt) using subnetwork w² (k-1) and
compute the (stochastic) gradient wrt. hyperparameters 1 (Eq. 4).
Update partition parameters v using an optimizer, such SGD or Adam.
The above sampling procedure yields an unbiased estimate of gradients in eq. 4.
The fact that we optimize hyperparameters with gradients based on data from single chunk time is again a consequence of the random weight- partitioning scheme for the partitioned networks. It is possible to compute gradients wrt. v for mini-batches with examples from multiple chunks at a time. With the random weight partitioning scheme, this would result in an increased memory overhead. Lastly, we could also accumulate gradients from different chunks, similarly to Immer et al. (2022), and this would likely result in nower-varianee estimate per update
is also possible to reduce the computational overhead of evaluating two mini-batches per iteration (one for updates to Wk. one for v as we do in our experiments by interleaving hyperparameter updates at less frequent intervals. We leave an exploration of these design choices to future work. Throughout all experiments, except those in the federated settings (see section J). we use the same batch-size for the hyperparameter udpates as for the regular parameter updates.
Weight-decay For partitioned networks, whenever using weight-decay, we scale the weight decay for earlier partitions with the reciprocal of the number of examples in chunks used to optimize them, following the diagonal Gaussian prior interpretation of weight-decay. This makes the training compatible with the variational interpretation in Appendix B.
I.2 PARTITIONED AFFINE TRANSFORMATIONS
In Appendix C we described how realize partitioned versions of fully-connected and convolutional layers. Design choices for other parameterized network layers used in our experiments are described below.
22

---

Published as a conference paper at ICLR 2023
Normalization layers It is common- common-place in most architectures to follow normalization layer (such as BatchNorm (loffe and Szegedy, 2015), GroupNorm (Wu and He, 2018)) with an element- wise or channel-wise, affine transformation. Namely, such transformation multiplies its input by a scale vectors and adds a bias vector b: o = a + b. For random weight-partitioned networks, we parameterize such affine transformations by defining separate vectors (81,...-7 and bb, for each partition; the actual scale and bias used in a given subnetwork w! 8 (R) (I...() 8; and b.a -iee()....))) respectively, This ensures that the final affine transformation for each subnetwork w² depends on the parameters in the previous partitions [l,...,k k I.. Doing increases the parameter count for the partitioned networks in architectures that use those normalization layers by a negligible amount.
Scale and bias in FixUp networks The FixUp paper (Zhang et al., 2019) introduces extra scales and biases into the ResNet architecture that transform the entire output of the layers they follow. We turn these into "partitioned" parameters using the same scheme as that for scales and biases of affine transformations following normalization layers.
For partitioned networks. through-out the paper, match the proportion of parameters assigned to each partition A in each layer to the proportion data examples in the corresponding chunk Dj.
1.3 ARCHITECTURE CHOICES
Input selection experiments We use fully- fully-connected feed-forward neural network with 2 hidden layers of size [256, 256]. and with GeLU (Hendrycks and Gimpel, 2016) activation functions. We initialise the weights using the Kaiming uniform scheme (He al., 2015). For partitioned networks, we use the random-weigh partitioning scheme.
Fixup Resnet For all experiments using FixUp ResNets we follow Immer et al. (2022); Zhang et al. (2019), and use 3-stage ResNet with channel-sizes (16, 32, 64) per stage, with identity skip- connections for the residual blocks as described in He et al. (2016). The residual stages are followed by average pooling and final linear layer with biases. We use 2D average pooling in the residual branches of the downsampling blocks. We initialize all the parameters as described in Zhang et al. (2019).
Wide ResNet For all experiments using a Wide-ResNet-N-D (Zagoruyko and Komodakis, 2016), with being the depth and the width multiplier, we use a stage ResNet with channel-sizes (16D, 32D, 64D). We use identity skip-connections for the residual blocks, described in He et al. (2016), also sometimes known as ResNetV2.
ResNet-50 We use the "V2" version of Wide ResNet as described in (Zagoruyko and Komodakis, 2016) and replace BatchNormallizttoon with GroupNormalization using 2 groups We use the 'standard' with with D and three stages of & layers for 50-layer deep ResNet.
We use ReLU activations for all ResNet experiments throughout.
MNIST CNN For the MNIST experiments, we use the same architecture as Schwöbel et al. (2021)
illustrated in the replicated Table 9.
Table 9: CNN architecture for MNIST experiments
Layer
Specification
2D convolution
channels=20, kernel size=(5, 5), padding=2, activation=ReLU
Max pooling
pool size=(2, 2). stride=2
2D convolution
channels=50, kernel size=(5,5), padding=2, activation=Rel.U
Max pooling
pool size= 2). stride=2
Fully connected
units=500, ReL
Fully connected
units=50, activation=ReLU
Fully connected
units=10, activation=SSftmaa
23

---

Published as a conference paper at ICLR 2023
1.4 TRAINING DETAILS
Learning affine augmentations For the parametrization of the learnable affine augmentation strategies, we follow prior works for a fair comparison. More specifically, for our MNIST based setup we follow the parametrization proposed in Schwöbel et al. (2021) whereas for our CIFARIO based setup we use the generator parametrization from Immer et al. (2022).
Input selection experiments For the model selection (non-differentiable) input selection exper- iments, we train all variants with Adam with learning rate of 0.001 and batch-size of 256 for 10000 iterations. For both Laplace and partitioned networks, we do early stopping based on the marginal likelihood objective (CML for partitioned networks). We use weight-decay 0.0003 in both cases. For the post-hoe Laplace method, we use the diagonal Hessian approximation, following the recommendation in (Immer et al., 2021). For partitioned networks, we divide the data and parameters into 8 chunks of uniform sizes. We plot results averaged across 3 runs.
Mask learning for input selection experiment We use the same optimizer settings as for the input selection experiment. We train for 30000 iterations, and optimize hyperparameters with Adam with a learning rate of 0.001. We divide the data and parameters into 4 uniform chunks.
MNIST experiments We follow Schwöbel et al. (2021), and optimize all methods with Adam with a learning rate of 0.001, no weight decay, and a batch-size of 200. For the partitioned net- works and Augerino results, we use 20 augmentation samples. We use an Adam optimizer for the hyperparameters with a learning rate of 0.001 (and default beta parameters).
For Augerino MNIST, we use the "safe" variant, as otherwise the hyperparameters and the loss diverge on every training run. We elaborate on this phenomenon in Appendix E. Otherwise, we follow the recommended settings from (Benton et al., 2020) and Immer al. (2022), namely, regularization strength of 0.01, and a learning rate for the hyperparameters of 0.05.
For both MNIST and CIFAR experiments, we found it beneficial to allocate more data to either the earlier, or the later, chunks. Hence, we use 3 chunks with [80%, 10%, 10%] split of examples for all MNIST and CIFAR experiments.
CIFAR variations experiments We again follow Immer et al. (2022), and optimize all ResNet models with SGD with a learning rate of 0.1 decayed by a factor of 100x using Cosine An- nealing, and momentum of 0.9 (as is standard for ResNet models). We use batch-size of 250. We again use Adam for hyperparameter optimization with a learning rate of 0.001 (and default beta parameters). We train our method for [2400, 8000, 12000, 20000, 40000) iterations subsets [1000, 5000, 10000, 20000, 50000] respectively for CIFAR- 10, just as in (Immer et al., 2022). For all methods, used weight-decay of le 4. For partitioned networks, we increase the weight decay for earlier partitions with the square root of the number of examples in chunks used to optimize them. following the diagonal Gaussian prior interpretation of weight-decay. We use 3 chunks with [80%, 10%, 10%) split of examples.
For RoICIFAR- results, we noticed our method hasn't fully converged (based on training loss) in this number of iterations, and so we doubled the number of training iterations for the RoMMNIST results. This slower convergence can be explained by the fact that, with our method, we only update a fraction of the network parameters at every iteration.
Tiny Imagenet experiments Our experiments with Tinylmagenet (Le and Yang, 2015) closely follow the setting for the CIFAR-10 experiments described above. Images are of size 64x64 pixels, to be classified into one of 200 classes. The training-se set consists of 100000 images and we compare our method against baselines on subset of [10000, 50000, 100000] datapoints. For the standard version of Tinylmagenet, we train for [80000, 80000, 40000) steps respectively and for the rotated version of Tinylmagenet we train for 120000 steps for all subset sizes. We tuned no other compared to the CIFAR- 10 setup and report method's 's result for partitioning with [80%, 20%] across chunks after finding it to perform slightly better than a (80%, 10%, 10%) split across chunks in preliminary comparison.
24

---

Published as a conference paper at ICLR 2023
Fine-tuning experiments For the fine-tuning experiments in table 2, we trained FixUp ResNet- 14 on a subset of 20000 CIFARIO examples, while optimizing affine augmentations (following affine augmentations parameterization in (Benton et al., 2020)). We used the same optimizer settings as for all other CIFAR experiments, and trained for 80000 iterations, decaying the learning rate with Cosine Annealing for the first 60000 iterations. For fine-tuning of validation-set optimization models, we used SGD with same settings, overriding only the learning rate to 0,01. We tried learning rate of 0.01 and 0.001, and selected the one that was most favourable for the baseline based on the test accuracy.
We also tried training on the full CIFAR-10 dataset, but found that all methods ended up within standard error of each other when more than 70% of the data was assigned to the first chunk (or training set, in the case of validation set optimization). This indicates that CIFAR-10 is sufficiently larger that, when combined with affine augmentation learning and the relatively small ResNet- architecture used, using the extra data in the 2nd partition (or the validation set) results negligible gains.
1.5 DATASETS
Input selection synthetic dataset For the input selection dataset, we sample 3000 datapoints for the training set as described section 5, and we use a fresh sample 1000 datapoints for the test set.
RoMMNIST
Sometimes in the literature, RoMNNIST referes a specific subset of 12000 MNIST examples, whereas in other works, the full dataset with 60000 examples is used. In this work, following (Benton et al., 2020; Immer al., 2022) we use the latter.
J FEDERATED PARTITIONED TRAINING
In this section, we explain how partitioned networks can be applied to the federated setting, as well the experimental details.
J.1 PARTITIONED NETWORKS IN FL
In order to apply partitioned networks to the federated setting, we randomly choose partition for each client such that the marginal distribution of partitions follows a pre-determined ratio. A given chunk D therefore corresponds to the union of several clients' datasets. Analogous to how "partitioned training" is discussed in the main text and Appendix I, we desire each partition wp to be updated on chunks P1:k- Equation in the main text explains which data chunks are used compute gradients wrt. parameter partition Wy. An analogous perspective to this objective is visualized by the exemplary algorithm in Figure and asks which partitions are influenced (i,e., updated) by data from chunk Da: A data chunk D& is used to compute gradients wrt. partitions W through subnetworks (k) !h to (C) respectively. Consequently, client whose dataset is assigned chunk D3, can compute gradients for all partitions Wx.C.
Updating network partitions Due to the weight- sratht-pptttiooeeed construction of the partitioned neural networks, it is not possible to compute gradients with respect to all partitions in single batched forward-pass through the network. Additionally, a change to the partition parameters wi directly influences subnetworks w(++11 (R++) to w!!) In order to avoid the choice of ordering indices à to C for the client's local update computation, we update each partition independently while keeping all other partitions initialised to the server-provided values that the client received in that round t: Denote Dj, as the dataset of client where we keep index i to emphasize the client's 's assignment to chunk k. ++1 Further denote wj. as the partition w after having been updated by client on dataset D
(+1 arg max logp (wi, w/. w/j+1.-....... +)+)), Vj E [k; C},
(25)
where the details of optimization are explained in the following section. We leave explo- ration for different sequential updating schemes to future work. The final update communi- cated by client to the server consists of the concatenation of all updated parameter partitions
25

---

Published as a conference paper at ICLR 2023
".t concat( k,s ..... w ). Note that partitions (wf:.... = have not been modified and need not be communicated to the server. The resulting communication reductions make partitioned networks especially attractive to FL as data upload from client to server poses a significant bottleneck. In practice, we expect the benefits of these communication reductions to outweigh the additional computation burden of sequentially computing gradients wrt., to multiple partitions.
f++
The server receives w from all clients that participates in round t, computes the delta's 's with the global model and proceeds to average them to compute the server-side gradient in the typical federated learning fashion (Reddi al., 2020).
Updating hyperparameters The computation of gradients on a client wrt. vb is straight-forward extension of equation 4 and the exemplary algorithm of Figure 1:
V VyL (Di,e: 1b) z V log (Drxl wo Wt (111)) +b)
(26)
where Di, corresponds to client i's local dataset which assigned to chunk k and ws ((+1).(k-1)
corresponds to the (I= 1)'th subnetwork after incorporating all updated partitions Wa,i (+++))(k-1)
concat(wf (wi,.................... - w ++ w(,,) ). Note that we compute a full-batch update to & in MNIST experiments and use a batch-size equal to the batch-size for the partitioned parameter updates for
CIFARIO.
Upon receiving these gradients from all clients in this round, the server averages them to form a server-side gradient. Conceptually, this approach updating &p corresponds to federated SGD.
J.2 FEDERATED SETUP
Non-i.id. partitioning For our federated experiments, we split the 50k MNIST and 45k CIFAR10 training data-points across 100 clients in a non-i.i.d. way create the typical challenge federated learning experiments. In order to simulate label-skew, we follow the recipe proposed in Reddi et al. (2020) with = 1.0 for CIFARIO and have data corresponding to only a single digit. For our experiments on rotated versions of CIFARIO and MNIST, we sample a degree of rotation per data-point and keep it fixed during training. In order to create non-i.i.d partitioning across the clients, we bin data-points according to their degree of rotation into 10 bins and sample using the same technique as for label-skew with a 0.1 for both datasets. Learning curves are computed using the 10k MNIST and 5k CIFARIO validation data-points respectively. For the rotated dataset experiments, we rotate the validation set in the same manner the training set.
0.1 for MNIST. Note that with a 0.1, most clients
Architectures and experimental setup We use the convolutional network provided at Schwöbel et al. (2021) for MNIST and the ResNet-9 (Dys) model for CIFARIO but with group normaliza- tion (Wu and He, 2018 instead of batch normalization. We include (learnable) dropout using the continuous relaxation proposed at Maddison al. (2016) between layers for both architectures. We select 3 chunks for MNIST with a [0.7, 0.2, 0.1] ratio for both, client-assignments and parameter- partition sizes. For CIFAR1O, we found a [0.9, 0.1] split across 2 sub-networks to be beneficial. In addition to dropout logits, v encompasses parameters for affine transformations, i.e., shear, trans- lation, scale and rotation. We report results after 2k and 5k rounds, respectively, and the expected communication costs as a percentage of the con-partitioned bascline.
Shared setting In order to elaborate on the details reproduce our results, we first focus on the settings that apply across all federated experiments. We randomly sample the corresponding subset of 1.25k, 5k data-points from the full training set and keep that selection fixed across experiments (i,e,, baselines and partitioned networks) as well as seeds. The subsequent partitioning across clients as detailed in the previous paragraph is equally kept fixed across experiments and seeds. Each client computes updates for one epoch of its local dataset, which, for the low data regimes of 1.25% data-points globally, results in single update per client using the entire local dataset. We averaged over 10 augmentation samples for the forward pass in both training and inference.
MNIST & RotMNIST For 5k data-points and correspondingly 50 data- aatatpoints on average per client, most clients perform single update step. A small selection of clients with more than 64 data-points
26

---

Published as a conference paper at ICLR 2023
performs two updates per round. For the experiments using the full dataset and mini-batch size of 64, each client performs multiple updates per round. After initial exploration on the baseline FedAvg task, select local learning-rate of 5e 2 and apply standard SGD. The server performs Adam Reddi et al. (2020) with a learning rate of le 3 for the model parameters. We keep the other parameters of Adam at their standard PyTorch values. We find this setting to generalize the partitioned network experiments but found higher learning rate of 3e - for the be helpful. We chose the convolutional network from Schwöbel al. (2021) with (leamed) dropout added between layers. The model's dropout layers are initialized to drop 10% of hidden activations. For the baseline model keep the dropout-rate fixed and found 10% be more stable than 30%.
to
CIFARIO & RoICIFAR10 We fix mini-batch size of 32, leading to multiple updates per client per round both, the full dataset regime as well as the 56 data-points setting. Similarly to the MNIST setting, performed initial exploration of hyperparameters
the baseline FedAvg task and use the same ones on partitioned networks. We used dropout on the middle layer of each block which was initialized 0.1 for both the baseline and partitioned networks and whereas partitioned networks optimized with LM and the concrete relaxation from Maddison et al. (2016), the baseline kept it fixed. For the server side optimizer used Adam with the default betas and learning rate of le 2. whereas for the hyperparameters we used Adam with the default betas and a learning rate of le In both cases used an Le For the local optimizer we used SGD with learning rate of 10-0.5 and no momentum.
J.3 MNIST LEARNING CURVES
In Figure 10 we show learning curves for the three considered dataset sizes on the standard MNIST task. Each learning curve is created by computing moving average across 10 evaluations, each of which is performed every 10 communication rounds, for each seed. We then compute the average and standard-error across sees and plot those values on the y-axis. On the x-axis we denote the total communication costs (up- and download) to showcase the partitioned networks reduction in communication overhead. We see that especially for the low dataset regime, training has converged yet and we expect performance to improve for an increased number of iterations.
Validatton accuracy
38-
Velidetion Accuretcy
Veldestio Pecurecy
38
90-
96-
04-
94
64-
Fedfug + PHML recfvo
redfug PNML
redfog PNNL
hes/wo
92
32
100
150
100
50
100
Figure 10: Learning curves for MNIST experiments on 1.25k, 5k and 50% data-points respectively.
27


Reviews:
[{"review_id": "tokeUwnL8_", "reviewer": "Reviewer_Na31", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper proposes an algorithm that partitions both the data and model parameters to efficiently approximate the log marginal likelihood. The key idea is to define the out-of-training-sample loss with the partitioned network and utilize this objective for hyperparameter optimization. Empirically, the authors show that the proposed method can adapt various hyperparameters in neural networks without a validation set. Moreover, the proposed methods can be utilized in federated learning, where cross-validation is difficult.\n", "summary_of_the_review": "Overall, I believe that the idea presented in this paper is interesting and novel and the paper is well-written. While I have concerns about the diversity of the baseline in the experiment section and the lack of computation and timing experiments. At the current state, I recommend 6.", "strength_and_weaknesses": "Strength\n- The paper is well-written and clearly motivated. I believe that the idea of explicit network partitioning is interesting and novel and the work is relevant to the ICLR community.\n- The technical details of the paper look correct. \n- The experiment setup is detailed making it easy to implement and reproduce the results.\n\nWeakness\n- While the paper argues that the proposed method is more efficient compared to previous methods such as Laplace approximation, it is not empirically shown in Section 5. I am happy to increase the score if this concern is addressed.\n- The proposed method is not compared with a competitive hyperparameter optimization baseline (that uses a validation set). While I understand that the proposed framework does not require partitioning the train and validation set, it would be helpful to understand how the method compares to other strong hyperparameter optimization baselines. \n- The authors do not sufficiently describe the limitation of the proposed approach. What are some additional hyperparameters and how sensitive are they? Can the method be more scalable (in terms of both parameter and hyperparameter) compared to the baseline methods?", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "jCHewcX1nkr", "reviewer": "Reviewer_ic1t", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The work starts with using marginal likelihood as an alternative objective to validation metric for HPO, which has advantages in the small data region and potentially better generalization performance. It then talked about how to compute marginal likelihood and proposed an efficient method that only requires a single training run and no validation set. The experiments on synthetic and a diverse set of real world experiments (data argumentation, classical HPO setting and Federated Learning) demonstrated its usefulness.", "summary_of_the_review": "The paper has many merits and they are mentioned already in the strength section. But the  weaknesses, so far, are also clear from my perspective. Thus, without more clarifications or justifications for the proposed method, I am not confident to accept the paper.", "strength_and_weaknesses": "On the strength of the work, using marginal likelihood as a HPO objective is very interesting and I can see practical settings that marginal likelihood can be useful, such as in small data regime, federated learning or potentially a better objective for general HPO. On the methodology, it is a bold exploration that partitioning the networks into subnetworks and training them separately over data chunks. On the empirical evaluations, it covers a wide range of tasks from different areas. It is also a well written paper, with clear structures and abundant details for reproducibility.\n\nOn the weakness, my first concern is motivation. The author mentioned partitioning the parameters, in the beginning of page 4, “can be a reasonable approximation in function space”. Why is that? Also, one thing I am not sure, if the goal is to approximate eq(2) and avoid the model sees every datapoint in the training set after 1 epoch, why one can’t create $C$ copies of the whole parameters and assign each copy to one $D_k$, without even partitioning the parameter (update all the parameter with related data chunk in each copy)? \n\nThen it is the support for the claims. The authors argued the proposed method “is more scalable than previous works that rely on marginal likelihood and Laplace approximations”. I can see why it is more scalable (avoid inverting a Hessian (Immer et al., 2021)), but I don’t see any empirical results supporting it. \n\nSince the main contribution is a new method to approximate the marginal likelihood, the comparison to relevant baselines, for example, Schwobel et al. (2021) and Immer et al., (2022), will be critical to demonstrate its usefulness. But they are not all the times included in all the experiments. For example, in the experiments of learning invariances through data-augmentations Schwobel et al. (2021) is not in Table 1 and Immer et al., (2022) not in Fig 3(a)? There are no baselines (Schwobel et al. (2021) and  Immer et al., (2022)) in the experiments for “Comparisons to traditional training / validation split”. It is hard to attribute the improvement to the partitioned network (a particular way to approximate learning speed) or the learning speed objective in general. Similarly, no other HPO baselines are included in the Federated Learning experiments. Are there some particular reasons why these results are missing? If so, it would be nice to explain them in the paper.\n\nIn the end, I am wondering about the impact of partitioning? How sensitive are the results to a partition?\n\n\n=== After reading authors' response ===\n\nThe authors' response answered all my questions and I increased the score to 6.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "9rRO-i7SxD", "reviewer": "Reviewer_Pdfa", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The authors present a method to tune hyperparameters which requires only a single training run and no validation data. This is achieved by partitioning the weights of a neural network and train each partition with respect to only its own training partition. Then, the remainder of the training data can be used for validation. The authors demonstrate the effectiveness of this approach for hyperparameter optimization in several scenarios including federated learning.", "summary_of_the_review": "This is a well-written paper on an important scientific problem with important practical applications. The authors discuss related work in detail and demonstrate the effectiveness of the method on a couple of problems. Overall, I think the paper could be a great addition for the conference.\n\n**After Rebuttal Phase**\nOverall the paper was interesting and seemed novel to me, and therefore I was positive about it. However, based on the comments of some other reviewers, it might not have been very novel (I am not familiar with that related work) and therefore I've updated my confidence and recommendation accordingly.", "strength_and_weaknesses": "**Strengths**\n- Good discussion of related work\n- Diverse set of experiments\n- Relevant problem, great motivation, well-written paper\n\n**Weaknesses**\n- Missing results to support the claim that the proposed method is computationally cheaper", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "1Uq8IjhD2n_", "reviewer": "Reviewer_Wybc", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The main contribution is to tune differentiable “hyperparameters” by splitting the trainable parameters and dataset in a way that separates the data seen for tuning hyperparameters, ensuring that the set of hyperparameters are generally good for a wide range of parameters. The splitting technique is inspired by the marginal likelihood.", "summary_of_the_review": "The main method proposed in the paper is limited in its novelty and shaky in its theoretical groundings. The paper does not clearly state the main benefits of using this method and such empirical benefits are not convincingly displayed in the experiments, especially in the \"toy model\" which are designed to clearly demonstrate the superiority of the new method over baselines.", "strength_and_weaknesses": "The paper mainly proposes a technique of splitting the training/validation dataset in an iterative manner, inspired by the marginal likelihood decomposition for cross validation, which is claimed to be an optimization objective that requires no validation data. \n\nWhile marginal likelihood is standardly used for Bayesian model selection, the choice of Bayesian prior/function class strongly influences the generalization capabilities of these methods since almost all loss functions, such as the L2/cross entropy can be derived by using Bayesian priors. In fact, a recent paper [https://arxiv.org/abs/2202.11678] demonstrated the pitfalls of using marginal likelihood without validation. Since the marginal likelihood is the main theoretical justification for the method design, I find the justification extremely weak so giving a theoretical example with concretely define Bayesian priors/assumptions would be better.\n\nThe novelty of the paper is also limited in the practical algorithm. The method of partitioning the model and parameters are reminiscent of cross validation and drop out techniques, which have both been readily used in the ML community. Adding a baseline of using dropout + cross validation would be better, as well as spending more time of using the toy example (of 1/2 spurious variables) to illustrate WHY your method does better than all the \"standard\" baselines. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster