Feature/ulrgw #410

michalk8 · 2023-08-03T18:19:01Z

Add unbalanced low-rank (fused) GW.

TODOs:

verify the impl.
check initialization
fix cost value in the fused case
~~[ ] decide on how to handle/whether to deprecate LR option in GW (would say yes)~~ will be done in a future PR
update the docs (e.g., remove mentions of LRSinkhorn)
fix primal_cost
check balanced case with the new LRGromovWasserstein class
update notebooks
tests

codecov · 2023-08-03T20:23:34Z

Codecov Report

Merging #410 (f04757f) into main (21d3627) will decrease coverage by 0.93%.
The diff coverage is 83.33%.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #410      +/-   ##
==========================================
- Coverage   91.40%   90.48%   -0.93%     
==========================================
  Files          55       56       +1     
  Lines        5943     6251     +308     
  Branches      866      887      +21     
==========================================
+ Hits         5432     5656     +224     
- Misses        375      453      +78     
- Partials      136      142       +6

Files Changed	Coverage Δ
src/ott/initializers/linear/initializers_lr.py	`91.30% <40.00%> (-1.08%)`	⬇️
src/ott/solvers/quadratic/gromov_wasserstein_lr.py	`81.47% <81.47%> (ø)`
src/ott/initializers/quadratic/initializers.py	`88.88% <100.00%> (+2.00%)`	⬆️
src/ott/math/utils.py	`94.54% <100.00%> (ø)`
src/ott/solvers/linear/lr_utils.py	`100.00% <100.00%> (ø)`
src/ott/solvers/linear/sinkhorn_lr.py	`97.33% <100.00%> (-1.04%)`	⬇️
src/ott/solvers/quadratic/gromov_wasserstein.py	`84.32% <100.00%> (-3.18%)`	⬇️

... and 2 files with indirect coverage changes

📢 Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ott-jax).

marcocuturi

Thanks Michal! let's chat :)

src/ott/solvers/quadratic/gromov_wasserstein_lr.py

marcocuturi · 2023-08-07T00:39:21Z

src/ott/solvers/quadratic/gromov_wasserstein_lr.py

+
+
+@jax.tree_util.register_pytree_node_class
+class LRGromovWasserstein(sinkhorn.Sinkhorn):


isn't this supposed to inherite from WasSolver? It's not clear either why there should be a link to Sinkhorn (I would understand LRSinkhorn)

For now, I think it makes sense, since if we were to use WasSolver with LRSinkhornSolver, we'd have:

3 loops (outermost GW loop, outer loop in LR Sinkhorn, Dykstra loop)

LRSinkhornSolver would have to be modified to accept a quadratic problem (or have a specific inner solver for unbalanced LR GW, which is more complicated than this)

We definitely need a abstraction for the solvers, this is just temporary; will come up with a better solution to this in the future.

I also think that our current LR GW (balanced) is a bit wrong when it comes to the convergence criterion (we exit the outermost GW loop when the successive costs are close; in the paper it should exist when the errors between succcessive Q/R/g are close [the convergence criterion for the linearized objective]).
This will be fixed in a future PR once the LR GW (balanced and unbalanced) are unified in this class.

Conceptually, I am not sure exactly what we're inheriting from Sinkhorn? It seems Sinkhorn is used as a simpler type of WassersteinSolver?

src/ott/solvers/quadratic/gromov_wasserstein_lr.py

review-notebook-app · 2023-09-04T08:15:29Z

Check out this pull request on

See visual diffs & provide feedback on Jupyter Notebooks.

Powered by ReviewNB

marcocuturi · 2023-09-05T12:51:18Z

docs/tutorials/notebooks/GWLRSinkhorn.ipynb

