Hyperparameter Optimization through Neural Network Partitioning
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.
Bruno Mlodozeniec, Matthias Reisser, Christos Louizos
# Hyperparameter Optimization

###### 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.

## 1 Introduction

Due to their remarkable generalization capabilities, deep neural networks have become the de-facto models for a 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 as 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 to optimize hyperparameters operates as follows; **1)** partition the dataset into training and validation data1, **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 a hyperparameter in itself. The conventional rule of thumb to use \(\sim 10\%\) of all data can result in significant overfitting, as pointed out by Lorraine et al. (2019), when one has a sufficiently large number of hyperparameters to tune. Furthermore, a validation set can be challenging
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 a 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 et al., 2018) and learning data-augmentations with schemes as in Benton et al. (2020); van der Wilk et al. (2018)). 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 (Schwobel 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 a 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 on marginal likelihood and Laplace approximations (which require computing or inverting a Hessian (Immer et al., 2021)) and is broadly applicable to any hierarchical modelling setup.

## 2 Marginal Likelihood and prior work

In Bayesian inference, the rules of probability dictate how any unknown, such as parameters \(\mathbf{w}\) or hyperparameters \(\psi\), should be determined given observed data \(\mathcal{D}\). Let \(p(\mathbf{w})\) be a prior over \(\mathbf{w}\) and \(p(\mathcal{D}|\mathbf{w},\psi)\) be a likelihood for \(\mathcal{D}\) with \(\psi\) being the hyperparameters. We are then interested in the posterior given the data \(p(\mathbf{w}|\mathcal{D},\psi)=p(\mathcal{D}|\mathbf{w},\psi)p(\mathbf{w})/p(\mathcal{D}|\psi)\). The denominator term \(p(\mathcal{D}|\psi)\) is known as the _marginal likelihood_, as it measures the probability of observing the data given \(\psi\), irrespective of the value of \(\mathbf{w}\): \(p(\mathcal{D}|\psi)=\int p(\mathbf{w})p(\mathcal{D}|\mathbf{w},\psi)d\mathbf{w}\).

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, SS 28). In the PAC-Bayesian literature, it has been shown that higher marginal likelihood gives tighter frequentist upper bounds on the generalization performance of a given model class (McAllester, 1998; Germain et al., 2016). It also has close links to cross-validation (see section 2.1) 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; Schwobel et al., 2021).

### "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 \(\mathcal{D}\) conditioned on some hyperparameters \(\psi\) can be written as:

\[\log p(\mathcal{D}|\psi)=\sum_{k}\log\mathbb{E}_{p(\mathbf{w}|\mathcal{D}_{1:k-1},\psi)}\left[p(\mathcal{D}_{k}|\mathbf{w},\psi)\right]\geq\sum_{k}\mathbb{E}_{p( \mathbf{w}|\mathcal{D}_{1:k-1},\psi)}\left[\log p(\mathcal{D}_{k}|\mathbf{w},\psi)\right] \tag{1}\]

where \((\mathcal{D}_{1},\ldots,\mathcal{D}_{C})\) is an arbitrary partitioning of the training dataset \(\mathcal{D}\) into \(C\) shards or chunks2, and \(p(\mathbf{w}|\mathcal{D}_{1:k},\psi)\) is the posterior over parameters of a function \(f_{\mathbf{w}}:\mathcal{X}\rightarrow\mathcal{Y}\), from the input domain \(\mathcal{X}\) to the target domain \(\mathcal{Y}\) after seeing data in shards \(1\) through \(k\). The right-hand side can be interpreted as a type of cross-validation in which we fix an ordering over the shards and measure the "validation" performance on each shard \(\mathcal{D}_{k}\) using a model trained on the preceding shards \(\mathcal{D}_{1:k-1}\)
Alternatively, it can be viewed as the _learning speed_ of a (probabilistic) model: _i.e._, a 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 \(\psi\), \(\mathbf{w}\) and \(f_{\mathbf{w}}\). Namely, such models can be expected to learn faster and generalize better after seeing fewer samples. For example, if the hypothesis space is constrained3 to 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 in 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.

Footnote 3: or if the learning algorithm is heavily biased towards returning hypotheses that satisfy a given invariance, _e.g._, through the use of a prior.

### Training speed for hyperparameter optimization

