One-step Optimal Transport via Regularized Distribution Matching Distillation

Unpaired domain translation remains a challenging task due to the need of finding a balance between faithfulness and realism. In this paper, we propose a method called Regularized Distribution Matching Distillation (RDMD) that combines the best properties of Optimal Transport (OT) and diffusion-based domain translation methods. Instead of the conventional adversarial training, RDMD utilizes diffusion-based distribution matching, addressing the common shortcomings of OT methods and providing a strong initialization for the trained models. RDMD provides efficient one-step inference, explicitly controls the input-output alignment via regularization of the transport cost and maintains high faithfulness similar to the OT methods. We prove that in theory RDMD approximates the OT map and demonstrate its empirical performance on several tasks, including unpaired image-to-image translation in pixel and latent space and unpaired text detoxification. Empirical results show that RDMD achieves a comparable or better faithfulness-realism trade-off compared to the diffusion and OT baselines.

Machine learning systems can generate realistic images and text, but changing an existing object into a related object from another domain is still hard. For example, we may want to turn a cat photo into a dog, make a face look like anime, or rewrite a toxic sentence in a safer style while preserving its meaning. The difficulty is that we usually do not have exact before-and-after pairs: there is no single “correct dog” for each cat. We introduce Regularized Distribution Matching Distillation (RDMD), a method that trains a model to balance two goals. The output should stay close to the input, preserving important structure and content. At the same time, it should look like a real example from the target domain. RDMD achieves this by combining a mathematical preference for small, meaningful changes with feedback from a pre-trained generative model that helps judge realism. The trained model performs translation in one step, making it faster than methods that repeatedly refine an output. Across image and text tasks, RDMD gives a strong trade-off between preserving the input and producing realistic outputs. This suggests that specialized translation models can remain useful even in the era of large general-purpose generative models.