对多模态扩散模型UNet架构的探索

UNet将文本嵌入作为条件信息，与图像特征信息融合，在去噪时遵循文本描述生成相关的图像。

而扩散模型的训练过程，则是一个对"噪声残差"进行预测和优化的循环过程。

这里结合Unet组件和伪代码尝试说明这一过程。

1 条件信息融合

文本到图像生成模型中，UNet不仅接收带噪图像，还必须理解文本提示。

其关键组件及文本信息的融入方式如下表所示。

1.1 时间步嵌入

告知模型当前在去噪过程中的位置，以使用正确的策略。

时间步嵌入与文本条件独立，共同作为模型的情境信息。

具体为时间步条件影响权重缩放和偏置。

1.2 条件嵌入

将文本提示词编码为模型可理解的向量序列。这里文本通过独立的文本编码器（如CLIP）转换为嵌入向量。

1.3 空间变换器

空间变换器模块负责将文本条件与图像视觉特征进行深度融合。

残差后引入，通过交叉注意力机制，图像特征作为查询，文本嵌入作为键和值，进行特征对齐。

1.4 编码器-解码器

标准的Unet型结构，用于捕获多尺度图像特征并通过跳跃连接保留细节。

在每个层级（尤其在设定的attention_resolutions层）插入空间变换器，实现多层次的条件控制。

2 处理过程伪码

这里展示上述流程的核心逻辑，首先指明文本处理核心机制，然后通过代码示例处理过程。

2.1 文本处理机制

UNet在Transformer框架下统一处理图文数据的核心机制，并非直接处理文本信息，而是将文本信息作为一组全局的、语义丰富的键值对（Key/Value），通过交叉注意力机制持续地引导和调制图像特征（Query）的生成过程。

这里展示了在扩散模型UNet中，文本条件如何被整合的过程。

2.2 处理伪码示例

为简化分析，这里使用一层Transformer，实际情况下可能为多层Transformer。

示例代码如下，细节参考注释。

复制代码

import torch
import torch.nn as nn
import torch.nn.functional as F

class SpatialTransformer(nn.Module):
    """空间变换器模块，实现图像特征与文本条件的交叉注意力。"""
    def __init__(self, in_channels, context_dim):
        super().__init__()
        self.norm = nn.GroupNorm(32, in_channels)  # 归一化图像特征
        self.proj_in = nn.Conv2d(in_channels, in_channels, 1)
        # 核心：多头交叉注意力层
        self.attn = CrossAttention(query_dim=in_channels, context_dim=context_dim)
        self.proj_out = nn.Conv2d(in_channels, in_channels, 1)

    def forward(self, x, context):
        """
        x: 图像特征图 [B, C, H, W]
        context: 文本条件嵌入 [B, L, D]
        """
        batch, channel, height, width = x.shape
        residual = x

        # 1. 对图像特征进行归一化和投影
        x = self.norm(x)
        x = self.proj_in(x)

        # 2. 重塑为序列以进行注意力计算
        x = x.view(batch, channel, height * width).permute(0, 2, 1)  # [B, N, C]

        # 3. 交叉注意力：图像特征为Query，文本嵌入为Key/Value
        x = self.attn(x, context)

        # 4. 重塑回图像特征图并与残差连接
        x = x.permute(0, 2, 1).view(batch, channel, height, width)
        x = self.proj_out(x)
        return x + residual

class UNetConditionalBlock(nn.Module):
    """集成了空间变换器的UNet条件残差块。"""
    def __init__(self, in_channels, time_emb_dim, context_dim):
        super().__init__()
        # 时间步信息融入
        self.time_emb_proj = nn.Linear(time_emb_dim, in_channels * 2)
        # 第一个卷积组
        self.conv1 = nn.Conv2d(in_channels, in_channels, 3, padding=1)
        # 空间变换器（条件注入点）
        self.transformer = SpatialTransformer(in_channels, context_dim)
        # 第二个卷积组
        self.conv2 = nn.Conv2d(in_channels, in_channels, 3, padding=1)

    def forward(self, x, time_emb, context):
        # 时间步条件影响权重缩放和偏置
        scale, shift = self.time_emb_proj(time_emb).chunk(2, dim=1)
        x = x * (1 + scale[:, :, None, None]) + shift[:, :, None, None]

        # 主干卷积路径
        residual = x
        x = F.silu(self.conv1(x))
        # 关键步骤：将文本条件context注入图像特征
        x = self.transformer(x, context)
        x = F.silu(self.conv2(x))

        return x + residual