清華大學(xué)發(fā)布ProteinOPD：面向蛋白質(zhì)設(shè)計(jì)高效多目標(biāo)偏好對(duì)齊框架

作者 | 論文團(tuán)隊(duì)

編輯丨ScienceAI

當(dāng)前，蛋白質(zhì)語(yǔ)言模型正從「生成合理蛋白序列」逐步走向「按需生成具有特定性質(zhì)的蛋白序列」。在合成生物學(xué)和藥物發(fā)現(xiàn)等場(chǎng)景中，研究者往往不只希望模型生成結(jié)構(gòu)合理的蛋白，還希望候選序列同時(shí)具備更高折疊性、更好溶解性、更強(qiáng)熱穩(wěn)定性等多種屬性。然而，多目標(biāo)蛋白質(zhì)偏好對(duì)齊長(zhǎng)期面臨核心挑戰(zhàn)：在提升目標(biāo)屬性的同時(shí)，模型容易遺忘預(yù)訓(xùn)練階段獲得的蛋白設(shè)計(jì)能力，并且難以在多個(gè)相互競(jìng)爭(zhēng)的偏好目標(biāo)之間取得穩(wěn)定平衡。

近日，清華大學(xué)高子琪團(tuán)隊(duì)聯(lián)合 IDEA AI4S 團(tuán)隊(duì)提出 ProteinOPD，該工作的主要貢獻(xiàn)如下：

1. 解決蛋白質(zhì)偏好對(duì)齊中的多目標(biāo)難點(diǎn)。ProteinOPD 是一個(gè)面向蛋白質(zhì)設(shè)計(jì)的高效多目標(biāo)偏好對(duì)齊框架，在提升目標(biāo)屬性的同時(shí)，保留蛋白語(yǔ)言模型原有的設(shè)計(jì)能力與生成分布質(zhì)量。

2. 提出多教師 OPD 的一種有效實(shí)現(xiàn)形式。該方法將預(yù)訓(xùn)練蛋白語(yǔ)言模型適配為多個(gè)偏好專屬教師，并通過(guò) normalized Product-of-Experts 構(gòu)造幾何共識(shí)分布，使學(xué)生模型能夠在自身生成軌跡上接受 token-level on-policy distillation，從而將 OPD 從單教師、單目標(biāo)設(shè)置擴(kuò)展到多教師、多目標(biāo)偏好對(duì)齊場(chǎng)景。

3. 提供較完整的開(kāi)源支持。項(xiàng)目已開(kāi)放論文、代碼、項(xiàng)目主頁(yè)和 Colab 體驗(yàn)入口，覆蓋無(wú)條件生成與條件生成兩類設(shè)置，并提供教師構(gòu)建、ProteinOPD 訓(xùn)練、序列生成和推理體驗(yàn)等關(guān)鍵流程，便于社區(qū)復(fù)現(xiàn)與后續(xù)擴(kuò)展。

論文地址：https://arxiv.org/abs/2605.10189

開(kāi)源地址：https://github.com/THU-AI4S/ProteinOPD

Colab 鏈接：https://colab.research.google.com/github/THU-AI4S/ProteinOPD/blob/main/notebooks/proteinopd_inference.ipynb

在模型表現(xiàn)方面，ProteinOPD 在多目標(biāo)場(chǎng)景下取得了最優(yōu)偏好對(duì)齊表現(xiàn)，同時(shí)對(duì)設(shè)計(jì)能力造成的損傷最小。以 ProtGPT2 為基礎(chǔ)模型時(shí)，ProteinOPD 分別將 foldability、solubility 和 thermostability 提升 14.8%、16.9% 與 54.2%；同時(shí)，其偏好對(duì)齊性能可在約 1/8 的訓(xùn)練時(shí)間內(nèi)達(dá)到 RL 方法的水平。

方法設(shè)計(jì)

ProteinOPD 的思路是將「偏好獲取」和「偏好組合」解耦。首先，針對(duì)折疊性、溶解性、熱穩(wěn)定性等目標(biāo)偏好，框架使用屬性 oracle 對(duì)蛋白序列進(jìn)行打分，并選取少量高分樣本構(gòu)建偏好專屬訓(xùn)練集。隨后，通過(guò)輕量化適配，將預(yù)訓(xùn)練蛋白語(yǔ)言模型轉(zhuǎn)化為多個(gè)偏好專屬教師。每個(gè)教師負(fù)責(zé)提供一種偏好的生成方向，而不是要求單個(gè)模型同時(shí)承擔(dān)全部目標(biāo)。

