TBQ:Trajectory-Aligned Branch-Aware Quantization for Few-Step MoE-Like Video Diffusion Models

Distillation objective.

Few-step distillation accelerates diffusion sampling by learning a student trajectory with fewer denoising evaluations, rather than simply truncating the teacher sampler. This follows a line of fast diffusion sampling methods, including progressive distillation and consistency models [28, 32]. Let C_step denote the cost of one denoising step and let p_S(x;θ) denote the student generation distribution. We evaluate 4-, 8-, 20-, and 40-step settings, using the 40-step setting as a long-budget reference.

Unlike truncation, distillation optimizes the intermediate states used by the reduced-step sampler. Performing distillation before quantization also provides a better calibration target: the quantizer can be fitted to the activation distribution of the deployed student trajectory, rather than to the long-step teacher trajectory.

Stage-aware distribution matching.

DMD aligns the student with the teacher distribution rather than enforcing a fixed one-to-one correspondence between teacher and student sampling paths [43]. This distinction is important for step selection. A short-budget DMD student is not merely memorizing a small set of denoising states; it learns a compressed sampling behavior that may still benefit from a moderately larger inference budget. Evaluating the same distilled model with more sampling steps can therefore refine the trajectory on a denser grid and recover quality lost at the shortest budget.

For MoE-like video diffusion model, this trajectory should also respect the high-noise and low-noise phases. The early phase establishes layout, subject positions, and motion, whereas the late phase refines texture, boundary quality, and local temporal coherence. Phased DMD extends distribution matching by matching score behavior within subintervals and is therefore well aligned with MoE-style video denoising [10]. In our pipeline, the high-noise and low-noise phases are used as routing-aware intervals for later calibration, rather than as an additional hand-tuned distillation loss. This keeps the distillation objective unchanged while preserving the stage distinction required by branch-aware quantization.

For deployment, the distilled student is treated as a new inference model with its own numerical behavior. A shorter trajectory changes which denoising states are visited, how long each expert remains active, and how activation magnitudes evolve. We therefore compare multiple step counts instead of assuming that the fewest steps are always optimal: 4 steps minimize latency, 40 steps expose long-budget behavior, and intermediate budgets test whether additional denoising evaluations improve the distilled sampling trajectory without excessive quantization-error accumulation.

Co-design with quantization.

At a given denoising time t, the activation distribution of the distilled student generally differs from that of the original high-step reference:

This co-design offers two direct benefits. First, few-step distillation reduces the number of quantized denoising evaluations, so the total inference cost can be approximated as

Thus, distillation preserves the teacher distribution under fewer denoising states, while branch-aware quantization preserves the student's phase-specific activations under those states. This order is more stable than quantizing the original long-step model first, because it avoids calibrating a distribution that will not be used after acceleration.

HiF4 is a 4-bit block floating-point format for low-bit inference [24]. Each block stores 64 signed 4-bit values and 32 bits of shared scaling metadata, giving an average cost of 4.5 bits per value. The metadata uses three scaling levels:

A diagram of the HiF4 block floating-point format, including the E6M2 base scale, binary micro-exponents, and signed low-bit values.

Global and local scaling. The first level stores an E6M2 base scale,

Format motivation. We choose HiF4 for its range behavior, not only its bitwidth. Early denoising requires coarse high-noise activations, while late denoising preserves fine details and small temporal corrections. We therefore use HiF4 as a specified wide-range low-bit format and focus our experiments on trajectory alignment, branch-aware calibration, and step selection rather than alternative 4-bit formats.

TBQ consists of five stages: distillation, trajectory-aligned calibration, branch-wise low-bit replacement, generation, and closed-loop evaluation. First, the original model is adapted to a short denoising trajectory through distribution-matching distillation. Trajectory-aligned calibration is then performed under the same deployment schedule, allowing the quantizer to observe the student activations encountered at inference time rather than the teacher activations before acceleration. Finally, the full-precision and compressed models generate videos using matched prompts and sampling settings. The automatic evaluation stream computes VBench scores under a unified protocol.

