跳转至

Disentangled Representation Learning via Flow Matching

会议: NeurIPS 2026
arXiv: 2602.05214
OpenReview: https://openreview.net/forum?id=bXNMPTOJ05
领域: 图像生成
关键词: 解耦表示学习, 流匹配, 连续时间动力学, 输出注意力, 正交正则化

一句话总结

将解耦表示学习建模为紧致潜空间内的因子条件流匹配过程,并通过输出注意力与非重叠正交正则化解耦速度场,实现兼顾高保真生成与精确语义干预的高质量解耦表征。

研究背景与动机

解耦表示学习(Disentangled Representation Learning)旨在从观测数据中分离并提取出底层的独立生成因子(如物体的形状、颜色、大小与空间姿态),从而使表征具有可解释性,并为下游控制生成、图像编辑与强化学习策略迁移提供结构化的归纳偏置。长期以来,生成模型在此方向占据主导地位:以 \(\beta\)-VAE 和 FactorVAE 为代表的变分自编码器通过对隐空间施加全相关惩罚或先验结构约束实现解耦,但始终受制于重构保真度与解耦强度之间的内在权衡;而以 InfoGAN 和 StyleGAN 为代表的对抗生成网络虽然生成质量卓越,却普遍缺乏从数据空间到隐空间的稳定可逆推理机制,导致表示学习与精细干预难以兼得。

近年来,以扩散模型(Diffusion Models)为代表的现代生成范式因其强大的模式拟合能力和训练稳定性受到广泛关注。然而,现有的扩散解耦方法通常依赖强归纳偏置来促使潜变量维度在统计上相互独立。在实际应用中,解耦不仅要求统计独立,更依赖严格的「语义对齐(Semantic Alignment)」准则,即每个潜变量单元或子空间必须唯一对应一个可解释的物理生成因子。由于扩散模型的迭代去噪轨迹具有固有的随机性,且缺少显式的几何结构约束,各潜变量维度往往在统计解相关的情况下依然互相混杂语义信息,导致跨因子泄露与模糊归因,极大地削弱了因子级干预的精确度。

面对上述瓶颈,连续时间流匹配(Flow Matching)的确定性常微分方程(ODE)表述为解决语义对齐提供了更具原则性的几何视角。流匹配直接回归概率流动场,不仅避免了离散时间噪声调度和迭代去噪的计算开销,还天然具备确定性的平滑动力学轨迹。核心 idea:将解耦表示学习转化为紧致潜空间内的因子条件流传输动力学,并利用输出注意力将总速度场显式解构为互不重叠的正交因子分量,通过非重叠正则化彻底切断跨因子干扰与信息泄漏。

方法详解

整体框架

本文提出的方法在由预训练 VQ-GAN 编码器构建的紧致潜空间内建模连续流匹配。输入图像首先通过可学习的因子编码器提取出一组独立的语义因子,并通过交叉注意力机制注入到 U-Net 主干的各层特征中;随后,U-Net 预测的整体传输速度场经过输出注意力机制被路由分解为多个因子专属性的速度分量,最后在全场匹配损失与正交非重叠正则化的联合约束下完成端到端优化。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入图像 I"] --> B["因子编码器与紧致潜编码映射<br/>S(I)={s_i} 与 z_1=E(I)"]
    B --> C["因子条件流匹配场<br/>线性插值路径与交叉注意力条件注入"]
    C --> D["输出注意力速度场解构<br/>单纯形权重路由分解 factor 专有速度"]
    D --> E["非重叠正交正则化约束<br/>最小化因子分量成对平方余弦相似度"]
    E --> F["确定性 ODE 采样与因子可控生成<br/>Dopri5 积分输出重构或属性干预图像"]

在推理阶段,可以通过标准高斯先验抽样初始隐变量,结合 Dopri5 自适应步长数值积分器求解常微分方程恢复目标潜变量,并通过替换特定因子嵌入实现精确的属性受控图像合成。

关键设计

1. 因子条件流匹配动力学:建立确定性平滑传输路径

