A Systematic Evaluation of Co-folding Model Representations for Small-Molecule Learning

This paper demonstrates that the Boltz2 co-folding model, which leverages protein-ligand interactions for pretraining, generates superior and complementary small-molecule representations that outperform existing standalone models across diverse tasks including ADMET prediction, generative modeling, and structure-guided optimization.

The Gist

Imagine you are trying to teach a robot how to understand the shape and behavior of tiny chemical building blocks (small molecules) to help design new medicines.

Usually, scientists teach these robots by showing them millions of isolated molecules, like looking at a single Lego brick in a vacuum and asking, "What does this do?" The robot learns the brick's shape, but it never sees how that brick actually snaps together with other pieces.

This paper introduces a new, smarter way to teach the robot. Instead of looking at the brick alone, they show the robot the brick while it is already snapped into a complex machine (a protein). This is called "co-folding."

Here is the breakdown of their discovery using simple analogies:

1. The Problem: Learning in a Vacuum

Most current AI models for chemistry are like students who only study flashcards of single words. They know what the word looks like, but they haven't seen it used in a sentence. They miss out on the context of how words interact with each other. In chemistry, this means the AI misses how a drug molecule interacts with the protein it's supposed to target.

2. The Solution: The "Co-Folding" Teacher

The researchers used a powerful new AI model called Boltz2. Originally, Boltz2 was trained to predict how a protein and a drug molecule fold together into a 3D shape. It's like a master architect who knows exactly how every brick fits into a wall.

The big question was: Can we take this architect, who is an expert at building walls, and use it just to understand the bricks themselves?

3. The Experiment: Reusing the Architect

The team took the "brain" of Boltz2 (its internal representations) and applied it to tasks where the protein isn't present. They treated Boltz2 as a general-purpose teacher for small molecules.

They tested this "teacher" in three main ways:

• The Exam (ADMET Prediction): They asked Boltz2 to predict if a drug would be safe or toxic (absorption, metabolism, etc.).
  • Result: Boltz2 scored just as well, or better, than models specifically trained for these exams, even though it was never explicitly taught for them. It learned the rules of chemistry just by watching how molecules interact with proteins.
• The Creative Writing Class (Molecular Generation): They tried to teach a robot to invent new molecules. Usually, this is slow and trial-and-error.
  • Result: By using Boltz2's "understanding" as a guide, the robot learned to create valid, high-quality molecules twice as fast. It was like giving the robot a map of the city instead of letting it wander blindly.
• The Treasure Hunt (Ligand Optimization): They asked the robot to find the perfect key (molecule) to fit a specific lock (protein).
  • Result: Instead of just getting a "Yes/No" score on how well the key fits, the robot used Boltz2's detailed "feelings" about the fit (intermediate representations) to learn much faster. It found better keys with fewer attempts.

4. The Secret Sauce: A New Perspective

The paper found that Boltz2 doesn't just memorize facts; it learns a different kind of map of the chemical world.

• Existing models are like a map of a city's streets (molecular structure).
• Boltz2 is like a map of the city's traffic and how cars interact with intersections (molecular interactions).

Because these maps are different, combining them (like using Boltz2 alongside an older model) creates a super-map that is better than either one alone.

5. The Conclusion

The paper concludes that watching molecules interact with proteins is a powerful way to teach AI about molecules.

They proved that you don't need to train a model from scratch on millions of isolated molecules to get a great result. You can take a model trained on complex protein-drug interactions, strip away the protein, and use its "knowledge" of the drug to solve problems on its own.

In short: They showed that an AI trained to build complex structures (protein-drug complexes) is actually a brilliant teacher for understanding the individual parts (small molecules) better than models that only study the parts in isolation. This makes Boltz2 a ready-to-use, "off-the-shelf" tool for drug discovery.

Technical Summary

Technical Summary: A Systematic Evaluation of Co-folding Model Representations for Small-Molecule Learning

Problem Statement

Small-molecule foundation models are typically pretrained on standalone molecular data, learning representations based on isolated molecules associated with molecule-level supervision (e.g., bioassay labels, quantum-chemical properties). This contrasts with vision and language domains, where strong representations often emerge from cross-modal or relational supervision. While protein-ligand co-folding structures provide a molecular analogue to such relational supervision by exposing models to atom-level interactions (e.g., hydrogen bonding, electrostatics) within a bound context, it remains unclear whether co-folding models can yield strong, transferable representations for standalone small-molecule tasks. Furthermore, the broader utility of modern co-folding models (like AlphaFold3 and Boltz) is debated, with some suggesting they rely heavily on evolutionary patterns rather than genuine atom-level understanding. This paper seeks to systematically evaluate whether repurposing a modern co-folding model for small-molecule representation learning is viable and effective.

Methodology

The authors investigate Boltz2, a modern 1B-parameter co-folding model designed for protein-ligand structure prediction, as a source of atom-level representations for small molecules. The core methodology involves extracting ligand representations from Boltz2 (omitting protein inputs for standalone tasks) and evaluating their transferability across three distinct downstream domains:

