🤖 machine learning

FORGE: Fragment-Oriented Ranking and Generation for Context-Aware Molecular Optimization

FORGE is a two-stage framework that reformulates molecular optimization as context-aware local editing by ranking and generating fragment replacements using automatically mined data, thereby outperforming existing language and graph-based methods while avoiding the hallucinations and scaling limitations of natural language training.

Original authors: Qingchuan Zhang, He Cao, Hao Li, Yanjun Shao, Zhiyuan Liu, Shihang Wang, Shufang Xie, Shenghua Gao, Xinwu Ye

Published 2026-05-12

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Qingchuan Zhang, He Cao, Hao Li, Yanjun Shao, Zhiyuan Liu, Shihang Wang, Shufang Xie, Shenghua Gao, Xinwu Ye

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are a master chef trying to improve a famous soup recipe. You want to make it taste better (optimize it), but you can't completely change the recipe; it still needs to taste like the original soup (maintain similarity).

Most recent AI methods try to do this by asking a giant language model: "Here is a soup. Make it tastier." The AI then tries to rewrite the entire recipe from scratch based on that vague instruction. The paper argues this approach is flawed for two main reasons:

Vague Instructions: In real science, the "taste" (the goal) is often a black box. You can't describe it perfectly in words (e.g., "make it bind to this specific protein"). The AI gets confused because the words don't match the reality.
Context Matters: A specific ingredient (like a pinch of salt) might make one soup delicious but ruin another. AI models that treat ingredients as isolated facts miss this crucial context.

Enter FORGE (Fragment-Oriented Ranking and Generation). Instead of rewriting the whole soup, FORGE acts like a surgical chef who knows exactly which small piece to swap and what to swap it with.

Here is how FORGE works, broken down into simple steps:

1. The "Atomic" Vocabulary (The Stable Ingredients)

Standard AI models read chemical formulas like text, where the same ingredient might be spelled differently depending on the sentence. FORGE uses a special "atomic" vocabulary. Think of this as giving every single atom and chemical group a unique, unchangeable ID card. No matter what soup it's in, the "salt" is always recognized as "salt." This prevents the AI from getting confused about what it's actually looking at.

2. Stage 1: The "Where" (Ranking the Weak Spots)

Before making changes, FORGE looks at the whole soup and asks: "Which specific ingredient is dragging the flavor down right now?"

The Old Way: The AI guesses based on general rules.
FORGE's Way: It looks at the entire context. It knows that a specific spice is bad only because of the other ingredients currently in the pot. It ranks the ingredients to find the "weakest link" that needs fixing.

3. Stage 2: The "How" (The Surgical Swap)

Once the weak link is found, FORGE doesn't rewrite the whole recipe. It performs a precise swap: "Take out this specific spice and replace it with this specific herb."

It learns this by looking at millions of "before and after" examples found in chemical databases (like a library of successful recipe tweaks).
Crucially, it learns that the best replacement depends on the specific context of the current soup, not just a generic rule.

4. Learning by Example (The "Show, Don't Tell" Approach)

When FORGE faces a brand-new goal it has never seen before (like optimizing for a new, unknown protein), it doesn't rely on a text description. Instead, it uses In-Context Learning.

Imagine you are playing a video game with a new, unknown boss. Instead of reading a manual, you watch a replay of a pro player beating a similar boss.
FORGE keeps a "replay buffer" of its own past attempts. It looks at a few examples of "Molecule A was tweaked to Molecule B, and the score went up." It uses these examples to figure out how to tweak the current molecule, without needing to understand the complex science behind the goal in words.

Why It Wins

The paper tested FORGE against much larger, more powerful AI models (some with 8 billion parameters, while FORGE uses only 0.6 billion).

The Result: FORGE consistently beat the giants.
The Reason: The giants were trying to "hallucinate" a perfect new molecule from a vague text prompt. FORGE succeeded because it focused on small, context-aware edits using a stable vocabulary. It treated molecular optimization as a series of precise, local fixes rather than a creative writing exercise.

In summary: FORGE is a tiny, highly efficient AI that doesn't try to rewrite the whole story. Instead, it finds the one sentence that needs fixing, understands the context of the paragraph, and swaps in the perfect word, learning from past examples rather than vague instructions.

Technical Summary: FORGE (Fragment-Oriented Ranking and Generation)

1. Problem Statement

Molecular optimization aims to improve a starting molecule's properties through small structural edits while maintaining high similarity to the original compound. Current state-of-the-art approaches typically frame this as a prompt-conditioned sequence generation task using large language models (LLMs). The authors argue this formulation is fundamentally misaligned with the local-edit nature of the problem for two reasons:

Semantic Mismatch: Real-world drug discovery objectives (e.g., proprietary binding scores, ADMET classifiers) often lack faithful natural language descriptions. Relying on text prompts leads to weak control, and training on LLM-distilled chain-of-thought data often inherits chemical inaccuracies or hallucinations.
Ignored Context Dependence: The contribution of a molecular fragment to a property is strongly dependent on its local chemical environment. Treating fragment contributions as global, context-free scalars collapses critical structural information and introduces label noise.

2. Methodology: FORGE Framework

