Pairwise Optimal Transports for Training All-to-All Flow-Based Condition Transfer Model
This work addresses a domain-specific problem in machine learning for handling continuous conditions in transfer learning, offering a novel approach but with incremental improvements over existing methods.
The paper tackles the challenge of learning all-to-all transfer maps among continuous conditional distributions with sparse data by proposing a flow-based method that approximates pairwise optimal transport, demonstrating effectiveness on synthetic, benchmark, and chemical datasets.
In this paper, we propose a flow-based method for learning all-to-all transfer maps among conditional distributions that approximates pairwise optimal transport. The proposed method addresses the challenge of handling the case of continuous conditions, which often involve a large set of conditions with sparse empirical observations per condition. We introduce a novel cost function that enables simultaneous learning of optimal transports for all pairs of conditional distributions. Our method is supported by a theoretical guarantee that, in the limit, it converges to the pairwise optimal transports among infinite pairs of conditional distributions. The learned transport maps are subsequently used to couple data points in conditional flow matching. We demonstrate the effectiveness of this method on synthetic and benchmark datasets, as well as on chemical datasets in which continuous physical properties are defined as conditions. The code for this project can be found at https://github.com/kotatumuri-room/A2A-FM