Generating Structurally Diverse Therapeutic Peptides with GFlowNet

This paper demonstrates that GFlowNet outperforms traditional reinforcement learning methods like GRPO in generating structurally diverse therapeutic peptides by inherently achieving uniform sequence sampling through proportional reward sampling, thereby eliminating the need for explicit diversity penalties and preventing mode collapse.

The Gist

The Big Problem: The "One-Trick Pony" of Drug Discovery

Imagine you are a chef trying to invent a new, delicious soup. You have a robot assistant (an AI) that can mix ingredients to create soup recipes. Your goal is to find the best soup possible.

However, the robot has a bad habit. It's a "perfectionist" who only wants to find the single best-tasting soup. So, it keeps making the exact same soup over and over again, just tweaking the salt by a tiny fraction. It ignores all the other delicious possibilities (spicy, sweet, creamy) because it's so obsessed with finding that one "perfect" recipe.

In the world of drug discovery, this is called Mode Collapse.

• The Goal: Find many different therapeutic peptides (tiny proteins) that could become drugs.
• The Problem: Traditional AI methods (like Reinforcement Learning) get stuck making the same few variations of a drug. Even if you tell them, "Hey, try to be diverse!" they usually ignore it and keep making the same thing. This is dangerous because if that one specific drug fails in clinical trials, you have no backup plan.

The Solution: GFlowNet (The "Taste-Tester" vs. The "Perfectionist")

The authors propose a new AI method called GFlowNet. To understand the difference, let's look at how the two robots think:

1. The Old Way (GRPO): The "Gold Digger"

• How it thinks: "I need to find the sequence with the highest score. I will ignore everything else."
• The Analogy: Imagine a gold digger who only cares about the biggest gold nugget. If they find a spot with a big nugget, they dig there forever. They ignore the smaller nuggets nearby, even though there might be 100 of them.
• The Flaw: If you try to force them to dig elsewhere by adding a "diversity penalty" (a rule saying "you must dig in other spots too"), they get confused. They fight against the rule. If you remove the rule, they immediately collapse back into digging only in one spot.

2. The New Way (GFlowNet): The "Proportional Explorer"

• How it thinks: "I will explore the whole map. If a spot has a big gold nugget, I'll visit it often. If a spot has a small nugget, I'll visit it sometimes. I won't ignore the small ones."
• The Analogy: Imagine a treasure hunter who maps the whole island. They don't just dig at the biggest pile of gold; they dig everywhere, but they spend more time digging where the gold is likely to be. They naturally visit many different spots because that's how they explore.
• The Magic: They don't need a rule telling them to be diverse. Diversity happens naturally because their strategy is to sample proportionally to the reward, not just maximize it.

The Experiment: Putting Them to the Test

The researchers put both robots to work designing therapeutic peptides. They tested them in two scenarios:

Scenario A: The "Strict" Chef (With Safety Rules)
Both robots were given a reward system that included a "diversity gate" (a rule that blocks repetitive, boring recipes).

• Result: Both robots looked good on the surface. They both produced a wide variety of soup names.
• The Catch: When the researchers looked closer (at the "ingredients"), the old robot (GRPO) was still sneaking in the same 3 ingredients over and over again. The new robot (GFlowNet) used a truly wide variety of ingredients.

Scenario B: The "Relaxed" Chef (No Safety Rules)
The researchers removed the "diversity gate" to see what happens when the rules are gone.

• The Old Robot (GRPO): Total disaster. It immediately collapsed. 100% of its recipes were the exact same repetitive pattern (like a soup that just says "Salt, Salt, Salt").
• The New Robot (GFlowNet): It kept working perfectly. It still produced a diverse, healthy mix of recipes. It didn't need the safety gate to stay diverse; it was built that way.

Why Does This Matter? (The "Structural Hedge")

Think of drug discovery like investing money.

• The Old Way: You put all your money into one stock because it looks like the best performer. If that stock crashes, you lose everything.
• The New Way (GFlowNet): You buy a portfolio of different stocks. Some are high-risk/high-reward, some are steady. If one fails, the others might succeed.