传统扩散模型由于迭代去噪路径的布朗运动随机性,隐变量动力学容易在时间步之间发生语义漂移。本文将生成任务设定为从标准高斯分布 \(p_0(\mathbf{z}) = \mathcal{N}(0, \mathbf{I})\) 向数据潜分布 \(p_1(\mathbf{z})\) 的最优速度场匹配。给定图像 \(I\),通过预训练 VQ-GAN 提取目标潜表示 \(\mathbf{z}_1 = E(I)\),并在源潜变量 \(\mathbf{z}_0 \sim \mathcal{N}(0, \mathbf{I})\) 与 \(\mathbf{z}_1\) 之间建立确定性线性插值桥梁: $\(\mathbf{z}_t = (1-t)\mathbf{z}_0 + t\mathbf{z}_1, \quad t \in [0, 1]\)$ 该桥梁在时间维度诱导出的真实目标速度为恒定矢量 \(\mathbf{u}_t = \mathbf{z}_1 - \mathbf{z}_0\)。网络使用 U-Net 参数化速度场 \(\mathbf{v}_\theta(\mathbf{z}_t, S_\gamma(I), t)\),并通过交叉注意力机制在特征层引入可学习因子编码器输出的 \(N\) 个解耦语义因子 \(S_\gamma(I) = \{s_i\}_{i=1}^N\)。流匹配回归损失约束整体场追踪线性轨迹: $\(\mathcal{L}_{\mathrm{FM}}(\theta) = \mathbb{E}_{t, \mathbf{z}_0, \mathbf{z}_1} \left[ \|\mathbf{v}_\theta(\mathbf{z}_t, S_\gamma(I), t) - (\mathbf{z}_1 - \mathbf{z}_0)\|_2^2 \right]\)$ 确定性的 ODE 轨迹消除了随机布朗项的扰动,使得从先验到数据的传输在几何上具有平直、可追踪的演化特性。

2. 基于输出注意力的速度场解构:构建因子专属动力学

仅仅约束总速度场并不能促成语义解耦,因为 \(\mathcal{L}_{\mathrm{FM}}\) 只要求所有因子共同贡献出的总速度与目标一致,各因子通道容易产生特征冗余或相互借位。为此,本文提出将总速度场显式解构为 \(N\) 个因子速度分量之和 \(\mathbf{v}_\theta = \sum_{i=1}^N \mathbf{v}_\theta^{(i)}\)。该分解通过输出注意力(Output Attention)路由掩码优雅实现:利用 U-Net 顶层隐藏特征 \(h\) 生成空间查询向量 \(\mathbf{Q} = \mathbf{W}_q h\),并与经过线性映射的因子键向量 \(\mathbf{K}_i = \mathbf{W}_k s_i\) 计算归一化注意力权重: $\(\mathrm{Attn}_i = \frac{\exp(\mathbf{Q}\mathbf{K}_i^\top / \sqrt{d})}{\sum_{j=1}^N \exp(\mathbf{Q}\mathbf{K}_j^\top / \sqrt{d})}\)$ 由于每个空间位置满足严格的单纯形约束 \(\sum_{i=1}^N \mathrm{Attn}_i = 1\),通过逐元素门控得到因子分量 \(\mathbf{v}_\theta^{(i)} = \mathrm{Attn}_i \odot \mathbf{v}_\theta\)。这种结构设计既保证了质量守恒(分解分量自然重组为总速度,完整保留流匹配的全局拟合监督),又赋予了每个因子沿空间维度的排他性动态归属权。

3. 非重叠正交正则化:阻断跨因子信息渗透

