lnp_ml/README_llm_encoders.md

LLM 分子编码器分支（LLM Encoder Branch）

本模块在 LNP 多任务预测模型上新增了一个可选的 LLM 分子编码器分支。它把可电离脂质的 SMILES 送入一个冻结的预训练分子模型，提取语义向量，作为一个额外的 token（V5）追加到模型的 fusion 序列中参与多任务预测。

支持四种分子编码器，可通过参数自由切换，用于横向对比：

llm_type	模型	年份 / 发表	参数量	hidden	输入
biot5	BioT5-plus	2024, arXiv	~250M	768	SMILES
molt5	MolT5-base	2022, EMNLP	109M	768	SMILES
moleculestm	MoleculeSTM GNN	2023, Nature MI	~2M (GNN)	300	分子图
selfiested	SELFIES-TED	2025, IBM	358M	1024	SELFIES

---

1. 设计说明

• 非侵入式：默认 use_llm=False，此时模型行为与原版完全一致，不影响既有功能。只有显式开启才会启用 LLM 分支。
• 统一接口：四个 encoder 的 forward(List[str]) -> [B, d_model] 接口完全一致，通过 models.build_llm_encoder() 工厂函数按 llm_type 实例化。
• 冻结主体：LLM / GNN 主体全程冻结，只有内部一层 projection（LayerNorm + Linear + ReLU + Dropout）参与训练，参数量增加极小。
• 惰性导入：encoders/__init__.py 对四个 LLM encoder 采用惰性导入，未安装 torch_geometric / selfies 等可选依赖时，只要不使用对应 encoder，导入不会报错。

---

2. 文件清单

新增文件（均在 lnp_ml/modeling/encoders/ 下）：

encoders/
├── llm_encoder.py            # BioT5-plus  (类名 LLMEncoder)
├── molt5_encoder.py          # MolT5       (类名 MolT5Encoder)
├── moleculestm_encoder.py    # MoleculeSTM (类名 MoleculeSTMEncoder)
├── molecule_gnn_model.py     # MoleculeSTM 的 GNN 网络定义（被上一个文件 import）
└── selfiested_encoder.py     # SELFIES-TED (类名 SELFIESTEDEncoder)

修改文件：

lnp_ml/modeling/models.py            # 新增 use_llm / llm_type 等参数 + build_llm_encoder 工厂 + V5 追加逻辑；desc 维度 210 -> 217
lnp_ml/modeling/encoders/__init__.py # 惰性导出四个 LLM encoder
verify_llm_encoder.py                # LLM 模块验证脚本（四项测试）

注：models.py 中 DEFAULT_INPUT_DIMS["desc"] 已从 210 更新为 217，以匹配新版 RDKit 的描述符维度。

---

3. 依赖安装

基础依赖（所有模型都需要）：

pip install rdkit loguru transformers sentencepiece optuna

按模型追加：

模型	额外依赖
biot5	无（基础依赖即可）
molt5	无
moleculestm	pip install torch_geometric torch_scatter torch_sparse ogb
selfiested	pip install selfies

---

4. 预训练权重

模型	是否需手动下载	说明
molt5	否	训练时自动从 HuggingFace 拉取 laituan245/molt5-base
selfiested	否	训练时自动从 HuggingFace 拉取 ibm/materials.selfies-ted
biot5	是	需下载 BioT5-plus 权重文件夹，见下方
moleculestm	是	需下载 MoleculeSTM GNN 权重，见下方

BioT5-plus 权重下载： 从 HuggingFace QizhiPei/biot5-plus-base 下载整个权重文件夹（含 config.json、tokenizer_config.json 等），放到本地任意路径，训练时用 --llm-model-path <该路径> 指定。也可直接传 HuggingFace ID QizhiPei/biot5-plus-base 在线加载。

MoleculeSTM GNN 权重下载：

from huggingface_hub import snapshot_download
snapshot_download(
    repo_id="chao1224/MoleculeSTM",
    repo_type="model",
    local_dir="<本地目录>/moleculestm_weights",
    allow_patterns="*Graph*",
)

下载后使用的 GNN 权重路径为：

<本地目录>/moleculestm_weights/pretrained_MoleculeSTM/SciBERT-Graph-3e-5-1-1e-4-1-InfoNCE-0.1-32-32/molecule_model.pth

训练时用 --llm-model-path <该 .pth 路径> 指定。

---

5. 用法

