Solved ASSIGNMENT III Generative modeling SH2150 Machine Learning in Physics

Diffusion models have become one of the most exciting ideas in modern generative AI, and at their core is a beautiful interplay between randomness and structure described by stochastic differential equations (SDEs). In this assignment, we dive into that world: from exploring classical processes like Brownian motion and the Ornstein–Uhlenbeck dynamics, to simulating trajectories with Euler–Maruyama, to building neural networks that learn how data is gradually corrupted and then reconstructed. By moving from simple 1D SDEs all the way to UNet-based image models, we get to experience how mathematical theory, numerical simulation, and deep learning come together to form the foundation of today’s powerful diffusion and flow-based generative models. Your tasks are to • Get familiar with SDEs • Set up a training/inference pipeline for generative models • Sample from a 2D distribution using a flow model • Define a UNet architecture for image generation • Generate images from a distribution of your choice 1. The theoretical basis for diffusion models is stochastic differential equations (SDE). An SDE is a differential equation symbolically of the form dXt = ut(Xt)dt + σtdWt, X0 = x0, (1) where Xt is the trajectory, ut is the vector field or drift coefficient, σt is the diffusion coefficient determining the amount of noise, and Wt is a Brownian motion. The term σtdWt is what gives rise to the stochasticity (randomness) in the system – without it the system reduces to an ordinary differential equation (ODE) dXt dt = ut(Xt). (a) Different choices of ut yield different stochastic processes that can be used to model (1p) many different physical phenomena. The Ornstein-Uhlenbeck (OU) process is given by ut(Xt) = −θXt, σt = σ, for some constants θ, σ. It can be used to model massive particles under the influence of friction. Show that the OU process is equivalent to Langevin dynamics defined by the drift coefficient ut(Xt) = 1 2 σ 2 t ∇ log p(Xt), σt = σ, when p(Xt) = N (0, σ 2 2θ ). 1 (b) Being able to simulate SDEs will be needed to sample new data points from gener- (1p) ative models. The simplest solver for SDEs is the Euler-Maruyama method defined as Xt+h = Xt + hut(Xt) + √ hσtϵt, ϵt ∼ N (0, I), (2) where h is the step size. This reduces to the Euler method for ODEs when σ = 0. Implement the Euler-Maruyama method. (c) Simulate 10 Brownian motions defined by σt = σ, ut = 0, X0 = 0. What happens (1p) for different values of σ? (d) Simulate the OU process defined above. Run it with a variety of different initial (2p) points X0. Pay attention to the ratio σ 2 2θ , and comment on the convergence behavior of the solutions. Are they approaching a particular point or a distribution? 2. A diffusion model is an SDE where the vector field ut is parametrized by a learnable neural network u θ t . If there is no diffusion term, and the SDE is reduced to an ODE, we get a so-called flow model. The network u θ t is trained on data points X1 ∼ pdata that have been corrupted with varying amounts of noise corresponding to different time steps Xt, t ∈ [0, 1], of the SDE/ODE. By simulating the SDE forwards from t = 0 to t = 1 with our trained vector field, starting with pure noise X0 ∼ pnoise, we can generate new data samples from pdata. In a common type of generative model, called denoising diffusion model, the conditional path (which describes how a data point z is corrupted) is given by xt = αtz + βtϵ, ϵ ∼ N (0, I) where αt, βt are continuously differentiable and monotonic noise schedulers with α0 = β1 = 0 and α1 = β0 = 1. This implies xt ∼ pt(· | z) = N (αtz, β2 t I). (a) Implement linear noise schedulers αt = t, βt = 1 − t. (1p) (b) Implement the conditional path that given a data point z and a time t returns the (1p) corrupted data point xt = αtz + βtϵ. (c) Define a 2D toy distribution (e.g. a Gaussian mixture) pdata. Simulate the corrup- (1p) tion process on 1000 samples from pdata for t = 0, 0.25, 0.50, 0.75, 1, and plot one 2D histogram per time point. (d) Modify the noise schedulers so that they are non-linear (while still satisfying the (1p) requirements stated above). Then, generate a plot analogous to panel (c) using the same data points z corrupted under this new schedule, and discuss how the change in scheduling influences the corruption process. 3. There are many ways to train generative models. In this case, we will want to minimize the conditional flow matching loss given by L(θ) = Et∼U[0,1],z∼pdata,x∼pt(·|z) [ u θ t (x) − ut(x | z) 2 ]. (a) The conditional vector field is given by (1p) ut(x | z) = α˙ t − β˙ t βt αt ! z + β˙ t βt x. Find a simplified expression for ut(xt | z) when xt is drawn from the conditional path pt(· | z), and implement it. Page 2 (b) Implement a MLP architecture that takes (x, t) ∈ R (1p) 3 as inputs and outputs the estimated vector field u θ t (xt) ∈ R 2 . (c) Implement a training loop according to Algorithm 1, and train an MLP with 4 (1p) hidden layers of dimension 64 on the 2D toy dataset. Choose noise schedulers αt, βt according to your liking (that fulfills the needed criteria), and state what you used. Does the loss converge? (d) Let the number of time steps nt = 1000, and sample 300 realizations from pdata (1p) using the Euler method (σt = 0) applied to the ODE parametrized by your trained vector field. Plot a 2D scatter plot of 300 corrupted data points according to the true conditional path for t = 0, 0.25, 0.50, 0.75, 1. Provide a similar plot for the 300 generated data points together with a plot of the simulated trajectories Xt. How does the choice of nt affect the sampling? Algorithm 1: Conditional flow matching Require: data set with samples pdata, vector field u θ t for each batch do Sample data point z ∼ pdata. Sample random time t ∼ U[0, 1]. Sample noise ϵ ∼ N (0, I). Set xt = αtz + βtϵ. Compute loss L(θ) = u θ t (xt) − ut(xt | z) 2 . Update model parameters θ using gradient step on L(θ). end 4. Let us now turn to image generation. To handle high-dimensional image data, we need another architecture than an MLP to parameterize our vector field. We will use the famous UNet architecture with some modification to allow for the time embedding (the network must be informed about the current time t and for CNNs it is not as straightforward to do this as for MLPs where we just fed it as an additional input). (a) Select a few data points z from your choice of image data set, and provide a plot of (1p) the corresponding corrupted data points xt for time points t = 0, 0.25, 0.50, 0.75, 1. (b) Implement the residual layer of the UNet architecture and add the pre-defined time (3p) embedding according to figure 1. Train a vector field over 5000 epochs on the image data with batch size of 250 using your UNet. (c) Sample a couple of images from the image distribution by solving the flow ODE (1p) dXt = u θ t (Xt)dt while using your trained UNet as the drift term. (d) Play around with different noise schedulers and comment on how they affect the (1p) sample quality. 5. So far we have only trained flow models (no diffusion term). Via the Fokker-Planck equation, it can be shown that the following SDE dXt =  u θ t (Xt) + σ 2 t 2 ∇ log pt(Xt)  dt + σtdWt (3) have the same probability paths pt as the flow ODE dXt dt = u θ t (Xt). Hence, to convert the flow model into a diffusion model, we have to learn the second drift term ∇ log pt(Xt) called the score function. This can be done by training a network s θ t using conditional score matching L(θ) = Et∼U[0,1],z∼pdata,x∼pt(·|z) [ s θ t (x) − ∇ log pt(x | z) 2 ]. Page 3 Initial Conv Encoder Encoder Encoder Bottleneck Decoder Decoder Decoder Final Conv 𝑥𝑡 𝑡 𝑢𝑡 𝜃 (𝑥𝑡) Embed Residual connection Residual connection Residual connection(s) SiLU + BatchNorm + Conv SiLU + BatchNorm + Conv MLP 𝑡 Residual layer Figure 1: UNet architecture together with a detailed overview of the residual layer. Each encoder consists of two residual layers followed by downsampling. Each decoder consists of upsampling followed by two residual layers. The bottleneck module consist of three residual layers. (a) Derive an explicit formula for the conditional score function ∇ log pt(x | z) from (3p) the conditional path and implement it. Then train a score network s θ t on the image dataset using the conditional score matching loss. (b) Sample new data using the diffusion model by combining the trained vector field (3p) u θ t and the score network s t θ and simulate the SDE in (3). How does the value of σt affect the samples? Final question: Did you use an AI tool (other than the machine learning models you trained in this exercise) for anything else than information searching, when solving this problem set? If so, please write a brief statement of how you used AI. Total number of points: 25 Motivate your answers wherever applicable. Good luck! Page 4