🐥

いろいろなタンパク質の言語モデルで遊んでみる

2025/12/26に公開

🐊はじめに🐊

タンパク質の言語モデル（protein language models: pLMs）は、自然言語処理 (Natural Language Processing: NLP) の技術をタンパク質のアミノ酸配列に適用して、タンパク質の機能予測やタンパク質同士の結合予測、タンパク質の構造予測など様々なタスクを行える機械学習モデルです。昨年の記事に代表的なタンパク質言語モデルESM-2とLoRAを用いてパラメタ数を抑えながら低コストにファインチューンする手法につして紹介しましたお家のパソコンでESM2のLoRAをする

今年はESM-2のようなアミノ酸配列を対象としたBERT型モデル以外のモデルで遊んでみたいと思います。

タンパク質の言語モデルと課題

昨年の記事で書いたとおり、ESM-2などのタンパク質の言語モデルは大規模なアミノ酸配列のデータを学習させることで、アミノ酸配列に含まれる特徴や進化的な文脈を学習したものでした。しかし、タンパク質の多くは正しく折りたたまれて3次元構造をとることによって機能を発現します。そのため、配列と構造を同時にモデル化することがより高品質な生成・予測などのタスクには必須でした^[1]。しかし従来モデルでは配列基盤と構造基盤が別々であり、最適なjointモデルが存在しませんでした。
そこで今回はBytedance社が作ったタンパク質言語モデルであるDPLM-2を紹介します。DPLM-2には次のような特徴があります。

拡散プロセスによる生成モデル
DPLM-2 は離散拡散（discrete diffusion）プロセスを用いた言語モデルのフレームワークに基づく generative model です。ESM-2ではBERTタイプの言語モデルだったものが離散拡散モデルになっています。
構造情報のトークン化（Lookup-Free Quantizer）
タンパク質の 3D 座標（連続値）は Lookup-Free Quantizer（LFQ）によって離散値トークンに変換され、言語モデルで扱えるようにします。その逆変換（デトークナイズ）もモデル内で可能です。ESM-2では取り入れられていなかった構造情報がトークン化されて取り込まれています。
事前学習のウォームアップ戦略
大規模な配列データで事前学習された DPLM を元に、進化的情報と構造の inductive bias（帰納的バイアス）を取り込みます。
これにより構造データの不足問題を軽減します。
Joint 配列・構造学習
配列と構造を同じモデルで学習・生成することで、二段階生成（生成 → 構造予測）が不要になりました。モデルは配列だけ／構造だけ／その両方の条件付きタスクに対応可能です。特徴量ベクトルを取得するところまでだったESM-2に比べるとたくさんのタスクに対応して設計されていることがわかります。

最後に、ESM-2との比較の簡単なまとめをLLMに作ってもらいました。

観点	ESM-2	DPLM-2
基本思想	配列言語モデル	配列＋構造の多モーダル拡散モデル
学習対象	アミノ酸配列のみ	配列 + 3D 構造
モデル形式	Masked Language Model (MLM)	Diffusion Language Model
構造の扱い	暗黙的（間接的）	明示的にトークン化して学習
表現の性質	進化・配列類似性中心	構造整合性・幾何学的制約を内包
Folding	別モデルが必要	単一モデルで対応
Inverse Folding	別モデルが必要	ネイティブ対応
Co-generation	不可	配列と構造を同時生成
設計用途	注釈・分類・解析	設計・生成・最適化

以上から、DPLM-2ではESM-2で行ったような予測タスク以外にも生成タスクを配列・構造をjointしたモデルで扱えることがわかります。そこで今回の記事ではDPLM-2を使って生成タスクをして遊んでみたいと思います。

Python環境

