在XGen-Image-1之中:Salesforce Research如何构建、训练和评估一个庞大的文本到图像模型
XGen-Image-1 How Salesforce Research builds, trains, and evaluates a large text-to-image model
有史以来实施的最高效的文本到图像模型训练过程之一。
我最近开始了一个以人工智能为重点的教育性通讯,已经有超过16万订阅者。TheSequence是一个不含废话(意思是没有炒作,没有新闻等等)以机器学习为导向的通讯,阅读时间只需5分钟。目标是让您了解机器学习项目、研究论文和概念的最新信息。请订阅下方的链接试试:
TheSequence | Jesus Rodriguez | Substack
机器学习、人工智能和数据发展的最佳信息来源…
thesequence.substack.com
Salesforce一直是新一波基础模型研究实验室中最活跃的研究部门之一。在最近几个月中,Salesforce Research在语言、编码和计算机视觉等不同领域发布了多种模型。最近,他们推出了XGen-Image-1,这是一个巨大的文本到图像模型,在不同的计算机视觉任务上展现了最先进的性能。与发布同时,Salesforce Research还提供了关于XGen-Image-1训练方法和最佳实践的详细信息。今天,我想深入探讨其中的一些细节。构建XGen-Image-1的过程涉及一系列非常棘手的决策,从战略设计选择和训练方法到关于他们首次图像生成模型的性能指标。
- 模型训练:从架构角度来看,XGen-Image-1是一个拥有8.6亿参数的潜在扩散模型。在训练方面,Salesforce Research利用了11亿公开可用的LAION文本-图像对数据集。• 分辨率管理:XGen-Image-1使用了潜在变量自编码器(VAE),配合现有的像素上采样器,有效地降低了低分辨率训练的计算成本。• 成本效益:在进行文本到图像模型训练时,成本始终是一个令人关注的变量。对于XGen-Image-1来说,Salesforce Research确定,使用Google的TPU堆栈培养出一种竞争力的图像生成模型,只需投资约75,000美元。• 性能平衡:值得注意的是,XGen-Image-1展示了与稳定扩散1.5和2.1相当的对齐性能,这些模型是图像生成能力的先锋之一。• 自动改进:一个显著的方面是通过修饰在特定区域进行自动增强。例如,像“脸”这样的区域经过改进,有效地提升了生成的图像。• 推理增强:通过在推理过程中集成拒绝采样,结果质量得到了显著提升,证明了深思熟虑的方法论增强的强大力量。
训练
大多数文本到图像模型都开创了自己的训练方法。XGen-Image-1则采取了不同的方法,重点不是在于条件、编码和上采样等直接的挑战,而是在于它们的可重用性的程度。这种独特的视角让他们尝试了高效训练的极限,通过预训练的自编码器和基于像素的上采样模型的结合,探索了低分辨率训练的边界。
正如管道所示,预训练的自编码器结合可选的基于像素的上采样器,在生成低分辨率(1024×1024)的同时产生了宏伟的图像。他们的展望还延伸到了分辨率的实用下限的进一步探索。定量评估发生在256×256的位置,直接在没有上采样的VAE阶段之后进行。定性评估利用了一个“重新上采样器”,从256到64到256的过程,镜像了SDXL的“Refiner”,以及一个从256到1024的上采样器。
训练基础设施
Salesforce Research在TPU v4上进行了模型训练。在整个训练过程中,团队利用Google Cloud Storage(GCS)保存模型检查点。此外,还利用Gloud-mounted驱动器存储庞大的数据集。在一台v4–512 TPU机器上进行了强大的训练,历时约9天,共进行了110万步,硬件成本约为73,000美元,与其他选择相比非常便宜!
这个过程并不是没有挑战。尽管批量大小很大,但模型的损失在每一步都表现出不稳定的波动。根本原因与所有工作器的统一播种有关,导致了这种波动性。通过随机播种来解决这个问题,与每个工作器的等级保持一致,以促进噪声步骤的平衡分散,并带来更平滑的损失曲线。
评估
Salesforce Research使用了一种方法来评估XGen-Image-1,该方法包括双轴表示,其中x轴表示CLIP分数-一种与提示对齐的体现,而y轴则包含FID(Fréchet Inception Distance),这是一个指示数据集上外观相似性的指标。这种评估涉及到15个引导尺度,涉及到初始图中的30,000个图像-提示对。对于检查点间的比较,这个范围缩小到1,000对。数据对是精心从COCO Captions数据集中抽样的,其中标题被丰富为“一张照片”,以减轻与不同图形风格相关的FID惩罚。
Salesforce Research还使用人工评估来对XGen-Image-1进行基准测试。这个评估是通过亚马逊机械土耳其进行的,参与者需要判断图像-提示的对齐情况。为了确保精确性,参与者需要在六个不同的试验中与所有1,632个提示进行互动,最终得到每个比较约10,000个响应的大量数据。
生成高质量图像
XGen-Image-1构建的一个重要部分是确保它生成高质量的图像。这个过程基于两个众所周知的技巧。初始策略是生成大量图像,然后挑选最优的图像。坚持保持一个提示对应一个输出的原则,Salesforce Research致力于自动化这种方法。对这种方法的探索使他们开始研究拒绝采样-一种包括生成多幅图像,然后自动选择最合适的候选图像的技术。最初考虑了一系列指标,包括美学分数和CLIP分数。最终,团队发现PickScore是一个强大的综合指标。正如文献所证实的那样,这个指标与人类偏好表现出了显著的一致性-这是一个关键的特征。按照这种方法,XGen-Image-1在大约5秒钟内在A100 GPU上生成32张图像(以4×8的网格形式),并选择得分最高的图像。
第二种技术作为图像改进的一个例子,通过修复呈现次优美学效果的区域来实现。
- 从提示中延伸出的分割掩码封装了目标对象。2. 基于这些分割掩码对对象进行了裁剪。3. 裁剪进行了扩展。4. 利用img2img操作对裁剪进行处理,将分割提示与匹配的标题(例如,“一张脸部照片”)并行,促进区域的优化。5. 分割掩码在将增强的裁剪区域无缝融合回原始图像中起到了关键作用。这两种方法的结合使XGen-Image-1能够创建高质量的图像。
XGen-Image-1在文本到图像模型中并不是一个突破,但它仍然提供了一些非常有用的贡献。最大的贡献是展示了可以重用许多预训练组件来实现这种模型的最新性能。X-Gen-Image-1在使用的训练计算资源的部分情况下实现了与StableDiffusion类似的性能。非常令人印象深刻。