PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Daniel Gilo, Sven Elflein, Ido Sobol, Or Litany

arXiv:2606.20404v1cs.CV

TL;DR

Conditional diffusion and flow models often fail to satisfy their conditioning constraints because existing supervised and guidance-based approaches do not train the model to use alignment error. FlowBender learns a closed-loop correction policy from inference-time feedback and consistently improves both conditional fidelity and sample plausibility across image translation, restoration, and 3D texturing.

  • Problem

    Existing conditional diffusion and flow models often produce samples whose measurements disagree with the conditioning signal, while leaving available alignment feedback unused at inference.

  • Method

    FlowBender trains a closed-loop two-pass policy that derives task-specific alignment error from an unguided look-ahead and uses it to correct the sampling update.

  • Results

    Across image translation, restoration, and 3D texturing, FlowBender consistently outperforms supervised, alignment-loss-augmented, and inference-time guidance baselines on fidelity and plausibility.

  • Takeaways & Limitations

    The results support learned use of alignment error as a feedback signal for improving conditional fidelity without sacrificing sample plausibility across these tasks.

  • Takeaways & Limitations

    FlowBender requires an additional model evaluation per training iteration, and performance sometimes improves further when combined with classifier-free guidance.

Abstract

from arXiv · show

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/

1 Introduction

FlowBender addresses conditional diffusion and flow models’ failure to maintain alignment by training a closed-loop correction policy that uses inference-time deviation feedback. It supports gradient-based and zero-order correction, improves fidelity and plausibility across several tasks, and avoids the trade-offs of hand-tuned guidance.

  • Motivation: Conditional diffusion and flow models must balance fidelity to an external signal y with plausibility on the target data manifold, but existing methods routinely fail alignment.Examples include text, depth, and geometric constraints.
  • Motivation: Supervised models treat conditioning as a static cue, while guidance methods use manually tuned linear updates that can under-satisfy constraints or leave the data manifold.Both paradigms fail to train the network to use its own alignment error.
  • Method: FlowBender trains a closed-loop correction policy that derives task-specific deviation feedback from an unguided look-ahead estimate and feeds it into a refinement pass.The framework treats alignment error as a first-class input rather than consulting it only through inference-time guidance.
  • Method: 80% of second-pass correction energy lies orthogonal to the gradient, indicating that FlowBender uses a nonlinear feedback policy beyond scalar-weighted guidance.The framework also includes a zero-order variant that directly uses raw measurement-space error for non-differentiable or black-box operators.
  • Results: FlowBender consistently outperforms supervised training, alignment-loss-augmented training, and state-of-the-art inference-time guidance across 3D texturing, JPEG restoration, and image translation.It improves plausibility, including FID, while enhancing fidelity rather than reproducing guidance’s fidelity–plausibility trade-off.

2 Related Work

Prior conditional models use static conditioning or Bayesian guidance, but open-loop training does not expose alignment deviation to the model. FlowBender instead frames explicit alignment error as feedback for closed-loop self-correction, drawing on learned iterative refinement.

  • Conditional Sampling via Open-Loop Training: Conditional diffusion and flow-matching models commonly use static signals, adapters, or paired data to impose controls such as depth and masks.These approaches parameterize scores or velocity fields from conditioning signals and often rely on adapters such as ControlNet or LoRA.
  • Conditional Sampling via Open-Loop Training: Enhanced conditional training schemes remain open-loop because they do not provide the model with an estimate of its deviation from the target.FlowBender makes alignment error an explicit input, turning generation into a closed-loop system capable of active self-correction.
  • Conditional Sampling via Bayesian Guidance: Bayesian guidance decomposes the conditional score into prior and likelihood terms, with likelihood approximations spanning classifier, classifier-free, and distance-based methods.Distance-based methods compare a predicted clean signal with the measurement, while the likelihood itself involves an intractable integral.
  • Conditional Sampling via Bayesian Guidance: These modular guidance approaches rely on heuristic weighting, creating a trade-off between conditional alignment and sample quality.The passage identifies heuristic weighting as the source of this trade-off across Bayesian and training-free guidance methods.
  • Learned Iterative Refinement: Learned iterative refinement trains networks to use error signals, extending from learned optimizers to inverse-problem reconstruction and computer-vision applications.The paradigm predicts update rules from gradient information or incorporates measurement-space reconstruction error as input.

