Tokenised Flow Matching for Hierarchical Simulation Based Inference

📄 arXiv: 2604.20723v1 📥 PDF

作者: Giovanni Charles, Cosmo Santoni, Seth Flaxman, Elizaveta Semenova

分类: cs.LG, cs.AI

发布日期: 2026-04-22

备注: 31 pages, 11 figures

💡 一句话要点

提出Tokenised Flow Matching (TFMPE)方法，加速分层模拟推断并降低计算成本。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 模拟推断 分层模型 似然分解 流匹配 神经代理模型

📋 核心要点

  1. 传统分层SBI方法在训练时需要模拟多个站点，计算成本高昂，限制了其应用。
  2. 论文提出Tokenised Flow Matching (TFMPE)，通过似然分解和token化流匹配，实现单站点模拟训练。
  3. 实验表明，TFMPE在分层SBI基准和实际模型上，能有效降低计算成本并获得校准良好的后验。

📝 摘要（中文）

模拟推断(SBI)中，模拟器评估的成本是关键瓶颈。在具有共享全局参数和可交换的站点级参数与观测的分层设置中，这种结构可用于提高模拟效率。现有的分层SBI方法分解了后验，但仍然为每个训练样本模拟多个站点；本文探索了似然分解(LF)，以从单站点模拟进行训练。在LF采样中，学习每个站点的模拟器神经代理，然后组装合成的多站点观测，以摊销完整分层后验的推断。在此基础上，本文提出Tokenised Flow Matching for Posterior Estimation (TFMPE)，一种通过似然分解支持函数值观测的token化流匹配方法。为了实现系统评估，本文引入了分层SBI的基准。在基准以及真实的传染病和计算流体动力学模型上验证了TFMPE，发现校准良好的后验，同时降低了计算成本。

🔬 方法详解

问题定义：在分层模拟推断(SBI)中，模拟器评估的计算成本是主要的瓶颈。现有的方法，即使分解了后验，仍然需要在训练过程中对每个样本模拟多个站点，导致计算量巨大，限制了其在复杂模型上的应用。尤其是在具有共享全局参数和可交换站点级参数的场景下，如何有效利用这种分层结构来降低计算成本是一个关键问题。

核心思路：论文的核心思路是利用似然分解(Likelihood Factorization, LF)将多站点模拟问题转化为单站点模拟问题。通过学习每个站点的神经代理模型，然后将这些代理模型组合起来，生成合成的多站点观测数据，从而实现对完整分层后验的推断。这种方法避免了直接模拟多站点的复杂过程，显著降低了计算成本。

技术框架：TFMPE方法主要包含以下几个阶段：1) 单站点模拟：对每个站点独立进行模拟，生成单站点的数据。2) 神经代理学习：使用单站点数据训练一个神经代理模型，用于近似该站点的似然函数。3) 似然分解与合成观测：利用似然分解，将多站点的似然函数分解为单站点似然函数的乘积。然后，使用训练好的神经代理模型生成合成的多站点观测数据。4) Tokenised Flow Matching：使用token化的流匹配方法，基于合成观测数据学习后验分布。

关键创新：TFMPE的关键创新在于将token化的流匹配方法与似然分解相结合，从而能够处理函数值观测，并实现高效的分层后验推断。传统的流匹配方法通常只能处理固定维度的观测数据，而TFMPE通过token化技术，将函数值观测转换为token序列，从而能够利用流匹配方法进行学习。此外，似然分解使得可以从单站点模拟中学习，极大地降低了计算成本。

关键设计：TFMPE的关键设计包括：1) Token化策略：选择合适的token化方法，将函数值观测转换为token序列。2) 流匹配网络结构：设计合适的流匹配网络结构，用于学习token序列的后验分布。3) 损失函数：使用合适的损失函数来训练流匹配网络，例如基于最优传输的损失函数。4) 神经代理模型：选择合适的神经代理模型，例如神经网络，来近似单站点的似然函数。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，TFMPE在分层SBI基准和实际的传染病和计算流体动力学模型上表现出色。与现有方法相比，TFMPE能够显著降低计算成本，同时获得校准良好的后验分布。具体来说，TFMPE在某些任务上可以将计算成本降低几个数量级，同时保持与现有方法相当甚至更好的推断精度。这些结果验证了TFMPE的有效性和实用性。

🎯 应用场景

TFMPE方法在多个领域具有广泛的应用前景，例如传染病建模、计算流体动力学、气候模型等。在这些领域中，通常需要对复杂的物理过程进行模拟，并且存在分层结构。TFMPE能够显著降低模拟成本，提高推断效率，从而加速科学发现和工程优化。该方法还可以应用于其他具有分层结构和函数值观测的领域。

📄 摘要（原文）

The cost of simulator evaluations is a key practical bottleneck for Simulation Based Inference (SBI). In hierarchical settings with shared global parameters and exchangeable site-level parameters and observations, this structure can be exploited to improve simulation efficiency. Existing hierarchical SBI approaches factorise the posterior yet still simulate across multiple sites per training sample; We instead explore likelihood factorisation (LF) to train from single-site simulations. In LF sampling we learn a per-site neural surrogate of the simulator and then assemble synthetic multi-site observations to amortise inference for the full hierarchical posterior. Building on this, we propose Tokenised Flow Matching for Posterior Estimation (TFMPE), a tokenised flow matching approach that supports function-valued observations through likelihood factorisation. To enable systematic evaluation, we introduce a benchmark for hierarchical SBI. We validate TFMPE on this benchmark and on realistic infectious disease and computational fluid dynamics models, finding well-calibrated posteriors while reducing computational cost.