Tokenization for Molecular Foundation Models

技术摘要：用于分子基础模型的分词技术

问题陈述
准确预测化学性质对于从储能到药物发现等各个行业都至关重要。虽然 Transformer 架构彻底改变了自然语言处理（NLP）领域，但其在分子基础模型中的应用面临着一个根本性的瓶颈：分词（Tokenization）。目前的分子模型主要依赖于“原子级”（Atom-wise）分词，即使用固定词表将简化分子线性输入规范（SMILES）字符串拆分为原子级标记。

这种方法的主要局限性在于无法完全覆盖 OpenSMILES 规范。原子级分词器将括号内的原子（这些括号编码了同位素、手性中心、电荷和显式氢计数）视为单一且不可分割的标记。为了覆盖这些特征的所有可能排列组合，词表将需要超过 28 万亿个标记。因此，现有的模型通常使用少于 3,000 个标记的词表，导致了显著的覆盖范围缺失。当遇到新型括号原子时，这些闭词表分词器会退而求其次使用通用的未知标记 [UNK]，这可能会掩盖关键的化学信息，例如手性或特定的同位素组成。此外，现有的开集词表替代方案（如基于 BPE 的方案）经常会出现歧义，例如将不同的化学实体（如硫-碳键与元素钪 Sc）混淆为同一个标记。

方法论
作者提出了一种新的分词框架 Smirk，以及其压缩变体 Smirk-GPE，旨在实现对 OpenSMILES 规范的完整覆盖，同时保持计算效率。

1. Smirk 分词： 该方案采用基于 OpenSMILES 规范定义的字形（Glyphs）对 SMILES 字符串进行两阶段字符级分解。

   • 第一阶段： 分解为原子（例如，OC[C@@H][OH] $\rightarrow$ O C [C@@H] [OH]）。
   • 第二阶段： 将括号内的原子分解为组成字形（例如，[C@@H] $\rightarrow$ [ C @ @ H ]）。
   • 这种方法通过将括号和内部符号视为不同的标记，区分了歧义序列（例如，作为键的 Sc 与作为钪元素的 [Sc]）。生成的词表固定为 165 个标记，无需训练，并保证任何 OpenSMILES 编码的分子都可以进行分词而无需使用 [UNK] 标记。

2. Smirk-GPE（字形对编码）： 为了解决由于完全分解括号原子而导致的序列长度（繁殖率/Fertility）增加问题，作者实现了 Smirk-GPE。该变体针对字形标记应用了一种类似于字节对编码（BPE）的压缩策略。与合并字符串的标准 BPE 不同，Smirk-GPE 在标记 ID 上学习合并规则，以确保具有化学意义的合并（例如，结合硫和碳的字形）不会与原子符号（例如，钪 Sc）产生歧义。

3. 评估框架：

   • 内在指标： 作者使用繁殖率（平均序列长度）、归一化熵（压缩效率）、标记不平衡度和 [UNK] 标记频率来评估分词器。
   • 低成本代理模型： 考虑到为每个分词器训练完整的 Transformer 模型计算成本过高，作者利用 n-gram 模型 作为代理。他们在 16 亿个 SMILES 字符串上训练了 n-gram 模型，并通过测量交叉熵损失和信息损失（通过 KL 散度）来估计下游性能。
   • 外部验证： 为了验证 n-gram 代理的有效性，作者使用 11 种不同的分词器和三种分子编码，从头开始预训练了 18 个仅编码器（Encoder-only）的 RoBERTa 模型。随后，这些模型在来自 MoleculeNet 和 tmQM 的六个回归任务和七个分类任务上进行了微调。

关键结果

• 覆盖范围： Smirk 和 Smirk-GPE 是评估的分词器中唯一实现 100% OpenSMILES 规范覆盖的方案，消除了使用 [UNK] 标记的情况。相比之下，现有的化学专用分词器（包括 SPE、APE 和各种 BPE 变体）都会以不可忽视的频率产生 [UNK] 标记（在 tmQM 数据集上高达 ~50%）。
• 信息损失： 覆盖范围有限的分词器表现出显著的信息损失，特别是在富含过渡金属和立体化学的数据集（如 tmQM）上。例如，由于未知标记的存在，MoLFormer 分词器在 tmQM 上产生了 40.3 nats/分子的损失，而 Smirk 变体减轻了这种退化。
• 性能相关性： 研究发现 n-gram 指标（交叉熵和信息损失）与 Transformer 模型下游性能之间存在强线性相关性。这验证了使用 n-gram 作为评估分词器质量的低成本代理模型的有效性。
• 下游影响：
  • Smirk 对 tmQM 数据集的预训练质量和下游性能产生了积极影响。
  • 在 MoleculeNet 任务（以小有机分子为主）中，Smirk 的表现与标准的原子级分词相当。
  • 覆盖范围较差的分词器（SPE/APE）相对于基准线，对预训练和下游性能均产生了负面影响。
  • 研究发现，分子编码（SMILES 与 SELFIES）的选择与分词器的选择相比，其影响可以忽略不计。

意义与主张
本文认为，化学基础模型必须能够编码整个化学空间的广度，以避免掩盖关键特征。作者声称，当前的分词器会无意中掩盖原子级信息（如顺铂中的手性或特定同位素），从而引发潜在的重大信息损失，这种损失不仅是理论上的，而且会影响临床和工业相关的分子。

这项工作的意义在于：

1. 鲁棒性： 证明了开集词表分词器（Smirk/Smirk-GPE）提供了鲁棒的化学空间覆盖，防止了与未知标记相关的的信息丢失。
2. 效率： 确立了 n-gram 模型可以作为评估分词器性能的可靠、低成本代理，从而减轻了超参数调优和模型选择的计算负担。
3. 可解释性： 强调了 Smirk 使研究人员能够直接操纵括号内富含信息的成分，在扩展原子级分词的可解释性优势的同时，消除了无法识别（Out-of-vocabulary）错误的风险。

作者总结道，虽然目前的基准测试（如 MoleculeNet）由于缺乏元素和立体化学的多样性，可能无法充分暴露有限覆盖分词器的缺陷，但转向能够编码整个化学空间的分词器对于可靠的分子基础模型是必要的。他们鼓励社区严格评估基准测试范围，并扩展包含多样化化学特征的数据集。