Mitigating mode collapse in normalizing flows by annealing with an adaptive schedule: Application to parameter estimation

📄 arXiv: 2505.03652v1 📥 PDF

作者: Yihang Wang, Chris Chi, Aaron R. Dinner

分类: cs.LG, physics.comp-ph, physics.data-an, q-bio.QM, stat.ML

发布日期: 2025-05-06

备注: 19 pages, 10 figures

💡 一句话要点

提出基于自适应退火策略的 Normalizing Flow 训练方法，缓解模式崩塌问题，加速参数估计。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: Normalizing Flows 模式崩塌 参数估计 自适应退火 有效样本量 MCMC 生化振荡器

📋 核心要点

  1. Normalizing Flows在参数估计中面临模式崩塌问题，导致采样效率低下，难以捕捉复杂分布。
  2. 论文提出基于有效样本量(ESS)的自适应退火策略，动态调整训练过程，缓解模式崩塌现象。
  3. 实验表明，该方法在生化振荡器模型参数估计中，计算效率优于传统MCMC方法十倍。

📝 摘要（中文）

Normalizing Flows (NFs) 能够从复杂分布中生成不相关的样本，使其成为参数估计的有力工具。然而，由于 NFs 容易崩塌到多峰分布的单一模式，其实际应用受到限制。本研究表明，采用基于有效样本量 (ESS) 的自适应退火策略可以缓解模式崩塌。我们证明，对于拟合时间序列数据的生化振荡器模型，我们的方法比广泛使用的集成马尔可夫链蒙特卡洛 (MCMC) 方法减少了十倍的计算时间即可收敛边际似然。我们还表明，ESS 可用于通过修剪样本来减少方差。我们期望这些进展对 NFs 的采样具有普遍用途，并讨论了进一步改进的潜在机会。

🔬 方法详解

问题定义：论文旨在解决 Normalizing Flows (NFs) 在应用于参数估计时，容易出现的模式崩塌问题。当目标分布是多峰时，NF 可能会陷入其中一个模式，而忽略其他模式，导致采样结果的偏差和不准确。现有的方法，如传统的 MCMC，计算成本高昂，而 NFs 的模式崩塌问题限制了其在复杂参数估计问题中的应用。

核心思路：论文的核心思路是利用退火算法，并结合自适应的退火策略来引导 NF 的训练过程，使其能够探索和覆盖目标分布的多个模式。自适应策略的关键在于使用有效样本量 (ESS) 作为指标，动态调整退火的强度。ESS 可以反映当前样本的多样性，当 ESS 较低时，表明发生了模式崩塌，需要加强退火，鼓励探索更多的模式。

技术框架：整体框架包含以下几个主要步骤：1) 初始化 NF 模型；2) 使用当前 NF 模型生成样本；3) 计算样本的 ESS；4) 根据 ESS 调整退火参数（例如，温度）；5) 使用调整后的退火参数训练 NF 模型；6) 重复步骤 2-5，直到收敛。退火过程通过在损失函数中引入一个温度参数来实现，该参数控制了目标分布的平滑程度。高温度对应于更平滑的分布，有助于 NF 逃离局部模式。

关键创新：最重要的技术创新点在于提出了基于 ESS 的自适应退火策略。与传统的固定退火策略相比，自适应策略能够根据 NF 的实际采样情况动态调整退火强度，从而更有效地缓解模式崩塌。此外，利用 ESS 进行样本修剪也是一个创新点，可以减少采样方差，提高参数估计的准确性。

关键设计：关键设计包括：1) ESS 的计算方法：论文采用了标准的 ESS 计算公式，用于衡量样本的多样性；2) 退火参数的调整策略：论文设计了一种基于 ESS 的退火参数调整函数，该函数将 ESS 映射到退火温度，使得当 ESS 较低时，退火温度较高，反之亦然；3) 样本修剪策略：论文根据样本的权重，选择性地保留一部分样本，从而减少方差。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，该方法在生化振荡器模型参数估计中，比传统的集成 MCMC 方法减少了十倍的计算时间即可收敛边际似然。此外，通过 ESS 进行样本修剪，可以有效地减少采样方差，进一步提高参数估计的准确性。这些结果表明，该方法在缓解模式崩塌和提高采样效率方面具有显著优势。

🎯 应用场景

该研究成果可广泛应用于需要进行参数估计的领域，例如生物化学、物理学、金融学等。在生物化学中，可以用于估计复杂的生化反应网络的参数；在物理学中，可以用于拟合实验数据，估计物理模型的参数；在金融学中，可以用于校准金融模型，进行风险评估。该方法有望提高参数估计的效率和准确性，加速相关领域的研究进展。

📄 摘要（原文）

Normalizing flows (NFs) provide uncorrelated samples from complex distributions, making them an appealing tool for parameter estimation. However, the practical utility of NFs remains limited by their tendency to collapse to a single mode of a multimodal distribution. In this study, we show that annealing with an adaptive schedule based on the effective sample size (ESS) can mitigate mode collapse. We demonstrate that our approach can converge the marginal likelihood for a biochemical oscillator model fit to time-series data in ten-fold less computation time than a widely used ensemble Markov chain Monte Carlo (MCMC) method. We show that the ESS can also be used to reduce variance by pruning the samples. We expect these developments to be of general use for sampling with NFs and discuss potential opportunities for further improvements.