Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

DeepChem Equivariant: SE(3)-Equivariant Support in an Open-Source Molecular Machine Learning Library

Jose Siguenza, Bharath Ramsundar

TL;DR

The paper addresses making SE(3)-equivariant neural networks readily usable for molecular machine learning by integrating SE(3)-Transformer and TFN-style components into the DeepChem ecosystem. It leverages spherical harmonics and Clebsch–Gordan-based irreducible representations to build rotation-equivariant kernels and implements an EquivariantGraphFeaturizer for seamless dataset integration, training pipelines, and testing. The approach demonstrates competitive QM9 performance against established equivariant baselines while prioritizing usability, maintainability, and open-source collaboration, with future work targeting efficiency improvements such as caching angular bases. Overall, the work lowers the barrier to applying symmetry-aware methods in drug discovery and materials science, enabling broader adoption and extended development of SE(3)-equivariant molecular modeling within a widely-used library.

Abstract

Neural networks that incorporate geometric relationships respecting SE(3) group transformations (e.g. rotations and translations) are increasingly important in molecular applications, such as molecular property prediction, protein structure modeling, and materials design. These models, known as SE(3)-equivariant neural networks, ensure outputs transform predictably with input coordinate changes by explicitly encoding spatial atomic positions. Although libraries such as E3NN [4] and SE(3)-TRANSFORMER [3 ] offer powerful implementations, they often require substantial deep learning or mathematical prior knowledge and lack complete training pipelines. We extend DEEPCHEM [ 13] with support for ready-to-use equivariant models, enabling scientists with minimal deep learning background to build, train, and evaluate models, such as SE(3)-Transformer and Tensor Field Networks. Our implementation includes equivariant models, complete training pipelines, and a toolkit of equivariant utilities, supported with comprehensive tests and documentation, to facilitate both application and further development of SE(3)-equivariant models.

DeepChem Equivariant: SE(3)-Equivariant Support in an Open-Source Molecular Machine Learning Library

TL;DR

The paper addresses making SE(3)-equivariant neural networks readily usable for molecular machine learning by integrating SE(3)-Transformer and TFN-style components into the DeepChem ecosystem. It leverages spherical harmonics and Clebsch–Gordan-based irreducible representations to build rotation-equivariant kernels and implements an EquivariantGraphFeaturizer for seamless dataset integration, training pipelines, and testing. The approach demonstrates competitive QM9 performance against established equivariant baselines while prioritizing usability, maintainability, and open-source collaboration, with future work targeting efficiency improvements such as caching angular bases. Overall, the work lowers the barrier to applying symmetry-aware methods in drug discovery and materials science, enabling broader adoption and extended development of SE(3)-equivariant molecular modeling within a widely-used library.

Abstract

Neural networks that incorporate geometric relationships respecting SE(3) group transformations (e.g. rotations and translations) are increasingly important in molecular applications, such as molecular property prediction, protein structure modeling, and materials design. These models, known as SE(3)-equivariant neural networks, ensure outputs transform predictably with input coordinate changes by explicitly encoding spatial atomic positions. Although libraries such as E3NN [4] and SE(3)-TRANSFORMER [3 ] offer powerful implementations, they often require substantial deep learning or mathematical prior knowledge and lack complete training pipelines. We extend DEEPCHEM [ 13] with support for ready-to-use equivariant models, enabling scientists with minimal deep learning background to build, train, and evaluate models, such as SE(3)-Transformer and Tensor Field Networks. Our implementation includes equivariant models, complete training pipelines, and a toolkit of equivariant utilities, supported with comprehensive tests and documentation, to facilitate both application and further development of SE(3)-equivariant models.
Paper Structure (21 sections, 25 equations, 1 figure, 3 tables, 4 algorithms)

This paper contains 21 sections, 25 equations, 1 figure, 3 tables, 4 algorithms.

Table of Contents

  1. Introduction
  2. Background and Related Work
  3. Equivariance and SE(3) symmetry
  4. Attention mechanism
  5. Method
  6. SE(3)-Transformers
  7. Residual equivariant attention
  8. Localized attention on graphs
  9. Self-attention mechanism and permutation invariance
  10. SE(3) equivariant graph convolutions
  11. Equivariant message passing and TFNs
  12. DeepChem Integration
  13. Experiments
  14. Conclusion
  15. Computation of Equivariant Features
  16. ...and 6 more sections

Figures (1)

  • Figure 1: DeepChem pipeline for SE(3)-Transformer model. Input molecular data undergoes 3D geometric featurization, dataset preparation, and splitting, followed by equivariant model implementation and training using DeepChem's API.