结构语言模型（SLM）用于蛋白质构象生成¶

概述¶

结构语言模型（Structure Language Models, SLM） 是一个 SE(3) 不变性学习框架，专门用于高效的蛋白质构象生成。该框架由 Jiarui Lu 等人在 ICLR 2025 上发表，核心思想是将蛋白质结构编码为紧凑的离散潜在空间，然后通过条件语言模型学习序列特异性构象分布。

SLM 框架的关键创新在于： 1. 离散变分自编码器（VQ-VAE）：将连续的三维蛋白质结构压缩为离散的 token 序列 2. 条件语言建模：在离散潜在空间上学习序列到构象的映射 3. ESMDiff：基于 ESM3 的掩码扩散微调方法，实现最先进的离散扩散性能

该框架开源了完整的代码库，提供了从数据处理、模型训练到推理评估的全流程工具。代码基于 PyTorch 和 Lightning 框架，使用 Hydra 进行配置管理。

核心知识点¶

1. SLM 框架架构¶

SLM 框架包含两个主要阶段：

编码阶段： - 输入：蛋白质的 3D 结构坐标（Cα 原子位置） - 使用 SE(3) 不变性编码器将结构映射到潜在空间 - 通过 VQ-VAE 将连续表示量化为离散 token - 输出：离散 token 序列（表示结构信息）

生成阶段： - 输入：蛋白质序列 + 部分结构 token - 使用条件语言模型（如 ESM3）预测完整的结构 token 序列 - 通过 VQ-VAE 解码器将 token 序列重建为 3D 结构

2. ESMDiff 模型¶

ESMDiff 是 SLM 框架中的核心生成模型，它基于 ESM3 进行掩码扩散微调：

基础模型：ESM3（EvolutionaryScale 发布的蛋白质语言模型）
微调方法：掩码扩散（Masked Diffusion）
采样策略：迭代去噪，逐步从掩码状态恢复完整结构 token
优势：结合了预训练语言模型的知识和扩散模型的生成能力

3. SE(3) 不变性¶

SLM 框架确保所有操作对旋转和平移保持不变： - 编码器使用 SE(3) 不变的特征提取 - 潜在空间表示不依赖于坐标系的绝对位置 - 生成的构象具有物理合理性

4. 离散潜在空间¶

VQ-VAE 将连续结构编码为离散 token 的优势： - 减少计算复杂度 - 便于与语言模型结合 - 支持条件生成和可控采样

代码实操¶

环境安装¶

# 创建 Python 3.10 环境
conda create -n py310 python=3.10 -y
conda activate py310

# 安装依赖
pip install -r requirements.txt

# 本地安装 SLM 包
pip install -e .

配置 ESM3 权重¶

# 从 HuggingFace 下载 ESM3 权重（需接受非商业许可）
from huggingface_hub import login
from esm.models.esm3 import ESM3
from esm.sdk.api import ESM3InferenceClient, ESMProtein, GenerationConfig

# 登录 HuggingFace（需要创建 API key，权限设为 "Read"）
login()

推理示例¶

# 使用 ESMDiff 进行迭代采样
# 参数说明：
#   --input: 输入蛋白质结构文件目录
#   --output: 输出目录
#   --num_steps: 扩散步数（默认25）
#   --num_samples: 采样数量（默认100）
#   --ckpt: 模型权重文件路径

# 基本采样
python slm/sample_esmdiff.py \
    --input data/targets/bpti \
    --output outputs/inference_esmdiff \
    --num_steps 25 \
    --num_samples 100 \
    --ckpt data/ckpt/release_v0.pt

# 局部修复（inpainting）：指定要重新生成的残基位置
python slm/sample_esmdiff.py \
    --input data/targets/bpti \
    --output outputs/inference_esmdiff \
    --num_steps 25 \
    --num_samples 100 \
    --mask_ids 1,2,3,4,5  # 掩码第1-5号残基

数据处理¶

# 1. 下载 PDB 数据库（可选）
bash scripts/download_pdb_mmcif.sh ./pdb_mmcif

# 2. 预处理 PDB 文件
# 注意：需要 biopython==1.79 版本
pip install biopython==1.79
python scripts/pdb/preprocess.py \
    --mmcif_dir pdb_data/pdb_mmcif \
    --output_dir pdb_data/processed_chains \
    --per_chain \
    --strip_array

# 3. 恢复 biopython 版本（主流程需要 1.84）
pip install biopython==1.84

