Bringing Atomistic Deep Learning to Prime Time
Abstract
Artiﬁcial intelligence has not yet revolutionized the design of materials and molecules.
In this perspective, we identify four barriers preventing the integration of atomistic deep learning, molecular science, and high-performance computing.
We outline focused research efforts to address the opportunities presented by these challenges.
Bottlenecks
Atomistic deep learning encompasses neural networks that learn representations of matter - from 0D molecules to 1D nanowires, 2D surfaces, 3D proteins, polymers, and crystalline materials [1].
Unlike other canonical application areas of DL like speech and images, data for atomistic systems is expensive to acquire and heterogeneous, with a practically inﬁnite space of emergent properties of interest.
AI revolutionized natural language processing and computer vision, but we are still largely waiting for AI to revolutionize the design of matter.
Atomistic systems underpin every physical and digital technology, and the potential impact of AI-enabled matter design cannot be overstated.
In this perspective, we identify four key challenges (Figure 1) to the widespread adoption of atomistic DL and outline promising paths forward.
The ﬁrst three challenges are practical difﬁculties related to robustness and scaling methods to be "production ready.
" They correspond to the three resources that DL methods traditionally leverage at scale to achieve superhuman success: data, model size, and compute.
The fourth challenge is a more fundamental issue related to AI’s place in the toolkit for natural scientists.
Data openness
In many application areas of atomistic DL, practitioners are data starved.
Data remains proprietary to the companies that generate it, or not easily accessible to research groups that did not generate the data.
Physical experiments and high-ﬁdelity computational simulations can be expensive and time-consuming, so acquiring more labeled data is often not practical.
There are a number of promising trends addressing this challenge.
Adoption of FAIR (ﬁndable, accessible, interoperable, reusable) data principles within the research community [2, 3] makes it easier for researchers to leverage data generated outside their own labs.
Databases like the
Materials Project [4] and automated experimental validation promise to alleviate the data bottleneck in many domains [5, 6, 7].
Active learning and Bayesian optimization strategies for intelligent sampling are shifting the paradigm from random or brute-force data collection to targeted data acquisition [8, 9].
However, it remains important to continue unfocused, high-throughput approaches that generate datasets to train general purpose models and to explore new applications in an open-ended fashion.
Datasets generated through intentional experimental design strategies may only be useful to the groups that generated them (and that have the capabilities to generate them).
Representation expressivity Geometric priors and therefore geometric DL are a natural choice for atomistic systems [10, 11], which can often be represented as graphs of atoms (nodes) and chemical bonds (edges).
However, graph neural networks often struggle to compete against human-crafted featurizations like chemical ﬁngerprint vectors and much simpler ML methods like random forests and logistic regression in low-data limits [12, 13].
To achieve widespread adoption, atomistic DL methods and molecular representations must be expressive enough to capture the truly staggering diversity of emergent behavior exhibited by matter [14], while also being robust and easy to use for non-DL experts.
Following the surprising performance of GPT-3 [15], one path forward is to leverage large language models trained on string-based representations of matter [16].
A natural extension is to combine the self-attention mechanism of transformers with the equivariance inherent to molecular graphs [17].
Identifying optimal pre-training tasks is also an ongoing and fruitful area of research [16, 18].
Scaling studies [19] are needed to investigate whether these architectures and representations are already sufﬁcient to handle generic molecular ML tasks, or if further methodological advances are needed.
Compute availability A common thread in headline-grabbing AI achievements is the need for massive compute resources, which are often only available to large corporations.
Recent reviews of large language models have argued that non-proﬁt research organizations should also be investigating models at scale [20].
Most atomistic DL methods are not tested beyond small benchmark datasets on a single GPU [18].
Even doing a thorough hyperparameter optimization for a moderately-sized dataset can be a burden.
As the business value of atomistic DL grows, large tech companies will become increasingly active in the ﬁeld.
To take two recent examples, industry AI research groups have made exciting contributions to catalyst discovery with the Open Catalyst Project [21] and protein folding with AlphaFold [22].
There is an opportunity to anticipate this trend in atomistic DL and prepare for it, to avoid the paradigm seen in ﬁelds like NLP, where fundamental methodological advances are largely left to university research groups and large-scale studies are done by companies.
An emphasis on open science [23] and closer integration between academic research groups and highperformance computing centers, cloud providers, and industry AI groups can alleviate this blocker.
Engineering advances like PyTorch Distributed [24] lower the barriers for operating at scale, and HPC centers like MIT Supercloud [25] and NERSC are increasingly prioritizing GPU resources, GPU acceleration, and interactive ML workﬂows.
As scalable model training becomes easier, researchers have the opportunity to investigate emergent behavior in larger models and critically examine where new methods are needed and where scaling existing methods is sufﬁcient.
Insight extraction Scientists and engineers need to trust atomistic DL methods to drive further adoption.
Human practitioners don’t have an intuition about the failure modes and so there is a resistance to conduct physical experiments based solely on model predictions.
Trust is built at a practical level when methods are useful, and at a more basic level when methods provide new, surprising, and relevant scientiﬁc insights.
Black-box models are useful when they automate a tedious task [26] and the bounds of the model’s applicability are well-understood.
A loftier goal is for atomistic DL to provide novel insights - discovering new functional forms of matter, uncovering causal or mechanistic processes underpinning empirical observations, etc.
Atomistic DL that discovers not only unexplored, but also synthetically accessible [27, 28, 29, 30], matter will build trust outside the core community of DL-practitioners.
Uncertainty quantiﬁcation of model predictions adds another necessary dimension that will clarify when models are operating outside of their domain of applicability, and uncertainty is a familiar metric for experimentalists to use to assess the trustworthiness of new methods [9, 31].
Conclusion The increasing availability of molecular modeling data and the heterogeneity of datasets and tasks present an exciting opportunity for real-world impact by integrating atomistic DL, highperformance computing, and molecular science.
Investigating the basic structure of atomistic DL models and their scaling behavior will accelerate the widespread adoption of these methods and uncover what, if anything, learned representations of matter can tell us about matter itself.