Rocky Ding

公众号：WeThinkIn

写在前面

【算法兵器谱】栏目专注分享AI行业中的前沿/经典/必备的模型&论文，并对具备划时代意义的模型&论文进行全方位系统的解析，比如Rocky之前出品的爆款文章Make YOLO Great Again系列。也欢迎大家提出宝贵的优化建议，一起交流学习💪

大家好，我是Rocky。

2022年，Stable Diffusion横空出世，成为AI行业从传统深度学习时代过渡至AIGC时代的标志模型，并为工业界和投资界注入了新的活力，让AI再次性感。

Stable Diffusion系列的第一篇文章中，Rocky已经详细讲解了Stable Diffusion的核心基础知识深入浅出解析Stable Diffusion完整核心基础知识 |【算法兵器谱】。本文作为Stable Diffusion系列的第二篇文章，Rocky将深入浅出的讲解Stable Diffusion中U-Net的核心知识，包括U-Net在传统深度学习中的核心价值与在AICG中的核心价值，让我们来看看U-Net是如何在两个时代中同时从容，并大放异彩的。同时，Rocky也希望我们能借助Stable Diffusion系列文章更好的入门Stable Diffusion及其背后的AIGC领域。

话不多说，在Rocky毫无保留的分享下，让我们开始学习吧！

So，enjoy：

正文开始

----【目录先行】----

传统深度学习时代的U-Net

【一】U-Net的“AI江湖”印象

【二】U-Net的核心结构与细节

Encoder-Decoder结构
U-Net结构细节挖掘
是什么让U-Net通向AIGC

Stable Diffusion中的U-Net

【一】U-Net在Stable Diffusion中扮演的角色

【二】U-Net在AIGC时代中的核心结构与细节

Time Embedding模块
Spatial Transformer(Cross Attention)模块
Stable Diffusion中U-Net的完整核心结构
GroupNorm

【三】U-Net在Stable Diffusion中的训练和推理

U-Net在Stable Diffusion中的训练过程
U-Net在Stable Diffusion中的推理过程

----【传统深度学习时代的U-Net】----

【一】U-Net的“AI江湖”印象

在2015年，传统深度学习时代的早期，U-Net: Convolutional Networks for Biomedical Image Segmentation（U-Net）正式发表，图像分割领域迎来了它的“ResNet”。

U-Net起初在生物医学图像这个细分领域取得了最佳的工业界效果，由于其简洁，高效，稳定的特性，随机被广泛的应用于图像分割的各个方向，比如智慧交通，智慧城市，工业检测等

可以说在传统深度学习时代，不管是实际业务、AI竞赛还是科研，U-Net都成为了当仁不让的图像分割通用Baseline。但是让人没想到的是，在8年后的AIGC时代到来后，U-Net顺应了时代的潮流，依旧爆发出了鲜活的生命力与价值。

【二】U-Net的核心结构与细节

Encoder-Decoder结构

U-Net最经典的特征是其Encoder-Decoder的结构，这样的结构简洁且高效，并且具备对称的“艺术”美感，也让U-Net具备了极强的生命力与适应性。

传统深度学习时代的U-Net

其中左半部分的Encoder模块负责进行特征的提取与学习，Encoder模块可以由ResNet、VGG、EfficientNet等一流特征提取模型担任，所以Encoder模块具备较强的工程潜力与科研势能。与此同时Encoder模块可以增加对扰动噪声的鲁棒性，减少过拟合的风险，降低运算量以及增加感受野的大小等作用。

而右半部分的Decoder模块则负责将feature map恢复到原始分辨率，并使skip-connection这个关键一招融合了浅层的位置信息与深层的语义信息。与此同时，Decoder模块和Encoder模块一样可以由ResNet、VGG、EfficientNet等一流模型担任，从而使得U-Net的变体非常繁荣，增加了工程“魔改”的可玩性。

U-Net结构细节挖掘

讲完Encoder-Decoder结构的整体框架，Rocky再向大家介绍一下Encoder-Decoder结构中的一些能够成为通用范式和经典Tricks的细节操作。

从上图的Encoder-Decoder结构中可以看到，U-Net是一个全卷积神经网络，网络最后一层使用了浅蓝色箭头，表示1*1卷积，其完全取代了全连接层，使得模型的输入尺寸不再受限制，极大增强了U-Net在各种应用场景的兼容性。

上图中的蓝色和白色框表示feature map，深蓝色箭头表示 3x3 卷积，padding=0 ，stride=1其用于特征提取。由于padding=0，所以每次经过卷积运算，feature map将有一定程度的下采样。深红色箭头表示max pooling，stride=2，用于降低维度。将卷积和max pooling两者结合，能够对feature map进行特征提取的同时从容进行下采样。

max pooling操作

上图中的绿色箭头表示Upsample操作，对feature map进行上采样从而恢复维度。

