Improving Deep Learning Optimization through Constrained Parameter Regularization

TL;DR

This work addresses the rigid nature of standard weight decay in deep learning by introducing Constrained Parameter Regularization (CPR), which enforces upper bounds on per-parameter regularization via an augmented Lagrangian framework. CPR yields adaptive, per-parameter training pressure by updating multipliers $\lambda^j_t$ for each parameter group, with four initialization strategies for the bounds $\kappa^j$, including a hyperparameter-free approach based on the first inflection point. Empirically, CPR improves performance over traditional weight decay across CIFAR-100, ImageNet, CLIP finetuning, OpenWebText language modeling with GPT-2 variants, and several medical segmentation tasks, while incurring only modest runtime overhead. The results also show CPR can reduce required training budgets and hyperparameter tuning, offering a practical, scalable alternative for large-scale pretraining and fine-tuning of deep models. Overall, CPR demonstrates that constraint-based, per-parameter regularization can enhance generalization and stability across diverse domains, suggesting wide applicability and avenues for future work with even larger models and adaptive bound mechanisms.

Abstract

Regularization is a critical component in deep learning. The most commonly used approach, weight decay, applies a constant penalty coefficient uniformly across all parameters. This may be overly restrictive for some parameters, while insufficient for others. To address this, we present Constrained Parameter Regularization (CPR) as an alternative to traditional weight decay. Unlike the uniform application of a single penalty, CPR enforces an upper bound on a statistical measure, such as the L2-norm, of individual parameter matrices. Consequently, learning becomes a constraint optimization problem, which we tackle using an adaptation of the augmented Lagrangian method. CPR introduces only a minor runtime overhead and only requires setting an upper bound. We propose simple yet efficient mechanisms for initializing this bound, making CPR rely on no hyperparameter or one, akin to weight decay. Our empirical studies on computer vision and language modeling tasks demonstrate CPR's effectiveness. The results show that CPR can outperform traditional weight decay and increase performance in pre-training and fine-tuning.

Figures (17)

• Figure 1: GPT2s training using Adam with weight decay or CPR (Kappa-IP). AdamCPR outperforms AdamW with the same budget and only requires 2/3 of the budget to reach the same score.
• Figure 2: Percentage of correct labels ($\uparrow$) of a ResNet18 trained on CIFAR100 with AdamW and AdamCPR with Kappa-IP or Kappa-WS. We use a learning rate warm-up of $500$ steps and the best Kappa-WS value is $2\times$ the warm-up steps. We report the mean of three runs with random seeds. We see that both CPR versions outperform weight decay
• Figure 3: Perplexity ($\downarrow$) $\pm$ std across three random seeds of GPT2s and GPT2m trained on OpenWebText with AdamW (left) and AdamCPR with Kappa-IP (middle) and AdamCPR with Kappa-WS (right). We use a learning rate warm-up of $5k$ steps. The CPR with the hyperparameter-free strategy Kappa-IP outperforms weight decay but also CPR with warm start.
• Figure 4: Percentage of performance change before and after fineuning Mistral 7B with pubmedQA artificial data ($\uparrow$) with the use of AdamW (left) and AdamCPR with Kappa-WS (right). We use a learning rate warm-up of $50$ steps. We see that CPR outperforms weight decay for each learning rate.
• Figure 5: The training dynamics of AdamW (blue) and AdamCPR with Kappa-IP (green) in a GPT2s training run. The upper plot shows the squared L2 norm of the first fully connected weight in the fifth layer. Below we see the gradient of the squared L2 norm regarding the training steps. After the inflection point (7400), Kappa-IP initializes kappa $\kappa^j \gets R({\bm{\theta}}^j_{t=i})$ and starts the regularization. The third plot shows CPR's lambda enforcing the constraint. At the bottom, we see the validation loss. AdamW converges faster in the beginning of the training but CPR leads to a more linear improvement and a better final performance.