FLUX: Geometry-Aware Longitudinal Flow Matching with Mixture of Experts¶

会议: NeurIPS 2026
arXiv: 2605.08648
OpenReview: https://openreview.net/forum?id=9kNHWMRtUj
领域: 图像生成
关键词: 流匹配、连续动力系统、非配对时间序列、流形学习、混合专家网络

一句话总结¶

针对非配对纵向时间序列中底层流形弯曲与动力学机制潜在跃迁的双重挑战，FLUX 结合数据驱动度量构建流形感知路径，并利用稀疏混合专家网络分解速度场，在实现高质量纵向传输的同时完成无监督动态机制发现。

研究背景与动机¶

许多复杂的生物与物理系统（如皮层钙成像跨天学习进程、胚胎发育过程中的单细胞转录组分化）本质上沿着连续的时间演化，但在观测上往往只能获取非配对的纵向时间切片。例如，在单细胞测序中细胞测量具有破坏性，在皮层宽场成像中神经元长期漂移，导致相邻时间点之间无法追踪单条轨迹，仅能得到边缘分布序列 $\mu_0, \mu_1, \ldots, \mu_{T-1}$。这类系统在纵向演化建模中面临双重挑战：其一，高维观测数据通常严格依附于弯曲的低维流形，欧氏空间中的直线插值或简单流匹配容易穿透低密度无效区域，违背内在物理或生物约束；其二，控制演化的动力学规则本身会随任务阶段、刺激上下文或发育分化发生机制跃迁，这些跃迁表现为局部输运机理的突变，而非静态观测空间中简单的几何空间聚集。

传统纵向生成模型与最优输运方法（如 Waddington-OT、TrajectoryNet 及 IMMFM）主要聚焦于多边缘分布之间的质量搬运，通常拟合全局共享的单一速度场，或者基于欧几里得测地假设插值，无法显式暴露输运速度场的机制突变。而在离散状态推断领域，诸如隐马尔可夫模型、循环切换线性动力系统（rSLDS）或开关循环神经网络（SRNN）等经典方法，又高度依赖时序连续的追踪轨迹，难以直接应用于非配对的群体快照序列。

本文的核心思路是：将流形度量学习与生成流匹配深度融合，并打破单一速度场假设。核心 idea：通过无监督流形学习构建保流形的弯曲条件路径，并利用 Straight-Through Gumbel-Softmax 门控路由将纵向速度场解耦为稀疏混合专家向量场，直接从局部速度场的动力学分异中无监督发现系统的潜在机制跃迁。

方法详解¶

整体框架¶

FLUX 的整体目标是从非配对边缘序列 $\mu_0, \ldots, \mu_{T-1}$ 中恢复全局连续神经 ODE 速度场 $v_\theta(t, \mathbf{x}(t))$，使得流动映射满足 $(\phi_{0\rightarrow t_k})_\# \mu_0 \approx \mu_k$。整个流程分为三个解耦训练阶段：首先在汇聚的观测样本上学习保流形的数据驱动度量；其次利用弯曲网络（Bend Network）在冻结度量下参数化保边界的流形感知路径；最后冻结几何模型，以流形路径样本和切向量为监督目标，联合训练路由网络与多个并行专家速度场。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["非配对纵向快照序列<br/>μ₀, μ₁, ..., μ_{T-1}"] --> B["RBF-MLP 数据驱动流形度量<br/>紧凑特征映射 + 径向基核中心建模"]
    B --> C["弯曲网络保流形路径构建<br/>边界约束非线性插值 γ_k(α) 与目标切线"]
    C --> D["混合专家速度场分解<br/>Straight-Through Gumbel-Softmax 稀疏路由"]
    D --> E["无监督机制发现与连续纵向输运<br/>确定性硬路由分配 + 神经 ODE 积分"]

关键设计¶

1. RBF-MLP 数据驱动流形度量：规避高维空间测度集中并惩罚偏离流形行为

对于高维复杂生物观测（如嵌入空间维度 $d \ge 5$），标准直接空间核函数往往面临维度灾难与邻域距离集中失效问题。FLUX 设计了 RBF-MLP 几何后端，先通过神经网络将环境样本映射至紧凑特征空间 $q_\theta(\mathbf{x}) \in \mathbb{R}^{d_z}$，并在该特征空间中维护 $M$ 个可学习核中心 $\{\mathbf{c}_m\}_{m=1}^M$。模型由此计算深层核流形置信分数： $$h_\theta(\mathbf{x}) = \sigma\left(a \log \sum_{m=1}^M \exp\left[ -\frac{\|q_\theta(\mathbf{x}) - \mathbf{c}_m\|_2^2}{2\sigma_k^2} \right] + b\right)$$ 通过取幂律倒数，诱导出各向同性空间度量张量 $\mathbf{G}(\mathbf{x}) = (h_\theta(\mathbf{x}) + \epsilon)^{-\alpha} \mathbf{I}$。当采样点偏离高维数据分布的主流形时，$h_\theta(\mathbf{x})$ 趋近于零，度量代价急剧发散，从而在底层几何上严厉阻断路径穿透低密度真空区。

