MRT: 一个 20B masked region transformer 把分层图生成的三个任务塞进同一个模型

MRT 四个任务的总览图
Fig. 1 — MRT 用同一个模型支持四种使用模式: ① Text-to-Layer Generation (从文本生成多层 RGBA 设计), ② Image-to-Layer Decomposition (把平面图分解成多个 RGBA 层), ③ Layer-to-Layer Addition (给现有设计添加新元素 + 自动协调位置/风格), ④ Layer-to-Layer Restylization (用户传一张图,自动改写成与现有设计风格一致的层)。

1. 出发点 (Motivation)

文生图模型(Qwen-Image / SD3 / FLUX 这些)主要生成的是flat raster image——一张已经合成好的位图。问题:用户在 Canva/Figma/Photoshop 里实际需要的是分层 PSD——每个元素一层 RGBA,可拖动、可换色、可重排。"先生成位图,再用分割模型抠出来当成层"这条路两个硬伤:

  1. 分割模型只能给可见像素,遮挡部分(被其他对象挡住的区域)没法补全 → 层一旦移动就露馅
  2. 超出画布边界的元素(比如 banner 一半在框外),裁剪后再也拼不回完整对象

所以这个领域近一两年有几条思路:

MRT 想一次性解决:

一句话主张:用大模型 + 大数据 + 单一架构 + adaptive masking,把分层生成做到 Photoshop 级别。

2. 方法 (Method)

2.0 具体输入输出 + Layout 如何进模型 (这节最重要)

抽象描述读完容易问"具体什么张量进、什么张量出、layout 怎么喂"。一次说清:

Layout 进模型的方式 = 2D RoPE 位置编码,没有别的。Paper §3.2 只在一句话里带过("we copy the RoPE positional encoding from the corresponding original layer token..."),但这是核心机制:

举例:画布 1024×1024,Layer i 的 bbox = (200, 300, 600, 700)。该 layer 经 WAN-VAE 编码后得到一个 50×50 的 latent grid(下采样 8 倍)。这 50×50 个 token 的 RoPE 2D 位置就分别是 (25, 37.5), (26, 37.5), ..., (74, 87.5) —— 对应全画布坐标系, 不是 layer 内部的 (0,0)..(50,50)。Transformer 在 full-attention 中自然就能看到 "这个 layer token 跟 composed image 哪个像素位置对齐"。

以 Image-to-Layers 为例走一遍张量

输入 (用户给):

编码阶段:

z_composed = WAN-VAE(I)                 # shape [16, 128, 128] (latent dim 16, 下采样 8×)
                                        # 16384 token, RoPE 位置 = 全画布 (i, j)

z_bg = noise                            # shape [16, 128, 128], 全画布 latent (待去噪)
z_i^fg = noise (i=1..K)                 # shape [16, H_i/8, W_i/8], bbox 大小 latent (待去噪)
                                        # 这些 token 的 RoPE 位置 = bbox 在全画布上的坐标

拼成一长串 token 序列 (一维序列, 不是空间张量):

[ z_composed (mask=CLEAN_CONDITION),    # 16384 token; 不加噪、不算 loss
  z_bg       (mask=NOISY_TARGET),       # 16384 token; 加噪 → 预测 velocity → 算 loss
  z_1^fg     (mask=NOISY_TARGET),       # ~3-5k token (按 bbox 大小)
  ...
  z_K^fg     (mask=NOISY_TARGET) ]

总 token 数 ≈ 16384 (composed) + 16384 (bg) + \(\sum_i\) (bbox\(_i\) 面积/64)。关键: K 层加起来的 fg token 数大致等于一张全画布(K 个 bbox 拼起来约覆盖画布一次),所以总 token ≈ 3× 全画布。Qwen-Image-Layered 用全画布表示每层,总 token = (K+2)× 全画布——20 层时差 ~10× token,这是 Fig. 19 100× 推理加速的根因。

Transformer:20B Qwen-Image,60 层,全 token full attention(都互相看)。mask 只控制"这个 token 要不要加噪 / 要不要算 loss",不限制 attention 谁能看谁。

采样:Flow matching Euler 50 步(或蒸馏 8 步)。每步:

v_pred = transformer(z_t, t, c_text, mask, RoPE_positions)
z_{t-Δt} = z_t - v_pred * Δt    # 仅更新 noisy 部分,clean condition token 不动

解码:

