PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Xin Xie, Fan Zhang, Dong Gong

arXiv:2609.05968v1cs.CV

TL;DR

Post-training must be repeated for each checkpoint, and existing test-time alignment derives its step-wise supervision from the reward. AlignGraft uses a source alignment pair to steer a larger frozen model, improving host-base results across image and video models and sometimes surpassing the source aligned model.

  • Problem

    Post-training must be repeated for each checkpoint, while reward-based test-time alignment requires reward-derived step-wise signals that can exclude non-differentiable or deployment-restricted rewards.

    Source: §1 · pp. 1–2

  • Method

    AlignGraft uses a source base/aligned model pair’s implicit reward as an additive velocity difference to steer a larger frozen host under a shared noising kernel.

    Source: §1 · pp. 2–3, §3.3 · p. 5

  • Results

    Across tested backbones spanning scale, model-family siblings, and image and video generation, AlignGraft improves host-base metrics, recovers most source alignment gain, and sometimes exceeds the source aligned model.

    Source: §1 · pp. 2–3, §5 · p. 10

  • Takeaways & Limitations

    One RL run on a small source model yields supervision reusable across a model family at test time, while leaving host weights untouched.

    Source: §1 · pp. 2–3, §5 · p. 10

    The method assumes shared latent space, conditioning encoders, and noising kernel; the source-pair value proxy is valid only on overlapping model support.

    Source: §3.3 · pp. 5–6

Abstract

from arXiv · show

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.

1 Introduction

AlignGraft transfers supervision from a weak source alignment pair to a frozen host model, avoiding reward-derived test-time signals. Across image and video flow models, it improves host-base results and can exceed the source aligned model.

1 Introduction

  • Post-training must be repeated for each checkpoint, while reward-based test-time methods require reward-derived per-step signals and may lack a KL constraint.Reward gradients, candidate evaluation, or per-model value functions create deployment costs and can exclude non-differentiable, slow, human, or proprietary rewards.Source: §1 · pp. 1–2
  • A source alignment pair stores the training reward as a step-wise density-ratio signal, already KL-anchored and expressed in the sampler’s coordinates.The pair supplies this signal with two forward passes, without evaluating the reward or its gradient.Source: §1 · p. 2
  • AlignGraft transfers the pair’s implicit reward to a frozen host by adding its velocity difference, with no schedule and one scalar strength.Under the shared noising kernel, the transport is exact; the scalar can extrapolate beyond the source pair’s alignment strength.Source: §1 · pp. 2–3
  • Transfer strength varies predictably by reward: coarse-spatial signals work at moderate strength, while fine low-noise structure needs stronger guidance.Source: §1 · p. 2
  • The host weights remain untouched, letting users choose the reward, strength, and target sample at test time and reuse one alignment run across a model family.The paper contrasts this reuse with distillation, which amortizes over deployment volume rather than many rewards, rotating checkpoints, and personalization.Source: §1 · pp. 2–3
  • Across Stable Diffusion 3.5, FLUX, and Wan, AlignGraft improves host models on preference, compositional, and text-rendering rewards, sometimes exceeding the source aligned model.The transfer spans image and video flow models and generalizes across scale and model-family siblings.Source: §1 · pp. 2–3

2 Related Work

Prior work aligns generators through reward-based post-training or test-time reward signals, while weak-to-strong generation methods transfer other kinds of guidance. AlignGraft instead transfers an RL-learned reward across models.

2.1 Diffusion Model Alignment via Post-Training

  • Post-training methods optimize rewards with policy gradients, preference pairs, or sampling gradients [35]; the run remains bound to its checkpoint and must be repeated.Source: §2.1 · p. 3

2.2 Test-Time Alignment of Diffusion Models

  • Test-time methods use reward gradients, particle evaluation [39], per-model value functions [23], or other reward access, leaving some rewards unavailable at deployment.AlignGraft instead uses a trained model pair, where the reward is pre-converted and pre-anchored.Source: §2.2 · p. 3

2.3 Weak-to-Strong Learning and Guidance

  • Decoding pairs transfer logit offsets or rerank candidates, while generation methods guide a model with its own pair [18] or transfer adapters and styles [16]; none transfers an RL-learned reward to a different, untuned model.In flow models, the log-ratio tilt becomes a closed-form drift for training-free, search-free reward transfer across scale.Source: §2.3 · pp. 3–4

3 Method

AlignGraft transfers reward supervision from a source alignment pair to a frozen host by adding the pair’s velocity difference during sampling. The method derives its step-wise signal from model probabilities rather than querying the reward, under a shared noising kernel.

