Tokenization for Molecular Foundation Models

技術要約：分子基盤モデルのためのトークン化

問題提起
化学的特性の正確な予測は、エネルギー貯蔵から医薬品探索に至るまで、あらゆる産業において極めて重要である。トランスフォーマー・アーキテクチャは自然言語処理（NLP）に革命をもたらしたが、その分子基盤モデルへの応用においては、トークン化という根本的なボトルネックに直面している。現在の分子モデルは、主に「原子単位（Atom-wise）」のトークン化に依存しており、これは簡略化された分子入力表記法（SMILES）文字列を、固定された語彙を用いて原子レベルのトークンに分割するものである。

このアプローチの主な限界は、OpenSMILES仕様を完全にカバーできないことにある。原子単位のトークナイザーは、括弧で囲まれた原子（同位体、キラル中心、電荷、および明示的な水素数をエンコードするもの）を、単一の分割不可能なトークンとして扱う。これらの特徴のあらゆる組み合わせをカバーするためには、28兆個を超えるトークンの語彙が必要となる。その結果、既存のモデルは3,000個未満の語彙を使用しており、その結果として重大なカバレッジの欠落が生じている。未知の括弧付き原子に遭遇した際、これらの閉じた語彙を持つトークナイザーは、汎用的な未知トークン [UNK] に頼ることになり、キラル性や特定の同位体組成といった重要な化学情報を隠蔽してしまう可能性がある。さらに、既存のオープン語彙型の代替案（BPEベースのスキームなど）は、しばしば曖昧さに苦しみ、異なる化学的実体（例：硫黄と炭素の結合と、元素のスカンジウム）が同一のトークンへと混同されることがある。

手法
著者らは、OpenSMILES仕様で定義されたグリフに基づき、完全なカバレッジの達成と計算効率の維持を両立させる新しいトークン化フレームワークである Smirk、およびその圧縮版である Smirk-GPE を提案している。

1. Smirkトークン化: このスキームは、OpenSMILES仕様で定義されたグリフに基づき、SMILES文字列を2段階の文字レベル分解を行う。

   • 第1段階: 原子への分解（例：OC[C@@H][OH] $\rightarrow$ O C [C@@H] [OH]）。
   • 第2段階: 括弧で囲まれた原子を構成要素となるグリフへと分解（例：[C@@H] $\rightarrow$ [ C @ @ H ]）。
   • このアプローチは、括弧と内部の記号を個別のトークンとして扱うことで、曖昧なシーケンス（例：結合としての Sc と、元素としての [Sc]）を区別する。結果として得られる語彙は165個のトークンに固定されており、学習を必要とせず、あらゆるOpenSMILESエンコードされた分子を [UNK] トークンを使用することなくトークン化できることを保証する。

2. Smirk-GPE (Glyph Pair Encoding): 括弧で囲まれた原子を完全に分解することによって生じるシーケンス長（繁殖性/fertility）の問題に対処するため、著者らはSmirk-GPEを実装した。このバリアントは、グリフトークンに対してByte-Pair Encoding（BPE）に似た圧縮戦略を適用する。標準的なBPEが文字列を結合するのとは異なり、Smirk-GPEはトークンID上でマージルールを学習し、化学的に意味のあるマージ（例：硫黄と炭素のグリフの結合）が原子記号（例：スカンジウム）との曖昧さを生じさせないように設計されている。

3. 評価フレームワーク:

   • 内的指標: 著者らは、繁殖性（平均シーケンス長）、正規化エントロピー（圧縮効率）、トークンの不均衡、および [UNK] トークンの頻度を用いてトークナイザーを評価した。
   • 低コストなプロキシ: すべてのトークナイザーに対してフルサイズのトランスフォーマーモデルを訓練することは計算コストが高いため、著者らはn-gramモデルをプロキシとして利用した。16億個のSMILES文字列を用いてn-gramモデルを訓練し、クロスエントロピー損失と情報損失（KLダイバージェンスによる）を測定することで、ダウンストリームの性能を推定した。
   • 外的検証: n-gramプロキシの妥当性を検証するため、著者らは11種類の異なるトークナイザーと3種類の分子エンコーディングを用いて、18個のエンコーダのみのRoBERTaモデルを（スクラッチから）事前訓練した。これらのモデルは、MoleculeNetおよびtmQMの6つの回帰タスクと7つの分類タスクに対してファインチューニングされた。

