Research
I want to make neural network training fast, efficient and scalable.
To achieve this, I primarily focus on second-order optimization methods,
with a specific interest in the Gauss-Newton approach
as well as diagonal Hessian approximation.
|
Software
I find JAX
to be a great tool for experimenting with new optimization algorithms.
During my PhD I've been inspired by specialty JAX repositories like
Patrick Kidger's
diffrax
and equinox,
as well as Robert Lange's
gymnax.
So, as part of my thesis I've been working on
somax -
a second-order complement to optax
(which is a library of mostly first-order solvers).
|
|
Somax: Stochastic Second-Order Optimization in JAX
Mikalai Korbit
code
/
arXiv
Somax is a library of stochastic second-order methods for machine learning
optimization written in JAX. Somax is based on the JAXopt StochasticSolver API,
and can be used as a drop-in replacement for JAXopt as well as Optax solvers.
|
How to Design Machine Learning Experiments [In Preparation]
Course Co-creator and Instructor
|
|