output_bg     = WAN-VAE-decode(z_bg_final)     # 全画布 RGBA
output_fg_i   = WAN-VAE-decode(z_i^fg_final)   # 每层 bbox 大小 RGBA,带 alpha 通道

输出 (给用户):

四个任务对照表

输入张量结构、输出张量结构、transformer 架构、loss 完全一致。只有 mask 模式不同:

Task composed bg fg layers 文本 \(c\) 用户额外输入
Text→Layers NOISY (生成) NOISY 全部 NOISY 全局 prompt 无
Image→Layers CLEAN (条件) NOISY 全部 NOISY (可选) 输入图 + layout
Layer Addition CLEAN CLEAN \(i \notin A\) CLEAN,\(i \in A\) NOISY <layer>$c_A$</layer> 现有 design + bbox
Layer Restylization CLEAN CLEAN non-target CLEAN; target NOISY; + append \(z_i^{\text{cond}}\) as CLEAN "Harmonize these layers" 现有 design + 用户图

Restylization 的小 trick: 用户上传的图也被 WAN 编码成额外 token,append 到 sequence 末尾,RoPE 位置复制 target layer 的 RoPE —— 让模型知道"这张用户图要替换 target layer 在那个位置的样子"。这就是 paper §3.2 末段说的 "copy RoPE positional encoding" 的实际作用。

2.1 数据表示: regional latents + overflow canvas

把一个多层设计表示为四类内容:

\[\{I_{\text{canvas}}, I_{\text{bg}}, \{I^{\text{fg}}_i\}_{i=1}^K\}\]

—— 翻译: 一个完整设计 = 1 张画布合成图(用于 layer coherence 监督) + 1 张半透明背景 + K 张前景 RGBA 层(可以超出画布)。每张都用 WAN-2.1-VAE 独立编码成 latent。

关键: token 数是按 region 算的。如果第 i 层的边界框是 100×200(占画布的 10% 面积),则该层的 token 数 ≈ 全画布 token 数 × 10%。这跟 Qwen-Image-Layered "每层用全画布 token"对比,K 层时 token 总数差 K 倍。这是论文 Fig. 19 显示 100× 推理速度的根本原因。

Overflow canvas layer:训练数据里 60% 的层会有超出画布的部分(比如 banner 拉到画框外)。MRT 引入一个全尺寸 canvas layer 专门承载这些"完整层 + 超出部分",其余 layer 仍按各自 region 编码。这样每个 fg layer 都能完整保留,后续编辑时拖动不丢内容。

2.2 Masked Region Transformer: 用 mask 切换任务

这是 paper 最核心的设计。三个任务(text→layers / image→layers / layers→layers)共用同一个 20B transformer,通过控制 "哪些 token 被 mask 当作干净条件、哪些 token 加噪做 diffusion 目标" 来切换任务。

Masked Region Transformer 框架图
Fig. 4 — 同一个 transformer 通过 mask 模式切换三类任务。 Text-to-Layers: 所有 layer latents 都加噪(纯生成); Image-to-Layers: 输入图作为z_mask(干净不加噪),layer latents 加噪去噪(分解); Layers-to-Layers Addition: 已有 layer 作为z_mask,新增 layer 加噪去噪(条件生成); Layers-to-Layers Restylization: 已有 layer + 用户上传的图作为z_mask,目标 layer 加噪去噪(风格迁移)。所有 mask token 和 noise token 在 transformer 内做 full attention。

记 \(z_0 = [z_{\text{composed}}; z_{\text{bg}}; z_1^{\text{fg}}; \ldots; z_K^{\text{fg}}]\) 为所有 latent token 的拼接,\(\epsilon \sim \mathcal{N}(0, I)\) 为噪声。Flow matching 标准的线性插值前向:

\[z_t = (1 - t) z_0 + t \epsilon\]

—— 翻译: 时刻 t (∈ [0, 1]) 时,latent 是干净的 $z_0$ 和纯噪声 $\epsilon$ 的线性插值。这跟 SD3 / FLUX / flow-opd-2026 / diffusion-opd-2026 全是同一套 rectified flow 框架。

Adaptive masking 的关键: 在拼接 \(z_0\) 之外,还有一组 \(z_{\text{mask}}\) 是"干净条件 token,不加噪、参与 attention 但不参与 loss"。每个任务定义 \(z_{\text{mask}}\) 不同:

