文生图 · 扩散模型 · 高效 AI

Moebius:0.22B 图像修复模型打平 11.9B FLUX

Moebius 是 0.22B 扩散修复模型,追平 11.9B 的 FLUX.1-Fill-Dev。CelebA-HQ FID 5.39 对 10.13,参数不到 2%,推理约省 15 倍。

Moebius:0.22B 图像修复模型打平 11.9B FLUX

快速答案

Moebius 是 0.22B 参数的扩散图像修复模型,在标准基准上追平甚至超过 11.9B 的 FLUX.1-Fill-Dev,代价却低得多。人像基准 CelebA-HQ(512)上它拿到 FID 5.39,FLUX 是 10.13;FFHQ(256)上 8.15 对 11.19。自然场景(Places2)更胶着:更难的 Test 子集 FLUX 领先(8.02 对 9.48),小孔洞子集 Moebius 反超(0.92 对 0.94)。而它用的参数不到 FLUX 的 2%(0.22B 对 11.9B),单步延迟 26.01 ms 对 FLUX 的 161.01,采样 20 步对 FLUX 的 50 步,叠起来总推理时间约省 15 倍。这个结果靠两块东西配合:重新设计的骨干模块,加上潜空间蒸馏配方。单用任何一块都不够。

难题:骨干一压缩,表征就塌

10B 量级的工业修复模型填得很好,但部署成本太高。直接的思路是训一个小的专用模型,而不是巨型通才。坑在于:参数砍狠了,模型表征上下文的能力被掐住,质量断崖式下跌。作者把这叫表征瓶颈,Table 2 给了实证:同一个紧凑骨干,不做蒸馏只能到 FID 33.21,加上他们的蒸馏配方就降到 24.73(Places2 同一检查点)。所以小架构是必要的,但不够,你还得把它教好。

Local-lambda Mix 模块怎么工作

骨干是潜扩散 U-Net,围绕一个重复单元重建,就是 Local-lambda Mix Interaction 模块。它有三部分。Local-lambda 把邻近的空间上下文压成一个固定大小的线性矩阵(lambda),这样混合上下文的开销不会像注意力那样随 token 数膨胀。Interactive-lambda 对全局语义先验做同样的事,把 Latent Categories Guidance 信号融进另一个紧凑 lambda。Mix-FFN 是前馈级。把上下文和全局先验都折进固定大小的矩阵,好处是保住了重注意力块能捕到的潜空间交互,却只付一个固定开销而非平方开销,参数和算力的节省大头就在这里。

把标准注意力换成 lambda 模块是模型变小的原因。标准全局注意力骨干(GLA-CA-FFN)用 526M 参数;完整 Local-lambda/Interactive-lambda 设计在 485M 上与之持平(未蒸馏 FID 33.21 对 32.75,基本打平),再加深度可分卷积和 Mix-FFN 压到出货的 226M、算力 154 GFLOPs。蒸馏下 485M 变体到 FID 24.73,226M 出货骨干到 26.43,最终 53% 的参数砍幅只换来约 1.7 的 FID 代价,是有意为之的取舍。

质量到底来自哪

让 0.22B 有竞争力的那个增量来自蒸馏,不是模块设计。教师是 PixelHacker(0.86B),Moebius 全程在潜空间对齐它的内部表征,跳过了像素空间蒸馏要做的昂贵解码。配方叠了四项损失,Table 5 的消融显示每一项都有用:只用粗粒度 KD 损失,FID 高达没法用的 74.20;加细粒度 KD 降到 36.17;加任务损失到 32.59;加感知损失落到最终的 26.43。最关键的对比是架构单独跑和架构加蒸馏:匹配设置下分别是 33.21 和 24.73。模块负责把模型做小,蒸馏负责把压出来的窟窿填上。

这和 DistilBERT 那种经典模型压缩是同一个杠杆,只是从 transformer 编码器换到了扩散骨干,多了一层在潜空间跨多个空间粒度对齐的花样。

一个容易看错的点

标题写”10B 级性能”,最省事的读法是 Moebius 全面碾压 FLUX。并没有。自然场景 Places2 Test 子集上,FLUX.1-Fill-Dev 的 FID 仍领先(8.02 对 9.48),256 子集也是(10.44 对 14.13)。Moebius 的明确胜场在人像(CelebA-HQ 和 FFHQ),以及小孔洞、大孔洞的自然子集,同时它紧贴自己的教师。准确的说法是”按子集不同,或追平或超过,且代价只有零头”。论文的用户研究支持这点:人类评分者在 10B 通才面前更偏好 Moebius,人脸上尤其明显,而它对自己的教师只是打平。把它当成把差距缩小的专才,不是通吃的赢家。

