Optimal Transport based Flow Matching

Below is a clear, practical, step-by-step procedure you can follow to implement Optimal Transport (OT)–based Flow Matching for graph generation. I give the mathematical core, numerical recipes (Sinkhorn / barycentric projection), implementation tips for graphs (node/edge matching and permutation invariance), pseudocode, and hyperparameter suggestions so you can implement end-to-end.

for epoch:
  for batch of target graphs {G1, ...}:
    # 1. Create or sample source noise graphs G0 (same sizes / padded)
    G0 = sample_noise_like(batch_shapes)

    # 2. For each pair (G0_k, G1_k):
    for k in batch:
      C_k = cost_matrix(G0_k.nodes, G1_k.nodes)
      pi_k = sinkhorn(C_k, eps, n_iters)
      X1_match_k = barycentric_projection(pi_k, G1_k.node_features)
      u_k = X1_match_k - G0_k.node_features

      # Form interpolation at random t in (0,1)
      t = Uniform(0,1)
      X_t = (1-t) * G0_k.nodes + t * X1_match_k
      E_t = (1-t) * G0_k.edges + t * G1_k.edges

    # 3. Stack X_t, E_t, t, and targets u
    v_pred = GNN_theta(X_t, E_t, t)
    loss = weighted_mse(v_pred, u, weights_from_pi)
    + regularizers

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

# assume batch of pairs (G0, G1)
for G0, G1 in dataloader:
  # 1. cost and Sinkhorn
  C = compute_cost_matrix(G0.nodes, G1.nodes)
  pi = sinkhorn(C, eps)

  # 2. barycentric projection -> matched targets
  X1_proj = barycentric(pi, G1.node_features)
  u = X1_proj - G0.node_features

  # 3. sample t and form x_t
  t = random_uniform()
  X_t = (1-t) * G0.nodes + t * X1_proj
  E_t = (1-t) * G0.edges + t * G1.edges

  # 4. predict velocity
  v_pred = GNN_theta(X_t, E_t, t)

  # 5. loss (weighted by mass)
  w = pi.sum(dim=0)   # per-source-node mass
  loss = sum_j w_j * ||v_pred_j - u_j||^2 + reg

  optimizer.step()