Computing the "learning speed", requires samples from the posterior \(p(\mathbf{w}|\mathcal{D}_{1:k},\psi)\). 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. 1 by looking at the _training speed_ during standard training of a neural network by SGD. Specifically, they define the training speed as the reduction in the training loss after a 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, \(\mathbf{w}\), have been updated with SGD steps using data from shards \(\mathcal{D}_{1:k}\), they can be approximately used in place of the sample from the posterior \(p(\mathbf{w}|\mathcal{D}_{1:k},\psi)\) 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 \(1\) 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.

For the sake of scalability and alignment with deep learning practice, we also focus on simple pointwise approximations \(q_{k}(\mathbf{w})=\delta(\mathbf{w}=\hat{\mathbf{w}}_{k})\) to the posteriors \(p(\mathbf{w}|\mathcal{D}_{1:k},\psi)\). However, in contrast to prior work, we explicitly parametrize the learning procedure such that, at any given training iteration, we have access to a model that is trained only on a subset of the data \(\mathcal{D}_{1:k}\). In doing so, we can approximate the objective in Eq. 1, and thus use it to optimize the hyperparameters during the entire training run.

## 3 Partitioned Neural Networks

Our goal is to optimize the objective

\[\mathcal{L}_{\mathrm{ML}}\left(\mathcal{D},\psi\right)=\sum_{k=1}^{C}\mathbb{ E}_{q_{k-1}(\mathbf{w})}\left[\log p(\mathcal{D}_{k}|\mathbf{w},\psi)\right] \tag{2}\]

wrt. \(\psi\), which is an approximation to the lower-bound presented in Eq. 1 above. In Appendix A, we show that the left-hand side is also a lower-bound on the marginal likelihood under some unobtrusive conditions. As mentioned in Section 2.2, our goal is to propose an architecture and a training scheme so that we can easily obtain models trained on only subsets of the data \(\mathcal{D}_{1:k}\) for all \(k\) throughout training. We propose that each \(\{q_{k}(\mathbf{w})\}_{k=1}^{C}\) optimizes a 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 \(\mathbf{w}\in\mathbb{R}^{N}\). We can define a partitioning \(((\mathbf{w}_{1},\dots,\mathbf{w}_{C}),P)\) of the parameters into \(C\) partitions, such that \(\mathbf{w}=P\operatorname{concat}(\mathbf{w}_{1},\dots,\mathbf{w}_{C})\) for a permutation matrix \(P\in\{0,1\}^{N\times N}\). For ease of exposition, we drop the dependence on \(P\), assuming that \(\mathbf{w}\) is already arranged such that \(P\) is identity, \(P=I_{N\times N}\).

Given the partitioning \((\mathbf{w}_{1},\dots,\mathbf{w}_{C})\) of the parameters, we then specify \(C\) subnetworks with weights \(\mathbf{w}_{s}^{(1)},\dots,\mathbf{w}_{s}^{(C)}\) such that \(\mathbf{w}_{s}^{(k)}=\operatorname{concat}(\mathbf{w}_{1},\dots,\mathbf{w}_{k},\hat{\mathbf{w} }_{k+1},\dots,\hat{\mathbf{w}}_{C})\), where \(\hat{\mathbf{w}}_{i}\) are some default
values not optimized during training4. More specifically, the \(k\)-th subnetwork, \(\mathbf{w}_{s}^{k}\), retains the first \(k\) partitions from the weight partitioning and sets the remaining parameters to \(\hat{\mathbf{w}}_{k+1:C}\). Note that, if each \(\mathbf{w}_{k}\) is only updated on chunks \(\mathcal{D}_{1:k}\), the subnetwork \(\mathbf{w}_{s}^{(k)}\) is only comprised of weights that have been updated on \(\mathcal{D}_{1:k}\). Thus, we can view the parameters of \(\mathbf{w}_{s}^{(k)}\) as an approximation to \(q_{k}(\mathbf{w})\). Although, given that a subset of the parameters in each \(\mathbf{w}_{s}^{(k)}\) is fixed, this would likely be a poor approximation to the true posterior over the weights given \(\mathcal{D}_{1:k}\), it could be, intuitively, a reasonable approximation in function space5.

Footnote 4: _e.g._, \(\hat{\mathbf{w}}_{i}\) could be the value of the weights at initialization, or \(\hat{\mathbf{w}}_{i}=\mathbf{0}\) corresponding to pruning those parameters and obtaining a proper subnetwork.