In drug discovery, we don't know exactly which chemical structure will work best in the human body. By generating a diverse portfolio of candidates (some stable, some sticky, some fast-acting), GFlowNet ensures that if one type of drug fails, we have other completely different types ready to try. This is called Structural Hedging.

The Takeaway

This paper shows that GFlowNet is a smarter way to design drugs.

• It doesn't just chase the "perfect" answer; it explores the whole landscape.
• It naturally produces a wide variety of candidates without needing complex rules to force it.
• It is more robust: even when the rules change or are removed, it doesn't break.

In short, while other AIs are like a dog chasing a single tennis ball, GFlowNet is like a dog exploring the whole park, finding balls, sticks, and leaves everywhere, giving us a much better chance of finding the next miracle drug.

Technical Summary

1. Problem Statement

Therapeutic peptide design relies heavily on generating a diverse set of structurally distinct candidates to ensure robust patent portfolios, comprehensive validation, and resilience against failure modes. However, standard Reinforcement Learning (RL) approaches for sequence generation suffer from mode collapse.

• The Core Issue: RL methods optimize to maximize expected reward ($E[R(x)]$), causing policies to converge on narrow regions of the sequence space (high-reward modes).
• Limitations of Current Solutions: Existing mitigations, such as explicit diversity penalties added to the reward function, address symptoms rather than the root cause. Fine-grained analysis reveals that even when coarse diversity metrics appear satisfactory, RL methods often exhibit hidden "mode-seeking" behavior (e.g., repetitive local motifs). When diversity mechanisms are weakened or removed, these methods collapse completely.

2. Methodology

The paper proposes using Generative Flow Networks (GFlowNet) as an alternative to reward-maximizing RL for peptide generation.

A. Theoretical Framework

• Reward Maximization (Standard RL): Optimizes $\max_\theta E_{x \sim \pi_\theta}[R(x)]$. This concentrates probability mass on the global maximum of the reward landscape, leading to mode collapse.
• GFlowNet Objective: Learns a policy $\pi_\theta$ that samples sequences proportionally to their reward:
  $$P(x) \propto R(x)$$
  Instead of seeking the single best mode, GFlowNet performs mode-covering, distributing probability mass across the reward landscape proportional to the reward magnitude. This inherently provides diversity without requiring explicit output diversity penalties.

B. Architecture and Training

• GFlowNet Model: A causal Transformer (4 layers, 256 hidden dim, 8 heads) trained from scratch.
• Training Objective: Uses Sub-Trajectory Balance (STB) loss. This variant of Trajectory Balance computes balance constraints on sub-trajectories to provide stable gradient signals for variable-length sequences.
  • Loss function aggregates over sub-trajectories, balancing the forward policy ($P_F$), backward policy ($P_B$), and a learnable partition function ($Z$).
• Key Hyperparameter: The learning rate for the log-partition function ($\log Z$) is critical. A $10\times$ multiplier relative to the policy learning rate was found necessary to prevent mode collapse (too slow) or undertraining (too fast).

C. Baseline: GRPO-D

To ensure a fair and rigorous comparison, the authors constructed a strong baseline called GRPO-D (Group Relative Policy Optimization with Diversity penalty).

• Architecture: Uses a pretrained ProtGPT2-distilled backbone (frozen embeddings), giving it a transfer learning advantage over the GFlowNet model trained from scratch.
• Objective: Maximizes expected reward but includes an explicit diversity penalty ($\lambda \cdot \text{div}(x)$) combining amino acid frequency rarity and pairwise Levenshtein distance.
• Significance: This setup tests whether GFlowNet's proportional sampling can outperform a reward-maximizing method even when the latter has architectural advantages and explicit diversity enforcement.

D. Reward Functions

The study evaluated methods under three configurations to test robustness:

1. ImprovedReward: Includes naturalness scoring + explicit entropy gating (penalizes low-entropy/repetitive sequences).
2. CompositeReward: Removes entropy gating; relies on stability, binding, and naturalness.
3. ESM2-PLL: A degenerate reward based purely on pseudo-log-likelihood (known to favor repetitive sequences).

