Re-implementation of "Sinkhorn Regression" (Lei Luo, Jian Pei, Heng Huang; IJCAI-2020) from scratch in Python.
ITU BLG 527E Machine Learning, Spring 2026.
No public source code exists for this paper. The Sinkhorn regression solver,
proximal operators, and alternating minimization are implemented from scratch
using NumPy/SciPy. Baselines are local implementations except ReliefF, which
uses skrebate when available.
The model solves a feature selection problem via Wasserstein distance minimization:
min_{Z, d} Σᵢ Wγ(h(ZᵀAᵢ), h(yᵢ)) + λ · Σᵢ Wγ(h(zᵢ), d)
- A ∈ ℝ^{m×l}: data matrix (m features, l samples)
- Z ∈ ℝ^{m×c}: coefficient matrix to be learned
- d ∈ ℝ^c: Wasserstein barycenter (regularizer target)
- Wγ: entropy-regularized unbalanced optimal transport distance
- Feature importance: row-wise L2-norm of Z
The augmented form introduces U = AᵀZ and alternates between sequential proximal coordinate updates for Z, a proximal U update, the unbalanced barycenter update for d, and Sinkhorn updates for P and P̂.
pip install numpy scipy scikit-learn pandas matplotlib tqdm skrebateDownload/cache the five ASU microarray datasets and the official NCI Prostate-MS archive:
python data/download_datasets.py# Reproduce Tables 1 & 2 (top-20 and top-80 SVM accuracy, all datasets x all methods)
python experiments/run_experiments.py
# Single dataset only
python experiments/run_experiments.py --dataset ALLAML
# Skip baseline comparisons (Sinkhorn only, faster)
python experiments/run_experiments.py --no_baselines
# Parameter sensitivity analysis (Table 3, lambda sweep)
python experiments/sensitivity_analysis.pyResults are saved to results/table_top20.csv and results/table_top80.csv.
# All tests
python -m pytest tests/ -v
# Individual test files
python -m pytest tests/test_sinkhorn.py -v
python -m pytest tests/test_regression.py -v
# Equation-by-equation paper compliance checks
python tests/verify_paper_compliance.pysrc/
sinkhorn_distance.py # Unbalanced Sinkhorn (log-domain), ground metric, histogram operator
sinkhorn_regression.py # Algorithm 1/2/3: alternating minimization, proximal updates
feature_selection.py # Z -> top-k feature indices (row-wise L2-norm)
baselines.py # 8 baseline methods: F-stat, ReliefF, mRMR, T-test, InfoGain,
# RFS (Nie et al. 2010), CGSSL (Li et al. 2014), RSR (Zhu et al. 2015)
utils.py # preprocess, one_hot, sklearn/LIBSVM-style SVM CV
data/
download_datasets.py # Loads five ASU datasets and downloads the NCI Prostate-MS archive
datasets/ # .npz files (not committed -- re-downloadable)
experiments/
run_experiments.py # Tables 1 & 2
sensitivity_analysis.py # Table 3
tests/
test_sinkhorn.py # Sinkhorn distance unit tests (marginals, symmetry, convergence)
test_regression.py # Proximal operator, Z/U/d update correctness tests
Unbalanced Sinkhorn — log-domain to avoid underflow at gamma=0.01. The positive marginals retain their original mass, as required by Eq.10-13:
u <- (a / (K*v))^(mu/(gamma+mu))
v <- (b / (K^T*u))^(mu/(gamma+mu))
P = diag(u) * K * diag(v)
Implemented in log-space using logsumexp for numerical stability.
Proximal operator (Lemma 1) — closed-form for eta(x) = x - a*log(x):
prox_{alpha*eta}(y) = 0.5 * [y - alpha + sqrt((y - alpha)^2 + 4*alpha*a)]
Step sizes — coordinate-wise Lipschitz constants:
- Z update: delta_i = 1 / (2rho||a_i||^2), applied sequentially as in Algorithm 1
- U update: delta_U = 1 / (2*rho)
Initialization — the paper does not specify Z/U initialization. The
experiment configuration records the deterministic choice (nnls by default);
sinkhorn_regression(..., initialization="ridge") is also supported.
After feature-wise z-score standardization, one global offset is added to meet the relaxed model's strict positivity assumption. A separate offset per feature would undo the common centering used by the paper.
Several details are under-specified or inconsistent in the paper:
- Unbalanced vs. balanced Sinkhorn: Eq.10-13 define unbalanced OT with
the
mu/(gamma+mu)exponent, but Algorithms 2-3 print standard balanced scaling updates. This implementation follows Eq.10-13. - Barycenter
d: Eq.9/13 optimisedas a variable, whereas Algorithm 3 uses a fixedepsilon * 1target. The experiment runner uses a hybrid schedule: raw closed-form Eq.13 update for the first 12 iterations, then simplex-projected. This resolves the paper's internal inconsistency. - Table 3 inconsistency: Table 3 is captioned as Top-20, but its
lambda=0.1entries match the Top-80 (Table 2) values for five of six datasets. - Unspecified details: Z/U initialisation,
rho, proximal step-size schedule, Prostate-MS preprocessing, and SVM fold construction are not reported. Our choices are documented inresults/reproduction_report.md.
Paper Section 5.3 specifies a linear SVM with C=1 and 5-fold cross-validation.
This implementation uses sklearn's StratifiedKFold(shuffle=False) with linear
SVC(C=1). Fold construction may differ from the LIBSVM-based evaluation in the
original paper, so absolute accuracy values can differ; within-table comparisons
remain valid. All features are re-standardised before SVM evaluation.
Five microarray datasets come from the ASU Feature Selection Repository. Prostate-MS is downloaded from the official NCI JNCI 7-3-02 archive:
| Dataset | Samples | Features | Classes |
|---|---|---|---|
| ALLAML | 72 | 7,129 | 2 |
| GLIOMA | 50 | 4,434 | 4 |
| LUNG | 203 | 3,312 | 5 |
| Carcinomas | 174 | 9,182 | 11 |
| Prostate-GE | 102 | 5,966 | 2 |
| Prostate-MS | 322 | 15,154 raw | 2 |
Generated tables: results/table_top20.csv and results/table_top80.csv.
Sinkhorn Regression — reproduced vs. paper (SVM 5-fold CV, linear, C=1):
| Dataset | Top-20 ours | Top-20 paper | Top-80 ours | Top-80 paper |
|---|---|---|---|---|
| ALLAML | 95.90% | 96.31% | 97.24% | 97.76% |
| GLIOMA | 86.00% | 74.00% | 72.00% | 70.00% |
| LUNG | 95.09% | 94.78% | 96.06% | 96.71% |
| Carcinomas | 90.18% | 93.14% | 92.52% | 94.48% |
| Prostate-GE | 93.14% | 95.16% | 92.19% | 95.89% |
| Prostate-MS | 87.87% | 98.71% | 84.14% | 100.00% |
| Average | 91.37% | 92.02% | 89.03% | 92.47% |
The Prostate-MS gap reflects an unresolved cohort ambiguity: the paper states 253 BPH + 63 no-evidence samples form the control class, yet no preprocessing specification matches the reported 98.71%/100.00%. All other datasets reproduce within 3 points.
Wasserstein Discriminant Score (WDS) — proposed contribution: WDS achieves 83.97% Top-20 and 90.51% Top-80 average, outperforming Sinkhorn Regression on the Top-80 setting. Notably, WDS on Carcinomas Top-80 (95.93%) exceeds the paper's own Sinkhorn result (94.48%).
See results/reproduction_report.md for
the full per-dataset comparison and discussion of reproduction limits.
Lei Luo, Jian Pei, Heng Huang. Sinkhorn Regression. Proceedings of the 29th International Joint Conference on Artificial Intelligence (IJCAI-2020). Official paper PDF