# Hyperparameter Optimization through Neural Network Partitioning

arXiv: https://arxiv.org/abs/2304.14766 · HTML: https://mlsanity.com/paper/2304.14766/

## Key facts

- Cited by 8 papers in the map (data through 14 Oct 2026).
- Builds on 36 of its 50 references that are in the map.
- First posted on arXiv 28 Apr 2023.

## Builds on (most cited first, 30 of 36)

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Communication-Efficient Learning of Deep Networks from Decentralized Data](https://mlsanity.com/paper/1602.05629/) (2016)
- [Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift](https://mlsanity.com/paper/1502.03167/) (2015)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet Classification](https://mlsanity.com/paper/1502.01852/) (2015)
- [Gaussian Error Linear Units (GELUs)](https://mlsanity.com/paper/1606.08415/) (2016)
- [Wide Residual Networks](https://mlsanity.com/paper/1605.07146/) (2016)
- [The Lottery Ticket Hypothesis: Finding Sparse, Trainable Neural Networks](https://mlsanity.com/paper/1803.03635/) (2018)
- [Identity Mappings in Deep Residual Networks](https://mlsanity.com/paper/1603.05027/) (2016)
- [DARTS: Differentiable Architecture Search](https://mlsanity.com/paper/1806.09055/) (2018)
- [Group Normalization](https://mlsanity.com/paper/1803.08494/) (2018)
- [The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables](https://mlsanity.com/paper/1611.00712/) (2016)
- [Adaptive Federated Optimization](https://mlsanity.com/paper/2003.00295/) (2020)
- [AutoAugment: Learning Augmentation Policies from Data](https://mlsanity.com/paper/1805.09501/) (2018)
- [Deep Kernel Learning](https://mlsanity.com/paper/1511.02222/) (2015)
- [Rotation Equivariant CNNs for Digital Pathology](https://mlsanity.com/paper/1806.03962/) (2018)
- [A Field Guide to Federated Optimization](https://mlsanity.com/paper/2107.06917/) (2021)
- [Optimizing Millions of Hyperparameters by Implicit Differentiation](https://mlsanity.com/paper/1911.02590/) (2019)
- [Monte Carlo Gradient Estimation in Machine Learning](https://mlsanity.com/paper/1906.10652/) (2019)
- [Concrete Dropout](https://mlsanity.com/paper/1705.07832/) (2017)
- [Fixup Initialization: Residual Learning Without Normalization](https://mlsanity.com/paper/1901.09321/) (2019)
- [Scalable Marginal Likelihood Estimation for Model Selection in Deep Learning](https://mlsanity.com/paper/2104.04975/) (2021)
- [PAC-Bayesian Theory Meets Bayesian Inference](https://mlsanity.com/paper/1605.08636/) (2016)
- [Federated Hyperparameter Tuning: Challenges, Baselines, and Connections to Weight-Sharing](https://mlsanity.com/paper/2106.04502/) (2021)
- [DADA: Differentiable Automatic Data Augmentation](https://mlsanity.com/paper/2003.03780/) (2020)
- [Bayesian Model Selection, the Marginal Likelihood, and Generalization](https://mlsanity.com/paper/2202.11678/) (2022)
- [Learning Invariances using the Marginal Likelihood](https://mlsanity.com/paper/1808.05563/) (2018)
- [Generalization Error of Invariant Classifiers](https://mlsanity.com/paper/1610.04574/) (2016)
- [Invariance Learning in Deep Neural Networks with Differentiable Laplace Approximations](https://mlsanity.com/paper/2202.10638/) (2022)
- [On the marginal likelihood and cross-validation](https://mlsanity.com/paper/1905.08737/) (2019)

## Built on by (most cited first, 8 of 8)

- [Learning Layer-wise Equivariances Automatically using Gradients](https://mlsanity.com/paper/2310.06131/) (2023)
- [Stochastic Marginal Likelihood Gradients using Neural Tangent Kernels](https://mlsanity.com/paper/2306.03968/) (2023)
- [A Generative Model of Symmetry Transformations](https://mlsanity.com/paper/2403.01946/) (2024)
- [Distributional Training Data Attribution: What do Influence Functions Sample?](https://mlsanity.com/paper/2506.12965/) (2025)
- [FedPop: Federated Population-based Hyperparameter Tuning](https://mlsanity.com/paper/2308.08634/) (2023)
- [Partitioned Neural Network Training via Synthetic Intermediate Labels](https://mlsanity.com/paper/2403.11204/) (2024)
- [Twin: Tuning Learning Rate and Weight Decay of Deep Homogeneous Classifiers without Validation](https://mlsanity.com/paper/2403.05532/) (2024)
- [Fusion Sampling Validation in Data Partitioning for Machine Learning](https://mlsanity.com/paper/2508.01325/) (2025)

## Abstract (arXiv)

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.