We propose to address this issue by treating the embedded patches of generated image and texts as vectors sampled from two probability distributions, and to use computational optimal transport (OT) [1]} to find the best matching between them.
As its name suggests, OT tries to find the minimal total effort required to “move” all patches towards texts, using the pairwise distance as the cost for measuring said effort. OT brings two advantages over simply taking the mean:
(1) Since patches are randomly sampled, it encourages the intrinsic diversity inside a single generated image.
(2) OT does not involve vector arithmetic in the latent space, sidestepping issues that may arise from the non-existing semantic of a mean prompt vector.
Concretely, we use Sinkhorn's Algorithm [2]}, [3]} for the matching,
in a way that is efficient and, most importantly, differentiable using OTT-JAX [4]}.
Such differentiability is crucial to allow the computation of gradient all the way back to \(z\) .