3 Preliminaries: Conditional Flow Matching Models

This section defines flow matching as modeling a probability path from Gaussian noise to the data distribution and uses conditional flow matching for models receiving task-specific external signals. It also introduces the learned vector field and its ability to estimate the clean signal during sampling.

  • Flow Matching: Flow Matching models a probability path interpolating between Gaussian noise p0(x0) ∼ N(0, I) and the target data distribution p1(x1).The path is denoted pt(xt).
  • Conditional Flow Matching: Conditional flow matching denotes flow models that receive an external task-specific conditioning signal c, such as a corrupted image or depth map.This usage differs from the framework-specific meaning of conditional flow matching based on paths constructed relative to data points x1.
  • Clean-Signal Estimation: The learned vector field provides a principled point estimate of the clean signal x̂1 at any step t.For optimal transport flow matching, at = t and σt = 1 − t.

4 Feedback-Aware Conditional Flows

FlowBender turns conditional generative sampling into a closed-loop process that feeds task-specific alignment error back into the network. It supports gradient-based and derivative-free feedback, trains unguided and feedback-aware modes jointly, and enables efficient two-pass correction with a prior-step shortcut.

  • Framework: FlowBender transforms generative sampling into a closed-loop system by making alignment feedback a first-class network input.The framework is designed to be architecture-agnostic through an auxiliary feedback input.
  • Two-Pass Execution: At each timestep, an unguided look-ahead estimates the clean signal, after which a refinement pass uses derived feedback to produce the final velocity.The feedback input is set to zero for the look-ahead pass, and the resulting signal is treated as constant during training for stability and efficiency.
  • First-Order Feedback: For differentiable operators, first-order feedback uses an alignment-loss gradient, with a shortcut gradient that omits the denoiser Jacobian and reduces memory overhead.The gradient is concatenated to x_t along the channel dimension to provide an explicit correction direction.
  • Zero-Order Feedback: Zero-order feedback supplies a task-specific measurement-space residual, avoiding gradients through H and supporting non-differentiable operators such as JPEG compression or physical simulations.The model learns to map measurement-space errors to signal-space updates, including in black-box systems where gradients are unavailable.
  • Training: Joint training supports unguided and feedback-aware modes while preserving look-ahead reliability, which is essential for deriving dependable inference-time feedback.Randomly replacing feedback with a null vector helps maintain accuracy in the unguided regime.
  • Efficient Sampling: The prior-step shortcut can reduce an N-step trajectory to N + 1 evaluations, nearly matching vanilla sampling efficiency while retaining closed-loop correction.The threshold controls the trade-off between train-test compatibility and sampling speed; t_thresh = 1 preserves two-pass execution, whereas t_thresh = 0 uses a single-pass loop after bootstrapping.

5 Experiments

Experiments across image-to-image translation and 3D mesh texturing show that FlowBender consistently improves fidelity and plausibility over supervised baselines and inference-time guidance. Ablations further demonstrate efficient shortcut correction and nonlinear use of feedback.

  • Image-to-Image Translation: FlowBender is evaluated on four image-to-image tasks spanning closed-form, neural, and non-differentiable forward operators, using Stable Diffusion 3.5 Large with ControlNet.Training uses Unsplash-25K, with 20k training and 5k test images, and 40-step Euler sampling.
  • Image-to-Image Translation: Across all four image-to-image tasks, FlowBender variants consistently outperform baselines, improving fidelity and plausibility simultaneously, whereas IT Guidance trades one against the other.IT Guidance can remain competitive on fidelity metrics but produces worse image quality, reflected by higher FID, and vice versa.
  • 3D Mesh Texturing: On 3D mesh texturing with TRELLIS-2, FlowBender consistently outperforms all baselines, with the ∇ˆx1 variant achieving the strongest performance and recovering fine-grained details.IT Guidance improves fidelity but does not match FlowBender in plausibility.
  • Efficiency Ablation: Reducing tthresh lowers computational cost toward n+1 NFEs while FlowBender remains ahead of Standard FT, despite a relative performance dip from the shortcut.The fidelity gap at tthresh = 0 is especially notable, showing corrective benefits with minimal overhead over open-loop sampling.
  • Hyperparameter Ablation: The best null-feedback setting is pun = 0.1, balancing reliable look-ahead estimation against training iterations available for second-pass refinement.Nonzero pun supports look-ahead estimates, but increasing it reduces iterations for learning refinement.
  • Correction Analysis: FlowBender’s learned correction is predominantly orthogonal to the gradient, while cos(∆v, sgrad_t) = 0.42 ± 0.11 confirms gradient utilization across noise levels.The parallel component contributes approximately 20% of correction energy, indicating nonlinear flow bending beyond additive guidance.

