Suiren-1.0 Technical Report: A Family of Molecular Foundation Models

The paper introduces Suiren-1.0, a family of open-source molecular foundation models that utilize spatial self-supervised pre-training and conformation compression distillation to bridge 3D geometry and 2D representations, achieving state-of-the-art performance in quantum property prediction and diverse organic system modeling.

The Gist

Imagine you are trying to understand a complex machine, like a car. You could look at a 2D blueprint (a flat drawing), or you could look at the actual 3D engine with all its moving parts.

For a long time, AI models trying to understand chemistry had to choose: either they were great at reading the flat blueprints (2D chemical formulas) but clueless about how the engine actually moves in 3D space, or they were experts at the 3D engine but couldn't easily apply that knowledge to real-world problems where you only have the blueprint.

Suiren-1.0 is a new family of AI models from the Shanghai Academy of AI for Science that solves this problem. Think of it as a "universal translator" that can understand both the flat blueprint and the 3D engine, and then teach a smaller, faster version of itself to do the job efficiently.

Here is how it works, broken down into three simple characters:

1. The Master Architect: Suiren-Base

Imagine a brilliant, super-smart architect who has studied 70 million detailed 3D blueprints of molecules.

• What they do: They don't just look at the atoms; they understand the physics of how those atoms vibrate, bend, and interact in 3D space. They know the "laws of the universe" (quantum mechanics) that govern these shapes.
• The Catch: This architect is huge. They require a massive supercomputer to run, and they need a 3D model of the molecule to work. If you only give them a flat 2D drawing (like a text string called SMILES), they can't help you.

2. The Specialized Expert: Suiren-Dimer

This is the Master Architect's twin, but they specialize in relationships.

• What they do: While the Master Architect focuses on how a single molecule holds itself together, Suiren-Dimer studies how two molecules interact with each other (like how a drug molecule grabs onto a protein in your body).
• Why it matters: This is crucial for drug discovery, where the "handshake" between two molecules is what makes a medicine work.

3. The Efficient Apprentice: Suiren-ConfAvg

This is the star of the show for everyday use. Imagine the Master Architect teaching a brilliant apprentice.

• The Magic Trick (Conformation Compression Distillation): The Master Architect takes all their deep, complex 3D knowledge and "distills" it into the apprentice's brain.
• The Result: The apprentice (Suiren-ConfAvg) is lightweight and fast. It doesn't need a 3D model to start working. You can hand it a simple 2D drawing or a text string, and it instantly "imagines" all the possible 3D shapes that molecule could take, averages them out, and gives you a prediction.
• Why it's cool: It's like having a genius who can look at a flat sketch of a house and instantly tell you exactly how much it will cost to build, how strong the walls are, and how the heat will flow through it, without needing to see the 3D construction plans first.

The Big Problem They Solved

In the past, scientists faced a "Multiscale Gap":

• Microscopic (The 3D World): To understand the laws of physics, you need expensive, slow computer simulations (like running a wind tunnel test for every single molecule).
• Macroscopic (The Real World): In the real world (like a pharmacy or a factory), we usually only have simple 2D lists of ingredients.

Existing AI models were either too slow (needing 3D data) or too dumb (ignoring the 3D physics). Suiren-1.0 bridges this gap. It learns the hard physics in the 3D world, then compresses that knowledge so it can be used instantly on simple 2D data.

What Can It Do?

The team tested Suiren-1.0 on over 50 different tasks, and it won almost every time. It can predict:

• Safety: Will this chemical explode? Is it toxic?
• Properties: At what temperature does it boil? How heavy is it?
• Drug Discovery: Will this new drug stick to a virus?
• Materials: Can this molecule be used to make a better battery?

The Bottom Line

Suiren-1.0 is like giving scientists a "crystal ball" for chemistry. Instead of running expensive, slow experiments in a lab for every new idea, they can use this AI to predict the outcome instantly with high accuracy.

The best part? The creators have open-sourced everything. They gave away the blueprints, the weights, and the benchmarks, allowing anyone in the world to use this "Master Architect" to build better medicines, greener materials, and safer chemicals.

Technical Summary

1. Problem Statement

The field of molecular modeling faces two primary challenges that hinder the development of universal foundation models:

• Complexity of Physical Priors: Molecular behavior is governed by intricate quantum mechanical laws (e.g., Schrödinger equation) and statistical thermodynamics (e.g., Boltzmann distributions). Capturing these mechanisms solely through data-driven learning is difficult due to the sparsity of high-fidelity labeled data.
• The Multiscale Gap: There is a disconnect between microscopic and macroscopic representations:
  • Microscopic: Requires explicit 3D conformations and electronic densities, typically modeled using Density Functional Theory (DFT). While data-rich, these models lack generalizability to broad chemical tasks.
  • Macroscopic: Relies on 1D SMILES or 2D molecular graphs (lacking explicit 3D info) for tasks like drug design and materials discovery. These tasks often suffer from data scarcity (requiring costly wet-lab experiments) and existing 2D models are "conformation-blind," limiting their predictive power for properties derived from ensemble averages.

Current approaches fail to bridge this divide: pure 3D models (e.g., UMA) lack generalizability, while pure 2D models (e.g., MoleBERT) cannot capture conformation-dependent physics.

2. Methodology

Suiren-1.0 introduces a three-stage framework designed to unify molecular scales by bridging 3D conformational geometry with 2D statistical ensemble spaces.

A. Model Architecture

