跳转至

Rethinking Vector Field Learning for Generative Segmentation

会议: NeurIPS 2026
arXiv: 2603.19218
OpenReview: https://openreview.net/forum?id=R12MIytQsT
代码: https://github.com/PKU-YuanGroup/FlowSeg
领域: 语义分割
关键词: 生成式分割, 流匹配, 向量场重塑, 类别编码, 像素神经场

一句话总结

针对传统流匹配在生成式语义分割中目标质心附近梯度消失以及缺乏类间排斥力导致的轨迹穿越难题,本文提出向量场重塑策略,引入基于距离感知的排斥/吸引校正项与拟随机 Kronecker 类别编码,显著缩小了生成式分割与判别式专家模型的性能差距。

研究背景与动机

近年来,扩散模型与流匹配技术在图像与视频生成领域取得了突破性进展,促使学术界探索将连续生成框架统一拓展至稠密感知任务。然而,语义分割在本质上是将每个像素分配给离散的语义类别,而扩散与流模型是在有界的连续空间中学习概率路径的流动。现有生成式分割方案要么将预训练扩散模型退化为特征提取器外挂分割头,要么依赖轻量级扩散模块精炼基础模型输出的掩码,要么通过 VAE 潜在空间将掩码编码为连续伪彩图像开展图像到图像的转换。在处理具有高类别基数(如 ADE20K 包含 150 类、COCO-Stuff 包含 171 类)的场景解析任务时,生成式模型与强判别式模型(如 SegFormer、MaskFormer)之间仍存在极为显著的性能鸿沟。

深入探究该瓶颈的根源,问题并非出在模型容量不足,而是源于连续流匹配回归目标与离散分类感知任务之间的优化动力学失配。标准的流匹配(Flow Matching)采用均方误差(MSE)回归速度场,其优化梯度大小直接正比于当前预测状态与目标质心之间的欧氏距离。当预测点逐步逼近目标类别质心时,梯度信号迅速衰减并趋于消失,致使训练后期的模型丧失精确收敛到目标质心的驱动力,诱发边界模糊与欠拟合。

更为严重的是,MSE 目标仅提供朝向正样本质心的单向吸引力,完全缺失对竞争类别的显式排斥机制。在高类别基数的紧凑表示空间中,各类别质心密集相邻,缺乏排斥力使得生成轨迹在输运过程中极易意外穿越邻近竞争类别的邻域,产生严重的类别混淆。核心 idea:通过构建基于语义距离的势函数推导显式排斥与吸引校正项,重塑流匹配的目标向量场以维持质心附近的非零梯度并惩罚轨迹穿越,同时结合拟随机 Kronecker 类别编码与像素神经场解码,实现纯像素级端到端生成式分割。

方法详解

整体框架

FlowSeg 将语义分割形式化为从条件输入图像 \(x_0 \sim p_0(x_0 | I)\) 到目标离散掩码嵌入 \(x_1 \sim p_1(x_1)\) 的确定性整流流动(Rectified Flow)。为摆脱 VAE 压缩失真对密集像素精度的破坏,框架采用 Transformer 骨干网络配合局部像素神经场动态解码,生成像素级连续速度场。首先利用基于代数无关增量的拟随机 Kronecker 序列,将 \(N\) 个语义类别确定性且均匀地散布到有界三维立方体 \([-1, 1]^3\) 中作为目标质心 \(\mathcal{P} = \{\mu_1, \dots, \mu_N\}\)。在训练过程中,根据网络预测的终点状态计算各类别语义分配概率,构造包含正类吸引与负类排斥的判别势函数梯度,并以截断梯度(stop-gradient)的方式注入目标速度场,驱动整流流动学习具有判别性决策边界感知的输运轨迹。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入图像与加噪状态<br/>I 与 x_t"] --> B["像素神经场端到端解码<br/>动态参数生成 + DCT坐标查询"]
    B --> C["Kronecker 拟随机类别编码<br/>确定性均匀嵌入 3D 质心"]
    C --> D["势函数判别校正<br/>计算软分配概率与距离势能梯度"]
    D --> E["向量场重塑与截断监督<br/>v_tilde = v_gt - grad_Phi"]
    E --> F["确定性欧拉采样输出分割掩码"]