Footnote 5: Since a) the mapping from parameters to functions is not bijective and b) neural networks are highly overparameterised and can be heavily pruned while retaining performance (Frankle and Carbin, 2018), obtaining a good fit to a subset of the training data with a 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 et al., 2020), thus it is reasonable for subnetworks fit to more data to have more learnable parameters.

**Partitioned training** Having partitioned the dataset \(\mathcal{D}\) into \(C\) chunks \((\mathcal{D}_{1},\ldots,\mathcal{D}_{k})\), we update each partition \(\mathbf{w}_{k}\) by optimising the negative log-likelihood6 on chunks \(\mathcal{D}_{1:k}\) using subnetwork \(\mathbf{w}_{s}^{(k)}\) by computing the following gradients:

Footnote 6: Optionally with an added negative log-prior regularization term \(\log p(\mathbf{w}_{s}^{(k)})\).

\[\nabla_{\mathbf{w}_{k}}\mathcal{L}\left(\mathcal{D}_{1:k},\mathbf{w}_{s}^{(k)}\right)= \sum_{(\mathbf{x},y)\in\mathcal{D}_{1:k}}\nabla_{\mathbf{w}_{k}}\log p\left(y\Big{|} \mathbf{x};\mathbf{w}_{s}^{(k)},\psi\right). \tag{3}\]

We interleave stochastic gradient updates of each partition of the weights with updating the hyperparameters \(\psi\) using \(\mathcal{L}_{\mathrm{ML}}\) in Eq. 2:

\[\nabla_{\psi}\mathcal{L}_{\mathrm{ML}}\left(\mathcal{D},\psi\right)\approx \sum_{k=2}^{C}\sum_{(\mathbf{x},y)\in\mathcal{D}_{k}}\nabla_{\psi}\log p\left(y \Big{|}\mathbf{x},\mathbf{w}_{s}^{(k-1)},\psi\right). \tag{4}\]

This can be seen as the sum of the _out-of-sample_ losses for each subnetwork \(\mathbf{w}_{s}^{(k)}\). 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 \(\mathbb{E}_{q_{0}(\mathbf{w})}[\log p(\mathcal{D}_{1}|\mathbf{w},\psi)]\) in Eq. 4, we chose to leave it out. Hence, the gradient of Eq. 4 is of an estimate that can be viewed as an approximation to the _conditional_ marginal likelihood \(\log p\left(\mathcal{D}_{2:C}|\mathcal{D}_{1},\psi\right)\). 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).

This procedure, inspired by the marginal likelihood, has several desirable properties compared to prior work. **1)** Our objective is computationally efficient, with a computational cost roughly corresponding to evaluating subnetworks on the training set. There is no need to compute nor invert a Hessian with

Figure 1: Best viewed in colour. Illustration of the partitioning scheme for a single hidden layer perceptron with \(C=3\) chunks.
respect to the weights, as in the Laplace approximation (Immer et 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 a 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 as 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 at random. For each subnetwork, for the weight partitions corresponding to future chunks, we use the values of the weights at initialisation. For a discussion of partitioning schemes, see Appendix C.

## 4 Related works

Hyperparameter optimization in deep learningMany 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, or other hyperparameters that affect the final neural network weights only through the training routine. Other works have proposed ways to parameterize and optimize data-augmentations (Cubuk et al., 2018; Li et al., 2020), search-spaces for neural network architectures, as 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 on optimizing hyperparameters on a separate validation set and 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; Schwobel et al., 2021; Nabarro et al., 2022; Immer et al., 2022). We compare against Benton et al. (2020) ("Augerino") and Immer et al. (2022) ("Differentiable Laplace") on this task in the experimental section.

Hyperparameter optimization without a validation setA limited number of works consider learning hyperparameters without a validation set in a deep learning context. Benton et al. (2020) propose a simple method for learning invariances without a 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 _a 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-augmentation transformations were to be parameterized by a neural network (as proposed in Lorraine et al. (2019)), it is 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. Schwobel et al. (2021) propose only marginalising-out the parameters in the last layer of the neural network by switching it out for a Gaussian Process. They treat the preceding layer effectively as a 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-10" (Schwobel et al., 2021, SS7) and highlighted the need to 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 on datasets like CIFAR-10. Immer et al. (2022) have adapted the Laplace approximation (Immer et al., 2021) to make it tractable for learning data augmentations. In contrast to Schwobel et al. (2021), they approximately marginalize out all the network parameters, 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 to a ResNet-14, whereas our method can easily scale to larger architectures.

