从混合专家MoE到混合注意力MoA，Snap提出全新个性化图像生成框架

个性化图像生成（Personalized Image Generation）是AIGC领域中一项重要的子课题，其旨在根据少量的输入图像，将用户特定的人物或目标嵌入到基础文本图像模型（Text-to-Image Models）的生成过程中。现有的方法需要在保留基础模型能力和有效嵌入个性化目标两方面进行性能权衡，导致生成的个性化图像偏离用户意图且缺乏高保真度。

本文介绍一篇来自Snap的研究论文，本文受大型语言模型（LLMs）中使用的专家混合机制（Mixture-of-Experts，MoE）启发，引入了一种用于文本到图像扩散模型个性化的新架构，称为混合注意力（Mixture-of-Attention，MoA）。MoA将生成任务分配到两条注意力路径中：个性化分支和非个性化的先验分支，前者会对生成过程进行小幅度的干预来将用户指定的主体嵌入到先验分支生成的布局和上下文中。后者则通过固定注意力层的参数来保留基础模型的先验生成能力。

为了实现个性化和通用内容的融合，作者设计了一种高效的路由机制来管理两个分支中的像素分布，经过训练后，MoA可以生成高质量的、包含多个主体的个性化图像。至关重要的是，MoA实现了更加清晰的主体-上下文解耦控制能力。

论文题目： MoA: Mixture-of-Attention for Subject-Context Disentanglement in Personalized Image Generation 论文链接： https://arxiv.org/abs/2404.11565 项目主页： https://snap-research.github.io/mixture-of-attention/

一、引言

在生成式人工智能蓬勃发展的今天，各种AIGC模型和应用层出不群，已从根本上改变了数字媒体创作的格局。普通的AIGC基础模型已经可以实现通用内容的高质量生成且学习成本较低，仅需要用户通过输入简单的文本提示就可以开始创作。但目前的技术在个性化图像生成方面仍存在较大的性能瓶颈，其要求模型高度理解用户的意图，实现图像中不同主体的组合和交互。例如下图生成两个人同时潜水的照片，需要模型能够保持整体的布局不变，而个性化编辑参与的主体，同时保持生成图像的多样性和自然性。

现有的个性化模型通常会面临以下三个严峻的挑战：

1. 局限于特定布局:现有方法生成的图像中，个性化目标通常被限制在特定布局和姿态中，缺乏多样性和与上下文的互动。

2. 背景内容缺乏多样性:基于微调的个性化方法，例如DreamBooth[1]需要微调整体模型的参数来学习个性化信息，这会导致模型丢失原始能力，生成的图像背景缺乏多样性。

3. 无法处理遮挡:现有方法难以处理个性化目标遮挡的情况，生成结果往往存在不连贯的问题。

为了应对上述挑战，作者引入了一种全新的混合注意力（MoA）机制，MoA包含两个注意力分支:一个固定的先验分支和一个可训练的个性化分支，以及一个路由网络用于在两个分支之间进行软路由，这种方式可以能够在注入个性化目标的同时，最大限度地保留原始模型的生成能力。MoA还具有一个关键特性，即将个性化信息转换为文本提示中的视觉词元，这可以使得MoA能够无缝地处理多个个性化目标，充分利用预训练模型对文本的理解能力，从而生成具有丰富互动和上下文关系的图像。

二、本文方法

本文的整体框架如下图所示，MoA仍然遵循文本到图像生成扩散模型的基础架构，即使用Diffusion U-Net，由序列化的Transformer块组成，每个块包含自注意力层和交叉注意力层，每个注意力层都有自己的投影参数，而本文将其中的每一层都替换为了MoA层。

2.1 Mixture-of-Attention Layer

MoA的设计灵感来源于大模型中经常使用的混合专家MoE，**MoE层通常包含多个专家网络(Expert)和一个路由网络(Router)**，用于软路由每个输入到不同的专家网络中：

其中 表示element-wise乘积， 是路由网络的输出注意力分布，而本文提出的MoA主要对MoE层进行了两处关键改进：

1. 每个专家(Expert)都是一个标准的注意力层，包括投影参数和注意力计算过程。

2. MoA层只有两个专家：一个参数固定的先验专家和一个可训练的个性化专家。

具体来说，MoA层的计算公式可以表示如下：