关键设计

1. 拟随机 Kronecker 类别编码:确定性且几何均匀的类别质心构造 针对高类别基数在低维连续嵌入空间中易发生质心聚集、流形塌缩或需要繁琐可学习参数的问题,本文引入基于素数平方根有理无关增量的准随机 Kronecker 序列。选取前三个素数的平方根集合 \(\mathcal{V} = \{\sqrt{2}, \sqrt{3}, \sqrt{5}\}\),其在有理数域 \(\mathbb{Q}\) 上线性无关,保证生成的序列永远不会塌缩至低维子空间或产生周期性相关。对任意类别索引 \(k \in \{0, \dots, N-1\}\),质心坐标按如下公式离散生成后归一化到 \([-1, 1]^3\) 空间: $\(\hat{c}_k = (k \cdot \mathcal{V}) \pmod 1, \quad \mu_k = 2 \cdot \text{Norm}(\hat{c}_k) - 1\)$ 该设计无需引入额外的参数优化即可获得各类别之间具有高最小几何间隔、近乎均匀分布的空间锚点,为后续向量场的学习提供了稳健且各向同性的几何基础。

2. 势函数判别校正:引入类间排斥与近质心梯度维系 标准流匹配优化目标仅对真实质心施加单一拉力,在质心附近梯度为 \(\nabla_{\hat{x}_1} \mathcal{L}_{FM} \propto (\hat{x}_1 - x_1)\),导致 \(\hat{x}_1 \to x_1\) 时梯度消失。为此,本文在嵌入空间构建全局语义势场 \(\Phi\)。给定预测状态 \(\hat{x}_1 = x_t + (1-t)v_\theta\) 以及与各质心的欧氏距离平方 \(d_k = \|\hat{x}_1 - \mu_k\|^2\),引入映射算子 \(\mathcal{T}(d) = -\log(d + \epsilon)\) 并通过温度系数 \(\tau\) 的 Softmax 转化为类别软分配概率 \(p_k\)。势函数定义为预测分布与 One-Hot 真实标签 \(y\) 的交叉熵,其对预测状态的解析梯度为: $\(\nabla \Phi = \frac{\partial \Phi}{\partial \hat{x}_1} = \frac{2}{\tau} \sum_{k=1}^N (y_k - p_k) \frac{\hat{x}_1 - \mu_k}{\|\hat{x}_1 - \mu_k\|^2 + \epsilon}\)$ 该校正项针对正样本 \(gt\) 产生标准化的强吸引力,即使在 \(\hat{x}_1\) 极度接近质心时,由于 \(p_{gt} < 1\),\((p_{gt} - 1)\) 项依然维持坚挺的非零梯度模长,彻底破解梯度消失难题;同时,当预测流偏向任意非目标类别 \(j \neq gt\) 时,项 \(p_j \frac{\hat{x}_1 - \mu_j}{d_j + \epsilon}\) 产生反向推力,显式将速度场偏转推离竞争区域,杜绝轨迹穿越。

3. 向量场重塑与截断监督:防范二阶失稳的整流重构 为了将判别校正项无缝融入扩散训练且不破坏流匹配原有的概率路径平滑性,本文通过减去势梯度构建重塑目标速度 \(\tilde{v}_t = v_{gt} - \nabla \Phi\)。关键在于,计算 \(\tilde{v}_t\) 时施加了梯度截断算子 \(\text{sg}[\cdot]\),即网络直接回归重塑后的向量场: $\(\mathcal{L}_{res} = \mathbb{E}_{t, x_0, x_1} \left\| v_\theta(x_t, t, I) - \text{sg}\left[ v_{gt} - \nabla \Phi \right] \right\|^2\)$ 这一操作避免了对非线性势函数进行二阶求导导致的数值爆炸与训练崩溃。网络参数在更新时接收到的有效监督目标由“重构速度”和“判别校正”两部分相加构成,使得学习出的速度场天然具备避开决策边界竞争区域的偏转流动特性。

