FlowBender: Feedback-Aware Training for Self-Correcting Conditional Flows

Abstract: 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/

• The paper introduces FlowBender, a closed-loop feedback framework that integrates alignment error into the model to improve condition fidelity and sample plausibility.
• It employs a two-pass training strategy—comprising a look-ahead pass and a refinement pass—with variants that support both differentiable and black-box operators.
• Experimental results in image restoration, super-resolution, and 3D texturing demonstrate significant improvements in metrics like PSNR and FID over traditional methods.

FlowBender: Feedback-Aware Training for Self-Correcting Conditional Flows

Motivation and Problem Statement

Conditional diffusion and flow models—such as those leveraging Flow Matching (FM) and ControlNet—are designed to produce samples aligned with external signals (e.g., depth maps, edge maps, or 3D geometry) while maintaining plausibility on the target manifold. A well-documented failure mode of these models is their inability to enforce consistency with the condition, even when the forward operator for measuring deviations is available during both training and inference.

Traditional approaches fall into two categories: supervised (open-loop) models that use the conditioning signal as a static cue and ignore active feedback, and guidance-based methods that perform hand-tuned linear updates at inference via measurement gradients. Both paradigms ultimately decouple the model from its alignment error, resulting in either drift from conditioning constraints or violation of the data manifold.

The FlowBender Framework

FlowBender reframes conditional generation as a closed-loop process. The core innovation is to inject the model's alignment error—quantified via a task-specific feedback signal derived from applying the forward operator to a clean-signal estimate—directly as an input to the network during both training and inference. This is achieved via a two-pass per-step protocol:

1. Look-ahead Pass: The model predicts an initial clean signal estimate with the feedback channel set to zero.
2. Feedback Computation: The feedback signal is computed as a function of the alignment error between the clean estimate and the condition (via the forward operator).
3. Refinement Pass: The model then processes the actual input, augmented by the feedback signal, to output the corrected velocity or update.

   Figure 1: FlowBender training and inference combine an unguided look-ahead and a feedback-aware refinement pass, enabling the model to internalize and leverage alignment error.

The authors propose first-order (gradient-based), zero-order (residual-based), and hybrid feedback formulations. Particularly, the zero-order variant supports non-differentiable or black-box operators, significantly broadening the domain of applicability.

For efficiency, a prior-step shortcut approximates the look-ahead pass using the cached clean estimate from the previous step, reducing the cost to only $N+1$ evaluations for $N$ ODE steps.

Comparative Analysis to Existing Methods

Open-loop methods (e.g., standard supervised fine-tuning, ControlNet, LoRA adapters) do not react to misalignment at inference. Guidance-based methods consult feedback only at test time and apply it through hand-tuned, usually scalar-weighted, updates, creating a misalignment between training and inference. Moreover, these approaches—especially as demonstrated in 2D toy problems—fail to satisfy both fidelity (condition satisfaction) and plausibility (sample realism) simultaneously.

Figure 2: Standard conditional generation and linear guidance either violate class constraints or diverge from the data manifold, whereas FlowBender aggregates alignment feedback non-linearly to satisfy both.

FlowBender's closed-loop, feedback-aware architecture, in contrast, facilitates simultaneous gains in both dimensions. Orthogonal decomposition shows the learned correction policy is not congruent with mere scalar reweighting of gradient guidance: over 80% of the correction energy in experiments lies orthogonal to the naïve guidance direction, signifying a nontrivial exploitation of alignment error.

Experimental Results

Image-to-Image Translation

Four tasks were considered: super-resolution, depth-to-RGB, edge-to-RGB, and JPEG restoration. FlowBender consistently outperformed baseline supervised and inference-time guidance methods across all fidelity and plausibility metrics. For example, on the super-resolution task, FlowBender's zero-order variant achieved a PSNR of 39.25 dB, compared to 34.35 dB for standard fine-tuning and 43.02 dB for IT guidance, the latter exhibiting severe deterioration in FID (18.96 vs. FlowBender's 3.36).

Qualitative results further highlight the ability of FlowBender to avoid the condition-drift and off-manifold failures seen in open-loop and guidance-based approaches, especially in tasks with neural or non-differentiable forward operators.

Figure 3: FlowBender eliminates condition drift and artifacts visible in both supervised and guidance-based baselines for Depth-to-RGB and Edge-to-RGB generation.

JPEG Restoration

This non-differentiable task demonstrates the utility of zero-order feedback. FlowBender yields significantly higher PSNR and SSIM, and lower LPIPS and FID, compared to supervised baselines.

Figure 4: FlowBender reduces both quantization artifacts and unwanted color shifts in JPEG restoration, outperforming baselines.

3D Mesh Texturing

Conditioned on rendered images, FlowBender achieves the best multi-view masked-PSNR, SSIM, and FID, as well as higher CLIP similarity to the conditioning, on both the Objaverse and Toys4K benchmarks. IT guidance marginally improves fidelity but at the expense of plausibility.

Figure 5: FlowBender produces 3D textures consistent with conditioning images and free of common inconsistencies found in baselines.

Guidance Scale and Efficiency Ablations

An increase in the guidance strength parameter in classifier-free guidance can, in FlowBender, further enhance fidelity without substantial plausibility deterioration, in contrast to baseline methods.

Figure 6: Optional classifier-free guidance can further improve adherence to the condition in FlowBender, visible in edge map recovery.

Abation studies demonstrate that the prior-step shortcut maintains much of the improvement of the full two-pass regime at a substantially reduced computational cost.

Figure 7: The shortcut policy exploits the rising similarity of feedback across successive steps, yielding efficient yet effective closed-loop inference.

Theoretical and Practical Implications

FlowBender demonstrates that closed-loop, feedback-aware training in flow-based generative models provides an effective mechanism for learning nuanced correction policies that cannot be expressed by linear guidance. This transition to self-correcting models is relevant not only for classical inverse problems and restoration, but also for tasks involving black-box measurement operators or complex controls, such as 3D rendering pipelines or unknown physical processes.

Practically, this architecture can be integrated into existing conditional generative models with modest modifications, requiring only the expansion of model input to accept feedback signals and the addition of a two-pass training loop. The zero-order correction opens new avenues for model deployment in settings where gradients are unavailable (e.g., proprietary APIs).

Future work can explore single-pass training regimes that cache intermediate clean-signal estimates, or more sophisticated architectures capable of leveraging additional modalities or history in their correction policies. Extending feedback-aware training to multi-stage or hierarchical generative processes is also a promising direction.

Conclusion

FlowBender formalizes and validates a closed-loop, feedback-driven training protocol for conditional flow and diffusion models. By internalizing alignment error as part of the model input, it enables the learning of correction policies that resolve the fundamental fidelity–plausibility trade-off of prior methods. This framework generalizes across tasks and operators, and empirically outperforms state-of-the-art alternatives, both in conditional satisfaction and manifold adherence, with minimal additional training or inference complexity. The approach establishes a scalable blueprint for self-correcting generative models across modalities and domains.