6 Conclusion and Limitations

FlowBender addresses the open-loop limitation of conditional flow models by using alignment errors as inputs to a learned two-pass correction policy. Its prior-step shortcut improves inference efficiency, but training still requires an extra model evaluation per iteration.

  • Conclusion: FlowBender treats alignment errors as first-class inputs, training conditional flow models to correct their initial predictions.The framework replaces the open-loop behavior of existing models with feedback-aware correction.
  • Conclusion: FlowBender replaces hand-tuned guidance with a learned, non-linear two-pass correction policy that enhances conditional fidelity and sample plausibility.The passage states that these gains improve over supervised and inference-time approaches, though the comparison is truncated.
  • Limitations: Training requires an additional model evaluation per iteration to derive feedback, increasing the computational budget for fine-tuning.The prior-step shortcut enables efficient inference, while future work could use cached prior-step predictions to restore single-pass efficiency during training.

A Implementation Details

FlowBender trains with an unguided look-ahead followed by feedback-conditioned refinement, optionally dropping the feedback signal. At inference, it reuses prior-step estimates when possible to reduce look-ahead computation.

  • Feedback-Aware Training: Training samples noisy states, estimates the clean signal with an unguided pass, computes feedback, and feeds stopped-gradient feedback into a refinement pass.Conditioning dropout instead replaces the feedback input with zero; refinement minimizes L_FA = ∥v_ref − u_t∥².
  • Feedback-Aware Inference: Inference computes feedback from an unguided look-ahead estimate, then uses it in a refinement pass to integrate the next state.The initial step and steps with t ≤ t_thresh perform the look-ahead explicitly.
  • Feedback-Aware Inference: Prior-step reuse supplies the previous clean-signal estimate when t > t_thresh and i ≠ 0, avoiding a new look-ahead pass.After refinement, the estimate is cached for reuse at the next step.

A.1 2D Toy Experiment · A.2 Image-to-Image Translation

The 2D toy experiment uses a thick, quadrant-conditioned Archimedean spiral to expose the fidelity–plausibility trade-off of guidance. Image-to-image translation extends FlowBender to Stable Diffusion 3.5 with zero-order and first-order feedback variants that refine ControlNet or DiT predictions.

  • A.1 2D Toy Experiment: The toy target is a finite-thickness Archimedean spiral, with isotropic Gaussian spread σ = 0.12 around its curve.Each point receives a quadrant label based on the spiral’s quarter-turn.
  • A.1 2D Toy Experiment: The toy model is a three-layer MLP with hidden dimension 64, SiLU activations, sinusoidal time embeddings, and learned quadrant embeddings.The scalar time input is lifted to a 16-dimensional sinusoidal embedding, while the quadrant embedding is added before the main network.
  • A.1 2D Toy Experiment: Training-free guidance exposes a fidelity–plausibility trade-off: small scales fail to enforce conditioning, whereas large scales push samples off the target manifold.The study therefore selects an intermediate guidance scale as a compromise.
  • A.2 Image-to-Image Translation: Image-to-image translation conditions Stable Diffusion 3.5 on an image-shaped signal y, such as edges, together with a text condition ctext.The experiments use the official SD3.5 checkpoint and train added ControlNets.
  • A.2 Image-to-Image Translation: The zero-order variant concatenates a normalized residual feedback signal with the image condition before ControlNet encodes the guidance signal.The residual uses H, the composition of the VAE decoder and pixel-space forward operator.
  • A.2 Image-to-Image Translation: The first-order variant doubles DiT input channels to inject standardized feedback, while an alternative uses ˜st = ∇ˆx1L without differentiating through vθ.The feedback is introduced during the unguided pass before computing the refined velocity.
  • A.2 Image-to-Image Translation: Training uses Unsplash-25K images resized to 1, 0242 resolution for 6, 250 optimizer steps, while only newly added parameters are trained and the SD3.5 base remains frozen.The setup uses AdamW with learning rate 10−5, weight decay 0.01, and 500 warmup steps.

A.3 3D Mesh Texturing

