README.md

---
license: mit
metrics:
- accuracy
pipeline_tag: image-to-text
tags:
- chemistry
---
# MS²-SMILES AlignNet 模型文档

## 📋 目录

1. [概述](#概述)
2. [数据处理流程](#数据处理流程)
3. [模型架构](#模型架构)
4. [损失函数设计](#损失函数设计)
5. [训练流程](#训练流程)
6. [配置参数](#配置参数)

---

## 概述

MS²-SMILES AlignNet 是一个基于**对比学习**的跨模态对齐模型，旨在学习质谱（MS²）数据与分子结构（SMILES）之间的语义对齐表示。该模型采用双分支架构，分别编码 MS 谱图和分子结构，并通过对比学习使得匹配的 MS-分子对在嵌入空间中靠近。

### 核心思想

```
┌─────────────────────────────────────────────────────────────────────────────┐
│                        MS²-SMILES AlignNet 架构                              │
├─────────────────────────────────────────────────────────────────────────────┤
│                                                                              │
│   SMILES ──→ [分子编码分支] ──→ 256维向量 ─┐                                 │
│                                            ├──→ 对比学习 ──→ 对齐损失        │
│   MS谱图 ──→ [MS编码分支]  ──→ 256维向量 ──┘                                 │
│                                                                              │
└─────────────────────────────────────────────────────────────────────────────┘
```

---

## 数据处理流程

### 2.1 数据加载入口

**文件**: `train.py` (第174-175行)

```python
train_loader = build_loaders(train_set, "train", cfg, 10)
valid_loader = build_loaders(valid_set, "valid", cfg, 10)
```

数据支持两种输入格式：
- **JSON/MGF 文件列表**: 使用 `PathDataset` 类
- **预处理字典列表**: 使用 `Dataset` 类

### 2.2 原始数据获取

**文件**: `dataset.py` (第103-128行)

```
PathDataset
    ├── proc_data(): 解析 JSON/MGF 文件
    │   └── 提取: ms (质谱峰列表), smiles (分子结构)
    │
    └── __getitem__(): 调用 calc_feats() 计算特征
```

**数据格式**:
- `ms`: List of (m/z, intensity) 元组
- `smiles`: 分子 SMILES 字符串

### 2.3 特征预处理

**文件**: `dataset.py` (第7-41行) - `calc_feats()` 函数

```
┌─────────────────────────────────────────────────────────────────────────────┐
│                           特征计算流程                                       │
├─────────────────────────────────────────────────────────────────────────────┤
│                                                                              │
│  输入: smi (SMILES), ms (质谱), nls (中性丢失)                               │
│                                                                              │
│  ┌────────────────────────┐  ┌────────────────────────────┐                 │
│  │   MS 特征处理          │  │   分子特征处理              │                 │
│  ├────────────────────────┤  ├────────────────────────────┤                 │
│  │ 1. ms_binner()         │  │ 1. mol_fp_encoder()        │                 │
│  │    → ms_bins           │  │    → mol_fps (指纹)        │                 │
│  │    (39600维稀疏向量)    │  │    (2048维)                │                 │
│  │                        │  │                            │                 │
│  │ 2. ms_feature_processor│  │ 2. mol_graph_featurizer()  │                 │
│  │    → ms_bins1 (100,29) │  │    → V (节点特征)          │                 │
│  │    → ms_bins2 (100,)   │  │    → A (邻接矩阵)          │                 │
│  │    (峰特征 + m/z值)     │  │    → mol_size             │                 │
│  └────────────────────────┘  └────────────────────────────┘                 │
│                                                                              │
└─────────────────────────────────────────────────────────────────────────────┘
```

#### 2.3.1 MS 特征处理详解

**方法一: `ms_binner()` - 稀疏分桶表示**

```python
# 参数配置
min_mz = 20, max_mz = 2000, bin_size = 0.05
# 计算: (2000-20)/0.05 = 39600 bins

# 流程:
1. 过滤 m/z 范围外的峰
2. 归一化强度 (0-1)
3. 分桶 (bin_idx = (m/z - min_mz) / bin_size)
4. 可选: 添加中性丢失特征 (额外 4000 bins)
```

**方法二: `ms_feature_processor()` - 序列化峰表示**

```python
输入: ms, precursor_mz, metadata_vec(135维)
输出: 
  - ms_bins1: (100,135) 峰特征矩阵
  - ms_bins2: (100,) m/z 值序列

## 135 维全局特征 = 仪器(20) + 加合物(21) + 碰撞能量(5) + 电荷电离(4) + 诊断离子(46) + 中性丢失(39)
线性变换 → 256 维模型输
```

#### 2.3.2 分子特征处理详解

**Morgan 指纹编码**:
```python
mol_fp_encoder(smiles, tp='morgan', nbits=2048)
# → 2048 维二进制向量
```

**分子图特征化**:
```python
mol_graph_featurizer(smiles)
# → V: 节点特征矩阵 (N_atoms, 74)
# → A: 邻接张量 (N_atoms, 6, N_atoms)
#      6 通道 = 4 键类型 + 环内 + 共轭
# → mol_size: 原子数量
```

**原子节点特征 (74维)**:
| 特征类型 | 维度 | 说明 |
|---------|------|------|
| 原子类型 | 12 | One-hot: H,C,N,O,S,F,Si,P,Cl,Br,I,B |
| 显式价态 | 7 | One-hot: 0-6 |
| 隐式价态 | 7 | One-hot: 0-6 |
| 氢原子数 | 5 | One-hot: 0-4 |
| 自由基电子 | 5 | One-hot: 0-4 |
| 总度数 | 7 | One-hot: 0-6 |
| 形式电荷 | 5 | One-hot: -2 to 2 |
| 杂化类型 | 5 | SP,SP2,SP3,SP3D,SP3D2 |
| 芳香性 | 2 | True/False |
| 环内原子 | 2 | True/False |
| 手性 | 4 | 4种类型 |
| CIP编码 | 2 | R/S |
| 范德华半径 | 1 | 连续值 |
| 手性可能性 | 1 | True/False |
| 原子序数 | 1 | 连续值 |
| 原子质量 | 1 | 连续值 × 0.01 |
| 度数 | 1 | 连续值 |
| 特征不变量 | 6 | 6位二进制编码 |

### 2.4 数据批处理聚合

**文件**: `train.py` (第36-83行) - `my_collate()` 函数

```
┌─────────────────────────────────────────────────────────────────────────────┐
│                         Batch 聚合流程                                       │
├─────────────────────────────────────────────────────────────────────────────┤
│                                                                              │
│  输入: [sample_1, sample_2, ..., sample_B]  (B = batch_size)                │
│                                                                              │
│  1. 过滤 None 样本                                                           │
│                                                                              │
│  2. 堆叠各特征:                                                              │
│     ├── ms_bins:  stack → (B, 39600+4000)                                   │
│     ├── ms_bins1: stack → (B, 100, 135)                                      │
│     ├── ms_bins2: stack → (B, 100)                                          │
│     ├── mol_fps:  stack → (B, 2048)                                         │
│     └── mol_fmvec: stack → (B, 10)  [可选: 分子式向量]                       │
│                                                                              │
│  3. 图数据 Padding:                                                          │
│     ├── 计算 max_n = max(所有分子的原子数)                                   │
│     ├── V: pad → (B, max_n, 74)                                             │
│     ├── A: pad → (B, max_n, 6, max_n)                                       │
│     └── mol_size: concat → (B,)                                             │
│                                                                              │
└─────────────────────────────────────────────────────────────────────────────┘
```

---

## 模型架构

**文件**: `modules.py` (第298-457行) - `FragSimiModelNew` 类

### 3.1 整体架构图

```
┌─────────────────────────────────────────────────────────────────────────────────────────┐
│                              FragSimiModelNew 双分支架构                                 │
├─────────────────────────────────────────────────────────────────────────────────────────┤
│                                                                                          │
│  ┌──────────────────────────────────┐    ┌──────────────────────────────────────────┐   │
│  │     分子编码分支 (LEFT)          │    │        MS 编码分支 (RIGHT)                │   │
│  ├──────────────────────────────────┤    ├──────────────────────────────────────────┤   │
│  │                                  │    │                                          │   │
│  │  SMILES                          │    │  MS 谱图                                 │   │
│  │     │                            │    │     │                                    │   │
│  │     ▼                            │    │     ▼                                    │   │
│  │  ┌─────────────┐                 │    │  ┌─────────────────────────────────────┐│   │
│  │  │ RDKit 解析   │                 │    │  │ 分支A: ms_bins (稀疏分桶)           ││   │
│  │  └──────┬──────┘                 │    │  │   │                                 ││   │
│  │         │                        │    │  │   ▼ ms_input_proj (Linear)          ││   │
│  │    ┌────┴────┐                   │    │  │   │                                 ││   │
│  │    ▼         ▼                   │    │  │   ▼ 6层 Transformer                 ││   │
│  │ ┌──────┐ ┌───────┐               │    │  │   │                                 ││   │
│  │ │ 分子图 │ │ Morgan│               │    │  │   ▼ Mean Pooling                   ││   │
│  │ │构建   │ │ 指纹  │               │    │  │   │                                 ││   │
│  │ └──┬───┘ └───┬───┘               │    │  │   ▼ ms_final_proj                   ││   │
│  │    │         │                   │    │  │   └───→ 256维                       ││   │
│  │    ▼         │                   │    │  └─────────────────────────────────────┘│   │
│  │ ┌──────────┐ │                   │    │                     │                    │   │
│  │ │MolGNN    │ │                   │    │                     ▼                    │   │
│  │ │Encoder   │ │                   │    │  ┌─────────────────────────────────────┐│   │
│  │ │(3层GCN   │ │                   │    │  │ 分支B: ms_bins1/2 (峰序列)          ││   │
│  │ │+Attention│ │                   │    │  │   │                                 ││   │
│  │ │+Readout) │ │                   │    │  │   ▼ SinusoidalMzEmbedding           ││   │
│  │ └────┬─────┘ │                   │    │  │   + feature_proj (135→256)           ││   │
│  │      │       │                   │    │  │   │                                 ││   │
│  │      ▼       ▼                   │    │  │   ▼ 6层 Transformer                 ││   │
│  │ ┌──────────────────┐             │    │  │   │                                 ││   │
│  │ │  MolFusionHead   │             │    │  │   ▼ Mean Pooling                   ││   │
│  │ │ (GNN→128, FP→256)│             │    │  │   │                                 ││   │
│  │ │ → Concat → 256   │             │    │  │   ▼ final_proj                      ││   │
│  │ └────────┬─────────┘             │    │  │   └───→ 256维                       ││   │
│  │          │                       │    │  └─────────────────────────────────────┘│   │
│  │          ▼                       │    │                     │                    │   │
│  │     256维向量                    │    │                     ▼                    │   │
│  │                                  │    │  ┌─────────────────────────────────────┐│   │
│  │                                  │    │  │    Concat (256+256=512)             ││   │
│  │                                  │    │  │           │                         ││   │
│  │                                  │    │  │           ▼                         ││   │
│  │                                  │    │  │    all_final_proj → 256维           ││   │
│  │                                  │    │  └─────────────────────────────────────┘│   │
│  └──────────────────────────────────┘    └──────────────────────────────────────────┘   │
│                  │                                           │                          │
│                  ▼                                           ▼                          │
│          mol_embeddings (256)                        ms_embeddings (256)                │
│                  │                                           │                          │
│                  └───────────────────┬───────────────────────┘                          │
│                                      ▼                                                  │
│                          ┌───────────────────────┐                                      │
│                          │    L2 Normalize       │                                      │
│                          └───────────┬───────────┘                                      │
│                                      ▼                                                  │
│                          ┌───────────────────────┐                                      │
│                          │   HybridAlignLoss     │                                      │
│                          │ (InfoNCE + Tanimoto)  │                                      │
│                          └───────────────────────┘                                      │
│                                                                                          │
└─────────────────────────────────────────────────────────────────────────────────────────┘
```

### 3.2 分子编码分支

#### 3.2.1 MolGNNEncoder

```python
class MolGNNEncoder(nn.Module):
    # 输入: V (B, N, 74), A (B, N, 6, N), mol_size (B,)
    # 输出: (B, mol_embedding_dim)
    
    结构:
    ├── block_layers: 3 × GConvBlockNoGF
    │   └── GraphCNNLayer: 图卷积 + BatchNorm + ReLU
    │       [74 → 256 → 256 → 256]
    │
    ├── attention_layer: MultiHeadGlobalAttention (4 heads)
    │   └── 对节点进行加权聚合，去除 padding 影响
    │       [256 → 256×4 = 1024]
    │
    └── readout_layers: 2 × Linear + GELU
        └── [1024 → 2048 → 2048]
```

#### 3.2.2 MolFusionHead

```python
class MolFusionHead(nn.Module):
    # 输入: gnn_feat (B, 2048), fps (B, 2048)
    # 输出: (B, 256)
    
    结构:
    ├── gnn_proj: Linear(2048 → 128)
    │
    ├── fp_proj: Linear(2048 → 256) + LayerNorm + GELU
    │
    └── fusion_layer: 
        ├── Concat: [128, 256] → 384
        └── Linear(384 → 256) + LayerNorm + Dropout + GELU
```

### 3.3 MS 编码分支

#### 3.3.1 分支A: 稀疏分桶编码

```python
# 输入: ms_bins (B, 43600)
# 输出: 256维

流程:
1. ms_input_proj: Linear(43600 → 256)
2. unsqueeze → (B, 1, 256) [假序列]
3. ms_transformer: 6层 Transformer (width=256, heads=8)
4. mean(dim=1) → (B, 256)
5. ms_final_proj: Linear + LayerNorm + Dropout → 256维
```

#### 3.3.2 分支B: 峰序列编码

```python
# 输入: ms_bins1 (B, 100, 29), ms_bins2 (B, 100)
# 输出: 256维

流程:
1. mz_embedder(ms_bins2) → (B, 100, 256)  # 正弦位置编码
2. feature_proj(ms_bins1) → (B, 100, 256)  # 峰特征投影
3. 相加融合 → (B, 100, 256)
4. transformer: 6层 Transformer
5. mean(dim=1) → (B, 256)
6. final_proj: Linear + LayerNorm + Dropout → 256维
```

#### 3.3.3 正弦 m/z 嵌入

```python
class SinusoidalMzEmbedding(nn.Module):
    # 将 m/z 值编码为 256 维向量
    # 类似 Transformer 的位置编码
    
    PE[pos, 2i]   = sin(pos / 10000^(2i/256))
    PE[pos, 2i+1] = cos(pos / 10000^(2i/256))
```

#### 3.3.4 双分支融合

```python
# 两个 MS 分支的输出融合
out1 = torch.cat([ms_embeddings_A, ms_embeddings_B], dim=-1)  # (B, 512)
ms_embeddings = all_final_proj(out1)  # Linear(512→256) + LayerNorm + Dropout
```

### 3.4 Transformer 结构

**文件**: `cliplayers.py` - 采用 CLIP 风格的 Transformer

```python
class ResidualAttentionBlock(nn.Module):
    # 单个 Transformer Block
    
    结构:
    ├── ln_1 → MultiheadAttention → 残差连接
    │
    └── ln_2 → MLP (Linear→QuickGELU→Linear) → 残差连接
    
    # QuickGELU: x * sigmoid(1.702 * x)
    # 比标准 GELU 更快
```

---

## 损失函数设计

**文件**: `modules.py` (第196-233行) - `HybridAlignLoss` 类

### 4.1 联合损失公式

$$\mathcal{L}_{total} = \beta \cdot \mathcal{L}_{InfoNCE} + \alpha \cdot \mathcal{L}_{Tanimoto-MSE}$$

**默认参数**: α = 0.5, β = 1.0

### 4.2 InfoNCE 对比损失

```python
# 目标: 让匹配的 MS-分子对相似度高，不匹配的对相似度低

logits = (ms_emb @ mol_emb.T) / temperature  # 相似度矩阵
labels = torch.arange(batch_size)            # 对角线为正样本

loss_i2t = CrossEntropy(logits, labels)      # MS → Mol
loss_t2i = CrossEntropy(logits.T, labels)    # Mol → MS

L_InfoNCE = (loss_i2t + loss_t2i) / 2
```

**温度参数**: τ = 0.07

### 4.3 Tanimoto MSE 损失

```python
# 目标: 让嵌入空间的相似度矩阵逼近真实的分子结构相似度

# 1. 预测的相似度矩阵 (余弦相似度)
pred_sim = ms_emb @ mol_emb.T  # 因为已 L2 normalize

# 2. 真实的结构相似度 (Tanimoto 系数)
#    Tanimoto(A,B) = (A·B) / (|A|² + |B|² - A·B)
target_sim = batch_tanimoto_sim(mol_fps, mol_fps)

# 3. MSE 损失
L_MSE = MSELoss(pred_sim, target_sim)
```

### 4.4 损失设计的意义

| 损失组件 | 作用 | 特点 |
|---------|------|------|
| InfoNCE | 跨模态对齐 | 区分正负样本对，对比学习核心 |
| Tanimoto MSE | 结构约束 | 保持分子结构相似性关系 |

```
┌─────────────────────────────────────────────────────────────────┐
│                    联合损失的几何意义                            │
├─────────────────────────────────────────────────────────────────┤
│                                                                  │
│   InfoNCE: "匹配的 MS-Mol 对要靠近"                             │
│                                                                  │
│       MS₁ ←──── 拉近 ────→ Mol₁                                 │
│       MS₂ ←──── 拉近 ────→ Mol₂                                 │
│        ↑                    ↑                                    │
│     推远                  推远                                   │
│        ↓                    ↓                                    │
│       MS₃ ←──── 拉近 ────→ Mol₃                                 │
│                                                                  │
│   Tanimoto MSE: "结构相似的分子在嵌入空间也要相似"               │
│                                                                  │
│     如果 Tanimoto(Mol₁, Mol₂) = 0.8                             │
│     则希望 CosSim(Emb₁, Emb₂) ≈ 0.8                             │
│                                                                  │
└─────────────────────────────────────────────────────────────────┘
```

---

## 训练流程

**文件**: `train.py`

### 5.1 训练入口

```python
# main() 函数

1. 数据划分: make_train_valid(data, valid_ratio=0.1)
2. 构建 DataLoader: build_loaders()
3. 初始化模型: FragSimiModelNew(cfg).to(device)
4. 优化器: AdamW(lr=1e-3, weight_decay=1e-3)
5. 学习率调度: ReduceLROnPlateau(patience=2, factor=0.5)
6. 训练循环: epochs=50
```

### 5.2 训练周期

```python
def train_epoch(model, train_loader, optimizer, lr_scheduler, step):
    for batch in train_loader:
        # 1. 数据移到 GPU
        batch = {k: v.to(device) for k, v in batch.items()}
        
        # 2. 前向传播
        loss = model(batch)  # 返回总损失
        
        # 3. 反向传播
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
```

### 5.3 模型保存策略

```python
# 保存最佳模型 (基于验证集损失)
if valid_loss < best_loss:
    checkpoint = {
        'state_dict': model.state_dict(),
        'optimizer': optimizer.state_dict(),
        'config': dict(CFG)
    }
    torch.save(checkpoint, best_model_fn)

# 只保留最新的 3 个最佳模型
keep_best_models_num = 3
```

---

## 配置参数

**文件**: `config.py`

### 6.1 核心参数

| 参数 | 默认值 | 说明 |
|------|-------|------|
| `batch_size` | 64 | 批次大小 |
| `lr` | 1e-3 | 学习率 |
| `epochs` | 50 | 训练轮数 |
| `dropout` | 0.1 | Dropout 比例 |
| `projection_dim` | 256 | 最终嵌入维度 |

### 6.2 MS 相关参数

| 参数 | 默认值 | 说明 |
|------|-------|------|
| `min_mz` | 20 | 最小 m/z |
| `max_mz` | 2000 | 最大 m/z |
| `bin_size` | 0.05 | 分桶大小 |
| `add_nl` | True | 是否添加中性丢失特征 |
| `binary_intn` | False | 是否二值化强度 |

### 6.3 分子编码参数

| 参数 | 默认值 | 说明 |
|------|-------|------|
| `mol_encoder` | 'gnn+fp' | 编码方式 (fp/gnn/gnn+fp) |
| `mol_embedding_dim` | 2048 | Morgan 指纹维度 |
| `fptype` | 'morgan' | 指纹类型 |
| `molgnn_n_filters_list` | [256,256,256] | GNN 各层维度 |
| `molgnn_nhead` | 4 | 注意力头数 |
| `molgnn_readout_layers` | 2 | Readout 层数 |

### 6.4 Transformer 参数

| 参数 | 默认值 | 说明 |
|------|-------|------|
| `tsfm_layers` | 6 | Transformer 层数 |
| `tsfm_heads` | 8 | 注意力头数 |
| `tsfm_in_ms` | True | MS 分支使用 Transformer |
| `tsfm_in_mol` | False | 分子分支使用 Transformer |

---

## 附录: 文件结构

```
train-001/
├── train.py          # 训练入口
├── dataset.py        # 数据集定义
├── modules.py        # 模型定义 (FragSimiModelNew)
├── config.py         # 配置参数
├── utils.py          # 工具函数 (特征处理等)
├── cliplayers.py     # CLIP 风格的 Transformer
└── GNN/
    ├── layers.py     # 图卷积层
    └── featurizer.py # 分子图特征化
```

---

## 参考

本模型设计参考了以下工作:
- CLIP (Contrastive Language-Image Pre-training)
- Graph Attention Networks (GAT)
- Morgan Fingerprints (Extended Connectivity Fingerprints)
- Tanimoto Coefficient for molecular similarity

---