2. 弯曲网络保流形路径构建：免去在线测地线求解并强制端点保真

在相邻时间边缘分布 $(\mu_k, \mu_{k+1})$ 之间，传统流匹配采用直线欧氏插值。FLUX 避免了耗时的测地线边值数值求解，而是训练一个轻量弯曲网络（Bend Network）$B_\psi(\mathbf{x}_k, \mathbf{x}_{k+1}, \alpha; \mathbf{G})$ 来生成非线性位移校正： $$\gamma_k(\alpha) = (1-\alpha)\mathbf{x}_k + \alpha\mathbf{x}_{k+1} + 4\alpha(1-\alpha) \Delta_\psi(\mathbf{x}_k, \mathbf{x}_{k+1}, \alpha)$$ 通过引入调制因子 $4\alpha(1-\alpha)$，位移补偿在两端点 $\alpha=0$ 与 $\alpha=1$ 处自然严格归零，从而完全继承了输入边缘对的端点真实性。对该几何感知路径求导，即可直接得到流匹配的速度场训练目标 $\mathbf{z}_{k,\alpha} = \gamma_k(\alpha)$ 与目标切线 $\dot{\mathbf{z}}_{k,\alpha} = \partial_\alpha \gamma_k(\alpha)$。

3. 混合专家速度场分解：以稀疏动力学门控解耦流场机理与发现机制

为了无监督捕捉演化机理的本质改变，FLUX 将纵向速度场表示为 $K$ 个并行专家向量场 $\{f_m(t, \mathbf{x})\}_{m=1}^K$ 与路由门控 $g(t, \mathbf{x})$ 的动态加权和 $v_\theta(t, \mathbf{x}) = \sum_{m=1}^K w_m(t, \mathbf{x}) f_m(t, \mathbf{x})$。训练阶段使用带有温度退火的 Straight-Through Gumbel-Softmax： $$w_m = \frac{\exp((\ell_m + \eta_m)/\tau_g)}{\sum_{q=1}^K \exp((\ell_q + \eta_q)/\tau_g)}, \quad \eta_m \sim \text{Gumbel}(0,1)$$ 在反向传播中保持连续梯度的同时，前向推理产生稀疏甚至离散的选择。在测试阶段，直接通过最大 logits 硬分配 $\hat{r}(t, \mathbf{x}) = \arg\max_m \ell_m(t, \mathbf{x})$ 提取时间与状态依赖的机制标签，实现了无需时序追踪轨迹即可从流动规律自身反推状态跃迁。

损失函数 / 训练策略¶

速度场与路由网络的总训练目标为 $\mathcal{L} = \mathcal{L}_{\mathrm{FM}} + \lambda_{\mathrm{vel}} \mathcal{L}_{\mathrm{vel}} + \mathcal{L}_{\mathrm{routing}}$。其中，流匹配回归损失在几何弯曲点与目标切向量上评估： $$\mathcal{L}_{\mathrm{FM}} = \mathbb{E}\left[ \| v_\theta(t_{k,\alpha}, \mathbf{z}_{k,\alpha}) - \dot{\mathbf{z}}_{k,\alpha} \|_2^2 \right]$$ 同时引入速度范数正则项 $\mathcal{L}_{\mathrm{vel}} = \mathbb{E}[\|v_\theta\|^2_2]$。路由正则项 $\mathcal{L}_{\mathrm{routing}}$ 综合了防止专家坍缩的批次多样性损失 $\mathcal{L}_{\mathrm{div}}$、负载均衡损失 $\mathcal{L}_{\mathrm{lb}}$、鼓励单样本明确选择的熵正则项 $\mathcal{L}_{\mathrm{sp}}$、抑制饱和的 Z-loss $\mathcal{L}_z$，以及维持时间维度平滑的一致性与全变差（TV）惩罚项 $\mathcal{L}_{\mathrm{TV}}$。

实验关键数据¶

主实验¶

评估在合成控制系统（Lorenz 动力学）、神经认知序列（小鼠皮层钙成像 NeuralTable）及发育生物学单细胞转录组（Embryoid Body）三个基准上展开。实验分别评测单步切片 Wasserstein 距离（$\text{WD}_1 \downarrow$）、全链神经 ODE 积分累积误差（$\text{WD}_{\mathrm{fc}} \downarrow$）以及机制发现精度（片段级调整兰德指数 seg-ARI $\uparrow$ 与归一化互信息 seg-NMI $\uparrow$）。

