useful diagram

Author: lucidrains

README.md

@@ -1,8 +1,10 @@
+<img src="./adam-atan2.png" width="400px"></img>
+
 ## Adam-atan2 - Pytorch
 
 Implementation of the proposed <a href="https://arxiv.org/abs/2407.05872">Adam-atan2</a> optimizer in Pytorch
 
-A multi-million dollar paper out of google deepmind basically proposes a small change to Adam update rule (using `atan2`) for greater stability
+A multi-million dollar paper out of google deepmind proposes a small change to Adam update rule (using `atan2`) to remove the epsilon altogether for numerical stability and scale invariance
 
 ## Install
 

adam-atan2.png

@@ -83,11 +83,11 @@ def step(
                 exp_avg_sq.lerp_(grad * grad, 1. - beta2)
 
                 # the following line is the proposed change to the update rule
-                # using atan2 instead of a division with epsilons - they also suggest hyperparameters `a` and `b` should be explored beyond its default of 1.
+                # using atan2 instead of a division with epsilon in denominator
 
-                update = a * atan2(exp_avg / bias_correct1, b * sqrt(exp_avg_sq / bias_correct2))
+                update = atan2(exp_avg / bias_correct1, b * sqrt(exp_avg_sq / bias_correct2))
 
-                p.add_(update, alpha = -lr)
+                p.add_(update, alpha = -lr * a)
 
                 # increment steps
 

adam_atan2_pytorch/adam_atan2.py

@@ -1,6 +1,6 @@
 [project]
 name = "adam-atan2-pytorch"
-version = "0.0.2"
+version = "0.0.3"
 description = "Adam-atan2 for Pytorch"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }