Lihan Lian

Lihan Lian

Robotics enthusiast.

Understanding Score-Based Diffusion and Flow Matching with PyTorch

1 minute read

Published:

Building on VAE and DDPM, this post explores learning scores and velocities to transform noise into data. The learning path moves from intuition to small PyTorch implementations, then compares the approaches.

Prerequisites: DDPM training and sampling, Gaussian distributions, gradients, and basic PyTorch. Review ordinary and stochastic differential equations (ODEs and SDEs) as needed.

Part I: Score-Based Diffusion

1. From DDPM Noise Prediction to Scores

  • Introduce the score of a probability density and its geometric meaning.
  • Connect DDPM noise prediction to estimating scores at different noise levels.

2. Learning Scores with Denoising Score Matching

  • Outline noise perturbations, time conditioning, training targets, and loss weighting.
  • Identify how these choices map to a PyTorch training loop.

3. Continuous-Time Diffusion and Sampling

  • Introduce forward SDEs and reverse-time generation.
  • Compare reverse-SDE sampling with the probability flow ODE.
  • State the time convention and distinguish shared marginal distributions from individual trajectories.

Part II: Flow Matching

4. From Scores to Velocity Fields

  • Use the probability flow ODE as the bridge to learning a velocity field directly.
  • Introduce probability paths and distinguish the training objective from the sampling procedure.

5. Conditional Flow Matching

  • Begin with independent Gaussian noise and data pairs, linear interpolation, and velocity regression.
  • Outline the training loop and ODE sampling.
  • Distinguish conditional training paths from learned sampling trajectories; place rectified flow and optimal transport as further study.

Part III: Implementation and Comparison

6. A Minimal 2D PyTorch Project

  • Train a score model and a flow-matching model on two moons or Gaussian mixtures.
  • Visualize learned vector fields, intermediate distributions, and sample trajectories.
  • Implement reverse-SDE, probability-flow-ODE, and flow-matching-ODE samplers.

7. Extending to Cartoon Images

  • Reuse the CartoonSet pipeline and adapt a time-conditioned U-Net.
  • Compare sample quality and diversity, runtime, and network function evaluations.
  • Keep model capacity and training budgets comparable; vary sampling steps.

8. Connections and Lessons Learned

  • Revisit the Gaussian diffusion–flow-matching connection using Diffusion Meets Flow Matching.
  • Separate the effects of paths, prediction targets, loss weighting, and numerical solvers.
  • Record implementation pitfalls, experimental observations, and open questions.

References

You May Also Enjoy

Implementation Details of Cartoon-VAE-DDPM Permalink

10 minute read

Published:

Variational Autoencoder (VAE) learns to compress data into a latent Gaussian space and reconstruct it in a single shot. Denoising Diffusion Probabilistic Model (DDPM) tackles the same evidence-lower-bound objective from another direction: it begins with pure noise and iteratively denoise through hundreds of steps, exchanging speed for high-fidelity, stable synthesis. Both frameworks connect random noise to data, yet VAE rely on an explicit encoder–decoder pair, whereas DDPM use a learned Markov chain that inverts a forward noising process. This blog traces the progression from VAE to DDPM, clarifying their shared principles, with code examples available at this repository.

From Bigram to Transformer Permalink

29 minute read

Published:

Character-level name generation is a simple but effective task for understanding modern generative sequence models. This blog explains a PyTorch-based implementation that progressively builds name generators from a count-based Bigram model to Neural Bigram, MLP, RNN, GRU, LSTM, and a tiny decoder-only Transformer. Inspired by Karpathy’s makemore, the project uses a names dataset where each model learns autoregressive next-character prediction with a special start/end token. This blog post explains concepts includes embeddings, hidden states, gated memory, and causal self-attention. Implementation can be found at repository.

Implementation Details of U-Net-Segmentation-Jax Permalink

22 minute read

Published:

Image segmentation is a fundamental computer vision task that assigns a class label to each pixel, enabling models to separate objects from their backgrounds with pixel-level precision. This blog provides the explanation of a JAX-based U-Net implementation for pet image segmentation task using the Oxford-IIIT Pet dataset. The project starts by building a U-Net from scratch, also covering evaluation metrics such as Dice score and Intersection over Union (IoU). It also implements ResNet-based transfer learning as an U-Net encoder alternative to compare with a fully trained-from-scratch U-Net segmentation model. Corresponding repository can be found here repository.

Interpreting LQR through Optimal Control and Reinforcement Learning Permalink

67 minute read

Published:

Linear Quadratic Regulator (LQR) is one the simplest yet most classical probelms in optimal control. While previous blog posts discussed the LQR optimal control law through Pontryagin’s Minimum Principle and Hamilton-Jacobi-Bellman Equation, This post approaches LQR from the perspective of Dynamic Programming and Reinforcement Learning (RL). In particular, it shows how LQR can be solved in a model-based setting using both policy iteration and value iteration, and how related ideas extend to model-free methods through temporal difference learning. Code accompanying the post is available in this repository.