跳转至

Protein Structure Tokenization: Benchmarking and New Recipe

会议: ICML 2025 arXiv: 2503.00089 代码: GitHub 领域: 蛋白质结构 / 分词评估 关键词: 蛋白质结构分词, VQ-VAE, codebook collapse, 基准测试, ESM3

一句话总结

提出 StructTokenBench——首个全面评估蛋白质结构分词器 (PST) 的基准框架，从下游有效性、敏感性、独特性和 codebook 利用效率四个维度评估现有方法，并提出 AminoAseed 策略通过 codebook 重参数化和 Pareto 最优配置显著改善 VQ-VAE 型分词器的质量（相比 ESM3 提升 6.31%、利用率提升 124%）。

研究背景与动机

蛋白质结构分词 (PST) 将蛋白质 3D 结构编码为离散或连续表示，是蛋白质语言建模和多模态模型的基础。现有 PST 方法分为两类：

VQ-VAE 型：FoldSeek、ProTokens、ESM3——通过重建目标压缩为离散 codebook

反折叠型 (IF)：ProteinMPNN、MIF——通过预测能折叠成目标结构的序列来压缩

但这些方法的能力和局限因缺乏统一评估框架而不明确。特别是 codebook collapse 问题严重——ESM3 的 4096 个 code 中最多 70% 在推理时不活跃。

方法详解

StructTokenBench 评估框架

从四个互补维度评估 PST 质量：

1. 下游有效性 (Downstream Effectiveness)：12 个监督任务（24 个测试分割），包含结合位点预测、催化位点预测、保守位点预测、重复 motif 预测、表位预测、结构柔性预测和远程同源性检测。

2. 敏感性 (Sensitivity)：检测高度相似蛋白质构象间的细微差异，通过表示相似度与 TM-score 的相关性衡量。

3. 独特性 (Distinctiveness)：codebook 向量对间的多样性，避免冗余的 token-子结构映射。

4. Codebook 利用效率：利用率 (UR)、困惑度 (Perplexity)、边际词汇化效用 (MUV)。

AminoAseed 方法

1. Codebook 重参数化

通过可学习线性变换作用于固定正交向量基：

\[\mathbf{Q} = \text{Linear}(\mathbf{C})\]

其中 \(\mathbf{C}\) 随机初始化为近似正交并在训练中固定。与 vanilla VQ-VAE 不同，线性层的所有参数都会收到梯度更新，避免未选中 code 的分布漂移。

2. Pareto 最优 Codebook 配置

在总容量 \(K \times D\) 约束下平衡 codebook 大小 \(K\) 和维度 \(D\)： - \(K > 2^{10}\)：下游性能下降 - \(K < 2^8\)：表达力不足 - \(K = 2^9 = 512\)：利用率和性能的最优平衡

VQ-VAE 优化目标

\[\mathcal{L} = \log p(\tilde{\mathbf{x}}|\mathbf{q}_k) + \|sg(\mathbf{z}) - \mathbf{q}_k\|_2^2 + \beta\|\mathbf{z} - sg(\mathbf{q}_k)\|_2^2\]

实验关键数据

下游有效性 (平均 AUROC %)

方法 类型 平均 AUROC
MIF IF 79.82
ProteinMPNN IF 75.92
ESM3 VQ-VAE 69.24
AminoAseed VQ-VAE 72.43 (+4.74%)
VanillaVQ VQ-VAE 68.30 (-0.86%)

24 个测试分割中 AminoAseed vs ESM3

在 24 个监督任务测试分割中，AminoAseed 平均提升 6.31%，部分任务（如 CatBio-SupFam）提升达 17.33%。

Codebook 利用效率

模型 利用率 UR (%)
FoldSeek 最高
ESM3 ~30% (4096 中活跃 ~1200)
AminoAseed ESM3 的 2.24× (提升 124%)

关键消融发现

  1. 向量量化影响模型表达力主要因为优化挑战，而非离散/连续格式差异
  2. 结构 token 保留了氨基酸 token 的大部分信息但对噪声更敏感
  3. 重建质量与 codebook 质量不一致相关——两者都需要单独评估
  4. 放大 VQ-VAE 编码器的收益呈亚指数递减

亮点与洞察

  1. 首个蛋白质结构分词统一基准：4 维度 + 12 任务 + 24 分割，填补了关键评估空白
  2. codebook collapse 的根因分析透彻：训练过程中未被选中的 code 不更新→分布漂移→进一步不被选中（正反馈）
  3. 重参数化解法简洁有效：一个线性层就让所有 code 都参与梯度更新
  4. \(K=512\) 的生物学合理性：与启发式方法 TERMs 的发现一致（~600 个子结构描述 50% PDB）
  5. IF 型 vs VQ-VAE 型的互补性：no single method dominates，各有擅长

局限性

  1. AminoAseed 在物理化学属性预测（FlexBFactor）上反而下降，说明 codebook 质量与所有下游任务的关系不总是正向的
  2. 评估聚焦于骨架结构输入，未覆盖全原子结构
  3. 仅用 10% 的过滤 PDB 数据预训练，可能限制模型上限
  4. 对 IF 型方法的 codebook 利用评估不适用（连续表示）

相关工作

  • VQ-VAE PST：FoldSeek, ProTokens, ESM3
  • IF PST：ProteinMPNN, MIF
  • Codebook collapse：EMA 更新、k-means 初始化
  • 蛋白质基准：TAPE, ProteinGLUE

评分

⭐⭐⭐⭐ (4/5)

基准框架设计出色，四维度评估全面且科学。AminoAseed 方法虽简单但有效性好。消融和缩放研究深入，提供了大量有价值的洞察。美中不足是 AminoAseed 仍未缩小 VQ-VAE 与 IF 方法的差距。

相关论文