LR-UGW on which marginal by default?

[EricZimmermann]

Hello,

I am implementing the semi-relaxed gw objective proposed in Dassot: https://openreview.net/pdf?id=TsMdJULS9o

Where:

• $' D_n: \textrm{Disconnected graph as cost matrix} $
• $' S_X: \textrm{Sample Gram matrix as cost matrix} $
• $' h = \frac{n}{N} 1_N: \textrm{Scaled marginal} $

I wanted to clarify two components:

• Using the full rank solver, what does tau_a, tau_b correspond to? Should I be setting 1 and inf, strictly to emulate the desired behaviour? I have found that I have to set the tau for Dn to 1 and tau for Sx to 0 to get something that at times works with a small range of epsilon, but this seems contrary to my understanding of the ULROT Solvers paper. gw_unbalanced_correction is also set to True.
• What is the correct way to implement this with the LRGW solver? Based on the tutorial, should be setting anything else explicitly?

Here is an example of the code I am using for LR=UGW:

def dassot_lr(X, n, eps):
    N, _ = X.shape
    Sx = jnp.array(X @ X.T)
    Dn = -jnp.ones((n, n))
    Dn = Dn.at[jnp.diag_indices(n)].set(1.0)
    p = jnp.ones(n) / n
    q = jnp.ones(N) / N
    Sx_geo = Geometry(Sx)
    Dn_geo = Geometry(Dn)
    problem = QuadraticProblem(
        Dn_geo, 
        Sx_geo, 
        a=p,
        b=q,
    )
    solver = jax.jit(LRGromovWasserstein(rank=int(N**0.5), epsilon=eps))
    coupling = solver(problem).matrix
    return coupling

Emulating the toy problem with the lrugw in the paper with a mixture of vmf yields selections collapsed on to similar points. Using POT entropic_semirelaxed_gromov_wasserstein got okay results with similar clustering issues (not sure if there is a penalty I could add to improve point spread beypund eps) that did not seem to be in the paper, but was way too slow and could not be scaled to anything reasonable, while OTT is much more efficient:

I am left wondering how I can potentially fix my implementation / understanding of the docs, or move on from this paper.

Thank you!
Eric