Mini Batch with OptaxSolver about jaxopt HOT 2 CLOSED

Hello jecampagne

Your code did not run properly when I tested it; some variables (such as theta) were undefined.
I fixed it here: https://colab.research.google.com/drive/1Z-TawjqdpSQONMdeIqw69C1rXuq4UD0m?usp=sharing

When using solver.update manually it's working fine and the error decreases quickly below 1e-8 even in float32 precision. When using run() you should tune the tol parameter (its default value was calibrated for neural networks and is quite high), otherwise the run() will finish prematurly:

solver = OptaxSolver(opt=optax.adam(1e-2), fun=ridgeless_reg_objective,maxiter=1000,tol=1e-8)

I would like to add that Adam was primarly designed in the context of neural networks so its default parameters work best for normalized data (which is not your case with radius r > 1); in your case you may want to increase the learning rate. You can also increase maxiter as long as you don't forget to decrease tol.

Let me kow if you have any question !

from jaxopt.

Closing for now. Feel free to re-open if there's a bug.

from jaxopt.