AnyFlow: Any-Step Video Diffusion Model with On-Policy Flow Map Distillation

Consistency Models (CMs) [6] accelerate diffusion sampling by directly learning a mapping from $\mathbf{z}_{t}$ to $\mathbf{z}_{0}$ along the probability-flow ODE (PF-ODE) defined by a teacher model. For multi-step generation, CMs repeatedly inject noise into intermediate states and perform iterative denoising. Subsequent studies improve CM training stability through annealed time schedules [7, 14] and segmented consistency objectives [15, 16, 17]. sCM [8] further simplifies these design choices and provides a strong practical baseline. Building on sCM, rCM [9] introduces score distillation as a regularizer, achieving strong performance in both image and video diffusion distillation. In contrast to prior CM-based approaches, we explore flow map distillation to enable any-step video generation.

Flow Map Models [12, 11, 18] offer a unified perspective on diffusion modeling and diffusion distillation by learning a generalized transition operator $\mathbf{f}_{\theta}(\mathbf{z}_{t},t,r)$ between arbitrary points along PF-ODE trajectories. From this perspective, the formulation recovers consistency modeling [6] when $r=0$ and reduces to standard flow matching [19, 20] when $t=r$. However, training flow map models is challenging because it requires accurate transition learning across arbitrary time pairs. MeanFlow [13] is a representative approach that connects instantaneous and average velocities, but it relies on Jacobian-vector products (JVPs), which are difficult to scale under Fully Sharded Data Parallel (FSDP). To overcome this limitation, subsequent studies either approximate JVP terms numerically [21] or derive JVP-free algebraic formulations [22, 23]. Building on these advances, AnyFlow leverages flow map distillation for scalable any-step video diffusion, enabling efficient and flexible inference for high-fidelity generation.

As shown in Tab. 1, recent video diffusion distillation methods typically follow a two-stage pipeline: (1) forward training for initialization and (2) on-policy distillation for refinement. The first stage improves optimization stability, while the second reduces rollout drift during few-step inference. For example, rCM [9] initializes from sCM [8] and then applies consistency backward simulation [24] with distribution matching [25]. In causal video generation, Self-Forcing [10] adopts data-free consistency ODE initialization [26] before performing a similar on-policy refinement stage. Our method follows the same high-level pipeline but replaces consistency modeling with a flow map formulation, enabling stronger any-step performance.

A concurrent work, TMD [27], also studies a flow map formulation for bidirectional video diffusion distillation. The main difference lies in how efficient rollout is achieved. TMD improves rollout efficiency from an architectural perspective by sharing the backbone and introducing an additional flow head. In contrast, our method improves simulation efficiency from the flow-trajectory perspective by decomposing a full trajectory into shortcut transitions based on the composition property of flow maps. As a result, our approach naturally supports arbitrary step budgets and generalizes to both bidirectional and causal architectures.

Causal Video Diffusion. A growing line of work studies causal video diffusion models [28, 29, 30, 31]. These models often suffer from exposure bias, which leads to error accumulation during autoregressive generation. Self-Forcing [10] mitigates this issue through on-policy distillation, but its consistency-based formulation is mainly tailored to few-step settings for rollout efficiency. Another line of work [32] focuses more on multi-step models by explicitly modeling test-time errors during pretraining, but it is not directly designed for efficient few-step sampling. By contrast, our method starts from a flow map formulation and uses flow map backward simulation to support both few-step and multi-step causal video generation within a unified framework.

On-policy distillation mitigates test-time errors by training the student on its own generated trajectories and supervising them with a stronger teacher. This approach has demonstrated benefits in both large language models [33, 34, 35] and diffusion models [25, 10, 24].

For diffusion models, on-policy distillation is typically implemented via Distribution Matching Distillation (DMD) [25, 24] or adversarial distillation [36, 37]. Self-Forcing [10] combines self-rollouts with distribution matching against a bidirectional teacher, while APT2 [37] evaluates one-step video rollouts using a discriminator. Our method follows the same on-policy distillation paradigm but introduces flow map backward simulation, specifically designed for the flow map formulation, enabling any-step video generation.

We first introduce the basic flow map formalism used in AnyFlow. Let $\mathbf{z}_{t}$ denote the latent state at continuous time $t\in[0,1]$ under the probability-flow ODE (PF-ODE):

where $\mathbf{v}$ is the velocity field.

