?作者 |?王建周

單位 |?來(lái)也科技AI團(tuán)隊(duì)負(fù)責(zé)人

研究方向 |?分布式系統(tǒng)、CV、NLP

前言

AIGC 目前是一個(gè)非常火熱的方向，DALLE-2，ImageGen，Stable Diffusion 的圖像在以假亂真的前提下，又有著腦洞大開(kāi)的藝術(shù)性，以下是用開(kāi)源的 Stable Diffusion 生成的一些圖片。

這些模型后邊都使用了 Diffusion Model 的技術(shù)，但是缺乏相關(guān)背景知識(shí)去單純學(xué)習(xí) Diffusion Model 門(mén)檻會(huì)比較高，不過(guò)沿著 AE、VAE、CVAE、DDPM 這一系列的生成模型的路線、循序?qū)W習(xí)會(huì)更好的理解和掌握，本文將從原理、數(shù)學(xué)推導(dǎo)、代碼詳細(xì)講述這些模型。

AE (AutoEncoder)

AE 模型作用是提取數(shù)據(jù)的核心特征（Latent Attributes），如果通過(guò)提取的低維特征可以完美復(fù)原原始數(shù)據(jù)，那么說(shuō)明這個(gè)特征是可以作為原始數(shù)據(jù)非常優(yōu)秀的表征。

AE 模型的結(jié)構(gòu)如下圖：

訓(xùn)練數(shù)據(jù)通過(guò) Encoder 得到 Latent，Latent 再通過(guò) Decoder 得到重建數(shù)據(jù)，通過(guò)重建數(shù)據(jù)和訓(xùn)練的數(shù)據(jù)差異來(lái)構(gòu)造訓(xùn)練 Loss，代碼如下（本文所有的場(chǎng)景都是 mnist，編碼器和解碼器都用了最基本的卷積網(wǎng)絡(luò)）：

class?DownConvLayer(tf.keras.layers.Layer):def?__init__(self,?dim):super(DownConvLayer,?self).__init__()self.conv?=?tf.keras.layers.Conv2D(dim,?3,?activation=tf.keras.layers.ReLU(),?use_bias=False,?padding='same')self.pool?=?tf.keras.layers.MaxPool2D(2)def?call(self,?x,?training=False,?**kwargs):x?=?self.conv(x)x?=?self.pool(x)return?xclass?UpConvLayer(tf.keras.layers.Layer):def?__init__(self,?dim):super(UpConvLayer,?self).__init__()self.conv?=?tf.keras.layers.Conv2D(dim,?3,?activation=tf.keras.layers.ReLU(),?use_bias=False,?padding='same')#?通過(guò)UpSampling2D上采樣self.pool?=?tf.keras.layers.UpSampling2D(2)def?call(self,?x,?training=False,?**kwargs):x?=?self.conv(x)x?=?self.pool(x)return?x#?示例代碼都是通過(guò)非常簡(jiǎn)單的卷積操作實(shí)現(xiàn)編碼器和解碼器 class?Encoder(tf.keras.layers.Layer):def?__init__(self,?dim,?layer_num=3):super(Encoder,?self).__init__()self.convs?=?[DownConvLayer(dim)?for?_?in?range(layer_num)]def?call(self,?x,?training=False,?**kwargs):for?conv?in?self.convs:x?=?conv(x,?training)return?xclass?Decoder(tf.keras.layers.Layer):def?__init__(self,?dim,?layer_num=3):super(Decoder,?self).__init__()self.convs?=?[UpConvLayer(dim)?for?_?in?range(layer_num)]self.final_conv?=?tf.keras.layers.Conv2D(1,?3,?strides=1)def?call(self,?x,?training=False,?**kwargs):for?conv?in?self.convs:x?=?conv(x,?training)#?將圖像轉(zhuǎn)成和輸入圖像shape一致reconstruct?=?self.final_conv(x)return?reconstructclass?AutoEncoderModel(tf.keras.Model):def?__init__(self):super(AutoEncoderModel,?self).__init__()self.encoder?=?Encoder(64,?layer_num=3)self.decoder?=?Decoder(64,?layer_num=3)def?call(self,?inputs,?training=None,?mask=None):image?=?inputs[0]#?得到圖像的特征表示latent?=?self.encoder(image,?training)#?通過(guò)特征重建圖像reconstruct_img?=?self.decoder(latent,?training)return?reconstruct_img@tf.functiondef?train_step(self,?data):img?=?data["image"]with?tf.GradientTape()?as?tape:reconstruct_img?=?self((img,),?True)trainable_vars?=?self.trainable_variables#?利用l2?loss?來(lái)判斷重建圖片和原始圖像的一致性l2_loss?=?(reconstruct_img?-?img)?**?2l2_loss?=?tf.reduce_mean(tf.reduce_sum(l2_loss,?axis=(1,?2,?3)))gradients?=?tape.gradient(l2_loss,?trainable_vars)self.optimizer.apply_gradients(zip(gradients,?trainable_vars))return?{"l2_loss":?l2_loss}

