ByteDanceのLi Hang氏：科学におけるAIの探究と進歩

近年、自然言語処理、コンピュータービジョン、音声処理など、人工知能のさまざまな分野が、ディープラーニングの強力な力の助けを借りて驚異的な進歩を遂げています。物理学、化学、生物学、医学などの伝統的な科学分野へのディープラーニング技術の応用、いわゆる「科学のためのAI」は、大きな可能性を秘めた新しい学際的なテーマとして徐々に浮上し、幅広い注目を集めています。

ByteDance Researchは、機械学習や量子化学、大規模量子化学コンピューティング、AI医薬品製造などの分野における課題の研究を含むAI for Scienceの研究も行っており、業界とともにこの分野の発展を促進したいと考えています。この記事では、過去 2 年間に私たちが成し遂げた進歩の一部を簡単に紹介します。これは議論の出発点としても機能し、業界とのさらなる交流と協力が進むことを期待しています。

機械学習や量子化学の分野において、私たちが提案したLapNetアルゴリズムは、代表的なFermiNetモデルよりも10倍の学習速度を持ち、計算できる化学系の規模と精度は現時点でこの分野で最大です。

大規模量子化学計算の分野では、周期系の計算に古典・量子ハイブリッドコンピュータ（実際には古典コンピュータ上のシミュレーション）を使用するPeriodic DMETアルゴリズムを開発しました。このアルゴリズムは、約10,000量子ビットを必要とした従来の方法と同じ精度を、わずか20量子ビットで実現します。

AI医薬品製造の分野において、当社が開発したLM-Designモデルは、大量のタンパク質配列データと一定量のタンパク質構造と配列対応データを用いて、タンパク質構造を配列に変換するモデルを学習し、これまでで最も高い精度のタンパク質配列設計を実現しました。

機械学習と量子化学

物理学者ディラックはかつてこう言いました。「物理学のほとんどと化学のすべてにおいて、数学的モデリングに必要な基本法則はすでに完全に明確です。難しいのは、これらの法則を適用することだけです。結果として得られる方程式は、一般的には解くには複雑すぎます。」

量子化学は、量子力学の原理に基づいた化学現象の研究です。重要な問題は、計算手法を使用して分子または周期系（固体など）の電子シュレーディンガー方程式を解き、分子または周期系の基底状態エネルギー、電気極性、およびその他の特性を推測することです。これはいわゆるアブイニシオ問題です。従来の方法には、密度汎関数理論 DFT、結合クラスター CCSD などがあります。計算精度が十分に高くないか、計算規模が十分に大きくありません。

近年、機械学習の手法を使用して第一原理計算問題を解決することが、大きな注目を集める新しい方向性となっています。基本的な考え方は、ディープラーニングの強力な表現と学習機能を活用して、第一原理計算の精度と規模を大幅に向上させることです。方法の 1 つは、NN-VMC (ニューラルネットワークベースの変分モンテカルロ) です。シュレーディンガー方程式の波動関数をニューラルネットワークで近似し、空間内の系内の電子のサンプルをランダムサンプリングで取得することで、シュレーディンガー方程式に基づく系全体のエネルギーを計算できます。エネルギーの上限を最小化し、ニューラルネットワークのパラメーターを最適化し、継続的に反復することで、最終的にほぼ最適なニューラルネットワーク (波動関数) とシステムのおおよその基底状態エネルギー (最小エネルギー) を得ることができます。（注：波動関数の2乗は、空間に現れる電子の確率密度関数です。波動関数を使用すると、空間内の電子のランダムサンプリングを行うことができます。）図1は、NN-VMCの基本原理を示しています。中心的な問題は、ニューラルネットワークと学習アルゴリズムをどのように設計するかです。

図1.NN-VMC法の基本原理

NN-VMCにおける代表的な手法としては、2019年にDeepMindとICLが提案したFermiNetがあります。その後、いくつかの研究機関が新しい方法を提案しました。 2021年以来、ByteDance Researchは北京大学と協力して一連の関連研究を実施し、いくつかの新しい方法を提案してきました。以下にこれらの方法について簡単に紹介します。

