Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Implementation of MoMo algorithm #721

Merged
merged 27 commits into from
Apr 23, 2024
Merged

Implementation of MoMo algorithm #721

Changes from 1 commit
Commits
Show all changes
27 commits
Select commit Hold shift + click to select a range
d35b1d5
initial momo implementation
fabian-sp Jan 18, 2024
a1cd60f
improve docstring
fabian-sp Jan 18, 2024
98d6601
test pass locally
fabian-sp Jan 19, 2024
47404ab
make if condition jittable
fabian-sp Jan 19, 2024
a2e6ad6
Update optax/contrib/momo.py
fabian-sp Jan 20, 2024
1fade90
Update optax/contrib/momo.py
fabian-sp Jan 20, 2024
23f0d91
add no step flag and momo adam
fabian-sp Jan 22, 2024
8a81f5c
variable naming
fabian-sp Jan 22, 2024
d653da9
merge files
fabian-sp Mar 26, 2024
900e841
merge main
fabian-sp Mar 26, 2024
269a828
add to common test
fabian-sp Mar 26, 2024
041f22b
remove common test for now
fabian-sp Mar 26, 2024
008785f
added to common tests
fabian-sp Mar 27, 2024
5728ce0
rename loss to value
fabian-sp Mar 28, 2024
63c2444
Update optax/contrib/momo.py
fabian-sp Apr 2, 2024
ea092c5
Update optax/contrib/momo.py
fabian-sp Apr 2, 2024
12ae93a
Update optax/contrib/momo.py
fabian-sp Apr 2, 2024
0273aef
code highlighting
fabian-sp Apr 2, 2024
a2d9c19
fix conflict
fabian-sp Apr 2, 2024
b295b3f
formatting
fabian-sp Apr 2, 2024
3ee7686
implement lower bound estimators
fabian-sp Apr 2, 2024
97128c1
change state init to array
fabian-sp Apr 2, 2024
4ad3708
merge
fabian-sp Apr 2, 2024
3bd4dfc
Merge pull request #3 from fabian-sp/momo-lb-adapt
fabian-sp Apr 2, 2024
d2511af
rename lb
fabian-sp Apr 4, 2024
766a8f8
Merge branch 'main' into momo
fabian-sp Apr 4, 2024
ce05763
Merge branch 'main' into momo
fabian-sp Apr 10, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
variable naming
  • Loading branch information
@fabian-sp
fabian-sp committed Jan 22, 2024
commit 8a81f5c7ae4b1eeeab68e150d3e3d35f95cd388a
8 changes: 4 additions & 4 deletions optax/contrib/momo_adam.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -105,14 +105,14 @@ def update_fn(
updates,
)
bc2 = 1-beta2**(count+1)
Dk = tu.tree_map(
precond = tu.tree_map(
lambda eas: eps + jnp.sqrt(eas/bc2),
exp_avg_sq
)
)
exp_avg_weighted = tu.tree_map(
lambda ea, prec: ea/prec,
exp_avg,
Dk
precond
)
exp_avg_norm = tree_utils.tree_vdot(exp_avg,exp_avg_weighted)
gamma = beta1*state.gamma + (1-beta1)*tree_utils.tree_vdot(updates, params)
Expand All @@ -133,7 +133,7 @@ def update_fn(
-(alpha*weight_decay)/(1+alpha*weight_decay)*p
- tau*ea/prec,
exp_avg,
Dk,
precond,
params
)
new_state = MomoAdamState(
Expand Down
Loading