任务 \(z_{\text{mask}}\)(干净条件) \(z_0\)(加噪去噪目标)
Text-to-Layers \(\emptyset\) (空) composed + bg + 全部 K 个 fg
Image-to-Layers composed(输入图) bg + 全部 K 个 fg
Layer Addition composed + bg + non-target fg target fg subset \(A \subseteq \{1...K\}\)
Layer Restylization composed (含 non-target) + bg + non-target fg + 用户图 \(z^{\text{cond}}_i\) target fg subset \(I\)

训练目标始终是同一个 flow matching MSE(论文 Eq. 2):

\[\mathcal{L}_{\text{flow}} = \mathbb{E}_{z_0, \epsilon \sim \mathcal{N}(0, I), t}\big[\|v_t - f_\theta(z_t, t, c)\|^2\big]\]

—— 翻译: 模型 $f_\theta$ 给定时刻 $t$ 的带噪 latent $z_t$、时刻 $t$ 本身、和文本条件 $c$,预测速度 $v_t = z_0 - \epsilon$(从干净指向噪声的方向)。loss 只在 $z_0$ 对应的 non-masked token 上计算。

为什么这设计很巧妙(我的理解):attention 是 set 操作,加什么 token 不加什么 token,结构是一样的。所以三个任务只需要改"input token 的组装方式"和"哪些 token 不参与 loss",权重一份就够。这本质上就是 LLM 里 prefix-LM / FIM (fill-in-the-middle) 那一套思路搬到 diffusion。

2.3 一些具体的实现选择

Layer Restylization 里的小 trick (paper §3.2 倒数第二段): 用户上传一张图想让它"长得跟现有设计风格一致",MRT 在 transformer 输入里把这张图编码后作为额外 condition token 加入,RoPE 位置编码从对应 target layer 复制过来——这样 condition token 和 target token 共享同一个空间位置先验,让模型知道"目标位置应该跟这个 condition 的风格对齐"。比单纯把 condition 放在 sequence 任意位置要稳。

Layer grouping augmentation (§3.2 末尾): 训练时随机把相邻或重叠的 layer 合并成一个,扩充 layout 多样性。原因:测试时用户用 layout detector 抠出来的边界框跟训练时的"ground-truth 干净边界框"差别可能很大,需要让模型学会处理"模糊/粗的边界框"。Table 5 ablation 显示这招在 I2L 任务上有持续提升。

Text condition 是 nice-to-have 不是 must (Table 4): I2L 任务给文本提示能涨 0.4 PSNR,但去掉也能跑——这跟"用 ART 那种纯文本驱动"的思路不一样。MRT 更像是"我看着这张图,加上画框告诉我每层在哪,我就能解出来"。

2.4 加速: DMD2 蒸馏到 8 步

50 步基线模型经过 Improved DMD (DMD2 [60, 61]) 蒸馏到 8 步,FID 仅从 16.02 升到 18.58(Table 6),换来 ~6× 推理加速。DMD 优化的是 KL 散度:

\[\mathcal{L}_{\text{DMD}} = \mathbb{E}_{z_0, t}\left[\mathrm{KL}\big(f_{\theta_T}(z_{t-1}|z_t) \,\Vert\, f_{\theta_S}(z_{t-1}|z_t)\big)\right]\]

—— 翻译: 让 student 模型 $f_{\theta_S}$ 在每个 timestep 上的反向 transition 分布近似 teacher $f_{\theta_T}$。这是当前 diffusion 蒸馏的标准做法,跟 SDXL-Turbo / Hyper-SD 一脉相承。

外加 CacheDiT(跨 timestep 复用注意力 KV)和多卡序列并行(把长 token 序列 shard 到 4×H100),最终 1K 分辨率 20 层在 4×H100 上 ~3s、单卡 H100 ~6s。

3. 结论 (Key Findings)

3.1 vs 同期 Qwen-Image-Layered (Table 1, Fig 19, Fig 20)

Fig 19: 效率对比
Fig. 19 — (a) Latency vs 层数: Qwen-Image-Layered 随层数线性增加(20 层 ≈ 540s),MRT 几乎恒定(~5s)。20 层时108.5× 加速。 (b) Latency vs token 数: MRT 在 H200/B200 上都呈线性 scaling,token 多时仍可控。 (c) Peak memory: MRT 内存 ~55-60 GB(随层数缓慢上升),Qwen-Image-Layered 从 78 GB 涨到 115 GB。MRT 在 20 层时省 23.6× 显存。

