Gaussian representation for image recognition and reinforcement learning of atomistic structure

Figure 1: (a-b) ASLA algorithm. (c) One-hot encoded ASLA representation, CNN & Q-map output. (d) Adapted ASLA algorithm with a learnable Gaussian-kernel

This paper builds on a previous project ASLA¹ that defined a reinforcement learning (RL) algorithm for optimisation of atomistic structures. This work was inspired by advancements in RL that at the time had gathered a lot of attention, namely super-human performance in Atari² and Go³. The networks used in both of these work were convolutional neural networks (CNN) which are excellent for learning based on images. This is quite a natural choice for both Atari, where the game is a literal stream of images, and Go where the board is a grid which is easily represented as an image.

Reinforcement learning algorithms generally learn a policy \(\pi(a | s)\) which is a probability distribution over actions \(a\) for a given state \(s\). For the game of chess the state could be the positions of all the game pieces and which player’s turn it is. An action, like in a normal game of chess, is a movement of one of the pieces. One way of learning such a policy is to learn the expected return \(Q(s, a)\) of an action. Without dwelling into the details \(Q(s, a)\) is really about learning how good an action is - in the chess analogy this could be the expected⁴ win percentage.

In order to apply a similar CNN-based methodology in ASLA the continouos coordinates of atoms are discretized to a grid - which like the Go board may be represented as an image. A simplified depiction of the algorithm is shown in Figure 1(a) and the original architecture in Figure 1(c). In the left-hand side of Figure 1(c) the discretized one-hot representation of an atomic structure is depicted, here black pixels indicate the presence of an atom. This image is fed through a CNN that produces Q-values. The algorithm then uses these Q-values to determine positions where additional atoms may be placed in order to produce stable a configuration.

For games like chess and Go moving adjacent pieces is unlikely to be strongly correlated - or said in an other way; just because moving a knight in a certain is a good move doesn’t mean that moving the tower next to it is. However, for atomic system this correlation does exist, if a particular position has been found to be favourable then neighbouring positions are likely to be too. The paper introduces an adapted network architecture that takes advantage of this correlation by inserting a convolutional layer that is restricted to Gaussian-kernels on the input side. We show that this improves the performance of the algorithm, allowing it to solve optimization problems more efficiently.

Footnotes

1. M. S. Jørgensen et al., Atomistic structure learning, J. Chem. Phys. 151, 054111 (2019)

2. Playing Atari with Deep Reinforcement Learning

3. Mastering the game of Go with deep neural networks and tree search

4. Here ‘expected’ refers to the expectation value. That is, it considers the probability and outcome of all subsequent actions.