Meaning of ProxGradState.error

[lucagrementieri]

Hello,

I'm using jaxopt.GradientDescent and the optimization stops too early because the error is below the tolerance, but in the meantime the value of the loss function is very high.

At the beginning I was assuming that the state attribute error was the value of the loss function with the latest parameters, but that's not the case.
Therefore, how is ProxGradState.error computed?

[Algue-Rythme]

The error field has different semantics for each solver in general, but as a general rule of thumb, since Jaxopt is thought to be used in the context of Implicit Differentiation, the error field is usually related to the value of optimality_fun, i.e in your case, the value of norm(step_size * grad(fun)(params)):

jaxopt/jaxopt/_src/proximal_gradient.py

Line 176 in 52cbdd4

def _error(self, x, x_fun_grad, hyperparams_prox):

So you may be far away from your optimum, but in this particular point the loss landscape is flat and this triggers early termination of the algorithm. For strongly convex problems (typically those for which Gradient Descent has convergence guarantees) this means you are close from the optimum. For other problems, well, it is not a great criterion.

Hence, I suggest you set tol to a smaller (eventually zero) value, and that you set maxiter to whatever you want. You can also drop the use of run method and use init_state and update manually, with your own custom criterion.

As a concluding remark, I would like to say that for constrained problems a huge (non zero) value of the loss is expected even at the optimum, so you may also want to do this sanity check.