论文笔记：Inf-DiT: Upsampling Any-Resolution Image with Memory-Efficient Diffusion Transformer
[ 2405.04312 ] Inf-DiT: 提出了高效图像上采样方法，并利用记忆高效的扩散变压器实现高分辨率图像重建 （arxiv.org）
现有图像扩散模型生产的图像通常受到1024×1024像素或更低分辨率的影响，在生产超分辨率图像（如4096×4096像素）时内存需求会显著提升
upsampling 达到更高的分辨率水平的主要困难是 显著 的 GPU 内存消耗。另一个主要问题是 完全地将整个图像输入到模型中 会导致内存占用。
的空间。
在此提出了一种单向块注意力机制（UniBA），该算法大幅降低了相比O(N²)至O(N)的生成空间复杂度，并显著提升了最高的可用分辨率。
Methodology
Unidirectional Block Attention (UniBA)
在UNet、DiT等模型架构中存在模块间的双向相互依存关系，在执行运算时需要确保图像中各模块的信息同步更新与处理。为了提高资源利用率，在设计相关算法时建议采用一种能够在不一次性构建完整特征图的前提下实现高效运算的技术方案。
许同一图像中的图像块被分割为多个批次进行处理。每一批次只需同步生成对应的图像块，并按照预设的批次顺序依次完成处理。
主要思想是将图片
划分为块
，其中B为块的大小。并提出了如下图所示的注意力实现：
左图：在单向块注意力机制中，每个计算单元主要受到其自身所在层三个关键位置的影响：位于其左上方的位置为左上角的区域；左边相邻位置以及上方相邻位置分别对应左侧和上面的部分。
如图所示： Inf-DiT 的推理过程。 Inf-DiT 该算法基于 内存容量 每次 生成 n × n 尺寸 的 block 。在整个过程中，在线 生成 每个 block 时 仅 需 考虑 其 后续依 走 所依附 的 block 对应 的 KV - cache 即可被 存储 在 内存 中 。
Inf-DiT 架构中, 块之间的相互关联是注意力机制的一种表现形式. 此外, 在 transformer 模型中, 单向块注意力机制的具体计算过程如下:
表示第n层i行j列的块的隐藏状态，
为块间相对位置编码。
然而该方法中的每一个block在计算时所依赖的区域缩小了；尽管如此，在特征传递的过程中仍保持了一定层次感；依然能够捕获远距离的信息
在上图中
变为
Basic Model Architecture
Inf-DiT 的体系结构采用了与DiT相似的设计部分，并将其Vision Transformer (ViT)技术整合到扩散模型中。相比于其他基于卷积神经网络的架构（例如UNet），Inf-DiT主要依赖于注意力机制来连接不同的patch。
交互过程可以较为容易地实现单向块注意力机制。基于此需求，在优化上采样效果方面做出了相应改进与调整。
Model input
鉴于颜色偏移和细节损失等压缩带来的影响，在 RGB 像素空间中完成 Inf-DiT 的重建而非潜在空间。当超分为 f 倍时，在低分辨率 RGB 图像基础上先进行 f 倍上采样，并将其与扩散噪声结合至特征维度后作为输入进入模型处理。
Position Encoding
基于RoPE旋转位置编码方案。
首先构建一个规模充足的的位置编码表。
为每个训练图像，在其左上角随机初始化一个坐标点(x, y)，而非传统的固定点(0,0)。
同时为了捕捉块内与块间交互的不同特征，在此基础之上引入了块级相对位置信息。
，它根据注意前的相对位置分配不同的可学习嵌入。
Global and Local Consistency
Global Consistency with CLIP Image Embedding
利用预训练的CLIP中的图像编码器从低分辨率图像中提取图像嵌入
由于CLIP是在大规模数据集上基于互联网进行训练的，在其图像编码器能够有效提取低分辨率图像的整体特征的基础上进行设计。将全局语义嵌入与时间嵌入进行整合，并将其作为每一层网络单元的输入使用，在这一过程中模型能够直接从高级别的语义信息中进行学习。
基于CLIP中的图像-文本潜在空间，在无论模型是否从未接触过任何图像-文本对的情况下，都可以通过文本来引导生成的方向。
给定一个正提示
和一个负提示
，就可以更新图像嵌入：
α用于控制语义的引导强度。在推理过程中，我们可以简单地使用
代替
作为全局语义嵌入来进行控制。
Local Consistency with Nearby LR Cross Attention
模型在学习LR与HR图像之间的局部对应关系时仍然可能存在连续性问题。为了缓解这一问题，在transformer的第一层中引入了Nearby LR Cross Attention。这一机制有助于提高模型的性能。
每个块对周围的3×3 LR块进行交叉注意力机制处理以捕捉附近区域的LR特征实验结果表明该方法显著降低了生成不连续图像的概率
Experiments
HPDV2数据集下超高分辨率的定量实验：
模型不仅具备生成高分辨率细节的能力，并且能够有效地协调全局信息。然而，在4096X4096分辨率下的FID值稍低于BSRGAN模型；相比之下，FIDcrop指标更能准确反映高分辨率特征的关键特性
FIDcrop采用从高分辨率图像中随机选取299×299像素的patch来进行FID评估，并未像传统的FID方法那样忽视高分辨率图像中的细节信息。由于传统的FID计算流程要求在特征提取之前对输入图像进行预处理
下采样到299 × 299的分辨率
下表是在DIV2K数据集下的超分定量实验：
Ablation Study
