Mamba2Poly

Efficient Polymer Foundation Model with Bidirectional Mamba-2 and 3D State Fusion

Overview

Mamba2Poly replaces the Transformer backbone in PolyConFM with a Bidirectional Mamba-2 (BiMamba2) state space model, achieving linear-complexity all-atom polymer modeling with explicit 3D spatial awareness.

Key Innovations

• BiMamba2 Encoder: Replaces the hierarchical Transformer with a bidirectional Mamba-2 SSM operating on flattened all-atom sequences, reducing complexity from O(n^2) to O(n).
• SE(3)-Invariant Atom Embedding: Three-channel fusion of element type, RBF-expanded neighbor distances (with optional bond angles), and centroid-centered coordinates.
• 3D State Fusion: Inspired by Spatial-Mamba, injects multi-shell RBF spatial neighborhood information after each BiMamba layer (covalent 1.5A / H-bond 3.5A / long-range 7.0A), alternating between 1D sequence and 3D spatial pathways with only +0.013% parameters overhead.
• Unified Atom-Level Sequence: Polymers are encoded as [CLS] atom_1 ... [SEP] atom_{n+1} ... [SEP], where [SEP] tokens mark monomer boundaries and carry spatially-interpolated coordinates for continuity.

Supported Tasks

Task	Description
Polymer Property Prediction	Band gap, electron affinity, ionization energy, dielectric constant, etc.
Polymer Conformation Generation	Predict stable 3D structures of polymers
Polymer Design	Generate novel polymers satisfying target conditions

Architecture

Input: SMILES + 3D Coordinates
         |
   FlattenedAtomDataset        # Atom-level tokenization with [CLS]/[SEP]
         |
   AtomEmbeddingSE3            # Element + RBF distances + centered coords
         |
   BiMamba2 Layer 1  <---->  StateFusion3D (multi-shell kNN RBF)
   BiMamba2 Layer 2  <---->  StateFusion3D
        ...                      ...
   BiMamba2 Layer L  <---->  StateFusion3D
         |
   Pooling (sep / cls / mean)
         |
   Task Head (PropertyPrediction / Denoising / Generation)

Installation

Prerequisites

• Python 3.10
• CUDA 11.8+
• Uni-Core (see its installation guide)
• Mamba (mamba-ssm and causal-conv1d)

Setup

# Install dependencies
pip install -r requirements.txt

# Install Mamba SSM (requires CUDA)
pip install mamba-ssm causal-conv1d

# Install mini-moses (for design task)
pip install git+https://github.com/igor-krawczuk/mini-moses

Datasets

Download datasets from Zenodo and place them under ./datasets:

mamba2poly/
├── datasets/
│   ├── pretrain_dataset/
│   ├── property_dataset/
│   └── design_dataset/

Quick Inference

Download model checkpoints from Zenodo into ./ckpts:

mamba2poly/
├── ckpts/
│   ├── pretrain_ckpt/
│   ├── property_ckpt/
│   └── design_ckpt/

Property Prediction

bash ./scripts/property_script/property_inference.sh
python ./scripts/property_script/property_eval.py

Conformation Generation

bash ./scripts/conf_script/conf_gen.sh
python ./scripts/conf_script/conf_eval.py

Polymer Design

bash ./scripts/design_script/design_inference.sh
python ./scripts/design_script/design_eval.py

Training

Pretraining

Self-supervised objectives: Masked Atom Prediction (15% masking) + Coordinate Denoising (sigma=0.1A).

bash ./scripts/pretrain.sh

Property Prediction Finetuning

bash ./scripts/property_script/property_train.sh

Design Finetuning

bash ./scripts/design_script/design_train.sh

Project Structure

mamba2poly/
├── mamba2poly/              # Core Python package
│   ├── models/
│   │   ├── mamba2poly.py            # BiMamba2 encoder + SE3 embedding + 3D state fusion
│   │   ├── mamba2poly_base.py       # Shared base model (atom tokenizer, score model, etc.)
│   │   ├── mamba2poly_property.py   # Property prediction model
│   │   ├── mamba2poly_conf_gen.py   # Conformation generation model
│   │   └── mamba2poly_design.py     # Polymer design model
│   ├── tasks/                       # Unicore task definitions
│   ├── losses/                      # Loss functions and metrics
│   └── datas/                       # Dataset loaders
├── scripts/                 # Training and inference scripts
├── data/                    # Data processing utilities
└── config/                  # Configuration files

Acknowledgments

This project is built upon PolyConFM. We also thank Uni-Core, Mamba, Spatial-Mamba, Uni-Mol, MolCLR, TorsionalDiff, FrameDiff, and GraphDiT for their excellent work.

Citation

If you use this code, please cite the original PolyConFM paper:

@article{wang2025unifying,
  title={Unifying Polymer Modeling and Design via a Conformation-Centric Generative Foundation Model},
  author={Wang, Fanmeng and Mei, Shan and Guo, Wentao and Wang, Hongshuai and Ou, Qi and Gao, Zhifeng and Xu, Hongteng},
  journal={arXiv preprint arXiv:2510.16023},
  year={2025}
}

License

This project is licensed under the GPL-3.0 License.