@@ -1,7 +1,6 @@
 {


if you are using the primal cost, then it's better to rename the titles to "Primal cost of LR Solution" and " ... of Entropic Solution" to avoid the ambiguity. I am also surprised that the plot of low rank is looks so much like a rank 1 or 2 matrix, can you double check?

Reply via ReviewNB

marcocuturi

This LGTM, but I'd like to chat about the inheritance of LRGromovWasserstein because I am not sure where this is headed! thanks!

marcocuturi · 2023-09-05T12:56:56Z

src/ott/solvers/quadratic/gromov_wasserstein_lr.py

+
+
+@jax.tree_util.register_pytree_node_class
+class LRGromovWasserstein(sinkhorn.Sinkhorn):


Conceptually, I am not sure exactly what we're inheriting from Sinkhorn? It seems Sinkhorn is used as a simpler type of WassersteinSolver?

michalk8 · 2023-09-05T14:11:59Z

This LGTM, but I'd like to chat about the inheritance of LRGromovWasserstein because I am not sure where this is headed! thanks!

Agreed, we sorely need a better solver hierarchy, let's talk about it! Will just write some tests and merge.

…ature/ulrgw

* Remove low-rank from GromovWasserstein solver * First skeleton loop * Add LRGW implementation * Add ULFGW * Revert change * Add a TODO * Fix `grad_g` in the fused case * Update docs * Remove duplicate citation * Fix cost for the fused case * Fix bugs in TI * Remove unused import * Change way array extraction in LR init works * Disallow LR in the old GW solver * Disallow LR in old GW class * Remove `is_entropic` property * Use `jnp.linalg.norm` * Simplify initializers in GW * Simplify initializer creation for low-rank * Remove temporary name * Fix norms * Fix linkcheck * Remove old initializers test * Fix more initializer tests * Remove `LRQuadraticInitializer`, `reg_ot_cost -> reg_gw_cost` * `host_callback` -> `io_callback` * Fix more initializers tests * Fix more tests * Remove initializer mention from the docs * Remove mention of LR initializer * Start incorporating GWLoss * Simplify reg GW cost computation * Finish `primal_cost` * Don't calculate unbal. grads in balanced case * Fix `primal_cost` in balanced case * Update GW LR notebook * Convert quad problem to LR if possible * Convert quad problem to LR if possible * Regenerate GWLR Sinkhorn * Regenerate `LRSinkhorn` * [ci skip] Fix linter * Fix convergence metric * Undo TODO * Fix factor * Regenerate notebooks * Add tests

michalk8 added 6 commits August 2, 2023 12:21

Remove low-rank from GromovWasserstein solver

19756c1

First skeleton loop

a079761

Add LRGW implementation

0d0c85f

Add ULFGW

1f7f10d

Revert change

07aeedf

Add a TODO

6d257ff

michalk8 added the enhancement New feature or request label Aug 3, 2023

michalk8 requested review from marcocuturi and meyerscetbon August 3, 2023 18:19

michalk8 self-assigned this Aug 3, 2023

Fix grad_g in the fused case

bf0a50d

marcocuturi reviewed Aug 7, 2023

View reviewed changes

michalk8 added 17 commits August 17, 2023 17:57

Update docs

7b8a0f4

Remove duplicate citation

5855a32

Merge branch 'main' into feature/ulrgw

01bd8e6

Fix cost for the fused case

e4a70c3

Fix bugs in TI

4ecdca0

Remove unused import

04a3e8b

Change way array extraction in LR init works

4fb9ee4

Disallow LR in the old GW solver

43d39eb

Disallow LR in old GW class

0747c56

Remove is_entropic property

83a1c36

Use jnp.linalg.norm

b783f3a

Simplify initializers in GW

0f661cc

Simplify initializer creation for low-rank

d79ce31

Remove temporary name

1be699f

Fix norms

87284ae

Fix linkcheck

30f275e

Remove old initializers test

f11e19f

michalk8 added 12 commits September 1, 2023 12:44

Fix more initializers tests

3714057

Fix more tests

0fa6be2

Remove initializer mention from the docs

3843609

Remove mention of LR initializer

62aad99

Start incorporating GWLoss

4e57972

Simplify reg GW cost computation

b312126

Finish primal_cost

a499102

Don't calculate unbal. grads in balanced case

4bcb71b

Fix primal_cost in balanced case

eea362f

Update GW LR notebook

3f14488

Convert quad problem to LR if possible

8207696

Convert quad problem to LR if possible

ba272f0

michalk8 and others added 2 commits September 5, 2023 10:26

Regenerate GWLR Sinkhorn

c565ee3

Regenerate LRSinkhorn

f508d5d

marcocuturi reviewed Sep 5, 2023

View reviewed changes

marcocuturi approved these changes Sep 5, 2023

View reviewed changes

michalk8 added 8 commits September 5, 2023 22:14

Merge remote-tracking branch 'upstream/main' into feature/ulrgw

f7f440f

Merge branch 'feature/ulrgw' of ssh://github.com/michalk8/ott into fe…

f703a55

…ature/ulrgw

[ci skip] Fix linter

ec45825

Fix convergence metric

68d3ce2

Undo TODO

afd8ea6

Fix factor

e006eae

Regenerate notebooks

e06f1b0

Add tests

f04757f

michalk8 merged commit a18c16c into ott-jax:main Sep 8, 2023

michalk8 deleted the feature/ulrgw branch September 8, 2023 07:46

michalk8 mentioned this pull request Sep 8, 2023

Adapt new low-rank GW solvers from OTT theislab/moscot#603

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Feature/ulrgw #410

Feature/ulrgw #410

michalk8 commented Aug 3, 2023 •

edited

Loading

codecov bot commented Aug 3, 2023 •

edited

Loading

marcocuturi left a comment

marcocuturi Aug 7, 2023

michalk8 Aug 17, 2023

marcocuturi Sep 5, 2023

review-notebook-app bot commented Sep 4, 2023

marcocuturi Sep 5, 2023 •

edited

Loading

marcocuturi left a comment

marcocuturi Sep 5, 2023

michalk8 commented Sep 5, 2023



		@jax.tree_util.register_pytree_node_class
		class LRGromovWasserstein(sinkhorn.Sinkhorn):

Feature/ulrgw #410

Feature/ulrgw #410

Conversation

michalk8 commented Aug 3, 2023 • edited Loading

codecov bot commented Aug 3, 2023 • edited Loading

Codecov Report

marcocuturi left a comment

Choose a reason for hiding this comment

marcocuturi Aug 7, 2023

Choose a reason for hiding this comment

michalk8 Aug 17, 2023

Choose a reason for hiding this comment

marcocuturi Sep 5, 2023

Choose a reason for hiding this comment

review-notebook-app bot commented Sep 4, 2023

marcocuturi Sep 5, 2023 • edited Loading

Choose a reason for hiding this comment

marcocuturi left a comment

Choose a reason for hiding this comment

marcocuturi Sep 5, 2023

Choose a reason for hiding this comment

michalk8 commented Sep 5, 2023

michalk8 commented Aug 3, 2023 •

edited

Loading

codecov bot commented Aug 3, 2023 •

edited

Loading

marcocuturi Sep 5, 2023 •

edited

Loading