The Forgetting-Learning Trade-off: Making Reinforcement Learning Work for Protein Language Models

Abstract

Reinforcement learning (RL) is increasingly applied to Protein Language Models (PLMs), yet its effectiveness varies across tasks, and standard metrics such as pass@k can rise even when the model's solvable problem set is shrinking. We introduce two capability-level diagnostics. The Expansion-Shrinkage Ratio (ESR) measures how RL shifts the set of problems a PLM can solve, separating genuine gain from probability redistribution. Dual-Reward ESR reports ESR under both the training reward and an orthogonal evaluator; the gap ΔESR quantifies reward hacking as a single observable number. Applied across four protein design tasks, three RL algorithms (DPO, PPO, GRPO), and two PLM architectures, the diagnostics reveal that RL on PLMs is governed by two reward properties: verifiability, whether the reward is a fixed environment or a learned surrogate vulnerable to distribution shift, and coverage, the fraction of sequence space giving an informative gradient. The two axes produce three regimes with distinct ΔESR signatures: well-covered verifiable rewards yield genuine expansion; sparse verifiable rewards induce a coverage bottleneck; predicted rewards induce reward hacking. Controlled analyses isolate these two factors as operative, letting practitioners predict an RL run's outcome before committing to costly fine-tuning.