GenMask：通过直接 Mask 生成将 DiT 适配到分割任务

用扩散 Transformer 直接生成分割 mask，一步推理，达到 SOTA

前置知识

Rectified Flow（修正流）

Rectified Flow 是一种生成式建模框架，学习从简单噪声到复杂数据的连续变换路径。与传统扩散模型使用复杂的随机路径不同，Rectified Flow 使用最简单的直线路径连接数据和噪声：给定图像 $x_0$ 和高斯噪声 $\epsilon \sim \mathcal{N}(0, I)$，定义线性路径 $x_t = t\epsilon + (1-t)x_0$，其中 $t \in [0,1]$。模型学习预测方向向量 $v = x_0 - \epsilon$，损失函数为 $\mathcal{L}(\theta) = \mathbb{E}_{x_0, \epsilon, t} \|(x_0 - \epsilon) - v_\theta(x_t, t)\|^2$。这种直线路径设计使生成过程更加高效稳定。

GenMask 基于 Rectified Flow 的 Flow Matching 框架构建，理解这一基础是理解整个方法的前提

DiT（Diffusion Transformer）

DiT 是将 Transformer 架构引入扩散模型的创新设计，替代了传统 U-Net 去噪器。DiT 接受带噪图像作为输入，通过交叉注意力机制整合条件信息（如文本描述），输出去噪后的图像。同时使用 AdaLN（自适应层归一化）操作将时间嵌入注入去噪过程。本文使用的 WAN-2.1 DiT 拥有 1.3B 参数，是开源的大规模视频生成模型。

GenMask 的核心创新是在不改变 DiT 架构的前提下，让它同时完成图像生成和分割任务

VAE（变分自编码器）潜在空间

VAE 将高维图像压缩到低维潜在空间（latent space），例如将 512×512 的 RGB 图像压缩为 64×64×16 的潜在表示。潜在扩散模型（LDM）在这个压缩空间中进行扩散过程，大幅提升生成效率。本文发现二值分割 mask 的 VAE 潜在表示具有独特的性质：它们是线性可分的、对噪声高度鲁棒的，这与自然图像的潜在表示完全不同。

发现并解决 VAE 潜在空间中 mask 与自然图像的分布差异，是本文的关键洞察

Timesteps Sampling（时间步采样）

在扩散模型训练中，时间步 $t$ 控制加噪程度：$t=0$ 表示干净数据，$t=1$ 表示纯噪声。不同的时间步采样策略会显著影响模型学习效果。Stable Diffusion 3 提出的 logit-normal 采样策略强调中间噪声水平，因为极端噪声（太强或太弱）提供的学习信号有限。本文针对分割任务提出了独特的极端长尾采样策略。

不同的采样策略是 GenMask 能够统一生成和分割任务的关键技术手段

研究动机

现有基于扩散模型的分割方法存在两个根本性问题。第一，表示层面的错位：扩散模型预训练的目标是建模 VAE 特征的低级分布（如纹理、颜色），而分割任务需要的是紧凑的语义级标签预测。这种表示层面的不匹配严重阻碍了有效的下游适配。第二，间接特征提取的复杂性：现有方法依赖精心设计的间接流水线从扩散模型中提取特征，常见方式包括扩散反演（diffusion inversion）和激活聚合（activation aggregation）。这些中间操作不仅使工作流程变得复杂，还限制了适配性能。例如，DiFFt 方法需要多步反演和复杂的特征聚合，RIS-CLIP 需要精心设计的特征提取模块。

本文的目标是本文的核心目标是将分割任务直接建模为生成任务，让 DiT 能够在原始架构和训练目标下同时完成图像生成和 mask 生成。具体而言，GenMask 要实现三个目标：（1）架构忠于原始 DiT，不需要结构改变或额外操作；（2）最大化对齐生成训练目标，消除隐式适配带来的优化差距；（3）无缝整合生成数据，让图像生成数据能够提升分割性能。最终目标是在 RefCOCO 系列基准上达到 SOTA 性能。

与已有工作不同的是，本文的独特切入角度在于发现了二值分割 mask 的 VAE 潜在表示与自然图像的根本差异。作者通过 PCA 分解实验发现，mask 的 VAE 表示可以仅用一个主成分就几乎完美重建，证明其线性可分性。更重要的是，只有在极高噪声强度下，这种线性可分性才会崩溃。基于这一洞察，本文提出针对不同任务使用完全不同的时间步采样策略：生成任务使用相对均匀的分布，分割任务使用极端长尾分布集中在高噪声区域。这种简单但关键的设计让单一模型能够同时学习两个任务。

核心方法