5.1 代码中直接构造模型

from lnp_ml.modeling.models import LNPModelWithoutMPNN

# 不启用 LLM（与原版行为一致）
model = LNPModelWithoutMPNN(d_model=256)

# 启用 LLM 分支（以 MolT5 为例）
model = LNPModelWithoutMPNN(
    d_model=256,
    use_llm=True,
    llm_type="molt5",                       # biot5 / molt5 / moleculestm / selfiested
    llm_model_path="laituan245/molt5-base", # 见上表，不同模型路径不同
    llm_device="cuda",
)

各模型对应的 llm_type 和 llm_model_path：

llm_type	llm_model_path 示例
biot5	QizhiPei/biot5-plus-base 或本地 BioT5 权重文件夹路径
molt5	laituan245/molt5-base
moleculestm	.../moleculestm_weights/.../molecule_model.pth
selfiested	ibm/materials.selfies-ted

5.2 验证 LLM 模块（训练前建议先跑）

python verify_llm_encoder.py \
    --model_path <llm_model_path> \
    --lnp_repo_path <仓库根目录>

四项测试（输出维度、维度对齐、可复现性、过拟合 sanity）全部 PASS 再训练。

5.3 训练

训练脚本（final_train_optuna_cv.py / pretrain.py）需要支持 --use-llm / --llm-type / --llm-model-path / --llm-device 参数。各模型的完整 Colab 操作步骤见 docs/ 下四份 quickstart 指南。

---

6. 实验结果

四个模型在 LNP 测试集上的对比（核心指标 delivery R²，baseline 为不含 LLM 的版本）：

任务	Baseline	MolT5	MoleculeSTM	BioT5-plus	SELFIES-TED
delivery R²	0.373	0.602	0.552	0.633	0.618
size R²	0.215	0.440	0.408	0.423	0.381
pdi acc	0.786	0.748	0.710	0.695	0.664
ee acc	0.679	0.702	0.679	0.679	0.679
toxic acc	0.980	0.980	0.951	0.960	0.941
biodist KL	0.293	0.727	0.646	0.703	0.643

主要结论：

1. 四个 LLM 编码器都显著提升了核心回归任务，delivery R² 从 baseline 的 0.37 最高提升到 0.63。
2. 领域专用预训练比堆参数更重要：BioT5-plus（针对生物分子预训练）效果最好，而参数最大、最新的 SELFIES-TED 仅排第二。
3. 序列模型整体优于图模型：三个基于 SMILES/SELFIES 的模型均优于基于图的 MoleculeSTM。
4. MolT5 性价比最高：仅 109M 参数，delivery R² 排第三且与第一仅差 0.03，分类任务表现最均衡。
5. 共同局限：biodist（分布预测）四个 LLM 均略低于 baseline，属多任务学习中的 trade-off。

各模型每个任务的详细指标见 results/<模型名>/ 目录。

---

附：训练脚本改动

本 PR 已经更新了两个训练脚本以支持 LLM 分支（MoE 相关代码未引入）：

• lnp_ml/modeling/final_train_optuna_cv.py
• lnp_ml/modeling/pretrain.py

两者都新增了四个命令行参数，并透传给模型：

--use-llm                                          # 是否启用 LLM 分支
--llm-type {biot5,molt5,moleculestm,selfiested}    # 选择哪个 encoder
--llm-model-path <路径或HF_ID>                     # 见第 4/5 节
--llm-device cuda                                  # LLM 推理设备

预训练保存的 checkpoint 的 config 中也会记录 LLM 设置，恢复模型时自动读取。

训练示例（以 MolT5 为例）

# 预训练
python -m lnp_ml.modeling.pretrain main \
    --train-path data/processed/train_pretrain.parquet \
    --val-path   data/processed/val_pretrain.parquet \
    --use-llm --llm-type molt5 \
    --llm-model-path laituan245/molt5-base \
    --llm-device cuda --epochs 50 --lr 1e-4 --device cuda

# 正式训练
python lnp_ml/modeling/final_train_optuna_cv.py \
    --init-from-pretrain models/pretrain_delivery.pt \
    --use-llm --llm-type molt5 \
    --llm-model-path laituan245/molt5-base \
    --llm-device cuda \
    --n-trials 20 --epochs-per-trial 30 --seed 42 \
    --device cuda --output-dir models/final

切换其他模型只需改 --llm-type 和 --llm-model-path（见第 5.1 节对照表）。