From Supervision to Exploration: What Does Protein Language Model Learn During Reinforcement Learning?

TL;DR

The paper evaluates how reinforcement learning (RL) interacts with protein language models (PLMs) to enhance protein design across four domains: antimicrobial peptides, kinase mutations, antibodies, and inverse folding. By applying multiple RL algorithms (DPO, PPO, GRPO) to domain-specific PLMs, the study demonstrates consistent improvements in sampling efficiency and high-reward region exploration when rewards are accurate and policy capacity is sufficient. However, RL can reduce diversity and focus exploration on high-fitness regions, particularly in rugged landscapes or with noisy rewards, highlighting a trade-off between exploitation and exploration. The results provide practical guidelines for RL in protein design, such as prioritizing reward fidelity and calibrating algorithm capacity to task difficulty, with an implementation available at the authors' GitHub.

Abstract

Protein language models (PLMs) have advanced computational protein science through large-scale pretraining and scalable architectures. In parallel, reinforcement learning (RL) has broadened exploration and enabled precise multi-objective optimization in protein design. Yet whether RL can push PLMs beyond their pretraining priors to uncover latent sequence-structure-function rules remains unclear. We address this by pairing RL with PLMs across four domains: antimicrobial peptide design, kinase variant optimization, antibody engineering, and inverse folding. Using diverse RL algorithms and model classes, we ask if RL improves sampling efficiency and, more importantly, if it reveals capabilities not captured by supervised learning. Across benchmarks, RL consistently boosts success rates and sample efficiency. Performance follows a three-factor interaction: task headroom, reward fidelity, and policy capacity jointly determine gains. When rewards are accurate and informative, policies have sufficient capacity, and tasks leave room beyond supervised baselines, improvements scale; when rewards are noisy or capacity is constrained, gains saturate despite exploration. This view yields practical guidance for RL in protein design: prioritize reward modeling and calibration before scaling policy size, match algorithm and regularization strength to task difficulty, and allocate capacity where marginal gains are largest. Implementation is available at https://github.com/chq1155/RL-PLM.

Figures (7)

• Figure 1: Reinforcement learning for protein design is akin to hill climbing. Task difficulty equates to mountain height, policy model capacity to the starting altitude, and reward accuracy to direction correctness. These three factors jointly determine the RL efficacy in protein design.
• Figure 2: Overview of the four biological systems based on PLM and RL. AR and MLM denote Auto-regressive and Masked Language Modeling, respectively.
• Figure 3: Pass@k results for four biological systems.
• Figure 4: Experimental results for inverse folding (A-D) and AMP design (E-H).
• Figure 5: Experimental results for kinase mutation (A-D) and antibody optimization (E-H).