NN-VMC+ECPはNN-VMCと擬ポテンシャルECP（有効コアポテンシャル）[1]を組み合わせた手法であり、計算の効率とシステムの規模をさらに向上させることができます。化学システムの特性を計算する場合、多くの場合、原子の外側の軌道にある電子のみに注目する必要があります。原子の内部軌道にある電子の位置エネルギーを定量的に表現することで、必要な計算量を大幅に削減できます。 ECP 技術を NN-VMC に適用し、良好な結果を達成したこの新しい方法を実現しました。

NN-DMCは、ニューラルネットワークと拡散モンテカルロ（DMC）を組み合わせた私たちが提案した別の手法です[2]。 DMC は、システムの基底状態エネルギーの上限を計算せず、虚数時間発展を使用してシステムの基底状態エネルギーを計算するという点で VMC とは異なります。この方法は、FermiNet などの既存の方法と比較して、計算の精度と規模を大幅に向上させることもできます。

最近開発されたLapNetもNN-VMC法[3]であり、ニューラルネットワークの学習時に順方向ラプラシアン演算子を使用するのが特徴です。シュレーディンガー方程式に基づいてシステムのエネルギー上限を計算する過程では、運動エネルギー部分を含むハミルトニアン演算子を計算する必要があります。これまでの方法では、関連するヘッセ行列を計算することで運動エネルギーを計算していましたが、アルゴリズムが非常に複雑で、学習のボトルネックとなっていました。学習の順方向伝播では、LapNet はラプラス演算子の計算を通じて運動エネルギーとハミルトニアン演算子を直接計算し、ヘッセ行列の計算を排除します。これにより、学習の計算効率が大幅に向上します。 FermiNet と比較すると、LapNet は平均で約 10 倍の高速化を実現します。

ECP、DMC、順ラプラスは、3 つの異なる技術的改善 (ポテンシャルエネルギー計算の簡素化、サンプリングの最適化、計算効率の向上) です。これら 3 つのテクノロジを組み合わせると、原理的には計算規模を大幅に拡大できます。これは、私たちが試みているアプローチでもあります。さらに、NN-VMC法を固体のシュレーディンガー方程式[4]、分子系の力場[5]、電気分極計算[6]などの問題にも適用し、NN-VMC法の実用性を実証しました。

図2は、現在のNN-VMC法における代表的な研究の精度と規模を示しています。縦軸は精度、円の大きさは規模を表しています。私たちが提案する LapNet メソッドは、より大きなシステムをより高い精度で計算できます。最大のシステムには 116 個の電子があります。

図2. NN-VMC法のスケールと精度

大規模量子化学計算

シュレーディンガー方程式を直接解いて化学システムの特性（基底状態エネルギーなど）を計算する方法は、規模がまだ限られています。量子埋め込み法はこの問題を解決する効果的な方法であると考えられています。基本的な考え方は、分割統治法と多精度コンピューティングを通じてスケーラビリティを実現することです。代表的な手法としては、密度行列埋め込み理論 (DMET) が挙げられます。システムはいくつかの部分（フラグメント）に分割され、各フラグメントとそれに対応する環境（バス）に対して高精度の計算が実行され、他の部分に対しては低精度の計算が実行されます。そして、各フラグメントは必要に応じて並列に処理されます。最後に、高精度の計算結果を組み合わせて継続的に反復し、元のシステムに近づきます。これにより、計算可能なシステムの規模を大幅に拡大できます。

図3 DMET法の直感的な説明

図3はDMET法のプロセスを示しています。まず、元のシステムを分割してフラグメントのセットを取得します。図の黄色い部分が、私たちが関心のあるフラグメント、たとえば 2 つの原子であると仮定します。画像の青い部分には環境やその他の部分が含まれています。関心のあるフラグメントとその環境は、CCSD などの高精度の方法を使用して計算され、他の部分は Hartree-Fock 法などの低精度の方法を使用して計算されます。すべてのフラグメントに対して同じ並列処理を実行します。

具体的なアルゴリズムは以下のとおりです。まず、低精度解法を実行して、全体（対象フラグメント、環境、その他の部分）の縮小密度行列を取得します。この低精度解法にはパラメータが含まれます。次に、この行列のフラグメントとその環境に対して特異値分解を実行し、射影演算子 P を構築し（射影演算子はフラグメントとその環境にのみ焦点を当てます）、射影演算子を使用して低次元システム（イメージ）を構築し、高精度で解きます。その後、すべてのフラグメントの計算結果がシステム全体の近似値として結合されます。最後に、反復処理を通じてパラメータが調整され、元の低精度のソリューションが徐々に結合された高精度のソリューション（L2 ノルムの意味で）に近づき、最終結果が得られます。

