跳转至

Learning Cell-Aware Hierarchical Multi-Modal Representations for Robust Molecular Modeling

会议: AAAI 2026 (Oral)
arXiv: 2511.21120v1
代码: https://github.com/limengran98/CHMR
领域: 药物发现 / 分子性质预测 / 多模态学习
关键词: 分子表示学习, 细胞感知, 层次化向量量化, 多模态缺失, 对比学习

一句话总结

提出CHMR框架，将分子结构(1D/2D/3D)与细胞形态/基因表达等生物模态联合建模，通过结构感知的模态增强解决>90%的外部生物模态缺失问题，用树状向量量化(Tree-VQ)捕获分子-细胞-基因的层次化依赖关系，在9个benchmark的728个任务上超越SOTA，分类平均AUC提升3.6%，回归MAE降低17.2%。

背景与动机

分子性质预测(活性、毒性、ADME)是药物发现的核心任务。现有方法主要基于分子结构(指纹、图、3D构象)做预测，忽略了分子与细胞相互作用后引发的生物学响应(细胞形态变化、基因表达改变)——这些响应蕴含着分子作用机制的关键信息。近期cell-aware方法(InfoAlign, InfoCORE)开始引入细胞表型和基因表达数据，但面临两个核心瓶颈：(1)模态缺失严重——外部生物模态(Cell Painting、L1000、CRISPR)的缺失率高达90-99%, 数据极不完整；(2)现有方法在扁平化的潜空间做对齐，无法捕获分子→细胞→基因的层次化因果依赖。

核心问题

如何在生物模态严重缺失的条件下，构建鲁棒的分子表示学习框架？需要同时解决：(1)90%+缺失率下的模态补全；(2)跨尺度(分子-细胞-基因)的层次化语义建模；(3)不同模态间的语义一致性对齐。

方法详解

整体框架

预训练阶段：129K分子的多模态数据(1D指纹+2D图+3D构象+Cell Painting+CRISPR/ORF+L1000) → 模态增强(MA, 图传播补全缺失模态) → 语义一致性对齐(SCA, InfoNCE+VICReg) → 树状向量量化(Tree-VQ, 捕获层次化语义) → 上下文传播重建(CPR, 随机游走+重建loss) → 冻结backbone，轻量MLP做下游任务。

关键设计

  1. 模态增强(MA): 基于Dirichlet能量最小化的迭代图传播。构建分子相似性图(top-K近邻)，缺失模态的特征通过邻居观测值迭代加权平均得到，收敛到图Laplacian的解。比简单的零填充/随机填充/均值填充效果显著更好(消融降5.3%/4.5%/2.9%)。
  2. 语义一致性对齐(SCA): 双层对齐——(a)样本级：InfoNCE loss对齐分子锚向量(1D+2D+3D聚合)和细胞模态表示；(b)分布级：VICReg loss对齐增强前后的细胞特征分布(消除传播引入的分布偏移)。两者互补，缺一性能降1.7-2.9%。
  3. 树状向量量化(Tree-VQ): 建立深度H=6的二叉树，所有模态共享这棵树。每个模态特征按cosine距离沿树路径做层次化量化——浅层对应粗粒度语义(分子指纹)，深层对应细粒度语义(细胞表型/基因表达)。核心定理证明：共享相同量化路径的两个输入，其cosine相似度下界为2α²-1，保证同一语义分区内的特征一致性。相比扁平VQ提升2.0%。
  4. 上下文传播重建(CPR): 在生物知识图(分子-基因-细胞关系图)上做随机游走，用沿路径的累积权重指导跨模态重建loss，增强缺失模态的补全质量。

损失函数 / 训练策略

  • L_total = L_CPR + λ1·L_SCA + λ2·L_TreeVQ
  • 最优：λ1=10, λ2=1, η=1, H=6
  • 预训练：RTX 3090, 50 epochs, Adam, lr=1e-4, batch=3072
  • 下游：冻结backbone，训练2层MLP；4个数据集(BACE/ClinTox/SIDER/HIV)额外ensemble随机森林(γ加权)

实验关键数据

数据集(类型) 指标 CHMR 前SOTA(InfoAlign) 提升
ChEMBL(分类) AUC 84.7 81.3 +3.4
ToxCast(分类) AUC 69.3 66.4 +2.9
Broad(分类) AUC 71.4 70.0 +1.4
Biogen(回归) MAE 40.9 49.4 -17.2%
5个OGB数据集均值(分类) AUC 82.2 79.1 +3.1

消融实验要点

  • MA最关键: 零填充→AUC降-5.3%，说明模态补全是基础
  • SCA不可or缺: 去掉→-3.6%，样本级对齐(-2.9%)比分布级(-1.7%)重要
  • Tree-VQ vs 扁平VQ: 层次化结构额外提升+2.0%
  • CPR中随机游走贡献: 去掉→-2.0%，说明生物知识图的上下文信息有用
  • 多模态协同: 仅用分子→-4.9%；加任何一种生物模态都有提升(~2%)

亮点

  • 90%+缺失率下仍有效 — 实际场景中生物模态极度缺失，MA+SCA组合让模型依然鲁棒
  • Tree-VQ的生物学可解释性 — 树的浅层对应化学指纹(粗粒度)，深层对应细胞/基因(细粒度)，与生物学因果链(分子→细胞→基因)自然对应
  • 理论保证 — 自定义的Tree-VQ有严格的cosine相似度下界证明(Theorem 1)
  • 9个benchmark 728个任务的大规模验证 — 覆盖活性、毒性、ADME等多种预测任务

局限性 / 可改进方向

  • InfoNCE中N²的复杂度限制了数据规模的进一步扩展
  • Tree-VQ的树深度H=6是手动设定，自适应深度可能更好
  • 仅在冻结backbone+MLP的设置下评估，全模型微调可能有不同结论
  • 4个数据集需要ensemble随机森林辅助，说明CHMR自身在某些分布上仍有不足

与相关工作的对比

与InfoAlign(ICLR 25)相比，核心差异是Tree-VQ的层次化建模(InfoAlign做扁平对齐)和基于图传播的模态增强(InfoAlign直接在缺失数据上对齐)。与CLOOME/InfoCORE等早期cell-aware方法相比，CHMR同时解决了模态缺失和层次化建模两个问题。与UniMol/MOL-Mamba等结构方法相比，CHMR额外利用了生物模态信息。

启发与关联

  • Tree-VQ的层次化量化设计 → 可迁移到其他有层次结构的多模态任务(如遥感多尺度理解)
  • 图传播做模态补全 → 可用于其他多模态严重缺失的场景(如多组学数据、临床多模态数据)
  • 分子-细胞-基因的跨尺度建模 → 药物发现和精准医疗的关键技术方向

评分

  • 新颖性: ⭐⭐⭐⭐ Tree-VQ在分子多模态中的应用有新意
  • 实验充分度: ⭐⭐⭐⭐⭐ 9个数据集728任务+20+baseline+完整消融+超参分析+可视化
  • 写作质量: ⭐⭐⭐⭐ 结构清晰，有理论推导
  • 价值: ⭐⭐⭐⭐ 药物发现场景的实用框架