前言

本篇文章记录 MIT S.184 作业 Lab 2: Flow Matching and Score Matching 的实现，和大家一起分享交流😄。

reference：https://github.com/eje24/iap-diffusion-labs

Lab Two: Flow Matching and Score Matching

欢迎来到 Lab Two！在本实验中，我们将通过直观且动手实践的方式，带你理解 flow matching 和 score matching。

Part 0: Miscellaneous Imports and Utility Functions

这部分没有问题需要回答，不过你可以自由阅读这些内容，以便熟悉这些辅助函数。其中大部分内容都是我们在 Lab One 中已经完成过的！

老样子，我们先为 Lab Two 准备基础依赖：

from abc import ABC, abstractmethod
from typing import Optional, List, Type, Tuple, Dict
import math

import numpy as np
from matplotlib import pyplot as plt
import matplotlib.cm as cm
from matplotlib.axes._axes import Axes
import torch
import torch.distributions as D
from torch.func import vmap, jacrev
from tqdm import tqdm
import seaborn as sns
from sklearn.datasets import make_moons, make_circles

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

导入的模块整体上和 Lab One 非常接近，但这里额外引入了一些和二维 toy 数据集、颜色映射相关的工具。

下面我们定义下后续实验会用到的分布对象：

class Sampleable(ABC):
    """
    Distribution which can be sampled from
    """
    @property
    @abstractmethod
    def dim(self) -> int:
        """
        Returns:
            - Dimensionality of the distribution
        """
        pass
        
    @abstractmethod
    def sample(self, num_samples: int) -> torch.Tensor:
        """
        Args:
            - num_samples: the desired number of samples
        Returns:
            - samples: shape (batch_size, dim)
        """
        pass

class Density(ABC):
    """
    Distribution with tractable density
    """
    @abstractmethod
    def log_density(self, x: torch.Tensor) -> torch.Tensor:
        """
        Returns the log density at x.
        Args:
            - x: shape (batch_size, dim)
        Returns:
            - log_density: shape (batch_size, 1)
        """
        pass

class Gaussian(torch.nn.Module, Sampleable, Density):
    """
    Multivariate Gaussian distribution
    """
    def __init__(self, mean: torch.Tensor, cov: torch.Tensor):
        """
        mean: shape (dim,)
        cov: shape (dim,dim)
        """
        super().__init__()
        self.register_buffer("mean", mean)
        self.register_buffer("cov", cov)

    @property
    def dim(self) -> int:
        return self.mean.shape[0]

    @property
    def distribution(self):
        return D.MultivariateNormal(self.mean, self.cov, validate_args=False)

    def sample(self, num_samples) -> torch.Tensor:
        return self.distribution.sample((num_samples,))
        
    def log_density(self, x: torch.Tensor):
        return self.distribution.log_prob(x).view(-1, 1)

    @classmethod
    def isotropic(cls, dim: int, std: float) -> "Gaussian":
        mean = torch.zeros(dim)
        cov = torch.eye(dim) * std ** 2
        return cls(mean, cov)

class GaussianMixture(torch.nn.Module, Sampleable, Density):
    """
    Two-dimensional Gaussian mixture model, and is a Density and a Sampleable. Wrapper around torch.distributions.MixtureSameFamily.
    """
    def __init__(
        self,
        means: torch.Tensor,  # nmodes x data_dim
        covs: torch.Tensor,  # nmodes x data_dim x data_dim
        weights: torch.Tensor,  # nmodes
    ):
        """
        means: shape (nmodes, 2)
        covs: shape (nmodes, 2, 2)
        weights: shape (nmodes, 1)
        """
        super().__init__()
        self.nmodes = means.shape[0]
        self.register_buffer("means", means)
        self.register_buffer("covs", covs)
        self.register_buffer("weights", weights)

    @property
    def dim(self) -> int:
        return self.means.shape[1]

    @property
    def distribution(self):
        return D.MixtureSameFamily(
                mixture_distribution=D.Categorical(probs=self.weights, validate_args=False),
                component_distribution=D.MultivariateNormal(
                    loc=self.means,
                    covariance_matrix=self.covs,
                    validate_args=False,
                ),
                validate_args=False,
            )

    def log_density(self, x: torch.Tensor) -> torch.Tensor:
        return self.distribution.log_prob(x).view(-1, 1)

    def sample(self, num_samples: int) -> torch.Tensor:
        return self.distribution.sample(torch.Size((num_samples,)))

    @classmethod
    def random_2D(
        cls, nmodes: int, std: float, scale: float = 10.0, x_offset: float = 0.0, seed = 0.0
    ) -> "GaussianMixture":
        torch.manual_seed(seed)
        means = (torch.rand(nmodes, 2) - 0.5) * scale + x_offset * torch.Tensor([1.0, 0.0])
        covs = torch.diag_embed(torch.ones(nmodes, 2)) * std ** 2
        weights = torch.ones(nmodes)
        return cls(means, covs, weights)

    @classmethod
    def symmetric_2D(
        cls, nmodes: int, std: float, scale: float = 10.0, x_offset: float = 0.0
    ) -> "GaussianMixture":
        angles = torch.linspace(0, 2 * np.pi, nmodes + 1)[:nmodes]
        means = torch.stack([torch.cos(angles), torch.sin(angles)], dim=1) * scale + torch.Tensor([1.0, 0.0]) * x_offset
        covs = torch.diag_embed(torch.ones(nmodes, 2) * std ** 2)
        weights = torch.ones(nmodes) / nmodes
        return cls(means, covs, weights)

Lab Two 会研究 flow matching 和 score matching，这两类方法都需要从某些分布中采样，有时需要计算分布密度，因此这里先把分布抽象成统一接口。

接着定义一组二维分布可视化工具函数：

# Several plotting utility functions
def hist2d_samples(samples, ax: Optional[Axes] = None, bins: int = 200, scale: float = 5.0, percentile: int = 99, **kwargs):
    H, xedges, yedges = np.histogram2d(samples[:, 0], samples[:, 1], bins=bins, range=[[-scale, scale], [-scale, scale]])
    
    # Determine color normalization based on the 99th percentile
    cmax = np.percentile(H, percentile)
    cmin = 0.0
    norm = cm.colors.Normalize(vmax=cmax, vmin=cmin)
    
    # Plot using imshow for more control
    extent = [xedges[0], xedges[-1], yedges[0], yedges[-1]]
    ax.imshow(H.T, extent=extent, origin='lower', norm=norm, **kwargs)

def hist2d_sampleable(sampleable: Sampleable, num_samples: int, ax: Optional[Axes] = None, bins=200, scale: float = 5.0, percentile: int = 99, **kwargs):
    assert sampleable.dim == 2
    if ax is None:
        ax = plt.gca()
    samples = sampleable.sample(num_samples).detach().cpu() # (ns, 2)
    hist2d_samples(samples, ax, bins, scale, percentile, **kwargs)

def scatter_sampleable(sampleable: Sampleable, num_samples: int, ax: Optional[Axes] = None, **kwargs):
    assert sampleable.dim == 2
    if ax is None:
        ax = plt.gca()
    samples = sampleable.sample(num_samples) # (ns, 2)
    ax.scatter(samples[:,0].cpu(), samples[:,1].cpu(), **kwargs)