为了防止不同的因子分量在传输空间中方向趋同,模型必须显式惩罚因子速度场之间的平行性。将各因子分量展平为一维向量 \(\mathbf{\alpha}_i = \mathrm{vec}(\mathbf{v}_\theta^{(i)})\) 后,计算两两分量之间的平方余弦相似度均值,构建正交非重叠损失: $\(\mathcal{L}_{\mathrm{orth}} = \frac{1}{N(N-1)} \sum_{i \neq j} \left( \frac{\mathbf{\alpha}_i^\top \mathbf{\alpha}_j}{\|\mathbf{\alpha}_i\|_2 \|\mathbf{\alpha}_j\|_2 + \varepsilon} \right)^2\)$ 该损失驱动不同因子的速度矢量在切空间保持近似正交,从而迫使每个语义因子独立控制一组不重叠的潜变量动力学自由度。总训练损失为两项的线性加权组合 \(\mathcal{L}(\theta, \phi) = \mathcal{L}_{\mathrm{FM}}(\theta) + \lambda_{\mathrm{orth}} \mathcal{L}_{\mathrm{orth}}\),在保证生成流平顺性的同时实现了各语义轴的硬解耦。

损失函数 / 训练策略

训练采用端到端方式同时优化流网络参数 \(\theta\) 与因子编码器参数 \(\gamma\)。批大小设置为 128,正交正则项系数在验证集搜索后固定设为 \(\lambda_{\mathrm{orth}} = 10^{-2}\)。在全部合成与真实数据集上,设定候选因子数量默认值为 \(N=10\)。采样过程采用自适应步长 Dopri5 积分器求解 ODE,兼顾求解精度与推理吞吐。

实验关键数据

主实验

在标准解耦评测基准 Cars3D、Shapes3D 和 MPI3D-toy 上,采用 FactorVAE 分数(衡量因子对齐度)与 DCI 解耦指标(衡量维度专一性)进行定量评测,结果如下表所示:

数据集 指标 本文 之前SOTA (EncDiff / DyGA) 提升幅度
Cars3D FactorVAE ↑ 0.964 ± 0.013 0.948 ± 0.017 (EncDiff) +0.016
Cars3D DCI ↑ 0.431 ± 0.034 0.357 ± 0.072 (EncDiff) +0.074 (+20.7%)
Shapes3D FactorVAE ↑ 1.000 ± 0.000 0.999 ± 0.001 (EncDiff) +0.001
Shapes3D DCI ↑ 0.973 ± 0.020 0.952 ± 0.028 (EncDiff) +0.021
MPI3D-toy FactorVAE ↑ 0.907 ± 0.017 0.862 ± 0.026 (EncDiff) +0.045 (+5.2%)
MPI3D-toy DCI ↑ 0.649 ± 0.024 0.629 ± 0.027 (EncDiff) +0.020

在复杂真实人脸数据集 CelebA 上的属性解耦可控性(TAD 指标)与生成图像质量(FID 指标)对比:

方法类别 代表方法 CelebA TAD ↑ CelebA FID ↓
VAE-based FactorVAE 0.148 ± 0.031 97.6 ± 1.8
Diffusion-based DisDiff 0.305 ± 0.010 18.2 ± 2.1
Diffusion-based FDAE 0.326 ± 0.044 19.0 ± 1.6
Diffusion-based EncDiff 0.638 ± 0.008 14.8 ± 2.3
Diffusion-based DyGA 0.954 ± 0.024 12.0 ± 1.2
Flow-based (本文) Ours 1.154 ± 0.089 8.1 ± 0.2

消融实验

在三项合成基准上验证正交非重叠正则化项 \(\mathcal{L}_{\mathrm{orth}}\) 的关键贡献:

配置 Cars3D (FactorVAE / DCI) Shapes3D (FactorVAE / DCI) MPI3D-toy (FactorVAE / DCI) 说明
纯流匹配 \(\mathcal{L}_{\mathrm{FM}}\) 0.944 ± 0.034 / 0.414 ± 0.029 0.999 ± 0.001 / 0.961 ± 0.026 0.745 ± 0.031 / 0.571 ± 0.062 基础模型,缺乏显式正交约束
完整模型 \(\mathcal{L}_{\mathrm{FM}} + \mathcal{L}_{\mathrm{orth}}\) 0.964 ± 0.013 / 0.431 ± 0.034 1.000 ± 0.000 / 0.973 ± 0.020 0.907 ± 0.017 / 0.649 ± 0.024 引入正交约束后显著抑制跨因子泄露

