U7.0 The Model Is Trained. Can We Still Change What We Want?

Builds on3 notes

Learning the data does not specify what we want this time

Earlier units learned distributions and transformed noise or source data into their samples.

Real requirements often arrive after training: a dog wearing a red hat, an image consistent with a blurry observation, a stable protein, or a robot trajectory avoiding a new obstacle. The pretrained model knows plausible data, but needs information about this particular request.

A navigation system already knows the roads. Asking for charging stops and less traffic should change route choices rather than require rebuilding the map.

The map corresponds to the pretrained model. Text, measurements, or a scoring rule provide the new requirement.

Can we guide sampling toward this task’s preferred region without retraining the generator?

These methods are called test-time guidance or test-time steering. We use steering for the broader problem and guidance for methods that add an intermediate direction.

When does the requirement enter?

Let $x_0$ denote clean data and $c$ a condition such as a class or text prompt.

Conditional training: the requirement is part of training

Training on paired examples $(x_0,c)$ can directly learn

p_\theta(x_0\mid c).

For continuous data this is a conditional density, rather than the probability of a single point. Given a red-hat-dog prompt, compatible images should receive higher density.

Both diffusion and Flow Matching retain their regression labels and also provide $c$ to the network.

Diffusion predicts noise from the noisy state, time, and condition:

\mathcal L_{\mathrm{diff,cond}}(\theta) =\mathbb E_{x_0,c,t,\epsilon} \left[ \left\|\epsilon-\epsilon_\theta(x_t,t,c)\right\|^2 \right].

The expectation averages training pairs, noise levels, and corruption noise. The target remains the sampled noise.

Flow Matching predicts the reference velocity while receiving the condition:

\mathcal L_{\mathrm{FM,cond}}(\theta) =\mathbb E_{t,(X_0,X_1,c)} \left[ \left\| v_\theta(I_t,t,c)-\partial_t I_t \right\|^2 \right].

Here $I_t$ connects a source endpoint to a data endpoint, and $\partial_t I_t$ is its derivative. The expectation includes endpoint pairing, condition, and sampled time.

In both cases, conditions enter the original training and affect the learned parameters. This is conditional training.

Fine-tuning: update parameters after pretraining

Suppose we later collect preferred generated examples. For an unconditional model, the preferred dataset contains clean samples; for a conditional model, it contains sample-condition pairs. These examples identify outcomes we want the model to generate more often.

A simple unconditional fine-tuning objective reuses denoising regression on preferred data:

\mathcal L_{\mathrm{align}}(\theta;\mathcal D_{\mathrm{pref}}) =\mathbb E_{x_0\sim\mathcal D_{\mathrm{pref}},t,\epsilon} \left[ \left\|\epsilon-\epsilon_\theta(x_t,t)\right\|^2 \right].

The loss averages preferred samples, times, and corruption noise. Conditional models additionally receive $c$ . Flow Matching can similarly fine-tune velocity regression. The data distribution driving parameter updates has changed.

Apply gradient descent:

\theta \leftarrow \theta-\eta\nabla_\theta \mathcal L_{\mathrm{align}} \bigl(\theta;\mathcal D_{\mathrm{pref}}\bigr).

The parameter gradient points uphill in loss, so its negative lowers the objective. The learning rate $\eta$ determines the update size.

Alignment methods may instead use comparisons or reward weighting. Their shared feature is writing preferences into model parameters. Updating pretrained parameters places the procedure outside this unit’s definition of test-time steering.

Test-time steering: change this sampling computation

Keep pretrained generator parameters fixed. Extra evaluators or scores may change intermediate computations or select completed candidates, but do not overwrite the generator’s parameters.

For a score sampler, temporarily use

\widetilde s_t(x_t) =s_\theta(x_t,t)+g_t(x_t), \qquad \theta\ \text{fixed}.

The base score is retained, while $g_t$ adds a task-dependent direction from a condition, observation, or reward. The sampler follows this combined field for the current request. The same unchanged model can use a different task signal next time.

Compare the three procedures:

Procedure	When does the requirement enter?	Update generator parameters?	Generation
Conditional training	Original paired training data	Yes, during training	Supply a learned condition
Fine-tuning	New preference data after pretraining	Yes	Use the updated model
Test-time steering	Current sampling request	No	Modify sampling or selection using a task signal

New requirements need a signal

The interpolant framework unifies transport by diffusion and flow. Additional requirements still need additional information. A robot trajectory generator cannot infer the location of a newly placed obstacle without receiving that location or a collision-related score.

We will study how to turn such signals into modifications of generation or selection.

A fixed pretrained model changes this generated route using information about a new obstacle.

Figure: task information changes this sampling run while pretrained parameters remain fixed.

DiscussionQ1

Why not generate many candidates and simply return the highest-scoring one? Do we need to change intermediate sampling at all?

Show the reasoning

This valid baseline is Best-of- $N$ . It needs neither reward gradients nor model internals. Selection changes the returned distribution even though each candidate was generated by the original procedure. We revisit it in U7.3.

The difficulty is a rare requirement. If one independent candidate succeeds with probability $p$ , the chance of at least one success among $N$ is

1-(1-p)^N.

For small $p$ , substantially increasing success can require $N$ on the order of $1/p$ . Near-successes are discarded without using their scores to improve them along the way.

Selection is useful, but can be expensive. Can intermediate feedback spend the same computation more effectively?

Check your understanding

References

Dhariwal, P., Nichol, A. Diffusion Models Beat GANs on Image Synthesis. NeurIPS 2021.
Ho, J., Salimans, T. Classifier-Free Diffusion Guidance. 2022.
Singhal, R., et al. A General Framework for Inference-time Scaling and Steering of Diffusion Models. ICML 2025.