Building symmetry and structure into models for molecules, materials, and simulation.
Scientific data usually comes with structure we already understand. A molecule has the same energy however you rotate it. A fluid behaves the same wherever you place the origin. A graph does not care in what order you happen to list its nodes. A model that has to rediscover these facts from data wastes capacity and examples on something we could simply tell it. The theme running through this work is to build that structure into the model instead, mainly through equivariant graph neural networks, so the symmetry holds by construction and the model learns more from less data.
In chemistry this pays off directly. Predicting NMR chemical shifts for carbohydrates is a hard, geometry-dependent problem, and an E(3) equivariant network cuts the error up to threefold over models that only see the flat 2D structure (Bånkestad et al., 2024). When labelled data is scarce, a related idea helps: pre-training a transformer on reaction data gives a representation that transfers to downstream property prediction (Broberg et al., 2022). The same geometric thinking extends to crystalline solids, where I supervise ongoing work on equivariant models for solid-state NMR.
The physics side looks different but uses the same principle. A fluid surrogate that is equivariant to rotations and translations gives more accurate, more stable rollouts than a plain network, and needs less training data to get there (Bånkestad et al., 2024).
Graphs bring their own structure. Rather than reducing a large graph in the dark, we learn to keep exactly the parts a downstream task needs, by placing an Ising model on the graph and learning its field with a neural network (Bånkestad et al., 2025). The same graph view reaches into numerical linear algebra: representing a matrix as a bipartite graph lets a graph neural network accelerate computations such as nonnegative matrix factorization (Sjölund & Bånkestad, 2022). And the same structure-aware thinking drives ongoing work on designing proteins that bind and modify cellulose.
Reducing a graph while preserving its overall structure is an important problem with many applications. Typically, the reduction approaches either remove edges (sparsification) or merge nodes (coarsening) in an unsupervised way with no specific downstream task in mind. In this paper, we present an approach for subsampling graph structures using an Ising model defined on either the nodes or edges and learning the external magnetic field of the Ising model using a graph neural network. Our approach is task-specific as it can learn how to reduce a graph for a specific downstream task in an end-to-end fashion. The utilized loss function of the task does not even have to be differentiable. We showcase the versatility of our approach on distinct applications, including image segmentation, explainability for graph classification, 3D shape sparsification, and sparse approximate matrix inverse determination.
@inproceedings{bankestad2025ising,title={Ising on the Graph: Task-Specific Graph Subsampling via the Ising Model},author={B{\aa}nkestad, Maria and Andersson, Jennifer R. and Mair, Sebastian and Sj{\"o}lund, Jens},booktitle={Proceedings of the Third Learning on Graphs Conference},series={Proceedings of Machine Learning Research},volume={269},publisher={PMLR},year={2025},}
2024
RSC Adv.
Carbohydrate NMR chemical shift prediction by GeqShift employing E(3) equivariant graph neural networks
Maria Bånkestad, Keven M. Dorst, Göran Widmalm, and 1 more author
Carbohydrates, vital components of biological systems, are well-known for their structural diversity. Nuclear Magnetic Resonance (NMR) spectroscopy plays a crucial role in understanding their intricate molecular arrangements and is essential in assessing and verifying the molecular structure of organic molecules. An important part of this process is to predict the NMR chemical shift from the molecular structure. This work introduces a novel approach that leverages E(3) equivariant graph neural networks to predict carbohydrate NMR spectra. Notably, our model achieves a substantial reduction in mean absolute error, up to threefold, compared to traditional models that rely solely on two-dimensional molecular structure. Even with limited data, the model excels, highlighting its robustness and generalization capabilities.
@article{bankestad2024carbohydrate,title={Carbohydrate NMR chemical shift prediction by GeqShift employing E(3) equivariant graph neural networks},author={B{\aa}nkestad, Maria and Dorst, Keven M. and Widmalm, G{\"o}ran and R{\"o}nnols, Jerk},journal={RSC Advances},volume={14},pages={26585--26595},year={2024},doi={10.1039/D4RA03428G},}
Flexible SE(2) graph neural networks with applications to PDE surrogates
This paper presents a novel approach for constructing graph neural networks equivariant to 2D rotations and translations and leveraging them as PDE surrogates on non-gridded domains. We show that aligning the representations with the principal axis allows us to sidestep many constraints while preserving SE(2) equivariance. By applying our model as a surrogate for fluid flow simulations and conducting thorough benchmarks against non-equivariant models, we demonstrate significant gains in terms of both data efficiency and accuracy.
@article{bankestad2024flexible,title={Flexible SE(2) graph neural networks with applications to PDE surrogates},author={B{\aa}nkestad, Maria and Mogren, Olof and Pirinen, Aleksis},journal={arXiv preprint arXiv:2405.20287},year={2024},}
2022
ICML-W
Pre-training Transformers for Molecular Property Prediction Using Reaction Prediction
Molecular property prediction is essential in chemistry, especially for drug discovery applications. However, available molecular property data is often limited, encouraging the transfer of information from related data. Transfer learning has had a tremendous impact in fields like computer vision and natural language processing signaling for its potential in molecular property prediction. We present a pre-training procedure for molecular representation learning using reaction data and use it to pre-train a SMILES transformer. We fine-tune and evaluate the pre-trained model on 12 molecular property prediction tasks from MoleculeNet within physical chemistry, biophysics, and physiology and show a statistically significant positive effect on 5 of the 12 tasks compared to a non-pre-trained baseline model.
@inproceedings{broberg2022pretraining,title={Pre-training Transformers for Molecular Property Prediction Using Reaction Prediction},author={Broberg, Johan and B{\aa}nkestad, Maria and Ylip{\"a}{\"a}, Erik},booktitle={ICML 2022 2nd AI for Science Workshop},year={2022},}
Graph-based neural acceleration for nonnegative matrix factorization
We describe a graph-based neural acceleration technique for nonnegative matrix factorization that builds upon a connection between matrices and bipartite graphs that is well-known in certain fields, e.g., sparse linear algebra, but has not yet been exploited to design graph neural networks for matrix computations. We first consider low-rank factorization more broadly and propose a graph representation of the problem suited for graph neural networks. Then, we focus on the task of nonnegative matrix factorization and propose a graph neural network that interleaves bipartite self-attention layers with updates based on the alternating direction method of multipliers. Our empirical evaluation on synthetic and two real-world datasets shows that we attain substantial acceleration, even though we only train in an unsupervised fashion on smaller synthetic instances.
@article{sjolund2022graph,title={Graph-based neural acceleration for nonnegative matrix factorization},author={Sj{\"o}lund, Jens and B{\aa}nkestad, Maria},journal={arXiv preprint arXiv:2202.00264},year={2022},}
