This work introduces a geometric framework and a novel network architecture for creating correspondences between samples of different conditions.
Under this formalism, the latent space is a fiber bundle stratified into a base space encoding conditions, and a fiber space encoding the variations within conditions. Furthermore, this latent space is endowed with a natural pull-back metric. The correspondences between conditions are obtained by minimizing an energy functional, resulting in diffeomorphism flows between fibers.
We illustrate this approach using MNIST and Olivetti and benchmark its performance on the task of batch correction, which is the problem of integrating multiple biological datasets.