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

Xin Xie; Fan Zhang; Dong Gong

cs.CVSep 5, 2026

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

Authors: Xin Xie, Fan Zhang, Dong Gong

Organizations: University of New South Wales (UNSW Sydney)

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.

Figures & tables

Figure 1: Sample images generated by AlignGraft on the SD 3.5-Large backbone. The generations align with the text prompt and human preference while keeping strong visual aesthetics.

Figure 3: Qualitative transfer of three reward sources on SD 3.5, in blocks (a) PickScore, (b) GenEval, and (c) OCR. Columns, left to right: source base (SD 3.5-Medium), source aligned (Flow-GRPO), host base (SD 3.5-Large), and guided host (SD 3.5-Large + AlignGraft). Each reward transfers its own behavior while the host base model keeps its native fidelity.

Figure 4: Qualitative comparison of text-to-video transfer on Wan-Fun, with (a) HPS v2.1 and (b) MPS reward sources. Rows, top to bottom: source base (Wan-Fun-1.3B), source aligned (Wan-Fun-1.3B RL), host base (Wan-Fun-14B), and guided host (Wan-Fun-14B + AlignGraft); each row shows four frames of a clip. The transfer improves per-frame appeal and prompt adherence while preserving motion (Table 4 ).

Figure 5: Qualitative comparison of cross-sibling transfer within the FLUX family.

Table 6

Method	Overall	Single obj.	Two obj.	Counting	Colors	Position	Color attr.
SD 3.5-M (source base)	0.63	0.98	0.82	0.55	0.81	0.28	0.52
Flow-GRPO (source aligned)	0.93	0.99	0.98	0.88	0.91	0.90	0.82
SD 3.5-L (host base)	0.68	0.97	0.85	0.67	0.84	0.28	0.61
SD 3.5-L + AlignGraft	0.90	0.99	0.97	0.84	0.87	0.91	0.79

Table 3: GenEval reward source on SD 3.5 (

w{=}2

). Every sub-task improves, most strongly position and counting.

Reward	Method	Aes	MPS	IR	HPS v2.1	HPSv3	VQ	MQ	TA
HPS v2.1	Wan-Fun-1.3B (source base)	4.942	0.088	$-$ 0.466	0.2235	2.887	$-$ 0.829	$-$ 0.500	$-$ 0.962
	Fun-Reward-HPS (source aligned)	5.508	0.104	0.394	0.2778	7.865	$-$ 0.403	$-$ 0.430	$-$ 0.430
	Wan-Fun-14B (host base)	5.029	0.095	$-$ 0.168	0.2292	4.470	$-$ 0.761	$-$ 0.339	$-$ 0.602
	Wan-Fun-14B + AlignGraft	5.533	0.111	0.549	0.2839	8.780	$-$ 0.314	$-$ 0.318	$-$ 0.147
MPS	Wan-Fun-1.3B (source base)	4.932	0.087	$-$ 0.495	0.2171	1.743	$-$ 0.818	$-$ 0.464	$-$ 0.894
	Fun-Reward-MPS (source aligned)	5.384	0.101	$-$ 0.020	0.2510	5.615	$-$ 0.468	$-$ 0.399	$-$ 0.580

Table 4: Text-to-video transfer on Wan-Fun with the HPS v2.1 and MPS reward sources (

w{=}1

). The guided host exceeds both the host base model and the source aligned model on different metrics, with motion quality preserved. The HPS v2.1 and MPS rows are sampled with different random seeds.

Figure 6: Guidance-strength search on SD 3.5: six reward metrics against

w\in[0.5,8]

, and

w\in[0.5,5]

for GenEval, with source base, source aligned, and host base as dashed references. Each metric is scored under its own source alignment pair (GenEval and OCR under theirs; Aes, Pick, IR, HPS v2.1 under PickScore). Five metrics follow an inverted U, while the OCR reward rises to a plateau; every metric peaks above the source aligned model except GenEval, which peaks below it.

Figure 7: AlignGraft under different guidance strength $w$ on SD 3.5 backbones.

Backbone (source $\to$ host)	Mod.	Resolution	Steps	CFG
SD 3.5 (M $\to$ L)	T2I	$1024{\times}1024$	40	4.5
FLUX (dev $\to$ Krea)	T2I	$1024{\times}1024$	28	3.5
Wan-Fun (1.3B $\to$ 14B)	T2V	$832{\times}480$ , 33 f	30	6.0

Table 5: Per-backbone sampling configuration. Steps is the number of sampler discretization steps

N

(Appendix B); CFG is the classifier-free guidance scale, shared by the source alignment pair and the host within each family; “33 f” denotes

33

frames per clip.