通過(guò) AE 模型可以看到，只要有有效的數(shù)據(jù)的 Latent Attribute 表示，那么就可以通過(guò) Decoder 來(lái)生成新數(shù)據(jù)，但是在 AE 模型中，Latent 是通過(guò)已有數(shù)據(jù)生成的，所以沒(méi)法生成已有數(shù)據(jù)外的新數(shù)據(jù)。

所以我們?cè)O(shè)想，是不是可以假設(shè) Latent 符合一定分布規(guī)律，只要通過(guò)有限參數(shù)能夠描述這個(gè)分布，那么就可以通過(guò)這個(gè)分布得到不在訓(xùn)練數(shù)據(jù)中的新 Latent，利用這個(gè)新 Latent 就能生成全新數(shù)據(jù)，基于這個(gè)思路，有了 VAE（Variational AutoEncoder 變分自編碼器）。

VAE

VAE 中假設(shè) Latent Attributes （公式中用 z）符合正態(tài)分布，也就是通過(guò)訓(xùn)練數(shù)據(jù)得到的 z 滿足以下條件：

因?yàn)?z 是向量，所 都是向量，分別為正態(tài)分布的均值和方差。有了學(xué)習(xí)得到正態(tài)分布的參數(shù) ，那么就可以從這個(gè)正態(tài)分布中采樣新的 z，新的 z 通過(guò)解碼器得到新的數(shù)據(jù)。

所以在訓(xùn)練過(guò)程中需要同時(shí)優(yōu)化兩點(diǎn)：

1. 重建的數(shù)據(jù)和訓(xùn)練數(shù)據(jù)差異足夠小，也就是生成 x 的對(duì)數(shù)似然越高，一般依然用 L2 或者 L1 loss；

2.? 定義的正態(tài)分布需要和標(biāo)準(zhǔn)正態(tài)分布的一致，這里用了 KL 散度來(lái)約束兩個(gè)分布一致；

Loss 公式定義如下，其中 和 為生成分布， 為編碼分布， 為從正態(tài)分布中采樣的先驗(yàn)分布：

Loss 的證明如下：

因?yàn)槲覀兊哪繕?biāo)是最大化對(duì)數(shù)似然生成分布 ，也就是最小化負(fù)的公式 15，也就是公式 1 的 Loss。

所以 VAE 的結(jié)構(gòu)如下：

注意的是在上圖中有一個(gè)采樣 z 的操作，這個(gè)操作不可導(dǎo)導(dǎo)致無(wú)法對(duì)進(jìn)行優(yōu)化，所以為了反向傳播優(yōu)化，用到重參數(shù)的技巧，也就是將 z 表示成 的數(shù)學(xué)組合方式且該組合方式可導(dǎo)，組合公式如下：

?

可以證明重參數(shù)后的模型 f 輸出期望是不變的（z 是連續(xù)分布）。

在計(jì)算 定義的正態(tài)分布和 定義的正態(tài)分布的 KL 散度時(shí)，用了數(shù)學(xué)推導(dǎo)進(jìn)行簡(jiǎn)化。

對(duì)公式 28 的 log 部分繼續(xù)簡(jiǎn)化：

令：

將公式 32 和 33 帶入公式 28 得到：

因?yàn)?#xff1a;

將公式 37、38、45 帶入公式 34 得到最終的 KL 散度 Loss 公式：

因?yàn)? 非負(fù)，所以我們通過(guò)神經(jīng)網(wǎng)絡(luò)來(lái)學(xué)習(xí) 。

有了前邊的鋪墊，所以 VAE 的實(shí)現(xiàn)上也比較簡(jiǎn)單，代碼如下：

