BioNeMo Framework で切り拓く創薬研究の未来

はじめに

人工知能 (AI) の急速な進化は、さまざまな科学分野を変革しており、その中でも TechBio は、生物学と計算科学の融合 により急成長している分野です。この分野の革新を推進する強力なフレームワークの 1 つが BioNeMo Framework です。

BioNeMo Framework は、大規模トランスフォーマー モデルを生体分子データに適用するために設計された AI プラットフォームであり、生物学的配列の解析、分子相互作用のモデリング、新規生体分子の生成を可能にします。これにより、創薬、タンパク質工学、ゲノム研究における発見が促進されます。

生物学研究 に AI を適用することで、研究者の データ解析、モデリング、予測 のアプローチが根本的に変わりました。従来の手法では計算コストが高く、シミュレーションや実験ベースのワークフローに依存していましたが、BioNeMo Framework は大規模データセットで訓練されたディープラーニング モデルを活用することで、以下のような利点を提供します。

• 生物学的配列のパターン抽出を高速化
• 分子構造と特性の精度の高い予測
• 生成 AI による生体分子設計の最適化
• タンパク質工学を加速し、突然変異や機能変化をシミュレート
• マルチオミクス統合を強化し、複雑な生物学的システムの理解を促進

これらの能力は、個別化医療、疾患モデリング、合成生物学などの分野での画期的な進歩を可能にし、生命科学の研究と応用の在り方を変革します。本記事では BioNeMo Framework (以下は BioNeMo) の全体像と概要を紹介します。

BioNeMo の主要な機能

図 1 で示すように、BioNeMo は、生体分子データを活用し、研究開発者が簡単にカスタマイズ可能な AI モデルを構築できるように設計されています。

• モジュール型の設計
  研究者は、自分の研究に最適な AI モデルを自由に組み立てることが可能です。例えば、分子の相互作用を予測するモデルや、化合物の最適化を行うモデルなど、目的に応じて柔軟にカスタマイズできです。
• AI モデルの学習/最適化
  研究者は、既存のモデルを再学習したり、新しいデータに適応させることで、より高精度な分子設計や、創薬に適した AI モデルを開発できます。これにより、特定の課題に最適化されたモデルの構築が可能になります。
• 予測、評価、実用化までのスムーズなワークフロー
  モデルを用いて分子の特性を予測し、その性能を評価した上で、最終的に実際の創薬プロセスに活用できる AI モデルを提供 できます。これにより、研究から実用化までの流れをスムーズに進めることが可能です。

BioNeMo の高速な処理性能

図 2 に示すように、GPU の数を増やすことで、BioNeMo は GitHub などで提供される一般的なオープン ソースの方法よりも圧倒的に高速に動作 します。グラフの緑の線は、BioNeMo のパフォーマンス向上を示しており、GPU を増やすことで単純に処理速度が向上するだけではく、それ以上のスピードアップが可能であることがわかります。

BioNeMo 1.x と BioNeMo 2.x の位置づけ

NVIDIA が主導して開発してきた BioNeMo 1.x は、生命科学分野における AI 活用を加速するための重要なプラットフォームです。しかし、AI 開発の進歩は非常に速く、NVIDIA だけで最新の AI モデルを適時提供し続けるには限界があるという課題がありました。そこで、より多くの研究者が BioNeMo の開発に貢献できるようにするため、NVIDIA は BioNeMo 2.x を GitHub でオープン ソース化し、誰でも自由に開発および拡張できるように移行しています。

BioNeMo 2.x への移行は、重点が「モデル中心の機能開発」から、「次世代AIモデル（例：LLMやStripedHyenaアーキテクチャ）を構築・スケーリングするためのツール開発」へと開発の軸足を移すことを意味しています。

本記事では、現時点でより豊富な事前学習モデルとサンプルコードを提供し、実用性が高いBioNeMo 1.x を前提に説明していきます。しかし、今後の開発は BioNeMo 2.x を中心に行われ、代表的なアーキテクチャや新機能は主にこのバージョンで導入されていく予定です。

BioNeMo 1.x の仕組み

BioNeMo 1.x の事前学習モデル

BioNeMo 1.x には、多様な事前学習 AI モデルが搭載されており、生体分子研究におけるさまざまなタスクをサポートします。

タンパク質配列モデル

• ESM-1nv & ESM-2nv – タンパク質機能予測と配列解析のためのトランスフォーマー モデル。
• ProtT5nv – 安定性予測、構造アノテーション、変異影響解析のためのタンパク質言語モデル。
• OpenFold – AlphaFold に基づいたディープラーニング モデルで、タンパク質の 3D 構造予測 を最適化。
• EquiDock – タンパク質-タンパク質ドッキングの AI モデル。