其中每个MoA层都对应一个路由器网络，作者使用预训练模型中的注意力层来初始化 MoA 层中的两个专家，先验专家的参数保持冻结状态，而个性化专家则进行微调。此外，在MoA的自注意力层中，两个专家接收相同的输入，而在MoA的交叉注意力层中，两个专家接受不同的输入，其中先验专家接收与原始模型相同的文本嵌入，以保留对文本的原始理解能力。个性化专家接收一个多模态提示，用于将图像特征注入到文本提示中。

2.2 多模态提示

本文的多模态提示模块设计如下图所示，首先给定一个主体图像 ，通过预训练的图像编码器(如CLIP)提取视觉特征 。然后将 与对应文本token(如"man")的嵌入 进行拼接，得到多模态嵌入 。

随后对该多模态嵌入根据扩散时间步和U-Net层进行条件编码得到：

其中 是一个可学习的位置编码函数，最后将 输入一个MLP进行投影，作为个性化分支的注意力输入。

2.3 模型优化

MoA模型的优化可以分为多个方面，对于路由网络，其主要设计目标是将背景像素的处理工作分配给先验分支，而将前景目标(人物等)分配给个性化分支。因此，路由网络的损失函数可以表示为：

其中 是前景目标的mask， 是第 层中路由到先验分支的路由权重， 是需要惩罚的U-Net层的集合(通常不包括最浅和最深的几层)。直观来说，该损失函数最小化了背景像素被分配到个性化分支的比重。

此外，作者也提到，在不同的扩散时间步和U-Net层中，路由网络根据不同的分辨率对前景目标和背景进行分离，来实现高效的主体-上下文解耦效果。为了更加清晰的进行解释，作者对路由网络的行为进行了可视化分析，如上图所示。在早期扩散步骤和浅层，个性化分支更多关注的是人物主体的轮廓和大致形状，在中期扩散步骤和中层，个性化分支开始关注人物的面部特征，而在后期扩散步骤和深层，个性化分支进一步学习了人物的细节，如身体部位、衣着等。与此同时，先验分支一直在处理背景部分的生成。这种动态分工机制使得MoA能够基于文本提示生成丰富多样的背景，同时又能注入精细的人物细节。

通常情况下，对扩散模型进行训练和微调时需要使用完整的图像重构损失，而对于个性化生成模型，需要考虑背景和前景主体之间的混淆情况，因此本文使用分割mask来构建对前景的重构损失，MoA的整体优化函数可以表示如下：

其中 是针对前景目标像素的重构损失， 是一个 L1 损失，用于约束个性化分支学习对象的可视性映射(object visibility mapping)，有助于更好地注入个性化目标， 和 是平衡因子。

2.4 MoA与其他图像生成技术的兼容性

由于本文提出的MoA仅在注意力层级别上进行了改动，并未修改U-Net的结构，因此其天然地兼容现有的基于扩散模型的图像生成和编辑技术。例如与ControlNet[3]进行结合，ControlNet是一种高效的条件引导方法，通过向U-Net中的每个残差块注入条件信息(如姿态等)来控制生成。由于MoA仅修改了注意力层，因此U-Net中的残差连接保持不变，自然而然地支持了ControlNet的条件引导，如下图所示。

除了生成之外，MoA还兼容基于扩散反转（DDIM Inversion）的真实图像编辑，DDIM Inversion可以通过优化扩散过程中的noise，来将真实图像拟合到潜在空间，从而实现图像编辑的效果。由于MoA在推理时与原始扩散模型共享相同的潜在语义，因此可以自然的实现真实图像中主体人物的替换，如下图所示。

由于MoA架构中包含有将人物视觉特征嵌入到文本提示中的设计，这使得不同人物特征可以进行插值混合，从而实现"主体变形(Subject Morphing)"的效果。如下图所示，作者将一名相扑力士与一个生成的人物进行了特征插值，生成了一系列外形、体型、肤色等逐渐过渡的人物形象。

三、实验效果

本文的实验主要在FFHQ和CelebA数据集上进行，其中训练集使用FFHQ中的70,000张人脸图像，以及使用BLIP-2生成的文本描述和MaskedFormer生成的人像分割掩码。测试集选取了CelebA数据集中的15个人物图像。基础的文本到图像模式使用StableDiffusion v1.5，图像编码器使用OpenAI的CLIP ViT-Large版本，训练设备使用4块 NVIDIA H100 GPU。

