A Maldição das Condições: Analisando e Melhorando o Transporte Ótimo para Geração Condicional Baseada em Fluxo
The Curse of Conditions: Analyzing and Improving Optimal Transport for Conditional Flow-Based Generation
March 13, 2025
Autores: Ho Kei Cheng, Alexander Schwing
cs.AI
Resumo
O acoplamento de transporte ótimo em minibatches endireita os caminhos no fluxo incondicional. Isso resulta em uma inferência computacionalmente menos exigente, pois menos etapas de integração e solucionadores numéricos menos complexos podem ser empregados ao resolver numericamente uma equação diferencial ordinária no momento do teste. No entanto, no cenário condicional, o transporte ótimo em minibatches se mostra insuficiente. Isso ocorre porque o mapeamento padrão de transporte ótimo ignora as condições, resultando em uma distribuição prévia condicionalmente distorcida durante o treinamento. Em contraste, no momento do teste, não temos acesso à distribuição prévia distorcida e, em vez disso, amostramos da distribuição prévia completa e não tendenciosa. Essa lacuna entre treinamento e teste leva a um desempenho inferior. Para preencher essa lacuna, propomos o transporte ótimo condicional C²OT, que adiciona um termo de ponderação condicional na matriz de custo ao calcular a atribuição de transporte ótimo. Experimentos demonstram que essa correção simples funciona tanto com condições discretas quanto contínuas em 8gaussians-to-moons, CIFAR-10, ImageNet-32x32 e ImageNet-256x256. Nosso método apresenta um desempenho geral melhor em comparação com as linhas de base existentes em diferentes orçamentos de avaliação de funções. O código está disponível em 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