Nubot Eval

nubot-eval

Neural Unbalanced Optimal Transport via Cycle-Consistent Semi-Couplings — Lubeck et al. (2022) (arXiv:2209.15621, 2022)

What this evaluates

Evaluates the ability of unbalanced optimal transport models to predict distributional shifts, mass creation (proliferation), and mass destruction (cell death) in heterogeneous populations under drug perturbation.

Datasets

• Synthetic Gaussian Mixture — total ?; splits: train (-1), test (-1)
• Single-Cell Perturbation Response (Melanoma) — total ?; splits: train (-1), test (-1)

Metrics

• weighted kernel MMD (primary) — range: [0, ∞)
  • Measures the distributional distance between predicted and observed perturbed cell states. For NUBOT and UBOT GAN, the MMD is weighted by the predicted transport weights; for other baselines, it uses uniform weights.
• weight-count correlation — range: [-1, 1]
  • Correlation between the sum of predicted weights per cell type and the observed population fractions (cell counts) at 8h and 24h post-treatment.

Input / output format

Input: Control cell population features (source distribution μ) and perturbed cell population features (target distribution ν) for each drug treatment and timepoint (8h, 24h).

Output: Predicted transport map (Monge map) T mapping source cells to target cells, plus predicted weights for each sample indicating proliferation (weight ≥ 1) or death (weight ≤ 1).

Scoring recipe

def weighted_kernel_mmd(X_pred, Y_obs, w):
    # X_pred: predicted target features, Y_obs: observed target features
    # w: predicted transport weights for X_pred
    # K: kernel function (e.g., RBF)
    K_xx = K(X_pred, X_pred) * np.outer(w, w)
    K_yy = K(Y_obs, Y_obs)
    K_xy = K(X_pred, Y_obs) * w[:, None]
    m, n = len(X_pred), len(Y_obs)
    return K_yy.mean() - 2 * K_xy.mean() + K_xx.sum() / (m * n)

Common pitfalls

• Ground truth point-to-point matching is destroyed by the measurement process, so evaluation relies on distributional fit and aggregate correlations rather than exact cell-level alignment.
• At 24h, severe treatment-induced cell death reduces observed cell counts, making accurate predictions and evaluation difficult for some drugs.
• Drug treatments can alter the abundance of cell type markers (MelA, Sox9), complicating cell type classification and weight interpretation.

Evidence (verbatim from paper)

For this, we measure the weighted version of kernel maximum mean discrepancy (MMD) between predictions and observations. More details on the evaluation metrics can be found in § D.1. The results are displayed in Fig. 4. NUBOT outperforms all baselines in almost all drug perturbations, showing its effectiveness in predicting OT maps and local variation in mass.

Citation

@misc{lubeck2022nubot,
  title={Neural Unbalanced Optimal Transport via Cycle-Consistent Semi-Couplings},
  author={Lubeck et al. (2022)},
  year={2022},
  note={arXiv:2209.15621}
}

• arXiv: 2209.15621