3.1 定性实验分析

基于MoA中的先验分支和个性化分支设计，MoA可以在单主体和多主体情况下均实现良好的主体-背景解耦控制效果，下图为单主体解耦效果。

主体-背景解耦要求模型在相同的随机种子(即相同背景)下，注入不同的人物图像，MoA可以实现只对人物的前景区域进行局部变换，而完全保留背景内容不变，下图为多主体解耦效果。

此外，作者还对MoA处理遮挡等复杂情况的能力进行了评估，由于MoA学习到了主体的整体模型(包括全身)，因此具有较强的处理遮挡情况的能力，即使主体被物体(如餐具、服装)遮挡，MoA也能生成肤色、身体部位一致且高质量的结果。此外对于主体的互动行为，例如舞蹈、做饭等行为（如下图所示），MoA也可以实现自然逼真的生成效果。

作者将MoA的生成图像与其他现有方法（如FastComposer[4]）的生成效果进行对比时发现，MoA生成的图像不仅在背景多样性和上下文关联度方面表现优异，在处理多主体场景时也展现出了更强的能力，如下图所示。可以看出，FastComposer生成的人物相互独立，缺乏互动，而MoA则能生成自然的互动姿态。

四、总结

本文引入了一种混合注意力机制（MoA），这是一种用于个性化生成的全新架构。MoA可以对现有的基础文本到图像模型进行增强，解决了现有方法在保留原始模型能力和有效嵌入个性化目标之间的权衡问题，能够生成高质量、多样且富有互动的个性化图像内容。MoA设计的一大亮点是设计了两个分离的专家分支和一个路由网络，通过注意力分离来实现模型能力的模块化，这使得我们能够使用 MoA 创建各种有趣的图像生成应用，例如主题交换、主题变形、风格转移等。此外，由于MoA架构的设计非常简洁，仅在注意力层进行改动，这使得它天然兼容现有的扩散模型技术，如ControlNet和DDIM Inversion等，进一步拓宽了模型的应用场景。

参考资料

[1] Nataniel Ruiz, Yuanzhen Li, Varun Jampani, Yael Pritch, Michael Rubinstein, and Kfir Aberman. 2023a. Dreambooth: Fine tuning text-to-image diffusion models for subject-driven generation. In CVPR. 22500–22510.

[2] William Fedus, Barret Zoph, and Noam Shazeer. 2022. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. The Journal of Machine Learning Research 23, 1 (2022), 5232–5270.

[3] Lvmin Zhang, Anyi Rao, and Maneesh Agrawala. 2023b. Adding conditional control to text-to-image diffusion models. In Proceedings of the IEEE/CVF International Conference on Computer Vision. 3836–3847.

[4] Guangxuan Xiao, Tianwei Yin, William T Freeman, Frédo Durand, and Song Han. 2023. FastComposer: Tuning-Free Multi-Subject Image Generation with Localized Attention. arXiv preprint arXiv:2305.10431 (2023).

llustration From IconScout By Delesign Graphics

-The End-

扫码观看！

本周上新！

“AI技术流”原创投稿计划

TechBeat是由将门创投建立的AI学习社区（

www.techbeat.net

投稿内容

// 最新技术解读/系统性知识分享 //

// 前沿资讯解说/心得经历讲述 //

投稿须知

稿件需要为原创文章，并标明作者信息。

我们会选择部分在深度技术解析及科研心得方向，对用户启发更大的文章，做原创性内容奖励

投稿方式

发送邮件到

chenhongyuan@thejiangmen.com

或添加工作人员微信（chemn493）投稿，沟通投稿详情；还可以关注“将门创投”公众号，后台回复“投稿”二字，获得投稿说明。

关于我“门”

将门是一家以专注于数智核心科技领域的新型创投机构，也是北京市标杆型孵化器。 公司致力于通过连接技术与商业，发掘和培育具有全球影响力的科技创新企业，推动企业创新发展与产业升级。

将门成立于2015年底，创始团队由微软创投在中国的创始团队原班人马构建而成，曾为微软优选和深度孵化了126家创新的技术型创业公司。

如果您是技术领域的初创企业，不仅想获得投资，还希望获得一系列持续性、有价值的投后服务，欢迎发送或者推荐项目给我“门”:

bp@thejiangmen.com

点击右上角，把文章分享到朋友圈