一位教授和他们的学生在使用PyTorch的MPS后端（针对Apple Silicon）时，与CUDA和CPU相比，在使用Stable Audio Open Small模型时遇到了结果不一致的问题。他们正在寻求Hacker News社区的帮助来诊断问题，并链接了一个相关的pull request，其中详细说明了MPS的不准确性。
该项目最初旨在为深度学习与AI伦理课程从头开始构建一个文本条件潜在流匹配模型，但由于资源限制，转而使用预训练模型进行指导调整——学生主要使用没有独立GPU的Mac电脑，并且云解决方案存在后勤挑战。
一位评论者提出了一个更广泛的问题，关于从非语言数据（如随机数）训练的模型中生成有用输出的意义，质疑这是否代表了超越基于语言学习的“智能”AI的新路径。 这篇文章既是对技术援助的请求，也是对AI模型中智能本质的讨论提示。
相关文章
原文
This tutorial demonstrates how to add inference-time controls to pretrained flow-based generative models to make them perform tasks they weren’t trained to do. We take an unconditional flow model trained on MNIST digits and apply two types of guidance: classifier guidance to generate specific digits, and inpainting to fill in missing pixels. Both approaches work by adding velocity corrections during the sampling process to steer the model toward desired outcomes. Since modern generative models operate in compressed latent spaces, we examine guidance methods that work directly in latent space as well as those that decode to pixel space. We also explore PnP-Flow, which satisfies constraints by iteratively projecting samples backward and forward in time rather than correcting flow velocities. The approaches demonstrated here work with other flow models and control tasks, so you can guide flows where you want them to go.
“When we put bits into the mouths of horses to make them obey us, we can turn the whole animal. Or take ships as an example. Although they are so large and are driven by strong winds, they are steered by a very small rudder wherever the pilot wants to go.” – James 3:3-4 (NIV)
In this tutorial, we’ll explore inference-time “plugin” methods for flow matching and rectified flow generative models like FLUX or Stable Audio Open Small. Unlike classifier-free guidance (CFG) [1], which requires training the model with your desired conditioning signal, these plugin guidance methods let you add controls at inference time—even for conditions the model never saw during training.
This tutorial assumes familiarity with flow-based generative models, by which we mean “flow matching” [2] and/or “rectified flows” [3]. See the blog post “Flow With What You Know”[4] for an overview, and/or my IJCNN 2025 tutorial [5] for further detail. The key insight is that flow models generate samples through iterative integration, and at each step we can add small velocity corrections to steer toward specific goals. This works for various objectives: generating specific classes, filling in missing regions, or satisfying other desired constraints.
Our discussion will bring us up to date on guidance methods for latent-space rectified flow models. While there’s an extensive literature on guidance for diffusion models [7] – see Sander Dieleman’s excellent blog post [8] for an overview — flow matching allows us to cast these in a more accessible and intuitive way. There’s some recent work unifying guidance for diffusion and flows [9], but in this tutorial we’ll focus on a simplified treatment for flows only.
The paradigm of latent generative models is covered in another superb Dieleman post [10], and combining latent-space models with flow-based guidance gives us powerful, flexible tools for adding flexible controls to efficient generation.
Let’s review the picture for flow-based generative modeling in latent space…
The Latent Flow-Matching Setup
The following diagrams illustrate the three key concepts:
a) A VAE compresses pixel-space images into compact latent representations. “E” is for encoder and “D” is for decoder:
b) The flow model operates in this latent space, transforming noise (“Source”, t=0) into structured data (“Target”, t=1) through iterative integration. The decoder then converts the final latents back to pixels.
c) While general flows can follow curved trajectories, some of our methods will focus on flows with nearly straight trajectories which allows for estimating endpoints without many integration steps:
These (nearly) straight trajectories can be obtained by “ReFlow” distillation of another model (covered in [4]) or by insisting during training that the models yield paths agreeing with Optimal Transport such as the “minibatch OT” method of Tong et al [11]. Even if the model’s trajectories aren’t super-straight, we’ll see that the guidance methods we use can be applied fairly generally anyway.
Intuitively, guidance amounts to “steering” during the integration of the flow model in order to end up at a desired end point. The following video provides a useful metaphor:
Ok, the analogy’s not quite right: you can’t just steer, you are going to have to paddle a little bit. In other words, you’re going to have to provide a bit of a extra velocity to correct where the “current” flow is taking you.
In flow matching, we go from a source data (distribution) at time \(t=0\) to target data at \(t=1\). Since this tutorial applies to latent space, we’ll use the letter \(z\) for position, such as \(z_t\) being the position at time \(t\).
When you’re “looking ahead” to estimate where you’ll end up, you project linearly along the current velocity \(\vec{v_t}\) for a duration of the remaining time. Let’s call this estimate \(\widehat{z_1}\), your projected endpoint :
…but perhaps that’s not where you want to go. Where you want to go is a distance \(\Delta \widehat{z_1}\) from \(\widehat{z_1}\), and to get there you’ll have to make a “course correction” \(\Delta \hat{v}\), as shown in the following diagram:
By similar triangles, \(\Delta \widehat{z_1} = (1-t)\Delta \vec{v}\), which means the course correction you want is
Since you’re going to see more math once you try to read the scholarly literature on these topics, let’s go a bit further into the math…
So \(\Delta \widehat{z_1}\) is a measure of the deviation from the desired endpoint. Now, in practical application we won’t actually use the “distance” \(\Delta \widehat{z_1}\), but we’ll use something that functions like a distance, such as a K-L divergence or Mean Squared Error (MSE).
When doing inference, this deviation serves as a “loss” – something we minimize via gradient descent, except we’ll vary the flow positions \(z\) instead of the model weights. More specifically, we’ll consider the “likelihood” \(p( \widehat{z_1} | y )\) of getting a \(z_1\) that matches a given control \(y\), and we’ll seek to maximize that likelihood, or equivalently to minimize the negative log-likelihood.
The expression \(-\nabla_{\widehat{z_1}} \log p( \widehat{z_1} | y )\) essentially answers the question, “In which direction should I adjust \(\widehat{z_1}\) so as to make \(p( \widehat{z_1} | y )\) more likely?” This gives us a direction and a magnitude, which we then multiply by a learning rate “guidance strength” \(\eta\) to turn it into a step size.
Applying this gradient-based approach, our expression for \(\Delta v\) will involve replacing \(\Delta \widehat{z_1}\) in (2) with \(- \eta \nabla_{\widehat{z_1}} \log p( \widehat{z_1} | y\):
where we used the fact that \(\nabla_{\widehat{z_1}} = \nabla_{z_t}\) (since \(\widehat{z_1} \propto z_t\)). The factor of \(1/(1-t)\) means small corrections suffice early on, but later times require larger adjustments—though other time scalings are possible, as we’ll see.
Now let’s apply this to a concrete example.
If we want our model to generate a member of a particular class, we can use an external classifier to examine the generated samples. The constraint to minimize will be the difference between the desired class and the argmax of the classifier output (or some similar relationship that enforces the class compliance).
Now let’s setup the (pretrained) classifier we’ll use for the guidance:
Let’s make a plot showing the classifier’s output probabilities (aka likelihoods) across all classes, for all 10 samples. The samples will be the rows, and the class-likelihoods outputs from the classifier will be the columns, where brightness is correlated with likelihood.
Code
def show_probs(probs, x=None):"""show probs as colormap intensities via imshow. have each row be a sample and each column be a class probability""" ncols =1if x isNoneelse2 fig, axs = plt.subplots(1, ncols, figsize=(8, 4))if ncols ==1: axs = [axs]if x isnotNone: # show a little version of the x image for each row axs[0].imshow(make_grid(x.unsqueeze(1).cpu(), nrow=1, padding=2, normalize=False).permute(1, 2, 0).cpu(), cmap='gray') axs[0].axis('off')# show probabilities as an intensity map im = axs[ncols-1].imshow(probs.cpu(), cmap='gray') axs[ncols-1].set_xlabel("Class") axs[ncols-1].set_ylabel("Sample #") plt.colorbar(im, ax=axs[ncols-1]) plt.tight_layout()show_probs(probs, x=x1)
From the “random” distribution of generated digits, we see that this is an unconditional generative model: there’s nothing determining the classes of the outputs – until we add guidance, below! ;-) In a short while, we’ll reproduce that diagram, but we’ll use guidance to get one class per sample, in order, along the diagonal.
To do that, we’re going to have to “break open” the generate_samples routine and even the integrate_path routine to allow us to add a correction to the velocity \(v_t\) generated by the flow model at time \(t\). That correction \(\Delta v\) will be based on the classifier’s output using the projected estimate\(\widehat{x_1}\) of the final data, which we’ll obtain via linear extrapolation.
In our latent space model, we flow with latents \(z\) which must be decoded using the VAE’s decoder \(D\) :
The correction \(\Delta v\) will generated from a constraint which in this case is just like regular “classifier loss” function in a supervised learning problem. The desired class label is the “target” and the classifier output of the projected estimate is the “prediction”.
Our code will follow this general layout:
loss_fn = torch.nn.CrossEntropyLoss()v_t = flow_model(z_t, t)z1_hat = z_t + (1-t)*v_t # projected destinationx1_hat = sub.vae.decoder(z1_hat) # decode it to pixel spaceprobs = classify(classifier, x1_hat) # classifer operates in pixel spaceloss = loss_fn(probs, target) # "supervised learning"delta_v = magic_function(loss,...???) # <--- here's the part we need to work outv_t = v_t + delta_v * guidance_strength # we can set the strength of the correction
So what magic function will turn that loss into a velocity? Well, we can turn to good old Gradient Descent! Except instead of taking the gradient with respect to the model weights like we used to, we’re going to take the gradient with respect to the flow coordinates \(z\) in the latent space, thereby generating a vector in the latent space.
The insight is that PyTorch lets us compute the gradient with respect to anything. We just need to tell it what we want. And we need to be careful to make sure that the VAE and flow models stay frozen, so the only thing that’s allowed to change are the latents \(z\).
The cleanest way to pull this off, code-wise, is to create a function called compute_v() which for starters will just call the flow model, but then we’ll add to it with guidance info:
# wherever we used to just call flow_model(), we'll now call compute_v() instead@torch.no_grad()def compute_v(flow_model, z, t, guidance_dict=None, **kwargs): v_t = flow_model(z, t)if guidance_dict isnotNone: v_t += compute_dv(v_t, z, t, guidance_dict, **kwargs)return v_t@torch.enable_grad() # <-- later, this will be a key for getting guidancedef compute_dv(v_t, z, t, g:dict, **kwargs):"placeholder for now, will add guidance math later"return torch.zeros_like(v_t).detach() # no correction yet; no gradients returned
We’ll to use our usual “boilerplate” flow integration code, except we’ll add “**kwargs” everywhere so we can pass controls “all the way in” to the compute_dv() guidance routine, and pair flow_model() as an arg to compute_v() via functools.partial.
Code
# flow integration routines, slightly modified from "flow with what you know" blogfrom functools import partial # use partial to package flow_model with compute_v@torch.no_grad()def rk4_step(f, y, t, dt, **kwargs): # regular rk4, + kwargs passthrough# f: callable (y, t) -> dy/dt k1 = f(y, t, **kwargs) k2 = f(y +0.5* dt * k1, t +0.5* dt, **kwargs) k3 = f(y +0.5* dt * k2, t +0.5* dt, **kwargs) k4 = f(y + dt * k3, t + dt, **kwargs)return y + (dt /6) * (k1 +2* k2 +2* k3 + k4)@torch.no_grad()def warp_time(t, dt=None, s=.5):"""Parametric Time Warping: s = slope in the middle. s=1 is linear time, s < 1 goes slower near the middle, s>1 goes slower near the ends s = 1.5 gets very close to the "cosine schedule", i.e. (1-cos(pi*t))/2, i.e. sin^2(pi/2*x)"""if s <0or s >1.5: raiseValueError(f"s={s} is out of bounds.") tw =4* (1- s) * t **3+6* (s -1) * t **2+ (3-2* s) * tif dt: # warped time-step requested; use derivativereturn tw, dt *12* (1- s) * t **2+12* (s -1) * t + (3-2* s)return tw@torch.no_grad()def integrate_path(model, initial_points, step_fn=rk4_step, n_steps=100, warp_fn=None, latent_2d=False, prog_bar=True, t0=0, **kwargs): p =next(model.parameters()) device, model_dtype = p.device, p.dtype current_points = initial_points.to(device=device, dtype=model_dtype).clone() model.eval() ts = torch.linspace(t0, 1, n_steps, device=device, dtype=model_dtype)if warp_fn: ts = warp_fn(ts)if latent_2d: t_batch = torch.empty((current_points.shape[0], 1), device=device, dtype=model_dtype) vel_model = partial(compute_v, model) # here's the secret sauce iterator =range(len(ts) -1)if prog_bar: iterator = tqdm(iterator, desc="Integrating Path")for i in iterator: t, dt = ts[i], ts[i +1] - ts[i]if latent_2d: t = t_batch.fill_(t.item()) current_points = step_fn(vel_model, current_points, t, dt, **kwargs)return current_pointsdef generate_samples(sub, n_samples: int, n_steps=15, z0=None, t0=0, **kwargs) -> torch.Tensor: z0 = torch.randn([n_samples, sub.latent_dim]).to(sub.device) if z0 isNoneelse z0 z1 = integrate_path(sub.flow_model, z0, n_steps=n_steps, step_fn=rk4_step, t0=t0, **kwargs) gen_xhat = F.sigmoid(sub.decode(z1).view(-1, 28, 28))return gen_xhat
Now that we know that works, let’s “supe up” compute_dv() to include the guidance correction. We’ll use the torch.autograd.grad() function to compute the gradient of the loss. First we have the guidance_dict that we’ll use to pass through our intentions through the various layers of routines to get to compute_dv():
guidance_dict =\ {'classifier': classifier, # the classifier model to use'decode': sub.decode, # how to decode to pixel space for classifier'loss_fn': torch.nn.CrossEntropyLoss(reduction='none'), # don't sum over batch dim'target': torch.arange(10).to(device), # desired class outcomes'strength': 5.0, # "guidance strength", you may vary this't_min': 0.01, 't_max': 0.99, # t range to apply guidance, may vary these }
Next we have the fully-equipped compute_dv(). This code is overly-commented to make it easy to follow each step. (We replaced guidance_dict with g locally for brevity.) No other changes to any preceding code are necessary. We’ll be ready to do guided inference after this definition!
@torch.enable_grad() # <-- Needed to compute gradients if calling code has @torch.no_grad()def compute_dv(v_t, z, t, g:dict, eps=1e-6, debug=False):"Compute the guidance correction to the flow velocity"if t < g['t_min'] or t > g['t_max']: return torch.zeros_like(v_t).detach() z.requires_grad_(True) # need to enable gradient tracking for z z1_hat = z + (1- t) * v_t # linear projection to estimated endpoint# Decoding to pixel space (if decoder provided) x1_hat = z1_hat if g['decode'] isNoneelse F.sigmoid(g['decode'](z1_hat)).view(-1, 28, 28) logits, probs = classify(g['classifier'], x1_hat) # run classifier loss = g['loss_fn'](logits, g['target'][:len(logits)]) # loss <-> "negative log likelihood"# Compute grad wrt z. "grad_outputs=": don't sum over over batch, keep unique to each datum grad_z = torch.autograd.grad(loss, z, grad_outputs=torch.ones_like(loss), retain_graph=False)[0] dv =-grad_z / (1- t + eps) # - minimizes, (1-t) makes it velocity, eps helps stability z.requires_grad_(False) # cleanup (z is a tensor so local changes could propagate)return g['strength'] * dv.detach() # detach so no gradients returned
Let’s go!
torch.manual_seed(0) # for reproducibility as we change other thingswith torch.no_grad(): x1 = generate_samples(sub, n_samples=10, guidance_dict=guidance_dict, debug=False) logits, probs = classify(classifier, x1)show_probs(probs, x=x1)
BLAM! Our desired goal: in order, along the diagonal. We see that guidance isn’t perfect, but that’s partially a function of the VAE and the flow model too.
To get a better survey of the capabilities, let’s make a 10x10 grid of outputs with classes along each column:
That worked fine, and on a GPU it’s pretty fast, but on a CPU, it’s painfully slow. So instead, let’s…
Train a Latent Classifier
We’ll train a model z_classifier that looks only in latent space, so we can use it as a guidance signal. This can be a very simple model consisting of a few Linear layers:
It will operate on latent encodings of the MNIST dataset. So let’s save the encoded latents to disk and load them into memory. These will be our training and test data.
Code
# You can probably skip this code block. It just runs MNIST# through the VAE's encoder and saves it to disk. from torch.utils.data import DataLoader, TensorDataset, Subsetfrom torchvision.datasets import MNISTfrom torchvision.transforms import ToTensorfrom glob import glob import math @torch.no_grad()def encode_dataset(vae, dataset, batch_size=512, chunk_size=10000, tag="train"):"""Encode dataset into VAE latents (z = mu), saving progress in temp chunk files. We use temp chunks in case execution gets interrupted, we can try again & resume. """ device =next(vae.parameters()).device total_chunks = math.ceil(len(dataset) / chunk_size) # check for existing chunk files basename =f"tmp_chunk_{tag}" chunk_files = glob(basename+'*.pt') existing_chunks =len(chunk_files)print(f"{tag}: Found {existing_chunks} of {total_chunks} expected chunks. Generating remaining...")for c inrange(existing_chunks, total_chunks): print(f"chunk {c+1}/{total_chunks}:",end="") indices =list(range(c*chunk_size, (c+1)*chunk_size )) data_subset = Subset(dataset, indices) loader = DataLoader(data_subset, batch_size=batch_size, shuffle=False) all_latents, all_labels = [], [] must_flatten =Nonewith torch.no_grad():for data, labels in loader:print("=",end="") # dumb progress bar x = data.to(device)# next bit is so it should work with linear layers or convif must_flatten isNoneor must_flatten==False:try: z, must_flatten = vae.encoder(x), FalseexceptRuntimeError: z = vae.encoder(x.view(x.size(0), -1)) must_flatten =Trueelse: z = vae.encoder(x.view(x.size(0), -1)) mu, logvar = z all_latents.append(mu.cpu()) all_labels.append(labels) chunk_latents, chunk_labels = torch.cat(all_latents), torch.cat(all_labels) tmp_filename =f"{basename}_{c+1}.pt"print(f"| Saving chunk to {tmp_filename}") torch.save({ 'latents':chunk_latents, 'labels':chunk_labels }, tmp_filename)# Assemble all the chunks from files and return tensors.print("Assembling", basename+'*.pt ...') all_latents, all_labels = [], []for chunk_file in glob(basename+'*.pt'): chunk_dict = torch.load(chunk_file, weights_only=True) all_latents.append(chunk_dict['latents']) all_labels.append(chunk_dict['labels']) latents, labels = torch.cat(all_latents), torch.cat(all_labels)#for f in glob(tmp_file_base+'*.pt'): os.remove(f) # clean up return latents, labels def encode_mnist(vae, filename=None, batch_size=512):print("Acquiring train & test MNIST image datasets...") train_ds = MNIST(root='./data', train=True, download=True, transform=ToTensor()) test_ds = MNIST(root='./data', train=False, download=True, transform=ToTensor())print(f"\nEncoding dataset to latents...") train_latents, train_labels = encode_dataset(vae, train_ds, batch_size=batch_size, tag='train') test_latents, test_labels = encode_dataset(vae, test_ds, batch_size=batch_size, tag='test')for f in glob('tmp_chunk_t*_c*.pt'): os.remove(f) # clean upif filename isnotNone:print(f"Saving to {filename} ...") torch.save({ 'train_z': train_latents, 'test_z': test_latents,'train_labels': train_labels, 'test_labels': test_labels }, filename)return train_latents, train_labels# Encode the datasetlatent_data_filename ='mnist_latents.pt'ifnot os.path.exists(latent_data_filename): train_latents, train_labels = encode_mnist(sub.vae, filename=latent_data_filename)def load_data(filename):if'MyDrive'in filename:from google.colab import drive drive.mount('/content/drive') data_dict = torch.load(filename, weights_only=True)return data_dictdata_dict = load_data(latent_data_filename)train_z, test_z = data_dict['train_z'], data_dict['test_z']train_z.shape, test_z.shape# Create datasets from the latent tensorstrain_latent_ds = TensorDataset(train_z, data_dict['train_labels'][:train_z.shape[0]])test_latent_ds = TensorDataset(test_z, data_dict['test_labels'])batch_size =512train_latent_dl = DataLoader(train_latent_ds, batch_size=batch_size, shuffle=True)test_latent_dl = DataLoader(test_latent_ds, batch_size=batch_size, shuffle=False)print(f"Train batches: {len(train_latent_dl)}, Test batches: {len(test_latent_dl)}")# print single latent sizeprint(f"Latent size: {train_latent_ds[0][0].shape}")
Train batches: 118, Test batches: 20
Latent size: torch.Size([16])
Let’s test our newly-trained latent classifier, to make sure it works before trying to use it for guidance. We’ll pull up data samples with known ground-truth “target” labels, and compare these to the predictions from the classifier. If the targets and predictions match up, we’re good to go:
Now that we have a trained classifier that operates in latent space, we can run basically the same code as before, only it will execute wayyyyy faster…
Code
guidance_dict ={'classifier': z_classifier,'decode': None, # no decoding, latent space only'loss_fn': torch.nn.CrossEntropyLoss(reduction='none'), # don't sum across batch dim'target': torch.arange(10).repeat(10).to(device),'strength': 5.0, # "guidance strength"'t_min': 0.01, 't_max': 0.99, }torch.manual_seed(42) # remove for new samples each timex1 = generate_samples(sub, n_samples=len(guidance_dict['target']), guidance_dict=guidance_dict)show_grid(x1, "Latent-Only Guidance")
Was that fast or WHAT?!
Since we no longer have to propagate gradients through the much larger VAE decoder model and pixel-space classifer, we can get answers a lot faster via our small latents-only classifier.
Let’s move on to another application of guidance, for which our guidance signal doesn’t depend on a separate trained (classifier) model at all: inpainting.
When inpainting, we have some “mask” inside which some of the data have been removed, and we want to use the model to fill in the missing part in a way that matches with the surrounding pixels. Let’s take a look at an example from MNIST, where we show an original image, the mask (where white=keep, black=remove) and the masked-out image:
Another example would be a picture of a face where you’ve blocked out the nose and you want the model to fill in a nose. Now, some of the “filling in” you can get “for free” because the model has only been exposed to data that satisfies the manifold or probability distribution of the training data – e.g. If it was trained on faces, then it only ever saw faces with noses and hence can only generate faces with noses – but the real trick is to do it “well” and have it be “good” in the end. ;-)
There’s a wealth of information on guidance as it was originally applied to diffusion models. Sander Dieleman’s blog post, “Guidance: a cheat code for diffusion models”[8], is a classic and should be read by all. Yet because of the stochastic/random nature of the diffusion path, there are several “complicating” aspects of diffusion guidance that we’re going to gloss over in this tutorial because in the case of deterministic, smooth flow-model trajectories, things become a lot more intuitive.
We’ll follow a method outlined in the paper “Training-free Linear Image Inverses via Flows”[12] by Pokle et al, a method that applies to general linear inverse problems of which inpainting is a particular case, and we’ll simplify their method to adapt it for just inpainting.
The method will be to try to generate an entire new image \(x_1\) that everywhere outside the mask matches up with the pixels in user-supplied (masked) image \(y\). So the constraint will be, given a 2D mask \(M\) (where \(M\)=1 means there’s an original pixel there, and \(M\)=0 is the masked-out region), to require that our estimate image \(\widehat{x_1}\) (i.e. the decoded/pixel version of the estimated latents \(\widehat{z_1}\) ) satisfies \(M*\widehat{x_1} = M* y\,\) , or in a “residual form”, we’ll just compute the Mean Squared Error (MSE) of \(M*(\widehat{x_1}-y)\):
\[ {\rm Constraint:} = M^2 * (\widehat{x_1}-y)^2 \] (and if we want, we can use the fact that \(M\) being a binary mask means \(M^2 = M\)).
If we want to do latent-only inpainting (which will be the fastest), then the same constraint applies just with the simplification \(\widehat{x_1} = \widehat{z_1}\)
The authors of the paper recommend only doing guidance from \(t=0.2\) onward because prior to that, it’s hard to make any meaningful estimate.. In fact, they don’t even integrate before \(t = 0.2\). They just interpolate between the source and the target data to get their starting point at \(t = 0.2\).
To use our constraint in the guidance equation (3) for computing \(\Delta v\,\), we’ll need to turn our constraint into a likelihood by raising it to an exponential power – so we get a Gaussian! But the guidance equation includes a logarithm that immediately undoes our exponentiation:
The gradient part is \[ \nabla_{z_t} M^2 *(\widehat{x_1}-y)^2 = 2M^2*(\widehat{x_1}-y) {\partial \widehat{x_1} \over \partial z_t } \]
If we’re inpainting in latent space and not using the decoder for the constraint, then \({\partial \widehat{x_1} / \partial z_t } = 1\). Otherwise that term will require evaluation via PyTorch’s autograd (=slow).
Our earlier time scaling was \(1/(1-t)\); turns out that doesn’t work very well in practice when it comes to inpainting. Instead, we’ll use a different time scaling that delivers good (albeit not perfect) results: \((1-t)/t\). Thus our full equation for the velocity correction will be:
\[ \Delta \vec{v} = -\eta {1-t\over t} M^2 *(\widehat{x_1} - y){\partial\widehat{x_1}\over\partial{z_t}}, \] where we absorbed the factor of 2 into \(\eta\), and the last partial derivative term can be one if we do latent-only inpainting.
Let’s implement this in code, using two different versions of the gradient calculation, depending on whether we can do it all in latent space or if we need to propagate gradients through the decoder:
Code
@torch.no_grad() # gradients computed analytically!def ip_latents_grad(v_t, z, t, g:dict, eps=1e-6, **kwargs):"inpainting gradients for latent-only guidance, fast" z1_hat = z + (1-t)*v_treturn g['M_sq'] * (z1_hat - g['y']) # x1_hat = z1_hat, dz1_hat/dz_t=1@torch.enable_grad()def ip_pixels_grad(v_t, z, t, g:dict, eps=1e-6, img_size=(28,28), **kwargs):"inpainting gradients for pixel-space guidance. needs decoder & tracks via autograd, = slow" z.requires_grad_(True) z1_hat = z + (1-t)*v_t x1_hat = F.sigmoid(g['decode'](z1_hat)).view(-1,1, *img_size) grad_x = g['M_sq'] * (x1_hat - g['y']) grad_z = torch.autograd.grad(x1_hat, z, grad_outputs=grad_x,retain_graph=False)[0] # mults grad_x by dx1_hat/dz1_hat z.requires_grad_(False)return grad_z.detach() # don't send gradients onwarddef t_timescale(t, choice='mine'):"our choice for adaptive time scale"if choice =='simple': return1/(1-t) # our earlier scale; doesn't workelif choice=='pokle': return (1-t)**2/ ((1-t)**2+ t**2) # from pokle et al; can't get it to workelif choice=='constant': return4# or any constant. The 4 is from Pokle et alelse: return (1-t)/t # Works pretty well! Strong guidance at start, zero at enddef compute_dv_inpainting(v_t, z, t, g:dict, **kwargs):"wrapper to call appropriate gradient-computation routine"if t < g['t_min'] or t > g['t_max']: return torch.zeros_like(v_t) grad_fn = ip_latents_grad if g['decode'] isNoneelse ip_pixels_grad grad = grad_fn(v_t, z, t, g, **kwargs) dv =-g['strength'] * t_timescale(t) * gradreturn dv.detach()
Do the Inpainting
Start with masked images \(y\)…
Code
# setup the datay = torch.stack([test_ds[i][0] for i inrange(50)])print(y.shape)y = M*yshow_grid(y.squeeze(), "Masked Images")
torch.Size([50, 1, 28, 28])
And now we run the inpainting code
Code
compute_dv = compute_dv_inpainting # register our new guidance routineinpainting_dict ={'decode': sub.decode, # how to decode to pixel space for classifier'M_sq': (M**2).to(device),'y': y.to(device),'strength': 1.0, # "guidance strength", you may vary this't_min': 0.2, 't_max': 0.999} # t range to apply guidance, may vary thesewith torch.no_grad(): torch.manual_seed(0) # for reproducibility as we change other things t0 =0.2# starting time as per Pokle et al z0 = torch.randn([len(y), sub.latent_dim]).to(sub.device) zy = sub.encode(y.to(device)) # encoded version of masked image z0 = z0 * (1-t0) + zy * t0 # interpolation init inpainting_dict['t_min'] = t0 x1 = generate_samples(sub, n_samples=len(y), t0=t0, z0=z0, guidance_dict=inpainting_dict, warp_fn=None, debug=False)show_grid(y.squeeze(), "Masked Images")show_grid(x1, "Inpainted Images")
We see that the generated images generally look great, although in some cases, the in-painting code has changed pixels even where the mask is 1. We can disallow this by just resetting those values to the pixels in \(y\).
Turning up the guidance strength would also enforce our constraint better, but turning up too high causes the whole thing to diverge and we get garbage out.
In order to experiment with other methods more easily, we should do inpainting only in latent space, and for that we will need a model that supports spatial latents….
Latent-Only Inpainting
To do inpainting in latent space, we’ll need to switch models to one where the latents preserve the spatial structure of the original images.
Get the spatial-latents model
# Get Spatial VAE & FLow DiT Model!wget -q --no-clobber=off https://raw.githubusercontent.com/dlaieburner/2025-leaderboard/refs/heads/main/sample_submission_dit.pytry:del SubmissionInterface # remove Marco's from earlier; make it reloadexceptNameError:pass# nevermindfrom sample_submission_dit import SubmissionInterfacesub = SubmissionInterface().to(device)
There are some limitations to what’s being produced. The problem is due to the low resolution of our latents. Inpainting algorithms typically assume higher resolution in order to work accurately. But this lesson was designed to run quickly on CPUs and thus there is a trade-off.
Now, if we want to enforce the original pixels where the mask is one, we can do that at every stage of the integration process. We just need to modify the integration process to overwrite \(z\) wherever Mz is 1.
Let’s move on and try a different method that can achieve similar results, albiet differently…
The term ‘guidance’ typically refers to velocity modifications, but PnP-Flow by Martin et al[13] achieves similar results by adjusting latent positions \(z\) directly.\(^,\)
PnP-Flow assumes straight-line trajectories, making the forward projection trivial: \(\widehat{z_1}\) is reached by simple linear extrapolation. Instead of incrementally moving \(z\) from \(t=0\) to \(t=1\), PnP-Flow projects forward to \(\widehat{z_1}\) and iterates on that estimate through a series of correction and projection steps. The first step applies our gradient correction:
where \(z_1^*\) (my notation) is our goal i.e. the endpoint of our projected course correction, and \(F(\widehat{z_1},y)\) is our (log-exp probability) constraint. For the time scaling, the PnP-Flow authors recommend \(\gamma_t = (1-t)^\alpha\) with \(\alpha \in [0,1]\) is a hyperparameter chosen according to the task – e.g., they use \(\alpha\)’s as large as 0.8 for denoising tasks, 0.5 for box inpainting, and 0.01 for random inpainting. This choice of \(\gamma_t\) is a bit different from our earlier one of \((1-t)/t\). Both go to zero as \(t \rightarrow 1\), but approach it differently and have different asymptotics as \(t\rightarrow 0\).
In the graph below, we show our earlier choice of \((1 - t)/t\) in green and \((1 - t)^\alpha\) in purple for various choices of \(\alpha\):
…where for “box inpainting” as we did above, they use \(\alpha\)=0.5.
But PnP-Flow doesn’t stop there! Two other key steps remain. We then project backward to overwrite\(z_t\) with a corrected value:
….and loop over Steps 1 to 3 for each value of \(t\) in our set of (discrete) integration steps, i.e. after Step 3, we let \(t := t+\Delta\,t\) and go back to Step 1. Our final value of \(\widehat{z_1}\) will be the output.
This image from the PnP-Flow paper may prove instructive, showing 3 different instances of the 3 PNP steps:
Image from PnP-Flow paper [12], slightly re-annotated by me.
This has a superficial resemblance to the “ping-pong” integration method used by the flow model Stable Audio Open Small (SAOS) [14], with a key distinction: the ping-pong integrator and updates the time-integrated latent variable \(z\) (called “\(x\)” in SAOS), whereas for PnP-Flow it is the projection \(\widehat{z_1}\) (called “denoised” in SAOS) that is the primary variable that is maintained between steps. This is a subtle distinction but worth noting.
To implement PnP-Flow in code, let’s replace our “integrator” with something specific to PnP-Flow:
The model we used earlier for latents-only inpainting was trained to have straight trajectories, so we should be able to use it again here, just calling sample_pnpflow (instead of integrate_path). The results are as follows:
…So that works, though it also seems to “take some liberties”: most of the letters in the second group look like “boldface” versions of the top ones. This could be because of the low spatial resolution of the latents, e.g. that they are encoding information about curvature or other aspects of the shape.
The results are similar to the previous inpainting method of Pokle et al, just a different way of doing it that may prove worthwhile.
PnP-Flow is a general guidance method not limited to inpainting or even “linear” image degradations. I recommend looking into it further for other tasks and datasets — and let me know what sort of results you find!
Training generative models can be expensive—lots of data, electricity, compute time. So what if you could take a pretrained model and add controls at inference time instead? That’s what this tutorial explored. While similar ideas are emerging for steering autoregressive models [16], here we focused on pretrained flow models.
The key idea is simple: at each integration step, you project forward to estimate where you’ll end up (\(\,\widehat{z_1}\,\)) , check how far that is from where you want to be, and add a small velocity correction to steer toward your goal. We applied this to an unconditional MNIST flow model for two tasks: generating specific digit classes via classifier guidance, and filling in masked-out regions via inpainting.
We looked at four approaches. First, standard classifier guidance in pixel space—it works but it’s slow because you’re propagating gradients through the VAE decoder. Second, we trained a simple latent-space classifier and did the same thing much faster. Third, we implemented the linear inpainting method from Pokle et al, which operates directly on latents. Fourth, we tried PnP-Flow, which achieves guidance not by correcting velocities but by iteratively projecting samples forward and backward in time.
The math here is much simpler than the corresponding diffusion methods because flow trajectories are smooth and deterministic. We’ve glossed over a lot of detail compared to the research papers, but hopefully this gives you enough to experiment with your own controls. There are limits to the effectiveness of guidance: small models that don’t generalize well won’t suddenly work miracles if you try to push them too far outside their training distribution. Nevertheless, these plugin methods are worth exploring as accessible ways to steer generative flows where you want them to go.
This work was supported by Hyperstate Music, Inc. who are awesome and you should invest in them! ;-) For providing feedback on early drafts of this tutorial: thanks to Raymond Fan, Alan Lockett, and my amazing students in PHY/DSC/BSA 4420, “Deep Learning and AI Ethics”! Also thanks to Danilo Comminiello and the signal processing group at Sapienza University of Rome for encouraging me to present a preceding lecture [5] at IJCNN 2025 – this blog post is the missing piece!
This document was partially prepared using the SolveIt platform of Answer.ai and via many discussions with Claude.ai; still overwhelmingly authored by the hands of this here human.
[2]
Y. Lipman, R. T. Q. Chen, H. Ben-Hamu, M. Nickel, and M. Le, “Flow matching for generative modeling,” in The Eleventh International Conference on Learning Representations (ICLR), 2023. Available: https://openreview.net/forum?id=PqvMRDCJT9t
[3]
X. Liu, C. Gong, and Q. Liu, “Flow straight and fast: Learning to generate and transfer data with rectified flow,” in 11th International Conference on Learning Representations (ICLR), 2023. Available: https://openreview.net/forum?id=XVjTT1nw5z
[7]
H. Ye et al., “TFG: Unified training-free guidance for diffusion models,”Advances in Neural Information Processing Systems, vol. 37, pp. 22370–22417, 2024, Available: https://arxiv.org/abs/2409.15761
[9]
Z. W. Blasingame and C. Liu, “Greed is good: Guided generation from a greedy perspective,” in Frontiers in probabilistic inference: Learning meets sampling, 2025. Available: https://openreview.net/forum?id=o4yQzZ5qCW
[11]
A. Tong et al., “Improving and generalizing flow-based generative models with minibatch optimal transport,”Transactions on Machine Learning Research, 2024, Available: https://openreview.net/forum?id=CD9Snc73AW
[12]
A. Pokle, M. J. Muckley, R. T. Q. Chen, and B. Karrer, “Training-free linear image inverses via flows,”Transactions on Machine Learning Research, 2024, Available: https://openreview.net/forum?id=PLIt3a4yTm
[13]
S. Martin, A. Gagneux, P. Hagemann, and G. Steidl, “PnP-flow: Plug-and-play image restoration with flow matching.” 2024. Available: https://arxiv.org/abs/2410.02423
[14]
Z. Novack et al., “Fast text-to-audio generation with adversarial post-training,” in IEEEWorkshop on Applications of Signal Processing to Audio and Acoustics (WASPAA), 2025. Available: https://arxiv.org/abs/2505.08175
[15]
Z. Novack, J. McAuley, T. Berg-Kirkpatrick, and N. J. Bryan, “DITTO: Diffusion inference-time t-optimization for music generation,” in Forty-first International Conference on Machine Learning (ICML), 2024. Available: https://openreview.net/forum?id=z5Ux2u6t7U
[16]
D. Zhao, D. Beaglehole, T. Berg-Kirkpatrick, J. McAuley, and Z. Novack, “Steering autoregressive music generation with recursive feature machines.” 2025. Available: https://arxiv.org/abs/2510.19127
If you found this tutorial helpful, I encourage you to cite it:
@misc{hawley2025flowwhere,title={Flow Where You Want}, author={Scott H. Hawley},date={Nov 27, 2025},howpublished={Scott Hawley's Blog},url={https://drscotthawley.github.io/blog/posts/FlowWhereYouWant.html}}
