P-Guide: Parameter-Efficient Prior Steering for Single-Pass CFG Inference

Abstract: Classifier-Free Guidance (CFG) is essential for high-fidelity conditional generation in flow matching, yet it imposes significant computational overhead by requiring dual forward passes at each sampling step. In this work, we address this bottleneck by introducing \textbf{P-Guide}, a framework that achieves high-quality guidance through a single inference pass by modulating only the initial latent state. We further show that, under a first-order approximation, P-Guide is equivalent to CFG in the sense that it steers generation from the prior space, without requiring explicit velocity field extrapolation during sampling. We consider both homoscedastic and \textbf{heteroscedastic} priors, and find that jointly modeling the mean and variance enables adaptive loss attenuation and improved robustness to data uncertainty. Extensive experiments demonstrate that P-Guide reduces inference latency by approximately 50\% while maintaining fidelity and prompt alignment competitive with standard dual-pass CFG baselines.

• The paper introduces a learned conditional prior that applies classifier-free guidance at the initial latent state, enabling one velocity-network evaluation per ODE step and reducing inference cost by about 50% with only 1.247 MB of added parameters.
• Experiments on MNIST, CIFAR-10, and ImageNet show competitive or improved FID near guidance scale 1.0, including ImageNet improvements from 27.78 to 22.33 with a U-Net and from 33.68 to 25.07 with DiT-B/2, although performance degrades at larger scales.
• The method is theoretically justified only to first order and works best as a complement to standard CFG, with open issues involving guidance-scale calibration, text conditioning, heteroscedastic priors, and reliance on pretrained initialization.

P-Guide addresses a well-known inefficiency in conditional generative modeling: classifier-free guidance (CFG) requires two forward passes of the velocity network at every ODE integration step, doubling inference cost relative to unguided sampling. The paper's central move is to relocate guidance from the sampling trajectory to its origin. Instead of extrapolating between conditional and unconditional velocity fields at each step, P-Guide applies CFG once, in the prior space, by modulating the initial latent state with a learned condition-dependent distribution. The result is single-pass guided sampling with approximately 50% lower inference cost and only 1.247 MB of additional parameters.

Method

The framework rests on replacing the standard Gaussian prior $p_0(z) = \mathcal{N}(0, \mathbf{I})$ with a learned conditional prior $p_0(z|y)$. Two variants are considered. In the homoscedastic case, the condition modulates only the mean: $z = \epsilon + \mu_\phi(y)$, where regression optimality implies $\mu_\phi(y)$ approximates $\mathbb{E}[x|y]$. In the heteroscedastic extension, the module outputs both $\mu_\phi(y)$ and $\sigma_\phi(y)$, so the initial state becomes $z = \mu_\phi(y) + \sigma_\phi(y)\odot\epsilon$. Training is decoupled into two stages: Stage 1 fits the prior module by minimizing Gaussian negative log-likelihood, which yields adaptive loss attenuation (gradients on the mean are scaled by $1/\sigma_\phi^2$, down-weighting noisy targets); Stage 2 freezes the prior and trains the velocity field $v_\theta$ with a Rectified Flow objective whose starting points are drawn from the learned conditional prior.

At inference, the guided initial state is constructed by linearly extrapolating both prior parameters:

$p_0(z|y)$0

and the probability flow ODE is then solved with a single velocity evaluation per step.

Trajectory-level approximation

The theoretical justification is a first-order argument rather than an exact equivalence. Under smoothness of the flow map $p_0(z|y)$1, shared-noise coupling between $p_0(z|y)$2 and $p_0(z|y)$3, and a local perturbation regime, Taylor expansion gives $p_0(z|y)$4, where $p_0(z|y)$5 is the flow Jacobian. Both state-space and velocity-space differences are thus governed by the same first-order direction $p_0(z|y)$6, and the paper argues that prior-space steering induces a trajectory perturbation aligned with velocity-space CFG up to scaling differences between $p_0(z|y)$7 and $p_0(z|y)$8. A supplementary proposition connects the prior shift to the conditional score via the standard identity $p_0(z|y)$9.