class?VAEModel(tf.keras.Model):def?__init__(self,?inference=False):super(VAEModel,?self).__init__()self.inference?=?inferenceself.encoder?=?Encoder(64,?layer_num=3)self.decoder?=?Decoder(64,?layer_num=3)#?mnist?的size是28，這里為了簡(jiǎn)單對(duì)齊大小，縮放成了32self.img_size?=?32#?z的維度self.latent_dim?=?64#?通過(guò)全連接來(lái)學(xué)習(xí)隱特征z正態(tài)分布的均值self.z_mean_mlp?=?tf.keras.Sequential([tf.keras.layers.Dense(self.latent_dim?*?2,?activation="relu"),tf.keras.layers.Dense(self.latent_dim,?use_bias=False),])#?通過(guò)全連接來(lái)學(xué)習(xí)隱特征z正態(tài)分布的方差的對(duì)數(shù)log(o^2)self.z_log_var_mlp?=?tf.keras.Sequential([tf.keras.layers.Dense(self.latent_dim?*?2,?activation="relu"),tf.keras.layers.Dense(self.latent_dim,?use_bias=False),])#?通過(guò)全連接將z?縮放成上采樣輸入適配的shapeself.decoder_input_size?=?[int(self.img_size?/?(2?**?3)),?64]self.decoder_dense?=?tf.keras.layers.Dense(self.decoder_input_size[0]?*?self.decoder_input_size[0]?*?self.decoder_input_size[1],activation="relu")def?sample_latent(self,?bs,?image):#?推理階段的z直接可以從標(biāo)準(zhǔn)正態(tài)分布中采樣，因?yàn)橛?xùn)練的decoder已經(jīng)可以從標(biāo)準(zhǔn)高斯分布生成新的圖片了if?self.inference:z?=?tf.keras.backend.random_normal(shape=(bs,?self.latent_dim))z_mean,?z_log_var?=?None,?Noneelse:x?=?imagex?=?self.encoder(x)x?=?tf.keras.layers.Flatten()(x)z_mean?=?self.z_mean_mlp(x)z_log_var?=?self.z_log_var_mlp(x)epsilon?=?tf.keras.backend.random_normal(shape=(bs,?self.latent_dim))'''實(shí)現(xiàn)重參數(shù)采樣公式17u?+?exp(0.5*log(o^2))*e=u?+exp(0.5*2*log(o))*e=u?+?exp(log(o))*e=u?+?o*e'''z?=?z_mean?+?tf.exp(0.5?*?z_log_var)?*?epsilonreturn?z,?z_mean,?z_log_vardef?call(self,?inputs,?training=None,?mask=None):#?推理生成圖片時(shí)，image為Nonebs,?image?=?inputs[0],?inputs[1]z,?z_mean,?z_log_var?=?self.sample_latent(bs,?image)latent?=?self.decoder_dense(z)latent?=?tf.reshape(latent,[-1,?self.decoder_input_size[0],?self.decoder_input_size[0],?self.decoder_input_size[1]])#?通過(guò)z重建圖像reconstruct_img?=?self.decoder(latent,?training)return?reconstruct_img,?z_mean,?z_log_vardef?compute_loss(self,?reconstruct_img,?z_mean,?z_log_var,?img):#?利用l2?loss?來(lái)判斷重建圖片和原始圖像的一致性l2_loss?=?(reconstruct_img?-?img)?**?2l2_loss?=?tf.reduce_mean(tf.reduce_sum(l2_loss,?axis=(1,?2,?3)))#?實(shí)現(xiàn)公式48kl_loss?=?-0.5?*?(1?+?z_log_var?-?tf.square(z_mean)?-?tf.exp(z_log_var))kl_loss?=?tf.reduce_mean(tf.reduce_sum(kl_loss,?axis=1))total_loss?=?kl_loss?+?l2_lossreturn?{"l2_loss":?l2_loss,?"total_loss":?total_loss,?"kl_loss":?kl_loss}@tf.functiondef?forward(self,?data,?training):img?=?data["img_data"]bs?=?tf.shape(img)[0]reconstruct_img,?z_mean,?z_log_var?=?self((bs,?img),?training)return?self.compute_loss(reconstruct_img,?z_mean,?z_log_var,?img)def?train_step(self,?data):with?tf.GradientTape()?as?tape:result?=?self.forward(data,?True)trainable_vars?=?self.trainable_variablesgradients?=?tape.gradient(result["total_loss"],?trainable_vars)self.optimizer.apply_gradients(zip(gradients,?trainable_vars))return?result

生成的圖片效果如下：

在我們大多數(shù)生成場(chǎng)景，都需要帶有控制條件，比如我們?cè)谏a(chǎn)手寫(xiě)數(shù)字的時(shí)候，我們需要明確的告訴模型，生成數(shù)字 0 的圖片，基于這個(gè)需求，有了 Conditional Variational AutoEncoder（CVAE）。

CVAE

CVAE 的改進(jìn)思路比較簡(jiǎn)單，就是訓(xùn)練階段的 z 同時(shí)由 x 和控制條件 y 決定，同時(shí)生成的 x 也是由 y 和 z 同時(shí)決定，Loss 如下：

而? q(z|y) 我們?nèi)匀黄谕蠘?biāo)準(zhǔn)正態(tài)分布，對(duì) VAE 代碼改動(dòng)非常少，簡(jiǎn)單的實(shí)現(xiàn)方法就是對(duì)條件 y 有一個(gè) embedding 表示，這個(gè) embedding 表示參與到 encoder 和 decoder 的訓(xùn)練，代碼如下：

