FlowBender: Treinamento Consciente de Feedback para Fluxos Condicionais Autocorretivos

Resumo

Modelos de difusão e fluxo condicionais frequentemente falham em satisfazer as próprias restrições que definem sua tarefa. Por exemplo, um modelo condicionado por profundidade frequentemente produz imagens cuja profundidade reextraída discorda da entrada, mesmo quando o operador direto — o preditor de profundidade que define a restrição — está disponível tanto durante o treinamento quanto na inferência. Abordagens existentes geralmente se enquadram em duas categorias: modelos supervisionados que tratam o sinal de condicionamento como uma dica estática e ignoram informações de alinhamento na inferência, e métodos baseados em orientação que o consultam por meio de atualizações lineares ajustadas manualmente, tipicamente trocando fidelidade à condição pela plausibilidade da amostra gerada. Argumentamos que a lacuna fundamental em ambos os paradigmas é que o modelo nunca é treinado para utilizar seu próprio erro de alinhamento. Apresentamos o FlowBender, uma estrutura de malha fechada que trata esse erro como uma entrada de primeira classe, treinando a rede para aprender uma política de correção condicionada ao feedback da inferência. A cada passo, uma passagem de antecipação não guiada estima o sinal limpo, um desvio específico da tarefa é calculado por meio do operador direto, e uma passagem de refinamento consome esse sinal para produzir uma velocidade corrigida. Propomos diversas variantes do FlowBender, incluindo uma formulação baseada em gradiente para operadores diferenciáveis e uma variante de ordem zero para cenários não diferenciáveis, como compressão JPEG. Para amostragem eficiente, introduzimos um atalho de etapa anterior que permite correção em malha fechada com custo computacional adicional mínimo. Em tradução de imagem para imagem, restauração e texturização de malhas 3D, o FlowBender supera consistentemente linhas de base supervisionadas padrão, treinamento aumentado com perda de alinhamento e orientação de inferência de última geração, melhorando simultaneamente fidelidade e plausibilidade, em vez de trocar uma pela outra. Página do projeto: https://flow-bender.github.io/

English

Conditional diffusion and flow models routinely fail to satisfy the very constraints that define their task. For instance, a depth-conditioned model often produces images whose re-extracted depth disagrees with the input, even though the forward operator--the depth predictor defining the constraint--is available during both training and inference. Existing approaches generally fall into two categories: supervised models that treat the conditioning signal as a static cue and ignore alignment information at inference, and guidance-based methods that consult it through hand-tuned linear updates, typically trading fidelity to the condition against the plausibility of the generated sample. We argue that the fundamental gap in both paradigms is that the model is never trained to utilize its own alignment error. We introduce FlowBender, a closed-loop framework that treats this error as a first-class input, training the network to learn a correction policy conditioned on inference-time feedback. At each step, an unguided look-ahead pass estimates the clean signal, a task-specific deviation is computed via the forward operator, and a refinement pass consumes this signal to produce a corrected velocity. We propose several variants of FlowBender, including a gradient-based formulation for differentiable operators and a zero-order variant for non-differentiable settings such as JPEG compression. For efficient sampling, we introduce a prior-step shortcut that enables closed-loop correction at a minimal additional computational cost. Across image-to-image translation, restoration, and 3D mesh texturing, FlowBender consistently outperforms standard supervised baselines, alignment-loss-augmented training, and state-of-the-art inference-time guidance, improving fidelity and plausibility simultaneously rather than trading them against each other. Project page: https://flow-bender.github.io/