0. TL; DR

本文提出了一种新的可控图像生成方法——Condition-Aware Neural Network（CAN）。CAN通过动态调整神经网络的权重来控制图像生成过程，引入了一个条件感知权重生成模块，基于输入条件生成卷积/线性层的条件权重。在ImageNet和COCO数据集上的实验表明，CAN能够显著提升扩散Transformer模型（如DiT和UViT）的性能，同时大幅减少计算量。与EfficientViT结合的CaT模型在ImageNet 512×512上实现了2.78的FID，超越了DiT-XL/2，且每步采样所需的MACs减少了52倍。

1. 背景介绍

近年来，大规模图像和视频生成模型在合成逼真图像和视频方面展现了惊人的能力。然而，要将这些模型转化为人类的生产工具，关键一步是加入控制，即让生成模型能够根据人类的指令（如类别标签、文本、姿态等）生成图像。以往的研究中，GANs等模型通过自适应归一化等技术动态调整中间特征图，或者使用交叉注意力、自注意力等方法将条件特征与图像特征融合，但这些方法都是在特征空间中添加控制，而神经网络的权重（卷积/线性层）对于不同条件保持不变。

本文提出了一种新的基于权重空间操作的条件控制方法——Condition-Aware Neural Network（CAN）。CAN通过引入一个权重生成模块，根据输入条件动态生成卷积/线性层的条件权重，从而实现对图像生成过程的控制。

2. CAN 结构

CAN的核心思想是让神经网络能够根据给定的条件动态调整其权重。与常规神经网络不同，CAN增加了一个条件感知权重生成模块。该模块的输入是条件嵌入（包括用户指令，如类别标签，以及扩散模型的时间步），输出是用于调整卷积/线性层静态权重的条件权重。在训练和推理过程中，条件权重与静态权重融合为一个单一的核，通过将权重值相加来实现。

理论上，可以将神经网络的所有层都设置为条件感知的，但在实践中，这可能并不是一个好的设计。一方面，过多的条件感知层可能会使模型优化变得困难；另一方面，从效率角度来看，为所有层生成条件权重虽然计算开销可以忽略不计，但会带来显著的参数开销。因此，本文选择性地将某些模块设置为条件感知，以平衡性能和效率。实验表明，将深度卷积层、patch嵌入层和输出投影层设置为条件感知可以显著提升性能。

另一种可能的方法是维护一组基础卷积核，并动态生成缩放参数来组合这些基础核，而不是直接生成条件权重。这种方法的参数开销比CAN小，但其性能无法与CAN相匹配，表明动态参数化本身并不是提升性能的关键，更好的条件感知适应能力才是至关重要的。

由于条件感知层对不同样本有不同的权重，因此不能进行批量训练和推理，而必须独立运行每个样本的核调用。这将显著降低GPU上的训练速度。为了解决这一问题，本文采用了一种高效的实现方法，将所有卷积核调用融合为一个分组卷积，其中分组数等于批量大小$B$。在运行分组卷积之前进行批量到通道的转换，以保留功能。操作完成后，添加通道到批量的转换，将特征图恢复为原始格式。理论上，这种高效实现与运行静态模型相比几乎没有训练开销。

3. 实验分析

CAN在各种UViT和DiT变体上都显著提升了图像质量和可控性，同时几乎不增加计算成本。训练曲线表明，CAN的改进并非由于更快的收敛速度，而是在训练更长时间后，性能提升仍然显著。

对于扩散模型，条件嵌入包含类别标签和时间步。实验表明，类别标签信息在权重生成过程中比时间步信息更重要。仅添加类别标签即可获得比仅添加时间步更低的FID和更高的CLIP分数。同时包含类别标签和时间步的条件嵌入能够实现最佳结果。

CAN可以单独作为有效的条件控制方法，也可以与其他条件控制方法结合使用以获得更好的结果。例如，在DiT-S/2上，单独使用CAN比自适应归一化获得了更好的性能；在UViT-S/2上，单独使用CAN比将条件作为标记的注意力机制获得了更好的性能。对于UViT模型，结合CAN和注意力机制（条件作为标记）会略微降低性能。

在ImageNet 256×256和512×512的类别条件生成任务中，CAN模型（如CaT-B0和CaT-L0）在FID上超越了DiT-XL/2和UViT-H/2等模型，同时在计算效率上具有显著优势。例如，CaT-B0的MACs比DiT-XL/2少了9.9倍，比UViT-H/2少了11.1倍。在COCO 256×256的文本到图像生成任务中，CaT-S0在FID上与UViT-S-Deep/2相当，但计算成本大幅降低。