3.1 Preliminaries

  • Flow matching transports Gaussian noise toward clean data with a learned velocity field, while alignment post-training maximizes reward under a KL constraint and encodes the optimum as a base/aligned density ratio.The velocity determines the marginal score, and the reward is evaluated on the clean output.Source: §3.1 · p. 4

3.2 Motivation and Problem Definition: From Reward to Model-Form Supervision

  • Large-model post-training must be repeated after checkpoint changes, while reward-derived test-time signals require gradients, many candidates, or model-specific value training.Reward estimates from predicted clean outputs can also be unreliable at high noise levels.Source: §3.2 · pp. 4–5
  • A base/aligned model pair stores the training reward as a log-density ratio available at every noise level, supplying dense supervision through forward passes without a reward, gradient, or surrogate training.The source pair can be a smaller model than the host, provided the models share the latent space, conditioning encoders, and forward noising kernel.Source: §3.1 · p. 4, §3.2 · p. 5, §3.3 · p. 5

3.3 AlignGraft: Weak-to-Strong Guidance

  • At each sampling step, AlignGraft adds the source pair’s alignment delta to the frozen host’s velocity and takes the ordinary sampler step.No weights are trained or merged, and the reward is not queried.Source: §3.3 · p. 5, Figure 2 · p. 6
  • The alignment strength w is AlignGraft’s only hyperparameter; exact cancellation means it uses no noise-level schedule, window, or filter.Larger w tilts the host distribution further toward the reward.Source: §3.3 · pp. 5–6
  • Under a shared forward kernel, tilting the host marginal shifts its velocity by exactly w times the source-pair delta; the stochastic one-step transition is exact when σt>0.The source value function stands in for the host’s on overlapping support, and the per-noise-level tilted family uses a product-of-experts construction.Source: §3.3 · pp. 5–6

4 Experiments

Experiments evaluate AlignGraft across image and video models, rewards, and evaluation dimensions. Results show reward gains on frozen hosts across these settings, with guidance strength and motion preservation providing important qualifications.

4.1 Experimental Setting

  • Experiments pair SD 3.5-Medium with frozen SD 3.5-Large, FLUX.1-dev with FLUX.1-Krea, and Wan-Fun-1.3B with Wan-Fun-14B, evaluating image and video rewards.The evaluation uses GenEval, HPD, OCR, and MovieGenBench prompts, with preference, text-rendering, compositional, aesthetic, and video-quality measures.Source: §4.1 · p. 7, §4 · p. 7

4.2 Qualitative Evaluation

  • For each SD 3.5 source reward, guidance shifts samples toward its learned behavior while retaining host fidelity; gains reflect success rates across prompts, not guarantees for individual seeds.PickScore yields richer appeal, GenEval improves composition, and OCR produces more legible in-image text.Source: §4.2 · p. 8

4.3 Quantitative Analysis

  • GenEval overall rises from 0.68 to 0.90, near the source-aligned model’s 0.93; PickScore transfer improves every metric and exceeds the source-aligned model on Aesthetic, PickScore, and HPS v2.1.The GenEval gains are largest on position (0.28→0.91) and counting (0.67→0.84).Source: §4.3 · p. 8, Table 1 · p. 9, Table 3 · p. 9

4.4 Generalization across Backbones and Modalities

  • At w=1, transfer to equally sized FLUX.1-Krea improves all four metrics and recovers 70% of the source HPS gain, with gains across prompt styles.This setting tests transfer between independently post-trained siblings sharing a latent space.Source: §4.4 · p. 9, Table 2 · p. 9
  • At w=1, Wan HPS v2.1 rises from 0.2292 to 0.2839, while MPS improves by 0.015 (107%); motion quality stays unchanged.All held-out judges improve together in both reward-source settings.Source: Table 4 · p. 10, §4.4 · p. 10

4.5 Guidance Strength

  • Five reward metrics peak and then decline as guidance strength rises, while OCR plateaus; preference and compositional rewards peak around w=2–4.At peak strength, transfer matches or exceeds the source-aligned model on most metrics while retaining host capability.Source: Figure 6 · p. 9, §4.5 · p. 10

5 Conclusion

AlignGraft transports a source alignment pair’s implicit reward into a frozen host sampler, exactly under a shared noising kernel. Across tested image and video backbones it improves reported metrics and can exceed the source aligned model; a model-independent score–velocity relation makes outputs from different models addable.

5 Conclusion

  • Across tested backbones, AlignGraft improves the host base model on all reported metrics, recovers most of the source alignment gain, and sometimes exceeds the source aligned model.The transfer covers differing scales, sibling models, and image and video generation.Source: §5 · p. 10
  • The source alignment pair supplies implicit reward supervision to a frozen host sampler without test-time reward access, exactly under a shared noising kernel.The method adds two source-pair forward passes and leaves transport beyond a shared latent space to future work.Source: §5 · p. 10

