measure-to-measure regression with transformers¶
会议: NeurIPS 2026
arXiv: 2605.28075
领域: 学习理论
关键词: 测度回归, 流匹配, Transformer, 均值场动力学, 单细胞摄动响应
一句话总结¶
本文形式化了非线性测度到测度（M2M）回归问题，并利用 Transformer 天生的均值场与测度依赖结构，提出了静态推前映射（M2M-Static）与动态速度场流匹配（M2M-TFM）两种高表达力且可扩展的算子学习方法。
研究背景与动机¶
在现代数据科学与自然科学中，许多实际问题处理的基本对象并非独立的向量样本，而是整个人口或粒子集合构成的概率分布，例如时间序列集合、点云数据以及单细胞高通量摄动筛选实验。在细胞生物学中，细胞在受到药物或微环境刺激时并非孤立演化，而是通过复杂的细胞间通讯与相互作用作为一个整体集体演化。因此，从控制组细胞群预测处理组细胞群的演化响应，本质上要求模型学习一个定义在概率测度空间上的未知映射算子。
然而，传统的统计回归方法通常假设样本之间相互独立，无法刻画群体层面的分布依赖行为；早期的分布回归方法大多基于核均值嵌入，或者仅能处理输入/输出单侧为分布的设定。近期基于最优传输（OT）的测度回归工作虽然将问题形式化为样本空间上的传输映射，但往往受限于有限的线性算子或马尔可夫核类，不仅计算扩展性差，而且在两个输入测度的支撑集发生重叠时，其推前像也必然产生重叠，面临根本性的表达力瓶颈。
为了突破这一瓶颈，本文切入的角度是：自注意力机制天然对输入样本集具有置换等变性与均值场结构，其本质就是一个由输入经验测度参数化的非线性算子。核心 idea：利用 Transformer 天然的经验测度依赖特性参数化非线性测度到测度映射，分别构建一步直接预测的静态推前映射，以及基于连续时间速度场的 Transformer 流匹配（M2M-TFM），实现无需分布编码器、可直接在测度空间连续外推的高表达力回归。
方法详解¶
整体框架¶
本文考虑的测度到测度（M2M）回归任务定义在概率测度空间 \(\mathcal{P}(\mathbb{R}^d)\) 上。给定训练数据集包含 \(n\) 对经验测度对 \(\{(\mu_i, \nu_i)\}_{i=1}^n\)，其中 \(\mu_i = \frac{1}{N}\sum_{j=1}^N \delta_{x_{ij}}\) 为输入源测度，\(\nu_i = \frac{1}{N}\sum_{j=1}^N \delta_{y_{ij}}\) 为目标测度。目标是学习一个未知的非线性推前算子 \(\mathcal{T}^\star(\mu) = (f_\mu)_\# \mu\)，使得模型能够泛化到未见过的测试测度 \(\mu_j\) 并预测其演化目标 \(\hat{\nu}_j\)。
针对这一问题，本文提出了两条互补的技术路径：其一是静态 M2M 映射（M2M-Static），将 Transformer 整体视作一步推前映射 \(F_\mu^\theta\)，直接利用可微分布损失（如 ED、MMD、Wasserstein）进行优化；其二是动态 M2M 映射（M2M-TFM），利用时间依赖的 Transformer 参数化测度依赖连续性方程中的速度场 \(\chi_t^\theta[\mu(t)](\cdot)\)，将分布层面的学习规约到欧氏空间上的条件流匹配回归，避免了高维分布距离的复杂优化。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["输入经验测度点集<br/>X = {x_1, ..., x_N} ~ mu"] --> B["空间傅里叶特征嵌入<br/>128维正弦特征 + 线性投影"]
B --> C{"模型范式选择"}
C -->|静态单步推前| D["静态 Transformer 映射<br/>自注意力编码均值场 + 可微分布损失"]
C -->|动态连续演化| E["Transformer 流匹配 M2M-TFM<br/>时间调制 AdaLN + 条件速度场回归"]
D --> F["单步输出预测测度<br/>nu_hat = (F_mu)# mu"]
E --> G["欧拉数值积分 ODE<br/>沿着速度场连续积分 t 从 0 到 1"]
G --> H["动态生成目标测度<br/>nu_hat = (f_mu)# mu"]关键设计¶
1. Transformer 的均值场算子表征：将自注意力机制直接视作测度依赖映射 传统的神经网络在处理集合或测度数据时，往往需要设计复杂的辅助分布编码器（如 DeepSets 或图神经网络）提取全局嵌入，这破坏了粒子级别的交互细节且优化繁琐。本文指出，标准的 Transformer 编码器天然具备置换等变性，对于输入测度 \(\mu_N = \frac{1}{N}\sum_{j=1}^N \delta_{x_j}\) 和任意查询粒子 \(x\)，单头自注意力的连续形式精确对应为： $$ \mathrm{Attn}(x; \mu) = \frac{V \int y \exp(\langle Qx, Ky \rangle) d\mu(y)}{\int \exp(\langle Qx, Ky \rangle) d\mu(y)} $$ 将自注意力层与逐点 MLP、归一化和残差连接堆叠后，Transformer 在数学上等价于一个由输入经验测度 \(\mu\) 参数化的非线性粒子交互系统 \(F_\mu^\theta(\cdot): \mathbb{R}^d \to \mathbb{R}^d\)。深层 Transformer 在弱假设下是非线性 M2M 算子的通用逼近器，无需外挂任何分布编码网络即可实现粒子的集体均值场交互。
2. 静态推前优化：结合可微分布几何损失的一步映射学习 为了直接以单步方式恢复 M2M 算子，静态方案将映射直接建模为 \(\mathcal{T}(\mu_N) = (F_{\mu_N}^\theta)_\# \mu_N = \frac{1}{N}\sum_{i=1}^N \delta_{F_{\mu_N}^\theta(x_i)}\)。模型使用 Adam 优化器最小化源目标对上的期望分布差异： $$ \mathcal{L}{\text{M2M}}(\theta) = \frac{1}{n} \sum}^n \mathbb{E{(x, y) \sim \gamma_i} [\mathbf{D}(F^\theta(x), y)] $$ 其中 \(\mathbf{D}\) 为几何可微分布损失，涵盖最大均值差异（MMD）、能量距离（ED）、基于 Sinkhorn 算法近似的 Wasserstein 距离（\(\mathcal{W}_1, \mathcal{W}_2\)），以及在预先计算最优传输耦合 \(\gamma_i\) 条件下的均方误差（OTMSE）。该方案结构直观、推理仅需一次前向传播，但高维下分布损失梯度的估计方差较大。
3. 动态连续流匹配 M2M-TFM：将测度连续性方程规约为点级速度场回归 为了规避高维空间中分布距离优化的维数灾难与高方差，本文借鉴 Wasserstein 空间上的 Otto 微积分，将测度的演化描述为测度依赖连续性方程 \(\partial_t \mu(t) + \nabla \cdot (\mu(t) \chi_t[\mu(t)]) = 0\)。模型利用以时间正弦编码为调制信号的 AdaLN-Transformer 参数化非梯度速度场 \(\chi_t^\theta[\mu(t)](\cdot)\)。在训练时，对于最优传输耦合 \((x, y) \sim \gamma_i\)，定义线性插值路径 \(\phi_t(x, y) = (1-t)x + ty\)，其真实演化速度恒为 \(y - x\)。训练目标退化为极简的点级条件流匹配均方误差： $$ \mathcal{L}{\mathrm{TFM}}(\theta) = \frac{1}{n} \sum \int_0^1 \left| \chi_t^\theta}^n \mathbb{E}_{(x, y) \sim \gamma_i\mu_{i,t} - (y - x) \right|^2 dt $$ 在推理阶段，给定未见过的测试初始测度 \(\mu_j\)，从 \(t=0\) 开始沿速度场进行常微分方程（ODE）数值积分（如欧拉法步进 100 步）至 \(t=1\)，即可获得高保真度的目标测度预测。
损失函数 / 训练策略¶
对于静态模型，采用 Adam 优化器，学习率设为 \(1\times 10^{-4}\)，并借助 GeomLoss 库通过 Sinkhorn 迭代稳定计算 Wasserstein 损失。对于动态 M2M-TFM，利用小批量最优传输（mini-batch OT）在线或离线计算粒子配对，采用线性学习率衰减策略（初始衰减因子 1.0 至最终 0.01），在插值时刻 \(t \sim \mathcal{U}[0, 1]\) 下直接监督预测速度与目标速度的 MSE。
实验关键数据¶
主实验¶
论文在三类极具挑战性的测度回归基准上进行了系统评测：（1）2D 合成多测度字母轮廓在正向扩散与核相互作用破坏下的反演重构（测试集涵盖未见字母 X/Y、数字、几何形状及复合多物体）；（2）高维 McKean-Vlasov 粒子随机动力系统（Kuramoto、FitzHugh-Nagumo、Mean-Field Atlas）；（3）真实世界结直肠癌患者来源类器官（PDO）单细胞质谱流式多标记物（44 维特征）治疗响应预测。
下表汇报了在合成多测度物体（Diffusion 与 Kernel Interaction 破坏逆转）平均指标以及结直肠癌类器官跨患者泛化（held-out PDO-21, PDO-27, PDO-75）上的主要表现：
| 评估任务 | 评估指标 | CFM (基线) | Meta-FM (基线) | M2M-ED (静态) | M2M-W1 (静态) | M2M-OTMSE (静态) | M2M-TFM (本文动态) |
|---|---|---|---|---|---|---|---|
| 多测度物体-Diffusion | \(\mathcal{W}_1 \downarrow\) | 0.2491 | 0.2306 | 0.2204 | 0.2092 | 0.1680 | 0.1300 |
| 多测度物体-Diffusion | \(\text{ED} \downarrow\) | 0.0071 | 0.0082 | 0.0087 | 0.0085 | 0.0044 | 0.0023 |
| 多测度物体-Kernel | \(\mathcal{W}_1 \downarrow\) | 0.2744 | 0.2292 | 0.2155 | 0.1663 | 0.1334 | 0.0780 |
| 多测度物体-Kernel | \(\text{ED} \downarrow\) | 0.0128 | 0.0130 | 0.0106 | 0.0069 | 0.0053 | 0.0020 |
| 类器官跨患者响应 | \(\mathcal{W}_1 \downarrow\) | 4.4099 | 4.4358 | 4.5184 | 4.1005 | 4.2772 | 4.3526 |
| 类器官跨患者响应 | \(\text{ED} \downarrow\) | 0.4828 | 0.5160 | 0.5689 | 0.8270 | 1.0085 | 0.4715 |
| 类器官跨患者响应 | \(r^2 \uparrow\) | 0.8961 | 0.8986 | 0.8882 | 0.7660 | 0.7211 | 0.9029 |
消融实验¶
在 100 维 McKean-Vlasov 复杂粒子动力系统（Kuramoto 系统）上，论文探究了不同演化模型架构对比，以及 M2M-TFM 在不同推理数值积分步数下的性能稳定性：
| 实验设置 / 模型架构 | Kuramoto 100D \(\mathcal{W}_1 \downarrow\) | Kuramoto 100D \(\text{ED} \downarrow\) | 说明 / 分析 |
|---|---|---|---|
| CFM (无测度依赖) | 11.312 | 6.978 | 忽略均值场相互作用，误差显著增加 |
| MF-Transformer (有限差分) | 13.657 | 21.019 | 基于离散时间差商近似，对时间跨度敏感不稳定 |
| NMKV (神经 McKean-Vlasov) | 9.978 | 14.130 | 为测度内泛化设计，跨测度回归能力受限 |
| Meta-FM (双网交替优化) | 7.303 | 6.231 | 仅依赖初始时刻嵌入，且双模型交替优化存在瓶颈 |
| M2M-W1 (静态 Transformer) | 3.684 | 4.786 | 静态一步映射具备良好表现，但在高维分布优化上受限 |
| M2M-TFM (1 step/marginal) | 3.219 | 2.633 | 单步常微分步进已具备极高准确率 |
| M2M-TFM (10 steps/marginal) | 3.226 | 2.586 | 增加积分步数微幅降低能量距离 |
| M2M-TFM (100 steps/marginal) | 3.229 | 2.584 | 完整模型性能最优，动力学轨迹极度平滑稳定 |
关键发现¶
- 动态流匹配大幅优于静态分布损失：在合成多测度反演与物理粒子系统中，M2M-TFM 相比静态模型误差减小 40% 以上。点级速度场回归绕过了分布距离的高方差估计，展现出压倒性的收敛质量与外推精度。
- 测度依赖是高维系统预测的决定性因素：无测度依赖的 CFM 在 McKean-Vlasov 动力学系统中误差高达 11.312（\(\mathcal{W}_1\)），而 M2M-TFM 降至 3.229，证明在群体互作系统中忽略分布依赖会导致致命偏差。
- 生物单细胞质谱数据中展现均衡稳健性：在仅有 10 例患者的真实生物实验中，静态 M2M-\(\mathcal{W}_1\) 虽然在其特定目标指标上占优，但在其他指标上出现退化；而 M2M-TFM 在 \(\text{ED}\)（0.4715）与决定系数 \(r^2\)（0.9029）上均斩获最佳，展示了极佳的综合泛化能力。
亮点与洞察¶
- 将自注意力优雅诠释为测度算子：摆脱了将 Transformer 局限于离散序列的狭隘视角，将其严格确立为由经验测度参数化的非线性均值场算子，理论与工程实现严丝合缝。
- 免除分布编码器的极简架构：相比 Meta-FM 需要单独维护图网络/嵌入模型并进行双重交替迭代，M2M-TFM 直接利用单体 Transformer 同时承担测度上下文感知与流场预测，架构大幅精简且端到端稳定。
- 测度空间外推的通用范式：不仅能处理点云粒子数量不一致的输入输出，更能无缝推广到任意维度的物理动力系统仿真与真实单细胞药物筛选。
局限与展望¶
- 输运动力学路径假设过于简化：目前动态方案主要依赖直线插值路径（OT 条件流匹配），无法显式建模生物细胞分裂、凋亡等质量不守恒动力学，亦未引入非梯度扩散或随机噪声。
- 尚未给出统计收敛速率理论保证：论文虽然证明了表达力，但对于 Transformer 在经验测度空间上学习未知速度场的样本复杂度与泛化误差界仍留作未来工作。
- 拓展至非平衡最优传输与单细胞全基因组筛选：未来可探索与 WFR-FM 结合以支持细胞生灭动力学，并在十亿级单细胞扰动图谱（如 Tahoe-100m）上进一步验证可扩展性。
相关工作与启发¶
- vs Meta-FM (Atanackovic et al., 2025): Meta-FM 需要外置分布编码网络提取单向量嵌入，并在双网络间交替优化，且分布依赖仅施加于 \(t=0\)；本文利用 Transformer 的全轨迹均值场注意力统一了向量场与分布依赖，单网即可完成端到端流匹配。
- vs WFM (Haviv et al., 2025): WFM 面向无配对测度的元生成对齐（无监督测度生成），需要执行测度间 Meta-OT；本文聚焦于严格的有监督测度到测度回归，直接对齐输入-输出条件映射，实验表明直接迁移 WFM 会导致严重欠拟合。
- vs CellOT (Bunne et al., 2023): CellOT 使用输入凸神经网络（ICNN）学习单一样本层面的静态最优传输势函数，忽略了细胞间的集体通信；本文构建的测度依赖算子能够精准捕捉细胞微环境诱导的非线性集体转变。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ [将 Transformer 自注意力等价于非线性经验测度算子，开创性地提出了静态分布回归与动态测度流匹配双重范式。]
- 实验充分度: ⭐⭐⭐⭐⭐ [覆盖几何形状多测度重构、高维物理粒子相互作用仿真以及真实世界临床类器官单细胞反应预测，评测维度完备。]
- 写作质量: ⭐⭐⭐⭐⭐ [数学表述严谨清晰，从连续测度视角到离散算法的推导逻辑一气呵成。]
- 价值: ⭐⭐⭐⭐⭐ [为单细胞摄动响应预测、宏观物理粒子演化等跨学科领域提供了高表达力且易于实现的全新基础架构。]