Upsampling常用的方式有两种：转置卷积和插值。两者的详细内容可以阅读Rocky之前的文章：【三年面试五年模拟】算法工程师的独孤九剑秘籍（前十二式汇总篇）V1版。而在U-Net中，使用了bilinear双线性插值。

在Encoder和Decoder两个模块之间，使用skip-connection作为桥梁，用于特征融合，将浅层的位置信息与深层的语义信息进行concat操作。图中用灰色箭头表示skip-connection，其中“copy”就是concat操作，而“crop”则通过裁剪使得两个特征图尺寸一致。

是什么让U-Net通向AIGC

讲完U-Net在传统深度学习时代的核心知识点与价值，接下来Rocky再阐述一下为何在AIGC时代，U-Net成为了Stable Diffusion这个划时代模型的关键结构。

在投资界有一句话，叫“在上个时代适应的越好的人，很有可能是下一个时代最大的失败者”。这个逻辑套用在技术上再合适不过了，有太多技术产生，也有太多技术消亡，而学习技术并从事技术行业的工人们就会背负更多的沉没成本与风险。

但是U-Net不这么认为，其同时成为了AIGC与传统深度学习这两个时代的弄潮儿，在AIGC时代，U-Net有了新的内涵和面貌，并且“文艺复兴”。

那么，是什么让U-Net能够通向AIGC，跨过周期呢？

主要有以下四个特质：

U-Net中Encoder模块的压缩特质。作为Encoder模块最初的应用，输入的图像经过下采样，抽取出比原图小得多的高维特征，相当于进行了压缩操作。这和Stable diffusion的latent逻辑不谋而合，随即在AIGC时代“文艺复兴”。
U-Net中Decoder模块的去噪特质，作为Decoder模块最初的应用，在AIGC时代“文艺复兴”。
U-Net整体结构上的简洁、稳定和高效，使得其在Stable Diffusion中能够从容的迭代去噪声，能够撑起Stable Diffusion的整个图像生成逻辑。
Encoder-Decoder结构的强兼容性，让U-Net不管是在分割领域，还是在生成领域，都能和Transformer等新生代模型的从容融合。

U-Net发表8年后的AIGC时代里，正是这些特质让U-Net顺应了时代的潮流，依旧爆发出了鲜活的生命力与价值。

----【Stable Diffusion中的U-Net】----

【一】U-Net在Stable Diffusion中扮演的角色

Stable Diffusion结构图，U-Net在最核心的位置

Stable Diffusion中的U-Net包含约860M的参数，在float32的精度下，约占3.4G的存储空间。

在上图中可以看到，U-Net是Stable Diffusion中的核心模块。U-Net主要在“扩散”循环中对高斯噪声矩阵进行迭代降噪，并且每次预测的噪声都由文本和timesteps进行引导，将预测的噪声在随机高斯噪声矩阵上去除，最终将随机高斯噪声矩阵转换成图片的隐特征。

在U-Net执行“扩散”循环的过程中，Content Embedding始终保持不变，而Time Embedding每次都会发生变化。每次U-Net预测的噪声都在Latent特征中减去，并且将迭代后的Latent作为U-Net的新输入。

总的来说，如果说Stable Diffusion是“优化噪声的艺术”，那么U-Net将是这个“艺术”的核心主导者。

【二】U-Net在AIGC时代中的核心结构与细节

Stable Diffusion中的U-Net，在Encoder-Decoder结构的基础上，增加了Time Embedding模块，Spatial Transformer(Cross Attention)模块和self-attention模块。

Time Embedding模块

首先，什么是Time Embedding呢？

Time Embedding（时间嵌入）是一种在时间序列数据中用于表示时间信息的技术。时间序列数据是指按照时间顺序排列的数据，例如股票价格、天气数据、传感器数据等。时间嵌入的目的是将时间作为一个特征进行编码，以便在深度学习模型中更好地学习时间相关性特征。

Time Embedding的基本思想是将时间信息映射到一个连续的向量空间，使得时间之间的关系可以被模型学习和利用。

Time Embedding的使用可以帮助深度学习模型更好地理解时间相关性，从而提高模型的性能。比如在Stable Diffusion中，将Time Embedding引入U-Net中，帮助其在扩散过程中从容预测噪声。

Stable Diffusion需要迭代多次对噪音进行逐步预测，使用Time Embedding就可以将time编码到网络中，从而在每一次迭代中让U-Net更加合适的噪声预测。

讲完Time Embedding的核心基础知识，我们再解析一下Stable Diffusion中U-Net的Time Embeddings模块是如何构造的：

Time Embeddings模块

可以看到，Time Embeddings模块 + Encoder模块中原本的卷积层，组成了一个Residual Block结构。它包含两个卷积层，一个Time Embedding和一个skip Connection。而这里的全连接层将Time Embedding变换为和Latent Feature一样的维度。最后通过两者的加和完成time的编码。