class?CVAEModel(VAEModel):def?__init__(self,?inference=False):super(CVAEModel,?self).__init__(inference=inference)#?定義label的Embeddingself.label_dim?=?128self.label_embedding?=?tf.Variable(initial_value=tf.keras.initializers.HeNormal()(shape=[10,?self.label_dim]),trainable=True,)self.encoder_y_dense?=?tf.keras.layers.Dense(self.img_size?*?self.img_size,?activation="relu")self.decoder_y_dense?=?tf.keras.layers.Dense(self.decoder_input_size[0]?*?self.decoder_input_size[0]?*?self.decoder_input_size[1],?activation="relu")def?call(self,?inputs,?training=None,?mask=None):#?推理生成圖片時(shí)，image為Nonebs,?image,?label?=?inputs[0],?inputs[1],?inputs[2]label_emb?=?tf.nn.embedding_lookup(self.label_embedding,?label)label_emb?=?tf.reshape(label_emb,?[-1,?self.label_dim])if?not?self.inference:#?訓(xùn)練階段將條件label的embedding拼接到圖片上作為encoder的輸入encoder_y?=?self.encoder_y_dense(label_emb)encoder_y?=?tf.reshape(encoder_y,?[-1,?self.img_size,?self.img_size,?1])image?=?tf.concat([encoder_y,?image],?axis=-1)z,?z_mean,?z_log_var?=?self.sample_latent(bs,?image)latent?=?self.decoder_dense(z)#?將條件label的embedding拼接到z上作為decoder的輸入decoder_y?=?self.decoder_y_dense(label_emb)latent?=?tf.concat([latent,?decoder_y],?axis=-1)latent?=?tf.reshape(latent,[-1,?self.decoder_input_size[0],?self.decoder_input_size[0],self.decoder_input_size[1]?*?2])#?通過(guò)特征重建圖像reconstruct_img?=?self.decoder(latent,?training)return?reconstruct_img,?z_mean,?z_log_var@tf.functiondef?forward(self,?data,?training):img?=?data["img_data"]label?=?data["label"]bs?=?tf.shape(img)[0]reconstruct_img,?z_mean,?z_log_var?=?self((bs,?img,?label),?training)return?self.compute_loss(reconstruct_img,?z_mean,?z_log_var,?img)def?train_step(self,?data):with?tf.GradientTape()?as?tape:result?=?self.forward(data,?True)trainable_vars?=?self.trainable_variablesgradients?=?tape.gradient(result["total_loss"],?trainable_vars)self.optimizer.apply_gradients(zip(gradients,?trainable_vars))return?result

生成 0~9 的圖片效果如下：

從 VAE 的原理可以看到，我們做了假設(shè) ，但是在大多數(shù)場(chǎng)景，這個(gè)假設(shè)過(guò)于嚴(yán)苛，很難保證數(shù)據(jù)特征符合基本的正態(tài)分布（嚴(yán)格意義上也做不到，嚴(yán)格分布的話說(shuō)明特征就是高斯噪聲了），因?yàn)檫@個(gè)缺陷，所以基本的 VAE 生成的圖像細(xì)節(jié)不夠，邊緣偏模糊。

為了解決這些問(wèn)題，又出現(xiàn) DDPM（Denoising Diffusion Probabilistic Model），因?yàn)?DDPM 相比 GAN，更容易訓(xùn)練（GAN 需要交替訓(xùn)練，而且容易出現(xiàn)模式崩塌，可以參考我們以前的文章），此外 DDPM 的多樣性相比 GAN 更好（GAN 因?yàn)樯傻膱D像要“欺騙”過(guò)鑒別器，所以生成的圖像和訓(xùn)練集合的真實(shí)圖像類似），所以最近 DDPM 成為最受歡迎的生成模型。

DDPM

DDPM 啟發(fā)點(diǎn)來(lái)自非平衡熱力學(xué)，系統(tǒng)和環(huán)境之間有著物質(zhì)和能量交換，比如在一個(gè)盛水的容器中滴入一滴墨水，最終墨水會(huì)均勻的擴(kuò)散到水中，但是如果擴(kuò)散的每一步足夠小，那么這一步就可逆。

所以主要流程上分兩個(gè)階段，前向加噪和反向去噪，原始數(shù)據(jù)為 ，每一步添加足夠小的高斯噪聲，經(jīng)過(guò)足夠的 step T 后，最終數(shù)據(jù) 會(huì)變成標(biāo)準(zhǔn)的高斯噪聲（下圖的 q），因?yàn)榍跋蚣釉肷鲜强尚械?#xff0c;所以我們假設(shè)反向去噪也是可行的，可以逐步的從噪聲中一點(diǎn)點(diǎn)的恢復(fù)數(shù)據(jù)的有用信息（下圖的 p）直到為 ，下邊將詳細(xì)介紹兩部分。

1. 前向加噪

假設(shè)前向加噪過(guò)程每一步添加噪聲的過(guò)程符合以下高斯分布，且整個(gè)過(guò)程滿足馬爾科夫鏈，即以下公式：

根據(jù)上文提到的重參數(shù)技巧，公式 50 可以寫(xiě)成（為了方便，寫(xiě)成標(biāo)量形式）：

其中 ，所以公式 52 可以理解為向原始的數(shù)據(jù)設(shè)中加非常小的高斯噪音，并且隨著t變大加的噪音逐漸變大，為了方便公式推導(dǎo)，令：

因?yàn)?#xff1a;