Two principles guide the implementation. First, calibration and replacement must follow the true high-noise and low-noise expert routes of MoE-like video diffusion model instead of treating the model as a dense trunk. This makes MoE routing a quantization boundary: statistics collected for global high-noise denoising are not pooled with statistics collected for late low-noise refinement. Second, step schedules must be compared under matched generation conditions. Although a 4-step setting minimizes the number of evaluations, moderately longer schedules such as 8 or 20 steps can reduce residual sampling error while keeping the number of quantized evaluations far below that of a conventional long sampler. The pipeline is therefore both a compression method and a step-selection procedure.

Distillation objective.

Few-step distillation accelerates diffusion sampling by learning a student trajectory with fewer denoising evaluations, rather than simply truncating the teacher sampler. This follows a line of fast diffusion sampling methods, including progressive distillation and consistency models [28, 32]. Let C_step denote the cost of one denoising step and let p_S(x;θ) denote the student generation distribution. We evaluate 4-, 8-, 20-, and 40-step settings, using the 40-step setting as a long-budget reference.

Unlike truncation, distillation optimizes the intermediate states used by the reduced-step sampler. Performing distillation before quantization also provides a better calibration target: the quantizer can be fitted to the activation distribution of the deployed student trajectory, rather than to the long-step teacher trajectory.

Stage-aware distribution matching.

DMD aligns the student with the teacher distribution rather than enforcing a fixed one-to-one correspondence between teacher and student sampling paths [43]. This distinction is important for step selection. A short-budget DMD student is not merely memorizing a small set of denoising states; it learns a compressed sampling behavior that may still benefit from a moderately larger inference budget. Evaluating the same distilled model with more sampling steps can therefore refine the trajectory on a denser grid and recover quality lost at the shortest budget.

For MoE-like video diffusion model, this trajectory should also respect the high-noise and low-noise phases. The early phase establishes layout, subject positions, and motion, whereas the late phase refines texture, boundary quality, and local temporal coherence. Phased DMD extends distribution matching by matching score behavior within subintervals and is therefore well aligned with MoE-style video denoising [10]. In our pipeline, the high-noise and low-noise phases are used as routing-aware intervals for later calibration, rather than as an additional hand-tuned distillation loss. This keeps the distillation objective unchanged while preserving the stage distinction required by branch-aware quantization.

For deployment, the distilled student is treated as a new inference model with its own numerical behavior. A shorter trajectory changes which denoising states are visited, how long each expert remains active, and how activation magnitudes evolve. We therefore compare multiple step counts instead of assuming that the fewest steps are always optimal: 4 steps minimize latency, 40 steps expose long-budget behavior, and intermediate budgets test whether additional denoising evaluations improve the distilled sampling trajectory without excessive quantization-error accumulation.

Co-design with quantization.

At a given denoising time t, the activation distribution of the distilled student generally differs from that of the original high-step reference:

This co-design offers two direct benefits. First, few-step distillation reduces the number of quantized denoising evaluations, so the total inference cost can be approximated as

Thus, distillation preserves the teacher distribution under fewer denoising states, while branch-aware quantization preserves the student's phase-specific activations under those states. This order is more stable than quantizing the original long-step model first, because it avoids calibrating a distribution that will not be used after acceleration.

Quantization method selection and formulation.

We use a rounding-based PTQ method and adapt it through granularity design, branch-wise calibration, and sensitive-layer protection. Prior PTQ methods such as GPTQ and SmoothQuant show that calibration design is critical for large transformer-style models [12, 39]; here, we extend the same principle to video diffusion.

For a tensor x and scale s, the quantized reconstruction can be written as

Branch-aware adaptation for dual-expert denoising.

The dual-expert routing at time step t can be written as

