Embed and Emulate: Contrastive representations for simulation-based inference

📄 arXiv: 2409.18402v1 📥 PDF

作者: Ruoxi Jiang, Peter Y. Lu, Rebecca Willett

分类: cs.LG, stat.ML

发布日期: 2024-09-27

💡 一句话要点

提出基于对比学习的Embed and Emulate方法，用于加速高维系统中的仿真推断。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 仿真推断 对比学习 高维数据 参数估计 低维嵌入

📋 核心要点

  1. 现有SBI方法在高维物理系统中面临挑战，因为训练高维模拟器成本高昂且复杂。
  2. E&E通过对比学习，学习数据的低维潜在嵌入和潜在空间中的快速模拟器，避免直接在高维空间进行模拟。
  3. 实验表明，E&E在处理高维混沌Lorenz 96系统时，参数估计性能优于现有方法。

📝 摘要（中文）

在科学建模和工程应用中，参数估计方法被广泛用于拟合物理模型和校准数值模拟。针对缺乏易处理似然函数的解析统计模型，现代基于仿真的推断(SBI)方法首先使用数值模拟器生成参数和模拟输出的数据集，然后利用该数据集近似似然函数，并估计给定观测数据的系统参数。一些SBI方法采用机器学习模拟器来加速数据生成和参数估计。然而，由于训练高维模拟器的成本和复杂性，将这些方法应用于高维物理系统仍然具有挑战性。本文提出了一种新的基于对比学习的SBI方法——Embed and Emulate (E&E)，它可以有效地处理高维数据和复杂的多模态参数后验。E&E学习数据的低维潜在嵌入（即摘要统计量）和潜在空间中相应的快速模拟器，从而消除了在推理过程中运行昂贵模拟或高维模拟器的需要。我们通过一个合成实验说明了学习到的潜在空间的理论性质，并通过使用高维混沌Lorenz 96系统进行真实的、非可识别的参数估计任务，证明了该方法优于现有方法。

🔬 方法详解

问题定义：论文旨在解决高维物理系统中，基于仿真的参数估计问题。现有方法，特别是依赖机器学习模拟器加速数据生成的SBI方法，在高维数据下训练模拟器的成本过高，导致推理效率低下。此外，对于具有复杂多模态后验分布的系统，现有方法难以准确估计参数。

核心思路：论文的核心思路是利用对比学习，将高维数据映射到低维潜在空间，并在该潜在空间中训练一个快速模拟器。这样，推理过程不再需要直接在高维空间进行昂贵的模拟，而是通过低维潜在空间的模拟器快速得到结果。对比学习能够有效地学习数据的结构信息，从而保证低维嵌入能够保留足够的信息用于参数估计。

技术框架：E&E方法包含两个主要阶段：嵌入阶段和模拟阶段。在嵌入阶段，使用对比学习训练一个编码器，将高维观测数据映射到低维潜在空间。对比学习的目标是使相似的数据点在潜在空间中距离更近，不相似的数据点距离更远。在模拟阶段，在低维潜在空间中训练一个模拟器，用于预测给定参数下的潜在嵌入。推理时，首先使用编码器将真实观测数据映射到潜在空间，然后使用模拟器在潜在空间中搜索与真实观测数据最匹配的参数。

关键创新：E&E方法的关键创新在于将对比学习引入到SBI框架中，从而能够有效地处理高维数据。与直接在高维空间训练模拟器相比，E&E方法通过学习低维嵌入，大大降低了模拟器的训练成本和推理时间。此外，对比学习能够更好地捕捉数据的结构信息，从而提高参数估计的准确性，尤其是在处理复杂多模态后验分布时。

关键设计：对比学习损失函数通常采用InfoNCE损失，用于最大化正样本对之间的互信息，最小化负样本对之间的互信息。编码器可以使用卷积神经网络或Transformer等结构，具体选择取决于数据的特点。低维潜在空间的维度需要根据具体问题进行调整，通常需要进行实验来确定最佳维度。模拟器可以使用神经网络或高斯过程等模型。参数搜索可以使用基于梯度的优化算法或无梯度优化算法。

📊 实验亮点

论文通过合成实验验证了E&E方法学习到的潜在空间的理论性质。在真实的Lorenz 96系统中，E&E方法在非可识别参数估计任务中表现出优于现有方法的性能。具体提升幅度未知，但摘要强调了其在处理高维数据和复杂后验分布方面的优势。

🎯 应用场景

E&E方法可广泛应用于科学建模和工程领域，例如气候模型校准、生物系统建模、材料设计等。通过降低高维仿真推断的计算成本，E&E方法使得复杂系统的参数估计成为可能，从而提高模型的预测精度和可靠性，并加速科学发现和工程创新。

📄 摘要（原文）

Scientific modeling and engineering applications rely heavily on parameter estimation methods to fit physical models and calibrate numerical simulations using real-world measurements. In the absence of analytic statistical models with tractable likelihoods, modern simulation-based inference (SBI) methods first use a numerical simulator to generate a dataset of parameters and simulated outputs. This dataset is then used to approximate the likelihood and estimate the system parameters given observation data. Several SBI methods employ machine learning emulators to accelerate data generation and parameter estimation. However, applying these approaches to high-dimensional physical systems remains challenging due to the cost and complexity of training high-dimensional emulators. This paper introduces Embed and Emulate (E&E): a new SBI method based on contrastive learning that efficiently handles high-dimensional data and complex, multimodal parameter posteriors. E&E learns a low-dimensional latent embedding of the data (i.e., a summary statistic) and a corresponding fast emulator in the latent space, eliminating the need to run expensive simulations or a high dimensional emulator during inference. We illustrate the theoretical properties of the learned latent space through a synthetic experiment and demonstrate superior performance over existing methods in a realistic, non-identifiable parameter estimation task using the high-dimensional, chaotic Lorenz 96 system.