How to decouple weight decay strength and model size?

Question

Consider the neural networks' loss function with the cross entropy term and the $L^2$ weight decay term, which are usually written as:

$$E = \frac{1}{N_{samples}} \sum_{i=1}^{N_{samples}} \text{cross_entropy}\left(x_i, y_i\right) + \lambda \sum_{j=1}^{N_{parameters}}\left(w_j\right)^2$$

The weight decay term can be written as either "sum square" or "mean square". They are equivalent by a scaling of $\lambda$ when the number of parameters is fixed, as discussed here and here.

However, the problem appears when the number of parameters increases, and we have to re-tune the weight decay strength $\lambda$. Let's consider the two options:

1. The "sum square" of parameters can become huge; thus, it can completely dominate the cross entropy loss, which is relatively unchanged in magnitude regardless of model size. This means the model is overly regularized and we need to decrease $\lambda$ to reduce bias. The good side of this option is that, the derivative of the weight decay term is $\lambda w_j$, meaning we reduce each parameter by a fixed amount $\lambda$ at each gradient update regardless of the model size. Thus, this option seems "bad" when considering the relative loss values, but seems "correct" when considering the gradient; how to unite this discrepancy?

2. For "mean square" weight decay, the weight decay term is relatively unchanged in magnitude regardless of model size; thus, the relative magnitude between the cross entropy loss and the weight decay loss is unchanged. So, $\lambda$ can stay the same (or set to slightly larger value to account for overfitting risk in larger model). However, the bad side of this option is that, the derivative is $\frac{\lambda}{N_{parameters}} w_j$, which become very small when the model size increases. Thus, this option seems "good" when considering the relative loss values, but seems "bad" (incorrect?) when considering the gradient; how to unite this discrepancy?

I cannot make my mind which option is better. Is it reasonable to use "mean square" weight decay to have a stable $\lambda$ regardless of the model size, or did I miss something?

Answer 1

The loss function should be either: $$ E = \sum_i^{N_{samples}} \text{cross_entropy}(x_i,y_i) + \lambda \sum_j^{N_{params}} w_j^2 $$ Or the averaged version: $$ \bar E = \frac{1}{N_{samples}} E \\ = \frac{1}{N_{samples}} \sum_i^{N_{samples}} \text{cross_entropy}(x_i,y_i) + \frac{1}{N_{samples}} \lambda \sum_j^{N_{params}} w_j^2 $$

If you only average over the cross entropy part, then your model won't be able to learn more from data when sample size is increased.

If you only average over the L2 part, then your model won't increase the penalty when the number of parameters is increased.