2 つのまったく異なる計算パラダイムに基づいて、DMET とその SIE のバリアントを実装し、大規模な量子化学システムの計算を実行します。 1つは古典的なコンピュータを使用することであり、もう1つは量子コンピュータを使用することです。本稿では主に後者の関連研究（前者の研究計画については今後大きな進展があった際に紹介する予定）を紹介し、量子計算化学とも呼ばれる。量子コンピュータへの実装も検討していますが、量子化学の発展に貢献することを目指して、シミュレーションは古典コンピュータ上でのみ実行しています。 DMET 方式は、コンピューティングパラダイムに関係なく、大規模なコンピューティングを可能にするため、私たちは DMET を実装して、コンピューティング可能なシステムを数桁増やすために懸命に取り組んでいます。

物理学者ファインマンはかつてこう言いました。「自然は古典力学ではない。自然をシミュレートしたいなら、量子力学を使うほうがよい。」量子コンピューティング技術開発の原動力は、量子レベルのコンピューティングデバイスを使用して量子現象をシミュレートすることです。言い換えれば、量子化学は量子コンピューティングに最も適した応用分野の 1 つです。

図4. 周期的DMET法の概略図

私たちは、量子と古典のハイブリッドコンピュータと DMET の特性を組み合わせた 2 つの量子計算化学手法を開発し、計算システムの精度と規模を大幅に向上させました。基本的な考え方は、量子コンピュータを使用して DMET の高精度計算部分を実装し、古典コンピュータを使用して DMET の低精度計算部分を実装することです。 DMET-ESVQEは分子系[7]を計算し、Periodic DMETは周期系[8]を計算します。前者は、以前の方法では 144 量子ビットが必要だった計算を、わずか 16 量子ビットで実現します。後者は、以前の方法ではほぼ 10,000 量子ビットが必要だった計算を、わずか 20 量子ビットで実現します。

図 4 は、ハイブリッドコンピュータに実装された DMET 周期法を示しています。入力は結晶であり、出力はシステムのエネルギーです。まず、システムを分割し、各フラグメントを並列に計算します。関心のあるフラグメントとその環境は、量子コンピュータ上の U-CCSD を使用して解決されます。残りの部分は、古典的なコンピュータ上でハートリー・フォック法を使用して解決されます。

VMCやDMCを含む量子モンテカルロ法は、量子化学において最も効果的な一連のアルゴリズムである[9]。また、量子コンピューティングと量子モンテカルロ法を組み合わせた新しい方法も提案しました。この方法は、量子化学に対する量子コンピューティングのいくつかの利点を実証することができます。具体的には、量子コンピューティングは量子モンテカルロ法における符号問題を部分的に解決できます。

AI製薬

AI 技術を活用した医薬品の発見支援は、業界で広く受け入れられている新しいパラダイムとなっています。近年多くの研究が行われており、いくつかの技術が実際のシナリオに適用されています。低分子医薬品、高分子医薬品（抗体医薬品）をはじめ、AI技術を基盤とした医薬品設計の研究開発を行っています。

小分子医薬品の設計プロセスには、タンパク質標的の発見、小分子医薬品候補の生成、候補と標的間の親和性の決定、候補の医薬品化可能性、および非毒性が含まれます。現在、これらの医薬品開発ステップを可能にする AI テクノロジーが利用可能です。我々は機械学習に基づいて低分子医薬品候補を生成する手法を開発しました。MARSはスコアリング関数を用いて候補を自動生成する方法[10]であり、DESERTはターゲットの形状に基づいて候補を自動生成する方法です[11]。

MARS はシード分子から開始し、最適な小分子薬剤候補が最終的に得られるまで分子を継続的に編集します。生成プロセスではマルコフ連鎖モンテカルロ法 (MCMC) が使用され、その安定分布は複数のスコアリング関数で構成された確率分布です。スコアリング関数は、小分子薬剤候補の親和性、薬剤化可能性、および非毒性を表します。提案された分布は、グラフニューラルネットワーク (MPNN) に基づいて編集する前と後の小分子医薬品候補の条件付き確率分布を表します。グラフニューラルネットワークは、ノードが原子、エッジが化学結合である小分子化合物の分子式を表します。小分子の編集には、新しいノードの追加と既存のノードの削除が含まれます。グラフニューラルネットワークは、小分子に対する可能な編集操作（追加または削除）を予測することができ、そのパラメータは学習を通じて取得されます。 MARS では、新しく多様な低分子医薬品候補を生成するために、スコアリング機能と低分子医薬品のデータベース (分子式) のみが必要です。現在、MARS は実用的な小分子医薬品の設計作業に使用されています。