Define the exact flow map $\Phi_{r\leftarrow t}$ of this ODE as the operator that transports states from time $t$ to time $r$, i.e., $\Phi_{r\leftarrow t}(\mathbf{z}_{t})=\mathbf{z}_{r}$ for $1\geq t\geq r\geq 0$. It satisfies two standard properties: identity $\Phi_{t\leftarrow t}(\mathbf{z})=\mathbf{z}$ and composition $\Phi_{q\leftarrow r}\circ\Phi_{r\leftarrow t}=\Phi_{q\leftarrow t}$ for $t\geq r\geq q$.

In practice, a neural flow map model learns an approximation

with boundary condition $\mathbf{f}_{\theta}(\mathbf{z}_{t},t,t)=\mathbf{z}_{t}$. Compared with endpoint-only mappings, this parameterization models transitions between arbitrary time pairs, which naturally supports variable step sizes and any-step inference.

MeanFlow Objective. MeanFlow [13] is a representative algorithm for training flow map models. It parameterizes the averaged transport velocity on $[r,t]$ and approximates the transition as $\mathbf{f}_{\theta}(\mathbf{z}_{t},t,r)=\mathbf{z}_{t}-(t-r)\,\mathbf{u}_{\theta}(\mathbf{z}_{t},r,t)$. Following this formulation, we optimize $\mathbf{u}_{\theta}$ by:

where $\mathbf{u}_{\text{tgt}}=\mathbf{v}(\mathbf{z}_{t},t)-(t-r)\,\frac{d\mathbf{u}_{\theta}(\mathbf{z}_{t},r,t)}{dt}$. Here, $\operatorname{sg}(\cdot)$ denotes stop-gradient.

Differential Derivation Equation. Computing the Jacobian-vector product (JVP) for the derivative term $\frac{d\mathbf{u}(\mathbf{z}_{t},r,t)}{dt}$ is expensive and not fully compatible with Fully Sharded Data Parallel (FSDP). Transition Model [21] proposes to approximate this derivative using the finite difference method:

This approximation requires only two forward passes per training step and is compatible with FSDP.

On-policy distillation is designed to reduce train–test mismatch by training the student on its own rollout states while being guided by a strong teacher. In diffusion models, this paradigm is mainly instantiated as Distribution Matching Distillation (DMD). It typically consists of two key components: backward simulation and a distribution-matching objective.

Backward Simulation. Given noise $\mathbf{z}$, the student first produces a sample through its inference trajectory, $\hat{\mathbf{z}}_{0}=\mathbf{f}_{\theta}(\mathbf{z})$. This self-generated trajectory is then used to compute the distillation loss, which helps mitigate exposure bias and discretization error at test time.

Distribution Matching Objective. We follow DMD-style reverse-KL training [25, 24]. First, DMD re-noises the self-rollout sample $\hat{\mathbf{z}}_{0}$ at $t\in[0,T]$ as $\mathbf{z}_{t}=(1-t)\hat{\mathbf{z}}_{0}+t\bm{\epsilon}$, where $\bm{\epsilon}\sim\mathcal{N}(\mathbf{0},\mathbf{I})$. It then estimates the DMD gradient as follows:

where $s_{\text{real}}$ and $s_{\text{fake}}$ are the real and fake score functions, respectively.

As summarized in Tab. 1, most recent few-step video distillation methods are built on consistency models. Although effective under very small sampling budgets, their performance often degrades as the number of sampling steps increases (Fig. 1). The underlying reason is that consistency-style sampling repeatedly performs endpoint projection and re-noising, which gradually drives the trajectory away from the target PF-ODE path during multi-step inference.

Motivation. Our goal is to preserve the strong few-step efficiency of distilled video diffusion models while enabling quality to improve, rather than deteriorate, as more sampling steps are used. To this end, AnyFlow shifts the distillation target from fixed-point endpoint mapping to flow map transition learning, i.e., $\mathbf{f}_{\theta}:(\mathbf{z}_{t},t,r)\mapsto\mathbf{z}_{r}$. Instead of only predicting $\mathbf{z}_{0}$, the model learns transitions between arbitrary time pairs. This formulation naturally supports both small time gaps for stable local refinement and large time gaps for efficient long-range jumps, making it well suited for any-step generation.

