Step-Conditioned Velocity Networks

• Step-conditioned velocity networks are architectures that condition on a continuous step parameter t to model nonstationary velocity fields in linear–Gaussian interpolants.
• Stable Velocity Matching reduces training variance by replacing single-reference targets with unbiased mixtures across n samples, leading to an O(1/n) variance reduction.
• The approach enables accelerated inference with closed-form integration and uses variance-aware auxiliary losses to improve sample quality in low-variance regimes.

Step-conditioned velocity networks are neural architectures used in modern flow-matching generative models where the network explicitly conditions on a continuous "step" parameter, typically denoted $t$, that indexes the progression from clean data to noise. This conditioning enables the network to parameterize nonstationary velocity fields associated with a linear–Gaussian interpolation—central to many recent advances in diffusion, flow matching, and score-based generative modeling. The step-conditioned velocity paradigm has become foundational for imposing temporal structure, improving expressivity, and enabling regime-aware training and inference procedures, such as variance-reduced objectives and accelerated sampling (Yang et al., 5 Feb 2026).

1. Linear–Gaussian Interpolants and Velocity Fields

The backbone of step-conditioned velocity networks lies in a one-parameter family of linear–Gaussian processes: $$x_t = \alpha_t x_0 + \sigma_t \varepsilon, \quad \varepsilon \sim \mathcal N(0,I), \quad t \in [0,1], \; \alpha_0=1, \sigma_0=0, \alpha_1=0, \sigma_1=1, \; \alpha_t^2+\sigma_t^2 > 0,$$ which smoothly interpolates between clean data $x_0$ ($t=0$) and pure noise ($t=1$). The velocity structure of these interpolants is governed by the instantaneous rate of change with respect to $t$:

• Conditional velocity: For each fixed reference $x_0$, the velocity field is