4. 像素神经场端到端解码:摆脱 VAE 伪影的高保真连续投影 传统基于 VAE 潜在空间的扩散模型受限于解码伪影与高分辨率优化脱节,难以保证像素级的精细语义轮廓。本文借鉴 PixNerd 思想,直接在图像像素空间构建神经场解码。Transformer 骨干网络为每个图像 Patch 动态预测局部紧凑 MLP 的权重矩阵 \(\{\mathbf{W}_1^n, \mathbf{W}_2^n\}\) 并执行 \(L_2\) 归一化;在 Patch 内任意连续空间坐标 \((i, j)\) 处,将离散余弦变换(DCT)位置编码与当前加噪像素状态 \(x_t(i, j)\) 拼接作为输入特征 \(h_{i, j}\),由动态 MLP 直接预测得到对应像素的瞬时速度向量 \(v^n(i, j) = \bar{\mathbf{W}}_2^n \sigma(\bar{\mathbf{W}}_1^n h_{i, j})\),实现了无压缩损耗的端到端细粒度像素建模。

损失函数 / 训练策略

训练采用两阶段递进分辨率策略:第一阶段在 \(256 \times 256\) 分辨率下以 Batch Size 64 进行基础表示训练,并引入表征对齐(REPA)辅助加速收敛;第二阶段在 \(512 \times 512\) 分辨率下以 Batch Size 16 进行微调。温度超参数设为 \(\tau = 1.0\)。训练损失采用带忽略区域 Mask 的截断重塑速度均方误差损失 \(\mathcal{L}_{res}\);推理阶段直接采用确定性一阶 Euler 数值积分求解器,仅需 10 步采样即可兼顾精度与速度,最终将像素预测向量通过最近邻欧氏距离映射回离散类别 ID。

实验关键数据

主实验

在具有高类别基数的两大经典稠密感知基准 ADE20K(150 类)和 COCO-Stuff(171 类)上进行评估。FlowSeg 仅依赖 ImageNet-1k 预训练,便大幅超越了现有基于大规模图文预训练的生成式分割模型,并赶超了主流强判别式模型。

数据集 模型类别 方法 骨干网络 预训练数据 mIoU (%)
ADE20K 判别式专家 DeeplabV3+ ResNet101 IN-1k 44.1
ADE20K 判别式专家 SegFormer MiT-B2 IN-1k 46.5
ADE20K 判别式专家 MaskFormer Swin-T IN-1k 46.7
ADE20K 生成扩散模型 InstructDiffusion† (SD1.5) 大规模图文 33.6
ADE20K 生成扩散模型 PixWizard (Lumina-Next) 大规模图文 32.8
ADE20K 生成扩散模型 FlowSeg (本文) (PixNerd) IN-1k 47.1
COCO-Stuff 判别式专家 DeeplabV3+ ResNet50 IN-1k 38.4
COCO-Stuff 判别式专家 OCRNet HRNet-W48 IN-1k 42.3
COCO-Stuff 判别式专家 SegFormer MiT-B2 IN-1k 44.6
COCO-Stuff 生成扩散模型 SymmFlow (SD2.1) 大规模图文 39.6
COCO-Stuff 生成扩散模型 FlowSeg (本文) (PixNerd) IN-1k 44.9

(注:† 表示其采用指代分割形式逐类别生成掩码)

消融实验

在 ADE20K 验证集上对目标速度构造形式、忽略区域掩码以及预测目标类型进行消融评估:

配置 预测类型 目标速度 \(\tilde{v}_t\) 机制 忽略区域 Mask mIoU (%) 说明
退化模型 A v-pred 时间退火衰减 (Annealing) ✗ 24.65 时间插值削弱了判别势能影响
基础模型 B v-pred 静态重塑 (Static) ✗ 28.15 静态校正带来明显增益 (+3.50)
完整模型 C v-pred 静态重塑 (Static) ✓ 29.75 过滤模糊未标注区域优化显著提升
变体模型 D x-pred 静态重塑 (Static) ✓ 29.06 预测目标形态可行,v-pred 略优