Hyperparameter optimization in FLImproving hyperparameter optimization is especially relevant 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
hyperparameters. FedEx relies on a training/validation split on the client level and uses a REINFORCE 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 to partitioned networks, which use standard, low-variance backpropagation for the hyperparameters and no separate validation set per client. To optimize the other hyperparameters, Khodak et al. (2021) wrapped FedEx with a traditional hyperparameter optimization strategy, the successive halving algorithm. This is orthogonal to our method and could be applied to partitioned networks as well. In Zhou et al. (2021), the authors perform a hyperparameter search independently on each client with some off-the-shelf methods and then aggregate the results of the search at the server once in order to identify the best hyperparameter setting. The main drawback of this method compared to partitioned networks is that when the local client datasets are small, a 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

Input SelectionTo demonstrate that \(\mathcal{L}_{\mathrm{ML}}\) 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\sim\mathrm{Bern}\left(\frac{1}{2}\right)\qquad\mathbf{x}=\big{[}\underbrace{y+ \epsilon_{1},\ldots,y+\epsilon_{15}}_{\text{Informative}},\underbrace{\epsilon_ {16},\ldots,\epsilon_{30}}_{\text{Spurious}}\big{]}^{\intercal}\qquad \epsilon_{1},\ldots,\epsilon_{30}\stackrel{{\text{iid}}}{{ \propto}}\mathcal{N}(0,1).\]

We specify a fixed mask over the inputs prior to training, where the first \(K\) 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) on a training set of \(1000\) examples using our method and compare the final values of the \(\mathcal{L}_{\mathrm{ML}}\) objective. The results are shown in Figure 1(a). \(\mathcal{L}_{\mathrm{ML}}\) correctly identifies \(15\) input features as the optimum, and correlates well with test accuracy and log-likelihood. Training loss and training accuracy, on the other hand, cannot alone disambiguate whether to use \(15\) or more input features.

Differentiable input selectionWe further show that we can learn the correct mask over the inputs in a differentiable manner using our method during a single training run. We parameterize a learnable mask over the inputs with a concrete Bernoulli distribution (Maddison et al., 2016) and treat the parameters of the mask distribution as a hyperparameter. We optimize them with respect to the proposed objective using our method. The evolution of the learned mask during training is shown in Figure 1(b), where we see that we can correctly identify the first 15 informative features.

Figure 2: (a) Demonstrating the ability of the marginal-likelihood inspired objective \(\mathcal{L}_{\mathrm{ML}}\) to identify the correct model on a toy input selection task. We plot the hyperparameter objective, train —
Learning invariances through data-augmentationsFollowing previous literature on learning soft invariances through learning data augmentations (Nabarro et al., 2022; van der Wilk et al., 2018; Benton et al., 2020; Schwobel et al., 2021; Immer et al., 2022), we show that we 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 a model hyperparameter rather than a training hyperparameter. For datasets, we consider MNIST, CIFAR10, TinyImagenet along with rotCIFAR10 and rotTinyImagenet, variants where the datapoints are randomly rotated at the beginning of training by angles sampled uniformly from \([-\pi,\pi]\)(Immer et al., 2022). Experimental setup details are provided in Appendix 1.

For the CIFAR10 and rotCIFAR10 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 \(\mathrm{\SIUnitSymbolMicro g}_{\mathrm{W}}\)ResNets (Zhang et al., 2019) for the architectures. The results can be seen in Table 1. There, we observe that partitioned networks outperform all baselines in the case of CIFAR10 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)TinyImagenet experiments we consider a ResNet-50 architecture with GroupNorm(2). In Table 1 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 (TinyImagenet) or are similar (rotTinyImagenet).

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 can 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 a baseline that does not learn augmentations, while performing favorably against prior works in most cases.

On MNIST, our method outperforms the last-layer marginal-likelihood method (last-layer ML) by Schwobel et al. (2021) in the large data regime but underperforms in the low-data regime. That is likely to be expected, as their work fits a Gaussian Process (GP) at the last layer (Wilson et al., 2016), which is better tailored for the low-data regime and results into a 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\) different networks where network \(k\) is fit on data \(\mathcal{D}_{1:k}\). 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 RotCIFAR10, 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 rotTinyImagenet experiments we observe that partitioned networks either outperform or are similar to the baselines on all data sizes considered.

