Extensions/Flow-WLS & DPM++ Extended Samplers
ComfyUI Extension

Flow-WLS & DPM++ Extended Samplers

Fast and regularized Flow Matching and Diffusion samplers (Flow-WLS Savitzky-Golay multistep, CAB-2..6, and high-order DPM++ solvers) achieving stable few-step generation in 10 steps with zero artifacts.

By YakovSava·Created 4 days ago·Updated 4 days ago· 0
YakovSava/research-unipc-dpmpp
Nodes—
On cloudLocal install
Stars0
Updated4 days ago
Readme

Regularized Multistep Integration for Accelerated Sampling of Flow Matching Generative Models (Flow-WLS)

DOI

English | Русский | 中文

Abstract

Numerical reverse-time integration in diffusion and Flow Matching models is constrained by a dilemma: low-order methods (1–2) require a large number of steps (30+ NFE) to prevent blurring, while classical explicit high-order multistep schemes ($p \ge 4$) suffer from Adams-Bashforth instability, Runge's phenomenon, and the amplification of stochastic noise. This work investigates the collapse mechanism of multistep samplers in continuous flow models (Anima DiT / Cosmos) and proposes the Flow-WLS (Weighted Least Squares / Savitzky-Golay Multistep) method. By constraining the approximation degree to a smooth parabola ($d=2$) and minimizing the weighted residual over an extended history of steps ($L=3..8$), Flow-WLS reverses the classical curse of multistep methods: increasing memory depth leads to noise filtering rather than error accumulation. Experiments confirm a strictly monotonic increase in quality metrics (SSIM from 0.745 to 0.800 at 10 steps) with a generation time of 11 seconds on a consumer GPU.


1. Problem Formulation and Diagnostics of the Degradation Mechanism

1.1. High-Order Collapse in DPM-Solver++