在采样步数探索中,Euler 求解器在不同步数下的 mIoU 表现如下: - 1 步: 43.47 mIoU - 2 步: 45.85 mIoU - 5 步: 46.82 mIoU - 10 步: 47.10 mIoU(最佳性能平衡点) - 20 步: 46.78 mIoU - 50 步: 46.13 mIoU - 100 步: 45.78 mIoU

关键发现

  • 收敛速度与性能双跃升:向量场重塑使模型从训练初期便维持显著更陡峭的收敛曲线。引入 REPA 进一步将基础阶段性能从 29.8 mIoU 提升至 32.7 mIoU。
  • 时间退火削弱判别力:尝试通过 \((1-t)v_t + t(v_t - \nabla\Phi)\) 在时间维度平滑过渡校正项反而导致性能重挫至 24.65 mIoU,证实全时段保持恒定的判别势能修正对塑造清晰的决策流至关重要。
  • 少步数确定性采样极佳:由于学习到的向量场具有强烈的类间排斥与定向流动特性,仅需 10 步 Euler 积分即可达成 47.1 mIoU 的峰值表现;过度增加步数(如 100 步)反而因累积微小积分近似偏差导致轻微性能衰减(45.78 mIoU)。

亮点与洞察

  • 从优化动力学视角切入生成式感知:跳出了以往仅在网络外挂分割头或启发式调参的浅层工程思维,一针见血地指出标准流匹配 MSE 损失在离散感知中面临的近质心梯度消失与缺乏排斥机制的根本矛盾。
  • 截断梯度的优雅向量场融合:无需引入双层优化或计算代价高昂的 Hessian 矩阵,巧妙通过将势函数梯度注入目标速度并通过 stop-gradient 截断,使普通回归损失隐式具备对比学习与边界推离的判别特性。
  • 极具迁移潜力的 Kronecker 拟随机编码:基于代数无理性构建离散类别的空间几何分布,计算开销近乎为零且保证了理论上的最大各向同性间隔,可直接迁移推广到任意高类别基数的多标签、实例级生成式稠密预测任务中。

局限与展望

  • 类别嵌入维度的扩展限制:目前质心编码被严格限制在与图像 RGB 对应的三维色彩空间 \([-1, 1]^3\) 中;虽然便于可视化且适配三通道架构,但当面对数百甚至数千类超大词表分割任务时,3D 空间的几何容量会显得捉襟见肘,未来可探索扩展至高维神经场潜在流空间。
  • 端到端训练计算开销较高:尽管去除了 VAE 编码器,但全像素级动态 MLP 权重生成仍带来了一定的显存开销,依赖两阶段训练方案。未来可探究更高效的轻量化神经隐式表示架构。

相关工作与启发

  • vs InstructDiffusion / PixWizard: 后者依赖庞大的外部基础模型作为多模态图像翻译先验,且在处理多类别时往往退化为单类别逐次提示分割;本文 FlowSeg 直接在标准单次前向流过程中通过向量场重塑实现端到端 150+ 类别多标签解析,效率与精度全面占优。
  • vs SymmFlow: SymmFlow 引入高斯噪声进行随机输运,输出掩码严重受随机种子扰动且依赖 VAE 空间导致色块混叠;FlowSeg 采用确定性整流流动与像素级神经场,保证了预测结果的一致性与超高轮廓边缘保真度。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次从向量场流形动力学视角揭示生成式分割的梯度消失与轨迹穿越缺陷,并给出优雅的势能梯度重塑方案。
  • 实验充分度: ⭐⭐⭐⭐☆ 在两大权威高基数数据集上深入对比了生成式与强判别式基准,涵盖收敛曲线、ODE采样步数及场可视化。
  • 写作质量: ⭐⭐⭐⭐⭐ 论证链条严密,数学公式推导扎实,图表直观且对比剖析切中要害。
  • 价值: ⭐⭐⭐⭐⭐ 成功将生成式分割推向媲美甚至超越判别式专家的崭新高度,为生成-理解统一范式提供了关键的理论与实践基石。