FORGE (Fragment-Oriented Ranking and Generation) reformulates molecular optimization as context-aware local editing. It decomposes the task into two explicit decisions: where to edit and how to edit. The framework is built on a compact 0.6B language model (Qwen3-0.6B) and operates in two supervised stages, followed by an inference mechanism adapted for black-box objectives.

Core Components

Atom-Level Tokenization (QwenAtom): To ensure stable fragment identity across different host molecules, the authors replace standard Byte-Pair Encoding (BPE) with a regex-based atomic parser for SMILES strings. This ensures the same chemical substructure is tokenized identically regardless of the surrounding molecule, a prerequisite for reliable fragment-level supervision.
Stage 1: Context-Aware Fragment Ranking (Discriminative):
- Goal: Determine where to edit.
- Mechanism: The model ranks fragments based on their contribution to a target property within the full molecular context.
- Supervision: Uses rule-based data (RDKit for properties like QED, logP) and GNN-based attribution (SME) for properties like hERG and ESOL. Crucially, it employs SME+, a context-conditioned extension of Substructure Mask Explanation that groups fragment occurrences by their local ECFP environment (radius $r=2$ ) rather than aggregating globally.
- Tasks: Decomposition, ranking, vulnerability prediction, and in-context learning (ICL) ranking.
Stage 2: Explicit Fragment Replacement (Generative):
- Goal: Determine how to edit.
- Mechanism: Generates explicit replacements for identified weak fragments.
- Supervision: Trained on verified low-to-high edit pairs derived from:
  1. SME+: Context-specific replacements that improve scores within a specific environment bin.
  2. ChEMBL Matched Molecular Pairs (MMPs): Real-world pairs sharing a scaffold but differing by one fragment, filtered to exclude benchmark targets (DRD2, JNK3, GSK3-β) to prevent data leakage.
- Output Format: Explicit modification syntax (e.g., frag_src >> frag_tgt) rather than full molecule rewriting.
Inference: In-Context Learning (ICL) for Black-Box Objectives:
- FORGE adapts to unseen, unnamed black-box objectives without gradient updates.
- It maintains a replay buffer of (source molecule, modification, result molecule, oracle score) tuples.
- During search, the model is prompted with demonstrations sampled from this buffer (biased toward high scores) to infer the target objective and propose edits.

3. Key Contributions

Data Strategy: Demonstrated that context-conditioned, rule-verified fragment edits are a viable, scalable, and hallucination-free alternative to natural language objectives or expensive LLM-distilled chain-of-thought data.
Framework Design: Proposed FORGE, a two-stage framework that decouples optimization into context-conditioned ranking (where) and explicit replacement generation (how), utilizing atom-level tokenization to stabilize fragment identity.
Performance Efficiency: Showed that a 0.6B model using this framework achieves state-of-the-art results, outperforming significantly larger models (8B+ LLMs, GPT-4o) and graph-based methods across multiple benchmarks.

4. Experimental Results

The authors evaluated FORGE on five benchmarks:

Prompt-MolOpt (Similarity-Constrained ADMET): FORGE consistently outperforms baselines, with the performance gap widening as similarity constraints tighten ( $\delta=0.6$ ). At $\delta=0.6$ , FORGE achieves a SUM score of 2.882 compared to 0.824 for the original Prompt-MolOpt, demonstrating that local fragment editing is superior to full SMILES rewriting under strict constraints.
PMO-1k (Black-Box Optimization): On 22 unnamed black-box tasks with a limited 1,000-call budget, FORGE achieves an aggregate score of 12.42, outperforming LICO (11.71) and other gradient-free methods.
QED-DRD2 (Explicit Modification): When comparing output formats, the explicit modification variant of FORGE outperforms the end-to-end generation variant by 3–5 percentage points, particularly at strict similarity thresholds.
Lead Optimization (Real-Target Docking): On docking tasks for five protein targets, FORGE achieves the best docking scores in 20 out of 30 configurations, significantly outperforming Graph GA and GenMol, especially under tight similarity constraints.
ChemCoTBench (Small Model vs. Large Models): A 0.6B FORGE model matches or exceeds the performance of Qwen-3-8B (SFT) and GPT-4o on six optimization tasks. Notably, on real-target tasks (DRD2, JNK3, GSK3-β), FORGE achieves success rates of 86–94%, while the 8B and GPT-4o baselines struggle with rates between 20–48%.

5. Significance and Claims

The paper claims that the primary value of FORGE lies in shifting the paradigm from natural language conditioning to explicit fragment-level supervision.

Context is King: The authors provide empirical evidence that fragment effects are highly context-dependent; conditioning fragment attribution on local environments (via SME+) reduces variance by 8–10% and is critical for performance under similarity constraints.
Prompt Limitations: Experiments show that changing the natural language description of a black-box objective in a prompt has negligible effect (<3% score change) on optimization behavior, indicating that natural language is a weak control signal for these tasks.
Scalability: The framework demonstrates that a small, specialized model (0.6B) can outperform massive generalist models when the training data and output format are aligned with the structural nature of the chemical problem (local edits rather than global generation).

The authors conclude that explicit fragment-level supervision is a more easily obtainable, scalable, and reliable alternative to natural language training for molecular optimization, particularly when dealing with black-box objectives and strict structural constraints.