Проклятие условий: анализ и улучшение оптимального транспорта для условной генерации на основе потоков

The Curse of Conditions: Analyzing and Improving Optimal Transport for Conditional Flow-Based Generation

March 13, 2025

Авторы: Ho Kei Cheng, Alexander Schwing

cs.AI

Аннотация

Оптимальное транспортное сопряжение с использованием мини-батчей выпрямляет пути в безусловном согласовании потоков. Это приводит к менее требовательным вычислительным затратам при выводе, так как можно использовать меньше шагов интегрирования и менее сложные численные решатели при численном решении обыкновенного дифференциального уравнения на этапе тестирования. Однако в условной постановке оптимальное транспортное сопряжение с мини-батчами оказывается недостаточным. Это связано с тем, что стандартное оптимальное транспортное отображение игнорирует условия, что приводит к условно смещенному априорному распределению во время обучения. Напротив, на этапе тестирования у нас нет доступа к смещенному априорному распределению, и вместо этого мы берем выборки из полного, несмещенного априорного распределения. Этот разрыв между обучением и тестированием приводит к неудовлетворительной производительности. Чтобы устранить этот разрыв, мы предлагаем условное оптимальное транспортное сопряжение C^2OT, которое добавляет условный весовой член в матрицу затрат при вычислении оптимального транспортного назначения. Эксперименты показывают, что это простое исправление работает как с дискретными, так и с непрерывными условиями в задачах 8gaussians-to-moons, CIFAR-10, ImageNet-32x32 и ImageNet-256x256. Наш метод демонстрирует лучшую общую производительность по сравнению с существующими базовыми подходами при различных бюджетах на вычисление функций. Код доступен по адресу https://hkchengrex.github.io/C2OT.

English

Minibatch optimal transport coupling straightens paths in unconditional flow matching. This leads to computationally less demanding inference as fewer integration steps and less complex numerical solvers can be employed when numerically solving an ordinary differential equation at test time. However, in the conditional setting, minibatch optimal transport falls short. This is because the default optimal transport mapping disregards conditions, resulting in a conditionally skewed prior distribution during training. In contrast, at test time, we have no access to the skewed prior, and instead sample from the full, unbiased prior distribution. This gap between training and testing leads to a subpar performance. To bridge this gap, we propose conditional optimal transport C^2OT that adds a conditional weighting term in the cost matrix when computing the optimal transport assignment. Experiments demonstrate that this simple fix works with both discrete and continuous conditions in 8gaussians-to-moons, CIFAR-10, ImageNet-32x32, and ImageNet-256x256. Our method performs better overall compared to the existing baselines across different function evaluation budgets. Code is available at https://hkchengrex.github.io/C2OT