DESERT は 2 つのステップで小分子薬剤候補を生成します。スケッチ: タンパク質ターゲットのポケット形状に相補的な薬物候補の形状をサンプリングします。生成: 薬物候補の形状に基づいて薬物候補の分子式を自動的に生成します。図5はこのプロセスを示しています。

図5. DESERT: 低分子医薬品候補の自動生成

候補とターゲットの組み合わせに必要な条件は、2 つの形状がうまくドッキングできることです。スケッチ段階では、タンパク質ターゲットの形状に基づいたヒューリスティックなアプローチを使用して候補形状が生成されます。生成段階では、以前に学習した形状から分子への生成モデル Shape2Mol を使用して、形状に基づいて分子式を自動的に生成します。この生成モデルは、分子ライブラリ内の多数の薬物の分子式と形状を使用して学習できます。図 6 に示すように、Shape2Mol では、エンコーダーが分子の 3D 形状をエンコードして中間表現を生成し、デコーダーが中間表現に基づいて分子式を生成します。 3D 形状はボクセルを使用して表され、分子式は記号のシーケンスを使用して表されます。 DESERT は、2022 年に低分子医薬品候補を生成するための最良の複合手法です。

図6. 形状から分子への生成モデルShape2Molの概略図

最近、私たちは高分子薬物設計、より一般的にはタンパク質設計に重点を置いています。タンパク質設計には抗体医薬品設計、ペプチド医薬品設計などが含まれます。タンパク質の配列（アミノ酸配列）がわかればその構造を予測でき、タンパク質の構造がわかればその機能を予測できます。これが有名な AlphaFold の機能です。タンパク質の設計は逆のプロセスとして考えることができます。一般的には、機能に基づいて対応するタンパク質構造が決定され、次にタンパク質構造に基づいて対応するタンパク質配列が決定されます。タンパク質構造からタンパク質配列を生成するためのモデル LM-Design を開発しました。

LM-Design の入力はタンパク質構造であり、出力は対応するタンパク質配列です。 LM-Design は、構造エンコーダーとシーケンスデコーダーで構成されています。このうち、構造エンコーダーはタンパク質構造を表現するトレーニング済みのグラフニューラルネットワークであり、シーケンスデコーダーは、BERT/Transformer Encoder（双方向セルフアテンションを使用）に似た、事前トレーニング済みの大規模タンパク質言語モデル（Protein Language Model）に基づいており、最後の層に構造アダプターが挿入されています。構造アダプタのパラメータを学習する必要があります。図 7 は LM-Design のモデルアーキテクチャを示しています。

図7. タンパク質構造から配列を生成するモデルLM-Designのアーキテクチャ

LM-Design の学習と予測はマスク言語モデリングであり、BERT モデルのトレーニングに似ています。その目的は、マスクされたシーケンス内のシンボル (アミノ酸) を複数回復元することです。つまり、トレーニングされたタンパク質言語モデルの情報と現在のタンパク質構造情報に基づいて、タンパク質配列が複数回書き換えられます。 LM-Design は、グローバルな配列情報に基づいて少数のシンボル (アミノ酸) を書き換えるため、長距離のタンパク質依存性を適切に表現および予測できます。注意: タンパク質が折り畳まれると、配列が離れているアミノ酸の構造が近くなることがあります。

現実には、配列されたタンパク質のデータは大量にあるものの、配列とタンパク質構造のアライメントに関するデータはわずかしかありません。 LM-Design の利点の 1 つは、膨大な量のタンパク質配列データを活用し、タンパク質配列間の進化関係を完全に学習して活用し、タンパク質構造から配列生成までの予測精度を大幅に向上できることです。さらに、事前トレーニング済みのタンパク質モデルのサイズを大きくすると、精度がさらに向上することがわかりました。図 8 に示すように、LM-Design は現在、最も優れたタンパク質配列生成モデルです。図中の円の大きさは、モデルのパラメータの数を表しています。

図8. タンパク質配列生成法の精度