# Enhancing selectivity using Wasserstein distance based reweighing

arXiv: https://arxiv.org/abs/2401.11562 · HTML: https://mlsanity.com/paper/2401.11562/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 12 of its 40 references that are in the map.
- First posted on arXiv 21 Jan 2024.

## Builds on (most cited first, 12 of 12)

- [Neural Message Passing for Quantum Chemistry](https://mlsanity.com/paper/1704.01212/) (2017)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [Joint Distribution Optimal Transportation for Domain Adaptation](https://mlsanity.com/paper/1705.08848/) (2017)
- [Molecular Graph Convolutions: Moving Beyond Fingerprints](https://mlsanity.com/paper/1603.00856/) (2016)
- [Optimization of Molecules via Deep Reinforcement Learning](https://mlsanity.com/paper/1810.08678/) (2018)
- [Stochastic gradient descent performs variational inference, converges to limit cycles for deep networks](https://mlsanity.com/paper/1710.11029/) (2017)
- [Information-Theoretic Generalization Bounds for Stochastic Gradient Descent](https://mlsanity.com/paper/2102.00931/) (2021)
- [High-dimensional limit theorems for SGD: Effective dynamics and critical scaling](https://mlsanity.com/paper/2206.04030/) (2022)
- [Formalizing Generalization and Robustness of Neural Networks to Weight Perturbations](https://mlsanity.com/paper/2103.02200/) (2021)
- [Machine learning on DNA-encoded libraries: A new paradigm for hit-finding](https://mlsanity.com/paper/2002.02530/) (2020)
- [Best-Effort Adaptation](https://mlsanity.com/paper/2305.05816/) (2023)
- [On Label Shift in Domain Adaptation via Wasserstein Distance](https://mlsanity.com/paper/2110.15520/) (2021)

## Abstract (arXiv)

Given two labeled data-sets S and T, we design a simple and efficient greedy algorithm to reweigh the loss function such that the limiting distribution of the neural network weights that result from training on S approaches the limiting distribution that would have resulted by training on T. On the theoretical side, we prove that when the metric entropy of the input datasets is bounded, our greedy algorithm outputs a close to optimal reweighing, i.e., the two invariant distributions of network weights will be provably close in total variation distance. Moreover, the algorithm is simple and scalable, and we prove bounds on the efficiency of the algorithm as well. As a motivating application, we train a neural net to recognize small molecule binders to MNK2 (a MAP Kinase, responsible for cell signaling) which are non-binders to MNK1 (a highly similar protein). In our example dataset, of the 43 distinct small molecules predicted to be most selective from the enamine catalog, 2 small molecules were experimentally verified to be selective, i.e., they reduced the enzyme activity of MNK2 below 50% but not MNK1, at 10μM -- a 5% success rate.