主な結果

• カバレッジ: SmirkおよびSmirk-GPEは、評価されたトークナイザーの中で唯一、OpenSMILES仕様の100%のカバレッジを達成し、[UNK] トークンの使用を排除した。対照的に、既存の化学特化型トークナイザー（SPE、APE、および様々なBPEバリアントを含む）は、非無視できる頻度（tmQMデータセットでは最大約50%）で [UNK] トークンを出力する。
• 情報損失: カバレッジが限定的なトークナイザーは、特に遷移金属や立体化学が豊富なデータセット（例：tmQM）において、大幅な情報損失を示す。例えば、MoLFormerのトークナイザーは、未知のトークンによりtmQMにおいて40.3 nats/moleculeの損失を引き起こすが、Smirkバリアントはこの劣化を軽減する。
• 性能相関: 本研究では、n-gram指標（クロスエントロピーと情報損失）と、トランスフォーマーベースのモデルのダウンストリーム性能との間に強い線形相関があることが示された。これにより、トークナイザーの品質を評価するための低コストなプロキシとしてn-gramを用いることが妥当であると検証された。
• ダウンストリームへの影響:
  • Smirkは、tmQMデータセットにおける事前訓練の質とダウンストリーム性能に対して、正の効果を示した。
  • 小さな有機分子が支配的なMoleculeNetのタスクにおいては、Smirkは標準的な原子単位のトークン化と同等の性能を示した。
  • カバレッジの低いトークナイザー（SPE/APE）は、ベースラインと比較して、事前訓練およびダウンストリームの両方の性能に悪影響を及ぼした。
  • 分子エンコーディング（SMILES対SELFIES）の選択は、トークナイザーの選択と比較して、無視できる程度のインパクトしか持たないことが判明した。

意義と主張
本論文は、化学の基盤モデルは、重要な特徴を隠蔽することを避けるために、化学空間の全領域をエンコードしなければならないと主張している。著者らは、現在のトークナイザーが、シスプラチンにおけるキラル性や特定の同位体のような原子レベルの情報を意図せず隠蔽しており、それが単なる理論的な問題ではなく、臨床的または産業的に重要な分子に影響を与える可能性のある重大な情報損失を引き起こしていると主張している。

本研究の意義は以下の点にある：

1. 堅牢性: オープン語彙型のトークナイザー（Smirk/Smirk-GPE）が化学空間の完全なカバレッジを提供し、未知のトークンに伴う情報の損失を防ぐことを実証した。
2. 効率性: n-gramモデルが、トークナイザーの性能を評価するための信頼できる低コストなプロキシとして機能することを確立し、ハイパーパラメータのチューニングやモデル選択の計算負荷を軽減した。
3. 解釈性: Smirkが、原子単位のトークン化の解釈性の利点を拡張しつつ、括弧付き原子の情報豊かな内容を直接操作することを可能にし、未知の語彙エラーのリスクを取り除いたことを強調している。

著者らは、現在のベンチマーク（MoleculeNetなど）は、元素や立体化学の多様性が欠如しているため、限定的なカバレッジを持つトークナイザーの欠陥を完全には露呈できていない可能性があると結論付けている。その上で、信頼できる分子基盤モデルを実現するためには、化学空間の全体をエンコードできるトークナイザーへの移行が必要であると述べている。彼らは、コミュニティに対し、ベンチマークの範囲を厳格に評価し、多様な化学的特徴を含むデータセットを拡張することを推奨している。