VAE简介
典型的自编码器是确定性的,而VAE则是概率模型,因为它将潜在空间建模为概率分布。VAE是无监督模型,它将输入数据x编码为潜在表示z,并从这个潜在空间中重建输入。从技术上讲,VAE并不需要用神经网络来实现,也可以由生成概率模型构建。然而,在当前深度学习的状态下,大多数VAE通常都是用神经网络来实现的。
简而言之,重参数化技巧之所以被使用,是因为我们无法对潜在空间的概率分布进行反向传播,但我们需要更新我们的编码分布。因此,我们定义了一个可微且可逆的函数,这样我们就可以对λ和x进行求导,同时仍然保留一个概率元素。
VAE(变分自编码器)使用由重构项和编码模型相对于先验分布的Kullback-Leibler散度(KLD)组成的ELBO(证据下界)损失进行训练。
向VAE添加条件输入
CVAE(条件变分自编码器)通过纳入如类别标签等额外信息作为条件变量来扩展VAE。这种条件设定使得CVAE能够生成受控的内容。条件输入特征可以在架构的不同位置添加,但通常会被插入到编码器和解码器中。带有条件输入的损失函数是传统VAE中ELBO损失的一种适应形式。
为了说明VAE(变分自编码器)和CVAE(条件变分自编码器)之间的区别,我们使用卷积编码器和解码器架构在Fashion-MNIST数据集上训练了这两种网络。图中展示了每个网络潜在空间的t-SNE(t-distributed Stochastic Neighbor Embedding)可视化结果。
原始的VAE显示出明显的聚类,而CVAE的分布则更加均匀。由于没有提供条件信号,原始的VAE会将类别和类别变化编码到潜在空间中。然而,CVAE不需要学习类别区分,潜在空间可以专注于类别内的变化。因此,CVAE可能能够学习更多信息,因为它不依赖于必须学习基本的类别条件。
CVAE的模型架构
为了测试图像生成,我们创建了两种模型架构。第一种架构是采用连接条件方法的卷积CVAE。所有网络都是为大小为28x28(总共784个像素)的Fashion-MNIST图像构建的。
class ConcatConditionalVAE(nn.Module):
def __init__(self, latent_dim=128, num_classes=10):
super().__init__()
self.latent_dim = latent_dim
self.num_classes = num_classes
# Encoder
self.encoder = nn.Sequential(
nn.Conv2d(1, 32, 3, stride=2, padding=1),
nn.ReLU(),
nn.Conv2d(32, 64, 3, stride=2, padding=1),
nn.ReLU(),
nn.Conv2d(64, 128, 3, stride=2, padding=1),
nn.ReLU(),
nn.Flatten()
)
self.flatten_size = 128 * 4 * 4
# Conditional embedding
self.label_embedding = nn.Embedding(num_classes, 32)
# Latent space (with concatenated condition)
self.fc_mu = nn.Linear(self.flatten_size + 32, latent_dim)
self.fc_var = nn.Linear(self.flatten_size + 32, latent_dim)
# Decoder
self.decoder_input = nn.Linear(latent_dim + 32, 4 * 4 * 128)
self.decoder = nn.Sequential(
nn.ConvTranspose2d(128, 64, 2, stride=2, padding=1, output_padding=1),
nn.ReLU(),
nn.ConvTranspose2d(64, 32, 3, stride=2, padding=1, output_padding=1),
nn.ReLU(),
nn.ConvTranspose2d(32, 1, 3, stride=2, padding=1, output_padding=1),
nn.Sigmoid()
)
def encode(self, x, c):
x = self.encoder(x)
c = self.label_embedding(c)
# Concatenate condition with encoded input
x = torch.cat([x, c], dim=1)
mu = self.fc_mu(x)
log_var = self.fc_var(x)
return mu, log_var
def reparameterize(self, mu, log_var):
std = torch.exp(0.5 * log_var)
eps = torch.randn_like(std)
return mu + eps * std
def decode(self, z, c):
c = self.label_embedding(c)
# Concatenate condition with latent vector
z = torch.cat([z, c], dim=1)
z = self.decoder_input(z)
z = z.view(-1, 128, 4, 4)
return self.decoder(z)
def forward(self, x, c):
mu, log_var = self.encode(x, c)
z = self.reparameterize(mu, log_var)
return self.decode(z, c), mu, log_var
CVAE编码器由3个卷积层组成,每个卷积层后面都跟着一个ReLU非线性激活函数。编码器的输出随后被展平。然后,类别编号通过一个嵌入层,并与编码器输出相加。接着,使用重参数化技巧,通过2个线性层获得潜在空间中的μ和σ。一旦采样完成,重参数化潜在空间的输出就被传递到解码器,现在与类别编号嵌入层输出相连接。解码器由3个转置卷积层组成。前两个包含ReLU非线性激活函数,最后一层包含sigmoid非线性激活函数。解码器的输出是一个28x28的生成图像。
另一种模型架构遵循相同的方法,但采用的是添加条件输入而不是连接的方式。一个主要的问题是,添加或连接哪种方式会导致更好的重建或生成结果。
class AdditiveConditionalVAE(nn.Module):
def __init__(self, latent_dim=128, num_classes=10):
super().__init__()
self.latent_dim = latent_dim
self.num_classes = num_classes
# Encoder
self.encoder = nn.Sequential(
nn.Conv2d(1, 32, 3, stride=2, padding=1),
nn.ReLU(),
nn.Conv2d(32, 64, 3, stride=2, padding=1),
nn.ReLU(),
nn.Conv2d(64, 128, 3, stride=2, padding=1),
nn.ReLU(),
nn.Flatten()
)
self.flatten_size = 128 * 4 * 4
# Conditional embedding
self.label_embedding = nn.Embedding(num_classes, self.flatten_size)
# Latent space (without concatenation)
self.fc_mu = nn.Linear(self.flatten_size, latent_dim)
self.fc_var = nn.Linear(self.flatten_size, latent_dim)
# Decoder condition embedding
self.decoder_label_embedding = nn.Embedding(num_classes, latent_dim)
# Decoder
self.decoder_input = nn.Linear(latent_dim, 4 * 4 * 128)
self.decoder = nn.Sequential(
nn.ConvTranspose2d(128, 64, 2, stride=2, padding=1, output_padding=1),
nn.ReLU(),
nn.ConvTranspose2d(64, 32, 3, stride=2, padding=1, output_padding=1),
nn.ReLU(),
nn.ConvTranspose2d(32, 1, 3, stride=2, padding=1, output_padding=1),
nn.Sigmoid()
)
def encode(self, x, c):
x = self.encoder(x)
c = self.label_embedding(c)
# Add condition to encoded input
x = x + c
mu = self.fc_mu(x)
log_var = self.fc_var(x)
return mu, log_var
def reparameterize(self, mu, log_var):
std = torch.exp(0.5 * log_var)
eps = torch.randn_like(std)
return mu + eps * std
def decode(self, z, c):
# Add condition to latent vector
c = self.decoder_label_embedding(c)
z = z + c
z = self.decoder_input(z)
z = z.view(-1, 128, 4, 4)
return self.decoder(z)
def forward(self, x, c):
mu, log_var = self.encode(x, c)
z = self.reparameterize(mu, log_var)
return self.decode(z, c), mu, log_var
从上式可以看出,所有 CVAE 都使用相同的损失函数。
def loss_function(recon_x, x, mu, logvar):
"""Computes the loss = -ELBO = Negative Log-Likelihood + KL Divergence.
Args:
recon_x: Decoder output.
x: Ground truth.
mu: Mean of Z
logvar: Log-Variance of Z
"""
BCE = F.binary_cross_entropy(recon_x, x, reduction='sum')
KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
return BCE + KLD
为了评估模型生成的图像,通常使用三种定量指标。均方误差(MSE)是通过逐像素计算生成图像与真实图像之间差值的平方和来得出的。结构相似性指数(SSIM)是一种通过比较两张图像的结构信息、亮度和对比度来评估图像质量的指标。SSIM可以用于比较任何大小的图像,而MSE则与像素大小相关。SSIM的得分范围从-1到1,其中1表示两张图像完全相同。弗雷谢特初始距离(FID)是一种量化生成图像的真实性和多样性的指标。由于FID是一种距离度量,因此较低的得分表明对一组图像的重建效果更好。
从Fashion-MNIST的简短文本到图像
在将规模扩大到从完整文本到图像的转换之前,我们先在Fashion-MNIST上测试了CVAE的图像重建和生成能力。Fashion-MNIST是一个类似于MNIST的数据集,包含60,000个训练样本和10,000个测试样本。每个样本都是一张28x28的灰度图像,与10个类别中的一个标签相关联。
我们创建了预处理函数,用于从输入的简短文本正则表达式匹配中提取包含类别名称的相关关键词。为了考虑每个类别中包含的类似时尚单品,我们为大多数类别使用了额外的描述符(同义词)(例如,外套和夹克)。
classes = {
'Shirt':0,
'Top':0,
'Trouser':1,
'Pants':1,
'Pullover':2,
'Sweater':2,
'Hoodie':2,
'Dress':3,
'Coat':4,
'Jacket':4,
'Sandal':5,
'Shirt':6,
'Sneaker':7,
'Shoe':7,
'Bag':8,
'Ankle boot':9,
'Boot':9
}
def word_to_text(input_str, classes, model, device):
label = class_embedding(input_str, classes)
if label == -1: return Exception("No valid label")
samples = sample_images(model, num_samples=4, label=label, device=device)
plot_samples(samples, input_str, torch.tensor([label]))
return
def class_embedding(input_str, classes):
for key in list(classes.keys()):
template = f'(?i)\\b{key}\\b'
output = re.search(template, input_str)
if output: return classes[key]
return -1
然后,将类别名称转换为其对应的类别编号,并作为条件输入与CVAE一起使用。为了生成图像,将从简短文本描述中提取的类别标签与从高斯分布中随机采样的样本一起传递到解码器中,以输入潜在空间中的变量。
在测试生成之前,先测试图像重建以确保CVAE的功能正常。由于创建的是处理28x28图像的卷积网络,因此可以在不到一个小时的时间内,用不到100个训练周期来完成网络训练。
重建图像包含了真实图像的大致形状,但图像中缺少了尖锐的高频特征。模型输出中的任何文本或复杂设计图案都是模糊的。输入包含Fashion-MNIST中任一类别的任何简短文本,都会生成与重建图像相似的输出。
生成的图像具有11的均方误差(MSE)和0.76的结构相似性指数(SSIM)。这些指标表明生成效果良好,意味着在简单、小尺寸的图像上,条件变分自编码器(CVAE)能够生成质量不错的图像。虽然生成对抗网络(GAN)和去噪扩散概率模型(DDPM)能够生成具有复杂特征的高质量图像,但CVAE在处理简单案例时同样表现出色。
使用CLIP和COCO从长文本到图像的转换
当将图像生成扩展到任意长度的文本时,除了正则表达式匹配之外,还需要更稳健的方法。为此,我们使用了OpenAI的CLIP模型,将文本转换为高维嵌入向量。嵌入模型采用其ViT-B/32配置,输出长度为512的嵌入向量。CLIP模型的一个限制是,它的最大标记长度为77,而研究表明,其有效长度甚至更小,仅为20。因此,在输入文本包含多个句子的情况下,我们将文本按句子拆分,并通过CLIP编码器进行处理。然后,将得到的嵌入向量进行平均,以创建最终的输出嵌入。
与Fashion-MNIST相比,长文本模型需要更复杂的训练数据,因此我们使用了COCO数据集。COCO数据集包含注释,这些注释可以传递给CLIP以获取嵌入向量。然而,COCO图像的尺寸为640x480,这意味着即使进行裁剪变换,也需要一个更大的网络来处理。对于从长文本到图像的生成,我们测试了添加和连接条件输入这两种架构,但在这里仅展示连接方法的结果。
class cVAE(nn.Module):
def __init__(self, latent_dim=128):
super().__init__()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.clip_model, _ = clip.load("ViT-B/32", device=device)
self.clip_model.eval()
for param in self.clip_model.parameters():
param.requires_grad = False
self.latent_dim = latent_dim
# Modified encoder for 128x128 input
self.encoder = nn.Sequential(
nn.Conv2d(3, 32, 4, stride=2, padding=1), # 64x64
nn.BatchNorm2d(32),
nn.ReLU(),
nn.Conv2d(32, 64, 4, stride=2, padding=1), # 32x32
nn.BatchNorm2d(64),
nn.ReLU(),
nn.Conv2d(64, 128, 4, stride=2, padding=1), # 16x16
nn.BatchNorm2d(128),
nn.ReLU(),
nn.Conv2d(128, 256, 4, stride=2, padding=1), # 8x8
nn.BatchNorm2d(256),
nn.ReLU(),
nn.Conv2d(256, 512, 4, stride=2, padding=1), # 4x4
nn.BatchNorm2d(512),
nn.ReLU(),
nn.Flatten()
)
self.flatten_size = 512 * 4 * 4 # Flattened size from encoder
# Process CLIP embeddings for encoder
self.condition_processor_encoder = nn.Sequential(
nn.Linear(512, 1024)
)
self.fc_mu = nn.Linear(self.flatten_size + 1024, latent_dim)
self.fc_var = nn.Linear(self.flatten_size + 1024, latent_dim)
self.decoder_input = nn.Linear(latent_dim + 512, 512 * 4 * 4)
# Modified decoder for 128x128 output
self.decoder = nn.Sequential(
nn.ConvTranspose2d(512, 256, 4, stride=2, padding=1), # 8x8
nn.BatchNorm2d(256),
nn.ReLU(),
nn.ConvTranspose2d(256, 128, 4, stride=2, padding=1), # 16x16
nn.BatchNorm2d(128),
nn.ReLU(),
nn.ConvTranspose2d(128, 64, 4, stride=2, padding=1), # 32x32
nn.BatchNorm2d(64),
nn.ReLU(),
nn.ConvTranspose2d(64, 32, 4, stride=2, padding=1), # 64x64
nn.BatchNorm2d(32),
nn.ReLU(),
nn.ConvTranspose2d(32, 16, 4, stride=2, padding=1), # 128x128
nn.BatchNorm2d(16),
nn.ReLU(),
nn.Conv2d(16, 3, 3, stride=1, padding=1), # 128x128
nn.Sigmoid()
)
def encode_condition(self, text):
with torch.no_grad():
embeddings = []
for sentence in text:
embeddings.append(self.clip_model.encode_text(clip.tokenize(sentence).to('cuda')).type(torch.float32))
return torch.mean(torch.stack(embeddings), dim=0)
def encode(self, x, c):
x = self.encoder(x)
c = self.condition_processor_encoder(c)
x = torch.cat([x, c], dim=1)
return self.fc_mu(x), self.fc_var(x)
def reparameterize(self, mu, log_var):
std = torch.exp(0.5 * log_var)
eps = torch.randn_like(std)
return mu + eps * std
def decode(self, z, c):
z = torch.cat([z, c], dim=1)
z = self.decoder_input(z)
z = z.view(-1, 512, 4, 4)
return self.decoder(z)
def forward(self, x, c):
mu, log_var = self.encode(x, c)
z = self.reparameterize(mu, log_var)
return self.decode(z, c), mu, log_var
另一个主要的调查点是不同尺寸图像的生成和重建。具体来说,是将COCO图像修改为64x64、128x128和256x256这三种尺寸。在网络训练完成后,应首先测试重建结果。
所有尺寸的图像都能重建出带有一些特征轮廓和正确颜色的背景。然而,随着图像尺寸的增加,能够恢复的特征也更多。这是有道理的,因为虽然使用更大尺寸的图像训练模型会花费更长的时间,但模型能够捕获和学习的信息也更多。
在图像生成方面,要生成高质量的图像是极其困难的。大多数图像都在一定程度上有背景,并且图像中的特征模糊不清。这对于使用条件变分自编码器(CVAE)进行图像生成来说是预料之中的。这种情况在条件输入的连接和相加方法中都会出现,但连接方法的表现更好。这可能是因为连接的条件输入不会干扰重要特征,并且能确保信息被清晰地保留下来。如果条件不相关,则可以忽略它们。然而,相加的条件输入可能会干扰现有特征,并在反向传播期间更新权重时完全扰乱网络。
所有COCO生成的图像的SSIM都远低于Fashion-MNIST上的SSIM,大约为0.4。MSE与图像大小成正比,因此很难量化差异。为了进一步证明COCO CVAE生成的图像不够稳健,COCO图像生成的FID值都在200左右。
使用CVAE进行图像生成的局限性
尝试使用CVAE进行图像生成时,最大的局限性恰恰在于CVAE本身。能够包含和重建/生成的信息量极大地依赖于潜在空间的大小。潜在空间太小将无法捕获任何有意义的信息,并且与输出图像的大小成正比。一个28x28的图像所需的潜在空间远小于一个64x64的图像(因为它与图像大小的平方成正比)。然而,一个比实际图像还大的潜在空间会添加不必要的信息,并且在这一点上只会创建一个1对1的映射。对于COCO数据集,至少需要512的潜在空间来捕获一些特征。虽然CVAE是生成模型,但卷积编码器和解码器是一个相当基础的网络。GAN的训练风格或DDPM的复杂去噪过程允许进行更复杂的图像生成。
图像生成的另一个主要局限性在于训练所用的数据集。虽然COCO数据集有注释,但这些注释并不详细。为了训练复杂的生成模型,应该使用不同的数据集进行训练。COCO没有提供背景细节的位置或额外信息。来自CLIP编码器的复杂特征向量无法有效地在COCO上用于CVAE。
尽管CVAE和在COCO上进行图像生成有其局限性,但它创建了一个可行的图像生成模型。