Adds a second biology example to the documentation notebooks.

The notebook presents an application of OTT's Gromov Wasserstein optimal transport to match single-cell points clouds from two different measurement spaces (e.g. mapping gene expressions measurements to chromatine accessibility measurements). It is adapted from Demetci et al, Gromov-Wasserstein optimal transport to align single-cell multi-omics data, ICML 2020 Workshop on Computational Biology, 2020 (pdf here).

This PR introduces some pre-commits, some configuration files, .gitignore and a missing RST template for class (related to docs formatting in #66) Per https://google.github.io/styleguide/pyguide.html#34-indentation and because black doesn't support 2 spaces as indentation, it's now 4 spaces. I've set the max line length to 120 (given current resoution; I always use this), but this can be adapted in pyproject.toml. flake8 pre-commits are failing, I just wanted to start a discussion whether the style is ok. Afterwards, I will start fixing these.

To try this out, run in the root of the project:

pre-commit install  # needs to be done only once; will run `pre-commit` before commiting/pushing
pre-commit run --all-files

closes #43

• Examples here: https://colab.research.google.com/drive/1vncmDEr3t6_OKfVC0PJin8PIRBViPyO0?usp=sharing

• Initialization methods stored in /core/initializers.py

  • Each initializer inherits from base class SinkhornInitializer
    • Class is more flexible in case of neural network initialisers in future
    • Currently added:
      • sorting initializer (n=m) for any cost
      • Gaussian initializer for squared ground cost added, only works for point cloud as need access to x
  • Each initializer has methods init_dual_a and init_dual_b

      def init_dual_a(
            self, ot_problem: LinearProblem, lse_mode: bool = True
        ) -> jnp.ndarray:
          """Initialzation for Sinkhorn potential f.
  

  • The base class also holds default behaviour, initialization for 0 and 1 depending on log lse_mode, also handling entries for 0 weights

• Modification to Sinkhorn api,

  • pass instantiated initializer to Sinkhorn

  Sinkhorn(
     lse_mode=lse_mode,
     threshold=threshold,
     norm_error=norm_error,
     inner_iterations=inner_iterations,
     min_iterations=min_iterations,
     max_iterations=max_iterations,
     momentum=momentum_lib.Momentum(start=chg_momentum_from, value=momentum),
     anderson=anderson,
     implicit_diff=implicit_diff,
     parallel_dual_updates=parallel_dual_updates,
     use_danskin=use_danskin,
     potential_initializer=potential_initializer,
     jit=jit
  )
  

  • Can be passed to sinkhorn functional wrapper

  gaus_init = init_lib.GaussianInitializer()
  
  @jax.jit
  def run_sinkhorn_gaus_init(x, y, a=None, b=None):
      sink_kwargs = {'jit': True, 
                  'threshold': 0.001, 
                  'max_iterations': 10**5, 
                  'potential_initializer': gaus_init}
  
      geom_kwargs = {'epsilon': 0.01}
      geom = PointCloud(x, y, **geom_kwargs)
      out = sinkhorn(geom, a=a, b=b, **sink_kwargs)
      return out
  

  • Initialization strategy will be overridden if init_dual_a/ init_dual_b vectors are passed to solver
  • Addressed this issue https://github.com/ott-jax/ott/issues/84#issue-1285013301 in tools wrapper

According to the different notebooks, there are two ways to run Sinkhorn, which seem to be equivalent.

For example, suppose we have two point clouds x and y for which we want to calculate the corresponding reg_ot_cost. The two ways to compute this reg_ot_cost are:

THRESHOLD = 1e-3
MAX_ITERATIONS = 2_000
EPSILON = 1e-3
cost_fn = SqEuclidean()

def get_reg_cost_1(x, y):
    "Functional."
    cost_matrix = cost_fn.all_pairs(x, y)
    geom = geometry.Geometry(cost_matrix, epsilon=EPSILON)
    return sinkhorn.sinkhorn(
        geom,
        **{"threshold": THRESHOLD,
           "max_iterations": MAX_ITERATIONS} # sinkhorn kwargs 
        ).reg_ot_cost
    
def get_reg_cost_2(x, y):
    "Via Sinkhorn class."
    cost_matrix = cost_fn.all_pairs(x, y)
    geom = geometry.Geometry(cost_matrix, epsilon=EPSILON)
    return sinkhorn.Sinkhorn(
            threshold=THRESHOLD,
            max_iterations=MAX_ITERATIONS,
        )(linear_problem.LinearProblem(geom)).reg_ot_cost

What I understand (and which may be wrong): The only thing that differs in these two ways of running the Sinkhorn algorithm, is that in the first one, we don't directly instantiate an element of the Sinkhorn class, but it will be done inside the sinkhorn.sinkhorn function, by the make function. Whereas in the second one, we instantiate an object of the Sinkhorn class and then we call this object to run the Sinkhorn algorithm. It seems that the first way is more "functional" than the second.

Both of these ways are used in notebooks, for example the first is used in the notebook "Meta OT and Sinkhorn Initializers" https://ott-jax.readthedocs.io/en/latest/notebooks/MetaOT.html and the second is used in the notebook "Point clouds" https://ott-jax.readthedocs.io/en/latest/notebooks/point_clouds.html (which seems to have been updated recently).

However, they give very different running times: the first seems to be significantly slower than the second.

If I run the following code on my machine:

# generate data
rng = random.PRNGKey(0)
rngs = random.split(rng, 2)
NUM_SAMPLES, DIM = 1024, 10
x = random.uniform(rng, minval=-1, maxval=1, shape=(NUM_SAMPLES, DIM))
y = random.normal(rng, shape=(NUM_SAMPLES, DIM))

%timeit -n100 -r3 get_reg_cost_1(x, y).block_until_ready()
%timeit -n100 -r3 get_reg_cost_2(x, y)

I get that get_reg_cost_1 is more than 70 times slower that get_reg_cost_2:

720 ms ± 4.1 ms per loop (mean ± std. dev. of 3 runs, 100 loops each)
10 ms ± 319 µs per loop (mean ± std. dev. of 3 runs, 100 loops each)

So I guess I'm making a mistake somewhere in the use of solvers? Or if not, how can we explain this difference?

Thank you very much!

This is an initial version of a Meta OT initializer along with a notebook demonstrating it's usage. Here is how the notebook renders in the website, and here are how the code docs render:

Quick edit: I believe my understanding was flawed, sorry. It could be correct as it currently is. I (may wrongly) assumed Y would denote the target distribution and we were looking for the map from X to Y (hence, source to target). However, it seems that Y is sampled from Q, which is stated as the source distribution in the paper (and not only in the pseudocode, but also at other locations). The goal is to learn the map from Q to P. I realized this after having a more detailed look at the terminology used there again. Now that I have thought about it, my assessment seems to have been overhasty. I would still be interested if someone could clarify my understanding, however, so I get it actually correct. :-)

Perhaps my understanding is also wrong (which I cannot rule out).

According to the original paper, and my understanding, the g potential should be applied to the target distribution and the f potential to the source (see, e.g., Eq. 5 or Eq. 9). In the current version, this is switched.

I suspect this mix-up comes from the fact that Algorithm 1 in the paper has the source and target also mixed up in the pseudocode. Here, the source distribution is indeed given to the g potential and the target distribution to f. However, following the equations, I think this is simply a mistake (although I am not 100% sure).

My suspicion is further backed by the original implementation of the authors, where the source samples are given to the f potential and the target to the g.

After changing this accordingly, I feel that this improved the convergence speed as well (from a few first test runs).

Hello,

It seems that the custom_vjp (defined with either _iterations_implicit.defvjp(_iterations_taped, _iterations_implicit_bwd) in sinkhorn.py for the implicit differentiation or fixpoint_iter_backprop.defvjp(fixpoint_iter_fwd, fixpoint_iter_bwd) in fixed_point_loop for the unrolling) does not seem to be used for the computation of higher order derivatives. I believe that this is also related to the error discussed previously in this issue. This is not evident with the computation of jax.hessian (which is implemented as jax.jacfwd(jax.jacrev), as the second derivation uses forward mode automatic differentiation, which is compatible with the jax.lax.while_loop. Therefore, no error is raised even if the custom_vjp is ignored the second time.

I think that this can be seen by adding a breakpoint in the _while_loop_jvp of JAX's control_flow.py. For instance, computation of jax.hessian passes through _while_loop_jvp first, then one time throughfixpoint_iter for implicit diff (or fixpoint_iter_fwd for unrolling) and then one time through _iterations_implicit_bwd for implicit diff (or fixpoint_iter_bwd for unrolling). In the case where a jax.lax.scan is forced by setting min_iterations equal to max_iterations (and both Jacobians are computed with reverse mode), instead of the initial pass through _while_loop_jvp, one gets in the end a pass through _scan_transpose of JAX's control_flow.py.

In either case, I think that if the custom_vjp was not ignored during rederivation, one should get two passes through _iterations_implicit_bwd (equivalently fixpoint_iter_bwd), right?

I am not sure how this could be fixed. The only relevant info that I could find in JAX's documentation for custom_vjp was this:

"Notice that f_jvp calls f to compute the primal outputs. In the context of higher-order differentiation, each application of a differentiation transform will use the custom JVP rule if and only if the rule calls the original f to compute the primal outputs. (This represents a kind of fundamental tradeoff, where we can't make use of intermediate values from the evaluation of f in our rule and also have the rule apply in all orders of higher-order differentiation.)"

I hope I am not missing something. If there is a solution to this, it would be great, as what I need to compute is jax.jacrev(jax.grad) for the sinkhorn divergence and the forced scan option causes a significant computational overhead, especially for the two autocorrelation terms Pxx, Pyy.

Many thanks!

The segmented approaches to pad efficiently point clouds rely on padding with "default" (e.g. zero) vectors. These padded vectors create spurious entries in cost matrices, but when running Sinkhorn, these entries are ignored thanks to the fact that these entries have 0 weights (in the respective a and b weight vectors). However, these entries would play a role when computing scalings (such as mean cost).

A possible approach would be to redefine some of those statistics (e.g. max, mean etc..) to only apply to entries that have a positive weight (e.g. C_ij is considered iff a_i b_j >0). For the mean cost, the mean might be weighted by a_i b_j as well.

TL;DR in this colab we provide an example for our failure in obtaining a valid mapping using low-rank.

problem setup: In this example data set we look into mapping spatial transcriptomics at single cell resolution from mouse embryonic tissues across two time-points. so the quadratic term accounts for distances in spatial coordinates and the linear captures distances in gene-expression.

evaluation of the mapping: As an initial sanity check we look at the cell-transition table, that is the transition matrix with entries grouped by cell types (we asses the row-stochastic, forward, setting). the naive assumption is that cells of the same type, e.g. brain will be mapped mainly mapped to themselves. Evaluating the regular FGW and FGW (unbalanced) this is indeed what we observe. However, for low-rank we get a matrix with constant columns. We observed a similar phenomena at different time-points. Comparing the results we can see hints for the constant columns as they are cell-types also favored in the regular regime.

hi, we noticed that one can accidentally pass online=True using gromov_wasserstein and this will propagate through **kwargs. Eventually mapping will fail at some point due to dimensionality assertion but we thought it may be useful to catch it earlier and warn. What do you think?

Currently, it is not possible to set parameters like max_iterations and min_iterations of the inner Sinkhorn loop when calling WassersteinSolver or GromovWasserstein because of the same argument names.

It would be great to rename the outer loop parameters to have more control about the algorithm.

Hi @marcocuturi @bunnech @michalk8, here is a WIP PR with the updates we've discussed to the neural dual solver. It improves the example from:

to this:

My paper has some more context on these updates adding some more options for fine-tuning and amortizing the conjugate approximation. I've marked this as a WIP because I'm still debugging the inverse map in the example, but otherwise, the rest of the code is ready for an initial review as it would be good to get your thoughts on the core updates and API/doc changes. I can also call and chat a little more about these soon if that would be helpful. Here are quick summaries of my updates to the code and notebook.

Updates to the core neural solver

1. swap the order of f and g potentials (\cc https://github.com/ott-jax/ott/issues/182#issuecomment-1344895638)
2. add an option to add a conjugate solver that fine-tunes g's prediction when updating f, along with a default solver using JaxOpt's LBFGS optimizer
3. add amortization modes to learn g: regression and objective (Makkuva et al.)
4. an option to update both potentials with the same batch (in parallel) in addition to the separate inner loop Makkuva et al. uses
5. the option to model the gradient of the g potential with a neural network and a simple MLP implementing this instead of the ICNN from Makkuva et al.
6. a callback so the training progress can be monitored (debugging the notebook/code was very difficult without this)

Updates to the notebook/example

(I'm still tweaking these around a little as it's still a little suboptimal)

1. Update f ICNN size from [64 ,64, 64, 64] to [128, 128] and modified the architecture
2. Use an MLP to model the gradient of g
3. Use the default JaxOpt LBFGS conjugate solver and update g to regress onto it
4. Increase batch size from 1000 to 10000
5. Sample the data in batches rather than sequentially
6. Tweak Adam (use weight decay and a cosine schedule)

I've also tried running the code with an approach closer to the older implementation, i.e., the approach from Makkuva et al. that models g with an ICNN and uses 10 inner objective-based updates, but can't get it to work as nicely as the version with conjugate fine-tuning+regression.

Remaining discussion and TODOs

• [ ] Decide on best format for the arguments. E.g., should the conjugate solver be a function and the amortization mode a string?
• [ ] Decide on the best defaults to use (a lightweight LBFGS conjugate solver for fine-tuning that regresses g onto it, or something closer to the approach in Makkuva et al.?)
• [ ] Improve the documentation and notebook to clarify these options
• [ ] Fix the inverse map in the notebook
• [ ] Think about any tests that would be useful to add for these updates
• [ ] Update other places using the Neural OT code (like the initialization schemes notebook)
• [ ] Decide on the default ICNN architecture
• [ ] Add ability to save/load the parameters and ability to re-initialize the solves with them

Without jitting, some properties as SinkhornOutput.primal_cost (using online point cloud geometry), as well as functions DualPotentials.Transport or k_means cause OOM issues. Previously, jitting was done by default, see #192. Question is whether to re-introduce the jit argument to solvers (would be again True by default), always jit the classes or try another solution.

in one way or another...

Use case: when sinkhorn is run, the marginal deviation is used as a stopping criterion.

By default the 1-norm is used: https://github.com/ott-jax/ott/blob/3ebac6369acfe848930e1e2abb69b61ecea7e5cc/src/ott/solvers/linear/sinkhorn.py#L455

and fed into error computation. https://github.com/ott-jax/ott/blob/3ebac6369acfe848930e1e2abb69b61ecea7e5cc/src/ott/solvers/linear/sinkhorn.py#L209

Typically, this involves comparing two marginal (probability) vectors of size n or m. Because those vectors are constrained to lie in the simplex of that dimension, the tolerance (what constitutes a small deviation) should scale graciously (if needed with n or m), and depend on the norm that is used.

Also, an alternative would be to add the Hellinger distance as a norm to control Sinkhorn...

Hi, for computations than can take several hours, it would be great if OTT-jax reported progress. I can see at least two types of progress reports:

1. In some computations, the number of steps is known in advance, for instance if using the optional batch_size parameter of PointCloud and using Sinkhorn or LRSinkhorn: each Sinkhorn iteration is a fixed set of submatrices operations.

2. In some other cases the number of steps is indeterminate because we're waiting for some convergence criterion. But OTT-jax could in these cases report the number of iterations performed so far and the value associated with the convergence criterion, same as you would report a loss in a training loop.

Because reporting progress is a side effect, jax has introduced the host_callback module. I think it's worth a try.

I don't think OTT-jax should depend on progress report libs; it would be instead the responsibility of the user to provide their own callback function, such as tqdm.

Would you be open to this?

For some transport problems, when searching for nearest neighbors, it is useful to analyze parts of the cost matrix. Until now, version 0.3.1, the API only offers a way to materialize and return the entire cost matrix, which is sometimes prohibitively expensive, or impossible (OOM).

I would like to API to offer a way to return only rows and columns of the cost matrix, or the documentation show how to do that.

The good news is that, for problems solved with Sinkhorn, extracting a few rows of the cost matrix is simple thanks to jax:

geom = pointcloud.PointCloud(...)  # n points in the input domain
lin_prob = linear_problem.LinearProblem(...)
out = sinkhorn.Sinkhorn()(lin_prob)

# method 1 - full materialization then extract row k
# C = geom.cost_matrix
# C_k = C[k, :]

# alternative - extract row k
u = jnp.zeros((1,n)).at[k].set(1.0)
C_k = out.apply(u)

Would it make sense for you?

Is it possible to run segment_sinkhorn with different sinkhorn_kwargs for each segment?

For example, with two PoinClouds x and y, each with two segments: (x[0:1024], y[0:1024]) and (x[1024:2048], y[1024: 2048]) can segment_sinkhorn be used to run Sikhorn in parallel on (x[0:1024], y[0:1024]) and (x[1024:2048], y[1024:2048]) but with different sinkhorn_kwargs?

If this is not possible, would you consider adding this feature?

In particular, it would be useful to compute the sinkhorn_divergence with segment_sinkhorn while keeping the benefits of the particular sikhorn_kwargs to speed up the computation of symmetric terms.

More precisely, the computation of the sinkhorn_divergence between two PoinClouds x and y involves the computation of 3 transports, namely between the 3 segments: (x,y), (x,x) and (y,y). With segment_sinkhorn one could compute these 3 transports in parallel. On the other hand, the terms induced by the symmetrical segments (x,x) and (y,y) can be computed faster using sinkhorn_kwargs={parallel_dual_updates=True, momentum=0.5, chg_momentum_from=0, anderson_acceleration=0}. So to be optimal, we should be able to run segment_sinkhorn with these particular sinkhorn_kwargs for the segments (x,x) and (y,y) but not for the segment (x,y), hence the usefulness of being able to use different sinkhorn_kwargs.

Thank you very much!