Challenge. Forward flow map training alone (e.g., MeanFlow [13]) is insufficient for strong test-time performance in video generation. As shown in Fig. 3 and Tab. 2, it still suffers from test-time errors, especially discretization error in few-step sampling and exposure bias in causal generation. This motivates an additional on-policy distillation stage to correct rollout mismatch. However, designing such a stage for flow map models is non-trivial. Unlike consistency models, which can naturally reach $\mathbf{z}_{0}$ at intermediate steps and therefore allow KL gradients to be computed directly at endpoint states, flow map models require a new simulation strategy that can accommodate transitions between arbitrary time pairs while remaining efficient for training.

Overview. As illustrated in Fig. 4, AnyFlow is trained in two complementary stages: forward flow map training (Sec. 4.2.1) and on-policy distillation (Sec. 4.2.2). In Stage 1, we fine-tune a pretrained video diffusion model on teacher-synthesized data to learn a stable transition operator $\mathbf{f}_{\theta}:(\mathbf{z}_{t},t,r)\mapsto\mathbf{z}_{r}$, which provides a strong initialization for any-step generation. In Stage 2, starting from the converged Stage-1 checkpoint, we jointly optimize forward flow map training and DMD-based on-policy distillation. This stage uses flow map backward simulation to generate student rollouts and apply reverse-divergence supervision from a strong teacher, thereby reducing discretization error and exposure bias while preserving any-step capability.

AnyFlow for Bidirectional Video Diffusion. For bidirectional video diffusion, we simply follow the established AnyFlow pipeline, as detailed in Alg. 1 and Alg. 2.

AnyFlow for Causal Video Diffusion. For causal video diffusion, we adopt the FAR [38] training pipeline, which uses context compression with asymmetric patchify kernels to encode videos efficiently. Specifically, we keep three full-token chunks with the standard patchify kernel size of 2, while compressing the remaining chunks with a larger kernel size of 4. This design substantially reduces teacher-forcing training cost and KV-cache size during sampling. To jointly support Image-to-Video (I2V) and Text-to-Video (T2V), we use a non-uniform chunk partition: the first chunk has size 1 for precise first-frame conditioning, and subsequent chunks use size 3 to balance motion modeling and throughput. The training objective remains the same as in the bidirectional setting, and we additionally cache and reuse KV states during rollouts to improve simulation efficiency for causal video diffusion models.

Continued Training on Downstream Datasets. A practical advantage of AnyFlow is that it preserves the pretrained model’s instantaneous flow field, making the distilled model compatible with continued training on downstream datasets. As illustrated in Fig. 7, AnyFlow enables continued training from a pretrained checkpoint, unlike Self-Forcing, where further training of the distilled model is difficult.

This property is especially useful in specialized domains whose motion patterns and visual statistics differ from those of the original training corpus. As shown in Fig. 8, the base distilled model, AnyFlow-FAR-Wan2.1-1.3B, still struggles with identity preservation in robotics videos and trajectory accuracy in driving scenes. After continued training on downstream data, these errors are substantially reduced.

We implement AnyFlow on top of Wan2.1 [1] in the Diffusers framework [45]. We train the model on a synthetic dataset of 256K prompt–video pairs generated from Wan2.1-T2V-14B, where each sample contains up to 81 frames at a resolution of $480\times 832$. Training is performed in two stages. In Stage 1, we optimize the forward flow map objective using AdamW [46] with a learning rate of $5\times 10^{-5}$ and a global batch size of 32 for the 1.3B model and 16 for the 14B model, for 6,000 and 4,000 iterations, respectively. In Stage 2, we perform on-policy distillation by jointly optimizing the forward flow map and on-policy objectives with AdamW at a learning rate of $2\times 10^{-6}$ for 800 iterations. In both stages, we use parameter-efficient fine-tuning with LoRA [47] of rank 256.

For text-to-video (T2V) evaluation, we adopt VBench [48], which evaluates generation quality across 16 fine-grained dimensions grouped into two high-level categories: Quality score and Semantic score. For image-to-video (I2V) evaluation, we use VBench-I2V [49] and report both the Quality score and the I2V score.

To ensure a fair comparison, we re-evaluate key counterparts under a unified protocol using the official VBench augmented prompts. These counterparts include the Wan2.1 base model [1], Self-Forcing [10], and rCM [9], as well as community-trained 14B causal models including FastVideo [42], Krea-Realtime-14B [43], and LightX2V [41]. Results for all other methods are taken directly from their original papers.

