FACET: A Fragment-Aware Conformer Ensemble Transformer

This repository contains the official PyTorch implementation of:
FACET: A Fragment-Aware Conformer Ensemble Transformer (ICLR 2026)
[Paper] [Poster]

Duy M. H. Nguyen^1,2,3, Trung Q. Nguyen³, Ha T. H. Le³, Mai Thanh Nhat Truong³, TrungTin Nguyen^4,5, Nhat Ho⁶, Khoa D Doan⁷, Duy Duong-Tran⁸, Li Shen⁸, Daniel Sonntag^3,9, James Zou¹⁰, Mathias Niepert^1,2, Hyojin Kim¹¹, Jonathan E Allen¹¹

¹Max Planck Research School for Intelligent Systems (IMPRS-IS)   ²University of Stuttgart   ³DFKI   ⁴ARC Centre of Excellence for the Mathematical Analysis of Cellular Systems   ⁵Queensland University of Technology   ⁶University of Texas at Austin   ⁷VinUniversity   ⁸University of Pennsylvania   ⁹Oldenburg University   ¹⁰Stanford University   ¹¹Lawrence Livermore National Labs

Figure 1: FACET processes 2D molecular graphs with fragment augmentation and multiple 3D conformers through a frozen fragment-aware Graph Transformer pre-trained to approximate FGW distances, enabling scalable geometry-aware molecular property prediction.

• News
• Highlights
• Architecture Overview
• Results
• Installation
• Data Preparation
• Training Pipeline
  • Quickstart
  • Stage 1: Train 2D and 3D MPNNs
  • Stage 2: Train Graph Transformer
  • Stage 3: End-to-End Fine-tuning
• Project Structure
• Citation
• Acknowledgments
• License

• [ICLR 2026] FACET is accepted at ICLR 2026!
• [04/2026] First release of the official codebase.

• Scalable geometry-aware conformer aggregation. FACET replaces the computationally expensive Fused Gromov-Wasserstein (FGW) alignment with a trainable Graph Transformer, enabling efficient, deterministic attention-based inference with theoretical approximation bounds.

• Fragment-level structural priors. We introduce a unified 2D-3D representation learning approach that embeds fragment-level structural priors (rings, paths, linkers) into both 2D message passing and 3D spatial self-attention, capturing multi-scale interactions between molecular topology and geometry.

• Over 6x faster with state-of-the-art performance. FACET delivers over 6x faster aggregation than prior geometry-aware FGW-based baselines and achieves state-of-the-art results across six benchmarks, including molecular property prediction (MoleculeNet) and Boltzmann-weighted ensemble tasks (MARCEL), scaling to datasets with up to 75,000 molecules.

Figure 3: One-epoch training time comparison between ConAN-FGW and FACET on MARCEL benchmark datasets Drugs-75K and Kraken. FACET achieves 5-6x speedup.

FACET employs a three-stage training pipeline:

Fragment-Enhanced 2D Representation. Molecules are decomposed into chemically meaningful substructures (rings, paths, junction nodes) using RingPath decomposition. A separate 2D-MPNN operates on the fragment-level graph, and fragment embeddings are fused back into atom-level representations.

Figure 6: RingsPaths decomposition on molecules from the Kraken dataset. Top: original 2D structures; Bottom: fragment decomposition where each color represents a distinct structural fragment (ring or path).

FGW-Guided Graph Transformer. The Graph Transformer is pre-trained to map conformer representations into a latent space where pairwise Euclidean distances approximate FGW distances. This enables fast, permutation-invariant conformer integration.

Figure 2: Strong correlation between true FGW distances and Graph Transformer embedding distances across four MoleculeNet datasets (ρ = 0.81-0.96).

Comparison with state-of-the-art methods on molecular property regression. FACET uses a SchNet backbone.

• Ubuntu 18.04+
• CUDA 11+
• Miniconda
• Python 3.9
• PyTorch 2.0+
• PyTorch Geometric

git clone https://github.com/duyhominhnguyen/FACET.git
cd FACET