# 4. 计算 VQ-VAE 编码
# 从预处理好的 pickle 文件生成编码
python scripts/dump.py \
    pdb_data/processed_chains \
    pdb_data/processed_chains_encoding \
    pkl

# 从原始 PDB 文件生成编码
python scripts/dump.py \
    pdb_data/raw_pdb \
    pdb_data/raw_pdb_encoding \
    pdb

模型训练¶

# 1. 测试训练流程（仅使用5个训练样本）
python slm/train.py \
    experiment=mdlm \
    paths.data_dir=data/dummy_train_data \
    data.batch_size=1 \
    logger=csv \
    trainer.devices=1 \
    data.train_val_split=[0.8,0.2]

# 2. 标准训练任务（使用 SLURM 提交）
# 联合语言模型（JLM）训练
sbatch train.sh \
    experiment=jlm \
    paths.data_dir=pdb_data/raw_pdb_encoding \
    data.batch_size=16 \
    logger=csv

# 条件语言模型（CLM）训练
sbatch train.sh \
    experiment=clm \
    paths.data_dir=pdb_data/raw_pdb_encoding \
    data.batch_size=16 \
    logger=csv

# 掩码扩散语言模型（MDLM）训练
sbatch train.sh \
    experiment=mdlm \
    paths.data_dir=pdb_data/raw_pdb_encoding \
    data.batch_size=16 \
    logger=csv

模型评估¶

评估脚本位于 ./analysis 目录下，包含： - 构象质量评估（RMSD、TM-score 等） - 生成多样性分析 - 物理合理性检查

常见问题¶

Q1: 如何获取 ESMDiff 的预训练权重？¶

A: 从 Google Drive 下载（链接在项目 README 中），将权重文件放在 data/ckpt 目录下。

Q2: 训练数据需要准备多少？¶

A: 建议使用完整的 PDB 数据库，但测试阶段可以使用 data/dummy_train_data 中的5个样本验证流程。

Q3: 为什么需要两个版本的 biopython？¶

A: 预处理阶段需要 biopython==1.79 来解析 mmCIF 文件，主训练流程需要 1.84 版本。建议使用虚拟环境管理不同版本。

Q4: 如何自定义采样参数？¶

A: 修改 slm/sample_esmdiff.py 中的配置，或通过命令行参数传入。关键参数包括： - num_steps: 扩散步数（影响生成质量） - num_samples: 采样数量（影响多样性） - mask_ids: 需要修复的残基位置

Q5: 训练时出现显存不足怎么办？¶

A: 尝试减小 batch_size 参数，或使用梯度累积。也可以考虑使用混合精度训练。

速查表¶

环境配置¶

组件	版本要求	说明
Python	3.10	推荐使用 conda 管理
PyTorch	最新稳定版	自动安装
Lightning	2.x	用于训练框架
Hydra	最新版	配置管理
ESM3	开源版	需从 HuggingFace 下载

关键命令速查¶

# 环境
conda create -n py310 python=3.10 -y
pip install -r requirements.txt
pip install -e .

# 推理
python slm/sample_esmdiff.py --input <INPUT> --output <OUTPUT> --ckpt <CKPT>

# 数据处理
python scripts/pdb/preprocess.py --mmcif_dir <DIR> --output_dir <DIR>
python scripts/dump.py <INPUT> <OUTPUT> <FORMAT>

# 训练
python slm/train.py experiment=<EXP> paths.data_dir=<DIR> data.batch_size=<N>

# 评估
cd analysis && python evaluate.py

文件结构¶

slm/
├── sample_esmdiff.py    # 推理采样脚本
├── train.py             # 训练入口
├── models/              # 模型定义
├── data/                # 数据处理
│   ├── targets/         # 目标蛋白质结构
│   └── ckpt/            # 预训练权重
├── scripts/             # 辅助脚本
│   ├── download_pdb_mmcif.sh
│   └── pdb/preprocess.py
└── analysis/            # 评估工具

许可证说明¶

ESMDiff 源码和模型：可自由用于非商业目的
ESM3 相关组件：需遵守 Biohub/ESM 的许可证（https://github.com/Biohub/esm/blob/main/LICENSE.md）

引用¶

@inproceedings{
    lu2025structure,
    title={Structure Language Models for Protein Conformation Generation},
    author={Jiarui Lu and Xiaoyin Chen and Stephen Zhewen Lu and Chence Shi and Hongyu Guo and Yoshua Bengio and Jian Tang},
    booktitle={The Thirteenth International Conference on Learning Representations},
    year={2025},
    url={https://openreview.net/forum?id=OzUNDnpQyd}
}