Cross-regularization: Adaptive Model Complexity through Validation Gradients

Conference: ICML2025
arXiv: 2506.19755
Code: To be confirmed
Area: Regularization
Keywords: Regularization, validation gradients, hyperparameter optimization, noise injection, uncertainty calibration, data augmentation

TL;DR

Proposes Cross-regularization, which directly optimizes regularization parameters (weight norm, noise scale, augmentation intensity) via validation set gradients, converging to the cross-validation optimal solution in a single training run, thereby eliminating the need for manual hyperparameter tuning.

Background & Motivation

• Limitations of Prior Work: Traditional regularization methods (such as weight decay and dropout) require manual tuning of the hyperparameter \(\lambda\), necessitating multiple training runs via cross-validation to identify the optimal value, which is computationally expensive and inflexible.
• Limitations of Existing Methods:
  • Methods like variational dropout optimize surrogate objectives on the training data, rather than directly targeting generalization performance.
  • Validation gradient-based methods (Maclaurin et al. 2015) require computing the inverse Hessian or maintaining the complete parameter trajectory, which is hard to scale to large networks.
  • Although Luketina et al. (2016) compute the validation gradient using only the most recent parameter update step, they train the regularization parameters on both data splits, which fails to guarantee convergence.
  • Population Based Training (PBT) requires training multiple models in parallel, resulting in high computational costs.
• Key Insight: Rather than controlling model complexity via an indirect hyperparameter \(\lambda\), it is better to directly optimize the regularization parameter \(\rho\) (e.g., weight norm, noise scale), leveraging the validation set gradient to provide continuous feedback on generalization.

Method

Overall Architecture: Double Separation of Parameters and Data

Model parameters are divided into two groups and optimized on different datasets:

• Feature learning parameters \(\theta\): Optimized via gradient descent on the training set.
• Regularization parameters \(\rho\): Optimized via gradient descent on the validation set (referred to as the "regularization set").

Alternating update rules:

\[\theta_{t+1} = \theta_t - \eta_\theta \nabla_\theta \mathcal{L}_{\text{train}}(\theta_t, \rho_t)\]

\[\rho_{t+1} = \rho_t - \eta_\rho \nabla_\rho \mathcal{L}_{\text{val}}(\theta_{t+1}, \rho_t)\]

L2 Regularization Instantiation

Reparameterizes weights into magnitude and direction: \(\rho = \|w\|_2\), \(\theta = w / \|w\|_2\). Thus, the training set optimizes the direction \(\theta\), while the validation set optimizes the magnitude \(\rho\), which is equivalent to seeking the optimal ridge regression solution.

Theorem 3.1 (L2 Cross-Validation Equivalence): Under smoothness and strong convexity conditions, cross-regularization converges to the same solution as the optimal ridge regression, i.e., \(\rho^* \theta^* = w_{\text{val}}(\lambda^*)\).

Extending Gradient Decomposition to Arbitrary Regularizers

For non-differentiable regularizers such as L1, decomposition is achieved by projecting the gradient onto the regularization direction and its orthogonal complement direction:

\[g = g_\rho + g_\perp, \quad g_\rho = \text{Proj}_{\nabla R}(g)\]

• Training updates proceed along the \(g_\perp\) direction (maintaining the current level of complexity).
• Validation updates proceed along the \(g_\rho\) direction (adjusting the regularization strength).

Stochastic Regularization (Noise Injection)

Learning-scale noise is injected into each layer of the neural network:

\[h_l = g(\hat{u}_l + \sigma_l \epsilon), \quad \epsilon \sim \mathcal{N}(0, I)\]

Key design: A single noise sample is used during training, whereas a Monte Carlo average (\(K\) samples) is used during validation:

\[f_{\text{val}}(x) = \frac{1}{K} \sum_{k=1}^K f(x, \epsilon_k)\]

Validation uses the averaged prediction rather than the deterministic prediction (\(\epsilon = 0\)); otherwise, the validation loss would be independent of the noise scale \(\sigma\), rendering optimization impossible.

Key Experimental Results

Classical Regularization Validation

Method	Task	Results
Cross-reg L2	Synthetic ridge data	Converges to the optimal ridge regression solution
Cross-reg L1	Diabetes prediction	Automatically discovers sparsity matching LASSO's optimal cross-validation
Cross-reg spline	Function fitting	Automatically learns the appropriate smoothness

Neural Network Noise Regularization (CIFAR-10)

Method	Accuracy	Remarks
Baseline (no regularization)	76.0%	—
Fixed noise (\(\sigma=1\), last 5 layers)	—	Hinders learning initially, still overfits later
PBT (multi-model evolutionary search)	83.7%	Requires training multiple models in parallel
Cross-reg	83.7%	Matches PBT performance in a single training run

