aMUSEd：高效文本到图像生成

Hugging Face 推出了 aMUSEd，这是一款高效的非扩散文本到图像模型。作为 Google MUSE 的公开复现，aMUSEd 利用遮罩图像建模（MIM）提供比传统潜在扩散模型更快、更具可解释性的替代方案。

遮罩图像建模架构

aMUSEd 采用遮罩图像建模（MIM）方法，与潜在扩散相比，它能够使用更少的推理步骤并具有更高的可解释性。模型架构由三个主要组件组成：用于标记化的 VQGAN、用于提示嵌入的 CLIP-L/14 文本编码器，以及用于预测遮罩补丁的 U-ViT 模型。

训练过程

在训练期间，模型遵循以下步骤：

  1. Tokenization：使用 VQGAN 将输入图像转换为图像标记。
  2. Masking：根据余弦遮罩调度对图像标记进行遮罩。
  3. Prediction：将带有 CLIP-L/14 文本编码器生成的提示嵌入条件的遮罩标记传递给 U-ViT 模型，以预测遮罩补丁。

推理过程

在推理期间，模型通过迭代过程生成图像：

  1. Embedding：使用 CLIP-L/14 文本编码器对输入提示进行嵌入。
  2. Iterative Prediction：过程从随机遮罩的标记开始。U-ViT 模型预测遮罩标记，仅保留最有信心的预测的百分比（由遮罩调度和步骤数 N 决定）。其余标记重新遮罩并再次传递给 U-ViT 模型。
  3. Decoding：最终输出通过 VQGAN 解码器生成最终图像。

与 MUSE 的关键区别

aMUSEd 与原始 MUSE 模型在三个关键方面不同：

  • Single-Stage Prediction：它不使用两阶段方法来预测最终的遮罩补丁。
  • Text Conditioning：它使用 CLIP L/14 而非 T5 来计算文本嵌入。
  • Micro-conditioning：遵循 SDXL，它加入了图像尺寸和裁剪等额外条件。

能力与性能

aMUSEd 旨在实现高效与快速。模型约有 800M 参数（包括文本编码器和 VQ-GAN），性能表现出色。

推理速度与延迟

在 A100 GPU 上进行的基准测试显示，aMUSEd 的推理延迟显著低于其他模型。这种效率使其成为设备端应用的合适候选。

零样本图像修复

由于其预训练目标，aMUSEd 能够进行零样本图像修复，这一特性在诸如 SDXL 的模型中并不存在。

风格迁移

继承 MUSE 的能力，aMUSEd 展示了使用单张图像进行风格迁移的能力，为个性化和特定风格的图像生成提供了潜力。

微调与可访问性

aMUSEd 在 OpenRAIL 许可证下发布，商业友好。Hugging Face 提供了用于在自定义数据集上微调的训练脚本。

微调的硬件需求

  • Standard Fine-Tuning：使用 8 位 Adam 优化器和 float16 精度，微调所需的 GPU 显存略低于 11GB。
  • LoRA：使用低秩适配（LoRA），显存需求进一步降低至 7GB。

局限性

虽然 aMUSEd 在原始图像质量方面并非最先进，但其发布旨在鼓励社区探索 MIM 在图像生成中的潜力。作者强调了 MIM 框架的多项优势，包括推理效率、适用于设备端的更小模型体积，以及无需昂贵微调即可进行任务迁移的能力。

Sources