This repository includes the official implementation of "An Efficient SE($p$)-Invariant Transport Metric Driven by Polar Transport Discrepancy-based Representation"

Brief introduction to directories and files:

• SEINT/: Core code for our implementation of the SEINT/ISEINT algorithm package.
• Experiments for validating metric properties:
  • SE(p) invariance/: Code for implementing point cloud classification tasks.
  • Cross Space tasks/: Code for implementing cross-space tasks.
  • High-dimensional data analysis/: Code for testing with high-dimensional data.
  • time_cost: Code for testing computational efficiency
  • Metric consistency/: Code for metric consistency detection.
• As a regularization term:
  • Molecule_Generation/: Core code for its use as a regularization term in E(3)-equivariant diffusion models.
  • Point-MAE/: Core code for its use as a regularization term in Point-MAE.

• python >= 3.8
• numpy
• scipy
• matplotlib
• sklearn
• pytorch >= 2.4.1
• pandas
• POT

1. Evaluate the results in the SEp_invariance.ipynb notebook.

1. Download the Mesh Data from Deformation Transfer for Triangle Meshes and place it in the Cross Space tasks/data/ directory.
2. Run the following script to compute the distance information:

   cd "Cross Space tasks"
   python cross_space_compare.py

3. Evaluate the results in the Cross_Space.ipynb notebook.

• Evaluate the results in the High-dimensional_Analysis.ipynb notebook.

• Evaluate the results in the Computational_Efficiency.ipynb notebook.

Backbone: EDM, UniGEM

1. Navigate to the Molecule_Generation directory and clone the original backbone repositories. Then, follow their instructions to download the QM9 and Drug datasets.

   cd Molecule_Generation
   git clone https://github.com/ehoogeboom/e3_diffusion_for_molecules.git
   git clone https://github.com/fengshikun/UniGEM.git

2. Replace the en_diffusion/en_diffusion.py file in the cloned EDM & UniGEM repository with the one provided in our Molecule_Generation directory.

3. To train EDM and UniGEM on the QM9 dataset, run the following script:

   bash SEINT_train_QM9.sh

4. To fine-tune UniGEM on the QM9 and DRUG datasets, first download the pre-trained checkpoints as instructed in the original UniGEM repository. Then, run the fine-tuning script:

   bash SEINT_ft_DRUG.sh

5. To evaluate the trained models, run the corresponding scripts:

   • On the QM9 dataset:

     bash SEINT_eval_QM9.sh

   • On the DRUG dataset:

     bash SEINT_eval_DRUG.sh

6. We provide pre-trained checkpoints for evaluation with SEINT-0.3 in the following directories: Molecule_Generation/EDM_QM9_ckpt, Molecule_Generation/UniGEM_QM9_ckpt, and Molecule_Generation/UniGEM_DRUG_ft_ckpt. You are welcome to use these for testing and reproducing our results. Checkpoints for training will be released soon.

1. Download the Mesh Data from Deformation Transfer for Triangle Meshes, and place the elephant-reference.obj, flam-reference.obj, and horse-01.obj files into the Metric consistency directory.
2. Run the experiments and view the results in the Metric consistency/metric consistency.ipynb notebook.

Base repository: Pang-Yatian/Point-MAE

We integrate SEINT as a regularization term into the Point-MAE architecture for enhanced point cloud reconstruction performance. To test our implementation:

1. Clone the original repository:

   git clone https://github.com/Pang-Yatian/Point-MAE.git

2. Replace the following three directories in the original codebase with the modified versions from this repository:

   • cfgs/
   • extensions/chamfer_dist/
   • models/

3. The modifications incorporate SEINT regularization during pretraining.

4. Follow the original training and evaluation instructions in the Point-MAE repository.

Hoogeboom Emiel, Satorras Vïctor Garcia, Vignac Clément and Welling Max. "Equivariant diffusion for molecule generation in 3d." International conference on machine learning. PMLR, 2022. [Github]

Pang Yatian, Wang Wenxiao, Tay Francis EH, Liu Wei, Tian Yonghong and Yuan Li. "Masked autoencoders for point cloud self-supervised learning." European conference on computer vision. Cham: Springer Nature Switzerland, 2022. [Github]

Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Z. Alaya, Aurélie Boisbunon, Stanislas Chambon, Laetitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, Léo Gautheron, Nathalie T.H. Gayraud, Hicham Janati, Alain Rakotomamonjy, Ievgen Redko, Antoine Rolet, Antony Schutz, Vivien Seguy, Danica J. Sutherland, Romain Tavenard, Alexander Tong, and Titouan Vayer. "POT Python Optimal Transport library." Journal of Machine Learning Research 22(78): 1-8, 2021. [Web]

Scetbon, Meyer, Gabriel Peyré, and Marco Cuturi. "Linear-time Gromov Wasserstein distances using low rank couplings and costs." International Conference on Machine Learning. PMLR, 2022. [Github]

Vayer Titouan, Flamary Rémi, Tavenard Romain, Chapel Laetitia and Courty Nicolas. "Sliced Gromov-Wasserstein." NeurIPS 2019-Thirty-third Conference on Neural Information Processing Systems. Vol. 32. 2019. [Github]

Shikun Feng and Yuyan Ni and Lu yan and Zhi-Ming Ma and Wei-Ying Ma and Yanyan Lan. "UniGEM: A Unified Approach to Generation and Property Prediction for Molecules." The Thirteenth International Conference on Learning Representations. [Github]

If you found this repository useful, please cite the following.

@inproceedings{lin2026seint,
      title={An Efficient {SE}(p)-Invariant Transport Metric Driven by Polar Transport Discrepancy-based Representation},
      author={Junyi Lin and Dunyao Xue and Jun Yu and Hongteng Xu and Cheng Meng},
      booktitle={The Fourteenth International Conference on Learning Representations},
      year={2026},
      url={https://openreview.net/forum?id=oyxExc7TEl}
}