def kdeplot_sampleable(sampleable: Sampleable, num_samples: int, ax: Optional[Axes] = None, **kwargs):
    assert sampleable.dim == 2
    if ax is None:
        ax = plt.gca()
    samples = sampleable.sample(num_samples) # (ns, 2)
    sns.kdeplot(x=samples[:,0].cpu(), y=samples[:,1].cpu(), ax=ax, **kwargs)

def imshow_density(density: Density, x_bounds: Tuple[float, float], y_bounds: Tuple[float, float], bins: int, ax: Optional[Axes] = None, x_offset: float = 0.0, **kwargs):
    if ax is None:
        ax = plt.gca()
    x_min, x_max = x_bounds
    y_min, y_max = y_bounds
    x = torch.linspace(x_min, x_max, bins).to(device) + x_offset
    y = torch.linspace(y_min, y_max, bins).to(device)
    X, Y = torch.meshgrid(x, y)
    xy = torch.stack([X.reshape(-1), Y.reshape(-1)], dim=-1)
    density = density.log_density(xy).reshape(bins, bins).T
    im = ax.imshow(density.cpu(), extent=[x_min, x_max, y_min, y_max], origin='lower', **kwargs)

def contour_density(density: Density, bins: int, scale: float, ax: Optional[Axes] = None, x_offset:float = 0.0, **kwargs):
    if ax is None:
        ax = plt.gca()
    x = torch.linspace(-scale + x_offset, scale + x_offset, bins).to(device)
    y = torch.linspace(-scale, scale, bins).to(device)
    X, Y = torch.meshgrid(x, y)
    xy = torch.stack([X.reshape(-1), Y.reshape(-1)], dim=-1)
    density = density.log_density(xy).reshape(bins, bins).T
    im = ax.contour(density.cpu(), origin='lower', **kwargs)

这些函数主要用于后面观察源分布、目标分布、概率路径中的中间分布以及模型生成样本的效果。

我们接着来定义两个抽象基类：ODE 和 SDE，它们的作用是为后续的连续时间动态系统提供统一接口：