PSNR/SSIM 数字(Table 1, 我重排了一下):

任务 Layers MRT (Ours) Qwen-Image-Layered
Merged PSNR ↑ [4, 8) 27.34 25.81
[8, 16) 25.91 23.06
[16, 32) 25.72 22.18
Merged SSIM ↑ [4, 8) 0.903 0.871
[8, 16) 0.876 0.832
[16, 32) 0.849 0.807

关键观察: 随层数增加,Qwen-Image-Layered 的质量下降很明显(PSNR 从 25.81 → 22.18),MRT 下降更慢(27.34 → 25.72)。这跟前面说的"regional latents 让 token 数线性 ≈ 单层而非 K×"是一致的——token 拥挤是质量下降的根因。

Fig 20: Image-to-Layers 定性对比
Fig. 20 — 三个测试样本,从左到右: 输入图 → MRT (Ours) → Lovart → RoboNeo → LayerD → Qwen-Image-Layered。 Lovart 抠不干净(背景残留),RoboNeo 有 artifact(边界毛刺), LayerD / Qwen-Image-Layered 倾向于把多个元素合并成一层(粒度太粗,失去编辑灵活性)。 MRT 的层数最多、边界最干净。

User study(Fig 8 那张未抠的 user study chart):MRT 对 Qwen-Image-Layered 的胜率: Quality 79.5%, Integrity 68.9%, Granularity 82.6%。

3.2 Text-to-Layers + Layers-to-Layers (Fig 6, Fig 22)

T2L user study: MRT 在 Overall / Aesthetics / Typography / Layout 四个维度对 ART [38] 的胜率分别是 63.0% / 69.0% / 61.0% / 54.3%(剩下是 tie 或 ART win,后者占比都不到 30%)。

Fig 22: Layers-to-Layers 定性对比
Fig. 22 — 两个编辑场景: Layer Addition (前两行) 给现有设计加两个新元素(banner + 文本); Layer Restylization (后两行) 用户上传一张图(model 截图 / hooded figure),MRT 把它改写成跟现有 disco-night 设计风格一致的层。 MRT 并行生成所有目标层,GPT-Image-1 迭代逐层,前者更协调。

3.3 关键消融

模型规模 + 数据规模都重要 (§4.4): FLUX.1-dev (13B) → Qwen-Image (20B) 在同 0.5M 数据上 FID 从 17.79 降到 16.15。Qwen-Image 上数据从 0.5M 扩到 10M FID 进一步降到 15.63。两者缺一不可。

Overflow 数据不损害性能 (Table 2): 加入 overflow 训练数据后,merged image 的 PSNR 从 21.81 涨到 22.75,SSIM 从 0.854 涨到 0.871。"支持完整层"和"主体合成质量"不是 trade-off。

多任务联合训练只略微伤性能 (Table 3): T2L 单任务 FID 16.15,T2L+I2L 混训 15.68(反而更好),T2L+I2L+L2L 三任务 17.06(略掉)。论文归因于 L2L 数据本身质量问题。

T2L 数据 60% 是 overflow layer,这点很反直觉但合理——专业设计师确实经常画"半个 banner 在框外、扩展裁切空间"。

3.4 失败案例 (Fig 18, paper §4.3.3 末尾)

Fig 24: 自然图像泛化
Fig. 24 — 把只在设计图上训练的 MRT 应用到自然场景(汉堡照片、蜜蜂照片、香蕉摊照片),它仍能产生合理的层分解(虽然没在自然图上专门训练过)。这暗示 MRT 学到的不只是"识别设计稿元素",而是更普适的"识别可分离视觉对象 + 补全遮挡"。

4. 实现细节 (Implementation Notes)

⚠️ 代码尚未发布: 论文 2026/05/26 上传,本博文撰写时(2026/06/02)Canva GitHub org 没有 MRT repo,Hugging Face 上也没有 MRT 权重(只有 Qwen/Qwen-Image-Layered 这个同期但不同方法的 baseline)。所以本节没有 repo/file.py:Lnn 代码引用,以下都是从论文 + 同类 baseline 推断的工程细节。

5. 批判性总结 (Critical Assessment)

5.1 优点

5.2 不足 / 疑点

5.3 适用 vs 不适用

5.4 进一步阅读

讨论 / Comments

评论托管在本仓库的 GitHub Discussions, 需 GitHub 账号。