分子相互作用 & 創薬モデル

• DiffDock – AI 駆動の分子ドッキング モデルで、小分子-タンパク質相互作用を解析。
• DSMBind – 創薬ターゲットの結合親和性予測を行うディープラーニング モデル。
• MegaMolBART – 分子設計、特性予測、リード最適化のための生成 AI モデル。
• MolMIM – 分子表現学習とドラッグライク化合物設計のための自己教師あり学習モデル。

ゲノミクス & DNA モデル

• DNABERT – DNA 配列解析、突然変異予測、ゲノム アノテーションに特化したトランスフォーマー モデル。
• Geneformer – マルチオミクス データで学習されたゲノムトランスフォーマー で、遺伝子発現モデリングと変異解析をサポート。

サンプル コードとデータの提供

BioNeMo の Docker コンテナーに入ると、最初に表示されるディレクトリは以下のような構成になっています。このディレクトリには、BioNeMo のコア フレームワーク、設定ファイル、テスト用コード、ライセンス情報などが含まれています。

├── README.md
├── artifact_paths.yaml
├── bionemo
├── bionemo.egg-info
├── ci
├── download_artifacts.py
├── examples
├── hydra_plugins
├── icla.pdf
├── launch.sh
├── license.txt
├── pyproject.toml
├── setup
├── setup.py
├── tests
├── third_party.txt
└── tokenizers

各ファイルおよびディレクトリの説明

その中でも、examples/ ディレクトリには、BioNeMo の主要なモデルや学習タスクのサンプル コードが格納されています。

├── conf
│   └── base_infer_config.yaml
├── dna
│   ├── dnabert
│   ├── enformer
│   └── nt_tokenizer.py
├── molecule
│   ├── diffdock
│   ├── dsmbind
│   ├── megamolbart
│   └── molmim
├── protein
│   ├── downstream
│   ├── equidock
│   ├── esm1nv
│   ├── esm2nv
│   ├── openfold
│   └── prott5nv
├── singlecell
│   └── geneformer
└── tests
    ├── README.md
    ├── conf
    ├── expected_results
    ├── test_connector.py
    ├── test_controlled_generation_integration.py
    ├── test_diffdock_preprocessing.py
    ├── test_flip_preprocessing.py
    ├── test_model_pretrain_and_downstream.py
    ├── test_physchem_preprocessing.py
    ├── test_scripts_e2e.py
    ├── test_stop_and_go.py
    ├── test_uniref_preprocessing.py
    ├── test_uspto50k_preprocess.py
    └── test_zinc_preprocessing.py

各ディレクトリの役割

• conf/: 設定ファイルを格納。モデルの推論やトレーニングに使用される YAML 形式の設定ファイルが含まれる。
• dna/: DNA 関連の AI モデルが格納されている。
• molecule/: 分子関連のモデルが格納され、創薬や分子特性予測などに使用される。
• protein/: タンパク質解析関連の AI モデルが含まれ、タンパク質機能予測/構造予測に特化。
• singlecell/: シングル セル RNA 解析のための Geneformer モデルを格納。
• tests/: BioNeMo のテスト コード。モデルの動作確認や検証に使用される。

例えば、molecule/molmim ディレクトリには、分子特性予測を行う MolMIM モデルのコードと設定ファイルが格納されています。MolMIM は、自己教師あり学習を利用して分子の表現を学習し、化学特性の予測や創薬応用に活用されます。

├── conf
│   ├── infer.yaml
│   ├── molmim_70m_24_3.yaml
│   ├── pretrain_base_canonicalized.yaml
│   ├── pretrain_base_equivariant.yaml
│   ├── pretrain_base_invariant.yaml
│   ├── pretrain_large_canonicalized.yaml
│   ├── pretrain_large_equivariant.yaml
│   ├── pretrain_large_invariant.yaml
│   ├── pretrain_small_canonicalized.yaml
│   ├── pretrain_small_canonicalized_logv.yaml
│   ├── pretrain_small_equivariant.yaml
│   ├── pretrain_small_invariant.yaml
│   ├── pretrain_xsmall_canonicalized.yaml
│   └── pretrain_xsmall_equivariant.yaml
├── dataset
│   └── ZINC-downloader-sample.txt
├── pretrain.py
└── scripts
    └── pretrain_quick.sh

各ファイルの役割

conf/:

• infer.yaml → MolMIM モデルを推論する際の設定。
• molmim_70m_24_3.yaml → パラメーター数 70M のモデルの学習設定。
• pretrain_*.yaml → 事前学習用の設定。分子の表現学習に適した各種モデルのパラメーター設定が含まれる。

dataset/:

• ZINC-downloader-sample.txt → ZINC データセットのサンプル (化合物の SMILES 表記が含まれる)。

pretrain.py:

• MolMIM モデルの事前学習を行うスクリプト。CUDA を活用し、大規模な分子データセットで学習する。

scripts/:

• pretrain_quick.sh → 事前学習を素早く実行するためのスクリプト。

データの前処理

• BioNeMo 今は以下のような前処理を実現できます。
• タンパク質、DNA、RNA の配列をトークン化。
• 学習、検証、テストデータを分割。
• トークンを高次元特徴空間に埋め込み、相互関係を学習。

モデルの学習

• 大規模トランスフォーマー モデルを活用し、配列の関係性を学習。
• 転移学習をサポートし、特定のバイオ医薬応用に適用可能。

推論

• 学習済みモデルを活用し、配列生成、創薬ターゲット予測、特性モデリングを実施。
• NVIDIA GPU に最適化され、高速でスケーラブルな推論が可能。

BioNeMo 1.10.1 の導入手順

BioNeMo の実行環境要件

ハードウェア要件

• GPU: NVIDIA Ampere 世代 (A100) 以降の GPU が推奨されています。特に、bfloat16 精度を使用する場合は、Ampere 世代以上の GPU が必要です。
• GPU サポート リスト:
  • 完全サポート: H100、A100、RTX A6000 など Ampere 世代以降の GPU
  • 部分サポート: V100、T4、Quadro RTX 8000、GeForce RTX 2080 Ti
  • 既知の問題あり: Tesla K80

ソフトウェア要件

• オペレーティング システム: x86 アーキテクチャの Linux システム
• 必要なソフトウェア:
  • GPU ドライバー: バージョン 535 以上
  • Docker: GPU サポートを有効にした Docker Engine 19.03 以上
  • NVIDIA Container Toolkit: Docker が GPU にアクセスできるようにするために必要

これらの要件を満たすことで、BioNeMo Framework を効果的に活用できます。詳細や最新情報については、公式ドキュメントをご参照ください。

BioNeMo を Docker コンテナーで実行

• NGC に登録
• NGC CLI を自分の環境に合わせてダウンロードし、インストールして、API キーを設定
• ターミナルで下記のように nvcr.io にログイン

docker login nvcr.io
Username: $oauthtoken
Password: <API キーを入力してください>

BioNeMo の現時点の最新コンテナー v1.10.1 を Pull

docker pull nvcr.io/nvidia/clara/bionemo-framework:1.10.1

BioNeMo のコンテナー v1.10.1 を実行

docker run -it --rm --gpus all -v /home/bionemo/:/workspace/bionemo/mydata nvcr.io/nvidia/clara/bionemo-framework:1.10.1 bash

-v: ホスト マシンのディレクトリをコンテナー内にマウントするためのものです。上記のコマンドはホストマシンの /home/bionemo/ を、コンテナー内の /workspace/bionemo/mydata にマウントしています。この設定により、コンテナー内からホストのデータにアクセスすることができるようになり、コンテナー内で作成または変更したファイルがホスト マシンにも反映され、コンテナーを削除してもデータが保持されるという利点があります。

では、次に MolMIM を例として使い方を説明します。

ダウンロード可能なモデル リストの確認

モデルリストは /workspace/bionemo/artifact_paths.yaml ファイルで確認できます。

# “artifact_paths.yaml”

…
    molmim_70m_24_3:
        # https://wandb.ai/clara-discovery/molmim_convergence_prod_March2024/runs/cOLZ4Eme
        ngc: "nvidia/clara/molmim:1.3"
        pbss: "s3://bionemo-ci/models/molmim_70m_24_3.nemo"
        symlink:
            source: "../../molmim_70m_24_3.nemo"
            target: "molecule/molmim/molmim_70m_24_3.nemo"
        md5sum: ede3edcfc90711eb915591b28b83eca2 # pragma: allowlist secret
…

モデルをダウンロード

ファインチューニングや推論のために、学習済みモデルのダウンロードする必要があります。

BioNeMo のコンテナー内には NGC CLI がインストールされていないため、モデルをダウンロードするには、コンテナー内で ngc を再インストールし、ngc config を設定する必要があります。インストール方法については、 NGC CLI をご参考ください。

