https://github.com/westlake-repl/ProTrek/blob/main/model/ProTrek/protein_encoder.py
这个蛋白质编码器的代码展示了如何将变长蛋白质序列投影到固定维度的统一空间。
1. 模型架构概览
class ProteinEncoder(torch.nn.Module):def __init__(self, config_path, out_dim, load_pretrained=True):config = EsmConfig.from_pretrained(config_path)if load_pretrained:self.model = EsmForMaskedLM.from_pretrained(config_path)else:self.model = EsmForMaskedLM(config)self.out = torch.nn.Linear(config.hidden_size, out_dim) # 关键投影层
核心组件：
- ESM模型：蛋白质语言模型，处理变长序列
- 投影层：
self.out线性层，将ESM输出投影到统一维度
2. 前向传播的详细流程
def forward(self, inputs: dict, get_mask_logits: bool = False):# 步骤1: 通过ESM模型获取隐藏状态last_hidden_state = self.model.esm(**inputs).last_hidden_state# 形状: [batch_size, seq_len, hidden_size] (如 [32, 512, 1280])# 步骤2: 使用[CLS] token作为序列表示reprs = last_hidden_state[:, 0, :] # 取第一个token ([CLS])# 形状: [batch_size, hidden_size] (如 [32, 1280])# 步骤3: 投影到统一维度空间reprs = self.out(reprs) # 线性投影# 形状: [batch_size, out_dim] (如 [32, 1024])
3. 关键投影步骤详解
3.1 从变长到固定的维度转换
输入序列的多样性：
# 不同长度的蛋白质序列示例
seq1 = "MALWMRLLPLLALLALWGPDPAAAFVNQHLCGSHLVEALYLVCGERGFFYTPKT" # 长度54
seq2 = "MVHLSREEKQ" # 长度10
seq3 = "MKWVTFISLLFLFSSAYSRGVFRRDAHKSEVAHRF" # 长度38
ESM模型处理后的输出：
# 经过ESM Transformer后的隐藏状态
last_hidden_state = self.model.esm(input_ids, attention_mask).last_hidden_state
# 对于batch_size=3的不同长度序列：
# 形状: [3, max_seq_len, hidden_size] 如 [3, 512, 1280]
# 实际有效部分由attention_mask标识
3.2 [CLS] Token池化策略
其实就是拿第一个Token!
reprs = last_hidden_state[:, 0, :] # 取每个序列的第一个token
为什么使用[CLS] token？
- 特殊标记：ESM在序列开头添加[CLS] token，专门用于表示整个序列
- 全局信息：通过自注意力机制，[CLS] token聚合了整个序列的信息
- 标准做法：这是BERT/ESM等Transformer模型的常用做法
3.3 维度投影的具体数学操作
# 投影层的具体计算
reprs = self.out(reprs) # 等价于: reprs = torch.matmul(reprs, self.out.weight.T) + self.out.bias
维度变换：
[batch_size, hidden_size] × [hidden_size, out_dim] = [batch_size, out_dim][32, 1280] × [1280, 1024] = [32, 1024]
4. 批量推理时的处理流程
def get_repr(self, proteins: list, batch_size: int = 64, verbose: bool = False):for i in iterator:# 步骤1: Tokenization和填充protein_inputs = self.tokenizer.batch_encode_plus(proteins[i:i + batch_size],return_tensors="pt",padding=True # 关键：自动填充到相同长度)# 生成attention_mask标识有效位置# 步骤2: 前向传播获取表示output, _ = self.forward(protein_inputs)# output形状: [batch_size, out_dim]# 步骤3: L2归一化protein_repr.append(output)protein_repr = torch.cat(protein_repr, dim=0)return normalize(protein_repr, dim=-1) # 最终L2归一化
5. 处理不同序列长度的关键技术
5.1 动态填充（Padding）
# tokenizer自动处理不同长度
protein_inputs = self.tokenizer.batch_encode_plus(proteins, padding=True, # 短序列填充到batch内最长序列return_tensors="pt"
)
# 生成attention_mask区分真实token和填充token
5.2 注意力掩码（Attention Mask）
# ESM模型内部使用attention_mask
last_hidden_state = self.model.esm(input_ids=inputs["input_ids"],attention_mask=inputs["attention_mask"] # 忽略填充位置
).last_hidden_state
5.3 [CLS] Token的智能性
由于自注意力机制，[CLS] token能够"看到"整个序列的所有有效token，而忽略填充部分。
6. 完整的维度变换流程
变长蛋白质序列↓ Tokenization + 填充
固定长度token序列: [batch_size, max_seq_len] ↓ ESM Transformer
序列隐藏状态: [batch_size, max_seq_len, hidden_size]↓ 取[CLS] token
序列级表示: [batch_size, hidden_size] ↓ 线性投影
统一空间表示: [batch_size, out_dim]↓ L2归一化
最终表示: [batch_size, out_dim] (模长为1)
7. 为什么这个设计有效？
7.1 长度不变性
- 无论蛋白质多长，最终都投影到相同维度
- [CLS] token通过自注意力学习序列的全局表示
7.2 语义保持
- ESM作为蛋白质语言模型，理解氨基酸的语义
- 投影层学习将语义信息映射到对比学习空间
7.3 计算效率
- 避免了复杂的池化操作
- 支持批量处理不同长度的序列
总结
这个蛋白质编码器通过ESM模型 + [CLS]池化 + 线性投影的三步策略，巧妙地将任意长度的蛋白质序列转换为固定维度的向量表示，为多模态对比学习提供了统一的表示空间。