\begin{table}
\begin{tabular}{l l l l l l} \hline \hline  & & & \multicolumn{3}{c}{Method} \\ Dataset & Architecture & Baseline & Augerino & Diff. Laplace & Partitioned \\ \hline RotCIFAR10 & \(\mathrm{\SIUnitSymbolMicro g}_{\mathrm{W}}\)ResNet-8 & \(54.2_{\pm 0.4}\) & \(75.4_{\pm 0.2}\) & \(\mathbf{79.5_{\pm 0.6}}\) & \(\mathbf{79.1_{\pm 0.0}}\) \\ \hline CIFAR10 & \(\mathrm{\SIUnitSymbolMicro g}_{\mathrm{W}}\)ResNet-8 & \(74.1_{\pm 0.5}\) & \(79.0_{\pm 1.0}\) & \(84.2_{\pm 0.8}\) & \(\mathbf{86.1_{\pm 0.4}}\) \\  & \(\mathrm{\SIUnitSymbolMicro g}_{\mathrm{W}}\)ResNet-14 & \(79.5_{\pm 0.3}\) & \(83.0_{\pm 0.1}\) & \(88.1_{\pm 0.2}\) & \(\mathbf{89.1_{\pm 0.8}}\) \\ \hline RotTinyImagenet & ResNet-50 & \(31.5_{\pm 0.6}\) & \(\mathbf{44.5_{\pm 0.2}}\) & OOM\({}^{\prime}\) & \(43.9_{\pm 0.3}\) \\ \hline TinyImagenet & ResNet-50 & \(44.2_{\pm 0.5}\) & \(41.1_{\pm 0.2}\) & OOM & \(\mathbf{48.6_{\pm 0.0}}\) \\ \hline \hline \end{tabular}
\end{table}
Table 1: Test accuracy with learning affine augmentations on (rot)CIFAR10 and (rot)TinyImagenet.
Comparisons to traditional training / validation splitWe further perform comparisons between partitioned networks and the more traditional training/validation split (denoted as validation set optimization) with additional finetuning to the task of learning data augmentations. This is realized as follows; we partition \(20k\) CIFAR10 examples into training and validation data of specific proportions. We then either train a partitioned network (along with the hyperparameters on \(\mathcal{L}_{\mathrm{ML}}\)) on 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 is finetuned on the whole dataset for \(20\) epochs. The results for varying chunk proportions are provided in Table 2.

We can see that partitioned networks (that do not employ additional finetuning) outperform validation set optimization with finetuning in all settings we tried. The gap does get smaller when we move to the more traditional \(90\)/\(10\) splits for training/validation: a \(10\%\) proportion for validation data is enough to optimize a handful of hyperparameters (just \(6\) scalars). To corroborate this claim, we set up an additional experiment; we use a Wide ResNet-20 on the full CIFAR10 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 is 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 to this setting is straightforward. To apply Augerino, one would have to come up with a metric that can be used to regularize the feature extractor towards "higher invariance".

Partitioned networks for federated learningWe consider federated learning (FL) (McMahan et al., 2017), a setting where data is distributed across many clients. In this setting, there are system properties that make hyperparameter optimization especially challenging (Wang et al., 2021). More specifically, obtaining a validation set and performing multiple training runs with different

\begin{table}
\begin{tabular}{l c c c c c} \hline \hline  & & \multicolumn{4}{c}{Chunk Proportions} \\ Method & \([0.3,0.7]\) & \([0.5,0.5]\) & \([0.7,0.3]\) & \([0.8,0.2]\) & \([0.9,0.1]\) \\ \hline Partitioned & \(\mathbf{82.9\%_{\pm 0.3}}\) & \(\mathbf{83.0\%_{\pm 0.01}}\) & \(\mathbf{83.7\%_{\pm 0.2}}\) & \(\mathbf{84.0\%_{\pm 0.6}}\) & \(\mathbf{84.6\%_{\pm 0.05}}\) \\ Validation set optim. & NaN & \(78.9\%_{\pm 0.04}\) & \(81.5\%_{\pm 0.2}\) & \(82.6\%_{\pm 0.1}\) & \(83.4\%_{\pm 0.1}\) \\ \(\blackleftarrow\)+Finetune & NaN & \(81.3\%_{\pm 0.09}\) & \(82.5\%_{\pm 0.2}\) & \(83.5\%_{\pm 0.1}\) & \(83.8\%_{\pm 0.3}\) \\ \hline \hline \end{tabular}
\end{table}
Table 2: Learning affine augmentations with \(\frac{\mathrm{fix}}{\mathrm{up}}\)ResNet-14 on subset of CIFAR-10 (\(20k\) examples). NaN denotes that a run crashed.