GenMask 的核心思路可以用一个类比来理解：想象一个画家既能画彩色风景画，也能画黑白轮廓画。传统方法是让画家先学风景画，然后通过复杂的步骤把风景画的技巧迁移到轮廓画。GenMask 则是直接让画家在同一个画室里同时练习两种绘画，但使用完全不同的颜料调配方式（采样策略）来处理两种画作的不同特性。技术路线如下：首先基于预训练的 WAN-2.1 DiT 架构，使用 Qwen2.5-VL-7B 作为视觉语言编码器来同时编码图像和文本指令。对于生成任务，遵循标准的扩散训练流程，使用 logit-normal 采样强调中间噪声水平。对于分割任务，使用极端长尾采样策略集中在高噪声区域，并在 DiT 输入中注入输入图像的 VAE 潜在表示作为低层信息捷径。

GenMask 的核心创新点在于发现了二值分割 mask 与自然图像在 VAE 潜在空间中的本质差异，并通过差异化的时间步采样策略来解决这一问题。具体而言，作者发现二值 mask 的 VAE 表示具有三个关键特性：（1）分布尖锐，占据潜在空间的狭窄区域；（2）对噪声高度鲁棒，即使加入极大噪声，mask 的全局位置和形状仍然保持完整；（3）线性可分，通过简单的 PCA 分解就能几乎完美重建。基于这些发现，作者提出对分割任务使用极端长尾采样策略，概率密度函数为 $p(t) = \frac{2a^2 t}{(t^2 + a^2)^2}$，其中 $a=0.05$ 使得分布极度集中在高噪声区域。这种策略让模型能够有效学习 mask 的去噪过程，同时避免了与图像生成任务的冲突。

方法步骤详情

GenMask 的完整方法包含以下步骤：（1）指令编码：使用 Qwen2.5-VL-7B 同时编码视觉和文本指令。对于分割任务，输入格式为 '[Image]. Please segment the {target} in the image.'，提取最后一层隐藏状态作为条件输入。（2）时间步采样：生成任务使用 logit-normal 分布 $\pi(t) = \frac{1}{2\pi t(1-t)} \exp\left(-\frac{1}{2}\left(\log\frac{t}{1-t}\right)^2\right)$，强调中间噪声；分割任务使用极端长尾分布，90% 的训练样本集中在 $t > 0.85$ 的高噪声区域。（3）VAE 捷径注入：将输入图像的 VAE 潜在表示与随机采样噪声拼接，作为 DiT 的输入，时间嵌入设为零表示无噪声。（4）训练目标：使用 MSE 损失 $\mathcal{L} = \mathbb{E}\|(x_0 - \epsilon) - v_\theta(x_t, t)\|^2$ 在 VAE 空间中监督，生成和分割任务以 1:1 比例混合训练。（5）推理：对于分割任务，固定 $t=1$，仅需一次前向传播：$x_{mask} = \epsilon + v(\epsilon, 1)$，然后通过 VAE 解码器得到最终 mask。

技术新颖性

GenMask 的技术新颖性体现在三个方面。第一，理论发现的新颖性：首次系统揭示了二值分割 mask 的 VAE 潜在表示具有线性可分性和噪声鲁棒性，这与自然图像的潜在表示有本质区别。第二，采样策略的新颖性：针对不同任务使用完全不同的时间步采样策略，分割任务的采样分布峰值概率（13%）是生成任务（1.6%）的 8 倍以上，这种极端差异化设计是前所未有的。第三，架构设计的新颖性：在不改变原始 DiT 架构的前提下，仅通过注入 VAE 捷径和调整采样策略，就实现了生成和分割的统一。与现有方法相比，GenMask 不需要额外的分割解码器、特征提取模块或复杂的反演过程，完全依靠生成式训练目标就能产生精确的分割结果。

实验结果

GenMask 在多个基准测试中达到了 SOTA 性能。在 Referring Image Segmentation（指代分割）任务上，GenMask 在 RefCOCO 验证集上取得了 mIoU 83.7%、oIoU 80.7% 的成绩，相比 GLaMM（mIoU 83.2%、oIoU 76.9%）提升了 0.5% 和 3.8%。在 RefCOCO+ 测试集 B 上，GenMask 取得了 mIoU 73.1% 的成绩，超越 C3VG（72.4%）0.7 个百分点。在 RefCOCO-g 测试集上，GenMask 以 mIoU 78.2%、oIoU 76.5% 大幅领先所有基线方法。在 ReasonSeg（推理分割）任务上，GenMask 在测试集上取得了 gIoU 52.3%、cIoU 45.8% 的成绩，相比 LISA 的 gIoU 47.3% 提升了 5 个百分点。消融实验表明：（1）采样策略至关重要，$a=0.05$ 比 $a=0.5$ 在 RefCOCO 上高出 16.2% mIoU；（2）MSE 损失优于 BCE，分别在 RefCOCO 上达到 82.2% 和 78.1% mIoU；（3）VAE 捷径对分割至关重要，移除后性能下降 8.1% mIoU；（4）加入生成数据以 1:1 比例混合训练能带来 1.2% 的 mIoU 提升。