Uncertainty Calibration (ECE)

Method	ECE ↓	Accuracy
Uncalibrated model	0.163	67.4%
Temperature Scaling	0.057	69.6%
Fixed Reg	0.175	74.7%
Deep Ensemble (5 models)	0.030	81.3%
Cross-reg	0.038	79.5%

Data Augmentation (SVHN)

• Test accuracy increases from 82.8% to 86.3%
• Generalization gap drops from 16.2% to 7.3%
• Automatically learns: translation of 1-2 pixels, rotation of 3°, and shear of approximately 0

Highlights & Insights

1. Astonishingly High Noise Tolerance: In certain layers of VGG-16, a noise scale of \(\sigma \approx 13\) (equivalent to a dropout rate of 99.87%) is learned, which far exceeds conventional understanding, yet the model continues to function normally. This aligns with findings from the Lottery Ticket Hypothesis, where VGG on CIFAR-10 can be pruned to 98%+ sparsity.
2. Architecture-Aware Regularization Patterns: The noise in ResNet is concentrated in the early and final layers that cannot be bypassed by skip connections (\(\sigma_2 = 10.4\)), revealing the information flow structure of residual networks.
3. Unified Framework: The same method can handle L2/L1 norm regularization, noise injection, data augmentation, and uncertainty calibration, needing only modifications to the definition of \(\rho\).
4. Computational Efficiency: Compared to the \(O(PT)\) forward plays of PBT, it requires only \(O(T(1+K/r))\), invoking approximately 10% additional overhead; moreover, the validation set can be as small as 1% of the training set.
5. Theoretical Guarantees: Proves linear convergence (Theorem 4.1), equivalence to cross-validation (Theorem 4.5), and that the statistical error depends solely on the dimension of the regularization parameters \(k\) (Theorem 4.4).

Limitations & Future Work

1. Limited Experimental Scale: Validated primarily on CIFAR-10/SVHN using VGG/WideResNet, lacking large-scale evaluations on datasets like ImageNet and model architectures like Transformers.
2. Dependency on a Validation Set: Requires setting aside an extra validation set for regularization optimization, which may affect the available size of training data in small-data scenarios (though the paper notes that 1% is sufficient).
3. Strong Constraints for Theoretical Convergence: Theorem 4.3 assumes that the validation loss is strongly convex with respect to \(\rho\) and that the gradient is Lipschitz continuous, assumptions that may not hold strictly in deep networks.
4. Restricted Noise Types: Standard Dropout (which is non-differentiable) is not supported; only continuous, differentiable stochastic regularizations like Gaussian noise are supported.
5. Interaction with Modern Training Techniques: Compatibility with modern training configurations like AdamW, cosine schedules, and mixed-precision training is not thoroughly explored.
6. Single Author/Lab: Work originated from NightCity Labs, potentially lacking comprehensive validation supported by large-scale computing resources.

Related Work & Insights

• Luketina et al. (2016): The most closely related prior work, which also optimizes hyperparameters using validation gradients. However, it trains regularization parameters on both data splits, leaving no convergence guarantees.
• Concrete Dropout (Gal et al. 2017): A variational inference method for learning dropout rates, but it optimizes a surrogate objective rather than directly targeting generalization.
• PBT (Jaderberg et al. 2017): Uses evolutionary search to find hyperparameters that yield the best validation performance, but requires training multiple models in parallel.
• Lottery Ticket Hypothesis: The high noise tolerance patterns discovered by Cross-reg align closely with pruning sparsity, providing mutually supporting evidence from a different perspective.

Rating

• Novelty: ⭐⭐⭐⭐ (The framework design of directly applying validation gradients to regularization parameters is simple and elegant)
• Experimental Thoroughness: ⭐⭐⭐ (Covers various types of regularization, but dataset and architecture scales are relatively small)
• Writing Quality: ⭐⭐⭐⭐ (The connection between theory and experiments is clear, and the noise dynamics analysis is profound)
• Value: ⭐⭐⭐⭐ (A unified framework addressing the pain points of regularization parameter tuning, offering strong practicality)

Related Papers

• [CVPR 2025] Uncertainty Weighted Gradients for Model Calibration
• [ICML 2025] Time-Aware World Model for Adaptive Prediction and Control
• [ACL 2025] ALGEN: Few-Shot Inversion Attacks on Textual Embeddings via Cross-Model Alignment
• [ICML 2025] Prediction-Powered Adaptive Shrinkage Estimation
• [ICML 2025] Regression for the Mean: Auto-Evaluation and Inference with Few Labels through Post-hoc Regression