PhD student @Oxford and very slow ultrarunner :)
Oxford/Malta
Joined April 2023
- Tweets57
- Following90
- Followers168
- Likes87
Cool stuff! Went ahead and implemented it in JAX as well using Optimistix:
github.com/packquickly/sched…
Screw it, here's a new JAX implementation in Optimsitix:
github.com/packquickly/sched…
Schedule-Free paper is up!
arxiv.org/abs/2405.15682
Joint work with collaborators @alicey_ang @HarshMeh1a @konstmish @akhaledv2 @AshokCutkosky
We have some strong small-scale experiments on Transformers, comparing to chinchilla-style cosine 10x reduction schedules.
Neat! Heuristics based on this crop up in stochastic optimisation. Often in the argument of stepsize * sum/sqrt(squared_sum) being "roughly" bounded by stepsize. Of course the independent and symmetric assumptions are both violated but that never stopped anyone...
Ultimately I think the question is "who do we want getting into wser?" Optimising the lottery is not the hard part, it's agreeing on who we should be optimising for that is.
Finally, a million thanks to @PatrickKidger, who supervised this whole project.
If you’re following me, chances are good you already follow him. If not, go give him a follow! (right after installing Lineax of course 😉)
4/4
The paper describes out how we achieved many of these things (such as differentiation through all our solvers,) and outlines some of the design choices we made when creating Lineax.
arxiv.org/abs/2311.17283
⭐ Lineax is now on arXiv! ⭐
If you’re doing linear solves or linear least-squares in JAX, give it a shot today!
Lineax
is fast ⚡️,
has new solvers (eg. QR, tridiagonal),
supports general linear Operators.
github: github.com/google/lineax
arXiv: arxiv.org/abs/2311.17283
1/n
They mentioned the choice between these two model functions made little difference in practice.
While this is believable, and indeed proved to be true, it struck me as an example of a claim which is very difficult to verify using existing optimisation software.
12/12
This example was taken from an offhand comment in "Training Deep and Recurrent Networks with Hessian-Free Optimisation" by Martens and Sutskever.
The first model function is minimised to produce the next iterate, this is the one with Tikhnov regularisation.
The second is used to predict how much the model was supposed to decrease, in order to adjust the step-size (trust-region radius.) This one is not regularised.
Tikhnov regularised trust-region methods (*cough* Levenberg-Marquardt) oddly use two different approximations to the objective function at each step.
One regularised, one not.
What if we just regularised both?
1/
Hurrah! 🎉🥳
Very excited this is for this release, after many months of work!
Lots of really cool things in Optimistix that I've wanted in optimisation software for a long time. I hope other researchers find as exciting as I do 😃
Announcing a new #JAX and #Equinox nonlinear optimisation library:
⭐️ Optimistix ⭐️
(GitHub: github.com/patrick-kidger/op…)
- Minimisation
- Nonlinear least-squares
- Root-finding
- Fixed-points
With blazing fast compile times and fast run times 🔥
1/7
Gives me Levenberg-Marquardt vibes: diagonal Hessian approx (quasi-Newton matrix) and clipping instead of LM parameter/trust region radius to handle the same numerical problems LM was introduced to handle.
Adam, a 9-yr old optimizer, is the go-to for training LLMs (eg, GPT-3, OPT, LLAMA).
Introducing Sophia, a new optimizer that is 2x faster than Adam on LLMs. Just a few more lines of code could cut your costs from $2M to $1M (if scaling laws hold).
arxiv.org/abs/2305.14342 🧵⬇️
This is quite cool! From a skim, it looks like the most significant insights here are: a) you only need to re-estimate the Hessian once every k steps and b) we really need to handle the saddle points and rapid curvature changes better. Does that sound right? 🧵