Square Root Regularization and High Loss

Question

I am testing out square root regularization (explained ahead) in a pytorch implementation of a neural network. Square root regularization, henceforth l1/2, is just like l2 regularization, but instead of squaring the weights, I take the square root of their absolute value. To implement it I penalize the loss as such in pytorch:

for p in model.parameters():
  loss += lambda * torch.sqrt(p.abs()).sum()

p.abs() is the absolute value of p, i.e the weights, torch.sqrt() is the square root and .sum() is the sum of the result for the individual weights. lambda is the penalization factor.
With no regularization, the loss settles around 0.4. With lambda=100, the loss settles somewhere around 0.45, if I use l2 or l1 regularization. Interestingly enough with lambda=0.001, the final value for the loss is around 0.44. Now if I use l1/2 with the same lambda it settles around 5000! This just does not make sense to me. If the regularization is such an overwhelming factor in the loss, then SGD has got to bring the (absolute value) of the weights down until, the penalty from the regularization is balanced with the actual classification from the Cross Entropy Loss I'm using - I know this is not the case, because the train and validation accuracies, are around the same of the original network (without regularization) at the end of training. I also know that the regularization is indeed happening in all three cases, as the loss in the early epochs differ significantly).
Is this to be expected or is this some error in my code or pytorch's SGD?

One more note If use l1/2 with a small lambda like 0.001, the loss comes down to around 0.5 and then becomes nan around epoch 70. For lambda=0.01, it becomes ~1.0 and then nan around the same epoch. For lambda=0.1, loss becomes 5 but no nan anymore for this lambda or any value higher (in 120 epochs total). For lambda=1.0, loss settles at ~50- as expected: apparently the weights settle down at a point where the sum of their square roots equals ~ 50, regardless of lambda...

Answer 1

The problem with using $L_p$ norm with $p < 1$ for regularization is the gradient.

Regularization term is used to force the parameters to be closer to zero. For this to work, when a parameter goes closer to zero, the gradient of regularization term, i.e. its contribution in updating the parameter, should decrease as well or at least remain constant.

However, by going closer to zero, the gradient increases unboundedly for $p < 1$, which causes radical changes in parameters, leading to your observations. Here is a visual plot for root $r$, absolute $a$, and squared $s$ regularization terms in one dimension (drawn here):

As you can see, for root regularization, the gradient explodes to infinity as we get closer to zero which is in conflict with our intention to bring the parameters close to zero.

For the sake of completeness, here is what a 1D parameter update with root regularization looks like: $$\theta_{n+1} = \theta_n + \alpha \frac{d L(\theta)}{d \theta}\Bigr|_{\theta=\theta_n} = \theta_n + \alpha \left(\frac{d l(\theta)}{d \theta}\Bigr|_{\theta=\theta_n} \pm \lambda \frac{1}{\color{red}{\sqrt{|\theta_n|}}}\right)$$ where $L(\theta) = l(\theta) + \lambda \sqrt{|\theta|}$.

You can spot the trouble! The more we get close to zero, the less we get close to zero!