# Create conda environment
conda create -n facet python=3.9
conda activate facet

# Install PyTorch (adjust CUDA version as needed)
pip install torch==2.0.0 torchvision==0.15.1 torchaudio==2.0.1

# Install remaining dependencies
conda env update -n facet --file environment.yml

We evaluate on four MoleculeNet datasets for molecular property regression:

For Boltzmann-weighted ensemble prediction:

Download the benchmark datasets from this link and place them in the data/ directory:

data/
├── esol/
│   ├── esol.csv
│   ├── train.csv, val.csv, test.csv
│   ├── conformers_train/, conformers_val/, conformers_test/
│   └── spd/
├── lipo/
├── freesolv/
└── bace/

Run the full FACET pipeline on ESOL with SchNet backbone:

bash conan_fgw/script/facet_transformer_single_stage/run_facet_transformer_individual.sh \
    esol 5 5 0 facet_esol 0.002

Arguments: <dataset> <num_conformers> <num_runs> <gpu_device> <exp_name> <learning_rate>

Supported datasets: esol, freesolv, bace, lipo

Train the 2D MPNN (GAT on atom + fragment graphs) and 3D MPNN (SchNet/ViSNet) jointly for 150 epochs. The 2D fragment features are fused into atom representations via RingPath decomposition.

export WORKDIR=$(pwd)
export PYTHONPATH="$WORKDIR:$PYTHONPATH"
export CUDA_VISIBLE_DEVICES=0

python conan_fgw/src/train_val.py \
    --config_path=${WORKDIR}/conan_fgw/config/schnet/property_regression/esol/esol_5.yaml \
    --cuda_device=0 \
    --data_root=${WORKDIR} \
    --number_of_runs=5 \
    --checkpoints_dir=${WORKDIR}/models \
    --logs_dir=${WORKDIR}/outputs \
    --run_name=schnet_esol_5_conan_fgw_pre \
    --stage=conan_fgw_pre \
    --model_name=schnet \
    --learning_rate=0.002

Note: For freesolv, use n_cfm=10 (10 conformers); for all other datasets, use n_cfm=5.

Extract 3D features, 2D fragments, shortest-path distances, and FGW distances from Stage 1, then train a Graphormer (12 layers, 8 heads, hidden dim 64, 372K params) to approximate pairwise FGW distances for 700 epochs.

# Step 2a: Extract features and compute FGW distances
python conan_fgw/src/save_fgw_distance.py \
    --config_path=${WORKDIR}/conan_fgw/config/schnet/property_regression/esol/esol_5.yaml \
    --data_root=${WORKDIR} \
    --model_name=schnet \
    --stage=conan_fgw_pre \
    --conan_fgw_pre_ckpt_dir=${WORKDIR}/models/schnet_esol_5_conan_fgw_pre/<run_id>

# Step 2b: Train Graph Transformer (uses external Graphormer codebase)
# See run_facet_transformer_stage2.sh for details
# Training: 700 epochs, lr=1e-4

Freeze the pre-trained Graph Transformer and train the full model end-to-end (2D/3D MPNNs + MLP adaptor layers) for 400 epochs with lr=5e-4.

export CUDA_VISIBLE_DEVICES=0,1,2,3

python conan_fgw/src/train_val.py \
    --config_path=${WORKDIR}/conan_fgw/config/schnet/property_regression/esol/esol_5_bc.yaml \
    --cuda_device=0 \
    --data_root=${WORKDIR} \
    --number_of_runs=5 \
    --checkpoints_dir=${WORKDIR}/models \
    --logs_dir=${WORKDIR}/outputs \
    --run_name=schnet_esol_5_facet \
    --stage=conan_fgw \
    --model_name=schnet \
    --conan_fgw_pre_ckpt_dir=${WORKDIR}/models/schnet_esol_5_conan_fgw_pre/<run_id>