根據(jù)正態(tài)分布的求和計(jì)算公式以及重參數(shù)技巧：

令 ，將公式 63 帶入 57 并推導(dǎo)到一般形式，得到如下前向公式：

公式 64 就是正向過(guò)程的最終公式，可以看到正向過(guò)程是不存在任何網(wǎng)絡(luò)參數(shù)的，而且對(duì)于給定的 t，無(wú)需迭代，通過(guò)表達(dá)式可以直接計(jì)算得到 。

2. 反向去噪

反向去噪期望從標(biāo)準(zhǔn)的高斯分布噪聲 逐步的消除噪音，每次只恢復(fù)目標(biāo)數(shù)據(jù)的一點(diǎn)點(diǎn)，最終生成目標(biāo)數(shù)據(jù) ，假設(shè)的反向去噪也是符合高斯分布和馬爾科夫鏈，可以用以下數(shù)學(xué)公式描述：

因?yàn)? 中 是依賴 的，所以單純的 是無(wú)法計(jì)算的，所以我們需要轉(zhuǎn)而計(jì)算 （上圖的粉色路徑），前者有帶學(xué)習(xí)的參數(shù)，我們假設(shè)：

接下來(lái)的目標(biāo)是需要寫(xiě)出 的表達(dá)式，主要是利用條件概率和貝葉斯公式（為了簡(jiǎn)化都用標(biāo)量的形式）。

帶入各自的表達(dá)式：

得到：

對(duì)比正態(tài)分布公式：

可以得到我們需要的 的表達(dá)式：

接下來(lái)我們需要推導(dǎo)下優(yōu)化的目標(biāo)，根據(jù)前邊公式 10 的推導(dǎo)有以下：

因?yàn)?#xff1a;

公式 101 代入公式 96 得到：

對(duì) 112 的 繼續(xù)推導(dǎo)：

其中 中 為直接計(jì)算出來(lái)等于常數(shù)，所以 為常數(shù)；而 為 的 t=1 的特殊表達(dá)式，故可以合并到 ，所以從公式 95 可以看出，我們最大化的對(duì)數(shù)似然 ，等價(jià)最小化公式 118，而根據(jù)公式 47，兩個(gè)正態(tài)分布的 KL 散度等于：

如果上述的 KL 離散度最小，我們希望 逼近 ，根據(jù)前邊公式93的推導(dǎo)，我們知道：

根據(jù)這個(gè)公式，對(duì)于已知 的情況下，如果能預(yù)測(cè)出 ，就可以解決我們的問(wèn)題，啟發(fā)我們?cè)O(shè)計(jì)以下目標(biāo)：

所以 KL 散度（公式 122）變成以下公式：

前邊的公式 130 的常數(shù) 在訓(xùn)練過(guò)程可以認(rèn)為被合并到學(xué)習(xí)率，所以可以被略掉，所以我們最終的優(yōu)化目標(biāo) Loss 為以下：

所以訓(xùn)練過(guò)程如下：

從公式 66 和 126，以及重參數(shù)技巧可以得知：

所以等待訓(xùn)練完成得到 后，循環(huán)執(zhí)行公式 132 就得到了最終的目標(biāo)數(shù)據(jù) ，過(guò)程如下：

經(jīng)過(guò)前邊較多的公式推導(dǎo)，最終得到 DDPM 的訓(xùn)練和生成過(guò)程確非常簡(jiǎn)單，從前邊能看到希望網(wǎng)絡(luò) 輸入輸出 shape 一致，所以常見(jiàn)的 DDPM 都是用 unet 來(lái)實(shí)現(xiàn)（下圖，核心是四點(diǎn)：下采樣、上采樣、上下采樣的特征拼接），在代碼上我們做了部分優(yōu)化。

1. 為了簡(jiǎn)化代碼，我們?nèi)サ舫Ｒ?jiàn)實(shí)現(xiàn)方式的 self-attention；

2. 一般時(shí)間步 t 也會(huì)采用 transformer 中基本的 sincos 的 position 編碼，為了簡(jiǎn)化編碼，我們的時(shí)間編碼直接采用可以學(xué)習(xí)網(wǎng)絡(luò)并只加入 Unet 的編碼階段，解碼階段不加入；

3. 相比前邊的 VAE 代碼，這里的代碼相對(duì)復(fù)雜，卷積模塊采用 Resnet 的殘差處理方式（經(jīng)過(guò)實(shí)驗(yàn)，前邊 VAE 基本的編碼器和解碼器過(guò)于簡(jiǎn)單，沒(méi)法收斂）；

4. 參照官方，用 group norm 代替 batch norm。