Figure 8: Additional PickScore transfers on SD 3.5 ( $w{=}3$ ). Columns, left to right: source base (SD 3.5-Medium), source aligned (Flow-GRPO, PickScore), host base (SD 3.5-Large), and guided host (SD 3.5-Large + AlignGraft); each row is one prompt. The transferred samples inherit the preference reward (richer lighting, texture, and subject appeal) while retaining the host base model’s native fidelity and layout.

Figure 9: More PickScore transfers on SD 3.5 ( $w{=}3$ ), continued from Figure 8 ; same column layout (source base, source aligned, host base, and guided host).

Figure 10: Additional OCR transfers on SD 3.5 ( $w{=}1$ ). Columns, left to right: source base (SD 3.5-Medium), source aligned (Flow-GRPO, OCR), host base (SD 3.5-Large), and guided host (SD 3.5-Large + AlignGraft); each row is one prompt with its intended in-image text highlighted. The host base model renders malformed or mirrored glyphs, whereas the transfer produces legible in-image text matching the prompt while keeping the host base model’s fidelity.

Figure 11: More OCR transfers on SD 3.5 ( $w{=}1$ ), continued from Figure 10 ; same column layout (source base, source aligned, host base, and guided host).

Figure 12: Additional GenEval transfers on SD 3.5 ( $w{=}1$ ). Columns, left to right: source base (SD 3.5-Medium), source aligned (Flow-GRPO, GenEval), host base (SD 3.5-Large), and guided host (SD 3.5-Large + AlignGraft); each row is one prompt testing counting, spatial relation, or color binding. The transfer corrects compositional errors in object count and left/right/above/below arrangement that the host base model otherwise misses.

Figure 13: More GenEval transfers on SD 3.5 ( $w{=}1$ ), continued from Figure 12 ; same column layout (source base, source aligned, host base, and guided host).

Figure 14: Additional cross-sibling transfers within the FLUX family ( $w{=}1$ ). Columns, left to right: source base (FLUX.1-dev), source aligned (DanceGRPO, HPS v2.1), host base (FLUX.1-Krea), and guided host (FLUX.1-Krea + AlignGraft); each row is one prompt.

Figure 15: Additional cross-sibling transfers within the FLUX family ( $w{=}1$ ). Columns, left to right: source base (FLUX.1-dev), source aligned (DanceGRPO, HPS v2.1), host base (FLUX.1-Krea), and guided host (FLUX.1-Krea + AlignGraft); each row is one prompt.

Figure 16: Additional text-to-video transfers on Wan-Fun with the HPS v2.1 source ( $w{=}1$ ). Rows, top to bottom: source base (Wan-Fun-1.3B), source aligned (Wan-Fun-1.3B, reward backprop, HPS v2.1), host base (Wan-Fun-14B), and guided host (Wan-Fun-14B + AlignGraft); each row shows six frames of one clip, and blocks are separated by prompt. The transfer improves per-frame appeal and prompt adherence while preserving the host base model’s motion.

Figure 17: More HPS v2.1 text-to-video transfers on Wan-Fun ( $w{=}1$ ), continued from Figure 16 ; same row layout (source base, source aligned, host base, and guided host), six frames per clip.

Figure 18: Additional text-to-video transfers on Wan-Fun with the MPS source ( $w{=}1$ ). Rows, top to bottom: source base (Wan-Fun-1.3B), source aligned (Wan-Fun-1.3B, reward backprop, MPS), host base (Wan-Fun-14B), and guided host (Wan-Fun-14B + AlignGraft); each row shows six frames of one clip, and blocks are separated by prompt. The transfer improves motion coherence and prompt adherence.

Figure 19: Effect of the guidance strength $w$ on SD 3.5. The leftmost column is the frozen SD 3.5-Large model ( $w{=}0$ ); the remaining columns apply AlignGraft at $w\in\{0.5,1,2,3,4,5\}$ . Blocks: (a) PickScore, (b) GenEval, and (c) OCR reward sources. Reward alignment strengthens with $w$ up to a broad peak and then degrades into over-optimization at large $w$ (saturated color, collapsed layout, distorted text), the qualitative counterpart of the inverted-U in Figure 6 .

Diffusion Model Guidance

Diffusion model guidance steers denoising toward desired conditions, rewards, constraints, or away from unwanted concepts.

See all 370 papers

May 8, 2026

Flow-OPD: On-Policy Distillation for Flow Matching Models