Figure 3: Learning affine data augmentations on subsets of data. (b) uses a \(\frac{\mathrm{fix}}{\mathrm{up}}\)ResNet-8 architecture whereas (c) a ResNet-50 architecture. (b,c) Top: normal dataset, bottom: rotated dataset.

\begin{table}
\begin{tabular}{l c c} \hline \hline Method & Chunk Proportions & Test accuracy \\ \hline Validation set optim. & \([0.9,0.1]\) & \(59.6\%_{\pm 0.6}\) \\ Partitioned & \([0.1,0.8,0.1]\) & \(\mathbf{87.3\%_{\pm 0.8}}\) \\ \hline \hline \end{tabular}
\end{table}
Table 3: Learning a feature extractor (first \(2\) out of \(3\) stages of a Wide ResNet-20) as a hyperparameter on CIFAR10.
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 a single run is therefore especially beneficial (Wang et al., 2021), and partitioned networks are a good fit for FL.

We extend our 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, a client belonging to chunk \(k\) sequentially optimizes partitions \(\mathbf{w}_{k:C}\) through sub-networks \(\mathbf{w}_{s}^{(k:C)}\) and computes a gradient wrt. the hyperparameters \(\psi\). Note that partitions \(\mathbf{w}_{1:k}\) 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 et 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 CIFAR10, 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-generating process and underlying source of non-i.i.d.-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 on the rotated datasets, partitioned networks learn to correctly increase the rotation angle.

## 6 Discussion

We propose partitioned networks as a new method for hyperparameter optimization inspired by the marginal likelihood objective. It provides a general and scalable solution to finding hyperparameters in a 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 our 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.

\begin{table}
\begin{tabular}{l c c c c c c} \hline \hline Dataset \& size & \(\uparrow\)MNIST & & \(\uparrow\)RotMNIST & \(\downarrow\)Upload \\ Method & \(1.25k\) & \(5k\) & \(50k\) & \(1.25k\) & \(5k\) & \(50k\) & \([\%]\) \\ \hline FedAvg & \(95.4\%_{0.1}\) & \(97.4\%_{0.1}\) & \(99.0\%_{0.1}\) & \(80.5\%_{0.0}\) & \(90.4\%_{0.5}\) & \(96.8\%_{\pm 0.1}\) & \(100\) \\ FedAvg + Augerino & \(94.2\%_{0.5}\) & \(96.4\%_{0.1}\) & \(99.1\%_{0.0}\) & \(79.5\%_{0.3}\) & \(89.0\%_{\pm 2.0}\) & \(95.3\%_{\pm 0.2}\) & \(100\) \\ FedAvg + Partitioned & \(\mathbf{97.0\%_{0.1}}\) & \(\mathbf{98.3\%_{0.0}}\) & \(99.2\%_{0.1}\) & \(\mathbf{85.7\%_{0.9}}\) & \(\mathbf{93.5\%_{0.6}}\) & \(\mathbf{97.8\%_{0.1}}\) & \(77\) \\ \hline \multicolumn{7}{c}{\(\uparrow\)CIFAR10} & \multicolumn{7}{c}{\(\uparrow\)RotCIFAR10} & \(\downarrow\)Upload \\  & \(1.25k\) & \(5k\) & \(45k\) & \(1.25k\) & \(5k\) & \(45k\) & \([\%]\) \\ \hline FedAvg & \(50.2\%_{0.4}\) & \(64.5\%_{0.3}\) & \(79.2\%_{0.7}\) & \(35.6\%_{0.3}\) & \(45.2\%_{0.1}\) & \(53.9\%_{\pm 1.1}\) & \(100\) \\ FedAvg + Augerino & \(49.9\%_{0.8}\) & \(65.0\%_{0.2}\) & \(79.9\%_{0.4}\) & \(36.1\%_{0.2}\) & \(45.0\%_{0.2}\) & \(56.4\%_{0.7}\) & \(100\) \\ FedAvg + Partitioned & \(50.8\%_{1.0}\) & \(64.8\%_{0.4}\) & \(\mathbf{81.5\%_{0.5}}\) & \(\mathbf{37.1\%_{0.2}}\) & \(45.3\%_{0.3}\) & \(\mathbf{60.6\%_{0.2}}\) & \(91\) \\ \hline \hline \end{tabular}
\end{table}
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.