• Suiren-Base (1.8B Parameters): A high-degree $SO(3)$-equivariant Graph Neural Network (GNN).
  • Core: Integrates EquiformerV2 with a dense Mixture-of-Experts (MoE) update block.
  • Experts: Each block utilizes both S2Activation and Equivariant Spherical Transformer (EST) experts.
  • Innovation: To mitigate discretization errors in spherical sampling, the authors propose a basis-rotation strategy. During training, the Fourier basis for each sample is pre-rotated by a random 3D rotation, exposing the model to diverse orientations and approximating continuous spherical Fourier behavior without significant computational overhead.
• Suiren-Dimer: A variant of Suiren-Base extended via continued pre-training on 13.5M intermolecular interaction samples to capture long-range dimer interactions.

B. Pre-Training Strategy

• Data: Trained on 70 million conformer samples generated via DFT (B3LYP/def2-SVP level) covering elements H, C, N, O, F, P, S, Cl, Br, I.
• Augmentation: Uses the EMPP (Enhanced Masked Atom Prediction) method where random atoms are deleted (not masked), and the model reconstructs coordinates conditioned on atom type and target energy. This doubles the training volume and enforces learning of local potential energy landscapes.
• Stages:
  1. Stage 1: Multi-task learning (energy, force, structure, and EMPP completion) on mixed-precision hardware.
  2. Stage 2: Full-precision fine-tuning on the original 70M samples to refine regression targets.
  3. Stage 3: Continued pre-training on dimer data to create Suiren-Dimer.

C. Conformation Compression Distillation (CCD)

To enable efficient downstream application using only 2D inputs (SMILES/Graphs), the authors propose CCD, a diffusion-based framework:

• Mechanism: A 3D diffusion model (dynamics network $\varphi_\theta$) is conditioned on the 2D molecular topology (encoded via GAT) and the conformer energy ($E$).
• Process: The model learns to reconstruct the 3D representation ($h_{3D}$) and coordinates ($c$) from noisy states, effectively distilling the 3D knowledge of Suiren-Base into a 2D latent space.
• Output: This yields Suiren-ConfAvg, a lightweight model that generates high-fidelity, conformation-averaged representations from 2D inputs. It implicitly learns the mapping from 2D graphs to 3D ensembles governed by the Boltzmann distribution.

D. Downstream Fine-Tuning

• Dual Graph Neural Network (DGNN): For specific tasks, a task-specific GNN is initialized alongside the frozen Suiren-ConfAvg. Latent representations from Suiren-ConfAvg are injected into the task-specific GNN layers to provide structural guidance, preventing catastrophic forgetting while adapting to new tasks.

3. Key Contributions

1. Unified Modeling Framework: Successfully bridges the microscopic (3D) and macroscopic (2D) gaps, allowing 2D inputs to access physics-aware 3D knowledge.
2. Suiren-Base & Suiren-Dimer: State-of-the-art 3D foundation models trained on 70M+ DFT samples with advanced equivariant architectures and basis-rotation strategies.
3. Conformation Compression Distillation (CCD): A novel diffusion-based distillation technique that compresses complex 3D structural knowledge into a robust 2D representation, enabling high-performance prediction without requiring 3D input.
4. MoleHB Benchmark: Introduction of a comprehensive benchmark comprising 40+ heterogeneous tasks across 9 scientific domains (safety, thermal, transport, etc.), sourced from rigorous wet-lab data (Yaws, 1999).
5. Open Science: Full release of model weights (Suiren-Base, Dimer, ConfAvg), code, and benchmarks to facilitate reproducible research.

4. Results

The models were evaluated on over 50 tasks across 9 domains using a unified fine-tuning pipeline.

• Pre-Training Performance:

  • Suiren-Base achieved a Molecular Energy MAE of 9.08 meV and Atomic Force MAE of 0.510 meV/Å, significantly outperforming baselines like EquiformerV2 and eSCN (which were trained on smaller subsets).
  • Compared to UMA-family models, Suiren-Base showed superior force prediction (0.510 vs. >2.90).

• Downstream Performance (MoleHB):

  • Suiren-ConfAvg achieved State-of-the-Art (SOTA) results on 41 out of 43 properties.
  • Key Improvements:
    • Energetic Properties: >30% improvement in Gibbs energy, internal energy, and enthalpy of formation.
    • Structural Properties: 45.5% improvement in liquid volume prediction.
    • Thermal Properties: 26.2% improvement in boiling point prediction.
    • Safety: Significant gains in flash point and explosive limits.
  • Robustness: The model maintained high performance even with a Scaffold Split (structural extrapolation), demonstrating strong generalization.

• TDC ADMET Benchmarks:

  • Evaluated on drug discovery tasks (ADMET) without complex hyperparameter search (using a single fixed config).
  • Achieved SOTA in 9 out of 18 metrics and ranked second in 4 others, validating its efficiency and reliability for real-world drug discovery pipelines.

5. Significance

• Paradigm Shift: Suiren-1.0 moves beyond the "3D vs. 2D" dichotomy by creating a model that learns 3D physics but operates efficiently on 2D inputs, making high-fidelity quantum-level predictions accessible for large-scale industrial applications (e.g., drug screening, materials discovery).
• Data Efficiency: By leveraging self-supervised pre-training on massive DFT datasets and distilling this knowledge, the model reduces the dependency on scarce, expensive experimental labels for downstream tasks.
• Reproducibility: The open-sourcing of the MoleHB benchmark and model weights sets a new standard for evaluating molecular foundation models, moving away from task-specific engineering toward unified, fair comparisons.
• Scalability: The architecture demonstrates that large-scale equivariant models (1.8B parameters) can be effectively trained and distilled, paving the way for even larger molecular foundation models in the future.

Limitations & Future Work:

• Current model size is limited by computational constraints.
• The MoE framework currently uses a dense expert strategy; future work may implement Top-K routing for inference speed.
• Specific downstream tasks may still benefit from further hyperparameter optimization.