提速这块也值得点一句:约 15 倍的总时间优势不全是架构带来的。一部分来自更低的单步延迟(26.01 对 161.01 ms),一部分来自只跑 20 步而非 FLUX 默认的 50 步。两者都合理,但如果把步数拉平,这个倍数会缩水。

关键结果

  • CelebA-HQ(512)人像,FID: Moebius 5.39,FLUX.1-Fill-Dev 10.13,SD3.5 Large-Inpaint 11.80,0.22B 的模型在人脸上几乎把 FLUX 的 FID 砍半。除自己的教师 PixelHacker(4.75)外,胜过表中所有学术和工业基线。
  • FFHQ(256)人像,FID: Moebius 8.15,FLUX 11.19,SD3.5 109.42,人像上对两个工业通才都明显领先。
  • Places2 Test(自然),FID: FLUX 8.02 领先 Moebius 9.48,这是 10B 通才仍占优的子集,所以”打平”的说法依子集而定。
  • 成本: 0.22B 参数对 FLUX 的 11.9B(不到 2%),154 GFLOPs,单步 26.01 ms 对 161.01,20 步对 50 步,叠起来总推理时间约省 15 倍。
  • 蒸馏消融(Places2,18K 步检查点): 只粗粒度 KD FID 74.20,加细粒度 KD 36.17,加任务损失 32.59,加感知损失 26.43,四项都贡献,粗粒度单用没法用。
  • 架构对蒸馏(匹配设置): 同一紧凑骨干,不蒸馏 FID 33.21,蒸馏后 24.73,把增量定位到教师信号而非模块本身。
  • 分布外泛化: LVIS(自然)和 DeepFakeFace(人像)零样本上,Moebius 拿 FID 17.81 / 15.32,SD3.5 是 114.21 / 81.91,且贴近 FLUX(14.52 / 12.24),说明专才没有过拟合训练分布。

局限与存疑

Moebius 是专才,不是通才。它针对自然和人像场景的掩码修复训练,不做文本提示编辑,所以和 FLUX 的对比只有在两者都能做的修复任务上才公平。它的天花板是教师:失败案例图显示,在上下文极稀薄的极小背景区域,它可能丢失细节或生成不太合理的纹理,这正是 0.22B 容量预算吃紧的地方。自然场景 Places2 Test 的结果也证实它在难的通用场景上还没完全追平。15 倍提速混了单步延迟优势和更少步数,逐步对齐去比的话倍数会更小(但仍很大)。蒸馏阶段没有公开的训练成本数据,结果也限于两个域、固定分辨率。

常见问题

Moebius 真的在所有场景都打败 FLUX.1-Fill-Dev 吗?

没有。Moebius 在人像上完胜(CelebA-HQ FID 5.39 对 10.13,FFHQ 8.15 对 11.19),小孔洞和大孔洞的自然子集也赢,但更难的 Places2 Test 子集 FLUX.1-Fill-Dev 仍领先(8.02 对 9.48),256 子集也是。准确的总结是:Moebius 按基准不同或追平或超过 FLUX,而参数不到 2%,不是全面碾压。

Moebius 的质量是来自新模块还是蒸馏?

主要是蒸馏。Local-lambda Mix Interaction 模块负责把模型做小做便宜,但匹配消融里,同一紧凑骨干不蒸馏是 FID 33.21,蒸馏后 24.73。从 PixelHacker 教师做的四项损失潜空间蒸馏,才是填上压缩造成的表征瓶颈的那块。

Moebius 的 15 倍提速是架构更快还是步数更少?

都有。Moebius 单步 26.01 ms 对 FLUX 的 161.01(架构层面的单步优势),采样用 20 步而非 FLUX 默认的 50 步。单步延迟乘步数,得到约 15 倍的总时间优势。把步数拉平,倍数会下降,不过 Moebius 每步仍便宜得多。

Moebius 里的 Local-lambda Mix Interaction 模块是什么?

它是替代标准注意力的重复骨干单元。Local-lambda 把局部空间上下文压成固定大小的线性矩阵,Interactive-lambda 用同样方式压全局语义先验,Mix-FFN 处理前馈级。把上下文折进固定大小矩阵,既保住丰富的潜空间交互,又只付固定而非平方的开销,这是参数和算力节省的主要来源。

一句话:Moebius 用基于 lambda 的紧凑骨干,加上从 0.86B 教师做的四项损失潜空间蒸馏,在 0.22B 上达到 FLUX 级修复效果、总推理时间约省 15 倍,人像上胜出,难的自然场景略逊。更多轻量化模型可看 高效 AI 主题,以及与修复相邻的 ControlNet。阅读 arXiv 原文