For 3D mesh texturing, FlowBender adapts TRELLIS-2 with operator-consistent feedback, gradient-conditioned LoRA fine-tuning, and closed-loop sampling. The setup uses differentiable PBR rendering and compares against matched fine-tuning, alignment-loss, and inference-time guidance baselines.

  • Data Preparation: FlowBender renders conditioning images with TRELLIS-2’s differentiable PBR renderer, ensuring the dataset observation and feedback use the same forward operator H.Lighting is fixed across assets to aid convergence.
  • Architecture and Feedback Signal: The adapted TRELLIS-2 texture flow model uses rank-128 LoRA adapters and expands its input from 2d to 3d channels to inject first-order gradient feedback.Here d = 32, and the added input block is zero-initialized so the model initially matches the unguided baseline.
  • Training: The model is fine-tuned on 7,500 Objaverse assets for 25,000 steps with batch size 16, AdamW learning rate 10^-4, weight decay 0.01, and EMA rate 0.9999.Training uses bfloat16 automatic mixed precision and adaptively clips gradients at maximum norm 1.0.
  • Inference: Inference uses 12 Euler steps with full two-pass execution at every step and no classifier-free guidance in the main comparisons.The guidance scale is 1.0, consistent with all other evaluated methods.
  • Baselines: Comparisons include matched Standard FT, FT + Lalign with λ = 5 × 10^-3 applied for t > 0.3, and FlowChef inference-time guidance using one SGD update per ODE step.Standard FT shares the LoRA architecture and training configuration, while FlowChef is applied to pretrained TRELLIS-2 without fine-tuning.

B Additional Results · B.1 Image-to-Image Translation

The extended analysis compares image-to-image translation methods across sampling budgets and guidance settings, finding that open-loop enhancements do not close fidelity gaps, whereas FlowBender improves fidelity while preserving plausibility. It also reports limitations of inference-time guidance for complex forward operators.

  • B.1 Image-to-Image Translation: Additional experiments provide qualitative comparisons for JPEG restoration, depth-to-RGB, and super-resolution, alongside tables covering sampling steps, classifier-free guidance, and hyperparameter choices.These results appear in Figures 8–10 and Tables 6–9.
  • B.1 Image-to-Image Translation: The extended analysis broadens the main-text evaluation beyond configurations marked with (∗) by comparing varying sampling budgets and guidance settings.This section is presented as a comprehensive comparison across these configurations.
  • B.1 Image-to-Image Translation: Doubling open-loop sampling steps yields negligible fidelity gains, while classifier-free guidance often degrades plausibility without significantly closing the fidelity gap.The analysis concludes that standard test-time enhancements cannot resolve open-loop failures.
  • B.1 Image-to-Image Translation: FlowBender’s classifier-free guidance often boosts fidelity, with w = 3.0 improving the zero-order variant’s PSNR by 6.0 dB versus 0.3 dB for the baseline.The comparison is reported in Table 6.
  • B.1 Image-to-Image Translation: Inference-time guidance exhibits a strict plausibility–fidelity trade-off across tasks, whereas FlowBender consistently achieves both stronger plausibility and fidelity.The trade-off depends on the chosen hyperparameters.
  • B.1 Image-to-Image Translation: FlowChef works for simple forward operators such as inpainting and super-resolution but does not generalize to complex forward operators such as neural networks.This limitation was evaluated through a dense sweep over key hyperparameters, including learning rate λ.

B.2 3D Mesh Texturing

In 3D mesh texturing, the failure patterns seen in image-to-image tasks persist: open-loop enhancements can hurt both fidelity and plausibility. The section reports extended quantitative, qualitative, and multi-view results for Objaverse and Toys4K.

  • 3D Mesh Texturing: Extended quantitative results are provided for 3D texturing on Objaverse and Toys4K, alongside qualitative comparisons and multi-view visualizations.These results appear in Tables 10–11 and Figures 12–13.
  • 3D Mesh Texturing: Open-loop enhancements such as doubling sampling steps or applying CFG are suboptimal or detrimental in 3D texturing.The passage states that these patterns persist from image-to-image tasks in the 3D domain.
  • 3D Mesh Texturing: 1.57 dB drop in M.PSNR occurs when CFG (w = 3.0) is applied to Standard FT on Toys4K.Plausibility also collapses, with FID increasing from 8.87 to 10.30.
Loading 2606.20404v1…