数据集	模型方法	1-hop WD (↓)	Full-chain WD (↓)	seg-ARI (↑)	seg-NMI (↑)
Lorenz 动力系统	IMMFM (多边缘欧氏)	1.020	1.020	–	–
Lorenz 动力系统	rSLDS (需追踪轨迹)	–	–	1.000	1.000
Lorenz 动力系统	FLUX 无流形学习	1.377	1.377	0.000	0.000
Lorenz 动力系统	FLUX (完整模型)	0.049	0.179	1.000	1.000
神经钙成像 (NeuralTable)	Vanilla CFM (单一速度场)	0.002	0.030+	–	–
神经钙成像 (NeuralTable)	GMM (静态高斯混合)	–	–	0.062	0.189
神经钙成像 (NeuralTable)	PCA + K-Means	–	–	0.487	0.702
神经钙成像 (NeuralTable)	FLUX 无流形学习	0.147	0.147	0.200	0.235
神经钙成像 (NeuralTable)	FLUX (完整模型)	0.002	0.030	0.461	0.589
胚胎体单细胞分化	IMMFM (多边缘欧氏)	0.128	0.368	–	–
胚胎体单细胞分化	FLUX 无流形学习	0.126	0.126	0.571	0.800
胚胎体单细胞分化	FLUX (完整模型)	0.210	0.450	0.571	0.800

消融实验¶

消融实验重点验证流形学习与 MoE 路由协同对于动力学机制识别的不可或缺性：

变体配置	Lorenz seg-ARI (↑)	Lorenz seg-NMI (↑)	NeuralTable seg-ARI (↑)	NeuralTable seg-NMI (↑)	说明
完整 FLUX (RBF-MLP + Gumbel MoE)	1.000	1.000	0.461	0.589	完整模型，机制与输运兼优
去掉流形学习 (FLUX wo Manifold)	0.000	0.000	0.200	0.235	欧氏路径下专家退化坍缩
传统聚类基准 (K-Means)	0.054	0.316	0.487	0.702	缺乏连续时间动力学建模
传统时序动力学 (rSLDS)	1.000	1.000	0.000	0.000	在高维非配对快照中彻底失效

关键发现¶

在 Lorenz 混沌/亚临界吸引子切换实验中，缺乏流形学习的欧氏 MoE 无法察觉流形局部切向曲率的变化，导致路由门控直接坍缩为单一专家（seg-ARI = 0.000）；而 FLUX 完美定位到 $\mu_3$ 与 $\mu_4$ 之间的参数分岔点（seg-ARI = 1.000）。
在小鼠皮层钙成像中，FLUX 识别出的动力学转换边界精确对齐了小鼠在任务学习中 CS++ 与 CS-- 舔水行为分化的发生期，证明皮层神经网络在大脑学习过程中重构了底层的动力学输运流形。
在胚胎体发育数据中，由于分化阶段在 PCA 空间中已具备明显的空间几何分离，欧氏路由与流形路由均能取得 0.571 的 seg-ARI；这从反面印证了流形感知机制最关键的作用场景在于：机制差异隐藏在局部动力学规则中，而非简单静态空间聚类。

亮点与洞察¶

流形感知与机制解耦的协同机制：证明了单纯堆叠 MoE 容量无法自发发现动力学机制；必须依靠保流形的真实速度切线约束，才能暴露不同演化阶段在相空间中的内在流向差异。
无需配对轨迹的时序机制反演：首次通过神经 ODE 速度场的稀疏门控分解，打破了传统切换状态模型（如 HMM、rSLDS）必须依赖追踪样本轨迹的前提，极大拓宽了纵向单细胞及无对准神经成像数据的分析方法。
可复用的三阶段管道设计：度量学习、弯曲路径参数化、速度场训练三阶段解耦，有效防止了流形几何优化与复杂 MoE 门控之间的数值梯度冲突。

局限与展望¶

计算开销较大：训练过程包含度量模型、弯曲网络以及多专家网络共 3 个阶段，整体流水线训练时间较标准欧氏流匹配显著增加。
专家数量预设：专家总数 $K$ 目前需要人工预先指定（如 Lorenz 设为 2，神经数据设为 3），缺乏根据动力学复杂性自适应增删专家的机制。
噪声敏感度：当低维流形周围存在强观测噪声或快照时间间隔跨度过大时，弯曲网络的曲率估计稳定性有待进一步加强。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次将流形感知流匹配与 MoE 速度场分解结合，实现非配对时间序列的无监督动力学机制发现。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖解析几何网格、经典物理动力系统、小鼠皮层钙成像与单细胞分化全方位验证。
写作质量: ⭐⭐⭐⭐⭐ 逻辑论证严谨，数学公式与生物背景结合紧密，实验消融切中核心论点。
价值: ⭐⭐⭐⭐⭐ 为发育生物学、计算神经科学及大规模纵向演化动力学建模提供了通用强力的基础工具。