1 of 20

Dogyun Park, Taehoon Lee, Minseok Joo, Hyunwoo J. Kim

Blockwise Flow Matching: Improving Flow Matching Models For Efficient High-Quality Generation

1

Korea University

MLV Lab

NeurIPS 2024

2 of 20

Diffusion/Flow Matching Model: State-of-the-art Generation

Diffusion/Flow Matching Models are state-of-the-art framework for image/video generation.

Images generated by nano-banana (https://gemini.google/overview/image-generation/)

MLV Lab

Korea University

MLV Lab

NeurIPS 2024

3 of 20

Diffusion/Flow Matching Model

Single model for all timesteps

Common design choice: Single model (e.g., Diffusion Transformers) for all timesteps

Problem:

  • Single model has to learn very different charateristics aross time.
  • Expensive and slow generation.
  • High noise level is irregular and low-frequency / low noise level is regular and high-frequency
  • e.g.) For sampling step T with model size M, the total inference complexity is O(MT)

MLV Lab

Korea University

MLV Lab

NeurIPS 2024

4 of 20

Diffusion/Flow Matching Model

Single model for all timesteps

Common design choice: Single model (e.g., Diffusion Transformers) for all timesteps

Problem:

  • Single model has to learn very different charateristics aross time.
  • Expensive and slow generation.
  • High noisy level is irregular and low-frequency, low noise level is regular and high-frequency
  • e.g.) For sampling step T with model size M, the total inference complexity is O(MT)

MLV Lab

Korea University

MLV Lab

NeurIPS 2024

5 of 20

Divide timesteps, conquer them

with Smaller, Specialized networks

Our core idea: Blockwise Flow Matching

MLV Lab

Korea University

MLV Lab

NeurIPS 2024

6 of 20

🚀 Our core Idea:

Blockwise Flow Matching

  • Divide timestep into multiple segments

 

Standard Flow Matching

 

Single Velocity Model

 

 

Velocity

Block 1

Velocity

Block 2

Velocity

Block 3

Velocity

Block 4

Velocity

Block 5

 

 

Blockwise Flow Matching (Ours)

Smaller &�Specialized

  • Assign smaller but specialized networks to each segment

Advantages:

  • Specialized over time
  • Each model focuses on a narrower regime, leading to more accurate prediction
  • Faster and cheaper sampling
  • At each step, only lightweight model runs, so inference is faster and less expensive

MLV Lab

Korea University

MLV Lab

NeurIPS 2024

7 of 20

🚀 Our core Idea:

Blockwise Flow Matching

  • Divide timestep into multiple segments
  • Assign smaller but specialized networks to each segment

Advantages:

  • Specialized over time
  • Each model focuses on a narrower regime, leading to more accurate prediction
  • Faster and cheaper sampling
  • At each step, only lightweight model runs, so inference is faster and less expensive

 

Standard Flow Matching

 

Single Velocity Model

 

 

Velocity

Block 1

Velocity

Block 2

Velocity

Block 3

Velocity

Block 4

Velocity

Block 5

 

 

Blockwise Flow Matching (Ours)

Smaller &�Specialized

MLV Lab

Korea University

MLV Lab

NeurIPS 2024

8 of 20

🚀 Our core Idea:

Blockwise Flow Matching

  • Divide timestep into multiple segments
  • Assign smaller but specialized networks to each segment

Advantages:

  • Specialized over time
  • Each model focuses on a narrower regime, leading to more accurate prediction
  • Faster and cheaper sampling
  • At each step, only lightweight model runs, so inference is faster and less expensive

 

Standard Flow Matching

 

Single Velocity Model

 

 

Velocity

Block 1

Velocity

Block 2

Velocity

Block 3

Velocity

Block 4

Velocity

Block 5

 

 

Blockwise Flow Matching (Ours)

Smaller &�Specialized

MLV Lab

Korea University

MLV Lab

NeurIPS 2024

9 of 20

Each specialized network = shared Feature Alignment Network + specialized Velocity Network

Architecture Design: Semantic Feature Guidance (SemFeat)

 

 

MLV Lab

Korea University

MLV Lab

NeurIPS 2024

10 of 20

SemFeat Improves Generation Quality

RDG [1] and REPA [2] show improving semantic representation improves the generation quality

  • Our SemFeat also significantly improves the generation quality.
  • REPA aligns internal hidden states of velocity model
  • SemFeat introduces dedicated network for alignment, whose output is explicitly conditioned on velocity model

