{'fraction_non_alphanumeric': 0.028169014084507043, 'fraction_numerical': 0.010449795547478418, 'mean_word_length': 5.245256251108352, 'pattern_counts': {'":': 0, '<': 0, '<?xml version=': 0, '>': 0, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
The Open MatSci ML Toolkit: A Flexible Framework for Machine Learning in Materials Science

We present the Open MatSci ML Toolkit: a flexible, self-contained, and scalable Python-based framework to apply deep learning models and methods on scientific data with a specific focus on materials science and the OpenCatalyst Dataset. Our toolkit provides: 1. A scalable machine learning workflow for materials science leveraging PyTorch Lightning, which enables seamless scaling across different computation capabilities (laptop, server, cluster) and hardware platforms (CPU, GPU, XPU). 2. Deep Graph Library (DGL) support for rapid graph neural network prototyping and development. By publishing and sharing this toolkit with the research community via open-source release, we hope to: 1. Lower the entry barrier for new machine learning researchers and practitioners that want to get started with the OpenCatalyst dataset, which presently comprises the largest computational materials science dataset. 2. Enable the scientific community to apply advanced machine learning tools to high-impact scientific challenges, such as modeling of materials behavior for clean energy applications. We demonstrate the capabilities of our framework by enabling three new equivariant neural network models for multiple OpenCatalyst tasks and arrive at promising results for compute scaling and model performance.


INTRODUCTION
Recent years have seen great advances in applying advanced machine learning methods, especially novel deep learning methods, to scientific challenges that rely on computational modeling in the development of new physical systems (Krenn et al., 2022;Axelrod et al., 2022;Jumper et al., 2021). Modern deep neural networks trained to reproduce physical calculations, which are used to understand and optimize the behavior of real systems, can operate with high accuracy and are often orders of magnitude faster compared to methods based solely on human-derived calculations (Friederich et al., 2021;Chen & Ong, 2022;Fung et al., 2021). However, these representational capabilities usually come with significant engineering costs: The large variety of physical systems studied in chemistry, as well as the methods to represent those physical systems along with the datadriven paradigms used to predict their behavior, can make it cumbersome, time-consuming and error-prone to adapt models and tools to new datasets and applications. To mitigate these frequent challenges, we present the Open MatSci ML Toolkit, a flexible framework to develop deep learning techniques for chemistry and materials science applications. Catalysts are essential components of chemical processes that help accelerate the rate of various chemical reactions. Catalytic materials design, especially the design of low-cost catalysts, remains an ongoing challenge that will continue to become more and more important for a variety of applications, including renewable energy and sustainable agriculture. The OpenCatalyst Project, jointly developed by Fundamental AI Research (FAIR) at Meta AI and Carnegie Mellon University's Department of Chemical Engineering, encompasses one of the first large-scale datasets to enable the application of machine learning (ML) techniques, with the full dataset containing over 1.3 million molecular relaxations of 82 adsorbates on 55 different catalytic surfaces (Chanussot* et al., 2021). The original release from 2019 has also been supplemented by subsequent updates in 2020 and 2022, along with an active leaderboard and annual competition (Tran et al., 2022). The notable effort of providing high-quality data for catalytic materials is a major step forward in enabling ML researchers and practitioners to innovate on materials design challenges at great computational scales, and has already enabled the development of new geometric deep learning architectures ( (Klicpera et al., 2020b) (Gasteiger et al., 2021)), some of which have been trained with nearly a billion parameters (Sriram et al., 2022).
While the toolkit of the original OpenCatalyst repository is very powerful, it incorporates a significant amount of complexity due to various interacting pieces of software: model definitions, functions for distributed training, and task ab- stractions. These components are often not self-contained, which makes it difficult for new ML researchers to navigate and interact with the repository, create new architectures or modeling methods, and run experiments on the dataset. To address these usability challenges, we introduce the Open MatSci ML Toolkit, a flexible and easy-to-scale framework for deep learning on the Open Catalyst Dataset. We designed the Open MatSci ML Toolkit with the following basic principles and compelling features: • Ease of use for new ML researchers and practitioners that want to get started with the OpenCatalyst dataset.
• Support for DGL (Wang et al., 2019) for rapid GNN development to complement the original repository's usage of PyTorch Geometric . In this work we showcase a set of models built on our toolkit for OCP-20 described in Section 4.

BACKGROUND & RELATED WORK
Geometric Deep Learning (GDL) generalizes neural networks to non-Euclidean domains such as graphs and manifolds (Bronstein et al., 2017;. Such domains are increasingly used to model systems in scientific applications, such as molecular and crystal structures like those found in the OpenCatalyst project. In particular, GDL provides a way to represent entity interactions and various invariances and equivariances under geometric transformations-both vital for modeling molecules and catalyst interactions. Graphs and point cloud data 1 are the most common representations of such molecule-catalyst systems with different GDL methods operating on inputs with various inductive biases. Graph Neural Networks, one of the earliest GDL applications to atomic systems, process a graph representation of molecular and solid-state material systems (Gilmer et al., 2017). These neural networks are able to take advantage of the natural representation of molecules as graphs, whereby message passing between nodes (atoms) resembles atomic interactions. Despite this, vanilla message passing GNNs are unable to distinguish between certain motifs and are bounded by the performance of the 1-Weisfeiler-Lehman (WL) isomorphism test (Xu et al., 2019). Even with recent advances in pushing GNNs beyond the WL-test (Morris et al., 2019), such networks do not utilize any of the known physical symmetries, thereby reducing their effectiveness on molecular data beyond their theoretical limitations.
Equivariant Neural Networks restrict the space of functions learnable by the network to obey symmetries found in the input data. Symmetries are a powerful inductive bias that drastically narrow the solution space that needs to be explored by the neural network, thereby making learning easier and improving generalization. Physical systems often exhibit many symmetries, either in their overall structure (e.g. water within the C 2v point group), within structural motifs (e.g. three-fold rotation symmetry in methyl rotors), and, when nuclear dynamics are considered, complete nuclear permutation groups (Bunker & Jensen, 1998;Williams & Eisfeld, 2020) (e.g. ethane and the G 36 group (Mellor et al., 2019)). Composing many layers that individually respect symmetries relevant to problems of interest has proven to be a powerful method to design deep neural network architectures for a range of modeling challenges. Exploiting such symmetries has emerged as a driving force for design-ing neural networks for physics models, for example as equivariance under the SO(3) group; that is, equivariance under translations and rotations in 3D space. Various works, such as Thomas et al. (2018); Anderson et al. (2019); Fuchs et al. (2020), leverage group representations to achieve full equivariance. While equivariant models can be more parameter efficient, they can also be computationally expensive or constrained in expressivity due to the model basis.
Equivariant/Directional Graph Neural Networks merge both graph and point cloud representations to leverage the benefits of both representations. One of these equivariant GNNs (E(n)-GNN by Satorras et al. (2021)) achieves this by separating node features into spatial components that are equivariant and invariant atomic features. These are then mixed together in a series of aggregation steps that preserve the spatial equivariance and feature invariance. Other GNNs take a more direct approach to equivariance by embedding angles, distance and dihedral angles in classic basis sets such as Legendre polynomials and spherical harmonics (GemNet (Gasteiger et al., 2021) and DimeNet (Klicpera et al., 2020a)), which possess well-known symmetry properties. These models can be viewed as directional GNNs since they typically expand their receptive field beyond the 1-hop neighborhood utilizing angular data to direct the message propagation. Directional models have high representational power, but come at a significant computational and memory cost. The performance of these models depends on the task and complexity of the data: EGNN can perform on-par with DimeNet and GemNet on simple classification tasks (Satorras et al., 2021); however, directional models typically perform better in more advanced applications such as conformer search (Ganea et al., 2021;Jing et al., 2022).

Short-Range Equivariant
Networks restrict the ability of signals to travel long distances as layers are composed. In a typical graph neural network, signals from one node are able to travel an additional neighbor hop away in the graph for every message passing layer in the network; in contrast, short-range networks avoid transmitting these signals. These short-range restrictions can help promote localized representations that are more efficient to learn. One such architecture uses geometric algebra-or Clifford algebra-to formulate rotation-invariant and -equivariant combinations of input vectors, which are transmitted in a permutation-equivariant way via an attention mechanism (Spellings, 2021). The Allegro framework (Musaelian et al., 2022) learns to produce localized representations from a series of equivariant tensor products, rather than using geometric algebra to achieve rotation equivariance.

SOFTWARE FRAMEWORK
The Open MatSci ML Toolkit software framework is designed with great emphasis on abstraction and inheritance in order to maximize reusability and agility for machine learning researchers. These ideas are achieved in part by presentday best practices in Python as an object-oriented language, and through modern, specialized frameworks such as Py-Torch Lightning and DGL. We believe these design choices make it significantly easier to apply novel model architectures and training techniques to scientific data, in particular the OpenCatalyst dataset. In the following sections, we will discuss changes in abstractions and refactoring steps from the original OpenCatalyst implementation.

PyTorch Lightning Refactor
In modern AI/ML workflows, the concept of "MLOps" comprises the lifecycle from model conception and implementation, training and testing in a variety of software/hardware environments, to drawing inferences on new data, and all of the iterative cycles in between. Thus, a non-negligible amount of time spent by researchers for new workloads is typically in engineering: interfacing data with new architectures, metric logging, performance profiling, and ensuring consistent functionality through the development process from testing locally on a laptop to distributed training on multiple computing nodes and across different types of accelerators. Because of the grand scale that OpenCatalyst aims for and successfully achieves, a large amount of the original codebase corresponds to performance and functionality; meaning that complexity is necessary to be able to take advantage of data parallelism, to perform hyperparameter optimization, and to support the various catalyst prediction tasks. This lays a significant amount of responsibility on both developers and users: the former must create a comprehensive suite of tests and rely heavily on CI/CD to ensure functionality, and the latter must navigate a maze of software dependencies and documentation, which are also maintained by the developer.
One half of the conceptual changes in the Open MatSci ML Toolkit-the other half being the primary graph frameworkis to offload MLOps related components to a well-designed and maintained framework: PyTorch Lightning (Falcon & The PyTorch Lightning Team, 2019). By reusing dataset and framework components from OpenCatalyst and relying on the pipeline abstractions from PyTorch Lightning, we can maintain a more compact codebase while providing more flexibility, extendibility, transparency, and functionality. Figure 2 illustrates the end-to-end pipeline/directed acyclic graph for the Open MatSci ML Toolkit, whose elements should be relatively familiar to those who have used OpenCatalyst and/or PyTorch Lightning.

Data abstraction
In order to support future neural network research, we expanded the scope of the original OpenCatalyst dataset to  Figure 2. Illustration of the Open MatSci ML Toolkit pipeline with concrete components using the Structure to Energy & Forces ("S2EF") task. Dataset/task splits and configurations are specified through LightningDataModules. Task specific LightningModules encode the logic for training, metric logging, and how data is passed from dataset to an underlying abstract deep learning model. The Trainer interface provides an the ability to control feedback (e.g. logging, progress bars), training flow (Callbacks), and XPU usage without the need to modify the pipeline source code. support graph and non-graph data structures, as well as implemented a number of quality of life improvements to the general developer workflow. We refer the reader to the Appendix (i.e. Figure A.5) for more details pertaining to the changes, and here we only briefly highlight the core differences in user experience.
One of the core goals of the Open MatSci ML Toolkit is to have continuity from developing and testing on local environments-such as laptops-to using the pipeline in high performance computing environments. In terms of data pipeline abstraction, on the one end the Open MatSci ML Toolkit provides preprocessed, miniature (∼100 graphs) development or "devset"s: this circumvents the need to download, extract, and preprocess the data on personal computers constrained by storage and by computational power, while allowing researchers to prototype on the full pipeline. The development sets are created by taking random subsplits of the smaller data splits from the OpenCatalyst dataset. We also provide the mechanism for creating other splits are provided with the Open MatSci ML Toolkit, facilitating further research into data efficiency. To use the devsets for development, there is a convenient mechanism for retrieving the DGL version of each task: On the other end of the spectrum, where one wishes to distribute the dataset across multiple workers on multiple compute nodes, users can leverage the same data modules as the miniature case: the DistributedDataParallel data sampling and loading is offloaded to PyTorch Lightning internals. We will discuss this further in Section 3.1.3.
As alluded to earlier in this section, the Open MatSci ML Toolkit data pipeline abstraction is designed to facilitate exploration of other data representations of materials systems, including algorithms less strongly linked to graph-based message passing. This includes recent approaches based on geometric algebra or Clifford algebras for point clouds by Spellings (2021), as well as methods based on tensor products, such as Musaelian et al. (2022). These algorithms operate locally on the bonds within a point cloud without typical message passing from one atom to another, while retaining the data efficiency imparted by model equivariance.

Model abstraction
The model abstraction, as seen in the bottom left nodes in Figure 2, comprises a neural network backbone and a task-specific LightningModule. In the concrete example described in Section 4, the EGNN model represents a subclass of an AbstractEnergyModel: a model that takes arbitrary input, and predicts the energy. For instance, a graph-based model will process nodes and perform some readout operation to reduce node and graph level features to a scalar value for the energy. At a higher level, the task specific S2EFLitModule is instantiated by passing an instance of EGNN, and implements the logic for training (i.e. forward-backward passes), validation and testing, as well as logging. By conceptually separating the model (i.e. the neural network itself) from the training mechanism, researchers only need to focus on architecture development by subclassing AbstractEnergyModel, as the rest of the pipeline stays the same barring changes in which kind of data is required by the model.

Training loop
For typical regression tasks such as the initial structure to relaxed energy ("IS2RE"), the training loop takes advantage of automatic optimization in PyTorch Lightning, meaning that the backpropagation and optimization steps are abstracted away from the user. In cases with specialized autograd steps, such as the force component of "S2EF", and in many physics inspired neural network architectures, backpropagation is manually performed. While this highlights a gap in coverage by MLOps frameworks such as PyTorch Lightning, it also illustrates their open-endedness: one can maintain the flexibility necessary for research, while keeping the many other benefits and best practices guided by PyTorch Lightning.
With regards to changes in the behavior of the pipeline, users can rely on modular PyTorch Lightning components that can trigger at certain parts of the training loop-such as at the beginning of a training batch, or at the end of the validation epoch-with the Callbacks abstraction.
The scope of such events include, but is not limited to: early stopping according to some validation metric, model checkpointing based on training intervals or metrics, as well as more specialized and flexible tasks like code profiling targeted to specific pipeline components-all without the need to inspect and modify the core source code.
Regarding logging/experiment tracking, the Open MatSci ML Toolkit relies on PyTorch Lightning Loggers, which are able to log locally to CSV or Tensorboard format, as well as to hosted services such as Weights and Biases and SigOpt; the latter of which can also be used for hyperparameter optimization (HPO), which makes up a valuable tool in a researcher's toolkit. Users are not confined to a single choice of logging or HPO platform, providing flexibility in tooling that fits a given researcher's workflow.
The final component relevant to the training process is the PyTorch Lightning Trainer class, which orchestrates the components mentioned above and executes training, validation, testing, and inference loops. The Trainer interface also configures more performance oriented settings such as accelerator usage, distributed compute, and mixed precision as shown in the example configuration below: The main advantage is being able to seamlessly navigate between development and training cycles: the core pipeline remains unchanged, yet with a simple change in configuration, the user is able to take advantage of computational resources as they become available. Under the hood, the Lightning abstractions handle data movement to devices, autocasting in correct contexts, and orchestrating workers.

Deep Graph Library (DGL) Refactor
The original OpenCatalyst repo leverages Pytorch Geometric (PyG) for implementing various neural networks described in Section 2. In the Open MatSci ML Toolkit, we chose to complement the original implementation by building on top of the Deep Graph Library (DGL). While both PyG and DGL are highly performant libraries for graph neural networks, and the decision to choose one over the other is often subjective, we motivate our choice of DGL for this library as follows: • Graph Abstraction: The DGL graph data structure dgl.DGLGraph offers more flexibility for storing diverse molecular data over the PyG structure torch_geometric.data.Data. This abstraction allows for more general data pipelines amicable to experimentation and ablation studies; two qualities that are important in scientific exploration.
• Cross-platform Optimization: While both DGL and PyG are well-optimized for single-node GPU deployment, DGL also supports additional platforms, in particular efficient sparse matrix algorithms for CPUs with documented compute benefits on datacenter CPUs.
• Support for Sampling-Based, and Distributed Training beyond Data-Parallel: many applications involve large graph data that does not fit onto a single GPU. Such cases require specialized sampling techniques to either shrink the graph size or distribute the storage across multiple nodes. While both PyG and DGL support sampling-based training, DGL is more mature when it comes to sampling graphs and running distributed training of GNNs (Mostafa, 2021;Zheng et al., 2021).
Next, we discuss each of the above points in more detail.

Graph Abstraction
PyG's Data has a graph data structure composed of the following fixed attributes: node feature matrix x, edges edge_index, edge attributes edge_attr, labels y, and node position matrix pos. Additional attributes can be set using keyword argument collection in the constructor. In contrast, DGL's DGLGraph provides a dictionary-style access to graph data through DGLGraph.ndata (node features) and DGLGraph.edata (edge features). While, these dictionaries do not impose any restrictions on feature names or number, the DGLGraph object will enforce that the feature dimension match the number of nodes and edges.
While PyG's Data covers a large part of the use cases in the GNN world, we argue that using the flexible DGLGraph as example representation decouples the dataset from the model-specific data pipeline, while retaining data consistency such as matching the number of features and edges. This lends itself to a wide variety of customized pipelines that can explore various model explorations and ablation studies. For example, consider a molecular data structure that contains various features on both atoms and bonds: QM9 (Ramakrishnan et al., 2014), a common molecular property prediction benchmark, includes atom features detailing atom coordinates x, atom type tp, atomic number z, number of hydrogens nH, and hybridization (Gilmer et al., 2017). A given set of molecular features would then contain the following fields: while in DGL's DGLGraph this would look like 1 src, dst = zip( * example_dict.pop('bonds')) 2 dgl_graph = dgl.graph((src, dst)) 3 dgl_graph.ndata = example._asdict() The above example demonstrates that the DGLGraph can accommodate diverse data without having to perform any special model-specific preprocessing while retaining graph consistency across assigned nodes and features. PyG Data, on the other hand, either requires feature preprocessing and assignment to the x field or generically attaching the features to the data object.

Cross-Platform Optimization
GPU acceleration plays a pivotal role in modern neural network training in general, including for GNNs. Recently, there has been an increasing interest in deploying GNNs, particularly in distributed computing settings, on other a diverse hardware platforms including datacenter CPUs and custom deep learning accelerators from nascent hardware vendors. While both DGL and PyG are well optimized for CUDA, DGL implements highly efficient kernels for sparse matrix operations central to graph processing through libxsmm, providing more flexibility in the type of hardware used for training and inference (Avancha et al., 2020).

Sampling and Distributed Training Beyond Data-Parallel
Increasingly, single device training of GNNs is reaching its limits as graphs and models scale in size and memory requirements. A common way to address these limits is to shrink the size of the graphs by sub-sampling the graph so that it fits into single device memory (Hamilton et al., 2017). Yet, this type of sampling can lead to performance degradation due to the discarded neighborhood information, prompting a surge in new distributed training for GNNs (Hamilton et al., 2017;Md et al., 2021;Mostafa, 2021;Zheng et al., 2021). DGL already supports various forms of graph distributed training for both sampled and full batch training, and this capability could prove invaluable for molecular design datasets as both model and data complexity grows. While these types of distributed training might not be directly relevant to OpenCatalyst, they could be useful in future applications with large system sizes, such as largescale molecular dynamics with billions of atoms (Musaelian et al., 2022;Guo et al., 2022;Shibuta et al., 2017).

Point Cloud Representation
Up to this point, our discussion in Section 3 has mainly highlighted graph representations of atomic data, and subsequently, graph neural network architectures. While graphs are a natural way to represent systems of bonded atoms, we also aim to support alternative data representations in our framework, in order to enable development of diverse deep learning architectures. One such example, which is already implemented in Open MatSci ML Toolkit, are atomcentered point cloud representations that are used to achieve translation and rotation equivariance in several recent architectures, including NequIP (Batzner et al.), Geometric Algebra Attention Networks (Spellings, 2021), and Allegro (Musaelian et al., 2022).
Atom-centered point cloud representations often incur a high memory cost, leading us to build our procedure on top of the data abstractions discussed earlier, while only constructing neighborhood point clouds for the most relevant atoms. In the case of OpenCatalyst, this means that we extract the molecular adsorbate (typically a few atoms) and the surface, which-according to prior chemical knowledgeshould contribute the most to describing the system as a whole. Furthermore, in each catalyst + molecule system, we randomly sample a number of additional bulk crystal atoms below the surface layer that constitute the substrate for a more holistic description.
We construct atomic features using symmetric one-hot embeddings based on the corresponding atomic numbers, which seeks to mirror the treatment for the positions and embed information about the central atom, as well as each noncentral atom in the point cloud. Finally, the point cloud data are batched by padding molecule and substrate dimensions with zeros that can subsequently be masked during computation. Given that the point cloud representation shares the same core components as graph datasets, the same pipeline can be used with minimal code changes, largely owing to inheritance patterns (see Figure 5 in Appendix C).

EXPERIMENTS & TESTING
We designed our experiments to showcase the capabilities of our toolkit highlighted in Section 1. Specifically, our experiments demonstrate: • Seamless Compute Scaling: We show the benefits of compute scaling across multiple devices in a single node and multi-node setting for different OpenCatalyst tasks described in Section 4.1. As shown in Table 1, the success of the new models relies on our compute scaling capabilities and our flexible MLOps data and model pipelines.
All of the models presented have the inductive bias of rotation equivariance; in general, equivariance for functions f (I), g(I) for an entity I is defined as: f (g(I)) = g(f (I)).
Intuitively this means that the features of an entity transform equally with a given manipulation, such as a rotation. This is particularly useful for property prediction in material compounds, such as the IS2RE and S2EF tasks, where rotation of the entire compound itself does not affect the properties of the compound. Each of the models, however, impart equivariance in different ways through their architecture. We refer the reader to Section 2 and the original papers for further details.

Hardware Scaling Capabilities
As an illustrative example, we deploy E(n)-GNN using the Open MatSci ML Toolkit to the OCP-20 S2EF task with 200K training samples and 1M validation samples, which is amendable to studying the compute performance one can achieve using our framework. We apply distributed dataparallel training based on the PyTorch library (Li et al., 2020) to scale training across multiple devices on single GPU compute node, as well as multiple GPU compute nodes. As seen in Figure 3, single node scaling to multiple GPUs for the S2EF and IS2RE tasks shows a decreasing benefit as more GPUs are added, asymptotically approaching the apparent limit of parallelizable computation benefits for this particular setting. Epoch training throughput in the multinode setting for two nodes suggests close to linear scaling. While the overall benefits of compute scaling increase with a more intensive task, the overarching compute time, both in core-time and in wall-clock time, also increases making the overall experiment more costly.   Table 2.

OCP-20 Task Performance
We leverage our toolkit to perform training experiments on two large-scale tasks within OCP-20: 1. IS2RE-All: Initial Structure to Relaxed Energy for all data contained in OCP-20 (∼500K training samples). IS2RE is trained on MAE (L1) loss: with f θ being the model parameterized by θ, y i being the energy labels and z i being the input data features.  We perform the IS2RE-All task for all three models and the S2EF tasks for E(n)-GNN and MegNet. We did not perform S2EF with Gala because the overall memory requirements were too large to fit onto GPU, mainly due to the significant additional memory cost of force predictions. Gala training on IS2RE was already strained in memory utilization given that we applied a batch size of 1 point cloud and performed gradient accumulation across 64 batches to compensate for the small batch size. The full list of relevant hyperparameters can be found in Appendix A.

IS2RE-All Task Performance
The IS2RE-All training results in Table 2 and Figures 4a, 4b & 4c show that Gala performs best on training loss and MegNet performs best on in-distribution validation loss. E(n)-GNN performs worst across all in-distribution and outof-distribution tasks. The trend of the training loss, whose deviation corresponds to the loss across various batches in a given epoch, indicates a downwards slope with recurring deviation pattern. This suggests that the networks may find particular data samples more challenging than others as it cycles through the same batches over different epochs. The learning curves for all models also show that the indistribution validation loss generally follows the training loss, which is expected for this setting. In the case of outof-distribution tasks, MegNet performs slightly better than Gala on OOD-Catalyst with Gala performing best on OOD-Adsorbates and OOD-Both. Gala's relatively consistent performance across ID and OOD tasks suggests that Gala may have greater generalization ability compared to E(n)-GNN and MegNet. It is also worth noting that all three models have the most difficulty on the OOD-Adsorbates task, indicating that it is more difficult to generalize for new molecules in the catalyst+molecule system compared to changing the underlying catalyst crystal structure.

S2EF Task Performance
The loss shown in Table 2 Table 1 also show the greater compute cost associated with the S2EF-2M task, both in the amount of hardware used as well as the related wall-clock time, driven by larger dataset size and greater task complexity.

Comparisons with existing models
One of the primary motivations of this paper, and the future work we hope to enable, is the search for new neural network architectures to accelerate materials discovery. We hence contextualize the results in Table 2 (Chanussot* et al., 2021)) and more recently GemNet-XL (0.38 eV, (Sriram et al., 2022)), suggest that our reported results are competitive in task performance.
In the case of S2EF, our comparisons have additional limitations given that our training set of 2M samples is at least an order of magnitude smaller than those applied in other works (Chanussot* et al., 2021;Sriram et al., 2022). Nev-ertheless, E(n)-GNN and MegNet both present promising performance when compared to GemNet-XL trained on the full S2EF training set (∼133M) by Sriram et al. (2022), who reported OOD-Both joint (sum of energy and forces) test errors of 0.38, compared with 0.43 [E(n)-GNN] and 0.38 (MegNet). We emphasize that, while these are not directly comparable, we believe that our results highlight the potential in discovering novel architectures-particularly parameter efficient ones like E(n)-GNN-and data representations, such as the atom-centered point cloud representation described in Section 3.3, which are facilitated by the Open MatSci ML Toolkit. As the framework matures, we look forward towards scaling up experiments further and enabling direct comparisons with the OCP-20 leaderboard, while also continuing to grow the application of machine learning models for materials discovery applications (Friederich et al., 2021;Chen & Ong, 2022).

Future Work
In this paper, we introduced the Open MatSci ML Toolkit and demonstrated how it can be applied to train advanced geometric deep learning models on different tasks within the OpenCatalyst dataset through seamless compute scaling and automated MLOps. While the primary aim of this paper was to showcase the capabilities of the framework, the challenge of training better and more effective machine learning models on the OpenCatalyst dataset remains. We hope that our results provide a convincing starting point to continue to build on top of the current version of the framework and enable future research in geometric deep learning applied to materials science. Moreover, while the OpenCatalyst remains the largest for materials science, further relevant datasets and benchmarks (Dunn et al., 2020;Jain et al., 2013;Kirklin et al., 2015;Calderon et al., 2015) could be integrated into the Open MatSci ML Toolkit to accelerate the development of advanced machine learning tools for materials science. As can be seen in the definition of trainer, this short script already performs training on two GPUs with users being able to change the devices variable to adjust the numbers of GPUs they want to leverage for distributed training on a single node.  Figure 5. Inheritance diagram for the data abstraction in Open MatSci ML Toolkit. The main user interaction layer is presented at the bottom, corresponding to subclasses of LightningDataModule. Arrows denote directional relationship between the classes; the dashed line indicates that the PointCloudDataset wraps the task specific datasets, whereby the user is provided with a sampled point cloud representation of the original graphs.