FACET/
├── conan_fgw/
│   ├── config/                          # Experiment configurations
│   │   └── <model>/<task>/<dataset>/    # e.g., schnet/property_regression/esol/
│   │       ├── esol_5.yaml              # Stage 1 config
│   │       └── esol_5_bc.yaml           # Stage 3 config
│   ├── script/                          # Training scripts
│   │   └── facet_transformer_single_stage/
│   │       ├── run_facet_transformer_individual.sh   # Main entry point
│   │       ├── run_facet_transformer_stage1.sh       # Stage 1
│   │       ├── run_facet_transformer_extract_feat.sh # Feature extraction
│   │       ├── run_facet_transformer_stage2.sh       # Stage 2 (Graph Transformer)
│   │       └── run_facet_transformer_stage3.sh       # Stage 3
│   └── src/
│       ├── train_val.py                 # Main training entry point
│       ├── trainer.py                   # PyTorch Lightning trainer
│       ├── experiments.py               # Experiment classes
│       ├── config_parser.py             # Config + CLI argument parsing
│       ├── save_fgw_distance.py         # FGW distance computation
│       ├── data/
│       │   ├── datasets.py              # Dataset classes + fragment graph construction
│       │   ├── datamodules.py           # PyTorch Lightning DataModule
│       │   └── splitters.py             # Train/val/test splits
│       └── model/
│           ├── common.py                # Base models, losses, metrics
│           ├── schnet_based_models.py   # SchNet/ViSNet model variants
│           ├── graph_embeddings/        # Backbone implementations (SchNet, ViSNet, GAT)
│           └── fgw/                     # Fused Gromov-Wasserstein module
├── data/                                # Benchmark datasets
├── models/                              # Saved checkpoints
├── outputs/                             # Training logs
├── environment.yml                      # Conda dependencies
├── Makefile                             # Build/lint/format commands
└── README.md

If you find this work useful, please cite our paper:

@inproceedings{nguyen2026facet,
  title={FACET: A Fragment-Aware Conformer Ensemble Transformer},
  author={Nguyen, Duy M. H. and Nguyen, Trung Q. and Le, Ha T. H. and Truong, Mai Thanh Nhat and Nguyen, TrungTin and Ho, Nhat and Doan, Khoa D and Duong-Tran, Duy and Shen, Li and Sonntag, Daniel and Zou, James and Niepert, Mathias and Kim, Hyojin and Allen, Jonathan E},
  booktitle={International Conference on Learning Representations (ICLR)},
  year={2026}
}

FACET extends our prior work ConAN-FGW (ICML 2024):

@inproceedings{nguyen2024structure,
  title={Structure-Aware E(3)-Invariant Molecular Conformer Aggregation Networks},
  author={Nguyen, Duy MH and Lukashina, Nina and Nguyen, Tai and Le, An T and Nguyen, TrungTin and Ho, Nhat and Peters, Jan and Sonntag, Daniel and Zaverkin, Viktor and Niepert, Mathias},
  booktitle={International Conference on Machine Learning (ICML)},
  year={2024}
}

This work was partially performed under the auspices of the U.S. Department of Energy by Lawrence Livermore National Laboratory under contract DE-AC52-07NA27344. This work was funded by the Defense Threat Reduction Agency (DTRA), HDTRA1242044 (HK) and HDTRA1036045 (JA). The project was also supported by Deutsche Forschungsgemeinschaft (DFG) under Germany's Excellence Strategy - EXC 2075 - 390740016, the DARPA ANSR program under award FA8750-23-2-0004, the DARPA CODORD program under award HR00112590089. The authors thank the International Max Planck Research School for Intelligent Systems (IMPRS-IS) for supporting Duy M. H. Nguyen. Duy M. H. Nguyen, Trung Q. Nguyen, Mai Thanh Nhat Truong, and Daniel Sonntag are also supported by the No-IDLE project (BMBF, 01IW23002), the MASTER project (EU, 101093079), and the Endowed Chair of Artificial Intelligence, Oldenburg University.

This project is licensed under the MIT License.

For academic use. For commercial inquiries, please contact Ho_Minh_Duy.Nguyen@dfki.de, Trung.Nguyen@dfki.de.