**Table 1**: Referring Segmentation 基准性能对比

查看结构化数据

任务	指标	本文	基线	提升
RefCOCO val	mIoU	83.7%	GLaMM 83.2%	+0.5%
RefCOCO val	oIoU	80.7%	GLaMM 76.9%	+3.8%
RefCOCO test A	mIoU	80.0%	C3VG 79.6%	+0.4%
RefCOCO test B	mIoU	73.1%	C3VG 72.4%	+0.7%
RefCOCO-g val	mIoU	77.2%	C3VG 76.3%	+0.9%
RefCOCO-g test	mIoU	78.2%	C3VG 77.1%	+1.1%
ReasonSeg test	gIoU	52.3%	LISA 47.3%	+5.0%
ReasonSeg test	cIoU	45.8%	LISA 34.1%	+11.7%

局限与改进

尽管 GenMask 取得了优异的性能，但仍存在一些局限性。首先，模型规模较大，DiT 有 1.3B 参数，VLM 有 7B 参数，总计约 8.3B 参数，这对部署和推理效率提出了挑战。其次，训练数据需求较高，需要同时使用语义分割数据（COCO-stuff、ADE20K、PASCAL）、指代分割数据（RefCOCO 系列）和生成数据（DiffusionDB、BLIP3o 等），数据准备较为复杂。第三，虽然一步推理已经足够，但推理时仍需要 VAE 解码器，增加了计算开销。第四，论文未报告推理速度和显存占用的具体数据，难以评估实际部署的可行性。第五，消融实验主要在 RefCOCO 系列上进行，未在 ReasonSeg 上进行系统的消融分析。

独立分析的弱点

从独立分析的角度来看，GenMask 存在以下弱点和改进方向。第一，模型效率问题：当前架构包含 8.3B 参数，可以考虑使用知识蒸馏或模型压缩技术减小规模，或者探索更轻量级的 VLM 替代方案。第二，采样策略的泛化性：极端长尾采样策略虽然有效，但其超参数 $a$ 的选择缺乏理论指导，目前是通过实验确定的。未来可以研究自适应采样策略，根据训练进度动态调整分布。第三，VAE 捷径的设计：当前方法直接拼接 VAE 潜在表示，但未考虑更复杂的融合方式，如交叉注意力或门控机制，可能进一步提升性能。第四，多类别分割：当前方法将语义分割转换为二值 mask，对于包含多个类别的场景可能效率较低，未来可以探索直接生成多类别 mask 的方法。第五，泛化能力验证：论文主要在标准基准上评估，未测试在域外数据（如医学图像、遥感图像）上的泛化能力。

未来方向

基于本文的成果，可以延伸出多个有前景的研究方向。首先，作者提到的扩展方向包括将方法应用到更大的 DiT 骨干网络（如 WAN-2.1 的更大版本）和更广泛的视觉任务（如医学图像分割、音视频分割）。其次，可以探索将 GenMask 的统一生成范式扩展到其他视觉感知任务，如实例分割、全景分割、目标检测等，这些任务同样可以从生成式训练中受益。第三，可以研究更高效的时间步采样策略，例如基于课程学习的自适应采样，让模型在训练初期学习简单样本，逐渐过渡到困难样本。第四，可以探索将 GenMask 与图像编辑模型结合，实现交互式分割和编辑的统一框架。第五，可以研究 GenMask 的零样本或少样本泛化能力，评估其在新类别上的表现。

复现评估

从复现评估的角度来看，GenMask 具有较好的可复现性。首先，代码和模型的开源情况良好：论文基于开源的 WAN-2.1 DiT（1.3B 参数）和 Qwen2.5-VL-7B 构建，这些模型都可以公开获取。其次，训练数据方面，论文使用了 COCO-stuff、ADE20K、PASCAL、RefCOCO 系列等标准数据集，以及 DiffusionDB、BLIP3o 等开源生成数据集，数据获取难度较低。第三，算力需求方面，训练使用全局 batch size 1024，约 8000 次迭代收敛，估计需要多卡 GPU 集群（如 8×A100）训练数天。第四，实现细节方面，论文提供了学习率（5e-5 到 1e-5 余弦衰减）、采样策略公式、训练数据比例等关键信息。第五，复现难度中等，主要挑战在于大规模数据的准备和多卡训练的工程实现。

GenMask：通过直接 Mask 生成将 DiT 适配到分割任务 GenMask: Adapting DiT for Segmentation via Direct Mask