A Score–Velocity Identity

  • Velocity and score are affine transforms with model-independent coefficients, making outputs from different models addable under the shared forward kernel.The appendix gives v(z,t) = −z + t s(z,t).Source: A Score–Velocity Identity · p. 14

B Sampler Coefficients

The appendix specifies the marginal-preserving SDE Euler step used in the experiments and notes that its coefficients depend only on the scheduler; setting η to zero gives a deterministic ODE sampler.

B Sampler Coefficients

C Proof of the Proposition

The proof shows that the tilted family’s velocity is the host velocity plus a scaled source-pair velocity difference. With shared Gaussian transition covariance, applying that velocity realizes the tilted one-step kernel exactly, without linearization.

C Proof of the Proposition

  • The tilted family’s velocity is the host velocity plus w times the source-pair velocity difference, with the time-dependent factors canceling exactly.Source: C Proof of the Proposition · p. 14
  • Under the shared covariance, the pair’s one-step log-ratio is affine, and exponentiating it shifts the transition mean without changing covariance.Source: C Proof of the Proposition · p. 14
  • Substituting the guided velocity into the unmodified sampler step realizes the reward-tilted one-step kernel exactly, with no linearization.The effective time weighting is inherited from the scheduler rather than designed.Source: C Proof of the Proposition · pp. 14–15

D The Pair Ratio Is the Source Value Function

The source alignment pair’s noisy-marginal ratio is the soft value function for its own process under the shared forward kernel. Thus the RL run amortizes value computation across noise levels.

D The Pair Ratio Is the Source Value Function

E Anchor Symmetry and Support

For w=1, the target has two equivalent factorizations; the source-aligned view grounds weak-to-strong transfer but depends on shared support.

  • At w=1, the target is either the host base tilted by the source pair’s implicit reward or the source aligned model tilted by the capability gap pl/ps.Source: E Anchor Symmetry and Support · p. 15
  • The source-aligned factorization makes the transfer weak-to-strong and assumes the host support contains the source-aligned model’s support.Source: E Anchor Symmetry and Support · p. 15
  • Outside shared support, the density ratio becomes unconstrained network extrapolation, marking the transfer boundary; on overlap, host-visited states remain in-distribution.Source: E Anchor Symmetry and Support · p. 15

F Implementation Details

AlignGraft adds the source pair’s velocity difference to the host velocity at each shared sampling state, using deterministic ODE sampling and backbone-specific settings.

  • At each step, the host and source pair are evaluated at the shared state, and the weighted alignment delta w ∆vs is added to the host velocity.Source: F Implementation Details · p. 15
  • The implementation uses a deterministic ODE sampler by setting the step variance coefficient σt to zero.Source: F Implementation Details · p. 15
  • Sampling settings vary by backbone: SD 3.5 uses 40 steps at 1024×1024, FLUX 28 steps at 1024×1024, and Wan 30 steps at 832×480 with 33 frames.Source: F Implementation Details · p. 15

G Additional Experimental Details and Results

Additional experiments show qualitative transfer across reward types, sibling backbones, and video, while guidance strength and responsible use bound interpretation.

G Additional Experimental Details and Results

H Social Impact

  • The authors flag objective-agnostic steering and source-reward bias as risks, recommending vetted rewards, provenance, watermarking, moderation, and anchored guidance.Source: H Social Impact · p. 16

I Limitations

The supplied evidence identifies a source-pair availability constraint and over-optimization at high guidance, alongside qualitative transfers across image and video settings.

  • The method requires the full source alignment pair because its signal is the aligned-minus-base velocity difference.Source: I Limitations · p. 17
  • The practical boundary is a post-RL checkpoint without its matching base; otherwise the base is usually public or the small source model is cheap to retrain.Source: I Limitations · p. 17
  • SD 3.5 examples show preference gains, legible prompt text, and corrected counting or spatial relations while retaining host fidelity.Source: Figure 8 · p. 18, Figure 10 · p. 20, Figure 12 · p. 22
  • Wan examples show improved per-frame appeal and prompt adherence under HPS v2.1, and improved motion coherence and prompt adherence under MPS.Source: Figure 16 · p. 26, Figure 18 · p. 28
  • At high guidance, reward alignment can over-optimize: reward alignment peaks broadly before large strengths degrade into saturated color, collapsed layout, or distorted text.Source: Figure 19 · p. 29
Loading 2609.05968v1…