From Transformers to State Spaces: GeoMamba-SE(3) for Fast and Accurate Molecular Learning

Jiayu Qin, Zhengquan Luo, Jian Chen, Xuhui Li, Jiayi Chen, Zhiqiang Xu
Proceedings of The 29th International Conference on Artificial Intelligence and Statistics, PMLR 300:2476-2484, 2026.

Abstract

Transformers play an important role in molecular representation learning, enabling unsupervised learning from large scale unlabeled molecule datasets. However, existing Transformer based methods suffer from heavy training computation and slow inference. To accelerate the computation and relieve the burdensome pre-training, we propose a Mamba-based framework that leverages selective state space models to learn molecular representations more efficiently. Unlike conventional methods, our model, GeoMamba-SE(3), offers streamlined computation with linear-time complexity. However, naively applying Mamba to molecules struggles with SE(3) symmetry, representations can drift under rotations/translations—leading to chemically inconsistent features. To address this, we introduce a geometry and statistics aware design: (i) complete local frames at atoms by converting geometric vectors into scalar channels suitable for SSMs; (ii) multi-stream Mamba blocks are modulated by SE(3)-invariant scalars to preserve geometric stability; and (iii) we impose statistical symmetry constraints via orbit-kernel losses and invariant risk minimization, treating SE(3) actions and conformers as environments. This yields practical SE(3) stability without heavy high-order tensor representations. Experiments show that our method achieves new state-of-the-art performance benchmarks on the MoleculeNet datasets, while using only one-sixth of the training computation and 57% less computation for inference.

Cite this Paper

BibTeX
@InProceedings{pmlr-v300-qin26b, title = { From Transformers to State Spaces: GeoMamba-SE(3) for Fast and Accurate Molecular Learning }, author = {Qin, Jiayu and Luo, Zhengquan and Chen, Jian and Li, Xuhui and Chen, Jiayi and Xu, Zhiqiang}, booktitle = {Proceedings of The 29th International Conference on Artificial Intelligence and Statistics}, pages = {2476--2484}, year = {2026}, editor = {Khan, Emtiyaz and Li, Yingzhen and Solin, Arno and Ramdas, Aaditya}, volume = {300}, series = {Proceedings of Machine Learning Research}, month = {02--05 May}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v300/main/assets/qin26b/qin26b.pdf}, url = {https://proceedings.mlr.press/v300/qin26b.html}, abstract = { Transformers play an important role in molecular representation learning, enabling unsupervised learning from large scale unlabeled molecule datasets. However, existing Transformer based methods suffer from heavy training computation and slow inference. To accelerate the computation and relieve the burdensome pre-training, we propose a Mamba-based framework that leverages selective state space models to learn molecular representations more efficiently. Unlike conventional methods, our model, GeoMamba-SE(3), offers streamlined computation with linear-time complexity. However, naively applying Mamba to molecules struggles with SE(3) symmetry, representations can drift under rotations/translations—leading to chemically inconsistent features. To address this, we introduce a geometry and statistics aware design: (i) complete local frames at atoms by converting geometric vectors into scalar channels suitable for SSMs; (ii) multi-stream Mamba blocks are modulated by SE(3)-invariant scalars to preserve geometric stability; and (iii) we impose statistical symmetry constraints via orbit-kernel losses and invariant risk minimization, treating SE(3) actions and conformers as environments. This yields practical SE(3) stability without heavy high-order tensor representations. Experiments show that our method achieves new state-of-the-art performance benchmarks on the MoleculeNet datasets, while using only one-sixth of the training computation and 57% less computation for inference. } }
Endnote
%0 Conference Paper %T From Transformers to State Spaces: GeoMamba-SE(3) for Fast and Accurate Molecular Learning %A Jiayu Qin %A Zhengquan Luo %A Jian Chen %A Xuhui Li %A Jiayi Chen %A Zhiqiang Xu %B Proceedings of The 29th International Conference on Artificial Intelligence and Statistics %C Proceedings of Machine Learning Research %D 2026 %E Emtiyaz Khan %E Yingzhen Li %E Arno Solin %E Aaditya Ramdas %F pmlr-v300-qin26b %I PMLR %P 2476--2484 %U https://proceedings.mlr.press/v300/qin26b.html %V 300 %X Transformers play an important role in molecular representation learning, enabling unsupervised learning from large scale unlabeled molecule datasets. However, existing Transformer based methods suffer from heavy training computation and slow inference. To accelerate the computation and relieve the burdensome pre-training, we propose a Mamba-based framework that leverages selective state space models to learn molecular representations more efficiently. Unlike conventional methods, our model, GeoMamba-SE(3), offers streamlined computation with linear-time complexity. However, naively applying Mamba to molecules struggles with SE(3) symmetry, representations can drift under rotations/translations—leading to chemically inconsistent features. To address this, we introduce a geometry and statistics aware design: (i) complete local frames at atoms by converting geometric vectors into scalar channels suitable for SSMs; (ii) multi-stream Mamba blocks are modulated by SE(3)-invariant scalars to preserve geometric stability; and (iii) we impose statistical symmetry constraints via orbit-kernel losses and invariant risk minimization, treating SE(3) actions and conformers as environments. This yields practical SE(3) stability without heavy high-order tensor representations. Experiments show that our method achieves new state-of-the-art performance benchmarks on the MoleculeNet datasets, while using only one-sixth of the training computation and 57% less computation for inference.
APA
Qin, J., Luo, Z., Chen, J., Li, X., Chen, J. & Xu, Z.. (2026). From Transformers to State Spaces: GeoMamba-SE(3) for Fast and Accurate Molecular Learning . Proceedings of The 29th International Conference on Artificial Intelligence and Statistics, in Proceedings of Machine Learning Research 300:2476-2484 Available from https://proceedings.mlr.press/v300/qin26b.html.

Related Material