[论文解读] Fast and Functional Structured Data Generators Rooted in Out-of-Equilibrium Physics

Alessandra Carbone, Aurélien Decelle+2 more|arXiv (Cornell University)|Jul 13, 2023

Machine Learning in Materials ScienceMaterials Science被引用 3

一句话总结

本文提出了一种用于受限玻尔兹曼机（RBMs）的新型非平衡训练方法，通过利用非平衡马尔可夫链动力学，实现了对基因组、蛋白质和RNA序列等结构化数据的快速、高质量条件生成。该方法仅需10步MCMC即可达到最先进水平的样本质量，在多样性和速度方面显著优于传统平衡训练方法，且在多种生物数据集上表现出更高的稳定性。

ABSTRACT

In this study, we address the challenge of using energy-based models to produce high-quality, label-specific data in complex structured datasets, such as population genetics, RNA or protein sequences data. Traditional training methods encounter difficulties due to inefficient Markov chain Monte Carlo mixing, which affects the diversity of synthetic data and increases generation times. To address these issues, we use a novel training algorithm that exploits non-equilibrium effects. This approach, applied on the Restricted Boltzmann Machine, improves the model's ability to correctly classify samples and generate high-quality synthetic data in only a few sampling steps. The effectiveness of this method is demonstrated by its successful application to four different types of data: handwritten digits, mutations of human genomes classified by continental origin, functionally characterized sequences of an enzyme protein family, and homologous RNA sequences from specific taxonomies.

研究动机与目标

• 为解决传统能量模型在生成多样化、标签特定的结构化数据时因MCMC混合性能差和收敛缓慢而导致的局限性。
• 克服在处理复杂、多模态生物数据集（如基因组和蛋白质序列）时，RBMs中平衡训练方法效率低下且不稳定的缺陷。
• 开发一种训练框架，实现在极少数采样步骤下快速、高保真地进行条件生成，即使对于罕见或亚家族特异性序列类别亦可实现。
• 证明RBMs的非平衡训练可作为类似扩散的生成器，提升样本多样性与生成速度，同时保持分类准确性。
• 在多种生物数据集上验证该方法，包括手写数字、人类基因组突变、酶蛋白家族和RNA序列，并通过生物上具有意义的结构验证。

提出的方法

• 该方法采用双梯度算法，利用非平衡马尔可夫链动力学训练RBMs，避免完全收敛至平衡态。
• 采用改进的对比分歧方法，其中梯度估计与采样均在非平衡状态下进行，从而提升训练的稳定性和速度。
• 通过将类别标签嵌入能量函数，使模型能够从随机初始化生成标签条件化的样本，实现条件生成。
• 每次生成仅需10步MCMC采样，显著缩短推理时间，同时保持高样本质量。
• 该方法适用于具有简单能量函数的RBMs，但该框架可扩展至更复杂的架构（如含卷积层的模型）。
• 通过多种指标评估：特征值谱误差、熵误差和对抗准确性，并通过ESMFold预测的pLDDT分数进行生物验证。

实验结果

研究问题

• RQ1 RBMs的非平衡训练能否在基因组学和蛋白质组学等复杂结构化数据集中生成高质量、多样化且标签特定的合成数据？
• RQ2 与标准平衡方法相比，RBMs的非平衡训练是否能提升生成速度和样本多样性？
• RQ3 即使在每种子家族数据量有限的情况下，该模型能否生成具有正确功能和结构特性的生物上合理的蛋白质序列？
• RQ4 F&F-10模型在生成质量与采样效率方面与平衡RBMs训练相比表现如何？
• RQ5 当仅有少量训练样本时，该模型在罕见或低频序列类别上的泛化能力如何？

主要发现

• F&F-10模型仅用10步MCMC采样即生成了高质量的合成样本，在MNIST、HGD、GH30和SAM等所有测试数据集上均与真实数据具有相近的多样性和结构特征。
• 该模型在所有数据集上均实现了高分类准确率，即使在子家族仅含100个训练样本时也表现出强大的泛化能力。
• 生成的蛋白质序列的pLDDT评分分布与真实测试数据几乎完全一致（由ESMFold预测），表明其具有高度的结构可靠性和生物合理性。
• 与标准持久对比分歧（PCD）训练相比，该方法显著减少了训练时间并提升了稳定性，后者在多模态分布中常因混合不良而失败。
• F&F-10模型在样本质量和生成速度方面均优于平衡RBMs训练，且仅需极少采样步数即可输出高保真结果。
• 双梯度非平衡训练策略实现了高效的条件生成，无需依赖长MCMC链，因此适用于现实世界中的生物数据生成任务。