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 ファイルで確認できます。

モデルをダウンロード

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

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 を使って最新のタンパク質言語モデルを作成する簡単な方法