class ODE(ABC):
    @abstractmethod
    def drift_coefficient(self, xt: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
        """
        Returns the drift coefficient of the ODE.
        Args:
            - xt: state at time t, shape (bs, dim)
            - t: time, shape (batch_size, 1)
        Returns:
            - drift_coefficient: shape (batch_size, dim)
        """
        pass

class SDE(ABC):
    @abstractmethod
    def drift_coefficient(self, xt: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
        """
        Returns the drift coefficient of the ODE.
        Args:
            - xt: state at time t, shape (batch_size, dim)
            - t: time, shape (batch_size, 1)
        Returns:
            - drift_coefficient: shape (batch_size, dim)
        """
        pass

    @abstractmethod
    def diffusion_coefficient(self, xt: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
        """
        Returns the diffusion coefficient of the ODE.
        Args:
            - xt: state at time t, shape (batch_size, dim)
            - t: time, shape (batch_size, 1)
        Returns:
            - diffusion_coefficient: shape (batch_size, dim)
        """
        pass

最后，我们来定义用于模拟 ODE / SDE 的通用框架：

class Simulator(ABC):
    @abstractmethod
    def step(self, xt: torch.Tensor, t: torch.Tensor, dt: torch.Tensor):
        """
        Takes one simulation step
        Args:
            - xt: state at time t, shape (bs, dim)
            - t: time, shape (bs,1)
            - dt: time, shape (bs,1)
        Returns:
            - nxt: state at time t + dt (bs, dim)
        """
        pass

    @torch.no_grad()
    def simulate(self, x: torch.Tensor, ts: torch.Tensor):
        """
        Simulates using the discretization gives by ts
        Args:
            - x_init: initial state at time ts[0], shape (batch_size, dim)
            - ts: timesteps, shape (bs, num_timesteps,1)
        Returns:
            - x_final: final state at time ts[-1], shape (batch_size, dim)
        """
        for t_idx in range(len(ts) - 1):
            t = ts[:, t_idx]
            h = ts[:, t_idx + 1] - ts[:, t_idx]
            x = self.step(x, t, h)
        return x

    @torch.no_grad()
    def simulate_with_trajectory(self, x: torch.Tensor, ts: torch.Tensor):
        """
        Simulates using the discretization gives by ts
        Args:
            - x_init: initial state at time ts[0], shape (bs, dim)
            - ts: timesteps, shape (bs, num_timesteps, 1)
        Returns:
            - xs: trajectory of xts over ts, shape (batch_size, num
            _timesteps, dim)
        """
        xs = [x.clone()]
        nts = ts.shape[1]
        for t_idx in tqdm(range(nts - 1)):
            t = ts[:,t_idx]
            h = ts[:, t_idx + 1] - ts[:, t_idx]
            x = self.step(x, t, h)
            xs.append(x.clone())
        return torch.stack(xs, dim=1)

class EulerSimulator(Simulator):
    def __init__(self, ode: ODE):
        self.ode = ode
        
    def step(self, xt: torch.Tensor, t: torch.Tensor, h: torch.Tensor):
        return xt + self.ode.drift_coefficient(xt,t) * h

class EulerMaruyamaSimulator(Simulator):
    def __init__(self, sde: SDE):
        self.sde = sde
        
    def step(self, xt: torch.Tensor, t: torch.Tensor, h: torch.Tensor):
        return xt + self.sde.drift_coefficient(xt,t) * h + self.sde.diffusion_coefficient(xt,t) * torch.sqrt(h) * torch.randn_like(xt)

def record_every(num_timesteps: int, record_every: int) -> torch.Tensor:
    """
    Compute the indices to record in the trajectory given a record_every parameter
    """
    if record_every == 1:
        return torch.arange(num_timesteps)
    return torch.cat(
        [
            torch.arange(0, num_timesteps - 1, record_every),
            torch.tensor([num_timesteps - 1]),
        ]
    )

Part 1: Implementing Conditional Probability Paths

回顾课堂和课程笔记中关于 conditional flow matching 的基本思想：我们希望描述一个 条件概率路径 $p_t(x|z)$ ，使得 $p_1(x|z)={\delta }_z(x)$ ，并且 $p_0(z)=p_{\text{simple}}$ ，例如高斯分布，同时 $p_t(x|z)$ 在 $p_0(x|z)$ 和 $p_1(x|z)$ 之间连续插值（这里我们不做严格表述）。

这样的条件路径可以被看作对应于某种 corruption process。反向来看，这个过程会把 $t=1$ 时刻的点 $z$ ，驱动成 $t=0$ 时刻服从 $p_0(x|z)$ 的分布。这样的 corruption process 可以由下面的 ODE 给出：

$$d{X}_t=u_t^{\text{ref}}(X_t|z)dt,X_0\sim p_{\text{simple}}.$$

其中漂移项 $u_t^{\text{ref}}(X_t|z)$ 被称为 条件向量场。通过对所有可能的 $z$ 对 $u_t^{\text{ref}}(x|z)$ 求平均，我们可以得到 边缘向量场 $u_t^{\text{ref}}(x)$ 。Flow matching 利用了这样一个事实：由边缘向量场 $u_t^{\text{ref}}(x)$ 生成的 边缘概率路径 $p_t(x)$ ，能够连接 $p_{\text{simple}}$ 和 $p_{\text{data}}$ 。由于条件向量场 $u_t^{\text{ref}}(x|z)$ 通常可以解析得到，因此我们可以不直接回归未知的边缘向量场 $u_t^{\text{ref}}(x)$ ，而是显式地回归条件向量场 $u_t^{\text{ref}}(x|z)$ ，从而隐式地学习边缘向量场。

这个构造中核心对象是 条件概率路径。它的接口在下面的 ConditionalProbabilityPath 类中实现。在本实验中，我们将实现两个子类：GaussianConditionalProbabilityPath 和 LinearConditionalProbabilityPath，它们分别对应课堂和笔记中同名的概率路径。

class ConditionalProbabilityPath(torch.nn.Module, ABC):
    """
    Abstract base class for conditional probability paths
    """
    def __init__(self, p_simple: Sampleable, p_data: Sampleable):
        super().__init__()
        self.p_simple = p_simple
        self.p_data = p_data

    def sample_marginal_path(self, t: torch.Tensor) -> torch.Tensor:
        """
        Samples from the marginal distribution p_t(x) = p_t(x|z) p(z)
        Args:
            - t: time (num_samples, 1)
        Returns:
            - x: samples from p_t(x), (num_samples, dim)
        """
        num_samples = t.shape[0]
        # Sample conditioning variable z ~ p(z)
        z = self.sample_conditioning_variable(num_samples) # (num_samples, dim)
        # Sample conditional probability path x ~ p_t(x|z)
        x = self.sample_conditional_path(z, t) # (num_samples, dim)
        return x

    @abstractmethod
    def sample_conditioning_variable(self, num_samples: int) -> torch.Tensor:
        """
        Samples the conditioning variable z
        Args:
            - num_samples: the number of samples
        Returns:
            - z: samples from p(z), (num_samples, dim)
        """
        pass
    
    @abstractmethod
    def sample_conditional_path(self, z: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
        """
        Samples from the conditional distribution p_t(x|z)
        Args:
            - z: conditioning variable (num_samples, dim)
            - t: time (num_samples, 1)
        Returns:
            - x: samples from p_t(x|z), (num_samples, dim)
        """
        pass
        
    @abstractmethod
    def conditional_vector_field(self, x: torch.Tensor, z: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
        """
        Evaluates the conditional vector field u_t(x|z)
        Args:
            - x: position variable (num_samples, dim)
            - z: conditioning variable (num_samples, dim)
            - t: time (num_samples, 1)
        Returns:
            - conditional_vector_field: conditional vector field (num_samples, dim)
        """ 
        pass

    @abstractmethod
    def conditional_score(self, x: torch.Tensor, z: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
        """
        Evaluates the conditional score of p_t(x|z)
        Args:
            - x: position variable (num_samples, dim)
            - z: conditioning variable (num_samples, dim)
            - t: time (num_samples, 1)
        Returns:
            - conditional_score: conditional score (num_samples, dim)
        """ 
        pass

ConditionalProbabilityPath 是所有条件概率路径的抽象基类，它的作用是统一描述下面这个对象：

$$p_t(x|z)$$

也就是在给定条件变量 $z$ 的情况下，时间 $t$ 上 $x$ 的条件分布。

Part 2: Gaussian Conditional Probability Paths

在这一节中，我们将通过 GaussianConditionalProbabilityPath 这个类来实现一个 高斯条件概率路径。然后我们会用它把一个简单的源分布 $p_{\text{simple}}=N(0,I_d)$ 变换为一个高斯混合分布 $p_{\text{data}}$ 。稍后我们还会实验一些更有趣的分布。

回顾一下，高斯条件概率路径定义为：

$$p_t(x|z)=N(x;{\alpha }_{t}z,{\beta }_t^2{I}_d),p_{\text{simple}}=N(0,I_d),$$

其中 ${\alpha }_t:[0,1]\to \mathbb{R}$ 和 ${\beta }_t:[0,1]\to \mathbb{R}$ 是单调、连续可微的函数，并满足：

$${\alpha }_1={\beta }_0=1,{\alpha }_0={\beta }_1=0.$$

换句话说，这意味着：

$$p_1(x|z)={\delta }_z$$

并且

$$p_0(x|z)=N(0,I_d)$$

是一个单位高斯分布。

在正式开始之前，我们先来看一看 $p_{\text{simple}}$ 和 $p_{\text{data}}$ 。

# Constants for the duration of our use of Gaussian conditional probability paths, to avoid polluting the namespace...
PARAMS = {
    "scale": 15.0,
    "target_scale": 10.0,
    "target_std": 1.0,
}


p_simple = Gaussian.isotropic(dim=2, std = 1.0).to(device)
p_data = GaussianMixture.symmetric_2D(nmodes=5, std=PARAMS["target_std"], scale=PARAMS["target_scale"]).to(device)

fig, axes = plt.subplots(1,3, figsize=(24,8))
bins = 200

scale = PARAMS["scale"]
x_bounds = [-scale,scale]
y_bounds = [-scale,scale]

axes[0].set_title('Heatmap of p_simple')
axes[0].set_xticks([])
axes[0].set_yticks([])
imshow_density(density=p_simple, x_bounds=x_bounds, y_bounds=y_bounds, bins=200, ax=axes[0], vmin=-10, alpha=0.25, cmap=plt.get_cmap('Reds'))


axes[1].set_title('Heatmap of p_data')
axes[1].set_xticks([])
axes[1].set_yticks([])
imshow_density(density=p_data, x_bounds=x_bounds, y_bounds=y_bounds, bins=200, ax=axes[1], vmin=-10, alpha=0.25, cmap=plt.get_cmap('Blues'))

axes[2].set_title('Heatmap of p_simple and p_data')
axes[2].set_xticks([])
axes[2].set_yticks([])
imshow_density(density=p_simple, x_bounds=x_bounds, y_bounds=y_bounds, bins=200, vmin=-10, alpha=0.25, cmap=plt.get_cmap('Reds'))
imshow_density(density=p_data, x_bounds=x_bounds, y_bounds=y_bounds, bins=200, vmin=-10, alpha=0.25, cmap=plt.get_cmap('Blues'))

可视化的结果如下图所示：

从上图我们可以看到：

• 左图 p_simple 是一个位于原点附近的单峰高斯，形状规则、对称、集中；
• 中图 p_data 是一个 5 模态的对称混合高斯，五个蓝色高密度区域均匀分布在中心周围；
• 右图叠加结果更直观地展示了二者差异：源分布在中心，目标分布在外围多峰结构上。

Problem 2.1: Implementing ${\alpha }_t$ and ${\beta }_t$

我们先从实现 ${\alpha }_t$ 和 ${\beta }_t$ 开始。可以把它们简单理解为两个可调用对象，它们需要满足基本约束：

$${\alpha }_1={\beta }_0=1$$

以及

$${\alpha }_0={\beta }_1=0$$

同时，它们还需要能够计算各自关于时间的导数 ${\dot{\alpha }}_t$ 和 ${\dot{\beta }}_t$ 。下面我们通过 Alpha 和 Beta 两个类来实现它们。

class Alpha(ABC):
    def __init__(self):
        # Check alpha_t(0) = 0
        assert torch.allclose(
            self(torch.zeros(1,1)), torch.zeros(1,1)
        )
        # Check alpha_1 = 1
        assert torch.allclose(
            self(torch.ones(1,1)), torch.ones(1,1)
        )
        
    @abstractmethod
    def __call__(self, t: torch.Tensor) -> torch.Tensor:
        """
        Evaluates alpha_t. Should satisfy: self(0.0) = 0.0, self(1.0) = 1.0.
        Args:
            - t: time (num_samples, 1)
        Returns:
            - alpha_t (num_samples, 1)
        """ 
        pass

    def dt(self, t: torch.Tensor) -> torch.Tensor:
        """
        Evaluates d/dt alpha_t.
        Args:
            - t: time (num_samples, 1)
        Returns:
            - d/dt alpha_t (num_samples, 1)
        """ 
        t = t.unsqueeze(1) # (num_samples, 1, 1)
        dt = vmap(jacrev(self))(t) # (num_samples, 1, 1, 1, 1)
        return dt.view(-1, 1)
    
class Beta(ABC):
    def __init__(self):
        # Check beta_0 = 1
        assert torch.allclose(
            self(torch.zeros(1,1)), torch.ones(1,1)
        )
        # Check beta_1 = 0
        assert torch.allclose(
            self(torch.ones(1,1)), torch.zeros(1,1)
        )
        
    @abstractmethod
    def __call__(self, t: torch.Tensor) -> torch.Tensor:
        """
        Evaluates alpha_t. Should satisfy: self(0.0) = 1.0, self(1.0) = 0.0.
        Args:
            - t: time (num_samples, 1)
        Returns:
            - beta_t (num_samples, 1)
        """ 
        pass 

    def dt(self, t: torch.Tensor) -> torch.Tensor:
        """
        Evaluates d/dt beta_t.
        Args:
            - t: time (num_samples, 1)
        Returns:
            - d/dt beta_t (num_samples, 1)
        """ 
        t = t.unsqueeze(1) # (num_samples, 1, 1)
        dt = vmap(jacrev(self))(t) # (num_samples, 1, 1, 1, 1)
        return dt.view(-1, 1)

这段代码定义了两个抽象类：Alpha 和 Beta。它们分别对应高斯条件概率路径中的两个时间调度函数：

$$p_t(x|z)=N(x;{\alpha }_{t}z,{\beta }_t^2{I}_d)$$

其中， ${\alpha }_t$ 控制目标数据点 $z$ 对当前分布均值的影响， ${\beta }_t$ 控制当前分布的噪声强度。

在这一节中，我们将使用：

$${\alpha }_t=t\text{and}{\beta }_t=\sqrt{1-t}.$$

不难验证，这两个函数在 $[0,1)$ 上都是连续可微且单调的，并且满足：

$${\alpha }_1={\beta }_0=1,$$

以及

$${\alpha }_0={\beta }_1=0.$$

我们的任务是补全下面 LinearAlpha 和 SquareRootBeta 类中的 __call__ 方法。

class LinearAlpha(Alpha):
    """
    Implements alpha_t = t
    """
    
    def __call__(self, t: torch.Tensor) -> torch.Tensor:
        """
        Args:
            - t: time (num_samples, 1)
        Returns:
            - alpha_t (num_samples, 1)
        """ 
        return t
        
    def dt(self, t: torch.Tensor) -> torch.Tensor:
        """
        Evaluates d/dt alpha_t.
        Args:
            - t: time (num_samples, 1)
        Returns:
            - d/dt alpha_t (num_samples, 1)
        """ 
        return torch.ones_like(t)

class SquareRootBeta(Beta):
    """
    Implements beta_t = rt(1-t)
    """
    def __call__(self, t: torch.Tensor) -> torch.Tensor:
        """
        Args:
            - t: time (num_samples, 1)
        Returns:
            - beta_t (num_samples, 1)
        """ 
        return torch.sqrt(1 - t)

    def dt(self, t: torch.Tensor) -> torch.Tensor:
        """
        Evaluates d/dt alpha_t.
        Args:
            - t: time (num_samples, 1)
        Returns:
            - d/dt alpha_t (num_samples, 1)
        """ 
        return - 0.5 / (torch.sqrt(1 - t) + 1e-4)

这两个函数的作用可以概括为：

• ${\alpha }_t=t$ ：随着时间增加，逐渐把均值推向数据点 $z$
• ${\beta }_t=\sqrt{1-t}$ ：随着时间增加，逐渐减小噪声强度

因此条件路径：

$$p_t(x|z)=N(x;{\alpha }_{t}z,{\beta }_t^2{I}_d)$$

会从：

$$p_0(x|z)=N(0,I_d)$$

逐渐过渡到：

$$p_1(x|z)={\delta }_z(x)$$

现在让我们转向实现 GaussianConditionalProbabilityPath 路径。

class GaussianConditionalProbabilityPath(ConditionalProbabilityPath):
    def __init__(self, p_data: Sampleable, alpha: Alpha, beta: Beta):
        p_simple = Gaussian.isotropic(p_data.dim, 1.0)
        super().__init__(p_simple, p_data)
        self.alpha = alpha
        self.beta = beta

    def sample_conditioning_variable(self, num_samples: int) -> torch.Tensor:
        """
        Samples the conditioning variable z ~ p_data(x)
        Args:
            - num_samples: the number of samples
        Returns:
            - z: samples from p(z), (num_samples, dim)
        """
        return self.p_data.sample(num_samples)
    
    def sample_conditional_path(self, z: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
        """
        Samples from the conditional distribution p_t(x|z) = N(alpha_t * z, beta_t**2 * I_d)
        Args:
            - z: conditioning variable (num_samples, dim)
            - t: time (num_samples, 1)
        Returns:
            - x: samples from p_t(x|z), (num_samples, dim)
        """
        raise NotImplementedError("Fill me in for Question 2.2!")
        
    def conditional_vector_field(self, x: torch.Tensor, z: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
        """
        Evaluates the conditional vector field u_t(x|z)
        Note: Only defined on t in [0,1)
        Args:
            - x: position variable (num_samples, dim)
            - z: conditioning variable (num_samples, dim)
            - t: time (num_samples, 1)
        Returns:
            - conditional_vector_field: conditional vector field (num_samples, dim)
        """ 
        raise NotImplementedError("Fill me in for Question 2.3!")

    def conditional_score(self, x: torch.Tensor, z: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
        """
        Evaluates the conditional score of p_t(x|z) = N(alpha_t * z, beta_t**2 * I_d)
        Note: Only defined on t in [0,1)
        Args:
            - x: position variable (num_samples, dim)
            - z: conditioning variable (num_samples, dim)
            - t: time (num_samples, 1)
        Returns:
            - conditional_score: conditional score (num_samples, dim)
        """ 
        raise NotImplementedError("Fill me in for Question 2.4!")

这个类实现的是高斯概率路径：

1. 初始化函数

def __init__(self, p_data: Sampleable, alpha: Alpha, beta: Beta):
    p_simple = Gaussian.isotropic(p_data.dim, 1.0)
    super().__init__(p_simple, p_data)
    self.alpha = alpha
    self.beta = beta

这里根据 p_data.dim 自动构造一个同维度的标准高斯分布：

$$p_{\text{simple}}=N(0,I_d)$$

然后把 p_simple 和 p_data 传给父类 ConditionalProbabilityPath。

alpha 和 beta 分别控制路径中的均值变化和噪声变化，也就是：

$${\alpha }_t,{\beta }_t$$

2. sample_conditioning_variable

这个函数的作用是采样条件变量：

$$z\sim p_{\text{data}}$$

也就是说，每条条件路径都以一个真实数据样本 $z$ 作为最终目标。

3. sample_conditional_path

这个函数的作用是采样条件路径，我们将在 Question 2.2 中实现。

4. sample_conditional_path

这个函数的作用是计算条件向量场，我们将在 Question 2.3 中实现。

5. conditional_score

这个函数的作用是计算条件 score，我们将在 Question 2.4 中实现。

Problem 2.2: Gaussian Conditional Probability Path

我们的任务是实现类方法 sample_conditional_path，用于从条件分布

$$p_t(x|z)=N(x;{\alpha }_{t}z,{\beta }_t^2{I}_d)$$

中采样。

提示：你可以使用下面这个事实：如果随机变量 $X\sim N(\mu ,{\sigma }^2{I}_d)$ ，那么它可以通过如下方式得到：$X=\mu +\sigma Z,Z\sim N(0,I_d)$ 。

实现代码如下：

def sample_conditional_path(self, z: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
    """
    Samples from the conditional distribution p_t(x|z) = N(alpha_t * z, beta_t**2 * I_d)
    Args:
        - z: conditioning variable (num_samples, dim)
        - t: time (num_samples, 1)
    Returns:
        - x: samples from p_t(x|z), (num_samples, dim)
    """
    epsilon = torch.randn_like(z)
    return self.alpha(t) * z + self.beta(t) * epsilon

这一题的目标是从下面这个条件高斯分布中采样：

$$p_t(x|z)=N(x;{\alpha }_{t}z,{\beta }_t^2{I}_d)$$

也就是说，在给定数据点 $z$ 和时间 $t$ 时， $x_t$ 服从一个高斯分布，其中均值为 ${\alpha }_{t}z$ ，协方差为 ${\beta }_t^2{I}_d$ 。

根据重参数采样公式，如果：

$$X\sim N(\mu ,{\sigma }^2{I}_d)$$

那么可以写成：

$$X=\mu +\sigma ϵ,ϵ\sim N(0,I_d)$$

在当前问题中：

$$\mu ={\alpha }_{t}z,\sigma ={\beta }_t$$

因此：

$$x_t={\alpha }_{t}z+{\beta }_tϵ,ϵ\sim N(0,I_d)$$

对应代码就是：

epsilon = torch.randn_like(z)
return self.alpha(t) * z + self.beta(t) * epsilon

其中 torch.randn_like(z) 会生成一个和 z 形状相同的标准高斯噪声，形状为：

(num_samples, dim)

而 self.alpha(t) 和 self.beta(t) 的形状是：

(num_samples, 1)

它们和 z 的形状 (num_samples, dim) 相乘时，会自动广播到每个维度上。

现在我们已经可以从条件概率路径中进行采样，因此也就可以将这个条件概率路径可视化出来了。

# Construct conditional probability path
path = GaussianConditionalProbabilityPath(
    p_data = GaussianMixture.symmetric_2D(nmodes=5, std=PARAMS["target_std"], scale=PARAMS["target_scale"]).to(device), 
    alpha = LinearAlpha(),
    beta = SquareRootBeta()
).to(device)

scale = PARAMS["scale"]
x_bounds = [-scale,scale]
y_bounds = [-scale,scale]

plt.figure(figsize=(10,10))
plt.xlim(*x_bounds)
plt.ylim(*y_bounds)
plt.title('Gaussian Conditional Probability Path')

# Plot source and target
imshow_density(density=p_simple, x_bounds=x_bounds, y_bounds=y_bounds, bins=200, vmin=-10, alpha=0.25, cmap=plt.get_cmap('Reds'))
imshow_density(density=p_data, x_bounds=x_bounds, y_bounds=y_bounds, bins=200, vmin=-10, alpha=0.25, cmap=plt.get_cmap('Blues'))

# Sample conditioning variable z
z = path.sample_conditioning_variable(1) # (1,2)
ts = torch.linspace(0.0, 1.0, 7).to(device)

# Plot z
plt.scatter(z[:,0].cpu(), z[:,1].cpu(), marker='*', color='red', s=75, label='z')
plt.xticks([])
plt.yticks([])

# Plot conditional probability path at each intermediate t
num_samples = 1000
for t in ts:
    zz = z.expand(num_samples, 2)
    tt = t.unsqueeze(0).expand(num_samples, 1) # (samples, 1)
    samples = path.sample_conditional_path(zz, tt) # (samples, 2)
    plt.scatter(samples[:,0].cpu(), samples[:,1].cpu(), alpha=0.25, s=8, label=f't={t.item():.1f}')

plt.legend(prop={'size': 18}, markerscale=3)
plt.show()

可视化结果如下图所示：

从上图中那我们可以清楚看到如下现象：

1. 当 $t=0$ 时

样本云集中在原点附近，而且分散程度较大。这与理论一致，因为：

$${\alpha }_0=0,{\beta }_0=1$$

所以：

$$x_0\sim N(0,I_d)$$

也就是标准高斯。

2. 随着 $t$ 增大

样本云整体逐渐朝着红色星号 $z$ 的方向移动，同时分布范围逐渐缩小。

这是因为：

• 均值部分 ${\alpha }_{t}z=tz$ 随时间逐渐把样本中心从 0 推向 $z$ ;
• 方差部分 ${\beta }_t^2=1-t$ 随时间减小，因此噪声逐渐变弱，样本云越来越集中。

3. 当 $t=1$ 时

样本云几乎完全坍缩到红色星号附近，理论上此时：

$${\alpha }_1=1,{\beta }_1=0$$

所以：

$$x_1=z$$

即：

$$p_1(x|z)={\delta }_z$$

这说明我们的 sample_conditional_path 实现是正确的：条件路径确实从标准高斯逐渐收缩到了指定的数据点。

Problem 2.3: Conditional Vector Field

根据课堂和课程笔记，我们知道条件向量场 $u_t(x|z)$ 可以写成：

$$u_t(x|z)=\left({\dot{\alpha }}_t-\frac{{\dot{\beta }}_t}{{\beta }_t}{\alpha }_t\right)z+\frac{{\dot{\beta }}_t}{{\beta }_t}x.$$

我们的任务是实现类方法 conditional_vector_field，用于计算条件向量场 $u_t(x|z)$ 。

提示：你可以使用已经为你实现好的 self.alpha.dt(t) 来计算 ${\dot{\alpha }}_t$ 。类似地，也可以用同样的方式计算 ${\dot{\beta }}_t$ 。

实现代码如下：

def conditional_vector_field(self, x: torch.Tensor, z: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
    """
    Evaluates the conditional vector field u_t(x|z)
    Note: Only defined on t in [0,1)
    Args:
        - x: position variable (num_samples, dim)
        - z: conditioning variable (num_samples, dim)
        - t: time (num_samples, 1)
    Returns:
        - conditional_vector_field: conditional vector field (num_samples, dim)
    """ 
    alpha_t = self.alpha(t)
    beta_t  = self.beta(t)
    dt_alpha_t = self.alpha.dt(t)
    dt_beta_t  = self.beta.dt(t)

    return (dt_alpha_t - dt_beta_t / beta_t * alpha_t) * z + (dt_beta_t / beta_t) * x

条件概率路径为：

$$p_t(x|z)=N(x;{\alpha }_{t}z,{\beta }_t^2{I}_d)$$

我们可以通过重参数化写成：

$$x_t={\alpha }_{t}z+{\beta }_tϵ,ϵ\sim N(0,I_d)$$

对时间 $t$ 求导：

$$\frac{d{x}_t}{dt}={\dot{\alpha }}_{t}z+{\dot{\beta }}_tϵ.$$

但是条件向量场 $u_t(x|z)$ 应该写成关于当前位置 $x$、条件变量 $z$ 和时间 $t$ 的形式，而不是关于噪声 $ϵ$ 的形式。

由：

$$x_t={\alpha }_{t}z+{\beta }_tϵ$$

可以得到：

$$ϵ=\frac{x-{\alpha }_{t}z}{{\beta }_t}.$$

代入前面的导数公式：

$$\frac{d{x}_t}{dt}={\dot{\alpha }}_{t}z+{\dot{\beta }}_t\frac{x-{\alpha }_{t}z}{{\beta }_t}.$$

整理后就是：

$$u_t(x|z)=\left({\dot{\alpha }}_t-\frac{{\dot{\beta }}_t}{{\beta }_t}{\alpha }_t\right)z+\frac{{\dot{\beta }}_t}{{\beta }_t}x.$$

代码中：

alpha_t = self.alpha(t)
beta_t  = self.beta(t)
dt_alpha_t = self.alpha.dt(t)
dt_beta_t  = self.beta.dt(t)

分别计算：

$${\alpha }_t,{\beta }_t,{\dot{\alpha }}_t,{\dot{\beta }}_t$$

最后：

return (dt_alpha_t - dt_beta_t / beta_t * alpha_t) * z + (dt_beta_t / beta_t) * x

就是直接按照题目给出的条件向量场公式实现。

现在我们可以对下面这个 ODE 对应的条件轨迹进行可视化：

$$d{X}_t=u_t(X_t|z)dt,X_0=x_0\sim p_{\text{simple}}.$$

class ConditionalVectorFieldODE(ODE):
    def __init__(self, path: ConditionalProbabilityPath, z: torch.Tensor):
        """
        Args:
        - path: the ConditionalProbabilityPath object to which this vector field corresponds
        - z: the conditioning variable, (1, dim)
        """
        super().__init__()
        self.path = path
        self.z = z

    def drift_coefficient(self, x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
        """
        Returns the conditional vector field u_t(x|z)
        Args:
            - x: state at time t, shape (bs, dim)
            - t: time, shape (bs,.)
        Returns:
            - u_t(x|z): shape (batch_size, dim)
        """
        bs = x.shape[0]
        z = self.z.expand(bs, *self.z.shape[1:])
        return self.path.conditional_vector_field(x,z,t)


# Run me for Problem 2.3!

#######################
# Change these values #
#######################
num_samples = 1000
num_timesteps = 1000
num_marginals = 3

########################
# Setup path and plot  #
########################

path = GaussianConditionalProbabilityPath(
    p_data = GaussianMixture.symmetric_2D(nmodes=5, std=PARAMS["target_std"], scale=PARAMS["target_scale"]).to(device), 
    alpha = LinearAlpha(),
    beta = SquareRootBeta()
).to(device)


# Setup figure
fig, axes = plt.subplots(1,3, figsize=(36, 12))
scale = PARAMS["scale"]
legend_size = 24
markerscale = 1.8
x_bounds = [-scale,scale]
y_bounds = [-scale,scale]

# Sample conditioning variable z
torch.cuda.manual_seed(1)
z = path.sample_conditioning_variable(1) # (1,2)

######################################
# Graph samples from conditional ODE #
######################################
ax = axes[1]

ax.set_xlim(*x_bounds)
ax.set_ylim(*y_bounds)
ax.set_xticks([])
ax.set_yticks([])
ax.set_title('Samples from Conditional ODE', fontsize=20)
ax.scatter(z[:,0].cpu(), z[:,1].cpu(), marker='*', color='red', s=200, label='z',zorder=20) # Plot z

# Plot source and target
imshow_density(density=p_simple, x_bounds=x_bounds, y_bounds=y_bounds, bins=200, ax=ax, vmin=-10, alpha=0.25, cmap=plt.get_cmap('Reds'))
imshow_density(density=p_data, x_bounds=x_bounds, y_bounds=y_bounds, bins=200, ax=ax, vmin=-10, alpha=0.25, cmap=plt.get_cmap('Blues'))


# Construct integrator and plot trajectories
sigma = 0.5 # Can't make this too high or integration is numerically unstable!
ode = ConditionalVectorFieldODE(path, z)
simulator = EulerSimulator(ode)
x0 = path.p_simple.sample(num_samples) # (num_samples, 2)
ts = torch.linspace(0.0, 1.0, num_timesteps).view(1,-1,1).expand(num_samples,-1,1).to(device) # (num_samples, nts, 1)
xts = simulator.simulate_with_trajectory(x0, ts) # (bs, nts, dim)

# Extract every n-th integration step to plot
every_n = record_every(num_timesteps=num_timesteps, record_every=num_timesteps // num_marginals)
xts_every_n = xts[:,every_n,:] # (bs, nts // n, dim)
ts_every_n = ts[0,every_n] # (nts // n,)
for plot_idx in range(xts_every_n.shape[1]):
    tt = ts_every_n[plot_idx].item()
    ax.scatter(xts_every_n[:,plot_idx,0].detach().cpu(), xts_every_n[:,plot_idx,1].detach().cpu(), marker='o', alpha=0.5, label=f't={tt:.2f}')
ax.legend(prop={'size': legend_size}, loc='upper right', markerscale=markerscale)


#########################################
# Graph Trajectories of Conditional ODE #
#########################################
ax = axes[2]

ax.set_xlim(*x_bounds)
ax.set_ylim(*y_bounds)
ax.set_xticks([])
ax.set_yticks([])
ax.set_title('Trajectories of Conditional ODE', fontsize=20)
ax.scatter(z[:,0].cpu(), z[:,1].cpu(), marker='*', color='red', s=200, label='z',zorder=20) # Plot z


# Plot source and target
imshow_density(density=p_simple, x_bounds=x_bounds, y_bounds=y_bounds, bins=200, ax=ax, vmin=-10, alpha=0.25, cmap=plt.get_cmap('Reds'))
imshow_density(density=p_data, x_bounds=x_bounds, y_bounds=y_bounds, bins=200, ax=ax, vmin=-10, alpha=0.25, cmap=plt.get_cmap('Blues'))

for traj_idx in range(15):
    ax.plot(xts[traj_idx,:,0].detach().cpu(), xts[traj_idx,:,1].detach().cpu(), alpha=0.5, color='black')
ax.legend(prop={'size': legend_size}, loc='upper right', markerscale=markerscale)


###################################################
# Graph Ground-Truth Conditional Probability Path #
###################################################
ax = axes[0]

ax.set_xlim(*x_bounds)
ax.set_ylim(*y_bounds)
ax.set_xticks([])
ax.set_yticks([])
ax.set_title('Ground-Truth Conditional Probability Path', fontsize=20)
ax.scatter(z[:,0].cpu(), z[:,1].cpu(), marker='*', color='red', s=200, label='z',zorder=20) # Plot z


for plot_idx in range(xts_every_n.shape[1]):
    tt = ts_every_n[plot_idx].unsqueeze(0).expand(num_samples, 1)
    zz = z.expand(num_samples, 2)
    marginal_samples = path.sample_conditional_path(zz, tt)
    ax.scatter(marginal_samples[:,0].detach().cpu(), marginal_samples[:,1].detach().cpu(), marker='o', alpha=0.5, label=f't={tt[0,0].item():.2f}')

# Plot source and target
imshow_density(density=p_simple, x_bounds=x_bounds, y_bounds=y_bounds, bins=200, ax=ax, vmin=-10, alpha=0.25, cmap=plt.get_cmap('Reds'))
imshow_density(density=p_data, x_bounds=x_bounds, y_bounds=y_bounds, bins=200, ax=ax, vmin=-10, alpha=0.25, cmap=plt.get_cmap('Blues'))
ax.legend(prop={'size': legend_size}, loc='upper right', markerscale=markerscale)

plt.show()

这段代码的核心目标是验证一个重要结论：

如果我们用前面解析得到的条件向量场

$$u_t(x|z)$$

去定义 ODE，并从 $p_{\text{simple}}$ 出发进行积分，那么在任意时间 $t$ 上得到的样本分布，应当与真实条件概率路径

$$p_t(x|z)$$

相一致。

换句话说，这段代码是在做一个 “条件路径 = 条件 ODE 生成分布” 的数值验证。

可视化结果如下图所示：

从上图中我们可以看出：

1. 左图与中图非常接近

• 左图：直接从真实条件路径 $p_t(x|z)$ 采样
• 中图：通过积分 conditional ODE 得到样本

两者在 $t=0.00,0.33,0.67,1.00$ 的样本云位置和扩散程度都非常接近。

这说明我们前面实现的 conditional_vector_field 是正确的：它确实生成了对应条件概率路径的流。

2. 右图显示出轨迹几何

右图给出了个体轨迹视角，可以看到：

• 样本从中心附近出发
• 随着时间推进逐渐向红色星号 $z$ 聚集
• 轨迹整体是平滑的、有方向性的

这正符合 ODE 轨迹的特点。

这一步验证了一个核心思想：

$$p_t(x|z)⟺d{X}_t=u_t(X_t|z)dt$$

也就是说，给定一条条件概率路径，我们不仅能写出它的分布形式，还能写出一个对应的条件向量场，使得从简单分布出发沿 ODE 积分，恰好能够在每个时间点重现这条路径。

这正是后面 Flow Matching 的基础。

注意：你可能已经注意到，对于高斯概率路径，由于 $z\sim p_{\text{data}}(x)$ ，因此方法 GaussianConditionalProbabilityPath.sample_conditioning_variable 实际上是在从数据分布中采样。

但是等等 — 我们一开始不正是想学习如何从 $p_{\text{data}}$ 进行采样吗？这确实是一个我们到目前为止刻意略过的细节。

答案是：在实际应用中，sample_conditioning_variable 返回的并不是真正从连续的 $p_{\text{data}}$ 中现采的点，而是来自一个有限的 训练数据集 的样本。形式上，我们假设这个训练集中的样本是从真实分布 $z\sim p_{\text{data}}$ 中独立同分布（IID）采样得到的。

Problem 2.4: The Conditional Score

和课堂中一样，现在我们可以可视化下面这个 SDE 所对应的条件轨迹：

$$d{X}_t=\left[u_t(X_t|z)+\frac{1}{2}{\sigma }_t^2{\nabla }_x\log p_t(X_t|z)\right]dt+\sigma ,d{W}_t,X_0=x_0\sim p_{\text{simple}},$$

它可以看作是在原来的 ODE 上加入了 Langevin dynamics 得到的结果。

我们的任务是实现类方法 conditional_score，用于计算条件分布的 score：

$${\nabla }_x\log p_t(x|z)$$

我们可以计算得到：

$${\nabla }_x\log p_t(x|z)={\nabla }_x\log N(x;{\alpha }_{t}z,{\beta }_t^2{I}_d)=\frac{{\alpha }_{t}z-x}{{\beta }_t^2}.$$

实现代码如下：

def conditional_score(self, x: torch.Tensor, z: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
    """
    Evaluates the conditional score of p_t(x|z) = N(alpha_t * z, beta_t**2 * I_d)
    Note: Only defined on t in [0,1)
    Args:
        - x: position variable (num_samples, dim)
        - z: conditioning variable (num_samples, dim)
        - t: time (num_samples, 1)
    Returns:
        - conditional_score: conditional score (num_samples, dim)
    """ 
    alpha_t = self.alpha(t)
    beta_t  = self.beta(t)
    return (alpha_t * z - x) / (beta_t ** 2)

条件概率路径为：

$$p_t(x|z)=N(x;{\alpha }_{t}z,{\beta }_t^2{I}_d)$$

这是一个均值为 ${\alpha }_{t}z$ ，协方差为 ${\beta }_t^2{I}_d$ 的高斯分布。

对于高斯分布：

$$N(x;\mu ,{\sigma }^2{I}_d)$$

它的 score 是：

$${\nabla }_x\log p(x)=-\frac{x-\mu }{{\sigma }^2}=\frac{\mu -x}{{\sigma }^2}.$$

在当前问题中，代入：

$$\mu ={\alpha }_{t}z,\sigma ={\beta }_t,$$

得到：

$${\nabla }_x\log p_t(x|z)=\frac{{\alpha }_{t}z-x}{{\beta }_t^2}.$$

对应代码就是：

alpha_t = self.alpha(t)
beta_t  = self.beta(t)
return (alpha_t * z - x) / (beta_t ** 2)

这里 alpha_t 和 beta_t 的形状是：

(num_samples, 1)

而 x 和 z 的形状是：

(num_samples, dim)

PyTorch 会自动广播，使每个样本在自己的时间 t 下计算对应的 score。

为了检查实现是否正确，我们会使用下面的代码，验证从条件 SDE 中采样得到的样本，是否与直接从条件概率路径解析采样得到的样本相匹配。

class ConditionalVectorFieldSDE(SDE):
    def __init__(self, path: ConditionalProbabilityPath, z: torch.Tensor, sigma: float):
        """
        Args:
        - path: the ConditionalProbabilityPath object to which this vector field corresponds
        - z: the conditioning variable, (1, ...)
        """
        super().__init__()
        self.path = path
        self.z = z
        self.sigma = sigma

    def drift_coefficient(self, x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
        """
        Returns the conditional vector field u_t(x|z)
        Args:
            - x: state at time t, shape (bs, dim)
            - t: time, shape (bs,.)
        Returns:
            - u_t(x|z): shape (batch_size, dim)
        """
        bs = x.shape[0]
        z = self.z.expand(bs, *self.z.shape[1:])
        return self.path.conditional_vector_field(x,z,t) + 0.5 * self.sigma**2 * self.path.conditional_score(x,z,t)

    def diffusion_coefficient(self, x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
        """
        Args:
            - x: state at time t, shape (bs, dim)
            - t: time, shape (bs,.)
        Returns:
            - u_t(x|z): shape (batch_size, dim)
        """
        return self.sigma * torch.randn_like(x)


# Run me for Problem 2.3!

#######################
# Change these values #
#######################
num_samples = 1000
num_timesteps = 1000
num_marginals = 3
sigma = 2.5

########################
# Setup path and plot  #
########################

path = GaussianConditionalProbabilityPath(
    p_data = GaussianMixture.symmetric_2D(nmodes=5, std=PARAMS["target_std"], scale=PARAMS["target_scale"]).to(device), 
    alpha = LinearAlpha(),
    beta = SquareRootBeta()
).to(device)


# Setup figure
fig, axes = plt.subplots(1,3, figsize=(36, 12))
scale = PARAMS["scale"]
x_bounds = [-scale,scale]
y_bounds = [-scale,scale]
legend_size = 24
markerscale = 1.8

# Sample conditioning variable z
torch.cuda.manual_seed(1)
z = path.sample_conditioning_variable(1) # (1,2)

######################################
# Graph Samples from Conditional SDE #
######################################
ax = axes[1]

ax.set_xlim(*x_bounds)
ax.set_ylim(*y_bounds)
ax.set_xticks([])
ax.set_yticks([])
ax.set_title('Samples from Conditional SDE', fontsize=20)
ax.scatter(z[:,0].cpu(), z[:,1].cpu(), marker='*', color='red', s=200, label='z',zorder=20) # Plot z

# Plot source and target
imshow_density(density=p_simple, x_bounds=x_bounds, y_bounds=y_bounds, bins=200, ax=ax, vmin=-10, alpha=0.25, cmap=plt.get_cmap('Reds'))
imshow_density(density=p_data, x_bounds=x_bounds, y_bounds=y_bounds, bins=200, ax=ax, vmin=-10, alpha=0.25, cmap=plt.get_cmap('Blues'))


# Construct integrator and plot trajectories
sde = ConditionalVectorFieldSDE(path, z, sigma)
simulator = EulerMaruyamaSimulator(sde)
x0 = path.p_simple.sample(num_samples) # (num_samples, 2)
ts = torch.linspace(0.0, 1.0, num_timesteps).view(1,-1,1).expand(num_samples,-1,1).to(device) # (num_samples, nts, 1)
xts = simulator.simulate_with_trajectory(x0, ts) # (bs, nts, dim)

# Extract every n-th integration step to plot
every_n = record_every(num_timesteps=num_timesteps, record_every=num_timesteps // num_marginals)
xts_every_n = xts[:,every_n,:] # (bs, nts // n, dim)
ts_every_n = ts[0,every_n] # (nts // n,)
for plot_idx in range(xts_every_n.shape[1]):
    tt = ts_every_n[plot_idx].item()
    ax.scatter(xts_every_n[:,plot_idx,0].detach().cpu(), xts_every_n[:,plot_idx,1].detach().cpu(), marker='o', alpha=0.5, label=f't={tt:.2f}')
ax.legend(prop={'size': legend_size}, loc='upper right', markerscale=markerscale)


##########################################
# Graph Trajectories of Conditional SDE  #
##########################################
ax = axes[2]

ax.set_xlim(*x_bounds)
ax.set_ylim(*y_bounds)
ax.set_xticks([])
ax.set_yticks([])
ax.set_title('Trajectories of Conditional SDE', fontsize=20)
ax.scatter(z[:,0].cpu(), z[:,1].cpu(), marker='*', color='red', s=200, label='z',zorder=20) # Plot z


# Plot source and target
imshow_density(density=p_simple, x_bounds=x_bounds, y_bounds=y_bounds, bins=200, ax=ax, vmin=-10, alpha=0.25, cmap=plt.get_cmap('Reds'))
imshow_density(density=p_data, x_bounds=x_bounds, y_bounds=y_bounds, bins=200, ax=ax, vmin=-10, alpha=0.25, cmap=plt.get_cmap('Blues'))

for traj_idx in range(5):
    ax.plot(xts[traj_idx,:,0].detach().cpu(), xts[traj_idx,:,1].detach().cpu(), alpha=0.5, color='black')
ax.legend(prop={'size': legend_size}, loc='upper right', markerscale=markerscale)


###################################################
# Graph Ground-Truth Conditional Probability Path #
###################################################
ax = axes[0]

ax.set_xlim(*x_bounds)
ax.set_ylim(*y_bounds)
ax.set_xticks([])
ax.set_yticks([])
ax.set_title('Ground-Truth Conditional Probability Path', fontsize=20)
ax.scatter(z[:,0].cpu(), z[:,1].cpu(), marker='*', color='red', s=200, label='z',zorder=20) # Plot z


for plot_idx in range(xts_every_n.shape[1]):
    tt = ts_every_n[plot_idx].unsqueeze(0).expand(num_samples, 1)
    zz = z.expand(num_samples, 2)
    marginal_samples = path.sample_conditional_path(zz, tt)
    ax.scatter(marginal_samples[:,0].detach().cpu(), marginal_samples[:,1].detach().cpu(), marker='o', alpha=0.5, label=f't={tt[0,0].item():.2f}')

# Plot source and target
imshow_density(density=p_simple, x_bounds=x_bounds, y_bounds=y_bounds, bins=200, ax=ax, vmin=-10, alpha=0.25, cmap=plt.get_cmap('Reds'))
imshow_density(density=p_data, x_bounds=x_bounds, y_bounds=y_bounds, bins=200, ax=ax, vmin=-10, alpha=0.25, cmap=plt.get_cmap('Blues'))
ax.legend(prop={'size': legend_size}, loc='upper right', markerscale=markerscale)

plt.show()

注意：你可能会发现，当 $\sigma$ 取较大值时，甚至当 $\sigma$ 不是特别大时，也会出现一些奇怪的现象。将

$${\nabla }_x\log p_t(x|z)=\frac{{\alpha }_{t}z-x}{{\beta }_t^2}$$

代入

$$d{X}_t=\left[u_t(X_t|z)+\frac{1}{2}{\sigma }^2{\nabla }_x\log p_t(X_t|z)\right]dt+\sigma d{W}_t$$

可以得到：

$$d{X}_t=\left[u_t(X_t|z)+\frac{1}{2}{\sigma }^2\left(\frac{{\alpha }_{t}z-X_t}{{\beta }_t^2}\right)\right]dt+\sigma d{W}_t.$$

当 $t\to 1$ 时，${\beta }_t\to 0$ ，因此漂移项中的第二项会爆炸，并且这种爆炸会随着 $\sigma$ 的平方增长。在有限数量的模拟步数下，我们无法准确模拟这种爆炸现象，因此会遇到数值问题。

在实践中，通常会通过设置例如 ${\sigma }_t={\beta }_t$ 来规避这个问题。这样一来，逐渐减小的噪声水平可以抵消这种爆炸效应。

可视化结果如下图所示：

从上图中我们可以知道：

左图不是通过 SDE 模拟得到的，而是直接从解析条件中采样：

$$p_t(x|z)=N({\alpha }_{t}z,{\beta }_t^2{I}_d)$$

因此它是 ground truth，用来对比 SDE 模拟结果是否正确。

中图展示的是通过 conditional SDE 模拟得到的样本。理论上，在数值积分足够精确的情况下，它应该和左图中的解析采样结果保持一致。

从图中可以看到，中图的样本云整体确实和左图比较接近：随着时间推进，样本从初始高斯逐渐移动并收缩到条件点 $z$ 附近。

右图展示了若干条单样本轨迹。和前面 conditional ODE 的平滑轨迹不同，这里轨迹明显带有随机抖动，这是因为 SDE 中加入了：

$$\sigma d{W}_t$$

所以每个样本不再是沿着确定性路径前进，而是在漂移项引导下，同时受到随机噪声扰动。

可视化图说明了直接解析采样得到的 conditional probability path 和通过 conditional SDE 模拟得到的样本分布整体上是匹配的。

这验证了我们的 conditional_score 实现是合理的。也就是说，将 score correction 加入原始条件向量场后，SDE 仍然能够保持同一条条件概率路径的边缘分布。

篇幅限制原因，Part 3 和 Part 4 的实现我们放在下篇文章中。