十分钟读懂StableDiffusion运行原理

点缀

作者：symonAIGC热潮正猛烈地席卷开来，可以说StableDiffusion开源发布把AI图像生成提高了全新高度，特别是ControlNet和T2I-Adapter控制模块的提出进一步提高生成可控性，也在逐渐改变一部分行业的生产模式。惊艳其出色表现，也不禁好奇其背后技术。本文整理了一些学习过程中记录的技术内容，主要包括StableDiffusion技术运行机制，希望帮助大家知其所以然。一背景介绍AI绘画作为AIGC（人工智能创作内容）的一个应用方向，它绝对是2022年以来AI领域最热门的话题之一。AI绘画凭借着其独特创意和便捷创作工具迅速走红，广受关注。举两个简单例子，左边是利用controlnet新魔法把一张四个闺蜜在沙滩边上的普通合影照改成唯美动漫风，右边是midjourneyv5最新版本解锁的逆天神技，只需输入文字“旧厂街风格，带着浓浓90年代氛围感”即可由AI一键生成超逼真！图1两个stablediffusion例子StableDiffusion，是一个2022年发布的文本到图像潜在扩散模型，由CompVis、StabilityAI和LAION的研究人员创建的。要提到的是，StableDiffusion技术提出者StabilityAI公司在2022年10月完成了1.01亿美元的融资，估值目前已经超过10亿美元。本文会在第二部分着重介绍StableDiffusion的技术思路，第三部分深入分析各个重要模块的运行机制，最后总结下AI绘画。二原理简介StableDiffusion技术，作为Diffusion改进版本，通过引入隐向量空间来解决Diffusion速度瓶颈，除了可专门用于文生图任务，还可以用于图生图、特定角色刻画，甚至是超分或者上色任务。作为一篇基础原理介绍，这里着重解析最常用的“文生图（texttoimage）”为主线，介绍stablediffusion计算思路以及分析各个重要的组成模块。下图是一个基本的文生图流程，把中间的StableDiffusion结构看成一个黑盒，那黑盒输入是一个文本串“paradise(天堂)、cosmic(广阔的)、beach(海滩)”，利用这项技术，输出了最右边符合输入要求的生成，图中产生了蓝天白云和一望无际的广阔海滩。图2StableDiffusion组成StableDiffusion的核心思想是，由于每张满足一定规律分布，利用文本中包含的这些分布信息作为指导，把一张纯噪声的逐步去噪，生成一张跟文本信息匹配的。它其实是一个比较组合的系统，里面包含了多个模型子模块，接下来把黑盒进行一步步拆解。stablediffusion最直接的问题是，如何把人类输入的文字串转换成机器能理解的数字信息。这里就用到了文本编码器textencoder（蓝色模块），可以把文字转换成计算机能理解的某种数学表示，它的输入是文字串，输出是一系列具有输入文字信息的语义向量。有了这个语义向量，就可以作为后续生成器imagegenerator（粉黄组合框）的一个控制输入，这也是stablediffusion技术的核心模块。生成器，可以分成两个子模块（粉色模块+黄色模块）来介绍。下面介绍下stablediffusion运行时用的主要模块：(1)文本编码器（蓝色模块），功能是把文字转换成计算机能理解的某种数学表示，在第三部分会介绍文本编码器是怎么训练和如何理解文字，暂时只需要了解文本编码器用的是CLIP模型，它的输入是文字串，输出是一系列包含文字信息的语义向量。(2)信息生成器（粉色模块），是stablediffusion和diffusion模型的区别所在，也是性能提升的关键，有两点区别：①信息生成器的输入输出均为低维向量（不是原始），对应上图里的粉色44方格。同时文本编码器的语义向量作为信息生成器的控制条件，把信息生成器输出的低维向量进一步输入到后续的解码器（黄色）生成。（注：原始的分辨率为512512，有RGB三通道，可以理解有RGB三个元素组成，分别对应红绿蓝；低维向量会降低到64*64维度）②Diffusion模型一般都是直接生成，不会有中间生成低维向量的过程，需要更大计算量，在计算速度和资源利用上都比不过stablediffusion；那低维空间向量是如何生成的？是在信息生成器里由一个Unet网络和一个采样器算法共同完成，在Unet网络中一步步执行生成过程，采样器算法控制生成速度，下面会在第三部分详细介绍这两个模块。StableDiffusion采样推理时，生成迭代大约要重复30~50次，低维空间变量在迭代过程中从纯噪声不断变成包含丰富语义信息的向量，信息生成器里的循环标志也代表着多次迭代过程。(3)解码器（黄色模块），输入为信息生成器的低维空间向量（粉色4*4方格），通过升维放大可得到一张完整。由于输入到信息生成器时做了降维，因此需要增加升维模块。这个模块只在最后阶段进行一次推理，也是获得一张生成的最终步骤。那扩散过程发生了什么？扩散过程发生在信息生成器中，把初始纯噪声隐变量输入到Unet网络后结合语义控制向量，重复30~50次来不断去除纯噪声隐变量中的噪声，并持续向隐向量中注入语义信息，就可以得到一个具有丰富语义信息的隐空间向量（右下图深粉方格）。采样器负责统筹整个去噪过程，按照设计模式在去噪不同阶段中动态调整Unet去噪强度。更直观看一下，如图3所示，通过把初始纯噪声向量和最终去噪后的隐向量都输到后面的解码器，观察输出区别。从下图可以看出，纯噪声向量由于本身没有任何有效信息，解码出来的也是纯噪声；而迭代50次去噪后的隐向量已经耦合了语义信息，解码出来也是一张包含语义信息的有效。图3可视化输出变化到这里，我们大致介绍了StableDiffusion是什么以及各个模块思路，并且简单介绍了stablediffusion的扩散过程。第三部分我们继续分析各个重要组成模块的运行机制，更深入理解StableDiffusion工作原理。三模块分析第二部分以零基础角度介绍了StableDiffusion技术思路，这部分会更细致地介绍下StableDiffusion文生图技术，训练阶段和采样阶段的总体框架如图4所示，可以划分成3个大模块：PART1-CLIP模型，PART2-Unet训练，PART3-采样器迭代。图4StableDiffusion文生图框架训练阶段，包含了图里PART1CLIP模型和PART2Unet训练，分成三步：1、用AutoEncoderKL自编码器把输入从像素空间映射到隐向量空间，把RGB转换到隐式向量表达。其中，在训练Unet时自编码器参数已经训练好和固定的，自编码器把输入张量进行降维得到隐向量。2、用FrozenCLIPEmbedder文本编码器来编码输入提示词Prompt，生成向量表示context，这里需要规定文本最大编码长度和向量嵌入大小。3、对输入图像的隐式向量施加不同强度噪声，再把加噪后隐向量输入到UNetModel来输出预估噪声，和真实噪声信息标签作比较来计算KL散度loss，并通过反向传播算法更新UNetModel模型参数；引入文本向量context后，UNetModel在训练时把其作为condition，利用注意力机制来更好地引导图像往文本向量方向生成；采样阶段，包含了图里PART1CLIP模型和PART3采样器迭代，分成三步：1、用FrozenCLIPEmbedder文本编码器把输入提示词Prompt进行编码，生成维度为[B,K,E]的向量表示context，与训练阶段的第2步一致；2、利用随机种子随机产出固定维度的噪声隐空间向量，利用训练好的UNetModel模型，结合不同采样器（如DDPM/DDIM/PLMS）迭代T次不断去除噪声，得到具有文本信息的隐向量表征；3、用AutoEncoderKL自编码器把上面得到的图像隐向量进行解码，得到被映射到像素空间的生成图像。上面对stablediffusion总体架构进行了介绍，那接下来进一步分析介绍下每个重要组成模块，分别是Unet网络、采样器和CLIP模型三个主要模块。1Unet网络StableDiffusion里采用的UNetModel模型，采用Encoder-Decoder结构来预估噪声，网络结构如图5：图5Unet网络结构示意图模型输入包括3个部分，(1)图像表示，用隐空间向量输入的维度为[B,Z,H/8,W/8]；(2)timesteps值，维度为[B,]；(3)文本向量表示context，维度为[B,K,E]。其中[B,Z,H,W]分别表示[batch_size数，C隐空间通道数，height长度，weight宽度]，K和E分别表示文本最大编码长度maxlength和向量嵌入大小。模型使用DownSample和UpSample进行样本的下上采样，在采样模块之间还有黑色虚线框的ResBlock和SpatialTransformer，分别接收timesteps信息和提示词信息（这里只画出一次作为参考）。ResBlock模块的输入有①来自上一个模块的输入和②timesteps对应的嵌入向量timestep_emb（维度为[B,4*M]，M为可配置参数）；SpatialTransformer模块的输入有①来自上一个模块的输入和②提示词Prompt文本的嵌入表示context，以context为注意力机制里的condition，学习提示词Prompt和图像的匹配程度。最后，UNetModel不改变输入和输出大小，隐空间向量的输入输出维度均为[B,Z,H/8,W/8]。ResBlock网络ResBlock网络，有两个输入分别是①来自上一个模块的输入和②timesteps对应的嵌入向量timestep_emb（维度为[B,4*M]，M为可配置参数），网络结构图如下所示。timestep_embedding的生成方式，用的是“AttentionisAllyouNeed”论文的Transformer方法，通过sin和cos函数再经过两个Linear进行变换。图6ResBlock模块和Timestep生成流程SpatialTransformer结构SpatialTransformer这里，包含模块比较多，有两个输入分别是①来自上一个模块的输入和②提示词Prompt文本的嵌入表示context作为condition，两者使用crossattention进行建模。其中，SpatialTransformer里面的注意力模块CrossAttention结构，把图像向量作为Query，文本表示context作为Key&Value，利用CrossAttention模块来学习图像和文本对应内容的相关性。注意力模块的作用是，当输入提示词来生成时，比如输入“一匹马在吃草”，由于模型已经能捕捉图文相关性以及文本中的重点信息，当看到“马”时，注意力机制会重点突出图像“马”的生成；当看到“草”时，注意力机制会重点突出图像“草”的生成，进而实现和文本匹配的生成。图7SpatialTransformer结构示意图Unet如何训练？StableDiffusion里面Unet的学习目标是什么？简单来说就是去噪。那在为去噪任务设计训练集时，就可以通过向普通照片添加噪声来得到训练样本。具体来说，对于下面这张照片，用random函数生成从强到弱的多个强度噪声，比如图8里0~3有4个强度的噪声。训练时把噪声强度和加噪后输入到Unet，计算预测噪声图和真正噪声图之间的误差损失，通过反向传播更新unet参数。图8去噪训练样本说明训练好Unet后，如图9所示，从加噪中推断出噪声后，就可以用加噪图减掉噪声来恢复原图；重复这个过程，第一步预测噪声图后再减去噪声图，用更新后的加噪图进行第二步去噪，最终就能得到一张很清晰的生成。由于使用了高斯分布的KL散度损失，Unet生成实际上是接近训练集分布的，和训练集有着相同像素规律。也就是说，使用真实场景的写实训练集去训练模型，它的结果就会具有写实风格，尽量符合真实世界规律。图9Unet网络采样阶段迭代2采样器迭代这部分介绍下采样阶段中扩散模型如何多次迭代去除噪声，进而得到生成的潜在空间表示。提到采样器，要从最基础的采样器DDPM（DenoisingDiffusionProbabilisticModels）进行介绍[4]。DDPM推导有点复杂，这里就用朴素一点的大白话结合几个关键公式来理清推导思路。1扩散模型的思路是，训练时先在上不断加噪来破坏，推理时对加噪后的去噪来恢复出原始。训练过程的T次迭代中，可推导出一个重要特性：任意时刻的Xt可以由X0和β表示，任意时刻的X0也可以由Xt和噪声z求得。其中，第一行里a和β可以描述噪声强度；第二行，X0为初始的干净，增加噪声z后生成加噪Xt，后个公式由前个公式变换而来，表示加噪减去一定强度噪声，得到X0。2问题变成，如何求逆向阶段的分布，即给定了一张加噪后，如何才能求得前一时刻没有被破坏得那么严重的略清晰。经过论文里的一顿推导，又得出两个重要结论：①逆向过程也服从高斯分布；②在知道原始清晰时，能通过贝叶斯公式把逆向过程转换成前向过程，进而算出逆向过程分布。在公式上体现如下：其中，第一行指的是给定X0情况下，逆向过程也服从高斯分布，并且利用贝叶斯公式把逆向过程转换成前向过程，前向过程是不断加噪的过程，可以被计算；第二行指的是，Xt和X0由于可以相互转换，从公式上看，均值也可以从Xt减去不同噪声得到。3算出逆向过程分布后，就可以训练一个模型尽量拟合这个分布，而且模型预估结果也应该服从高斯分布：其中，求均值公式里只包含Xt和噪声，由于Xt在训练时已知，那只需要得到模型输出的预估噪声，该值可由模型用Xt和t预估得到。4把逆向过程分布（也就是Label值）和模型的预估分布做比较，由于①KL散度可以用来描述两个分布之间的差异和②多元高斯分布的KL散度有闭式解，经过一番推导发现损失函数变成计算两个高斯分布的KL散度。其中，q分布是逆向过程分布，p是模型预估分布，训练损失是求两个高斯分布的KL散度，即两个分布之间的差距。5DDPM训练过程和采样过程的伪代码如下图所示。由于DDPM的高质量生成依赖于较大的T（一般为1000以上），导致Diffusion前向过程特别缓慢，因此后续进一步有了DDIM、PLMS和EulerA等一些优化版采样器。其中，训练阶段实际上是求真实噪声和模型预估噪声的MSE误差，再对Loss求导反向传播来训练模型；采样阶段，求得均值和方差后，采用重参数技巧来生成样本。总结下，扩散模型采样阶段是对加噪后去噪来恢复出原始，基于①任意时刻的均可以由原始和噪声表示；②逆向过程的参数符合高斯分布，优化目标转化为计算逆向分布和预估分布的KL散度差异，并在采样阶段使用重参数技巧来生成。3CLIP模型在前面有提到，提示词Prompt文本利用文本模型转换成嵌入表示context，作为Unet网络的condition条件。那问题来了，语义信息和信息属于两种模态，怎么用attention耦合到一起呢？这里介绍下用于提取语义信息的CLIP模型。语义信息的好坏直接影响到了最终生成的多样性和可控性，那像CLIP这样的语言模型是如何训练出来的？是如何结合文本串和计算机视觉的呢？首先，要有一个具有文本串和计算机视觉配对的数据集。CLIP模型所使用的训练集达到了4亿张，通过从网络上爬取及相应的标签或者注释。图9CLIP模型训练用到的和标签示例CLIP模型结构包含一个encoder和一个文字encoder，类似于推荐场景常用到的经典双塔模型。训练时，从训练集随机取出一些样本（和标签配对的话就是正样本，不匹配的话就是负样本），CLIP模型的训练目标是预测图文是否匹配；取出文字和后，用encoder和文字encoder分别转换成两个embedding向量，称作embedding和文字embedding；用余弦相似度来比较两个embedding向量相似性，并根据标签和预测结果的匹配程度计算损失函数，用来反向更新两个encoder参数。在CLIP模型完成训练后，输入配对的和文字，这两个encoder就可以输出相似的embedding向量，输入不匹配的和文字，两个encoder输出向量的余弦相似度就会接近于0。推理时，输入文字可以通过一个textencoder转换成textembedding，也可以把用imageencoder转换成imageembedding，两者就可以相互作用。在生成的采样阶段，把文字输入利用textencoder转换成嵌入表示textembedding，作为Unet网络的condition条件。图10CLIP模型结构示意图四本文小结AI绘画各种应用不断涌现，目前有关StableDiffusion的文章主要偏向应用介绍，对于StableDiffusion技术逻辑的介绍还是比较少。这篇文章主要介绍了StableDiffusion技术结构和各个重要组成模块的基本原理，希望能够让大家了解StableDiffusion是如何运行的，才能更好地控制AI绘画生成。AI绘画虽然还面临一些技术挑战，但随着技术不断迭代和发展，相信AI能够在更多领域发挥出惊喜生产力。（本文参考了stablediffusion官方仓库以及一些解读Blog，结合个人在其他ML领域经验的一些解读。如有不合理的地方，欢迎在评论区指出。）参考链接1GitHub-CompVis/stable-diffusion:Alatenttext-to-imagediffusionmodel2TheIllustratedStableDiffusion:TheIllustratedStableDiffusion–JayAlammar–Visualizingmachinelearningoneconceptatatime.3由浅入深了解DiffusionModel4DenoisingDiffusionProbabilisticModels