迷你版 DALL-E:模型缩小 27 倍,训练成本仅 200 美元,在线可玩

自从今年 1 月份 OpenAI 祭出 120 亿参数魔法模型 DALL・E,众多科研工作者开始对其进行多方位的剖析,希望也能做出给定“自然语言文字描述直接生成对应图片”的程序或软件。

120 亿参数……… 这意味着只有“顶级玩家”才能入局。如果缩小模型参数,程序还能跑出高质量的结果么?

今日,谷歌和 Hugging Face 的研究员们开放了一款 DALL・E mini 的小程序。顾名思义,作为 DALL・E 的复刻版,模型规模比原版小了 27 倍,且只在单个 TPU v3-8 上训练了 3 天。

mini 版的 DALL・E 创造思路非常简单:在官方开源代码以及预训练模型的基础上,简化体系结构和减少模型内存需求。

效果如何?作者在博客中说:

“在硬件资源受限的情况下,虽然质量比较低,但结果仍然是令人印象深刻的。”

AI 科技评论将其与正版 DALL・E 进行了对比,结果如下:

文本提示:穿着芭蕾舞短裙的萝卜宝宝在遛狗。

AI 生成的图像:

图注:左边为 DALL・E mini 输出结果,右边为 OpenAI 官方博客中的示范

文本提示:写着“OpenAI”字样的店面。

AI 生成的图像:

图注:左边为 DALL・E mini 输出结果,右边为 OpenAI 官方博客中的示范

文本提示:鳄梨形状的扶手椅。

AI 生成的图像:

图注:左边为 DALL・E mini 输出结果,右边为 OpenAI 官方博客中的示范

通过以上对比,mini 版的 DALL・E 只在“鳄梨形状的扶手椅”交出了令人满意的答卷,在“遛狗”和“店面”测试中,前者生成的是抽象派画面,后者只生成了店面…… 果然,魔法还是念完完整咒语,才能发挥最完美的效果。

目前,已经开放 demo,在线可玩。

https://huggingface.co/spaces/flax-community/dalle-mini

考虑到较小规模的模型架构,以及只训练了三天,因此,其模型结构和训练方法对我们也有借鉴意义。

模型架构 

图注:模型训练流程

1. 图像通过 VQGAN 编码器进行编码,该编码器将图像转换为一系列 tokens。

2. 文字描述通过 BART 编码器进行编码。其中,BART 解码器是一个自回归模型,其目标是预测下一个 tokens。

3. 损失函数来自 VQGAN 中真实图像编码和预测值之间的 softmax cross entropy。

图注:模型推理流程

在推理时候,只有标题(captions)可用,图像生成过程如下:

1. 标题通过 BART 编码器编码

2.token 通过通过 BART 解码器传送

3. 基于解码器在下一个 token 的预测分布,对图像 token 进行顺序采样

4. 图像 token 序列通过 VQGAN 解码器解码

5.Clip 用于选择最佳生成的图像

前面也提到,对于图像的编码和解码,作者使用 VQGAN。VQGAN 的目标是将一幅图像编码成一系列离散的 token,这些 token 可用于 transformers 模型。

图注:训练 transformers 以实现高分辨率图像合成

值得一提的是,为了充分利用 TPU,该模型用 JAX 编程。此外,为了加快数据加载速度,作者使用图像编码器对所有图像进行了预编码。

其中,每个 TPU 处理的 Batch size 为 56,这可以最大限度地增加每个 TPU 的可用内存;采用 Adafactor 为优化器,提高内存效率,使我们可以使用更高的 Batch size。作者们还花了半天的时间,通过启动超参数搜索来为模型找到一个好的学习率(0.005)。

与 OpenAI DALL・E 相比 

1.Dall・E 使用 120 亿参数版本的 GPT-3。相比之下,Dall・E mini 的模型小 27 倍,参数约为 4 亿。

2.Dall・E mini 使用大量预先训练好的模型 (VQGAN、BART 编码器和 CLIP),而 OpenAI 从头开始训练。

3. 在图像编码方面,表 (vocabulary) 的对比是 8192 vs 16384,tokens 的对比是 1024vs256。Dall・E 使用 VQVAE,而作者使用 VQGAN。

4. 在文本编码方面,表 (vocabulary) 的对比是 16384 vs 50264,tokens 的对比是 256 vs 1024。

5.Dall・E 通过自回归模型读取文本,而 Dall・E mini 使用双向编码器。

6.Dall・E 接受了 2.5 亿对图像和文本的训练,而 Dall・E mini 只使用了 1500 万对。

综上,以上的设置差异使得 Dall・E mini 的训练时间大大缩小,即只在单个 TPU v3-8 上训练了 3 天。据作者计算,模型的训练成本不到 200 美元,即使算上在 TPU 和超参数搜索上的实验,额外的开销也不超过 1000 美元。

最后,作者写道:

基于此,Dall・E 生成的图像的质量虽然比 Dall・E mini 模型高得多,但调用的资源非常少。不得不承认这个模型在复现 OpenAI 几个官方示例中确实失败了。

同时,作者也放出了一些实验结果,如下所示:

3

差在哪?

作者也承认了 Dall・E mini 的几个缺陷,在文末,他们总结到:

1. 生成的图片经常有水印;

2. 人物和动物的生成往往有缺陷;

3. 很难预测该模型的优势和缺陷,例如该模型很擅长生成“鳄梨形状的扶手椅”,但不能生成与“计算机 logo ”相关的任何内容;

4. 目前仅接受英文输入,在其他语言中表现不佳。

经过分析,作者认为之所以出现缺陷,可能是因为数据集,毕竟数据偏见一直是研究社区所讨论的;另外,模型本身以及训练的流程、推理的流程都有可能是出问题的地方。

Published by

风君子

独自遨游何稽首 揭天掀地慰生平

发表回复

您的电子邮箱地址不会被公开。 必填项已用 * 标注