用向量量化编码分子结构，提升扩散模型生成质量
VQ-SAD: Vector Quantized Structure Aware Diffusion For Molecule Generation
- 将原子和键类型作为潜变量，用VQ-VAE编码为离散代码
- 在QM9和ZINC250k上略超现有扩散模型性能
- 适合需要精确分子结构生成的研究者使用
许多基于扩散的分子生成方法忽略分子的符号信息，将原子和键类型表示为独热编码。基于摩根指纹的方法存在哈希冲突，难以无损嵌入连续空间，随机指纹则可能对应无效分子。为此，我们采用新范式，将原子和键编码视为VQ-VAE的潜变量。提出VQ-SAD：先训练VQ-VAE，冻结预训练模型，并将原子与键的码本作为下游扩散过程的分词器。VQ-SAD是神经符号模型，结合符号与神经结构信息，支持可学习的前向过程。大而离散的代码空间使原子与键类型分布更均衡，增强去噪能力。在QM9和ZINC250k数据集上，其性能略优于当前最先进扩散模型。
原文摘要 · Abstract (English)
Many diffusion based molecule generation methods ignore the symbolic information of molecules and represent the atom and bond type as one hot representation. Methods based on Morgan fingerprints produce hash collisions and are hard to embed into a continuous space without information loss and random fingerprints correspond to no valid molecule. To circumvent this issue we use another paradigm and consider atom and bond codes as latent variables of VQ-VAE. We introduce VQ-SAD which first trains a VQ-VAE and uses the frozen pretrained VQ-VAE model and considers the codebooks for both atom and bond types as tokenizers for the downstream diffusion process. VQ-SAD is a neuro-symbolic model that utilizes both symbolic and neural structural information for a diffusion based model with learnable forward process. The large discrete code space provides a more balanced atom and bond types which enhances the denoising process. VQ-VAE slightly outperforms SOTA models for diffusion based molecule generation on QM9 and ZINC250k datasets.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。