1. Predictive Transfer (ADMET): The authors apply "probing" to evaluate Boltz2's atom-level representations on the Therapeutics Data Commons (TDC) ADMET benchmark. They extract pair representations from specific layers ({16, 32, 48, 64}) of the 64-layer Pairformer trunk, apply a hybrid pooling strategy (concatenating diagonal, bonded, and all entries), and feed the pooled vector into a Multi-Layer Perceptron (MLP) for property prediction. They compare these results against specialized foundation models (MolE, MiniMol, KPGT, QIP, UniMol2, MolGPS, UniQSAR).
2. Generative Transfer: The authors employ representation alignment-based distillation to transfer Boltz2's representations into state-of-the-art molecular generative models (specifically GruM). They train the generative model to maximize the cosine similarity between its hidden representations of noisy molecules and the corresponding clean representations from Boltz2. This is evaluated on the ZINC250k dataset using metrics such as validity, Fréchet ChemNet Distance (FCD), and novelty.
3. Optimization Transfer (Structure-Guided Discovery): The authors extend representation alignment to online reinforcement learning for ligand discovery. Using the SynFlowNet-Boltz pipeline, they treat Boltz2 not just as a black-box scalar reward oracle (binding affinity) but as a "white-box" teacher. In addition to scalar rewards, the policy network is guided by dense representation-level supervision (alignment with Boltz2's intermediate ligand representations) to improve credit assignment and sample efficiency in discovering high-affinity ligands.

Key Contributions

The paper presents four primary findings that position modern co-folding models as strong baselines for small-molecule foundation models:

• Predictive Performance: Boltz2 representations match or outperform specialized small-molecule foundation models on the extensive TDC ADMET benchmark. Notably, Boltz2 achieves top performance on 9 of 22 tasks without any standalone small-molecule pretraining, suggesting that co-folding supervision provides signals not captured by architecture scaling or data volume alone.
• Generative Acceleration: Representation alignment with Boltz2 significantly accelerates the training of molecular generative models (2× faster) and yields higher-quality molecules across multiple metrics (validity, FCD, NSPDK) compared to alignment with other foundation models. This indicates that co-folding models provide superior structural supervision for generation.
• Sample Efficiency in Optimization: In online structure-guided ligand discovery, incorporating dense representation-level supervision from Boltz2 improves the sample efficiency of finding high-reward molecules compared to using scalar rewards alone. This marks the first application of representation alignment to improve online reinforcement learning in molecular discovery.
• Representation Complementarity: Boltz2 yields a representation space that is distinct from and complementary to existing molecular foundation models. The authors demonstrate that ensembling Boltz2 with a less aligned model (MolE) yields greater performance gains on certain benchmarks than ensembling with a more aligned model (MiniMol), suggesting a practical strategy for representation fusion.

Results

• ADMET Benchmark: Boltz2 achieved the best average performance on 9 of 22 tasks among non-ensemble methods. When combined with MiniMol (Boltz2Mini.), it also led on 9 tasks. Boltz2 significantly outperformed UniMol2 (which uses the same 1B Pairformer architecture but is pretrained on standalone 3D structures), highlighting the value of co-folding supervision.
• Molecular Generation: Models distilled with Boltz2 representations achieved the best scores across validity, FCD, and NSPDK simultaneously, whereas baseline diffusion models typically excelled on only a subset of these metrics.
• Ligand Discovery: In structure-guided discovery tasks across seven target proteins, the representation-aligned policy discovered more high-reward molecules and explored more diverse modes (distinct high-scoring molecules) than the vanilla policy under the same reward evaluation budget.
• Complementarity Analysis: CKNNA (Centered Kernel Nearest-Neighbor Alignment) scores showed Boltz2 has relatively weak alignment with existing models. However, ensembling Boltz2 with MolE (low alignment) outperformed ensembling with MiniMol (higher alignment) on specific metabolism tasks, confirming the complementary nature of the representations.

Significance and Claims

The paper positions protein-ligand co-folding as a promising pretraining paradigm for small-molecule representation learning. It argues that co-folding models like Boltz2 learn atom-level interaction patterns (e.g., binding conformations) that transfer effectively to standalone tasks, offering a strong, off-the-shelf foundation model for drug discovery.

The authors claim to be the first to demonstrate that:

1. Informative molecular representations from co-folding models can improve molecular generation and optimization.
2. Representation alignment-based distillation is effective for online reinforcement learning, providing dense signals that complement scalar rewards.
3. Combining weakly aligned representations (e.g., Boltz2 + MolE) can be a more effective strategy for representation fusion than combining strongly aligned ones.

Limitations Acknowledged:
The authors modestly note that co-folding models may lack complete atom-level understanding of local mechanisms, evidenced by their underperformance on metabolism substrate prediction tasks (a weakness shared by UniMol2, suggesting a broader limitation of 3D structure-based learning). They also note that their experiments use a frozen-backbone setting via probing and distillation, leaving backbone fine-tuning for future work.