Existing Flow Matching (FM) text-to-image models suffer from two critical bottlenecks under multi-task alignment: the reward sparsity induced by scalar-valued rewards, and the gradient interference arising from jointly optimizing heterogeneous objectives, which together give rise to a 'seesaw effect' of competing metrics and pervasive reward hacking. Inspired by the success of On-Policy Distillation (OPD) in the large language model community, we propose Flow-OPD, the first unified post-training framework that integrates on-policy distillation into Flow Matching models. Flow-OPD adopts a two-stage alignment strategy: it first cultivates domain-specialized teacher models via single-reward GRPO fine-tuning, allowing each expert to reach its performance ceiling in isolation; it then establishes a robust initial policy through a Flow-based Cold-Start scheme and seamlessly consolidates heterogeneous expertise into a single student via a three-step orchestration of on-policy sampling, task-routing labeling, and dense trajectory-level supervision. We further introduce Manifold Anchor Regularization (MAR), which leverages a task-agnostic teacher to provide full-data supervision that anchors generation to a high-quality manifold, effectively mitigating the aesthetic degradation commonly observed in purely RL-driven alignment. Built upon Stable Diffusion 3.5 Medium, Flow-OPD raises the GenEval score from 63 to 92 and the OCR accuracy from 59 to 94, yielding an overall improvement of roughly 10 points over vanilla GRPO, while preserving image fidelity and human-preference alignment and exhibiting an emergent 'teacher-surpassing' effect. These results establish Flow-OPD as a scalable alignment paradigm for building generalist text-to-image models. The codes and weights will be released in: https://github.com/CostaliyA/Flow-OPD .

Zhen Fang, Wenxuan Huang, Yu Zeng +8

University of Science and Technology of China · University of California, Los Angeles · 3The Chinese University of Hong Kong +1

Diffusion Model Alignment Flow Matching RL for Image Generation Flow-Based Generative Modeling On-Policy Distillation T2I Diffusion Models

Ask PDF

Apr 16, 2026

LeapAlign: Post-Training Flow Matching Models at Any Generation Step by Building Two-Step Trajectories

This paper focuses on the alignment of flow matching models with human preferences. A promising way is fine-tuning by directly backpropagating reward gradients through the differentiable generation process of flow matching. However, backpropagating through long trajectories results in prohibitive memory costs and gradient explosion. Therefore, direct-gradient methods struggle to update early generation steps, which are crucial for determining the global structure of the final image. To address this issue, we introduce LeapAlign, a fine-tuning method that reduces computational cost and enables direct gradient propagation from reward to early generation steps. Specifically, we shorten the long trajectory into only two steps by designing two consecutive leaps, each skipping multiple ODE sampling steps and predicting future latents in a single step. By randomizing the start and end timesteps of the leaps, LeapAlign leads to efficient and stable model updates at any generation step. To better use such shortened trajectories, we assign higher training weights to those that are more consistent with the long generation path. To further enhance gradient stability, we reduce the weights of gradient terms with large magnitude, instead of completely removing them as done in previous works. When fine-tuning the Flux model, LeapAlign consistently outperforms state-of-the-art GRPO-based and direct-gradient methods across various metrics, achieving superior image quality and image-text alignment.

Zhanhao Liang, Tao Yang, Jie Wu +2

1The Australian National University · 2ByteDance Seed

Diffusion Model Alignment Flow Matching Image Generation Diffusion Model Fine-Tuning Flow-Based Generative Modeling

Ask PDF

Apr 29, 2026

How to Guide Your Flow: Few-Step Alignment via Flow Map Reward Guidance

In generative modeling, we often wish to produce samples that maximize a user-specified reward such as aesthetic quality or alignment with human preferences, a problem known as guidance. Despite their widespread use, existing guidance methods either require expensive multi-particle, many-step schemes or rely on poorly understood approximations. We reformulate guidance as a deterministic optimal control problem, yielding a hierarchy of algorithms that subsumes existing approaches at the coarsest level. We show that the flow map, an object of significant recent interest for its role in fast inference, arises naturally in the optimal solution. Based on this observation, we propose Flow Map Reward Guidance (FMRG): a training-free, single-trajectory framework that uses the flow map to both integrate and guide the flow. At text-to-image scale, FMRG matches or surpasses baselines across inverse problems and reward-guided generation with as few as 3 NFEs, giving at least an order-of-magnitude speedup in comparison to prior state of the art. Code is available at https://github.com/jrrhuang/fmrg.

Jerry Y. Huang, Justin Lin, Sheel Shah +2

Carnegie Mellon University

Generative Modeling Flow-Based Generative Modeling Conditional Generative Modeling

Ask PDF

Abstract

Diffusion Model Guidance

Explore similar work

Flow-OPD: On-Policy Distillation for Flow Matching Models

LeapAlign: Post-Training Flow Matching Models at Any Generation Step by Building Two-Step Trajectories

How to Guide Your Flow: Few-Step Alignment via Flow Map Reward Guidance