在學(xué)生訓(xùn)練階段，ProteinOPD 不采用傳統(tǒng)離線模仿方式，而是在學(xué)生模型自己的生成軌跡上進(jìn)行 on-policy distillation。具體而言，學(xué)生模型先生成當(dāng)前會(huì)訪問(wèn)的蛋白序列前綴，教師模型再基于這些前綴給出下一 token 的概率分布。相比只在固定數(shù)據(jù)集上訓(xùn)練，這種機(jī)制能夠讓學(xué)生在自身實(shí)際訪問(wèn)的狀態(tài)空間中獲得糾正信號(hào)，從而緩解訓(xùn)練階段與生成階段的分布不一致問(wèn)題。同時(shí)，token-level 分布監(jiān)督也比序列級(jí)獎(jiǎng)勵(lì)更密集，有助于提升訓(xùn)練效率和穩(wěn)定性。

對(duì)于多目標(biāo)偏好對(duì)齊，ProteinOPD 的關(guān)鍵設(shè)計(jì)是多教師 OPD。研究團(tuán)隊(duì)將多個(gè)偏好教師的 next-token 分布通過(guò) normalized Product-of-Experts 進(jìn)行組合，形成一個(gè)幾何共識(shí)目標(biāo)。與算術(shù)平均不同，幾何共識(shí)更強(qiáng)調(diào)多個(gè)教師共同支持的 token：只有當(dāng)多個(gè)偏好教師都傾向于某個(gè)氨基酸選擇時(shí)，該 token 才會(huì)在共識(shí)分布中獲得更高權(quán)重。這一設(shè)計(jì)使模型更傾向于尋找多種偏好之間的交集，而不是被某一個(gè)目標(biāo)單獨(dú)牽引。

這一機(jī)制還帶來(lái)一個(gè)重要性質(zhì)：歸一化項(xiàng)可以反映教師之間的分歧程度。當(dāng)多個(gè)教師在某個(gè)序列前綴上意見(jiàn)一致時(shí)，目標(biāo)偏好之間沖突較小；當(dāng)教師分布明顯不一致時(shí)，沖突程度增大。因此，ProteinOPD 不僅提供了多教師偏好組合的訓(xùn)練目標(biāo)，也為觀察多屬性之間的內(nèi)部競(jìng)爭(zhēng)提供了一個(gè)無(wú)需額外計(jì)算的信號(hào)。

模型性能

1. 多目標(biāo)偏好對(duì)齊取得帕累托最優(yōu)。(1) 相比評(píng)估中最強(qiáng)的跨范式參考方法 MoMPNN，ProteinOPD 的 HV 提升 34.8%。(2) 與基礎(chǔ)模型 ProtGPT2 相比，ProteinOPD 分別將蛋白可折疊性、溶解性和熱穩(wěn)定性提升 14.8%、16.9% 與 54.2%。

論文主表：無(wú)條件設(shè)置下多目標(biāo)偏好對(duì)齊性能對(duì)比。

2. 單目標(biāo)實(shí)驗(yàn)顯示 ProteinOPD 能夠轉(zhuǎn)移 SFT 教師的偏好能力同時(shí)緩解直接微調(diào)造成的蛋白設(shè)計(jì)能力的下降。(1) 在無(wú)條件生成中，ProteinOPD 在保留大部分屬性收益的同時(shí)，將新穎性損失控制在更低水平。(2) 在條件生成中，ProteinOPD 還提升了 ProTrek Score，說(shuō)明模型在偏好對(duì)齊過(guò)程中沒(méi)有破壞原有條件一致性。

3. 訓(xùn)練成本大幅降低。(1) ProteinOPD 相比 RL 類方法可以更快提升熱穩(wěn)定性，并取得約 8 倍訓(xùn)練加速。 (2) 教師構(gòu)建只需要少量經(jīng) oracle 篩選的高質(zhì)量樣本，使整體流程更適合實(shí)際蛋白設(shè)計(jì)中的快速迭代需求。

樣例研究

論文進(jìn)一步給出了 ProteinOPD 與 ASPO 的生成樣例對(duì)比。為確保比較集中在多屬性對(duì)齊能力上，作者選取熱穩(wěn)定性均超過(guò) 0.95 的樣例，并要求生成序列與 UniRef 的最大序列一致性低于 5%。在可視化結(jié)果中，ProteinOPD 生成樣例的 pLDDT 達(dá)到 0.73，高于 ASPO 的 0.49；溶解性得分達(dá)到 0.69，高于 ASPO 的 0.43。這表明 ProteinOPD 不僅能提升單一屬性，還能在保持新穎性的同時(shí)實(shí)現(xiàn)更有效的多屬性對(duì)齊。

ProteinOPD 為蛋白質(zhì)偏好對(duì)齊提供了一條高效路徑：用輕量教師學(xué)習(xí)單一偏好，用多教師幾何共識(shí)處理目標(biāo)沖突，再通過(guò)學(xué)生自身軌跡上的 token-level OPD 完成穩(wěn)定遷移。該工作將 OPD 從語(yǔ)言模型偏好對(duì)齊拓展到蛋白質(zhì)設(shè)計(jì)場(chǎng)景，并進(jìn)一步推進(jìn)到多目標(biāo)、多教師設(shè)置，為生成式蛋白設(shè)計(jì)中的高效偏好對(duì)齊提供了新的技術(shù)參考。