class?ConvResidualLayer(tf.keras.layers.Layer):def?__init__(self,?filter_num):super(ConvResidualLayer,?self).__init__()self.conv1?=?tf.keras.layers.Conv2D(filter_num,?kernel_size=1,?padding='same')#?import?tensorflow_addons?as?tfaself.gn1?=?tfa.layers.GroupNormalization(8)self.conv2?=?tf.keras.layers.Conv2D(filter_num,?kernel_size=3,?padding='same')self.gn2?=?tfa.layers.GroupNormalization(8)self.act2?=?tf.keras.activations.swishdef?call(self,?inputs,?training=False,?*args,?**kwargs):residual?=?self.conv1(inputs)x?=?self.gn1(residual)x?=?tf.nn.swish(x)x?=?self.conv2(x)x?=?self.gn2(x)x?=?tf.nn.swish(x)out?=?x?+?residualreturn?out?/?1.44class?SimpleDDPMModel(tf.keras.Model):def?__init__(self,?max_time_step=100):super(SimpleDDPMModel,?self).__init__()#?定義ddpm?前向過(guò)程的一些參數(shù)self.max_time_step?=?max_time_step#?采用numpy?的float64，避免連乘的精度失準(zhǔn)betas?=?np.linspace(1e-4,?0.02,?max_time_step,?dtype=np.float64)alphas?=?1.0?-?betasalphas_bar?=?np.cumprod(alphas,?axis=0)betas_bar?=?1.0?-?alphas_barself.betas,?self.alphas,?self.alphas_bar,?self.betas_bar?=?tuple(map(lambda?x:?tf.constant(x,?tf.float32),[betas,?alphas,?alphas_bar,?betas_bar]))filter_nums?=?[64,?128,?256]self.encoders?=?[tf.keras.Sequential([ConvResidualLayer(num),tf.keras.layers.MaxPool2D(2)])?for?num?in?filter_nums]self.mid_conv?=?ConvResidualLayer(filter_nums[-1])self.decoders?=?[tf.keras.Sequential([tf.keras.layers.Conv2DTranspose(num,?3,?strides=2,?padding="same"),ConvResidualLayer(num),ConvResidualLayer(num),])?for?num?in?reversed(filter_nums)]self.final_conv?=?tf.keras.Sequential([ConvResidualLayer(64),tf.keras.layers.Conv2D(1,?3,?padding="same")])self.img_size?=?32self.time_embeddings?=?[tf.keras.Sequential([tf.keras.layers.Dense(num,?activation=tf.keras.layers.LeakyReLU()),tf.keras.layers.Dense(num)])for?num?in?filter_nums]#?實(shí)現(xiàn)公式?64?從原始數(shù)據(jù)生成噪音圖像def?q_noisy_sample(self,?x_0,?t,?noisy):alpha_bar,?beta_bar?=?self.extract([self.alphas_bar,?self.betas_bar],?t)sqrt_alpha_bar,?sqrt_beta_bar?=?tf.sqrt(alpha_bar),?tf.sqrt(beta_bar)return?sqrt_alpha_bar?*?x_0?+?sqrt_beta_bar?*?noisydef?extract(self,?sources,?t):bs?=?tf.shape(t)[0]targets?=?[tf.gather(source,?t)?for?i,?source?in?enumerate(sources)]return?tuple(map(lambda?x:?tf.reshape(x,?[bs,?1,?1,?1]),?targets))#?實(shí)現(xiàn)公式?131，從噪聲數(shù)據(jù)恢復(fù)上一步的數(shù)據(jù)def?p_real_sample(self,?x_t,?t,?pred_noisy):alpha,?beta,?beta_bar?=?self.extract([self.alphas,?self.betas,?self.betas_bar],?t)noisy?=?tf.random.normal(shape=tf.shape(x_t))#?這里的噪聲系數(shù)和beta取值一樣，也可以滿足越靠近0，噪聲越小noisy_weight?=?tf.sqrt(beta)#?當(dāng)t==0?時(shí)，不加入隨機(jī)噪聲bs?=?tf.shape(x_t)[0]noisy_mask?=?tf.reshape(1?-?tf.cast(tf.equal(t,?0),?tf.float32),?[bs,?1,?1,?1])noisy_weight?*=?noisy_maskx_t_1?=?(x_t?-?beta?*?pred_noisy?/?tf.sqrt(beta_bar))?/?tf.sqrt(alpha)?+?noisy?*?noisy_weightreturn?x_t_1#?unet?的下采樣def?encoder(self,?noisy_img,?t,?data,?training):xs?=?[]for?idx,?conv?in?enumerate(self.encoders):noisy_img?=?conv(noisy_img)t?=?tf.cast(t,?tf.float32)time_embedding?=?self.time_embeddings[idx](t)time_embedding?=?tf.reshape(time_embedding,?[-1,?1,?1,?tf.shape(time_embedding)[-1]])#?time?embedding?直接相加noisy_img?+=?time_embeddingxs.append(noisy_img)return?xs#?unet的上采樣def?decoder(self,?noisy_img,?xs,?training):xs.reverse()for?idx,?conv?in?enumerate(self.decoders):noisy_img?=?conv(tf.concat([xs[idx],?noisy_img],?axis=-1))return?noisy_img@tf.functiondef?pred_noisy(self,?data,?training):img?=?data["img_data"]bs?=?tf.shape(img)[0]noisy?=?tf.random.normal(shape=tf.shape(img))t?=?data.get("t",?None)#?在訓(xùn)練階段t為空，隨機(jī)生成成tif?t?is?None:t?=?tf.random.uniform(shape=[bs,?1],?minval=0,?maxval=self.max_time_step,?dtype=tf.int32)noisy_img?=?self.q_noisy_sample(img,?t,?noisy)else:noisy_img?=?imgxs?=?self.encoder(noisy_img,?t,?data,?training)x?=?self.mid_conv(xs[-1])x?=?self.decoder(x,?xs,?training)pred_noisy?=?self.final_conv(x)return?{"pred_noisy":?pred_noisy,?"noisy":?noisy,"loss":?tf.reduce_mean(tf.reduce_sum((pred_noisy?-?noisy)?**?2,?axis=(1,?2,?3)),?axis=-1)}#?生成圖片def?call(self,?inputs,?training=None,?mask=None):bs?=?inputs[0]x_t?=?tf.random.normal(shape=[bs,?self.img_size,?self.img_size,?1])for?i?in?reversed(range(0,?self.max_time_step)):t?=?tf.reshape(tf.repeat(i,?bs),?[bs,?1])p?=?self.pred_noisy({"img_data":?x_t,?"t":?t},?False)x_t?=?self.p_real_sample(x_t,?t,?p["pred_noisy"])return?x_tdef?train_step(self,?data):with?tf.GradientTape()?as?tape:result?=?self.pred_noisy(data,?True)trainable_vars?=?self.trainable_variablesgradients?=?tape.gradient(result["loss"],?trainable_vars)self.optimizer.apply_gradients(zip(gradients,?trainable_vars))return?{"loss":?result["loss"]}def?test_step(self,?data):result?=?self.pred_noisy(data,?False)return?{"loss":?result["loss"]}