[1] Li, Tianhong, Dina Katabi, and Kaiming He. "Return of unconditional generation: A self-supervised representation generation method.”, NeurIPS 2024

[2] Yu, Sihyun, et al. "Representation alignment for generation: Training diffusion transformers is easier than you think.”, ICLR 2025

PCA Visualization over timestep

MLV Lab

Korea University

MLV Lab

NeurIPS 2024

11 of 20

SemFeat Improves Generation Quality

  • Our SemFeat also significantly improves the generation quality.
  • REPA aligns internal hidden states of velocity model
  • SemFeat introduces dedicated network for alignment, whose output is explicitly conditioned on velocity model

[1] Li, Tianhong, Dina Katabi, and Kaiming He. "Return of unconditional generation: A self-supervised representation generation method.”, NeurIPS 2024

[2] Yu, Sihyun, et al. "Representation alignment for generation: Training diffusion transformers is easier than you think.”, ICLR 2025

PCA Visualization of aligned feature over timestep

RDG [1] and REPA [2] show improving semantic representation improves the generation quality

MLV Lab

Korea University

MLV Lab

NeurIPS 2024

12 of 20

 

Observation from Feature Alignment Network

  • Outputs within segments show high similarity

 

Lightweight Feature Residual Network can replace feature alignment network

Start point

Intermediate point

MLV Lab

Korea University

MLV Lab

NeurIPS 2024

13 of 20

Observation from feature alignment network output:

Observation from Feature Alignment Network

  • Outputs within segments show high similarity

 

 

 

Feature Residual Network

Residual Approximation

Lightweight Feature Residual Network can replace feature alignment network

MLV Lab

Korea University

MLV Lab

NeurIPS 2024

14 of 20

Faster inference with Feature Residual Network

During inference, lightweight Feature Residual Network can replace feature alignment network for intermediate points

Velocity model

Feature Alignment Network

Velocity model

Feature Alignment Network

Feature

Residual Network

MLV Lab

Korea University

MLV Lab

NeurIPS 2024

15 of 20

For same training iterations:

Analysis: Network Specialization

  • Achieves a lower loss than SiT for each temporal segments
  • Achieves a smaller Fréchet distance from real images in terms of spectral fidelity than SiT, indicating that BFM is more faithful to the natural image frequencies

Timestep

Loss

BFM shows 3% lower loss than SiT on average!

Fourier power spectrum analysis

MLV Lab

Korea University

MLV Lab

NeurIPS 2024

16 of 20

For same training iterations:

Analysis: Network Specialization

  • Achieves a lower loss than SiT for each temporal segments
  • Achieves a smaller Fréchet distance from real images in terms of spectral fidelity than SiT, indicating that BFM is more faithful to the natural image frequencies

Timestep

Loss

BFM shows 3% lower loss than SiT on average!

Fourier power spectrum analysis

MLV Lab

Korea University

MLV Lab

NeurIPS 2024

17 of 20

For same sampling steps:

Analysis: Inference Efficiency

  • Achieves 1.1x (BFM-SF) and 2.3x (BFM-SF-RA) faster inference latency in seconds
  • Achieves more favorable Pareto-frontier between inference complexity and performance

MLV Lab

Korea University

MLV Lab

NeurIPS 2024

18 of 20

For same sampling steps:

Analysis: Inference Efficiency

  • Achieves 1.1x (BFM-SF) and 2.3x (BFM-SF-RA) faster inference latency in seconds
  • Achieves more favorable Pareto-frontier between inference complexity and performance

MLV Lab

Korea University

MLV Lab

NeurIPS 2024

19 of 20

Results on ImageNet 256x256

Comparison to state-of-the-art models

  • BFM-SF achieves 1.75 FID with lower inference FLOPs compared to SiT and REPA
  • BFM-SF-RA significantly reduces inference GFLOPs to 37.8 without degrading performance

MLV Lab

Korea University

MLV Lab

NeurIPS 2024

20 of 20

Conclusion

  • Diffusion/flow matching model with Blockwise Flow Modeling can improve the inference efficiency and generation quality by network specialization
  • Architecture design such as Semantic Feature Guidance and Feature Residual Network can further push the performance and efficiency

Github

Paper

Project page

MLV Lab

Korea University

MLV Lab

NeurIPS 2024