Concretely, let $\mathcal {D}_{\mathrm{cal}}$ be the calibration prompt set, and let $\mathcal {A}^{h}_{\ell }$ and $\mathcal {A}^{l}_{\ell }$ denote the activation sets captured for layer ℓ in the high-noise and low-noise experts, respectively. The calibration objective is computed on branch-specific statistics rather than on a pooled activation set:

Entrance-layer protection and enhanced calibration.

The protected prefix controls where low-bit error first enters the latent trajectory. Because early errors are repeatedly propagated through denoising, entrance modules are kept in higher precision or calibrated more aggressively. Let $\mathcal {K}$ denote the protected set. We define $\mathcal {K}$ using a rule-based prefix criterion: modules before the first quantized transformer block on each active expert route are protected, including the latent/input projection, timestep and text-conditioning projections, and the first denoising block of the high-noise and low-noise experts. All remaining transformer blocks are quantized layer by layer under the branch-specific calibration procedure. The effective quantized model is

Evaluation metrics. We evaluate five VBench dimensions [16]: Subject Consistency (SC), Aesthetic Quality (AQ), Imaging Quality (IQ), Overall Consistency (OC), and Motion Smoothness (MS). SC uses DINO features [2]; AQ uses CLIP and a LAION aesthetic predictor [26, 29]; IQ uses MUSIQ [17]; OC uses ViCLIP/InternVideo [36]; and MS uses AMT-based interpolation [20]. These dimensions separate common compression failures—subject drift, frame artifacts, semantic loss, and temporal jitter—and we therefore report both aggregate and per-metric scores.

Implementation protocol. We instantiate TBQ on Wan2.2-T2V-A14B, a video diffusion model with separate high-noise and low-noise expert branches. All experiments run on one 80 GB NVIDIA H100 GPU. We compare three variants: the original full-precision Wan2.2-T2V-A14B model, its few-step distilled full-precision student, and the corresponding TBQ quantized model. Sampling budgets are 4, 8, 20, and 40 denoising steps. All variants share VBench prompt indices, generation settings, output naming, and evaluation dimensions. The original-versus-student comparison isolates the effect of distillation, while the student-versus-TBQ comparison isolates the effect of low-bit quantization. The workflow follows VBench prompt metadata, organizes generated videos by dimension, and includes all successfully generated videos without manual selection.

Memory-aware execution. TBQ performs calibration and module replacement layer by layer to avoid out-of-memory errors. At each stage, only the current block and required temporary buffers remain on the GPU; auxiliary distillation modules are moved to the GPU only when needed. The same procedure is applied independently to the high-noise and low-noise branches of Wan2.2-T2V-A14B. For each block, inputs are captured, the block is quantized and validated, temporary buffers are released, and processing continues to the next block. This is slower than fully resident calibration but enables reliable single-H100 execution. Verification uses batch generation rather than selected examples: the distilled full-precision and TBQ models share prompts, resolutions, output formats, and sampling schedules so that quality changes can be attributed to the selected VBench dimensions.

After a few-step distillation, the number of denoising steps jointly governs inference efficiency and the extent of quantization-error accumulation. Table 2 compares the original full-precision model, the distilled full-precision student evaluated at the same sampling budget, and the TBQ quantized model across four denoising-step configurations. The 40-step setting is included as a long-budget reference.

Effect of sampling budget. At four steps, TBQ marginally outperforms the same-step full-precision student, though both fall below the original full-precision baseline. At eight and 20 steps, TBQ surpasses both the original baseline and the corresponding full-precision student. At 40 steps, TBQ still exceeds the original baseline but falls 0.0007 below the same-step full-precision model.

Performance improves substantially from four to 20 steps and largely saturates beyond that. The marginal gap at 40 steps is consistent with the hypothesis that low-bit perturbations accumulate along longer sampling trajectories, as each additional denoising evaluation introduces an independent quantization error. We therefore select 20 steps as the representative operating point, which captures most of the quality gain from longer schedules while requiring only half the denoising evaluations of the 40-step configuration.