Two caveats deserve emphasis. First, the equivalence is explicitly first-order: higher-order terms are neglected, and scaling differences between prior-space and velocity-space guidance scales remain unresolved. Second, the appendix itself shows that a closed-form "distribution-level" CFG (geometric interpolation of Gaussians) breaks the shared-noise coupling assumption and empirically fails on MNIST, producing distorted samples — evidence that the coupling structure, not merely the marginal density, is what makes prior steering work. This also implies the heteroscedastic variant is theoretically less well-grounded than the homoscedastic one, although the paper notes the inconsistency claim holds only in the general heteroscedastic case.

Empirical results

Experiments cover MNIST, CIFAR-10, and ImageNet-1k ($z = \epsilon + \mu_\phi(y)$0), all trained for 400K steps under identical budgets against dual-pass CFM baselines, evaluated with FID, sFID, IS, class accuracy, and GFLOPs over 50 steps.

On MNIST, P-Guide with learnable variance achieves FID 1.40 at $z = \epsilon + \mu_\phi(y)$1 versus 5.87 for single-pass CFM, while maintaining roughly constant compute (~417.6 GFLOPs) across all guidance scales; dual-pass CFM costs ~835.2 GFLOPs. On CIFAR-10, P-Guide reaches FID 6.51 at $z = \epsilon + \mu_\phi(y)$2 (versus 9.44 for CFM), but the picture is more mixed: dual-pass CFM attains substantially better peak quality (FID 2.54 at $z = \epsilon + \mu_\phi(y)$3), and P-Guide degrades sharply at $z = \epsilon + \mu_\phi(y)$4 (FID 38.46). On ImageNet-1k with U-Net and DiT-B/2 backbones (the latter initialized from a pretrained latent flow model and fine-tuned for ~4% of the original training budget), FID improves from 27.78 to 22.33 (U-Net, $z = \epsilon + \mu_\phi(y)$5) and from 33.68 to 25.07 (DiT-B/2, $z = \epsilon + \mu_\phi(y)$6), again with constant single-pass cost and only 1.247 MB added parameters (<0.5% storage overhead).

Ablations reinforce the mechanism's interpretation. Scaling the prior module tenfold (to 12.5 MB) leaves FID essentially unchanged (33.68 → 34.34), indicating the prior shift acts as a low-dimensional directional signal rather than requiring high-capacity function approximation. Applying guidance to the variance contributes only marginal gains over mean-only steering. Joint use of P-Guide with standard CFG is stable and complementary — e.g., U-Net FID drops to 13.26 at $z = \epsilon + \mu_\phi(y)$7 — supporting compositionality of prior-space and trajectory-level control.

Limitations and open questions

The paper is candid about several constraints. The most significant is the narrow effective guidance range: unlike standard CFG, which remains stable up to $z = \epsilon + \mu_\phi(y)$8 on CIFAR-10, P-Guide degrades beyond roughly $z = \epsilon + \mu_\phi(y)$9–$\mu_\phi(y)$0, meaning it cannot reach the peak fidelity of dual-pass guidance when used alone. The first-order equivalence leaves the relationship between prior-space scale $\mu_\phi(y)$1 and velocity-space scale $\mu_\phi(y)$2 unspecified, and the theory relies on assumptions (smooth flow map, shared-noise coupling, local perturbation regime) whose validity at large shifts or on high-dimensional text-conditioned manifolds is not established. All experiments use class conditioning; whether the approach transfers to free-form text prompts in systems like SD3 or FLUX remains untested. Finally, the ImageNet results depend on initializing Stage 2 from a pretrained flow model — from-scratch training under the same budget yields notably worse FID (49.03 vs. 33.68) — so the reported efficiency figures partly reflect transfer learning rather than the method alone.

Conclusion

P-Guide offers a parameter-efficient route to halving CFG inference cost by injecting guidance at the trajectory origin through a learned conditional prior, supported by a first-order trajectory-level approximation and heteroscedastic uncertainty modeling. Its empirical profile is that of a strong complement rather than a replacement for standard CFG: competitive quality near $\mu_\phi(y)$3 at half the compute, but a restricted guidance range and dependence on pretrained initialization. The open questions it leaves — exact scale correspondence with velocity-space CFG, behavior under text conditioning, and robustness beyond the local perturbation regime — define the natural next steps for source-space control in continuous-time generative models.