Skip to content

[MRG] Faster gromov-wasserstein linesearch for symmetric matrices #607

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.

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

Merged
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
3 changes: 3 additions & 0 deletions RELEASES.md
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,9 @@

## 0.9.3

#### New features
+ `ot.gromov._gw.solve_gromov_linesearch` now has an argument to specifify if the matrices are symmetric in which case the computation can be done faster.

#### Closed issues
- Fixed an issue with cost correction for mismatched labels in `ot.da.BaseTransport` fit methods. This fix addresses the original issue introduced PR #587 (PR #593)
- Fix gpu compatibility of sr(F)GW solvers when `G0 is not None`(PR #596)
Expand Down
15 changes: 11 additions & 4 deletions ot/gromov/_gw.py
Original file line number Diff line number Diff line change
Expand Up @@ -171,7 +171,7 @@ def line_search(cost, G, deltaG, Mi, cost_G, **kwargs):
return line_search_armijo(cost, G, deltaG, Mi, cost_G, nx=np_, **kwargs)
else:
def line_search(cost, G, deltaG, Mi, cost_G, **kwargs):
return solve_gromov_linesearch(G, deltaG, cost_G, hC1, hC2, M=0., reg=1., nx=np_, **kwargs)
return solve_gromov_linesearch(G, deltaG, cost_G, hC1, hC2, M=0., reg=1., nx=np_, symmetric=symmetric, **kwargs)

if not nx.is_floating_point(C10):
warnings.warn(
Expand Down Expand Up @@ -479,7 +479,7 @@ def line_search(cost, G, deltaG, Mi, cost_G, **kwargs):
return line_search_armijo(cost, G, deltaG, Mi, cost_G, nx=np_, **kwargs)
else:
def line_search(cost, G, deltaG, Mi, cost_G, **kwargs):
return solve_gromov_linesearch(G, deltaG, cost_G, hC1, hC2, M=(1 - alpha) * M, reg=alpha, nx=np_, **kwargs)
return solve_gromov_linesearch(G, deltaG, cost_G, hC1, hC2, M=(1 - alpha) * M, reg=alpha, nx=np_, symmetric=symmetric, **kwargs)
if not nx.is_floating_point(M0):
warnings.warn(
"Input feature matrix consists of integer. The transport plan will be "
Expand Down Expand Up @@ -647,7 +647,7 @@ def fused_gromov_wasserstein2(M, C1, C2, p=None, q=None, loss_fun='square_loss',


def solve_gromov_linesearch(G, deltaG, cost_G, C1, C2, M, reg,
alpha_min=None, alpha_max=None, nx=None, **kwargs):
alpha_min=None, alpha_max=None, nx=None, symmetric=False, **kwargs):
"""
Solve the linesearch in the FW iterations for any inner loss that decomposes as in Proposition 1 in :ref:`[12] <references-solve-linesearch>`.

Expand Down Expand Up @@ -676,6 +676,10 @@ def solve_gromov_linesearch(G, deltaG, cost_G, C1, C2, M, reg,
Maximum value for alpha
nx : backend, optional
If let to its default value None, a backend test will be conducted.
symmetric : bool, optional
Either structures are to be assumed symmetric or not. Default value is False.
Else if set to True (resp. False), C1 and C2 will be assumed symmetric (resp. asymmetric).

Returns
-------
alpha : float
Expand Down Expand Up @@ -708,7 +712,10 @@ def solve_gromov_linesearch(G, deltaG, cost_G, C1, C2, M, reg,

dot = nx.dot(nx.dot(C1, deltaG), C2.T)
a = - reg * nx.sum(dot * deltaG)
b = nx.sum(M * deltaG) - reg * (nx.sum(dot * G) + nx.sum(nx.dot(nx.dot(C1, G), C2.T) * deltaG))
if symmetric:
b = nx.sum(M * deltaG) - 2 * reg * nx.sum(dot * G)
else:
b = nx.sum(M * deltaG) - reg * (nx.sum(dot * G) + nx.sum(nx.dot(nx.dot(C1, G), C2.T) * deltaG))

alpha = solve_1d_linesearch_quad(a, b)
if alpha_min is not None or alpha_max is not None:
Expand Down