The same-step full-precision student provides the most direct control for isolating the effect of quantization, since the distilled model and sampling trajectory are held fixed. Across all moderate budgets (8 and 20 steps), TBQ introduces no measurable quality degradation. The small improvements observed after quantization are attributable to a combination of branch-aware calibration, quantization-induced regularization, and evaluation variance.

Per-metric analysis at 20 steps. Table 3 provides a per-metric breakdown at the selected 20-step operating point. TBQ obtains an aggregate score of 0.7074, compared with 0.7051 for the same-step full-precision student and 0.6976 for the original full-precision model.

Relative to the original full-precision model, TBQ improves Subject Consistency (SC), Aesthetic Quality (AQ), Imaging Quality (IQ), and Motion Smoothness (MS). Overall Consistency (OC) is slightly lower than the original but remains comparable to the same-step full-precision student, indicating that few-step distillation primarily enhances spatial and temporal quality while preserving semantic alignment.

Compared with the same-step full-precision model, Tab obtains consistent improvements across all five reported dimensions: SC increases from 0.9597 to 0.9629, IQ from 0.6828 to 0.6872, and MS from 0.9821 to 0.9825. These results confirm that TBQ's branch-aware calibration preserves the spatial and temporal characteristics of the distilled model without introducing measurable degradation under the representative sampling budget.

Comparison with long-video generation methods. We further evaluate the 20-step TBQ model against recent long-video generation methods on a common set of VBench dimensions. To ensure a fair comparison across methods with partially overlapping evaluation protocols, Table 4 focuses on four metrics shared across all compared systems—Subject Consistency, Motion Smoothness, Aesthetic Quality, and Imaging Quality—with Common Average defined as their arithmetic mean.

TBQ achieves the highest Common Average (82.24) among all reported methods, improving over LongLive-2.0 (NVFP4) (79.13) by 3.11 points despite employing low-bit weights and activations. The gain is primarily driven by substantially higher AQ and IQ, while SC and MS remain competitive.

This comparison should be interpreted as a shared-dimension reference rather than a strictly controlled benchmark. The external results are generated under the VBench-Long protocol for 60-second videos, whereas TBQ is evaluated under the standard VBench setting with 20 denoising steps. Differences in video duration, base model, and prompt distribution preclude a direct end-to-end comparison. Nevertheless, the shared-metric results provide complementary evidence that TBQ maintains strong spatial and temporal quality at the 20-step operating point, corroborating the controlled comparisons in Tables 2 and 3.

We estimate the theoretical compute benefit of HiF4 under a representative generation setting with a spatial resolution of 480 × 848, 81 frames, and four denoising steps. Following the standard convention, one multiply–accumulate operation is counted as two floating-point operations. Table 5 reports the dominant matrix-computation cost of the DiT backbone.

Quantization does not reduce the strict FLOP count, since the network architecture, token sequence length, and sampling schedule remain unchanged. We therefore estimate the potential benefit using a bitwidth-normalized equivalent cost. In our configuration, all Linear layers are quantized to W4A4 HiF4 except those in the first Transformer block of each expert, while attention matrix multiplications remain in BF16.

Under this approximation, the normalized cost of the Linear layers decreases from 6.403 to 0.550, resulting in a total equivalent cost of 7.969. Relative to the 13.822 PFLOPs required by the W16A16 baseline, this corresponds to an idealized matrix-compute speedup of

This value should be interpreted as an upper bound on the matrix-computation benefit rather than an end-to-end latency improvement. Our evaluation stack does not provide native HiF4 matrix kernels, and currently available HiF4 implementations primarily support numerical simulation rather than fully optimized W4A4 DiT inference. We therefore report theoretical normalized cost instead of measured runtime.

In practice, the achievable speedup is further affected by operations not captured by the matrix-compute estimate, including softmax, normalization, scheduler execution, VAE decoding, dequantization, memory movement, and kernel-launch overhead.