Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 12 additions & 12 deletions labs/lab_one.ipynb
Original file line number Diff line number Diff line change
Expand Up @@ -82,7 +82,7 @@
" @abstractmethod\n",
" def drift_coefficient(self, xt: torch.Tensor, t: torch.Tensor) -> torch.Tensor:\n",
" \"\"\"\n",
" Returns the drift coefficient of the ODE.\n",
" Returns the drift coefficient of the SDE.\n",
" Args:\n",
" - xt: state at time t, shape (batch_size, dim)\n",
" - t: time, shape ()\n",
Expand All @@ -94,7 +94,7 @@
" @abstractmethod\n",
" def diffusion_coefficient(self, xt: torch.Tensor, t: torch.Tensor) -> torch.Tensor:\n",
" \"\"\"\n",
" Returns the diffusion coefficient of the ODE.\n",
" Returns the diffusion coefficient of the SDE.\n",
" Args:\n",
" - xt: state at time t, shape (batch_size, dim)\n",
" - t: time, shape ()\n",
Expand All @@ -118,7 +118,7 @@
"metadata": {},
"source": [
"# Part 1: Numerical Methods for Simulating ODEs and SDEs\n",
"We may think of ODEs and SDEs as describing the motion of a particle through space. Intuitively, the ODE above says \"start at $X_0=x_0$\", and move so that your instantaneous velocity is given by $u_t(X_t)$. Similarly, the SDE says \"start at $X_0=x_0$\", and move so that your instantaneous velocity is given by $u_t(X_t)$ plus a little bit of random noise given scaled by $\\sigma_t$. Formally, these trajectories traced out by this intuitive descriptions are said to be *solutions* to the ODEs and SDEs, respectively. Numerical methods for computing these solutions are all essentially based on *simulating*, or *integrating*, the ODE or SDE. \n",
"We may think of ODEs and SDEs as describing the motion of a particle through space. Intuitively, the ODE above says \"start at $X_0=x_0$\", and move so that your instantaneous velocity is given by $u_t(X_t)$. Similarly, the SDE says \"start at $X_0=x_0$\", and move so that your instantaneous velocity is given by $u_t(X_t)$ plus a little bit of random noise scaled by $\\sigma_t$. Formally, these trajectories traced out by this intuitive descriptions are said to be *solutions* to the ODEs and SDEs, respectively. Numerical methods for computing these solutions are all essentially based on *simulating*, or *integrating*, the ODE or SDE. \n",
"\n",
"In this section we'll implement the *Euler* and *Euler-Maruyama* numerical simulation schemes for integrating ODEs and SDEs, respectively. Recall from lecture that the Euler simulation scheme corresponds to the discretization\n",
"$$d X_t = u_t(X_t) dt \\quad \\quad \\rightarrow \\quad \\quad X_{t + h} = X_t + hu_t(X_t),$$\n",
Expand Down Expand Up @@ -243,7 +243,7 @@
"metadata": {},
"source": [
"# Part 2: Visualizing Solutions to SDEs\n",
"Let's get a feel for what the solutions to these SDEs look like in practice (we'll get to ODEs later...). To do so, we we'll implement and visualize two special choices of SDEs from lecture: a (scaled) *Brownian motion*, and an *Ornstein-Uhlenbeck* (OU) process."
"Let's get a feel for what the solutions to these SDEs look like in practice (we'll get to ODEs later...). To do so, we'll implement and visualize two special choices of SDEs from lecture: a (scaled) *Brownian motion*, and an *Ornstein-Uhlenbeck* (OU) process."
]
},
{
Expand Down Expand Up @@ -271,7 +271,7 @@
"id": "62ba5c9e-f5cc-41a9-a850-43e46d79b3fb",
"metadata": {},
"source": [
"**Your job**: Fill in the `drift_coefficient` and `difusion_coefficient` methods of the `BrownianMotion` class below."
"**Your job**: Fill in the `drift_coefficient` and `diffusion_coefficient` methods of the `BrownianMotion` class below."
]
},
{
Expand All @@ -287,7 +287,7 @@
" \n",
" def drift_coefficient(self, xt: torch.Tensor, t: torch.Tensor) -> torch.Tensor:\n",
" \"\"\"\n",
" Returns the drift coefficient of the ODE.\n",
" Returns the drift coefficient of the SDE.\n",
" Args:\n",
" - xt: state at time t, shape (bs, dim)\n",
" - t: time, shape ()\n",
Expand All @@ -298,7 +298,7 @@
" \n",
" def diffusion_coefficient(self, xt: torch.Tensor, t: torch.Tensor) -> torch.Tensor:\n",
" \"\"\"\n",
" Returns the diffusion coefficient of the ODE.\n",
" Returns the diffusion coefficient of the SDE.\n",
" Args:\n",
" - xt: state at time t, shape (bs, dim)\n",
" - t: time, shape ()\n",
Expand Down Expand Up @@ -397,7 +397,7 @@
"id": "12325951-709c-4486-9ea7-f4c22b3cc1ef",
"metadata": {},
"source": [
"**Your job**: Fill in the `drift_coefficient` and `difusion_coefficient` methods of the `OUProcess` class below."
"**Your job**: Fill in the `drift_coefficient` and `diffusion_coefficient` methods of the `OUProcess` class below."
]
},
{
Expand All @@ -414,7 +414,7 @@
" \n",
" def drift_coefficient(self, xt: torch.Tensor, t: torch.Tensor) -> torch.Tensor:\n",
" \"\"\"\n",
" Returns the drift coefficient of the ODE.\n",
" Returns the drift coefficient of the SDE.\n",
" Args:\n",
" - xt: state at time t, shape (bs, dim)\n",
" - t: time, shape ()\n",
Expand All @@ -425,7 +425,7 @@
" \n",
" def diffusion_coefficient(self, xt: torch.Tensor, t: torch.Tensor) -> torch.Tensor:\n",
" \"\"\"\n",
" Returns the diffusion coefficient of the ODE.\n",
" Returns the diffusion coefficient of the SDE.\n",
" Args:\n",
" - xt: state at time t, shape (bs, dim)\n",
" - t: time, shape ()\n",
Expand Down Expand Up @@ -804,7 +804,7 @@
" \n",
" def drift_coefficient(self, xt: torch.Tensor, t: torch.Tensor) -> torch.Tensor:\n",
" \"\"\"\n",
" Returns the drift coefficient of the ODE.\n",
" Returns the drift coefficient of the SDE.\n",
" Args:\n",
" - xt: state at time t, shape (bs, dim)\n",
" - t: time, shape ()\n",
Expand All @@ -815,7 +815,7 @@
" \n",
" def diffusion_coefficient(self, xt: torch.Tensor, t: torch.Tensor) -> torch.Tensor:\n",
" \"\"\"\n",
" Returns the diffusion coefficient of the ODE.\n",
" Returns the diffusion coefficient of the SDE.\n",
" Args:\n",
" - xt: state at time t, shape (bs, dim)\n",
" - t: time, shape ()\n",
Expand Down
34 changes: 20 additions & 14 deletions labs/lab_three.ipynb
Original file line number Diff line number Diff line change
Expand Up @@ -76,7 +76,7 @@
"id": "f5900b74-1960-44ed-b837-664d6daa1a96",
"metadata": {},
"source": [
"As we will see shortly, a dataset like MNIST contains both images (in this case handwritten digits), as well as class labels (a value from 0-9 indicating). We will therefore generalize our notion of `Sampleable` to accommodate these labels as well. Whereas the old, `OldSampleable.sample` method returned only `samples: torch.Tensor`, we will now have it return both `samples: torch.Tensor` *and* `labels: Optional[torch.Tensor]`. In this way, we are formally realizing every such `Sampleable` instance as sampling from a *joint distribution* over data and labels. We implement our new `Sampleable` below."
"As we will see shortly, a dataset like MNIST contains both images (in this case handwritten digits), as well as class labels (a value from 0-9 indicating the label). We will therefore generalize our notion of `Sampleable` to accommodate these labels as well. Whereas the old, `OldSampleable.sample` method returned only `samples: torch.Tensor`, we will now have it return both `samples: torch.Tensor` *and* `labels: Optional[torch.Tensor]`. In this way, we are formally realizing every such `Sampleable` instance as sampling from a *joint distribution* over data and labels. We implement our new `Sampleable` below."
]
},
{
Expand Down Expand Up @@ -130,7 +130,7 @@
" self.std = std\n",
" self.dummy = nn.Buffer(torch.zeros(1)) # Will automatically be moved when self.to(...) is called...\n",
" \n",
" def sample(self, num_samples) -> Tuple[torch.Tensor, torch.Tensor]:\n",
" def sample(self, num_samples) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:\n",
" return self.std * torch.randn(num_samples, *self.shape).to(self.dummy.device), None"
]
},
Expand All @@ -140,7 +140,7 @@
"metadata": {},
"source": [
"Next, we make two updates in adding `ConditionalProbabilityPath` (and `GaussianConditionalProbabilityPath`):\n",
"1. We adjust to handle the addition of labels to `Sampleable`. Recall earlier that our called our conditioning variable `z` with $z \\sim p_{\\text{data}}(z)$. Now, we sample both `z`, as well as a label `y`, with $(z,y) \\sim p_{\\text{data}}(z,y)$.\n",
"1. We adjust to handle the addition of labels to `Sampleable`. Recall earlier that we called our conditioning variable `z` with $z \\sim p_{\\text{data}}(z)$. Now, we sample both `z`, as well as a label `y`, with $(z,y) \\sim p_{\\text{data}}(z,y)$.\n",
"2. We ensure that the logic is compatible with shapes of size `(batch_size, c, h, w)`, rather than `(batch_size, dim)`. While the latter was sufficient for 2D data of shape `(batch_size, 2)`, we will now be working with images which, when batched, have shape `(batch_size, c, h, w)`. Here `c`, `h`, and `w`, denote the number of channels, the height, and the width, respectively.\n",
"3. To avoid any unfortunate broadcasting issues, we will maintain our time variable `t` in the shape `(batch_size, 1, 1, 1)`."
]
Expand Down Expand Up @@ -245,13 +245,16 @@
"class Alpha(ABC):\n",
" def __init__(self):\n",
" # Check alpha_t(0) = 0\n",
" t0 = torch.zeros(1, 1, 1, 1)\n",
" assert torch.allclose(\n",
" self(torch.zeros(1,1,1,1)), torch.zeros(1,1,1,1)\n",
" )\n",
" self(t0), torch.zeros_like(t0)\n",
" ), \"Alpha(0) must equal 0\"\n",
" \n",
" # Check alpha_1 = 1\n",
" t1 = torch.ones(1, 1, 1, 1)\n",
" assert torch.allclose(\n",
" self(torch.ones(1,1,1,1)), torch.ones(1,1,1,1)\n",
" )\n",
" self(t1), torch.ones_like(t1)\n",
" ), \"Alpha(1) must equal 1\"\n",
" \n",
" @abstractmethod\n",
" def __call__(self, t: torch.Tensor) -> torch.Tensor:\n",
Expand Down Expand Up @@ -279,13 +282,16 @@
"class Beta(ABC):\n",
" def __init__(self):\n",
" # Check beta_0 = 1\n",
" t0 = torch.zeros(1, 1, 1, 1)\n",
" assert torch.allclose(\n",
" self(torch.zeros(1,1,1,1)), torch.ones(1,1,1,1)\n",
" )\n",
" self(t0), torch.ones_like(t0), torch.ones_like(t0)\n",
" ), \"Beta(0) must equal 1\"\n",
" \n",
" # Check beta_1 = 0\n",
" t1 = torch.ones(1, 1, 1, 1)\n",
" assert torch.allclose(\n",
" self(torch.ones(1,1,1,1)), torch.zeros(1,1,1,1)\n",
" )\n",
" self(t1), torch.zeros_like(t1)\n",
" ), \"Beta(1) must equal 0\"\n",
" \n",
" @abstractmethod\n",
" def __call__(self, t: torch.Tensor) -> torch.Tensor:\n",
Expand Down Expand Up @@ -751,7 +757,7 @@
"source": [
"**Guidance**: Whereas for unconditional generation, we simply wanted to generate *any* digit, we would now like to be able to specify, or *condition*, on the identity of the digit we would like to generate. That is, we would like to be able to say \"generate an image of the digit 8\", rather than just \"generate an image of a digit\". We will henceforth refer to the digit we would like to generate as $x \\in \\mathbb{R}^{1 \\times 32 \\times 32}$, and the conditioning variable (in this case, a label), as $y \\in \\{0, 1, \\dots, 9\\}$. If we imagine fixing our choice of $y$, and take our data distribution as $p_{\\text{simple}}(x|y)$, then we have recovered the unconditional generative problem, and we can construct a generative model using e.g., a conditional flow matching objective via $$\\begin{align*}\\mathcal{L}_{\\text{CFM}}^{\\text{guided}}(\\theta;y) &= \\,\\,\\mathbb{E}_{\\square} \\lVert u_t^{\\theta}(x|y) - u_t^{\\text{ref}}(x|z)\\rVert^2\\\\ \\square &= z \\sim p_{\\text{data}}(z|y), x \\sim p_t(x|z)\\end{align*}$$\n",
"We may now then allow $y$ to vary by simply taking our conditional flow matching expectation to be over $y$ as well (rather than fixing $y$), and explicitly conditioning our learned approximation on $u_t^{\\theta}(x|y)$ on the choice of $y$. We therefore obtain the the *guided* conditional flow matching objective $$\\begin{align*}\\mathcal{L}_{\\text{CFM}}(\\theta) &= \\,\\,\\mathbb{E}_{\\square} \\lVert u_t^{\\theta}(x|y) - u_t^{\\text{ref}}(x|z)\\rVert^2\\\\ \\square &= z,y \\sim p_{\\text{data}}(z,y), x \\sim p_t(x|z)\\end{align*}$$\n",
"Note that $(z,y) \\sim p_{\\text{simple}}(z,y)$ is obtained in practice by sampling an image $z$, and a label $y$, from our labelled (MNIST) dataset. This is all well and good, and we emphasize that if our goal was simply to sample from $p_{\\text{data}}(x|y)$, our job would be done (at least in theory). In practice, one might argue that we care more about the *perceptual quality* of our images. To this end, we will a derive a procedure known as *classifier-free guidance*."
"Note that $(z,y) \\sim p_{\\text{simple}}(z,y)$ is obtained in practice by sampling an image $z$, and a label $y$, from our labelled (MNIST) dataset. This is all well and good, and we emphasize that if our goal was simply to sample from $p_{\\text{data}}(x|y)$, our job would be done (at least in theory). In practice, one might argue that we care more about the *perceptual quality* of our images. To this end, we will derive a procedure known as *classifier-free guidance*."
]
},
{
Expand Down Expand Up @@ -780,7 +786,7 @@
"id": "d8728902-21ce-4b00-b4f1-bb3573542206",
"metadata": {},
"source": [
"**Training and CFG**: We must now amend our conditional flow matching objective to account for the possibility of $y = \\varnothing$. Of course, when we sample $(z,y)$ from MNIST, we will never obtain $y = \\varnothing$, so we must introduce the possibliity of this artificially. To do so, we will define some hyperparameter $\\eta$ to be the *probability* that we discard the original label $y$, and replace it with $\\varnothing$. In practice, we might set $\\varnothing = 10$, for example, as it is sufficient to distinguish it from the other digit identities. When we go and implement our model, we need ony be able to index into some embedding, such as via `torch.nn.Embedding`. We thus arrive at our CFG conditional flow matching training objective:\n",
"**Training and CFG**: We must now amend our conditional flow matching objective to account for the possibility of $y = \\varnothing$. Of course, when we sample $(z,y)$ from MNIST, we will never obtain $y = \\varnothing$, so we must introduce the possibliity of this artificially. To do so, we will define some hyperparameter $\\eta$ to be the *probability* that we discard the original label $y$, and replace it with $\\varnothing$. In practice, we might set $\\varnothing = 10$, for example, as it is sufficient to distinguish it from the other digit identities. When we go and implement our model, we need only be able to index into some embedding, such as via `torch.nn.Embedding`. We thus arrive at our CFG conditional flow matching training objective:\n",
"$$\\begin{align*}\\mathcal{L}_{\\text{CFM}}(\\theta) &= \\,\\,\\mathbb{E}_{\\square} \\lVert u_t^{\\theta}(x|y) - u_t^{\\text{ref}}(x|z)\\rVert^2\\\\\n",
"\\square &= z,y \\sim p_{\\text{data}}(z,y), x \\sim p_t(x|z),\\,\\text{replace $y$ with $\\varnothing$ with probability $\\eta$}\\end{align*}$$\n",
"In plain English, this objective reads:\n",
Expand Down Expand Up @@ -897,7 +903,7 @@
"metadata": {},
"source": [
"# Part 3: An Architecture for Images\n",
"At this point, we have discussed classifier free guidance, and the necessary considerations that must be made on the part of our model and in training our model. What remains is to actually discuss the choice of model. In particular, our usual choice of an MLP, while fine for the simple distributions of the previous lab, will no longer suffice. To this end, we will a new convolutional architecture - the **U-Net** - which is specifically tailored toward images. A diagram of the U-Net we'll be using is shown below. ![image.png](attachment:bd703834-9239-4ed3-b8c1-9639fc971575.png)"
"At this point, we have discussed classifier free guidance, and the necessary considerations that must be made on the part of our model and in training our model. What remains is to actually discuss the choice of model. In particular, our usual choice of an MLP, while fine for the simple distributions of the previous lab, will no longer suffice. To this end, we will use a new convolutional architecture - the **U-Net** - which is specifically tailored toward images. A diagram of the U-Net we'll be using is shown below. ![image.png](attachment:bd703834-9239-4ed3-b8c1-9639fc971575.png)"
]
},
{
Expand Down
29 changes: 20 additions & 9 deletions labs/lab_two.ipynb
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@
"metadata": {},
"source": [
"### Part 0: Miscellaneous Imports and Utility Functions\n",
"No questions here, but free to read through to familiarize yourself with these helper functions. Most of this is what you already completed in lab one!"
"No questions here, but feel free to read through to familiarize yourself with these helper functions. Most of this is what you already completed in lab one!"
]
},
{
Expand Down Expand Up @@ -201,6 +201,11 @@
"source": [
"# Several plotting utility functions\n",
"def hist2d_samples(samples, ax: Optional[Axes] = None, bins: int = 200, scale: float = 5.0, percentile: int = 99, **kwargs):\n",
" if isinstance(samples, torch.Tensor):\n",
" samples = samples.detach().cpu().numpy()\n",
" else: \n",
" samples = np.array(samples)\n",
" \n",
" H, xedges, yedges = np.histogram2d(samples[:, 0], samples[:, 1], bins=bins, range=[[-scale, scale], [-scale, scale]])\n",
" \n",
" # Determine color normalization based on the 99th percentile\n",
Expand Down Expand Up @@ -574,13 +579,16 @@
"class Alpha(ABC):\n",
" def __init__(self):\n",
" # Check alpha_t(0) = 0\n",
" t0 = torch.zeros(1, 1, 1, 1)\n",
" assert torch.allclose(\n",
" self(torch.zeros(1,1)), torch.zeros(1,1)\n",
" )\n",
" self(t0), torch.zeros_like(t0)\n",
" ), \"Alpha(0) must equal 0\"\n",
" \n",
" # Check alpha_1 = 1\n",
" t1 = torch.ones(1, 1, 1, 1)\n",
" assert torch.allclose(\n",
" self(torch.ones(1,1)), torch.ones(1,1)\n",
" )\n",
" self(t1), torch.ones_like(t1)\n",
" ), \"Alpha(1) must equal 1\"\n",
" \n",
" @abstractmethod\n",
" def __call__(self, t: torch.Tensor) -> torch.Tensor:\n",
Expand Down Expand Up @@ -608,13 +616,16 @@
"class Beta(ABC):\n",
" def __init__(self):\n",
" # Check beta_0 = 1\n",
" t0 = torch.zeros(1, 1, 1, 1)\n",
" assert torch.allclose(\n",
" self(torch.zeros(1,1)), torch.ones(1,1)\n",
" )\n",
" self(t0), torch.ones_like(t0), torch.ones_like(t0)\n",
" ), \"Beta(0) must equal 1\"\n",
" \n",
" # Check beta_1 = 0\n",
" t1 = torch.ones(1, 1, 1, 1)\n",
" assert torch.allclose(\n",
" self(torch.ones(1,1)), torch.zeros(1,1)\n",
" )\n",
" self(t1), torch.zeros_like(t1)\n",
" ), \"Beta(1) must equal 0\"\n",
" \n",
" @abstractmethod\n",
" def __call__(self, t: torch.Tensor) -> torch.Tensor:\n",
Expand Down
Loading