Quantitative Comparison. As shown in Tab. 3, AnyFlow achieves strong few-step performance and continues to improve as the sampling budget increases. On the 14B bidirectional backbone, AnyFlow-Wan2.1-T2V-14B obtains 84.04 at 4 NFEs, outperforming rCM-Wan2.1-T2V-14B (83.73 at 4 NFEs). On the 14B causal backbone, AnyFlow-FAR-Wan2.1-14B reaches 84.05 at 4 NFEs and further improves to 84.41 at 32 NFEs, outperforming strong community counterparts.

For image-to-video results in Tab. 4, AnyFlow-FAR-Wan2.1-14B achieves a VBench-I2V score of 87.87 using only 4 NFEs. This result is comparable to Wan2.1-I2V-14B with 50$\times$2 NFEs (87.71) and surpasses FastVideo-CausalWan2.2-A14B-Preview (86.82), showing that AnyFlow improves both visual quality and I2V faithfulness while retaining strong sampling efficiency.

Qualitative Comparison. As shown in Fig. 9, we provide qualitative comparisons at the 14B scale for both T2V and I2V generation. In causal T2V examples, FastVideo-Wan2.2-A14B-Preview exhibits blurry details, LightX2V-Wan2.1-14B-CausVid shows flickering artifacts, and Krea-Realtime-14B produces unrealistic motion in some clips. By comparison, AnyFlow-FAR-Wan2.1-14B delivers the best visual quality and motion with 4 NFEs. In bidirectional T2V examples, AnyFlow-Wan2.1-T2V-14B shows slightly better motion than rCM-Wan2.1-T2V-14B.

For I2V generation, AnyFlow reuses the causal generator through our non-uniform chunk partition and maintains good first-frame faithfulness with smooth motion transitions. Compared with Wan2.1-I2V-14B, AnyFlow-FAR-Wan2.1-14B shows comparable temporal stability and motion quality, while FastVideo-Wan2.2-A14B-Preview exhibits visible flickering and inconsistencies with the input image. Overall, these qualitative results are consistent with the quantitative trends and further support the effectiveness of AnyFlow in few-step sampling.

Time Sampler. We investigate different choices of the time sampler, with their induced densities over $t$ visualized in Fig. 10(a). The uniform weighting yields the worst test-time scalability, as reflected by its 32-step performance in Tab. 5(a), likely because it mismatches the logit-normal time weighting used in pretraining. Among the other variants, assigning more probability mass to high-noise regions leads to better overall performance. Therefore, we fix $w(t)=\text{Beta}(2,1.5)$ in all experiments.

Interpolated Timestep Conditioning. Incorporating an additional timestep embedding during post-training is non-trivial. A baseline design is to add a new embedding projection whose final layer is zero-initialized, as in TMD [27]. However, under zero initialization, the model needs to amplify the influence of the $r$ embedding, which causes its norm to keep increasing during training and eventually leads to over-saturated results, as shown in Fig. 11. In contrast, our interpolated timestep conditioning preserves the boundary case $t=r$ for stability while avoiding an overly cold start when $t\neq r$. As a result, the embedding norm remains stable after roughly 2K training steps and stays much closer to that of the pretrained model, effectively suppressing the over-saturation effect.

Comparison to ODE-Init. We compare AnyFlow with Consistency ODE-Init for causal video generation in Fig. 12. For a fair comparison, we fine-tune the teacher on the same synthetic dataset, since we do not have access to the original training data of the Wan base model. After flow map training, AnyFlow produces clearer results than Consistency ODE-Init at both 4 and 32 NFEs, indicating that it better preserves the knowledge inherited from the flow-matching-pretrained base model. After on-policy distillation, AnyFlow further reduces test-time errors in few-step sampling and improves the overall sampling trajectory. We also observe that our method preserves the style and content of the pretrained model more faithfully.

Training Cost Breakdown. Tab. 5 reports the training cost of each component in AnyFlow, measured in seconds per iteration on 8$\times$H100 GPUs. In the forward training stage, guidance-fused training and the differential derivation equation introduce additional overhead compared with standard flow matching, but the overall cost remains practical for large-scale training. In the on-policy distillation stage, our flow map backward simulation incurs a slightly higher cost than consistency backward simulation at 4 steps, specifically 15.7% more for the causal model and 22.5% more for the bidirectional model, because we backpropagate through the full rollout chain. However, our method becomes significantly more efficient when simulating larger step counts. At 16 steps, it reduces the training cost by 43.4% for the causal model and 47.0% for the bidirectional model compared with consistency backward simulation, thanks to the shortcut transitions learned by the flow map.