そして、下記のように ngc config set をしてください。

ngc config set 
<順番にAPI キー、CLI output format, org, team, aceを入力してください。API キーはNGCのAPIキーを入力してください。Org は ’no-org’ 以外のものを選択して、その他は「Enter」でデフォルト値を入力すれば大丈夫です>

python download_artifacts.py --models molmim_70m_24_3 --model_dir /workspace/bionemo/models

ダウンロードしたモデルは /workspace/bionemo/models/ に保存されています。

# ダウンロードしたモデルの保存場所
/workspace/bionemo/models/
├── molecule/
└── molmim_70m_24_3.nemo

モデルの事前学習

事前学習 (Pretraining) とは、大量のデータを用いて汎用的な特徴を学習し、その後、特定のタスクに適応させるための土台を作るプロセスです。BioNeMo のようなバイオ分子 AI モデルにおいて、事前学習は特に重要な役割を果たします。

• データ不足を補うため
  生命科学の分野では、高品質なラベル付きデータ (正解データ) が限られていることが多く、新しい AI モデルをゼロから学習させることは困難です。事前学習を行うことで、以下のような大量の無ラベル データ (未分類のデータ) を活用できます。
• ゼロから学習するよりも効率的
  AI モデルをゼロから学習させるには、膨大な計算資源と時間が必要になります。しかし、事前学習済みのモデルを活用することで、すでに学習した知識を活かしながら、新しいデータに素早く適応できます。
• 汎用的な特徴を学習し、応用範囲を広げる
  事前学習されたモデルは、特定のタスクに依存しない汎用的な特徴を学習しています。事前学習されたモデルをベースにすることで、新しいタスクに対して素早く適応し、より少ないデータで高精度な予測が可能になります。

データの前処理

モデルの学習や推論を行う前に、データを適切なフォーマットに変換し、学習に適した形へ処理する必要があります。

cd /workspace/bionemo
python examples/molecule/molmim/pretrain.py\
 --config-path conf\
 --config-name molmim_70m_24_3\
 ++do_training=False\
 ++trainer.devices=1 \
 ++trainer.num_nodes=1 \
 ++model.data.dataset_path=/workspace/bionemo/examples/molecule/molmim/data \
 ++model.data.links_file=/workspace/bionemo/examples/molecule/megamolbart/dataset/ZINC-downloader-sample.txt \
 ++exp_manager.create_wandb_logger=False \
 ++exp_manager.exp_dir=/workspace/bionemo/results/logs

前処理はサンプル データのダウンロード、配列のトークン、学習、検証、テスト データを分割するなどをします。実行が完了したら、下記のログが表示されます。

[NeMo I 2025-02-05 02:48:02 preprocess:253] Download complete.
[NeMo I 2025-02-05 02:48:02 preprocess:255] Now splitting the data into train, val, and test sets with 10050000, 100, 50000 samples per file, respectively.
[NeMo I 2025-02-05 02:48:02 preprocess:199] The data has been be split into 1 files.
[NeMo I 2025-02-05 02:48:02 preprocess:152] Splitting file /workspace/bionemo/examples/molecule/molmim/data/split_data/x000 into train, validation, and test data

モデルの事前学習

データの前処理が完了したら、モデルの事前学習を行います。BioNeMo では、事前学習用のスクリプトが用意されており、データセットを指定することで、簡単に学習を開始できます。

モデルが正しく動作するかどうかを確認するため、サンプル データを 1 つだけ使用して動作確認を行うことができます。

export TRAIN_FILE_RANGE="x000"
export TEST_FILE_RANGE="x000"
export VAL_FILE_RANGE="x000"

事前学習を実行するためには、適切な設定ファイル (YAML 形式) を使用する必要があります。例えば、/workspace/bionemo/examples/molecule/molmim/conf/molmim_70m_24_3.yaml には、以下のような学習設定が記述されています。

model:
  encoder:
    num_layers: 6
    hidden_size: 512
    ffn_hidden_size: ${multiply:${.hidden_size}, 4}
    num_attention_heads: 8
    init_method_std: 0.02
    hidden_dropout: 0.1
    attention_dropout: 0.1
    ffn_dropout: 0.0

...

do_training: true
do_testing: false
seed: 42
restore_from_path: null
trainer:
  devices: 8
  num_nodes: 4
  precision: 32
  accelerator: gpu
  max_epochs: null
  max_steps: 1000000

...

この設定では、以下のようなパラメーターを指定しています。