生成的圖片如下：

類似 CVAE，使用 DDPM 的時(shí)候，我們依然希望可以通過(guò)條件控制生成，如前邊提到的 DALLE-2，Stable Diffusion 都是通過(guò)條件（文本 prompt）來(lái)控制生成的圖像，為了實(shí)現(xiàn)這個(gè)目的，就需要采用 Conditional Diffusion Model。

Conditional Diffusion Model

目前最主要使用的 Conditional Diffusion Model 主要有兩種實(shí)現(xiàn)方式，Classifier-guidance 和 Classifier-free，從名字也可以看出，前者需要一個(gè)分類器模型，后者無(wú)需分類器模型，下邊講簡(jiǎn)單推導(dǎo)兩種的實(shí)現(xiàn)方案，并給出? Classifier-free Diffusion Model 的實(shí)現(xiàn)代碼。

1. Classifier-guidance

參考前邊的推導(dǎo)公式在無(wú)條件的模型下，我們需要優(yōu)化；而在控制條件 y 下，我們需要優(yōu)化的是，可以用貝葉斯進(jìn)行以下的公式推導(dǎo)：

從以下公式推導(dǎo)可以看出，我們需要一個(gè)分類模型，這個(gè)分類模型可以對(duì)前向過(guò)程融入噪音的數(shù)據(jù)很好的分類，在擴(kuò)散模型求梯度的階段，融入這個(gè)分類模型對(duì)當(dāng)前噪音數(shù)據(jù)的梯度即可。

2. Classifier-free

通過(guò) classifier-guidance 的公式證明，我們很容易得到以下的公式推導(dǎo)：

取值 0~1 之間，從公式 140 可以看出，只要我們?cè)谀Ｐ洼斎肷?#xff0c;采樣性的融入 y 就可以達(dá)到目標(biāo)，所以在前邊的 DDPM 代碼上改動(dòng)比較簡(jiǎn)單，我們對(duì) 0~9 這 10 個(gè)數(shù)字學(xué)習(xí)一個(gè) embedding 表示，然后采樣性的加入 unet 的 encoder 的階段，代碼如下：

