Test-Time Weak-to-Strong Alignment: Transferring Implicit Rewards from Weak to Strong Flow Models

• URL: http://arxiv.org/abs/2609.05968v1
• Date: Sat, 05 Sep 2026 08:12:49 GMT
• Title: Test-Time Weak-to-Strong Alignment: Transferring Implicit Rewards from Weak to Strong Flow Models
• Abstract summary: Test-time alignment steers a frozen model during sampling, allowing task-specific and sample-specific guidance.<n>We propose changing the supervision source: let a pair of weak models, not a reward function, supply the supervision.<n>A source aligned model, kept together with its base as a source alignment pair, stores its training reward as an implicit, step-wise, KL-anchored signal.
• Score: 19.4335297564654
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: Aligning a text-to-image generation flow model with a reward makes it follow objectives that the training data alone does not provide. Alignment fine-tuning delivers this by reinforcement learning (RL) or preference optimization, but it must be repeated for every checkpoint and returns a model fixed at the reward and strength it was trained with. Test-time alignment instead steers a frozen model during sampling, allowing task-specific and sample-specific guidance. Existing methods obtain this only by drawing the per-step signal from the reward function itself, through its gradient, or through a separately trained value function. We propose changing the supervision source: let a pair of weak models, not a reward function, supply the supervision. A source aligned model, kept together with its base as a source alignment pair, stores its training reward as an implicit, step-wise, KL-anchored signal expressed in the sampler's own coordinates. We explore whether this model-form supervision can cross scale, and show that it does: our method, AlignGraft, aligns a larger, frozen, never-tuned model by adding the pair's velocity difference during sampling. The transport is exact under a shared noising kernel and needs neither the reward nor its gradient at test time. The method has no schedules, only a single scalar that controls the alignment strength and can extrapolate it beyond that of the source alignment pair. Across image and video flow models (Stable Diffusion 3.5, FLUX, and Wan), the transfer lifts the frozen large model on preference, compositional, and text-rendering rewards, can exceed the source aligned model itself, and preserves the large model's fidelity at a small constant sampling overhead. Extensive experiments show that one alignment run on a weak model produces supervision that the whole model family can reuse at test time.

Related papers

• FlowBender: Feedback-Aware Training for Self-Correcting Conditional Flows [18.746098727212544]
  FlowBender is a new diffusion and flow model for 3D reconstruction.<n>We show that FlowBender consistently outperforms standard supervised alignment training.<n>We propose several variants of FlowBender for differentiable operators.
  arXiv  Detail & Related papers  (2026-06-18T15:56:07Z)
• Bootstrap Your Generator: Unpaired Visual Editing with Flow Matching [79.46676932242569]
  We propose a framework for unpaired training of flow matching editing models.<n>ByG pairs instruction-following cues extracted from the frozen model with cycle-consistency for structure preservation.<n>We demonstrate state-of-the-art results on challenging data-scarce image and video editing scenarios.
  arXiv  Detail & Related papers  (2026-06-02T17:07:08Z)
• Aligning Few-Step Generative Models by Amortizing Sample-based Variational Inference [60.82875447631739]
  We introduce FAV, Few-step Generative Models Alignment via Sample-based Variational Inference.<n>FAV cast alignment as sampling from a reward-tilted distribution anchored to a reference distribution.<n>We evaluate FAV on two domains: robotics manipulation and image generator alignment.
  arXiv  Detail & Related papers  (2026-05-26T05:02:49Z)
• ParetoSlider: Diffusion Models Post-Training for Continuous Reward Control [14.748861060405973]
  Reinforcement Learning (RL) post-training has become the standard for aligning generative models with human preferences.<n>The prevailing practice of early scalarization'' collapses rewards into a fixed weighted sum when multiple criteria matter.<n>This commits the model to a single trade-off point at training time, providing no inference-time control over inherently conflicting goals.
  arXiv  Detail & Related papers  (2026-04-22T17:44:56Z)
• Designing Instance-Level Sampling Schedules via REINFORCE with James-Stein Shrinkage [46.743961849406844]
  We reschedule the sampling timeline of a frozen text-to-image sampler.<n>We introduce a novel reward baseline based on a principled James-Stein estimator.<n>We position our scheduling framework as an emerging model-agnostic post-training lever.
  arXiv  Detail & Related papers  (2025-11-27T07:33:05Z)
• Learning an Image Editing Model without Image Editing Pairs [83.03646586929638]
  Recent image editing models have achieved impressive results while following natural language editing instructions.<n>They rely on supervised fine-tuning with large datasets of input-target pairs.<n>Current workarounds use synthetic training pairs that leverage the zero-shot capabilities of existing models.<n>We present a new training paradigm that eliminates the need for paired data entirely.
  arXiv  Detail & Related papers  (2025-10-16T17:59:57Z)
• Temporal Score Rescaling for Temperature Sampling in Diffusion and Flow Models [30.36841165328262]
  We present a mechanism to steer the sampling diversity of flow matching models.<n>We show that rescaling these allows one to effectively control a local' sampling temperature.
  arXiv  Detail & Related papers  (2025-10-01T17:59:51Z)
• Learn to Guide Your Diffusion Model [84.82855046749657]
  We study a technique for improving quality of samples from conditional diffusion models.<n>We learn guidance weights $omega_c,(s,t)$, which are functions of the conditioning $c$, the time $t$ from which we denoise, and the time $s$ towards which we denoise.<n>We extend our framework to reward guided sampling, enabling the model to target distributions tilted by a reward function.
  arXiv  Detail & Related papers  (2025-10-01T12:21:48Z)
• Aligning Text-to-Image Diffusion Models with Reward Backpropagation [62.45086888512723]
  We propose AlignProp, a method that aligns diffusion models to downstream reward functions using end-to-end backpropagation of the reward gradient. We show AlignProp achieves higher rewards in fewer training steps than alternatives, while being conceptually simpler.
  arXiv  Detail & Related papers  (2023-10-05T17:59:18Z)
• Test-Time Adaptation with CLIP Reward for Zero-Shot Generalization in Vision-Language Models [76.410400238974]
  We propose TTA with feedback to rectify the model output and prevent the model from becoming blindly confident. A CLIP model is adopted as the reward model during TTA and provides feedback for the VLM. The proposed textitreinforcement learning with CLIP feedback(RLCF) framework is highly flexible and universal.
  arXiv  Detail & Related papers  (2023-05-29T11:03:59Z)
• Consistent Diffusion Models: Mitigating Sampling Drift by Learning to be Consistent [97.64313409741614]
  We propose to enforce a emphconsistency property which states that predictions of the model on its own generated data are consistent across time. We show that our novel training objective yields state-of-the-art results for conditional and unconditional generation in CIFAR-10 and baseline improvements in AFHQ and FFHQ.
  arXiv  Detail & Related papers  (2023-02-17T18:45:04Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.