Inventing Cheap Catalysts by Modeling Atomic Interactions using Graph Neural Nets. How ML can help save the world from catastrophe.
Computational-chemist - the new breed of scientist
In today's rapidly evolving tech landscape, a new breed of scientist/engineer is emerging: the computational chemist, along with their counterpart in biology and physics. Their mission? To harness the power of Machine Learning (ML) to expedite the discovery of revolutionary new compounds. These compounds may serve as catalysts in planetary-scale CO2 removal machines, or as key components in ground-breaking battery chemistry that might someday allow commercial airliners to go electric! Alternatively, they could form the basis for innovative drugs or vaccines in the fight against plague and disease.
In this post, I discuss how Graph Neural Nets (GNNs) can be used to accelerate the discovery of new materials critical for fighting climate change.
At the recent CVPR conference, the premier gathering for scientists and engineers in Computer Vision and Machine Learning, we were told by Larry Zitnick, the co-architect of the celebrated MS-COCO dataset for object detection, that our community had a lot to offer to the advancement of computational chemistry. The specific problem he was tackling, and encouraged us to consider, was the speeding up of the invention of cheap, new catalysts for fuel cells.
A significant hurdle, according to him, in the acceptance of solar and wind as renewable energy sources in much of the world is their irregularity of supply. Consider solar providing the majority power for a place like Vancouver (the host of this year's CVPR), where it might only be generated in the summer months, but peak demand for electricity is in the chilly depths of winter. How do we transport the electricity over time from summer to winter, in the quantities that flow through the national grid? While batteries excel at short-term storage, they have limitations stemming from the high cost of large-scale battery storage, insufficient capacity, degradation over time and environmental impacts from their production and disposal. Therefore alternative long-term storage solutions need to be considered, even as a supplement to battery storage.
 |
The cathode in a fuel cell uses expensive platinum to accelerate the oxygen reduction reaction and generate electricity. One of the tasks at hand for computational chemists, is to invent cheaper catalysts that can be used to mass-produce fuel cells, key to long-term energy storage and solving climate change (Image courtesy Green Car Congress). |
One such method available to us is through the use of fuel cells: electricity (through electrolysis) can be converted to hydrogen, which in turn can be (using fuel cells) converted back to electricity. In other words, electricity generated by solar during summer can be used to electrolyze water into hydrogen and oxygen. Hydrogen can be stored under pressure in tanks till winter, and then combined with oxygen in a fuel cell to generate electricity during peak demand.
Within a fuel cell, the chemical reaction of $2H_{2}\rightarrow 4H^{+}+4e^{-}$ happens at the anode (one of the electrodes), where a catalyst such as platinum promotes the splitting up of the hydrogen gas fed into the fuel cell into hydrogen ions (protons) and electrons. At the cathode (the other electrode), we have $ O_{2}+4H^{+}+4e^{-}\rightarrow 2H_{2}O$, called an oxygen reduction reaction, where oxygen gas from the air combines with the incoming electrons and protons to produce water. The motion of electrons from the anode to the cathode creates a flow of electricity that can be used to drive an external circuit, with a load like an electric motor, or as desired.
Let's zoom into the reaction at the cathode. The oxygen molecule with two oxygen atoms comes into contact with the platinum (catalyst) surface. The platinum weakens the bond between the two oxygen atoms in the oxygen molecule so that two neighboring hydrogen atoms can strip away one of the oxygen atoms from the other, combine with it, and form a water molecule.
The problem is that the platinum catalyst is super expensive and the main driver of the current high cost of fuel cells. The price of platinum today is about 40 times that of silver and about half that of gold. The challenge at hand, is to quickly discover cheap new catalysts, that might be used produce fuel cells economically in the scale required for long term electricity storage. Note that electrolysis (the opposite of the reactions discussed above, where water is split into oxygen and hydrogen using electricity) also requires the use of expensive catalysts - iridium, at the anode and platinum at the cathode. We already know how to store hydrogen at scale. So; how might we go about inventing new and cheaper catalysts for electrolysis and fuel cells?
Accelerating the invention of catalysts using ML
Simulation can be used to directly model the interactions of prospective new catalytic compounds with the oxygen molecule. One would simulate a surface of the catalytic molecule, bring it in close proximity with an oxygen molecule, calculate the atomic forces at play, move the atoms according to the forces and repeat until equilibrium (achieved when a local energy minimum is obtained). The energy affiliated with this minimum at the end of this process (called a Relaxation), is directly correlated with the reaction rate, and can hence be used to determine the efficacy of the catalyst.
 |
| Atomic arrangement and forces at the surface of a prospective catalyst in a fuel cell. The task of DFT - and the ML algorithm that approximates it, is to predict these forces (and energies) given a particular arrangement of atoms. Image adapted from Larry Zitnick's invited talk at CVPR 2023. |
Central to this, is the process of computation of atomic forces, and is done using something called Density Functional Theory (DFT). DFT is a quantum mechanical model that can be used to screen potential molecules, with promising candidates later assessed with experimental testing in a lab. Unfortunately, DFT is slow; of the order $O(n^3)$, where n is the number of electrons being simulated. In practice, the simulation of a potential catalyst with 100s of atoms takes an entire day. How would we go about reducing this compute time from a day to a second, thereby enabling the simulation of millions of molecules for catalytic function?
With this in mind, Larry and others have launched the Open Catalyst (OC) Challenge, featuring datasets with 1.3 million molecular relaxations with results from over 260 million DFT calculations. The inputs are a set of atoms (comprising both the catalyst molecule (like Platinum) - as a surface and the molecules being catalyzed, the adsorbate (like Oxygen), their atomic numbers and 3D positions. The outputs are the atomic forces and a single system energy number.
One of the leading ways in which ML can accelerate the simulation and successive screening of these reactions is to come up with Graph Neural Nets trained on this data to approximate the output of the computationally complex DFT model, but 1000x faster.
 |
| Left: Oxygen atom (blue) approaching a lattice of platinum atoms (orange). Right: Inter-atomic forces represented as a single cell and expressed as a graph with atomic nodes and edges encoding the relative orientations and distances between atoms. |
The Graph Neural Network (GNN)
One way to model the forces and energies of a particular configuration of atoms, is a Graph Neural Net (GNN). Unlike a Convolutional Neural Net that operates on a 2D image (Computer Vision people are most familiar with data in this form), a GNN can be used to make inferences out of unstructured data that can be expressed in a graphical form. Our particular problem can be visualized by treating the molecular configuration in a single cell (a cell being the smallest volume of the structure that repeats itself to form the entire molecular lattice) as a graph. Each node in this graph represents an atom and edges depict interactions between neighboring atoms. For each node, we would encode properties of the atom, such as atomic number, and learned features/embeddings thereof. Interactions between nodes are encoded via edge features, which would encode the distance and angle from the nodal atom to neighboring atoms/nodes.
This spatial relationship is encoded using the Equivariance property of Spherical Harmonics (more on that in a later post maybe) to allow the molecular ensemble to be viewed from any orientation and the output - atomic forces and system energy - would change accordingly.
Message passing, a fundamental operation in GNNs would allow us to propagate and aggregate feature information between a central atom (node) and its neighbors across the layers of the network. As the depth of the network increases, each node gathers information from a progressively larger 'receptive field' of atoms, thereby allowing the network to learn and capture longer range interactions between atoms. With k layers, each node in the graph will have information from nodes that are at most k steps away in the graph.
 |
| GNN used for predicting the inter-atomic forces and the overall energy of the system. This GNN has two hidden layers indicated by the vertical columns, where messages are passed between the atoms and a set of MLPs to aggregate those messages. We have two fully connected layers (FC1 and 2) at the output to predict the forces and energy. |
The message passing update can be expressed as
$$x_{i}^{k}=\gamma^{k}(x_{i}^{k-1},\bigoplus_{j\varepsilon N(i)}\phi^{k}(x_{i}^{k-1},x_{j}^{k-1},e_{j}^{i}))$$
where $x_{i}^{k-1}$ =denotes features for node i in layer k-1 and $e_{j}^{i}$ are features between nodes j and i. $\bigoplus$ denotes a differentiable, permutation invariant function like sum, mean or max (that does not depend on the order of the messages) and $\gamma$ and $\phi$ are MLPs.
The function $\phi^{k}(x_{i}^{k-1},x_{j}^{k-1},e_{j}^{i})$ is applied to the features of node $i$, the features of each neighboring node $j$, and the edge features $e_{j}^{i}$. This function, implemented as a MLP, calculates the message that node $j$ sends to node $i$.
The $\bigoplus$ operator aggregates these messages from all neighboring nodes of node $i$.
The function $\gamma^{k}(x_{i}^{k-1}, \bigoplus_{j\varepsilon N(i)}\phi^{k}(x_{i}^{k-1},x_{j}^{k-1},e_{j}^{i}))$ then combines the aggregated messages with the features of node $i$ itself. This function, again implemented as a MLP, updates the features of node $i$.
This is repeated for each hidden layer of the network, resulting in n feature vectors for n nodes in the graph. These are then aggregated using a pooling operation to obtain a graph-level feature vector. Finally, this would be passed through a fully connected final layer to produce the final graph level output, in this case - energy of the system.
If we wanted to produce an edge-level output, that would predict the forces between every pair of atoms (nodes) in the graph, then instead of aggregating the node features after the final hidden layer, we would concatenate the node features of every pair of atoms and pass that through a final fully connected layer to predict the atomic forces along that edge.
The network would be trained using supervised learning using the data available from DFT, using say - Mean Average Error on the energy (in electron volts) and forces (electron volts per Angstrom) for the configuration of the system.
Comments
Post a Comment