3. Key Contributions

1. Fine-Grained Diversity Analysis: The authors introduce metrics beyond standard sequence identity to detect hidden mode-seeking behavior, including:
   • Dipeptide Concentration: Percentage of dipeptides belonging to the top-10 most frequent.
   • Dipeptide Entropy: Normalized Shannon entropy of dipeptide frequencies.
   • Consecutive Repeats: Frequency of identical consecutive amino acids.
2. Robustness Characterization: Systematic ablation studies showing that GFlowNet maintains diversity even when diversity-enforcing mechanisms (entropy gating or diversity penalties) are removed, whereas GRPO collapses.
3. Architectural Disadvantage Test: Demonstrating that GFlowNet (trained from scratch) outperforms a pretrained baseline (GRPO-D) in diversity metrics, suggesting the training objective is the dominant factor.

4. Results

A. Coarse vs. Fine-Grained Metrics

Under the ImprovedReward (with entropy gating):

• Coarse Metrics: Both methods appeared equivalent (Sequence Diversity $\approx 0.95$, 100% unique ratio).
• Fine-Grained Metrics: GFlowNet significantly outperformed GRPO-D:
  • Dipeptide Concentration: GFlowNet's top-10 dipeptides accounted for only 4.0% of samples, compared to 21.7% for GRPO-D (a 5.4$\times$ difference).
  • Reward Consistency: GFlowNet had a 1.9$\times$ lower coefficient of variation in rewards.
  • Quality Floor: GFlowNet's 5th percentile reward was 3.6% higher, meaning its "worst" samples were better than GRPO-D's worst.
  • Repetitive Motifs: GFlowNet had 3.9$\times$ fewer sequences with consecutive repeats.

B. Ablation Studies (Stress Testing)

When diversity mechanisms were removed, the difference became catastrophic for RL:

• Removing Entropy Gating (CompositeReward):
  • GRPO-D: Suffered complete mode collapse. 100% of samples contained the repetitive tripeptide pattern RMMRMMRMM.
  • GFlowNet: Maintained natural diversity (0.937) with no degenerate patterns.
• Removing Diversity Penalty (Vanilla GRPO):
  • GRPO: Top-10 dipeptides accounted for 52.5% of all dipeptides (13$\times$ worse than GFlowNet), despite coarse diversity metrics appearing acceptable (0.863).

C. Reward Robustness

GFlowNet consistently outperformed or tied GRPO-D across all reward configurations. The only scenario where both failed was the ESM2-PLL (degenerate) reward, which directly incentivized repetition, proving that proportional sampling offers robustness but not immunity to pathological reward functions.

5. Significance and Implications

• Inherent Robustness: GFlowNet's proportional sampling objective is inherently robust to reward function design. It does not require explicit diversity penalties, which often create a trade-off between reward quality and diversity.
• Structural Hedging in Drug Discovery: By covering the reward landscape rather than seeking a single mode, GFlowNet generates structurally distinct families of candidates. If one family fails due to unforeseen properties (e.g., toxicity, poor bioavailability), other structurally diverse families may succeed. This transforms a fragile single-candidate pipeline into a resilient portfolio.
• Theoretical Insight: The paper clarifies that diversity penalties in RL are merely counter-forces that resist the gradient toward high-reward modes, whereas GFlowNet changes the optimization target itself to distributional matching.
• Practical Impact: This approach accelerates therapeutic peptide discovery by providing a reliable method to generate diverse, high-quality candidates without the hyperparameter sensitivity and collapse risks associated with standard RL.

Conclusion

The study demonstrates that GFlowNet fundamentally outperforms reward-maximizing RL (even with strong baselines and diversity penalties) for therapeutic peptide generation. By sampling proportionally to reward, GFlowNet achieves superior structural diversity and robustness, eliminating the need for explicit diversity penalties and offering a more reliable framework for drug discovery pipelines.