1. model セクションではモデルのアーキテクチャに関する設定が定義されています。
2. 学習の制御 (do_training, do_testing, seed)
   • do_training: true に設定すると、事前学習を実行する。
   • do_testing: true にすると、学習後にテストを実施する。ここではテストは無効 (false)。
   • seed: 乱数のシード値 (42)。この値を固定することで、同じデータと設定で null の場合、新しく学習を開始する。途中から再開したい場合は、保存済みのモデルのパスを指定する。再現性のある学習を行うことが可能。
   • restore_from_path: null の場合、新しく学習を開始する。途中から再開したい場合は、保存済みのモデルのパスを指定する。
3. トレーニング環境 (trainer セクション)
   • devices: 各ノードで使用する GPU の数。
   • num_nodes: 分散学習のための計算ノード数
   • precision: 計算精度 (32-bit)。これを 16 にすると、メモリ消費を抑えつつ高速化できる (FP16学習)。
   • accelerator: gpu を指定すると、学習時に GPU を使用。
   • max_epochs: エポック数 (ここでは null で指定されていない)。代わりに max_steps で制御。
   • max_steps: 最大学習ステップ数 (100 万ステップ)。

また、設定を変更したい場合は、設定ファイル (YAML 形式) のパラメーターを直接修正することもできますが、コマンドライン上で ++ 記号を使って設定を上書きする方法もあります。設定ファイルを準備したら、pretrain.py を使用して事前学習を開始します。

cd /workspace/bionemo
python examples/molecule/molmim/pretrain.py\
 --config-path conf\
 --config-name molmim_70m_24_3\
 ++trainer.devices=1 \
 ++trainer.num_nodes=1 \
 ++trainer.max_steps=5000 \
 ++trainer.accumulate_grad_batches=1 \
 ++trainer.val_check_interval=5 \
 ++trainer.limit_val_batches=1.0 \
 ++trainer.precision=32 \
 ++model.micro_batch_size=128 \
 ++model.global_batch_size=128 \
 ++model.dwnstr_task_validation.enabled=False \
 ++model.data.dataset_path=examples/molecule/molmim/data \
 ++model.data.dataset.train=$TRAIN_FILE_RANGE \
 ++model.data.dataset.val=$VAL_FILE_RANGE \
 ++model.data.dataset.test=$TEST_FILE_RANGE \
 ++model.data.index_mapping_dir=/results/data_index/ \
 ++model.seq_length=128 \
 ++exp_manager.exp_dir=/results/logs/ \
 ++exp_manager.create_wandb_logger=False

学習中は、以下のようなログが表示されます。

[NeMo W 2025-02-05 02:55:12 nemo_logging:349] /usr/local/lib/python3.10/dist-packages/nemo/collections/nlp/modules/common/megatron/clip_grads.py:86: UserWarning: The torch.cuda.*DtypeTensor constructors are no longer recommended. It's best to use methods such as torch.tensor(data, dtype=*, device='cuda') to create tensors. (Triggered internally at /opt/pytorch/pytorch/torch/csrc/tensor/python_tensor.cpp:78.)
      dummy_overflow_buf = torch.cuda.IntTensor([0])

Epoch 0:   4%| | 177/5000 [00:24<11:14,  7.15it/s, v_num=, reduced_train_hiddens_mim_loss=0.525, reduced_train_hiddens_mim_weight_loss=0.175, reduced_train_hiddens_loss=0.0919, red

事前学習が完了すると、学習済みモデルは /workspace/bionemo/results/ ディレクトリに保存されます。このディレクトリには、モデルのチェックポイントや学習ログが格納されています。なお、max_steps が少ない場合、チェックポイントが作成されない可能性があります。

最後に

BioNeMo は 創薬研究を変革する AI フレームワークであり、深層学習と生体分子科学を融合させることで、次世代の研究を加速します。本記事では以下のポイントを紹介しました。

• BioNeMo の概要と主要な機能
• BioNeMo 1.x と BioNeMo 2.x の主な違い
• BioNeMo 1.x の仕組み
• BioNeMo 1.10.1 の導入手順

BioNeMo は、柔軟でスケーラブルな AI モデル開発を可能にし、研究のスピードと精度を飛躍的に向上させます。今後も、生命科学分野での AI 活用がますます進展し、新たなブレイクスルーを生み出すことが期待されます。

関連情報

• BioNeMo 1.x ドキュメント
• BioNeMo 2.x ドキュメント
• 技術ブログ: BioNeMo Framework を使って最新のタンパク質言語モデルを作成する簡単な方法