class?SimpleCDDPMModel(SimpleDDPMModel):def?__init__(self,?max_time_step=100,?label_num=10):super(SimpleCDDPMModel,?self).__init__(max_time_step=max_time_step)#?condition?的embedding和time?step的一致self.condition_embedding?=?[tf.keras.Sequential([tf.keras.layers.Embedding(label_num,?num),tf.keras.layers.Dense(num)])for?num?in?self.filter_nums]#?unet?的下采樣def?encoder(self,?noisy_img,?t,?data,?training):xs?=?[]mask?=?tf.random.uniform(shape=(),?minval=0.0,?maxval=1.0,?dtype=tf.float32)for?idx,?conv?in?enumerate(self.encoders):noisy_img?=?conv(noisy_img)t?=?tf.cast(t,?tf.float32)time_embedding?=?self.time_embeddings[idx](t)time_embedding?=?tf.reshape(time_embedding,?[-1,?1,?1,?tf.shape(time_embedding)[-1]])#?time?embedding?直接相加noisy_img?+=?time_embedding#?獲取?condition?的embeddingcondition_embedding?=?self.condition_embedding[idx](data["label"])condition_embedding?=?tf.reshape(condition_embedding,?[-1,?1,?1,?tf.shape(condition_embedding)[-1]])#?訓(xùn)練階段一定的概率下加入condition，推理階段全部加入if?training:if?mask?<?0.15:condition_embedding?=?tf.zeros_like(condition_embedding)noisy_img?+=?condition_embeddingxs.append(noisy_img)return?xs#?生成圖片def?call(self,?inputs,?training=None,?mask=None):bs?=?inputs[0]label?=?tf.reshape(tf.repeat(inputs[1],?bs),?[-1,?1])x_t?=?tf.random.normal(shape=[bs,?self.img_size,?self.img_size,?1])for?i?in?reversed(range(0,?self.max_time_step)):t?=?tf.reshape(tf.repeat(i,?bs),?[bs,?1])p?=?self.pred_noisy({"img_data":?x_t,?"t":?t,?"label":?label},?False)x_t?=?self.p_real_sample(x_t,?t,?p["pred_noisy"])return?x_t

最終生成的圖片如下：

參考文獻(xiàn)

[1] https://www.jarvis73.com/2022/08/08/Diffusion-Model-1/

[2]?https://blog.csdn.net/qihangran5467/article/details/118337892

[3] https://jaketae.github.io/study/vae/

[4] https://pyro.ai/examples/cvae.html

[5] https://lilianweng.github.io/posts/2021-07-11-diffusion-models/

[6] https://spaces.ac.cn/archives/9164

[7] https://zhuanlan.zhihu.com/p/575984592

[8] https://kxz18.github.io/2022/06/19/Diffusion/

[9] https://zhuanlan.zhihu.com/p/502668154

[10]?https://xyfjason.top/2022/09/29/%E4%BB%8EVAE%E5%88%B0DDPM/

[11] https://arxiv.org/pdf/2208.11970.pdf

更多閱讀

#投 稿?通 道#

?讓你的文字被更多人看到?

如何才能讓更多的優(yōu)質(zhì)內(nèi)容以更短路徑到達(dá)讀者群體，縮短讀者尋找優(yōu)質(zhì)內(nèi)容的成本呢？答案就是：你不認(rèn)識(shí)的人。

總有一些你不認(rèn)識(shí)的人，知道你想知道的東西。PaperWeekly 或許可以成為一座橋梁，促使不同背景、不同方向的學(xué)者和學(xué)術(shù)靈感相互碰撞，迸發(fā)出更多的可能性。?

PaperWeekly 鼓勵(lì)高校實(shí)驗(yàn)室或個(gè)人，在我們的平臺(tái)上分享各類優(yōu)質(zhì)內(nèi)容，可以是最新論文解讀，也可以是學(xué)術(shù)熱點(diǎn)剖析、科研心得或競(jìng)賽經(jīng)驗(yàn)講解等。我們的目的只有一個(gè)，讓知識(shí)真正流動(dòng)起來(lái)。

📝?稿件基本要求：

? 文章確系個(gè)人原創(chuàng)作品，未曾在公開(kāi)渠道發(fā)表，如為其他平臺(tái)已發(fā)表或待發(fā)表的文章，請(qǐng)明確標(biāo)注?

? 稿件建議以?markdown?格式撰寫(xiě)，文中配圖以附件形式發(fā)送，要求圖片清晰，無(wú)版權(quán)問(wèn)題

? PaperWeekly 尊重原作者署名權(quán)，并將為每篇被采納的原創(chuàng)首發(fā)稿件，提供業(yè)內(nèi)具有競(jìng)爭(zhēng)力稿酬，具體依據(jù)文章閱讀量和文章質(zhì)量階梯制結(jié)算

📬?投稿通道：

? 投稿郵箱：hr@paperweekly.site?

? 來(lái)稿請(qǐng)備注即時(shí)聯(lián)系方式（微信），以便我們?cè)诟寮x用的第一時(shí)間聯(lián)系作者

? 您也可以直接添加小編微信（pwbot02）快速投稿，備注：姓名-投稿

△長(zhǎng)按添加PaperWeekly小編

🔍

現(xiàn)在，在「知乎」也能找到我們了

進(jìn)入知乎首頁(yè)搜索「PaperWeekly」

點(diǎn)擊「關(guān)注」訂閱我們的專欄吧

·

·

總結(jié)

以上是生活随笔為你收集整理的AIGC基础：从VAE到DDPM原理、代码详解的全部?jī)?nèi)容，希望文章能夠幫你解決所遇到的問(wèn)題。

如果覺(jué)得生活随笔網(wǎng)站內(nèi)容還不錯(cuò)，歡迎將生活随笔推薦給好友。