关键发现

  • 正交正则化项对最具挑战性的现实物理渲染基准 MPI3D-toy 提升最为显著:使 FactorVAE 分数从 0.745 跃升至 0.907(提升超 21.7%),DCI 从 0.571 提升至 0.649,这证明在因子间相互纠缠严重的复杂场景下,显式的速度矢量去相关约束是实现有效解耦的关键保障。
  • 在真实人脸数据集 CelebA 上,流匹配模型不仅将属性可控性指标 TAD 从以往最优扩散模型 DyGA 的 0.954 大幅推升至 1.154,更将生成保真度 FID 从 12.0 骤降至 8.1,打破了传统 VAE 和扩散解耦框架下「解耦度与生成清晰度不可兼得」的妥协局面。
  • 在下游少样本分类任务(GBT 分类器评测)中,当仅使用 100 个带标签样本时,本文学到的表征在 Shapes3D 上的分类准确率达到完整样本准确率的 90.8%(基线仅在 73.2%~80.0%),体现出纯粹因果属性解耦对低数据量场景的极高统计效率。

亮点与洞察

  • 将流匹配的平滑概率路径引入解耦表征:突破了扩散模型布朗随机去噪带来的语义漂移瓶颈,借助确定性 ODE 动力学建立了因子到潜轨迹的一对一平滑映射。
  • 输出注意力与质量守恒的解构机制:通过空间单纯形归一化分解总速度,既保证了全场严格遵循流匹配拟合目标,又为每个因子分配了互不冲突的专属性动态通道。
  • 极佳的下游迁移与少样本泛化力:在 GBT 下游属性预测测试中,极少样本(100 个样本)即可激活接近全监督的预测性能,为通用表征预训练提供了强有力的几何解耦范例。

局限与展望

  • 潜在因子数目先验依赖:当前方法需在训练前显式指定因子的数量 \(N\)(默认设为 10)。尽管敏感度分析表明增加 \(N\) 可以增强容量,但在面对未知未知生成维度的开集场景时,无法自动推断真实的有效因子总数。
  • 依赖预训练潜空间质量:模型以预训练 VQ-GAN 潜表示为目标,若第一阶段编码器自身丢失细微高频几何纹理,流匹配过程无法从根本上恢复这些丢失的信息。
  • 未来方向:探索可根据数据复杂度动态修剪或增生因子维度的自适应稀疏流机制,并将框架推广至视频与具身物理交互中的动作解耦任务。

相关工作与启发

  • vs VAE 系列 (FactorVAE / \(\beta\)-TCVAE):VAE 依靠在潜变量先验上强化因子独立假设,但受限于 ELBO 的重构-正则化博弈;本文基于连续流匹配并在确定性 ODE 动力学中引入正交解构,彻底化解了生成质量与解耦强度的对立。
  • vs 扩散解耦方法 (DisDiff / EncDiff / DyGA):现有扩散解耦模型需要设计复杂的随机时间步调度或动态高斯针刺机制;本文采用平直流匹配线段与输出注意力门控,实现更具几何确定性的显式语义对齐,显著改善跨因子信息泄露。

评分

  • 新颖性: ⭐⭐⭐⭐☆ [首次系统性将流匹配连续时间动力学与显式输出注意力速度解构结合用于通用解耦表示学习]
  • 实验充分度: ⭐⭐⭐⭐⭐ [覆盖三大经典合成评测集与真实复杂人脸数据集,包含严谨消融、敏感度分析、属性交换定性分析及下游少样本迁移实验]
  • 写作质量: ⭐⭐⭐⭐⭐ [数学推导清晰严密,动力学物理图景直观,逻辑链条层层递进]
  • 价值: ⭐⭐⭐⭐☆ [为基于现代生成模型的因果解耦表征与精确可控图像合成提供了极具实用价值的新路径]