library_name: transformers tags: []

重要说明

此前Huggingface权重绑定存在一个bug，导致ESM++的输出逻辑与ESMC不一致。该问题现已修复。

ESM++模型

ESM++是对ESMC（非商业许可协议）的忠实实现，支持批处理操作且完全兼容Huggingface生态，无需依赖ESM官方Python包。 其中large版本对应ESMC的6亿参数模型。

使用🤗 transformers调用

from transformers import AutoModelForMaskedLM
model = AutoModelForMaskedLM.from_pretrained('Synthyra/ESMplusplus_large', trust_remote_code=True)
tokenizer = model.tokenizer

sequences = ['MPRTEIN', 'MSEQWENCE']
tokenized = tokenizer(sequences, padding=True, return_tensors='pt')

# 如需进行MLM训练，需正确掩码input_ids并将未掩码的labels设为-100：
# tokenized['labels'] = tokenized['input_ids'].clone() 

output = model(**tokenized) # 获取所有隐藏状态需设置output_hidden_states=True
print(output.logits.shape) # 语言建模logits，(批次大小, 序列长度, 词表大小)，示例输出(2, 11, 64)
print(output.last_hidden_state.shape) # 最终隐藏状态，(批次大小, 序列长度, 隐藏层维度)，示例(2, 11, 1152)
print(output.loss) # 传入labels时可获取语言建模损失
#print(output.hidden_states) # 设置output_hidden_states=True可获取所有隐藏状态(元组形式)

与ESM2类似，ESM++也支持序列级和token级分类任务。初始化时指定标签数量即可：

from transformers import AutoModelForSequenceClassification, AutoModelForTokenClassification

model = AutoModelForSequenceClassification.from_pretrained('Synthyra/ESMplusplus_large', num_labels=2, trust_remote_code=True)
logits = model(**tokenized).logits
print(logits.shape) # (批次大小, 标签数量)，示例(2, 2)

ESM++默认加载fp32精度权重，支持fp16/bf16加载方式：

import torch
model = AutoModelForMaskedLM.from_pretrained('Synthyra/ESMplusplus_large', trust_remote_code=True, torch_dtype=torch.float16) # 或torch.bfloat16

零代码实现数据集嵌入

通过embed_dataset方法可快速嵌入蛋白质序列列表。系统会自动按长度排序减少填充token，初始进度条预估时间通常远长于实际耗时。

示例：

embedding_dict = model.embed_dataset(
    sequences=[
        'MALWMRLLPLLALLALWGPDPAAA', ... # 蛋白质序列列表
    ],
    tokenizer=model.tokenizer,
    batch_size=2, # 根据GPU显存调整
    max_len=512, # 按需设置
    full_embeddings=False, # True时不进行池化
    embed_dtype=torch.float32, # 指定输出精度
    pooling_types=['mean', 'cls'], # 多池化类型将拼接输出
    num_workers=0, # 多CPU核心环境下建议设为4加速处理
    sql=False, # True时嵌入结果存入SQLite数据库
    sql_db_path='embeddings.db',
    save=True, # True时嵌入结果保存为.pth文件
    save_path='embeddings.pth',
)
# 返回字典结构：序列->嵌入张量(.pth)或numpy数组(sql)

model.embed_dataset()参数说明：
    sequences: 蛋白质序列列表
    batch_size: 处理批次大小
    max_len: 最大序列长度
    full_embeddings: 返回完整残基级嵌入(True)或池化结果(False)
    pooling_type: 池化类型('mean'或'cls')
    num_workers: 数据加载线程数，0表示主进程处理
    sql: 是否存入SQLite数据库(强制float32精度)
    sql_db_path: 数据库路径
    
返回：
    序列到嵌入向量的映射字典(sql=True时返回None)

注意：
    - SQL模式适合需要实时流式处理超大规模训练数据的场景
    - save模式适合能全量加载嵌入字典到内存的情况
    - 当sql和save同时设置时优先使用SQL存储
    - 会自动检测已有数据库/pth文件中的缓存序列
    - 序列将按长度降序排列并截断至max_len以优化处理速度

使用🤗 peft微调

model = AutoModelForSequenceClassification.from_pretrained('Synthyra/ESMplusplus_large', num_labels=2, trust_remote_code=True)
# 以下模块适配ESM++和ESM2的注意力层结构
target_modules = ["layernorm_qkv.1", "out_proj", "query", "key", "value", "dense"]

lora_config = LoraConfig(
    r=8, # 可自定义LoRA参数
    lora_alpha=16,
    lora_dropout=0.01,
    bias="none",
    target_modules=target_modules,
)

# 应用LoRA适配器
model = get_peft_model(model, lora_config)

# 解冻分类头参数
for param in model.classifier.parameters():
    param.requires_grad = True

完整微调示例参见脚本。

获取注意力图谱

默认使用F.scaled_dot_product_attention加速计算，但无法返回注意力图谱。设置output_attentions=True将切换为手动计算模式（速度较慢，仅需时启用）。

output = model(**tokenized, output_attentions=True)
att = output.attentions
len(att) # 33个注意力矩阵(每层一个)，形状均为(批次大小, 头数, 序列长度, 序列长度)

精度与实现对比

我们测量了fp32权重在fp16/bf16下的隐藏状态差异，发现fp16更接近fp32输出。需注意ESM官方包默认将ESMC转为bf16，各有优劣，可按需选择半精度类型。

FP16平均MSE: 0.00000003

BF16平均MSE: 0.00000122

在1000条随机序列上对比ESM++与ESMC(均为bfloat16)输出，确保与官方包一致性：

最终隐藏状态平均MSE: 2.46e-09

可通过.from_pretrained_esm('esmc_600m')直接加载ESM官方包权重。

模型探针分析

我们采用类似先前研究的线性探针方法评估隐藏状态与生物特性的关联性。ESMC（及ESM++）表现优异。

下图显示各模型在标准化指标下的表现（分类任务取MCC与F1平均值，回归任务取Spearman ρ与R²平均值）：

推理速度基准

在H100显卡上测试不同ESM模型的吞吐量。相比ESMC，ESM++的批处理优化带来显著加速，单批次推理也更快。ESM++ small处理长序列时甚至快于ESM2-35M！Linux系统搭配PyTorch > 2.5可获得最佳性能。

引用声明

使用本实现或相关研究成果时请引用（同时需引用ESMC预印本）：

@misc {ESMPlusPlus,
	author       = { Hallee, L. and Bichara, D. and Gleghorn, J, P. },
	title        = { ESMPlusPlus },
	year         = 2024,
	url          = { https://huggingface.co/Synthyra/ESMplusplus_small },
	doi          = { 10.57967/hf/3726 },
	publisher    = { Hugging Face }
}