跳转至

Steering Protein Language Models

会议: ICML 2025
arXiv: 2509.07983
代码: 无(Tencent AI Lab)
领域: 蛋白质工程 / 计算生物学
关键词: 蛋白质语言模型, Activation Steering, 蛋白质优化, 突变位点识别, 无训练控制

一句话总结

首次将LLM领域的Activation Steering技术迁移到蛋白质语言模型(PLM),通过在推理时编辑模型内部激活来引导蛋白质序列生成和优化朝向目标属性(如热稳定性、溶解度),完全无需重新训练,并提出基于steering vector相异度的突变位点识别算法(ASPO),在溶菌酶和GFP优化任务上大幅超越传统方法。

研究背景与动机

领域现状: 蛋白质语言模型(PLM)通过在数十亿进化序列上预训练,已成为蛋白质设计的核心工具。PLM分为自编码(AE)型如ESM2/ESM3和自回归(AR)型如ProLLaMA,在突变效应预测、结构推断等任务上表现优异,但难以精确控制输出序列使其具有特定功能属性。

现有痛点: 控制PLM输出的方法存在明显缺陷:(1) 微调需要大量标注数据和计算资源,且有遗忘预训练知识的风险;(2) 基于关键词标签的prompt控制缺乏灵活性,受限于预训练阶段使用的标签;(3) 搜索/采样方法(如AdaLead、PEX)效率低下,需要大量序列生成后筛选,还依赖适应度预测器。

核心矛盾: PLM内部已编码了丰富的蛋白质属性知识(如热稳定性、溶解度信息),但这些知识并不总是在输出中得到体现。如何在不修改模型权重的情况下,精准地"释放"这些隐含知识?

本文目标: 探索一种推理时干预方法,在不需要训练或权重更新的前提下,控制PLM生成具有目标属性的蛋白质序列,并扩展到蛋白质优化场景中的突变位点识别和定向突变。

切入角度: 从LLM中已成功应用的Activation Steering技术出发,将其适配到PLM——蛋白质序列反映生物物理功能而非语言语义,激活空间受进化和结构约束塑造,需要特殊处理。

核心 idea: 通过计算具有/不具有目标属性的蛋白质在PLM内部表示的平均差异作为steering vector,在推理时将其加到激活上即可引导生成方向,并利用steering vector与token表示的余弦相异度来识别突变位点。

方法详解

整体框架

方法分为两个层次:(1) Activation Steering用于蛋白质生成控制——在PLM每层的激活上添加steering vector来偏置生成方向,适用于AR-PLM(自回归生成)和AE-PLM(掩码位置预测);(2) ASPO用于蛋白质优化——先通过steering vector识别需要突变的位点,再用activation steering引导突变预测,多轮迭代逐步提升目标属性。