$$v_t(x_t \mid x_0) = \frac{\sigma_t'}{\sigma_t}(x_t - \alpha_t x_0) + \alpha_t' x_0,$$

where prime denotes differentiation with respect to $t$. This field depends explicitly on both $x_t$ and the interpolation time $t$.

• Marginal velocity: When $x_0$ is randomized according to the data distribution, the marginal or true velocity field is

$$u(x_t, t) = \mathbb E_{x_0 \sim p_t(\cdot \mid x_t)}[v_t(x_t \mid x_0)],$$

where $p_t(x_0 \mid x_t)$ is the posterior induced by the interpolant.

A step-conditioned velocity network, denoted $v_\theta(x_t, t)$, is trained to match $u(x_t, t)$. Conditioning on the step $t$ is essential: it enables the network to model the nonstationary transition from the highly structured, low-variance data regime to the high-variance, noise-dominated regime.

2. Variance Analysis in Single-Sample Conditional Targets

Training by directly matching the conditional velocity using single reference samples from the data introduces high variance into the learning objective. Though the conditional velocity is an unbiased estimator of the true velocity field, its variance with respect to the data posterior is generally non-negligible and can be expressed as: $\Var_{\text{CFM}}(t) = \mathbb{E}_{x_t \sim p_t} \| v_t(x_t \mid x_0) - u(x_t, t)\|^2 = \mathbb{E}_{p_t(x_t)}[\Tr(\Cov_{p_t(x_0 \mid x_t)}(v_t(x_t \mid x_0)))].$ For Gaussian posteriors $x_0 \mid x_t \sim \mathcal N(\bar x_0, \Sigma_t)$, the conditional-velocity variance admits a closed-form: $\Var_{\text{CFM}}(t) = \left(\frac{\sigma_t'}{\sigma_t}\right)^2 \alpha_t^2 \Tr(\Sigma_t).$ In the linear interpolant ($\alpha_t=1-t$, $\sigma_t=t$), this variance is vanishingly small for small $t$ (data regime) and grows towards $t \to 1$ (noise regime). This establishes a regime dichotomy: training is stable and low-variance near the data, but optimization becomes unstable as the trajectory approaches the prior.

3. Stable Velocity Matching: Unbiased Variance Reduction

To mitigate the high-variance bottleneck, Stable Velocity Matching (StableVM) replaces the single-reference target with a mixture constructed from $n$ independent samples $\{x_0^i\}_{i=1}^n$:

• The mixture path is defined as

$$p_t^{\text{GMM}}(x_t \mid \{x_0^i\}) = \frac{1}{n}\sum_{i=1}^n p_t(x_t \mid x_0^i).$$

• The target is the self-normalized weighted average

$$\hat v_{\text{StableVM}}(x_t; \{x_0^i\}) = \frac{\sum_{k=1}^n p_t(x_t \mid x_0^k)v_t(x_t \mid x_0^k)}{\sum_{j=1}^n p_t(x_t \mid x_0^j)}.$$

• The objective is

$$\mathcal L_{\text{StableVM}} = \mathbb{E}_{t, \{x_0^i\}, x_t \sim p_t^{\text{GMM}}}\lambda_t\,\|v_\theta(x_t, t) - \hat v_{\text{StableVM}}(x_t; \{x_0^i\})\|^2.$$

This construction remains unbiased by design ($$\mathbb{E}[\hat v_{\text{StableVM}} \mid x_t]=u(x_t,t)$$), but crucially, the variance decreases as $O(1/n)$. Formally,

$\Var_{\text{StableVM}}(t) \leq \Var_{\text{CFM}}(t) / n,$

allowing for substantial variance reduction without introducing bias (Yang et al., 5 Feb 2026).

4. Regime-Gated Auxiliary Objectives: Variance-Aware Representation Alignment

Auxiliary representation losses, such as feature alignment with a pretrained encoder, can improve sample quality, particularly in the low-variance (data-proximate) regime. However, these signals are generally ill-posed in the high-variance regime near $t \approx 1$, where the posterior is diffuse and the data signal is negligible.

Variance-Aware Representation Alignment (VA-REPA) addresses this by modulating auxiliary losses with a regime-dependent weight $w(t) \in [0,1]$ that vanishes in the high-variance regime. The total objective is: $$\mathcal L = \mathcal L_{\text{StableVM}} + \lambda_{\mathrm{RA}} \frac{\mathbb{E}_{t, x_t}[w(t)\,\ell_{\mathrm{RA}}(x_t)]}{\mathbb{E}_t[w(t)]},$$ where $\ell_{\mathrm{RA}}(x_t)$ is the auxiliary representation alignment loss.

Common functional forms for $w(t)$ include:

• Hard cutoffs: $w(t) = \mathbf{1}[t < \xi]$
• Smooth sigmoids: $w(t) = \sigma(k(\xi-t))$
• SNR-based: $$w(t) = \operatorname{SNR}(t)/(\operatorname{SNR}(t) + \operatorname{SNR}(\xi))$$

The result is robust auxiliary supervision confined to the statistically reliable, low-variance region (Yang et al., 5 Feb 2026).

5. Closed-Form Accelerated Inference: Stable Velocity Sampling

In the low-variance regime ($t \leq \xi$), the conditional posterior $p_t(x_0 \mid x_t)$ collapses and the true velocity field approximates the conditional velocity: $u(x_t, t) \approx v_t(x_t \mid x_0).$ This admits closed-form integration of the reverse dynamics, permitting large-step, finetuning-free inference in both SDE and ODE formulations:

• SDE Posterior (DDIM-style): For $\tau < t$,

$$p_\tau(x_\tau \mid x_t, v_t(x_t)) = \mathcal N(\mu_{\tau|t}, \beta_t^2 I),$$

with

$$\mu_{\tau|t} = (\rho_t - \lambda_t \frac{\sigma_t'}{\sigma_t}) x_t + \lambda_t v_t(x_t),$$

and deterministic updates for $\beta_t = 0$.

• ODE Probability-Flow: Using the integral factor $\Psi_{t,\tau}$ and $C(s)$, the update becomes

$$x_\tau = \sigma_\tau \left[\left(\sigma_t - \frac{\sigma_t'}{\sigma_t} \Psi_{t,\tau}\right)^{-1} x_t + \Psi_{t,\tau} v_t(x_t)\right].$$

For the linear interpolant, these reduce to simple Euler steps ($x_\tau = x_t + (\tau-t) v_t(x_t)$), supporting arbitrarily large steps in the low-variance regime. This results in empirical sampling speedups exceeding $2\times$ in practice without quality degradation (Yang et al., 5 Feb 2026).

6. Architectural Encodings of Step $t$

In neural implementations, $t$ is encoded analogously to prior diffusion or flow models: either via sinusoidal ("Fourier") embeddings or learnable transformations, which are then concatenated or injected (e.g., using MLPs) into each block of a ResNet or Transformer backbone. This explicit step-conditioning is critical for capturing the nonstationary structure of the velocity field. Empirically, ablative studies demonstrate that without proper $t$-embedding, the model fails to learn the desired temporal dynamics; enhanced embeddings provide marginal benefit in the high-variance regime but exhibit diminishing returns once variance-reduction is in place.

7. Summary and Broader Implications

Step-conditioned velocity networks formalize and unify three variance-centric enhancements for flow-matching models: Stable Velocity Matching for unbiased, low-variance training; Variance-Aware Representation Alignment for regime-confined auxiliary supervision; and Stable Velocity Sampling for drastic acceleration in low-variance inference. Together, these techniques address both optimization stability and sampling efficiency, notably on high-dimensional datasets and large-scale pretrained models, as demonstrated across ImageNet and leading text-to-image/video architectures (Yang et al., 5 Feb 2026). A plausible implication is that regime-separated strategies may further generalize to non-Gaussian or multimodal flows, stimulating developments in expressive generative modeling under uncertainty.