技术摘要：FLUX – 基于混合专家模型的几何感知纵向流匹配

1. 问题陈述

许多生物系统（如神经回路和分化中的细胞群）通过连续的局部动力学演化，同时在由学习、刺激背景或发育阶段定义的潜在机制之间切换。然而，这些系统通常仅作为非配对的纵向快照被观测到。在宽场钙成像或单细胞 RNA 测序等场景中，无法跨时间点追踪相同的个体细胞或神经元；相反，研究人员观测到的是一系列边缘分布（$\mu_0, \mu_1, \dots, \mu_{T-1}$）。

这种设定带来了两个相互耦合的挑战：

1. 几何约束：高维生物测量数据通常集中在低维、弯曲的非线性流形附近。非配对端点之间的标准欧几里得插值可能会穿过低密度区域，从而违反生物物理约束，并未能尊重状态空间的内在几何结构。
2. 机制发现：底层输运机制（即“流”）通常会随时间变化。例如，神经活动可能从广泛的刺激 - 奖励预期转变为特定的辨别，或者细胞可能从多能性状态过渡到谱系承诺程序。这些是局部动力学（机制）的变化，而不仅仅是空间聚类。现有方法通常侧重于输运精度，但未能分解速度场以识别这些潜在的动力学机制。

2. 方法：FLUX

作者提出了FLUX（用于非配对纵向数据的混合专家流匹配），这是一个联合建模连续输运和无监督机制发现的框架。该方法按三个顺序阶段运行：

A. 几何感知条件路径

FLUX 将度量流匹配（Metric Flow Matching）扩展到纵向场景。FLUX 不使用相邻边缘分布之间的直线欧几里得路径，而是从合并的观测数据（包括未标记样本）中学习数据依赖的度量。

• 度量学习：一个几何后端（使用 RBF 或 RBF-MLP 深度核度量）估计到流形的距离分数，定义一个度量 $G(x)$，对离开数据流形的路径施加惩罚。
• 弯曲网络（Bend Network）：一个网络 $B$ 参数化了端点 $x_k$ 和 $x_{k+1}$ 之间近似的几何感知路径 $\gamma_k(\alpha)$。该路径发生弯曲以保持停留在流形上，提供尊重局部几何结构的训练点 $z_{k,\alpha}$ 和目标切线 $\dot{z}_{k,\alpha}$。

B. 机制切换速度场

FLUX 将全局速度场 $v_\theta$ 分解为混合专家（MoE）。

• 专家向量场：模型学习 $M$ 个稀疏的专家向量场 $f_m$。
• 路由器：一个路由器网络 $g(t, x)$ 将时间和状态映射到专家对数几率（logits）。
• 路由机制：在训练期间，采用直通 Gumbel-Softmax机制为每个样本选择一个活跃专家，从而允许对离散路由决策进行可微优化。在推理阶段，机制由确定性的硬分配（argmax）决定。
• 目标：速度场被训练以匹配几何感知路径的切线。路由器受到正则化，以鼓励稀疏的、时间上连贯的分配，并防止专家坍缩。

C. 训练流程

训练分三个阶段进行以确保稳定性：

1. 度量学习：在合并的样本上训练几何后端。
2. 弯曲网络训练：使用冻结的度量训练弯曲网络，以生成几何感知插值。
3. 速度场训练：冻结几何和弯曲网络；在生成的路径和切线上训练速度专家和路由器。

3. 主要贡献

1. 扩展到纵向边缘分布：作者将几何感知流匹配扩展到有序的边缘分布链，强调了在轨迹标签稀疏但存在未标记观测数据时，无监督流形学习的作用。
2. 通过速度分解进行机制发现：他们引入了一种机制切换速度场，其中稀疏专家路由直接从局部输运结构中揭示潜在的动力学机制，而不是在事后对静态观测进行聚类。
3. 几何与机制的相互作用：该工作证明，几何学习对于机制发现至关重要，特别是当机制由局部动力学的变化而非静态空间分离定义时。
4. 生物学验证：该框架在两个不同的生物学场景中得到了验证：
   • 神经数据：在联想学习中，将早期/中期训练与晚期训练分离，这与行为分歧相吻合。
   • 发育生物学：恢复了与胚胎体中多能性与分化细胞群相关的表达 - 演化机制。

4. 结果

该论文在四个数据集上评估了 FLUX：Stanford Bunny（受控流形）、Lorenz 吸引子（机制切换动力学）、宽场钙成像（神经学习）和胚胎体分化。

• Stanford Bunny：FLUX 在高维嵌入中保持了表面输运，而欧几里得方法无法保持在流形上。
• Lorenz 吸引子：FLUX 通过识别混沌机制和亚临界机制之间的切换，实现了完美的机制恢复（ARI=1.0, NMI=1.0）。关键在于，不带流形学习的 FLUX（使用欧几里得路径）坍缩为单个专家，尽管具有相同的 MoE 架构，却未能发现机制。这证明了几何感知路径揭示了专家专业化所需的曲率和速度特征。
• 神经数据：FLUX 识别出皮层动力学的转变，该转变与行为辨别（CS+ 与 CS−）的出现相一致，其表现优于欧几里得路由和切换动力学基线。
• 胚胎体：虽然欧几里得路由在此处表现良好（因为阶段在空间上是可分离的），但带有几何学习的 FLUX 仍保持竞争力，这加强了当机制编码在局部动力学中时几何学习最为关键的结论。

消融研究结果：论文明确指出，仅靠混合专家路由是不够的。如果没有几何学习，模型可以拟合局部输运，但当这些机制编码在局部动力学中而非静态聚类中时，无法恢复机制。

5. 意义与主张

作者声称，几何感知速度分解提供了一种通用策略，用于从非配对的纵向快照中发现潜在的生物状态转变。

• 科学影响：FLUX 将重点从单纯优化输运精度（最小化 Wasserstein 距离）转移到恢复可解释的动力学结构。它表明，生物系统在学习和分化等过程中，往往会重组为不同的动力学机制。
• 局限性：作者承认，由于需要训练三个不同的网络，FLUX 的计算成本高于欧几里得流匹配。他们还指出，专家的数量是预先固定的，当局部动力学存在噪声或采样较弱时，推断出的机制应谨慎解读。
• 结论：该论文得出结论，FLUX 提供了一个框架，用于对非配对生物快照进行可解释的动力学建模，弥合了几何流形学习与机制切换动力学系统之间的差距。