SLM

🔧 总体框架

结构语言建模（SLM） 是一种用于高效 蛋白质构象生成 的两级生成框架，采用 潜在空间中的语言建模。

---

✅ 输入

• 氨基酸序列 ( c \in \mathcal{S}^L )（其中 (\mathcal{S}) 是包含 20 种氨基酸的集合）。
• （可选）用于修复任务的部分或遮蔽构象。

---

⚙️ 两级计算

第一阶段：离散变分自编码器（dVAE）训练

• 目标：学习结构到潜在编码和解码。

1. 编码器：
   将 3D 蛋白质主链结构 ( x \in \mathbb{R}^{L \times 4 \times 3} ) 编码为离散潜在标记 ( z \in \mathcal{V}^L )，其中 (\mathcal{V}) 是一个学习到的码本（例如，512 个标记）。

2. 解码器：
   根据潜在标记 ( z ) 和序列 ( c ) 重建结构 ( x )，即 ( p_\phi(x | z, c) )。

3. 目标：最大化证据下界（ELBO）：

   [
   \log p(x | c) \geq \mathbb{E}_{q(z|x)}[\log p(x | z, c)] - \text{KL}(q(z|x) \parallel p(z|c))
   ]

   （先验 ( p(z|c) ) 固定为均匀分布）。

---

第二阶段：结构语言建模（SLM）

• 目标：使用语言建模学习结构标记的条件分布 ( p_\theta(z|c) )。

选项：

1. 编码器 - 解码器（T5 风格）：

   [
   p(z | c) = \prod_{l=1}^L p(z_l | z_{<l}, c)
   ]

2. 仅解码器（GPT 风格）：

   [
   p(y) = p([c, z]) = \prod_{l=1}^{2L} p(y_l | y_{<l})
   ]

---

🧊 实例化：ESMDiff

• 基于ESM3（一种蛋白质 BERT）构建的掩码离散扩散模型。

• 通过前向掩码过程添加噪声：( z_0 \to z_t \to \text{[MASK]} )

• 使用 Transformer 进行去噪：从掩码输入中恢复标记，条件为 ( c )。

• 使用损失函数：

  [
  \mathcal{L} = \mathbb{E}{c,z_0}\left[\int \frac{1}{1-\alpha(t)} \frac{\partial \alpha(t)}{\partial t} \lambda_M(z_t) \log \langle u\theta(t, z_t, c), z_0 \rangle dt \right]
  ]

---

🧾 输出

• 一组采样的 3D 蛋白质构象 ( x^{(1)}, x^{(2)}, \dots, x^{(N)} )，通过以下方式解码：

  1. 潜在标记 ( z \sim p(z|c) )
  2. 通过 ( x \sim p(x|z, c) ) 解码

---

🧠 关键优势

• 比传统扩散模型快 20 到 100 倍（运行时间见图 5）。
• 高效捕捉多样化构象。
• 支持 修复任务（例如，填充缺失区域，如 CDR 环）。
• 兼容大型预训练语言模型（例如，ESM3）。
• 在 旋转和平移不变空间 运行 → 避免昂贵的 3D 等变性。

---

需要图示风格的可视化拆解，还是实现步骤的代码风格伪代码？

pipline

🔁 完整流程总结

✅ 输入：

• 氨基酸序列 ( c = {c_1, c_2, \ldots, c_L} \in \mathcal{S}^L )
• （仅用于训练）3D 主链坐标 ( x \in \mathbb{R}^{L \times 4 \times 3} )

---

🧩 第一阶段：dVAE – 结构标记化

目标：将 3D 结构压缩为离散潜在标记。

• 输入：结构 ( x )

• 编码器 ( q_\psi(z | x) ) ：将 3D 坐标映射到潜在标记 ( z \in \mathcal{V}^L )

• 解码器 ( p_\phi(x | z, c) ) ：从 ( z ) 重建 3D 坐标，可选地基于序列进行条件生成

• 损失函数：VAE 风格的 ELBO：

  [
  \log p(x|c) \geq \mathbb{E}_{z \sim q(z|x)}[\log p(x|z, c)] - \mathrm{KL}(q(z|x) \parallel p(z|c))
  ]

---

🧠 第二阶段：SLM – 潜在结构语言建模

目标：学习从氨基酸序列 ( c ) 生成结构标记 ( z )。

• 输入：序列 ( c )
• 模型：语言模型（例如，GPT、T5 或掩码扩散模型）
• 输出：采样标记 ( z \sim p_\theta(z | c) )

---

🧬 最终阶段：解码为 3D 结构

目标：从标记生成实际的 3D 坐标。

• 输入：采样的潜在标记 ( z ) 和序列 ( c )
• 解码器 ( p_\phi(x | z, c) )：在第一阶段训练得到
• 输出：蛋白质构象 ( x \in \mathbb{R}^{L \times 4 \times 3} )

---

⛓️ 流程总结

            [氨基酸序列]         +  [结构 (x)]  →  编码器 q(z|x)
                     ↓                                  →  [潜在标记 z]
        ┌─────────────────────────────┐
        │    离散 VAE（第一阶段）     │
        └─────────────────────────────┘
                     ↓
     [潜在标记 z，固定 φ, ψ] + [序列 c]
                     ↓
        ┌─────────────────────────────┐
        │ 语言模型（第二阶段）       │
        └─────────────────────────────┘
                     ↓
       [采样 z ∼ p(z | c)] → 解码器 → [输出 3D 结构 x]

---

需要将此流程作为带标记模块（编码器、解码器、语言模型）的图示进行清理吗？