Skip to content

Toy implementations of some popular ML optimizers using Python/JAX

Notifications You must be signed in to change notification settings

shreyansh26/ML-Optimizers-JAX

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

9 Commits
 
 
 
 
 
 
 
 

Repository files navigation

ML Optimizers from scratch using JAX

Implementations of some popular optimizers from scratch for a simple model i.e., Linear Regression on a dataset of 5 features. The goal of this project was to understand how these optimizers work under the hood and try to do a toy implementation myself. I also use a bit of JAX magic to perform the differentiation of the loss function w.r.t to the weights and the bias without explicitly writing their derivatives as a separate function. This can help to generalize this notebook for other types of loss functions as well.

Kaggle Open In Colab

The optimizers I have implemented are -

  • Batch Gradient Descent
  • Batch Gradient Descent + Momentum
  • Nesterov Accelerated Momentum
  • Adagrad
  • RMSprop
  • Adam
  • Adamax
  • Nadam
  • Adabelief

References -