关键设计

  1. Steering Vector计算与激活编辑

    • 功能:计算代表目标属性方向的向量,并在推理时注入模型各层激活
    • 核心思路:收集目标属性的正集\(\mathcal{P}\)和负集\(\mathcal{N}\)(各100个序列),对每层\(l\)计算 \(\mathbf{v}_l = \frac{1}{|\mathcal{P}|}\sum_{x_p \in \mathcal{P}} \mathbf{h}_l^{avg}(x_p) - \frac{1}{|\mathcal{N}|}\sum_{x_n \in \mathcal{N}} \mathbf{h}_l^{avg}(x_n)\)(AE-PLM用全token平均,AR-PLM用最后token)。推理时每层执行 \(\tilde{\mathbf{h}}_l = \mathbf{h}_l + \alpha \mathbf{v}_l\),并rescale到原始范数
    • 设计动机:均值差异提取的是属性在表示空间中的线性方向,rescale避免激活幅度偏移破坏模型内部动态
  2. 基于Steering Vector相异度的突变位点识别

    • 功能:从蛋白质序列中自动识别与目标属性最不相关的氨基酸位点作为突变候选
    • 核心思路:计算每个token表示与steering vector的余弦相似度 \(s^k = \frac{\mathbf{v}_l^\top \mathbf{h}_l^k}{||\mathbf{v}_l|| \cdot ||\mathbf{h}_l^k||}\),选择得分最低的\(T\)个位点。通过在正/负集上训练线性分类器验证各层判别力,选最优层计算得分
    • 设计动机:相似度低意味着该位点的当前氨基酸与目标属性方向相悖,是最值得改变的位置。比随机选择更有针对性,且不需要额外的适应度预测器
  3. ASPO多轮迭代优化

    • 功能:通过多轮mask-then-predict流程逐步将蛋白质优化到目标属性
    • 核心思路:每轮执行:计算所有token的相异度得分→选\(T\)个最低分位点→mask这些位点→用activation steering引导PLM预测新氨基酸。重复\(R\)轮(热稳定性实验\(R=8, T=4\);溶解度/GFP实验\(R=4, T=2\)
    • 设计动机:分轮优化避免一次性改变过多位点导致结构崩溃,每轮重新计算得分使选择自适应更新

损失函数 / 训练策略

本方法完全不涉及模型训练或权重更新。所有操作发生在推理时。微调基线:AE-PLM微调最后一层,AR-PLM使用LoRA(rank=4, alpha=16)在所有层微调。ASPO基线方法(AdaLead、PEX、GWG)需要训练适应度预测器。

实验关键数据

蛋白质生成实验(溶菌酶,1000条序列)

基础模型 方法 热稳定性↑ 多样性↑ 新颖性↑ 溶解度↑
ProLLaMA Original 56.18 (8.05) 0.931 0.767 0.230
ProLLaMA Fine-tuning 57.24 (8.64) 0.958 0.798 0.241
ProLLaMA Act. Steering 67.68 (12.86) 0.927 0.807 0.276
ESM2 Original 56.48 (12.04) 0.954 0.591 0.289
ESM2 Fine-tuning 63.56 (14.87) 0.953 0.585 0.356
ESM2 Act. Steering 82.20 (12.92) 0.971 0.739 0.494
ESM3 Original 55.20 (11.14) 0.952 0.573 0.257
ESM3 Fine-tuning 62.82 (14.72) 0.949 0.568 0.318
ESM3 Act. Steering 82.06 (12.06) 0.954 0.614 0.582

蛋白质优化实验(热稳定性 + GFP亮度)

方法 热稳定性-Medium↑ 热稳定性-Hard↑ GFP-Medium↑ GFP-Hard↑
AdaLead 63.56 55.16 1.179 1.255
PEX 66.80 48.95 1.426 1.320
GWG 68.25 47.73 1.683 1.510
ESM2+ASPO 84.34 74.69 3.862 3.907
ESM3+ASPO 88.42 86.43 3.739 3.687

关键发现

  • Activation Steering在ESM2上将热稳定性从56.48提升到82.20(+46%),远超微调的63.56
  • ESM3+ASPO在hard难度热稳定性优化中达86.43,是GWG(47.73)的1.8倍
  • GFP亮度优化中ESM2+ASPO达3.862,比最佳基线GWG的1.683翻倍以上
  • 正/负集大小100时AE-PLM达≥95%最大性能,AR-PLM在10个样本时即最优
  • 推荐默认α=1.0,覆盖90-98%最大性能,溶解度任务在α>5时性能崩溃
  • 生成的序列保持甚至提升了多样性和新颖性,说明方法不是简单记忆正集

亮点与洞察

  • 跨领域迁移的优雅成功:NLP→蛋白质工程的迁移非常自然。t-SNE和线性探针证实PLM激活空间中已编码属性信息,activation steering只是释放了这些知识
  • 零训练成本的巨大实用价值:仅需100对正/负样本、一个超参α,即可在任意PLM上实现属性导向生成,无需训练
  • 突变位点识别算法的巧妙:利用steering vector本身定位突变位点,取消了对适应度预测器的依赖,减少了一整个组件
  • ASPO在GFP上的惊人表现:亮度从1.3-1.5优化到3.7-3.9,翻倍提升,展示了在真实蛋白质工程任务上的强大潜力

局限与展望

  • Steering vector基于均值差异,是线性方向,可能无法捕获复杂的非线性属性-表示关系
  • 属性评估依赖计算预测器(非湿实验),预测器本身误差可能影响结论可靠性
  • 多属性同时steering可能产生向量冲突,论文仅在附录简要提及
  • 仅在溶菌酶和GFP两个蛋白质家族上验证,更广泛的蛋白质空间覆盖待测试
  • AR-PLM(ProLLaMA)从大对比集反而下降,暗示last-token表示的局限性

相关工作与启发

  • Activation Addition (Turner et al., 2023): 本文的直接灵感来源,使用对比提示计算steering vector控制LLM
  • CAA (Panickssery et al., 2023): 从数百对对比中聚合steering vector降低噪声
  • GGS (Kirjner et al., 2023): 用能量模型平滑适应度景观+Gibbs采样做蛋白质优化,ASPO的主要对比基线
  • ESM2/ESM3/ProLLaMA: 三个被steering的基线PLM,分别代表AE和AR架构
  • 启发:PLM激活空间的线性结构比预想中更规整,简单的线性操作即可实现精确控制

评分

  • 新颖性: ⭐⭐⭐⭐(首次跨领域迁移,突变位点识别算法是蛋白质特异性创新)
  • 实验充分度: ⭐⭐⭐⭐⭐(3架构×3属性×2任务×2难度,超参敏感性分析详尽)
  • 写作质量: ⭐⭐⭐⭐(方法清晰,图示直观,流程算法伪代码完整)
  • 价值: ⭐⭐⭐⭐(为蛋白质设计提供即插即用的零成本控制方案,实用性强)

相关论文