背景

方法大意

• 对于训练$$\epsilon_{\theta}(x_t, t)$$来估计$$x_t$$在时间$$t$$上添加的噪声，再根据采样公式推出$$x_{t-1}$$，从而实现图片生成。训练数据只需要准备图片即可。

• 对于classifier-guided DM是在普通DM的基础上，额外再训练一个Classifier来获得当前时间步生成的图片类别概率分布，从而实现特定类别的图片生成。

• 对于classifier-free DM将类别信息（或语义信息）集成到diffusion model的训练过程中，训练$$\epsilon_{\theta}(x_t, y,t)，\epsilon_{\theta}(x_t, y=\empty,t)(\text{即}\epsilon_{\theta}(x_t,t))$$。训练的时候也会加入无类别信息（或语义信息）的图片进行训练。

1. 模型如何融入类别信息（或语义信息）

2. 如何训练$$\epsilon_{\theta}(x_t, y,t)$$与$$\epsilon_{\theta}(x_t, y=\empty,t)$$

3. 如何进行采样生成

模型如何融入类别信息（或语义信息）

采用交叉注意力机制融入

class SpatialCrossAttention(nn.Module):

    def __init__(self, dim, context_dim, heads=4, dim_head=32) -> None:
        super(SpatialCrossAttention, self).__init__()
        self.scale = dim_head ** -0.5
        self.heads = heads
        hidden_dim = dim_head * heads

        self.proj_in = nn.Conv2d(dim, context_dim, kernel_size=1, stride=1, padding=0)
        self.to_q = nn.Linear(context_dim, hidden_dim, bias=False)
        self.to_k = nn.Linear(context_dim, hidden_dim, bias=False)
        self.to_v = nn.Linear(context_dim, hidden_dim, bias=False)
        self.to_out = nn.Conv2d(hidden_dim, dim, 1)

    def forward(self, x, context=None):

        x_q = self.proj_in(x)
        b, c, h, w = x_q.shape
        x_q = rearrange(x_q, "b c h w -> b (h w) c")
        if context is None:
            context = x_q
        if context.ndim == 2:
            context = rearrange(context, "b c -> b () c")
        q = self.to_q(x_q)
        k = self.to_k(context)
        v = self.to_v(context)
        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=self.heads), (q, k, v))
        sim = einsum('b i d, b j d -> b i j', q, k) * self.scale

        # attention, what we cannot get enough of
        attn = sim.softmax(dim=-1)

        out = einsum('b i j, b j d -> b i d', attn, v)
        out = rearrange(out, '(b h) n d -> b n (h d)', h=self.heads)
        out = rearrange(out, "b (h w) c -> b c h w", h=h, w=w)
        out = self.to_out(out)
        return out

基于channel-wise attention融入

# mixture time-embedding and label embedding
t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False)
emb = self.time_embed(t_emb)
if self.num_classes is not None:
    assert y.shape == (x.shape[0],)
    emb = emb + self.label_emb(y)
while len(emb_out.shape) < len(h.shape):
    emb_out = emb_out[..., None]
emb_out = self.emb_layers(emb).type(h.dtype)  # h is image feature
scale, shift = th.chunk(emb_out, 2, dim=1)  # One half of the embedding is used for scaling and the other half for offset
h = h * (1 + scale) + shift

如何训练$$\epsilon_{\theta}(x_t, y,t)$$与$$\epsilon_{\theta}(x_t, y=\emptyset,t)$$

• 在实践中，为了统一$$y=\empty$$和$$y \neq \empty$$ 两种情形，通常会给定一个$$y=\empty$$的embedding（可以随机初始化，也可以人为给定），来统一两种情形的建模。

• 即使所有的数据都有图片对也没有关系，只需在每一个batch中随机将某些数据的标签编码替换为$$y=\empty$$的embedding即可。另外

如何进行采样生成

结语

参考文献

附录