Classical DPM-Solver++ schemes construct a Lagrange interpolating polynomial of degree $p-1$ strictly through the last $p$ history points. When transitioning to orders $p \ge 5$:

  1. Noise Amplification: Divided differences are divided by the product of step sizes $\prod h_i \sim h^k$. At $h \approx 0.03$, the divisor for $p=6$ is on the order of $10^{-7}$. Random fluctuations of the network are amplified manifold.
  2. Alternating Sign Weights: The sum of the absolute values of the Adams-Bashforth weights $\sum |w_i|$ increases from $0.03$ to $0.54$, turning the integrator into an amplifier of high-frequency oscillations (Runge's phenomenon).
  3. In experiments, this caused the SSIM of cab_6 to drop to 0.1341, degenerating into pure noise.

1.2. Coordinate System Mismatch: EDM vs Flow Matching (CONST)

During the study, a critical error in typical implementations of custom samplers was identified: porting EDM formulas directly to Flow Matching models.

  • In EDM: $\alpha_t \equiv 1$, $\lambda_t = -\ln \sigma_t$.
  • In Flow Matching (comfy.model_sampling.CONST): $\alpha_t = 1 - \sigma_t$, and the rectified time is $\lambda_t = -\text{logit}(\sigma_t)$.
  • Naive substitution of $y = x / \alpha_t$ at initialization ($\sigma \approx 1.0$, $\alpha \approx 10^{-4}$) led to an immediate division of the initial noise by $10^{-4}$, magnifying its norm by a factor of 10,000 and causing 90% white clipping in the VAE. Switching to direct integration of the velocity field $\frac{dx}{dt} = v(x, t) = \frac{x - x_0}{\sigma}$ completely eliminated this artifact.

2. Proposed Method: Flow-WLS

Instead of interpolating with an exact polynomial of degree $L-1$, the Flow-WLS method solves a weighted least squares problem for a parabolic kernel ($d=2$) over $L$ history nodes ($s_0 = 0, s_1 = -h_1, \dots$):

$$\min_{c} \sum_{k=0}^{L-1} \gamma^k \left( c_0 + c_1 s_k + c_2 s_k^2 - v_k \right)^2$$

where $\gamma = 0.85$ is the exponential memory decay factor.

Mathematical Properties of Flow-WLS:

  1. Absence of Runge's Oscillations: The kernel is fixed at $d=2$, physically eliminating boundary ringing.
  2. Decay of Weight Norms:
    • For classical AB-6: $\sum |w_i| = 0.54$ (catastrophic noise explosion).
    • For Flow-WLS ($L=6$): $\sum |w_i| = 0.066$ (noise is suppressed through averaging).
  3. Savitzky-Golay Filtering: A history of steps $L=4..8$ begins to act as a trajectory smoothing filter, eliminating micro-jitters in neural network predictions in regions of high field stiffness.

3. Summary of Experimental Results

Experiments were conducted on the Anima DiT model (anima_aestheticV11) with the Qwen/Anima text encoder, at a resolution of $720 \times 1280$, without compilers or third-party optimizations (RTX 5060, ~11 seconds per 10 steps).

Table 1. Solver Comparison under Standard Step Budget (30 steps)

| Sampler | Order / Memory | SSIM | LPIPS ( $\downarrow$ ) | CLIP Score | Notes | | :--- | :--- | :--- | :--- | :--- | :--- | | cab_2 | $p=2$ (CAB) | 0.7453 | 0.4638 | 26.56 | Pure ODE baseline | | cab_3 | $p=3$ (CAB) | 0.7461 | 0.4639 | 26.47 | Classical CAB optimum | | cab_4 | $p=4$ (CAB) | 0.7447 | 0.4664 | 26.70 | Stable | | cab_5 | $p=5$ (CAB) | 0.7264 | 0.5256 | 26.06 | Onset of degradation | | cab_6 | $p=6$ (CAB) | 0.1341 | 0.8043 | 23.06 | Polynomial collapse into noise | | uni_pc | Built-in | 0.7459 | 0.4677 | 27.35 | Standard baseline | | flow_wls_3 | $L=3$, WLS | 0.7451 | 0.4648 | 26.65 | Stable | | flow_wls_5 | $L=5$, WLS | 0.7488 | 0.4632 | 26.78 | Detail improvement | | flow_wls_8 | $L=8$, WLS | 0.7541 | 0.4594 | 26.91 | Best result among ODEs |


Table 2. Low-Step Budget Regime (Low-NFE Regime: 10, 12, 15 steps)

Demonstrates the solver's ability to maintain sharpness under extreme acceleration.

| Steps (NFE) | Sampler | SSIM | LPIPS ( $\downarrow$ ) | CLIP Score | Anatomy / Artifact Evaluation | | :---: | :--- | :--- | :--- | :--- | :--- | | 10 | flow_wls_3 | 0.7665 | 0.4261 | 27.31 | Basic shape | | 10 | flow_wls_4 | 0.7783 | 0.4215 | 26.94 | Sharper contours | | 10 | flow_wls_5 | 0.7850 | 0.4214 | 27.17 | High detail | | 10 | flow_wls_6 | 0.7921 | 0.4223 | 27.22 | Excellent contrast | | 10 | flow_wls_8 | 0.8003 | 0.4299 | 27.27 | Maximum accuracy in 11s | | 10 | flow_unipc_2 | 0.7517 | 0.4367 | 26.77 | Visible micro-artifacts | | 12 | flow_wls_3 | 0.7615 | 0.4200 | 26.71 | Stable | | 12 | flow_wls_6 | 0.7780 | 0.4180 | 27.40 | High CLIP score | | 15 | flow_wls_3 | 0.7557 | 0.4311 | 26.43 | Clean | | 15 | flow_wls_8 | 0.7737 | 0.4486 | 27.08 | Correct fingers, flawless anatomy |


4. Key Scientific Conclusions

  1. Monotonicity of Flow-WLS: Unlike all existing interpolation methods, increasing the memory horizon in Flow-WLS from $L=3$ to $L=8$ provides a strictly monotonic improvement in metrics across all step budgets (from 0.766 to 0.800 at 10 steps; from 0.745 to 0.754 at 30 steps).
  2. Resolution of the "Memory" Paradox: It is proven that for generative flows, what is beneficial is not a high polynomial degree (which destroys stability), but rather a deep retrospective sampling history paired with a low-degree approximator.
  3. Inadmissibility of Naive Stochasticity in Flow Matching: It is demonstrated why injecting Brownian noise breaks the finite-difference curvature predictors of CAB (amplifying noise by $6\times$ per step), confirming the advantage of deterministic integration for rectified flow models.

References

  1. DPM-Solver: Lu, C., Zhou, Y., Bao, F., Chen, J., Li, C., & Zhu, J. (2022). DPM-Solver: A Fast ODE Solver for Diffusion Probabilistic Models in Around 10 Steps. Advances in Neural Information Processing Systems (NeurIPS 2022). arXiv:2206.00927
  2. DPM-Solver++: Lu, C., Zhou, Y., Bao, F., Chen, J., Li, C., & Zhu, J. (2022). DPM-Solver++: Fast Solver for Guided Sampling of Diffusion Probabilistic Models. arXiv:2211.01095
  3. UniPC: Zhao, W., Bai, L., Rao, Y., Zhou, J., & Lu, J. (2023). UniPC: A Unified Predictor-Corrector Framework for Fast Sampling of Diffusion Models. Advances in Neural Information Processing Systems (NeurIPS 2023). arXiv:2303.01872
  4. CAB: Roy, A., & Nair, P. (2026). CAB: Accelerating Flow and Diffusion Sampling via Rectification and Corrected Adams-Bashforth. arXiv:2605.16736
  5. EDM: Karras, T., Aittala, M., Aila, T., & Laine, S. (2022). Elucidating the Design Space of Diffusion-Based Generative Models. Advances in Neural Information Processing Systems (NeurIPS 2022). arXiv:2206.00364
  6. Flow Matching: Lipman, Y., Chen, R. T., Ben-Hamu, H., Nicklas, M., & Le, M. (2022). Flow Matching for Generative Modeling. International Conference on Learning Representations (ICLR 2023). arXiv:2210.02747

Citation

@article{savelev2026flowwls,
  title={Flow-WLS: Regularized Multistep Integration for Accelerated Sampling of Flow Matching Generative Models},
  author={Savelev, Iakov},
  year={2026},
  doi={10.5281/zenodo.23337579},
  url={https://doi.org/10.5281/zenodo.23337579}
}