こちらのレポジトリ(https://github.com/bytedance/dplm/tree/main) に従ってインストールを行います。

# clone project
git clone --recursive https://url/to/this/repo/dplm.git
cd dplm

# create conda virtual environment
env_name=dplm

conda create -n ${env_name} python=3.9 pip
conda activate ${env_name}

# automatically install everything else
bash scripts/install.sh

condaで作られる環境はPython3.9であることに注意が必要です。install.shで依存関係が同時にインストールされるのですが、依存関係の中にopenfoldが入っておりこちらが少し重た目です。少し依存が重いので環境によっては難しいかもしれません。参考までにわたしのPC環境はこんな感じです。

環境	バージョン等
GPU	RTX-3060
Driver	535.183.01
CUDA	12.2

補足で、より最新に近い安定環境であるCUDA 12.8でも環境構築をしてみました。Tipsかんたんにまとめておきます。
重めの依存：Pytorch,torch_XX系列, deepspeed (CUDA12.8版が必要、特にTorch_scatterが自動では入らない)
install.sh内の --index-urlをhttps://download.pytorch.org/whl/cu128にします。
また、別途torch_scatterを

pip install torch-scatter -f https://data.pyg.org/whl/torch-2.8.0+cu128.html

でインストールします。ビルド済のpipのcudaとtorchの対応一覧はこちらにあります：https://data.pyg.org/whl

build環境：pip install時にbuild環境が作られるが、この中にはtorchなどが入っていないためsetup.py内で乱雑にimportされたtorchがないことで怒られてしまいます。
vendor/openfold/pyproject.tomlを作成して下記の内容を書きましょう。

pyproject.toml

[build-system]
requires = [
    "setuptools",
    "wheel",
    "torch"
]
build-backend = "setuptools.build_meta"

推論を行う際にfairscaleが要求されるので

pip install fairscale

でインストールしておきます。
これらの修正で概ね新しいバージョンのcudaで環境構築が整うはずです。

実際に動かしてみる

無事にインストールできたらまずは無事にインポートしてモデルのチェックポイントをロードできるか確認してみます。python上で下記のコマンドを試してみます。

python

from byprot.models.dplm import DiffusionProteinLanguageModel as DPLM
from byprot.models.dplm2 import MultimodalDiffusionProteinLanguageModel as DPLM2
from byprot.models.dplm2 import DPLM2Bit

dplm = DPLM.from_pretrained("airkingbd/dplm_650m").cuda()
dplm2 = DPLM2.from_pretrained("airkingbd/dplm2_650m").cuda()
dplm2_bit = DPLM2Bit.from_pretrained("airkingbd/dplm2_bit_650m").cuda()

エラーが出ずに上の動かしてみたいと思います。テストにまずはdplm2の前進モデルであるdplmで配列生成を行ってみます。

python

from generate_dplm import initialize_generation

input_tokens = initialize_generation(
  length=200,
  num_seqs=5,
  tokenizer=dplm.tokenizer,
  device=next(dplm.parameters()).device
)
samples = dplm.generate(
  input_tokens=input_tokens,
  max_iter=500,
)
print([''.join(seq.split(' ')) for seq in dplm.tokenizer.batch_decode(samples, skip_special_tokens=True)])

数分で生成が終わり、出力結果はこんな感じになりました。

MINYQHLQMINYQHLQMINYQHLQMINYQHLQMINYQHLQMINYQHLQMINYQHLQMINYQHLQ
MINYQHLQMINYQHLQMINYQHLQMINYQHLQMINYQHLQMINYQHLQMINYFTLY

MKLVLASRNAGKLREFRALLADLPLEVVSPEELGLAEPEETGETYGANALLKARAAARASGL
AALADDSGLEVDALQGAPGVYSARFAGAHDDARNNAKLLAALAGVPPERRGGRFVCVLALAE
PDGREQLFEGRVEGRIALAPRGAGGFGYDPIFVPEGETRTFAEMPPEEKNALSHRARALRA
LRPWLARRGWLQEPA

GHTDWVNAVAFSPDGRTLASGSSDETVRLWDTATGRELGPLTGHTDWVNAVAFSPDGRTLAS
AAFSPDGRTLASGSSDETVRLWDTATGRELGPLTGHTDWVNAVAFSPDGRTLASGSSDETV
RLWDTATGRELGPPLTGHTDWVNAVAFSPDGRTLAS

MGRARRRGRRPARRPRGWLRPRARRRRRRPRRRRRVRRGRGRPRALRRRRRSAVSESYTRG
DYRISTDKTRLDVPAIHAFLSANAYWAKGIPLEVVQRSIDNSLCFGVYAGERQVGFARVVT
DRATFAYVADVFVVESERGKGLGKWLIETVVAHPDLQGLRRIALIADPGKEAFYAQFGFGPR
PLPENGMEIYTPSVYR

MEQEAAAPRRRGRPRSAAVERAIHDATWALLAEVGYDRLTFEAVAARAGTSKPVLYRRWRS
KAELVLDAYLDRVGPEVPAPDTGTLRGDLLELLAHVARKLRTPVGTGLIAAAQTDPELARA
FRERFVERRRAATRAVVERAVARGELDPARVTPRVVDVPAALLRHRLLVTGEPVDDAFLAE
VVDGVLLPLAAAPPGAR

一番目の配列をcolabfoldでフォールドするとこんな感じの不思議なシート構造が一番pLDDTが高いモデルとして出力されました。配列をみると極端なリピート配列になっているので、確かにこうなりそうという構造になっています。一方で4番目や5番目はRがかなり多いArgリッチな配列にみえ、核局在や核酸結合タンパクの配列から影響を受けているのかもしれません。全体を通してみると配列に多様性があり拡散モデルのいい面が出ていそうですね！

続いて、DPLM-2の構造・配列同時生成を試してみます。
DPLM-2とより細分化された構造表現を取り入れたDPLM-2 bitで構造・配列の同時生成を行っています。

python

from generate_dplm2 import initialize_generation, save_results

input_tokens = initialize_generation(
  task="co_generation",
  length=200,
  num_seqs=5,
  tokenizer=dplm2.tokenizer,
  device=next(dplm2.parameters()).device
)[0]

samples = dplm2.generate(
  input_tokens=input_tokens,
  max_iter=500,
)
save_results(
    outputs=samples,
    task="co_generation",
    save_dir="./generation-results/dplm2_generation",
    tokenizer=dplm2.tokenizer,
    struct_tokenizer=dplm2.struct_tokenizer, save_pdb=True
)

samples = dplm2_bit.generate(
  input_tokens=input_tokens,
  max_iter=500,
)
save_results(
    outputs=samples,
    task="co_generation",
    save_dir="./generation-results/dplm2_bit_generation",
    tokenizer=dplm2_bit.tokenizer,
    struct_tokenizer=dplm2_bit.struct_tokenizer
)

生成結果の一部を可視化するとこんな感じになりました。DPLMで配列生成だけして構造予測した場合よりもより、ロスマンフォールドのようにβシートとαヘリックスが調和した複雑でタンパクらしい構造が生成できていると思います。

おわりに

今回は構造と配列をjointした拡散言語モデルであるDPLM/DPLM-2で遊んでみました。特にDPLM-2で配列と構造を同時生成したモデルでは従来の言語モデルによる配列生成に比べて高次でタンパクらしい構造が得られているように感じました。今後は条件付けでscaffoldを指定した生成や予測タスクなどでも遊んでみたいと思います。また、時間があったら他の言語モデルでも遊んで記事にできたらいいなと思います。

脚注

アンフィンセンのドグマの立場としてはアミノ酸配列だけから立体構造に関する情報も得られると考えられています。実際に、AlphaFoldなどの立体構造予測モデルはアミノ酸配列からの立体構造予測で大きな成功を収めています。しかし、依然としてアミノ酸配列だけから学習した言語モデルでは生成・予測などのタスクで表現力が不十分で、立体構造をはじめとする他の情報をjointすることでその表現力を高めることができるとされています。 ↩︎

🐊はじめに🐊

タンパク質の言語モデルと課題

Python環境

実際に動かしてみる

おわりに

Discussion