# Balancing Robustness and Fairness via Partial Invariance

Moulik Choraria
University of Illinois at Urbana-Champaign
&Ibtihal Ferwana
University of Illinois at Urbana-Champaign
&Ankur Mani
University of Minnesota, Twin Cities
&Lav R. Varshney
University of Illinois at Urbana-Champaign
Corresponding author<EMAIL_ADDRESS>
###### Abstract

The Invariant Risk Minimization (IRM) framework aims to learn invariant
features from a set of environments for solving the out-of-distribution (OOD)
generalization problem. The underlying assumption is that the causal
components of the data generating distributions remain constant across the
environments or alternately, the data “overlaps” across environments to find
meaningful invariant features. Consequently, when the “overlap” assumption
does not hold, the set of truly invariant features may not be sufficient for
optimal prediction performance. Such cases arise naturally in networked
settings and hierarchical data-generating models, wherein the IRM performance
becomes suboptimal. To mitigate this failure case, we argue for a _partial
invariance_ framework. The key idea is to introduce flexibility into the IRM
framework by partitioning the environments based on hierarchical differences,
while enforcing invariance locally within the partitions. We motivate this
framework in classification settings with causal distribution shifts across
environments. Our results show the capability of the partial invariant risk
minimization to alleviate the trade-off between fairness and risk in certain
settings.

##  1 Introduction

In standard machine learning frameworks, models can be expected to generalize
well to unseen data if the data follows the same distribution [29]. However,
it has been noted that distribution shifts during test time (due to data being
from different sources, locations or times) can severely degrade model
performance [18, 24]. For instance, in a computer vision task to classify
camels and cows, the authors ([6]) showed that during testing, the model with
perfect training loss misclassified cows as camels during test time, if the
image background was that of a desert. The issue arose due to the model
picking up a strong but spurious correlation: during training, most cow images
had green pastures, whereas camel images had deserts. Consequently, the model
began to associate cows to green pastures and camels to deserts, leading to
misclassification during test time. The key observation is that the vanilla
classifiers are often prone to exploiting such spurious features during
training (in this case, the image background), which can then hamper its
performance in out-of-distribution (OOD) tasks. To fix this problem, multiple
lines of research have explored the idea of enforcing invariance during
learning [4, 14, 3, 17], to find invariant, causal, features that are
independent of the input distribution which may vary across training
environments [28]. With respect to the vision task considered before, the
features specific to the cow in the image (color, shape, etc.) would
correspond to invariant features. Therefore, we can obtain a robust predictor
by capturing only invariant features, which inevitably capture some causal
relationships about the data and thus generalize well in OOD tasks.

The Invariant Risk Minimization (IRM) framework is a promising step to address
OOD generalization [4], which consists of learning a data representation
$\phi$ that captures invariant features, followed by training a classifier on
the invariant features such that it is close to optimal for all environments.
The IRM classifier can perform well in OOD tasks if the conditional
expectation of the target given the representation features, does not vary
across environments [4, 3], which intuitively means that the data
distributions overlap across environments. Due to the additional constraint
that the IRM classifier should be optimal for all environments, there is an
implicit notion of group fairness in IRM that encourages similar performance
across environments, which has been explicitly studied in follow-up works [17,
1]. However, it is unclear how IRM performance is affected if environments do
not overlap adequately. Nevertheless, it is essential to overcome any
potential shortcomings of IRM in this setting, since minimally overlapping
environments often occur in real-world applications. For instance, in
discussions in ethnic community forums, while some linguistic features are
shared among communities, there are additional features that are specific only
to certain intra-community comments but not outside the community [12, 23].
Thus, one would expect the performance in a downstream task such as hate-
speech detection would degrade if such intra-community (non-invariant)
features are discarded by the IRM classifier.

In this work, we propose a _partial invariance_ framework that creates
partitions within environments, thus introducing the flexibility within IRM to
learn non-invariant features locally within the partitions. We show that the
partially invariant solution can help improve risk-fairness trade-offs in
different distribution shift settings. The organization of this paper is as
follows: we describe the related literature in section 2, and in section 3, we
explain the details pertaining to the proposed partial invariance framework .
In section 4, we conduct experiments to evaluate the efficacy of our
framework, and we provide some the subsequent discussion in section 5.
Finally, we highlight possible applications and future work in section 6.

##  2 Related Work

Many recent approaches aim to learn deep invariant feature representations:
some focus on domain adaptation by finding a representation whose distribution
is invariant across source and target distributions [7, 31], while others
focus on conditional domain-invariance [13, 20]. However, there is evidence
that domain adaption approaches are insufficient when the test distribution is
unknown and may lie outside the convex hull of training distributions [19, 11,
25].

The focus of this work is the Invariant Risk Minimization (IRM) framework,
proposed in [4]. The framework relates to domain generalization where access
to the test distribution is not assumed. It is rooted in the theory of
causality [28] and links to invariance, which is useful for generalization
[26, 15]. In [2], the authors reformulate IRM via a game-theoretic approach,
wherein the invariant representation corresponds to the Nash equilibrium of a
game. While the IRM framework assumes only the invariance of the conditional
expectation of the label given the representation, some follow-up works rely
on stronger invariance assumptions [30, 22].

Several follow-up works comment on the performance of the invariant risk
framework, and compare against standard Empirical Risk Minimization (ERM). It
has been noted that carefully tuned ERM outperforms state-of-the-art domain
generalization approaches, including IRM, across multiple benchmarks [14]. The
failure of IRM may stem from the gap between the proposed framework and its
practical “linear” version (IRMv1), which fails to capture natural invariances
[16]. The authors of [27] demonstrate that a near-optimal solution to the
IRMv1 objective, which nearly matches IRM on training environments, does no
better than ERM on environments that differ significantly from the training
distribution. The authors of [3] argue that IRM has an advantage over ERM only
when the support of the different environment distributions have a significant
“overlap”. However, unlike previous work which aim to minimize OOD risk, our
primary focus is to understand the trade-off between risk minimization and
fairness across environments in the setting when such a significant “overlap”
may not exist. Consequently, we analyze how IRM behaves in these settings and
explore alternatives via the notion of partial invariance. The notion of
fairness considered in this paper relates to _group sufficiency_ [21], which
is intricately linked with the IRM objective. In fact, several works build
upon the IRM framework by explicitly minimizing a quantity resembling this
notion of fairness for achieving invariance [17, 10].

##  3 The Partial Invariance Framework

In this section, we describe a simple experiment using Gaussian mixture models
(GMMs) to clearly highlight some of the pitfalls of IRM that are encountered
in more complicated practical settings. For completeness, we first state the
IRM formulation.

#### IRM Framework

Consider a set of training environments $e\in\mathcal{E}_{tr}$ and datasets
$\\{X^{e}_{i},Y^{e}_{i}\\}_{i=1}^{n_{e}}$ sampled in each environment from
distribution $P_{e}(X_{e},Y_{e})$, with number of samples $n_{e}$, with
$X^{e}_{i}\in\mathcal{X}$, $Y^{e}_{i}\in\mathcal{Y}$ and consider a predictor
function $f:\mathcal{X}\rightarrow\mathcal{Y}$. The IRM framework aims to find
a predictor $f$ which generalizes well for unseen environments. The goal is to
find a data representation $\Phi:\mathcal{X}\rightarrow\mathcal{Z}$ such that
the predictor $w:\mathcal{Z}\rightarrow\mathcal{Y}$ on top of $\Phi$ is
invariant across all environments $\mathcal{E}_{tr}$.111Formally, for all
$z\in\mathcal{Z}$ that lies in supports of both
$\Phi(X^{e}),\Phi(X^{e^{\prime}})$, we want that
$E_{P_{e}}[Y^{e}|\Phi(X^{e})=z]=E_{P_{e^{\prime}}}[Y^{e^{\prime}}|\Phi(X^{e^{\prime}})=z],\
\forall\ (e,e^{\prime})\in\mathcal{E}_{tr}$. Such an invariant $\Phi$ allows
for the existence of a predictor $w$ which is simultaneously optimal for all
environments, and this $w$ is the one that best approximates the conditional
expectation of the label. To achieve this, the training objective for IRM can
be summarized as:

$\displaystyle\min_{\Phi,w}\sum_{e\in\mathcal{E}_{tr}}\mathcal{R}_{e}(w\circ\Phi)$
(IRM) s.t. $\displaystyle\
w\in\text{arg}\min_{\tilde{w}}{R}_{e}(\tilde{w}\circ\Phi)\ \forall\
\tilde{w}\mbox{,}$

where $R_{e}(f)=E_{X_{e},Y_{e}}[l(f(X_{e}),Y_{e})]$ represents the risk in
environment $e$, pertaining to a loss function $l$ like the mean-squared or
the cross-entropy loss. To simplify the bi-level optimization, the authors of
[4] propose the following simplification (IRMv1):

$\min_{\Phi,w}\sum_{e\in\mathcal{E}_{tr}}\mathcal{R}_{e}(w\circ\Phi)+\lambda||\nabla_{w}\mathcal{R}_{e}(w\circ\Phi)||^{2}$

with regularizer $\lambda$. In the subsequent discussion, we will describe our
experimental setup. As mentioned previously, the pitfalls of IRM, including
the shortcomings of (IRMv1) in finding predictors with better OOD risks have
been studied before [16, 27]. We take an alternate approach where we highlight
that even if IRM (IRMv1) works as intended, the underlying data-generating
distribution may not allow finding a ‘good’ solution, from both classification
risk and fairness considerations. Simultaneously, we will demonstrate how the
partial invariance approach can help in this regard.

#### Problem Setup

Consider a binary classification task on univariate Gaussian mixtures on an
interval of $\mathbb{R}$. Each GMM is parametrized as
$g(\mu_{1},\mu_{2},\sigma)$, i.e. two Gaussians with class means
$(\mu_{1},\mu_{2})$ and equal variance $\sigma$. In our setup, we consider $N$
environments wherein each environment in $\\{e_{i}\\}_{i=0}^{N-1}$ is
characterized by exactly one Gaussian mixture as
$g(\mu_{1}+i\delta,\mu_{2}+i\delta,\sigma)$. Therefore, $e_{0}$ corresponds to
$g(\mu_{1},\mu_{2},\sigma)$ and $e_{N}$ corresponds to
$g(\mu_{1}+(N-1)\delta,\mu_{2}+(N-1)\delta,\sigma)$, see Fig. 1 for an
illustration. Then $\delta$, along with $\sigma$, controls how far apart the
individual environments are, i.e. the cross-environmental overlap. For each
environment, we consider a threshold-based classifier. The loss is the (0-1)
classification loss, which yields a risk as a continuous function of the
threshold $t$. Intuitively, for $N=2$, if $\delta$ is very large, any
predictor that does well on the first environment will incur the a very high
classification risk in the second environment.

/

Figure 1: Illustration of the Gaussian mixture distributions for $N=2$,
highlighting the role of cross-environmental overlap parameter $\delta$. Here,
$(\mu_{1},\mu 2)=(1,2)$, $\sigma=1$, and $\delta=0.3$.

#### Motivation

One would argue that such a problem setting is naturally unsuitable for
studying the nature of the IRM solution. Indeed with a single IRM predictor,
it follows that with $\Phi(x)=x$, any predictor (threshold) cannot be
simultaneously optimal for all $N$ environments, unless $|t|\rightarrow\infty$
and thus violates the assumption on the existence of a ‘meaningful’ IRM
solution. However, we argue such situations can naturally arise in standard
high-dimensional learning tasks, where existence of an IRM solution that is
both fair and meaningful may not be guaranteed a priori. Hence the learned
$\Phi$, either due to the nature of the data distribution or lack of
representation power or simply a failure of the learning algorithm, may yield
conditional distributions across environments that show little overlap.
Therefore, it is important to analyze and develop countermeasures against this
failure case.

#### Partial Invariance

It is clear from the previous discussion that in certain settings, IRM may not
work as intended. Looking at the GMM problem in particular, the issue seems to
arise from forcing the invariant predictor to be optimal in non-overlapping
environments. One intuitive proposal is to partition the environments such
that the environments within a particular partition are sufficiently “close”.
Consequently, we can now enforce invariance within each partition. And that is
exactly the idea of partial invariance. For now, we stick to the GMM settings
as described above and we enforce partial invariance via optimizing the IRMv1
objective, within each partition. We roughly characterize the Partial-IRM
approach below.

Algorithm 1 Partial-IRM algorithm (rough)

partitions $\\{p_{i}\\}_{i=1}^{k}$, $p_{i}\subseteq\mathcal{E}_{tr}$,
$\sum_{i}|p_{i}|=|\mathcal{E}_{tr}|$, $p_{i}\cap p_{j}=\phi\ \forall\ \ (i,j)$

for $\\{p_{i}\\}_{i=1}^{k}$ do

$w_{i},\Phi_{i}\leftarrow\text{IRM}(p_{i})$ $\triangleright$ Impose IRM
objective over $e$ $\in$ $p_{i}$, obtain $(w_{i},\Phi_{i})$ for partition
$p_{i}$

end for

It follows from the prior discussion that partial invariance will be most
effective when the environments within each partition are sufficiently close
in terms of the distribution, and therefore may be inferred from the
environments by minimizing some suitable divergence criterion. For our
Gaussian mixtures however, fixing the size of each partition makes the
environment assignment intuitive i.e. we can simply group together consecutive
environments. In later sections, we will discuss how this formulation could
fit in the context of causality and invariant features.

#### Performance Metrics

As mentioned previously, our focus is not on IRM and learning invariant
features for improving OOD performance. In fact, within our GMM setup, we
assume access to all environments of interest (i.e.
$\\{e_{i}\\}_{i=0}^{(N-1)}$). Consequently, our performance metrics depend
only on the different $e_{i}$’s. Within this framework, we analyze and compare
the IRM and Partial-IRM solutions on a) classification risk and b) fairness.
Given a threshold $t$, we assume a uniform measure on the set of environments
and thus, the classification risk may be computed as follows:

$R_{g}(t)=\frac{1}{N}\sum_{i=0}^{(N-1)}R_{e_{i}}(t)\mbox{.}$ (1)

Here $R_{e_{i}}(t)$ denotes the classification risk due to the choice of
threshold $t$ for environment $e_{i}$ (i.e. the risk on the GMM parametrized
as g($\mu_{1}+i.\delta,\mu_{2}+i.\delta,\sigma$)). Next, we estimate fairness
based on the notion of "Group-Sufficiency" –satisfied when
$E[y|S(x),e]=E[y|S(x),e_{0}]$ for all $e,e_{0}$, where $e$ denotes membership
of some sensitive class and $S(x)$ represents a score metric. It roughly
equates to equalizing the risks across environments. For a choice of threshold
$t$, we measure fairness using the V-REx penalty[17], which roughly estimates
the variance of the risks across environments (lower variance implies higher
fairness):

$F_{V}(t)=\frac{1}{2N^{2}}\sum_{i=0}^{(N-1)}\sum_{j=0}^{(N-1)}(R_{e_{i}}(t)-R_{e_{j}(t)})^{2}\mbox{.}$
(2)

##  4 Experiments and Results

Our goal is to test the effect of creating partitions on the performance of
invariant predictors, characterized by both risk and fairness. Given 24
environments, we test with different number of partitions
$p\in\\{1,2,3,4,6,12,24\\}$, with equal sized partition for simplicity. The
IRM case corresponds to $p=1$ with a single classifier for all environments,
whereas the ERM case is $p=24$ with a classifier for each environment. Based
on the Gaussian mixture model, with $\mu_{1}=1$ and $\mu_{2}=2$, we experiment
with four combinations of intra-environment separation $\sigma$ (0.1 or 1),
and cross-environments overlap $\delta$ (0.1 or 1), noting that the highest
overlap between environments is achieved with large $\sigma$ ($\sigma$ = 1)
and a small $\delta$ ($\delta$ = 0.1). We set the regularization $\lambda$ to
be 10, which for our problem suffices for IRMv1 to mirror the true IRM
objective. We learn threshold classifiers, by minimizing the IRMv1 objective
for the environments within that partition. In the case of partition size
being equal to 1, we consider the ERM objective by setting $\lambda=0$.

#### High overlap (small $\delta)$

A small value of $\delta$ means that environments show a stronger degree of
overlap. In this setting, we observe that the IRM solution prioritizes
fairness over risk. Herein, partial invariance allows us to trade-off
performance for reduced fairness. However, in certain settings wherein the
classes in the environment are not well separated (large $\sigma$), one can
improve significantly improve the performance by paying a much lower relative
cost in terms of fairness (Fig. 2(b)). This also demonstrates that P-IRM
introduces flexibility in terms of balancing risk and fairness.

#### Low Overlap (large $\delta)$

Next, we consider a larger value of $\delta$, meaning a low environmental
overlap. In both cases, we observe that the degree of change in the
risk/fairness values is lower than the previous case. However, we point out
the curious case of well separated environments with low overlap, wherein
P-IRM can seemingly improve both risk and fairness simultaneously (Fig. 2(c)).
Since we expect such settings to arise in real world problems, our results
offer us some intuition towards how P-IRM can improve performance in such
settings.

(a) Large separation, high overlap

(b) Minimal separation, high overlap

(c) Large separation, low overlap

(d) Minimal separation, low overlap

Figure 2: Experiments with Gaussian mixtures for different $\delta,\sigma$
values. Lower values of risk and fairness variance indicate better performance
and fairness respectively.

###  4.1 The Fairness-Robustness Trade-off

The previous experiments give us insights into the situations in which partial
invariance can be beneficial in controlling the fairness-risk trade-off.
Specifically, in settings with a smaller $\delta$ and consequently more
overlap, the fairness-risk values change significantly as one increases the
number of partitions. However, one could argue that such a trade-off could be
achieved by simply varying the $\lambda$ for the IRM solution. But we claim
that introducing the notion of partial invariance allows us easier access to a
richer solution set on this trade-off, while still retaining desirable
properties.

To show this, we consider the high overlap setting in which the individual
Gaussians within an environment are easily separable, i.e.
$\delta=0.1=\sigma$. The rest of the setup is exactly the same as the previous
experiment. But here, we fix our data generating distribution and instead vary
$\lambda$ on the log scale. Then, for each value of $\lambda$ and given
partition size, we measure the corresponding risk and fairness values for the
obtained optimal threshold. The obtained plots are presented in Fig.3.

(a) Risk of different partitions

(b) Fairness of different partitions

Figure 3: Effects of adjusting the regularizing parameter $\lambda$ on risk
and fairness of IRM, P-IRM, and ERM solutions.

We notice that in the regions corresponding to sufficiently high values of
$\lambda$, the Partial-IRM solutions allow us to operate on different curves
corresponding to risk-fairness trade-offs, remarkable similar to receiver
operating characteristic (ROC) curves in classical decision theory.
Additionally, we can switch between different levels of risk/fairness by
switching between the number of partitions. Note our predefined constraints on
equal-sized partitions restricts the choice of partition size to integer
factors of $N$. We speculate that we might achieve even more flexibility in
certain other settings, such as when both the environments and partitions
being defined as continuous intervals. And the key observation here is that
while IRM may be able to find similar trade-offs by reducing $\lambda$ (the
sharp slope region for the blue curve, Fig. 3(a)), it does so by compromising
on the basic essence of invariance. In the discussion, we will elaborate on
how this might not be desirable in certain settings.

##  5 Discussion

The experiments in the previous section offered us some interesting insights.
We saw that when the cross-environments overlap is low, Partial-IRM offers
only marginal benefits over standard IRM and ERM. This is expected because a
lower similarity between environments implies that there is less common
information that may be exploited. Additionally, there may exist some
interesting cases wherein we can simultaneously improve upon on both fairness
and classification risk, by adopting the partial invariance approach. This is
somewhat surprising and certainly merits additional investigation. Next, we
considered the high overlap settings. Here, the role of partial invariance in
navigating the risk-fairness trade-off comes to fore. We elaborate upon this
idea in the subsequent experiment, wherein we observe that by changing the
size of the partition, we can switch between different levels of the risk-
fairness trade-off. To achieve the same flexibility within the IRM framework,
we need to operate within regions where the IRMv1 solution shows large
deviations from the original IRM formulation (by reducing $\lambda$). We
believe this can have additional detrimental effects in situations where we
would implicitly need invariance. We demonstrate one such setting below.

###  5.1 The Case for Partial IRM

In section 4, we observed that the notion of partial invariance allowed us to
explore the inherent trade-off between risk minimization and fairness and in
certain cases, simultaneously improve upon both risk and fairness as compared
to the IRM solution. However, it is not yet clear as to how partial invariance
fits in the “causal” picture of invariant risk minimization. We attempt to
shed some light on this by slightly modifying the motivating regression task
in the original IRM formulation [4].

Consider the following data generation model, where
$\sigma(e)\in[0,\sigma^{2}_{max}]$. The target $Y$ is to be regressed based on
observed values of $X=(X_{1},X_{2},X_{3})$. We consider two independent
environment variables $(E,e)$ at play, where $e$ affects the individual
distributions of the set of observables $X$, whereas $E$ controls the marginal
distribution of $Y$ given $X_{2}$ via the random variable $P$:

$\begin{gathered}X_{1}\leftarrow N(0,\sigma(e)^{2}),X_{2}\leftarrow
N(0,\sigma(e)^{2})\\\ E\leftarrow\textbf{U}(\\{1,2,3\\}),P(E)=f(E)\\\
Y\leftarrow X_{1}+P(E).X_{2}+N(0,\sigma(e)^{2})\\\ X_{3}\leftarrow
Y+N(0,1)\end{gathered}$

We consider the prediction model to estimate $Y$ as
$\hat{Y}=\alpha_{1}X_{1}+\alpha_{2}X_{2}+\alpha_{3}X_{3}$. Then. it follows
that if $\alpha_{3}>0$, the regression coefficients will explicitly depend on
$\sigma(e)$ and the OOD risk of the model can grow without bound, as noted in
[4]. Within the IRM framework, the only feasible representation $\Phi$ that
yields invariant predictors across $(E,e)$ is $\Phi(X)=X_{1}$ and the
corresponding regression coefficients are $(1,0,0)$ respectively. While this
solution works and generalizes for OOD tasks w.r.t. $e$, we lose in terms of
performance by discarding $X_{2}$. In addition, deviating from the IRM
solution via IRMv1 (as in the previous section) is not feasible since we
require invariance w.r.t $e$ for a finite risk in worst-case environments.

Now note that the notion of partial invariance offers an intuitive solution;
We could first create partitions across $E$ and then apply IRM to obtain
invariance w.r.t e, potentially allowing us to get the best of both worlds.
However, as in the case with GMMs, the idea implicitly assumes that the P-IRM
framework can indeed partition across $E$. In general, it is quite hard to
even achieve this distinction between causal and non-causal features, without
prior information on the data-generating distributions. Fortunately, it seems
that it is possible to surmise this information under certain settings that
arise naturally and a study of precisely these settings will be the focus of
our future work.

##  6 Future Work/Applications

Distribution shifts naturally arise in real-world settings, particularly when
environments with low overlap interact so that even though sub-groups share
some common features, they can differ substantially in some other locale-
specific features [23]. In such settings, the partial invariance framework can
better tackle the OOD problem. Similar notions can also be found in recent
literature pertaining to equivariant neural networks, in relaxing global
symmetry requirements to local gauge symmetries[9].

Next, we describe applications where partial invariance could prove to be
useful. In online forums, identity-related discussions in which hatred or
humiliation is expressed towards certain out-groups may be considered toxic,
but similar discussions in a completely in-group discussion may not [12].
Also, in forums that usually welcome diversity, explicitly mentioning
identities would be acceptable and not toxic, unlike in forums that routinely
make discriminatory remarks against minorities. Existing systems to detect
toxicity in online forums [5, 32] do not, however, consider the environmental
differences between forums and tend to perform poorly in more diverse settings
[8]. As such, a model that can distinguish the type/purpose of a forum and
then infer toxic mentions from this additional context, can be expected
improve performance.

We note that in the current form however, it is unclear how partial invariance
can be applied to real data wherein firstly, we need to decide how to
partition environments and secondly, we need to infer the partition membership
during test time, so as to use the corresponding classifier for optimal
prediction for that sample. Therefore, in future work, we want to develop
methods to automatically infer partitions from environments, by introducing
appropriate notions of distance between environments.

## References

  * [1] David Madras Adragna Robert, Elliot Creager and Richard Zemel.  Fairness and robustness in invariant learning: A case study in toxicity classification.  arXiv:2011.06485 [CS.LG], December 2020.
  * [2] Kartik Ahuja, Karthikeyan Shanmugam, Kush R. Varshney, and Amit Dhurandhar.  Invariant Risk Minimization Games.  In Proceedings of the 37th International Conference on Machine Learning (ICML’20), volume 119, pages 145–155. PMLR, July 2020.
  * [3] Kartik Ahuja, Jun Wang, Amit Dhurandhar, Karthikeyan Shanmugam, and Kush R. Varshney.  Empirical or Invariant Risk Minimization? A Sample Complexity Perspective.  In Proceeding of the 8th International Conference on Learning Representations (ICLR’20), May 2020.
  * [4] Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz.  Invariant Risk Minimization.  arXiv:1907.02893 [stat.ML], 2019.
  * [5] Pinkesh Badjatiya, Shashank Gupta, Manish Gupta, and Vasudeva Varma.  Deep learning for hate speech detection in tweets.  In Proceedings of the 26th International Conference on World Wide Web Companion, pages 759–760, April 2017.
  * [6] Sara Beery, Grant V. Horn, and Pietro Perona.  Recognition in Terra Incognita.  In Proceedings of the European Conference on Computer Vision (ECCV), pages 456–473, September 2018.
  * [7] Shai Ben-David, John Blitzer, Koby Crammer, Alex Kulesza, Fernando Pereira, and Jennifer Vaughan.  A theory of learning from different domains.  Machine Learning, 79:151–175, 2010.
  * [8] Daniel Borkan, Lucas Dixon, Jeffrey Sorensen, Nithum Thain, and Lucy Vasserman.  Nuanced metrics for measuring unintended bias with real data for text classification.  In Companion Proceedings of the 2019 World Wide Web Conference, pages 491–500, May 2019.
  * [9] Taco Cohen, Maurice Weiler, Berkay Kicanaoglu, and Max Welling.  Gauge equivariant convolutional networks and the icosahedral cnn.  In Proceedings of the 36th International Conference on Machine Learning (ICML’19), pages 1321–1330. PMLR, June 2019.
  * [10] Elliot Creager, Jörn-Henrik Jacobsen, and Richard S. Zemel.  Exchanging lessons between algorithmic fairness and domain generalization.  2020\.
  * [11] John Duchi, Peter Glynn, and Hongseok Namkoong.  Statistics of Robust Optimization: A Generalized Empirical Likelihood Approach.  Mathematics of Operations Research, 46(3), January 2021.
  * [12] John D. Gallacher.  Leveraging cross-platform data to improve automated hate speech detection.  arXiv:2102.04895 [CS.CL], 2021.
  * [13] Mingming Gong, Kun Zhang, Tongliang Liu, Dacheng Tao, Clark Glymour, and Bernhard Schölkopf.  Domain Adaptation with Conditional Transferable Components.  In Proceedings of the 33rd International Conference on Machine Learning (ICML’16), volume 48, pages 2839–2848, June 2016.
  * [14] Ishaan Gulrajani and David Lopez-Paz.  In Search of Lost Domain Generalization.  In Proceeding of the 8th International Conference on Learning Representations (ICLR’20), May 2020.
  * [15] Christina Heinze-Deml, Jonas Peters, and Nicolai Meinshausen.  Invariant Causal Prediction for Nonlinear Models.  Journal of Causal Inference, 6(2), September 2018.
  * [16] Danica Sutherland Kamath Pritish, Akilesh Tangella and Nathan Srebro.  Does Invariant Risk Minimization Capture Invariance?  In Proceedigns of the International Conference on Artificial Intelligence and Statistics, pages 4069–4077. PMLR, April 2021.
  * [17] David Krueger, Ethan Caballero, Joern-Henrik Jacobsen, Amy Zhang, Jonathan Binas, Dinghuai Zhang, Remi Le Priol, and Aaron Courville.  Out-of-Distribution Generalization via Risk Extrapolation.  In Proceedings of the 38th International Conference on Machine Learning (ICML’21), pages 5815–5826. PMLR, 2021.
  * [18] Brenden M. Lake, Tomer D. Ullman, Joshua B. Tenenbaum, and Samuel J. Gershman.  Building machines that learn and think like people.  Behavioral and Brain Sciences, 40:e253, 2017.
  * [19] Jaeho Lee and Maxim Raginsky.  Minimax Statistical Learning with Wasserstein Distances.  In Proceedings of the 32nd International Conference on Neural Information Processing Systems (NIPS’18), pages 2692–2701, December 2018.
  * [20] Ya Li, Mingming Gong, Xinmei Tian, Tongliang Liu, and Dacheng Tao.  Domain Generalization via Conditional Invariant Representation.  In Proceedings of the 32nd Association for the Advancement of Artificial Intelligence (AAAI’18), volume 31, July 2018.
  * [21] Lydia T. Liu, Max Simchowitz, and Moritz Hardt.  The implicit fairness criterion of unconstrained learning.  In Proceedings of the 36th International Conference on Machine Learning (ICML’19), volume 97, pages 4051–4060. PMLR, June 2019.
  * [22] Divyat Mahajan, Shruti Tople, and Amit Sharma.  Domain Generalization using Causal Matching.  In Proceedings of the 38th International Conference on Machine Learning (ICML’21), volume 139, pages 7313–7324. PMLR, June 2021.
  * [23] Ankur Mani, Lav R. Varshney, and Alex Pentland.  Quantization Games on Social Networks and Language Evolution.  IEEE Transactions on Signal Processing, 69:3922–3934, 2021.
  * [24] Gary Marcus.  Deep Learning: A Critical Appraisal.  arXiv:1801.00631 [CS.AI], 2018.
  * [25] Mehryar Mohri, Gary Sivek, and Ananda Theertha Suresh.  Agnostic Federated Learning.  In Proceedings of the 36th International Conference on Machine Learning (ICML’19), volume 97, pages 4615–4625. PMLR, June 2019.
  * [26] Jonas Peters, Peter Bühlmann, and Nicolai Meinshausen.  Causal inference by using invariant prediction: identification and confidence intervals.  Journal of the Royal Statistical Society. Series B (Statistical Methodology), 78(5):947–1012, 2016.
  * [27] Elan Rosenfeld, Pradeep Kumar Ravikumar, and Andrej Risteski.  The Risks of Invariant Risk Minimization.  In Proceeding of the 8th International Conference on Learning Representations (ICLR’20), May 2020.
  * [28] B. Schölkopf, D. Janzing, J. Peters, E. Sgouritsa, K. Zhang, and J. Mooij.  On causal and Anticausal Learning.  In Proceedings of the 29th International Conference on Machine Learning (ICML’12), pages 1255–1262, June 2012.
  * [29] Vladimir Vapnik.  The Nature of Statistical Learning Theory.  Springer Science and Business Media, 2013.
  * [30] Chuanlong Xie, Haotian Ye, Fei Chen, Yue Liu, Rui Sun, and Zhenguo Li.  Risk Variance Penalization.  arXiv:2006.07544 [cs.LG], April 2021.
  * [31] Kun Zhang, Mingming Gong, and Bernhard Schoelkopf.  Multi-Source Domain Adaptation: A Causal View.  In Proceedings of the 29th Association for the Advancement of Artificial Intelligence Conference (AAAI’15), January 2015.
  * [32] Ziqi Zhang and Lei Luo.  Hate speech detection: A solved problem? the challenging case of long tail on twitter.  Semantic Web, 10(5):925–945, 2019.


authorpersonFelipe Bañados Schwerter,
personRonald Garcia, and personÉric Tanter.
A Theory of Gradual Effect Systems. In
booktitleProceedings of the 19th ACM SIGPLAN International
Conference on Functional Programming (seriesICFP '14).

authorpersonNick Benton.
Embedded Interpreters.
journalJournal of Functional Programming
volume15, number04 (year2005),

authorpersonJohn Tang Boyland.
The Problem of Structural Type Tests in a
Gradual-Typed Language. In booktitleWorkshop on Foundations
of Object-Oriented Languages (FOOL), informal proceedings.

[Cimini and SiekCimini and Siek2016]
authorpersonMatteo Cimini and
personJeremy G. Siek. year2016.
The Gradualizer: A Methodology and Algorithm for
Generating Gradual Type Systems. In booktitleProceedings of
the 43rd Annual ACM SIGPLAN-SIGACT Symposium on Principles of Programming
Languages (seriesPOPL '16).

[Cimini and SiekCimini and Siek2017]
authorpersonMatteo Cimini and
personJeremy G. Siek. year2017.
Automatically Generating the Dynamic Semantics of
Gradually Typed Languages. In booktitleProceedings of the
44th ACM SIGPLAN Symposium on Principles of Programming Languages
(seriesPOPL 2017). pages789–803.

[Dagand, Tabareau, and TanterDagand
et al2016]
authorpersonPierre-Evariste Dagand,
personNicolas Tabareau, and personÉric Tanter.
Partial Type Equivalences for Verified Dependent
Interoperability (seriesICFP 2016).

[Dagand, Tabareau, and TanterDagand
et al2018]
authorpersonPierr-Évariste Dagand,
personNicolas Tabareau, and personÉric Tanter.
Foundations of dependent interoperability.
journalJournal of Functional Programming
volume28 (year2018), pagese9.

[(Favonia), Benton, and Harper(Favonia)
et al2017]
authorpersonKeun-Bang Hou (Favonia),
personNick Benton, and personRobert Harper.
Correctness of compiling polymorphism to dynamic
journalJournal of Functional Programming
volume27 (year2017).

[Felleisen and HiebFelleisen and
authorpersonMatthias Felleisen and
personRobert Hieb. year1992.
A Revised Report on the Syntactic Theories of
Sequential Control and State.
journalTheor. Comput. Sci. volume103,
number2 (year1992), pages235–271.

[Findler and BlumeFindler and Blume2006]
authorpersonRobby Findler and
personMatthias Blume. year2006.
Contracts as Pairs of Projections. In
booktitleInternational Symposium on Functional and Logic
Programming (FLOPS).

[Findler and FelleisenFindler and
authorpersonRobert Bruce Findler and
personMatthias Felleisen. year2002.
Contracts for higher-order functions. In
booktitleInternational Conference on Functional
Programming (ICFP), Pittsburgh, Pennsylvania. pages48–59.

authorpersonCarsten Führmann.
Direct models of the computational
journalElectronic Notes in Theoretical Computer
Science  volume20 (year1999),

[Garcia, Clark, and TanterGarcia
et al2016a]
authorpersonRonald Garcia, personAlison M.
Clark, and personEric Tanter.
Abstracting Gradual Typing. In
booktitleACM Symposium on Principles of Programming
Languages (POPL), St. Petersburg, Florida.

[Garcia, Clark, and TanterGarcia
et al2016b]
authorpersonRonald Garcia, personAlison M.
Clark, and personÉric Tanter.
Abstracting Gradual Typing. In
booktitleProceedings of the 43rd Annual ACM SIGPLAN-SIGACT
Symposium on Principles of Programming Languages
(seriesPOPL '16).

[Gronski, Knowles, Tomb, Freund, and
FlanaganGronski et al2006]
authorpersonJessica Gronski, personKenneth
Knowles, personAaron Tomb, personStephen N. Freund,
and personCormac Flanagan. year2006.
Sage: Hybrid Checking for Flexible Specifications.
In booktitleScheme and Functional Programming Workshop
(Scheme). pages93–104.

authorpersonFritz Henglein.
Dynamic Typing: Syntax and Proof Theory.
journalScience of Computer Programming
volume22, number3 (year1994),

[Herman, Tomb, and FlanaganHerman
et al2010]
authorpersonDavid Herman, personAaron Tomb,
and personCormac Flanagan. year2010.
Space-efficient Gradual Typing.
journalHigher Order Symbol. Comput.
volume23, number2 (dateJune

[Igarashi, Thiemann, Vasconcelos, and
WadlerIgarashi et al2017b]
authorpersonAtsushi Igarashi, personPeter
Thiemann, personVasco Vasconcelos, and personPhilip
Wadler. year2017b.
Gradual Session Types. In
booktitleInternational Conference on Functional
Programming (ICFP), Oxford, United Kingdom.

[Igarashi, Sekiyama, and IgarashiIgarashi
et al2017a]
authorpersonYuu Igarashi, personTaro
Sekiyama, and personAtsushi Igarashi.
On Polymorphic Gradual Typing. In
booktitleInternational Conference on Functional
Programming (ICFP), Oxford, United Kingdom.

[Ina and IgarashiIna and Igarashi2011]
authorpersonLintaro Ina and personAtsushi
Igarashi. year2011.
Gradual typing for generics. In
booktitleProceedings of the 2011 ACM international
conference on Object oriented programming systems languages and
applications (seriesOOPSLA '11).

[Lehmann and TanterLehmann and
authorpersonNico Lehmann and
personÉric Tanter. year2017.
Gradual Refinement Types. In
booktitleProceedings of the 44th ACM SIGPLAN Symposium on
Principles of Programming Languages (seriesPOPL 2017).

[New and LicataNew and Licata2018]
authorpersonMax S. New and personDaniel R.
Licata. year2018.
Call-by-Name Gradual Type Theory. In
booktitle3rd International Conference on Formal Structures
for Computation and Deduction (FSCD 2018) (seriesLeibniz
International Proceedings in Informatics (LIPIcs)),
editorpersonHélène Kirchner (Ed.),
Vol. volume108.

authorpersonDana Scott.
Continuous lattices.
In booktitleToposes, algebraic geometry and logic.

[Siek, Vitousek, Cimini, and BoylandSiek
et al2015]
authorpersonJeremy Siek, personMicahel
Vitousek, personMatteo Cimini, and personJohn Tang
Boyland. year2015.
Refined Criteria for Gradual Typing. In
booktitle1st Summit on Advances in Programming Languages
(seriesSNAPL 2015).

[Siek and TahaSiek and Taha2006]
authorpersonJeremy G. Siek and personWalid
Taha. year2006.
Gradual Typing for Functional Languages. In
booktitleScheme and Functional Programming Workshop
(Scheme). pages81–92.

[Siek and WadlerSiek and Wadler2010]
authorpersonJeremy G. Siek and
personPhilip Wadler. year2010.
Threesomes, with and without blame. In
booktitleACM Symposium on Principles of Programming
Languages (POPL), Madrid, Spain. pages365–376.

[Swamy, Fournet, Rastogi, Bhargavan, Chen,
Strub, and BiermanSwamy et al2014]
authorpersonNikhil Swamy, personCédric
Fournet, personAseem Rastogi, personKarthikeyan
Bhargavan, personJuan Chen, personPierre-Yves Strub,
and personGavin M. Bierman. year2014.
Gradual typing embedded securely in
JavaScript. In booktitleACM Symposium on
Principles of Programming Languages (POPL), San Diego, California.

authorpersonSatish Thatte.
Quasi-static typing. In
booktitleACM Symposium on Principles of Programming
Languages (POPL). pages367–381.

[Tobin-Hochstadt and FelleisenTobin-Hochstadt
and Felleisen2006]
authorpersonSam Tobin-Hochstadt and
personMatthias Felleisen. year2006.
Interlanguage Migration: From Scripts to Programs.
In booktitleDynamic Languages Symposium (DLS).

[Tobin-Hochstadt and FelleisenTobin-Hochstadt
and Felleisen2008]
authorpersonSam Tobin-Hochstadt and
personMatthias Felleisen. year2008.
The Design and Implementation of Typed Scheme. In
booktitleACM Symposium on Principles of Programming
Languages (POPL), San Francisco, California.

[Wadler and FindlerWadler and
authorpersonPhilip Wadler and
personRobert Bruce Findler. year2009.
Well-typed programs can't be blamed. In
booktitleEuropean Symposium on Programming (ESOP).

[Wolff, Garcia, Tanter, and AldrichWolff
et al2011]
authorpersonRoger Wolff, personRonald
Garcia, personÉric Tanter, and personJonathan
Aldrich. year2011.
Gradual Typestate. In
booktitleProceedings of the 25th European Conference on
Object-oriented Programming (seriesECOOP'11).


# Latent Conservative Objective Models for
Data-Driven Crystal Structure Prediction

Han Qi1,∗, Xinyang Geng1,∗, Stefano Rando2, Iku Ohama2, Aviral Kumar1, Sergey
Levine1
1UC Berkeley, 2Panasonic North America (∗ Equal contribution)
<EMAIL_ADDRESS><EMAIL_ADDRESS><EMAIL_ADDRESS>
###### Abstract

In computational chemistry, crystal structure prediction (CSP) is an
optimization problem that involves discovering the lowest energy stable
crystal structure for a given chemical formula. This problem is challenging as
it requires discovering globally optimal designs with the lowest energies on
complex manifolds. One approach to tackle this problem involves building
simulators based on density functional theory (DFT) followed by running search
in simulation, but these simulators are painfully slow. In this paper, we
study present and study an alternate, data-driven approach to crystal
structure prediction: instead of directly searching for the most stable
structures in simulation, we train a surrogate model of the crystal formation
energy from a database of existing crystal structures, and then optimize this
model with respect to the parameters of the crystal structure. This surrogate
model is trained to be conservative so as to prevent exploitation of its
errors by the optimizer. To handle optimization in the non-Euclidean space of
crystal structures, we first utilize a state-of-the-art graph diffusion auto-
encoder (CD-VAE) to convert a crystal structure into a vector-based search
space and then optimize a conservative surrogate model of the crystal energy,
trained on top of this vector representation. We show that our approach,
dubbed LCOMs (latent conservative objective models), performs comparably to
the best current approaches in terms of success rate of structure prediction,
while also drastically reducing computational cost.

###  1 Introduction

Data-driven optimization problems arise in many areas of science and
engineering. In these settings, we have an unknown function that we would like
to optimize with respect to its inputs, provided only with a dataset of input-
output pairs. Examples include drug design, where inputs might be molecules
and outputs are the efficacy of a drug, protein design, where inputs
correspond to protein sequences and outputs are some metric such as
fluorescence (Sarkisyan et al., 2016) or, as in our experiments, prediction of
crystal structures, where inputs consist of crystal structures and outputs
correspond to their formation energy. Such data driven optimization problems
present several challenges. First, naïvely training a predictive model to
predict the output from the input and then optimizing against such a model may
lead to exploitation: a sufficiently strong optimizer can typically discover
inputs that lead any learned model to extrapolate erroneously, and then
exploit these errors to find inputs that “fool” the model into making the
desired predictions. Second, even if a model can be suitably robustified, many
of the most important design and optimization problems in science and
engineering, including crystal structure prediction, require optimizing over
complex sets and non-Euclidean manifolds, such that naïvely applying gradient-
based methods in the input space is unlikely to result in a meaningful
improvement.

In this paper, we study these challenges in the context of crystal structure
prediction (CSP) Woodley and Catlow (2008). Crystals are a class of solid-
state materials characterized by the periodic placement of atoms. These
structures form the basis of a wide variety of applications such as designing
super-conductors, batteries Yamashita et al. (2016), and solar cells Walsh et
al. (2012). Computationally identifying stable crystal geometries given a
particular chemical formula typically involves minimizing (an estimate of) the
crystal’s formation energy to find the minimal energy structure. Conventional
approaches to this problem rely on slow and compute-intensive DFT simulators
Chermette (1998), but more recent machine learning approaches dispense with
DFT-based simulators and use databases of structures and their corresponding
energies to train models that estimate crystal formation energy directly
Gasteiger et al. (2021), Klicpera et al. (2020). However, the CSP problem
suffers from both issues outlined above: crystal structures typically exhibit
highly complex geometries characterized by periodicity of the lattice that
forms the crystal and discrete (e.g., number and types of atoms in the
chemical compound) and continuous features (e.g., positions of atoms in 3D
space), which make it hard to produce reliable estimates of energies across
the entire manifold of possible structures. Optimizing the structure using
such inaccurate models then bears the risk of the optimization procedure
“exploiting” these inaccuracies, resulting in structures that erroneously
appear promising in this learned model but are not stable.

In this paper, we aim to develop a data-driven optimization approach to
overcome these challenges. First, to avoid the complexities associated with
optimization over the complex manifold of crystal structures that consists of
both discrete and continuous objects, our optimization procedure utilizes a
crystal diffusion variational auto-encoder (CD-VAE) (Xie et al., 2021) to
convert crystal structures into latent representations, which are much more
amenable to simple gradient-based optimization. Second, to prevent the
optimizer from getting “fooled” by the errors in the learned surrogate model,
we extend conservative objective models (Trabucco et al., 2021b), a
robustification procedure, to our surrogate energy prediction model. This
procedure explicitly pushes down over-estimated out-of-distribution designs in
the latent space. Using a combination of these techniques, we develop a method
for finding stable crystal structures that alleviates the time and compute
costs associated with using DFT simulators, while also addressing the
inaccuracies in a purely offline approach for designing crystal structures.

Our main contribution is a data-driven optimization approach, that we call
latent conservative objective models (LCOM), for the problem of crystal
structure prediction for solid materials. Our method leverages both advances
in generative modeling over periodic solid-state materials for latent space
learning (Xie et al., 2021) and recent advances in model-based optimization
for robustifying the learned model to make it amenable to direct optimization
of formation energies (Trabucco et al., 2021b). We summarize the approach in
Figure 1. We instantiate our approach, latent conservative objective models
(LCOMs), using crystal diffusion variational auto-encoders (CD-VAE) Xie et al.
(2021) for learning the latent space and conservative objective models (COMs)
Trabucco et al. (2021b) for optimization. Empirically, we demonstrate that
LCOMs are able to match the performance of the best prior method from Cheng et
al. (2022) while significantly reducing the total wall-clock time needed for
optimization by 40 times. In particular, a single optimization cycle in our
framework takes an average of 2 seconds. This allows our model to provide
predictions for more than 100 compounds in 4 minutes, much faster than prior
works.

Figure 1: Overview of LCOMs. We train a graph-based CD-VAE to construct a
latent space the represents crystal structure, conditioned on the molecular
structure of the compound. Different points in this latent space correspond to
different crystal structures, and we then optimize over the structure with
simple gradient-based optimization methods operating on this latent space. To
do so, we train a surrogate energy prediction model on the learned latent
space via conservative training (Trabucco et al., 2021b) to make it robust on
out-of-distribution inputs, thus preventing the optimizer from discovering
latent space points far from the training data for which the energy
predictions yield erroneously low energies. The optimized latent vector is
then decoded into a structure. Since the entire optimization is performed in
the latent space, the comparatively complex encoder and decoder only need to
be used once during optimization (to encode the initial structure and decode
the final one).

###  2 Background on Crystal Structures and Materials

In this section, we present the background definitions associated with
crystals and solid-state materials. A crystal is a solid-state material
characterized by a periodic placement of its constituents, which are chemical
elements. The stoichiometry or the composition of a crystal, like NaCl,
consists of the elements that make up the solid-state material (i.e., Na and
Cl in this case) and in what ratio. In real-world applications of solid-state
materials it is not enough to develop a material with a suitable chemical
composition, but we must also account for the crystalline periodic structure
of the solid and the atoms’ positions with respect to it to assess the
stability of a given compound.

Mathematically, we can describe the periodic structure of a chemical by
defining its lattice $L$ in 3D space, which repeats periodically. To
characterize a lattice, we define its base vectors
$\textbf{v},\textbf{w},\textbf{z}$. Every point in the lattice is a linear
combination of these vectors using only integer coefficients.

$\displaystyle p\in L\iff\exists
n,m,k\in\mathbb{Z}\,|\,p=n\textbf{v}+m\textbf{w}+k\textbf{z}.$ (1)

Given a lattice $L$, the unit cell is the volume of 3D space contained between
the base vectors, defined formally as follows:

$\displaystyle\left\\{p\in\mathbb{R}^{3}\,|\,\exists
x,y,z\in\left[0,1\right],\,p=x\textbf{v}+y\textbf{w}+z\textbf{z}\right\\}.$
(2)

We can obtain the entire lattice of a crystal by periodically repeating this
unit cell in 3D space. Given a lattice in 3D space, a crystal is additionally
characterized by how many atoms $n$ are in the unit cell. We observe that the
number $n$ is always a multiple of the number of elements in the chemical
composition of the material. For example, for a formula MgO3, which consists
of 4 atoms, we can have 4, 8, 12, or generally $4k$ elements in a unit cell
(one element corresponding to one atom), but not 3 or 6 elements, which are
not a multiple of 4.

Finally, we can describe atoms’ types and positions with two matrices
$A\in\mathbb{R}^{n\times 128}$, $X\in\mathbb{R}^{n\times 3}$. The matrix $A$
identifies different atoms in the unit cell using a one-hot representation.
Specifically, $A_{i}$ is a vector with a $1$ at position $Z$ corresponding to
the atomic number of the $i$-th element, and $0$ everywhere else. The matrix
$X$ provides fractional coordinates for the atoms. These are coordinates
between $0$ and $1$ with respect to the basis defined by the lattice base
vectors. More specifically, the vector $X_{i}$ tells us that the $i$-th
element is at $X_{i}^{1}\textbf{v}+X_{i}^{2}\textbf{w}+X_{i}^{3}\textbf{z}$ in
the unit cell.

To summarize, a crystal is defined by three quantities: (1) A $3\times 3$
matrix $L$ representing the lattice, the rows of which corresponding to the
base vectors of the lattice; (2) number ($n$) and types
($A\in\mathbb{R}^{n\times 128}$) in the chemical; and (3) atoms positions
$X\in\mathbb{R}^{n\times 3}$, which is specified in terms of fractional
coordinates between $0$ and $1$. In the following sections, we will denote a
crystal with the variable $\mathbf{x}$ and we will use subscripts to refer to
lattice parameters $\mathbf{x}_{L}$, atoms’ types $\mathbf{x}_{A}$, and
fractional coordinates $\mathbf{x}_{X}$. We finally remark that the majority
of crystal and lattice configurations for a given chemical compound are
“unstable” and would collapse to a different configuration when synthesized.

###  3 Problem Statement, Dataset, and Evaluation

Crystal structure prediction (CSP) is the problem of finding a crystal of a
given chemical composition (e.g. NaCl, or MgO3) that attains the global
minimum of the crystal formation energy. Such a crystal is sythesizable and
can be utilized for various downstream applications.

######  Problem 3.1 (Crystal structure prediction).

Given a chemical composition $c$, find the crystal
$\mathbf{x}^{*}=(L^{*},A^{*},X^{*})$ with lattice matrix $L^{*}$, atom types
$A^{*}$, and atom coordinates $X^{*}$ such that $\mathbf{x}^{*}$ minimizes the
formation energy function $E$ for the chemical composition.

$\mathbf{x}^{*}=\mathrm{argmin}_{\mathbf{x}}\,E\left(\mathbf{x},c\right).$

Why is solving CSPs hard? Only very few crystal structures are actually stable
and only one of these stable structures is at a global minimum, which makes
crystal structure prediction equivalent to searching for a “needle in a
haystack”. The difficulty of solving a CSP is further compounded by the fact
that the search space over all possible crystal structures for a given
chemical composition is quite complex and non-Euclidean. This is because there
is no one-to-one correspondence between the matrices $(A,X)$ and lattices $L$.
Given a lattice matrix $L$, every other matrix that is rotationally equivalent
or permutation equivalent represents the same lattice. Moreover, reducing the
design space from all possible structures to only stable ones changes the
manifold of designs considerably.

In principle, we could always attempt to find the globally optimal crystal
structure by evaluating many possible candidate structures against a
simulator, but simulators for CSP are typically based on density functional
theory (DFT), and generally these are extremely slow in terms of the wall-
clock time. Therefore, we intend to solve this problem using only existing
static datasets (OQMD (Saal et al., 2013) and MatBench Dunn et al. (2020)).
that contain several (sub-optimal) crystal structures for a variety of
chemical compounds along with their corresponding formation energies. With no
access to the simulator, our goal is to find a globally optimal crystal
structure given a new chemical compound. We describe our procedure for
constructing the dataset and our evaluation protocol next.

Datasets for training. To robustly evaluate our method, we consider two
training scenarios with different datasets: the OQMD dataset (Saal et al.,
2013) and the MatBench dataset Dunn et al. (2020), both of which consist of
the crystal structures and formation energies for obtained via DFT
simulations. Every sample in the dataset represents a stable crystal structure
$\mathbf{x}$ (i.e., a crystal at a local minimum of energy) computed via
numerical relaxation (Hafner, 2008), together with its chemical composition
$c$ and formation energy $E\left(\mathbf{x},c\right)$. We chose these datasets
because of their large size (OQMD has more than 1 million examples), and the
availability of more than one stable structure per chemical, all of which are
not at their global optimum.

Held-Out evaluation datasets. We evaluate our offline optimization approach in
terms of its efficacy in discovering the globally optimal structure for a
given chemical compound. To this end, following the protocol of Cheng et al.
(2022), we construct a held-out test set consisting of some compounds and the
associated globally optimal crystal structures (Table 1) and utilize this set
for evaluations.

Evaluation protocol. Following the evaluation protocol in prior works Xie et
al. (2021), Cheng et al. (2022), we test our method in terms of its efficacy
in recovering the globally optimal structure on 26 of the 29 compounds in
(Cheng et al., 2022), where the 3 remaining compounds are omitted because they
cannot be simulated with GPAW to compute their ground truth energy values.
These compounds are listed in Table 1. For each of these compounds, we run our
optimization process to convergence and check the final energy of the
optimized design. We compare the optimized energy against the known global
minimum energy. We consider it a success if the final energy of the optimized
structure produced by our approach recovers the value of the known global
minimum, up to a predefined noise threshold of $20\%$ of the ground truth
minimum energy.

To stress-test our gradient-based approach, we seed the optimizer with a
random stable initial crystal structure. We compute this initial stable
structure by running simulations in the GPAW (Enkovaara et al., 2011)
simulator, an open-source DFT simulator fully integrated into python packages
for chemistry like ase and pymatgen. Concretely, we utilized the following
protocol for obtaining this initial crystal structure: (i) for every compound,
we first select the number of atoms corresponding to the optimal compound
design as listed in the materials project database, (ii) we then randomly
initialize the lattice matrix and the atom coordinates, and (iii) we then run
the process of structure relaxation in the simulator to obtain the closest
local minimum (i.e., a closeby structure that is stable).

###  4 LCOMs: Latent Conservative Objective Models

To design crystal structures with the lowest possible energy entirely from an
existing dataset of only sub-optimal structures, we utilize techniques from
offline model-based optimization. Directly applying these techniques (Trabucco
et al., 2021b, Yu et al., 2021, Qi et al., 2022) for optimizing crystals is
non-trivial as these methods typically employ optimization procedures that
iteratively make local changes to the design (e.g., gradient-based updates or
random mutations) to optimize a “surrogate” estimate of the objective
function. Such iterative procedures fall short when optimizing over non-smooth
geometries and non-Euclidean manifolds like that of crystals. To alleviate
this issue, we propose an approach for offline optimization that first learns
a latent vector representation for a crystal structure, then performs data-
driven optimization in this vector space, and finally maps back the resulting
outcome to a valid crystal structure. We outline each part of this procedure
below.

####  4.1 Transforming Crystal Structures to a Latent Representation

Which sort of a latent representation space is especially desirable for CSP?
Since one of the central challenges in our problem is the abundance of invalid
or infeasible structures in the space of all possible crystals, it is very
desirable to learn latent representations that only encode valid and feasible
structures. Once we learn such a space, we can directly perform offline
optimization in this latent space. If we can ensure that every possible latent
vector corresponds to some feasible crystal structure, then we are guaranteed
to at least prune out the possibility of finding infeasible structures during
the optimization process. To this end, we train a crystal diffusion
variational auto-encoder (CD-VAE) on our training dataset for various chemical
compositions, and then run offline optimization in the latent space of this
auto-encoder. Since our training dataset only consists of stable structures,
the decoder of a well-trained CD-VAE should map latent vectors to the manifold
of stable crystal structures only, which would greatly benefit optimization by
enabling the optimizer to move in a much smaller manifold. Below we describe
the architecture and the training objective for the CD-VAE.

CD-VAE. A CD-VAE (Xie et al., 2021) is composed of three parts: a graph neural
network (GNN) encoder $\textsc{PGNN}_{\textsc{Enc}}$ that takes a crystal
$\mathbf{x}$ as input and outputs a latent vector representation, an NN
predictor $\textsc{MLP}_{\textsc{Agg}}$ that outputs lattice parameters
$\mathbf{x}_{L}$ from its encoded representation
$\textsc{PGNN}_{\textsc{Enc}}\left(\mathbf{x},c\right)$, and a GNN diffusion
denoiser $\textsc{PGNN}_{\textsc{Dec}}$ that takes a random noisy crystal
$\tilde{\mathbf{x}}$ and a latent encoding
$\textsc{PGNN}_{\textsc{Enc}}\left(\mathbf{x},c\right)$ as inputs, and outputs
forces to apply on the atoms coordinates $\tilde{\mathbf{x}}_{X}$ to build the
original crystal $\mathbf{x}$ via a diffusion process. Following Xie et al.
(2021), the encoder $\textsc{PGNN}_{\textsc{Enc}}$ uses a DimeNet++ Klicpera
et al. (2020) architecture. Likewise, the decoder
$\textsc{PGNN}_{\textsc{Dec}}$ uses a GemNet-dQ Gasteiger et al. (2021)
architecture. During decoding, CD-VAE initializes a structure with random
lattice and coordinates and utilizes Langevin dynamics Song and Ermon (2019)
to gradually recover the stable structure represented by the latent vector.

To make the notation compact, we will refer to the
$\textsc{PGNN}_{\textsc{Enc}}$ as $\phi$, and thus the latent representation
of a crystal $\mathbf{x}$ with chemical composition $c$ will be denoted by
$\phi(\mathbf{x},c)$. We will use the notation
$\textsc{PGNN}_{\textsc{Dec}}\left(z\right)$ to denote the structure obtained
after applying the denoising process with latent vector $z$. Akin to a
variational auto-encoder (Kingma and Welling, 2013), CD-VAE (Xie et al., 2021)
is also trained to maximize the likelihood of crystral structures seen in the
dataset, agnostic of the energy objective that we wish to optimize.

Training objective for the latent representation. We follow the training
objective utilized by the CD-VAE (Xie et al., 2021). The first term in this
objective is the reconstruction error over lattice parameters, formally
defined as:

$\displaystyle\mathcal{L}_{\textsc{Agg}}\left(\textsc{MLP}_{\textsc{Agg}}\left(\phi(\mathbf{x},c)\right),\mathbf{x}_{L}\right)={\lVert\mathbf{x}_{L}-\phi(\mathbf{x},c)\rVert}^{2}.$
(3)

Akin to a VAE (Kingma and Welling, 2013), we also include a loss term
minimizes the KL-divergence between a normal distribution over the latent
representation induced by the encoder (with mean $\phi(\mathbf{x},c)$ and a
learned standard deviation) and a standard multi-dimensional normal
distribution.

The decoder of the CD-VAE is a denoising diffusion model (Ho et al., 2020)
that attempts to transform an input latent vector into the corresponding
crystal structure, starting from a random structure $\tilde{\mathbf{x}}$,
which is iteratively refined via the diffusion process. Succinctly, the
objective for training this term is

$\displaystyle\mathcal{L}_{\textsc{Dec}}\left(\tilde{\mathbf{x}},\mathbf{x},\phi(\mathbf{x},c)\right)=\frac{1}{2L}\sum_{j=1}^{L}\left[\mathbb{E}_{\sigma_{j}}\left\lVert\textsc{PGNN}_{\textsc{Dec}}\left(\tilde{\mathbf{x}},\phi(\mathbf{x},c)\right)-\frac{d\left(\mathbf{x}_{X},\tilde{\mathbf{x}}_{X}\right)}{\sigma_{j}}\right\rVert\right],$
(4)

where $\left\\{\sigma_{j}\right\\}_{j=1}^{L}$ are noise schedule scalars for
the diffusion process and are in a geometric sequence with common ratio
greater than $1$. Finally, we remark that we do not utilize terms for
reconstructing types or the number of atoms (i.e., $\mathbf{x}_{A}$ or
$\mathbf{x}_{n}$) because our optimization procedure only aims to optimize
over other parameters of the crystal lattice so the number and types of atoms
are fixed.

####  4.2 Conservative Optimization in Latent Space

Once the encoder of the CD-VAE is trained, we can then train a surrogate
model, $\widehat{E}_{\theta}(\phi(\mathbf{x},c),c)$ to estimate the formation
energy $E$ of a crystal structure $\mathbf{x}$ for a given chemical
composition, $c$ via standard supervised regression. Then, we can simply
optimize the crystal structure to maximize the outputs of this surrogate
model. However, as prior works (Kumar et al., 2020, Trabucco et al., 2021b)
note, this simple strategy often fails at finding optimized designs due to the
exploitation of errors in the learned surrogate model by the optimizer. To
address this issue, we extend the conservative objective models (COMs)
technique for optimizing crystals in the learned latent space.

Training latent space conservative models. To prevent the optimization
procedure from exploiting inaccuracies in this learned surrogate model, we
apply an additional regularizer from Trabucco et al. (2021b), Kumar et al.
(2021) to robustify the surrogate model. This regularizer mines for
adversarial vectors in the latent space $z^{+}$ that appear to have very low
energies $\widehat{E}_{\theta}(z^{+},c)$ under the learned surrogate model,
and then explicitly pushes up the predicted energy
$\widehat{E}_{\theta}(z^{+},c)$ on such adversarial $z^{+}$. Following the
COMs approach (Trabucco et al., 2021b), we interleave the training of the
learned surrogate model $\widehat{E}_{\theta}$ with an optimization procedure
$\mathsf{Opt}(\widehat{E}_{\theta},c)$ that seeks to find the aforementioned
adversarial vectors $z^{+}$ that optimize the current snapshot for the
surrogate model, for a given chemical composition $c$. After these adversarial
vectors are found, the training procedure explicitly pushes up the energy
output of the surrogate model on such points. To compensate for the effect of
increasing the learned energy values in an unbounded manner on all latent
vectors, we additionally balance the push up term by pushing down the energy
values on the latent representations induced by crystal structures in the
data. This idea can be formalized into the following loss for training
$\widehat{E}_{\theta}$:

We will discuss the formulation for $\mathsf{Opt}$ below. Crucially, unlike
COMs (Trabucco et al., 2021b), which directly runs gradient descent in the
input space, our approach operates in the latent space.

Optimizing in the latent space. Once a conservative surrogate model
$\widehat{E}_{\theta}(z,c)$ is obtained using the above training procedure, we
must now optimize this model to obtain the best possible structures. The
optimization procedure $\mathsf{Opt}$ that was used to obtain adversarial
latent vectors in the training objective above can then be repurposed to
obtain optimized latent vectors once the latent conservative model is trained.
Since the latent space $z$ is a continuous Euclidean vector space, for any
given chemical composition $c$, our choice of $\mathsf{Opt}$ is to run $T$
rounds of gradient descent on the surrogate energy $\widehat{E}_{\theta}(z,c)$
with respect to the latent vector $z$, starting from the latent vector $z_{0}$
corresponding to a random initial crystal structure. For a given $c$, this
procedure can be formalized as follows:

$\displaystyle z_{k+1}$ $\displaystyle\leftarrow
z_{k}-\alpha\nabla_{z}\widehat{E}_{\theta}(z,c),$ (5) where
$\displaystyle~{}~{}z_{0}\sim\phi(\mathbf{x}_{0},c),~{}~{}\mathbf{x}_{0}\sim\mathcal{D}.$

Once this optimization procedure is run for $T$ steps, we pass the final
latent vector $z_{T}$ to the decoder of the pre-trained CD-VAE to obtain the
optimized crystal structure:
$\widehat{\mathbf{x}}^{*}=\textsc{PGNN}_{\textsc{Dec}}(z_{T})$.

####  4.3 Implementation Details

For obtaining the latent space, we train a CD-VAE identically to Xie et al.
(2021) on our datasets, following their implementation details for the encoder
and the decoder. After training the CD-VAE, we encode molecule structures from
the dataset into vectors, and these vectors are then used as inputs for
training the optimization model. For training LCOMs, we represent the
conservative objective model $\widehat{E}_{\theta}(\phi(\mathbf{x},c),c)$ as a
neural network with two hidden layers of size 2048 each and leaky ReLU
activations. For computing $z^{+}$, we perform one gradient descent step on
the vector $z$ from input latent space. We perform 50 gradient steps and get
an optimized vector in the latent space. We then decode these latent vectors
into an optimized crystal structure. For reporting statistically robust
results, we repeat the optimization procedure for three seeds and average over
the resulting energy value.

###  5 Related Work

One widely studied optimization-based approach to CSP utilizes evolutionary
algorithms Lonie and Zurek (2011), Oganov et al. (2011). For instance, USPEX
Glass et al. (2006) is an algorithm that uses a variety of heuristic
strategies for iteratively evolving structures directly in the space of the
crystal parameters. Each of these intermediate structures need to be evaluated
against the simulator, which repeatedly involves running relaxation to the
nearest stable structure. This extensive use of simulation makes such an
approach computationally impractical, necessitating offline learning-based
approaches like our method that do not require any simulation.

Due to the availability of large public datasets such as the materials project
database and the open catalyst project (Chanussot et al., 2021), recent works
develop learning-based approaches for solving CSPs. Another subset includes
methods that use different types of evolutionary optimizers to optimize a GNN-
based surrogate energy model instead of the ground-truth energy function. This
includes methods based on random search (Cheng et al., 2022), particle swarm
optimization (Clerc, 2010), and Bayesian optimization (Pelikan et al., 1999).
In contrast, our method prescribes the use a conservative surrogate model of
the energy function, that takes as input a latent representation of the
crystal. As we show in our experimental results, both of these aspects are
crucial for CSP.

Another line of prior work learns generative models for graph data including,
but not limited to crystal structures. This includes methods that leverage
variational auto-encoders (Simonovsky and Komodakis, 2018) normalizing flows
(Satorras et al., 2021), generative adversarial networks (Kim et al., 2020),
recurrent neural networks Grisoni et al. (2020), reinforcement learning
techniques (Pereira et al., 2020) or a combination of auto-encoders and
diffusion models, for example, the CD-VAE (Xie et al., 2021), that we build
upon. While these approaches aim to model the manifold the graph-structured
data and our approach utilizes the latent space learned by one such approach,
CD-VAE (Xie et al., 2021), our goal of optimizing the structure is distinct
from the goal of modelling the data.

Model-based optimization (MBO) refers to the problem of optimizing an unknown
function by constructing a surrogate model. Bayesian optimization represents
one of the most widely known classes of MBO methods (Snoek et al., 2015; 2012,
Ghavamzadeh et al., 2015), but classically MBO requires iteratively sampling
new function values, which can be very expensive when evaluating a crystal
structure’s energy requires an expensive simulation process. More recently,
offline MBO methods, sometimes referred to as data-driven optimization, have
been proposed to optimize designs based entirely on previously collected
_static_ datasets (Brookes et al., 2019, Trabucco et al., 2021a, Kumar and
Levine, 2019, Trabucco et al., 2021b, Qi et al., 2022). Our work builds on
these methods, and is most closely related to the COMs algorithm proposed by
Trabucco et al. (2021b). However, while prior offline MBO methods focus on
robustifying the surrogate model in the design space directly, we integrate
these approaches with latent space optimization that makes it possible to
optimize over the manifold of only stable structures, while still using a
simple gradient optimizer.

###  6 Experimental Evaluation

The goal of our experimental evaluation is to evaluate the efficacy of LCOMs
for crystal structure prediction by answering the following questions: (1) Can
LCOMs successfully optimize in the latent representation space? (2) Do LCOMs
manage to effectively recover the optimal energy structure up to a pre-defined
threshold of accuracy? and (3) Does LCOM drastically reduce simulation wall-
clock time compared to prior methods? To answer these questions, we evaluate
LCOMs against prior methods following the protocol from Section 3 and then
perform some diagnostic experimental studies that we will discuss in this
section.

Comparing LCOMs with baselines and prior methods. We compare LCOMs to two
methods from prior work (Cheng et al., 2022): particle-swarm optimization
(PSO), and Bayesian optimization (BO). We also study two baseline methods: a
method that does not run any optimization in the latent space and simply
constructs a stable structure for a given chemical compound via the decoder
(“CD-VAE”), and a method where the crystal is optimized with a naïve
supervised learning model in the latent space of the CD-VAE via gradient
descent (”Supervised learning; SL”). Note that the latter is similar to LCOMs,
but the surrogate energy prediction model is not trained with any
conservatism, but rather with only standard supervised regression. We evaluate
these methods on the 26 chemical compounds in our evaluation dataset, and
present the results in Table 1. During evaluation, we mark a crystal structure
successful if the energy of the optimized structure is close to the energy of
the best known globally optimal structure up to a certain threshold. This
threshold is defined as an upper bound of 0.2 on the quantity
$\left(E(\mathbf{x},c^{*})-E(\mathbf{x},\hat{c})\right)/|E(\mathbf{x},c^{*})|$
to account for imprecision in the simulator, where $c^{*}$ is the ground truth
optimal crystal and $\hat{c}$ is the optimized crystal discovered by the
optimizer.

OQMD | Baselines | MatBench
---|---|---
Compounds | RAS* | PSO* | BO* | LCOMs | SL | CD-VAE | RAS* | PSO | BO | LCOMs
LiF |  | ✓ | ✓ | ✓ |  |  | ✓ | ✓ |  | ✓
NaF | ✓ | ✓ | ✓ |  | ✓ |  |  |  |  |
KF | ✓ |  | ✓ | ✓ |  | ✓ |  |  |  | ✓
RbF |  |  |  | ✓ | ✓ | ✓ | ✓ |  |  | ✓
CsF |  |  |  | ✓ |  | ✓ | ✓ |  | ✓ | ✓
LiCl |  |  |  |  |  |  |  | ✓ |  |
NaCl | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓
KCl | ✓ |  |  | ✓ |  |  | ✓ |  | ✓ | ✓
RbCl | ✓ |  |  | ✓ |  |  | ✓ | ✓ |  | ✓
CsCl | ✓ |  | ✓ | ✓ |  |  | ✓ |  |  | ✓
BeO |  | ✓ | ✓ | ✓ |  | ✓ |  | ✓ | ✓ | ✓
MgO | ✓ |  | ✓ |  |  |  | ✓ |  | ✓ | ✓
CaO | ✓ |  | ✓ | ✓ |  |  | ✓ |  |  | ✓
SrO | ✓ |  | ✓ | ✓ |  |  | ✓ | ✓ |  | ✓
BaO | ✓ |  | ✓ |  |  |  | ✓ |  |  | ✓
ZnO |  | ✓ | ✓ | ✓ |  |  |  | ✓ | ✓ | ✓
CdO |  |  |  |  |  | ✓ | ✓ |  |  |
BeS | ✓ | ✓ | ✓ | ✓ |  |  | ✓ | ✓ |  | ✓
MgS | ✓ |  |  |  | ✓ |  | ✓ | ✓ |  | ✓
CaS | ✓ |  | ✓ | ✓ |  |  | ✓ | ✓ | ✓ | ✓
SrS | ✓ |  |  | ✓ | ✓ |  | ✓ | ✓ | ✓ | ✓
BaS | ✓ |  | ✓ |  |  |  | ✓ | ✓ | ✓ |
ZnS | ✓ |  | ✓ | ✓ |  |  |  |  |  | ✓
CdS |  |  | ✓ |  |  |  |  | ✓ | ✓ |
C | ✓ |  |  |  |  |  | ✓ |  |  |
Si |  |  |  |  |  |  |  |  |  |
Accuracy | 17/26 | 6/26 | 16/26 | 16/26 | 5/26 | 6/26 | 18/26 | 13/26 | 10/26 | 19/26

Table 1: Evaluation of LCOMs and other prior crystal structure prediction
methods in terms of the accuracy of discovering the globally optimal structure
for 26 compounds. Check marks indicate successful discovery as per our
criterion. Note crucially that while we utilize a threshold on optimization
energy to determine success, baseline methods marked with ∗ in this table
follow prior works and utilize a manual inspection protocol as discussed in
the footnote in Section 6.

The results in Table 1 show that when trained on OQMD, our method improves
significantly over naïvely optimizing in the CD-VAE latent space without
conservative training (supervised learning; SL), and also that it is
competitive with the prior state-of-the-art methods RAS, PSO, and BO,
exceeding the performance of the PSO baseline and matching BO and RAS, without
needing any simulations (we will quantify the benefits on wall-clock time
soon). A similar trend also holds for the MatBench dataset in Table 1,
indicating that LCOMs is performant for different choices of the training
data. On MatBench, LCOMs outperforms both PSO and BO methods by a large
margin. We also note that BO and RAS perform comparably to LCOMs when trained
on the OQMD dataset, when results for these prior comparisons are reported
using the evaluation metric in (Cheng et al., 2022)111Our evaluation protocol
for determining the optimality of a structure uses an energy threshold. This
is a contrast to Cheng et al. (2022), which adopts a manual inspection
approach to check for equality between optimized and optimal structures as
their evaluation criterion. As such, comparisons between our work and that of
Cheng et al. (2022) marked with an asterisk (∗) should be made with an
understanding of this fundamental difference in evaluation methodology.. Due
to the difference in evaluation criteria, these comparisons to LCOMs are not
as exact like in MatBench, where both BO and PSO are evaluated using the same
criteria as LCOMs.

Does LCOM improve over prior methods in terms of wall-clock time? Our method
optimizes the crystal structure in the latent space of the CD-VAE, using
gradient-based optimization. One advantage of this approach is computational
efficiency, since the complex graph-based component of the pipeline is only
used during the encoding and decoding stages at the beginning and end of the
optimization, rather than at each optimization step. Hence, we measure the
wall-clock time needed to run optimization with LCOMs comparatively against
other prior methods in Table 2. Observe that while utilizing a graph neural
network (GNN-BO) reduces the wall-clock time needed by about 875$\times$
compared to DFT-PSO that queries the simulator for every design, LCOMs further
reduces the wall-clock time 40$\times$ by running optimization in the latent
space of a CD-VAE, which does not require running expensive message passing
loops of a graph neural network encoder but rather runs relatively faster
forward passes through small MLPs. This indicates that LCOMs not only
discovers are more optimal crystal structure, but it does so 40$\times$ faster
than the best prior method.

| DFT-PSO | GN-BO | LCOMs
---|---|---|---
Optimization time per structure (seconds) | 70000 | 80 | 2

Table 2: Comparing wall-clock time for different methods. Observe that not
using a simulator reduces the wall-clock time from 70000 seconds to 80 seconds
per structure, and further utilizing a latent space surrogate model in LCOMs
instead of a graph neural network model cuts down the total time further by
40$\times$ to 2 seconds. Figure 2: Comparing the energy of intermediate
structures observed over the course of optimization with LCOMs (blue) and non-
conservative models (orange). Note that while the non-conservative model gets
exploited as more steps of gradient-based optimization are performed,
structures discovered by conservative LCOMs attain lower energies after
gradient descent, and the final structures are close to the global optimum
(marked as red in the plot above).

How does conservative training influence optimization? In the next set of
experiments, we aim to understand how conservative training of the surrogate
model influences the behavior of gradient descent optimization in the CD-VAE
latent space. If the energy model is trained naïvely (i.e., with standard
supervised regression), it will make arbitrarily erroneous predictions when
queried on adversarial, out-of-distribution crystals that differ too much from
the training data. Some of these errors will be under-estimation errors, and
therefore, a strong optimizer would be able to exploit these errors to find
points in the latent space for which the model erroneously predicts
arbitrarily low energies. Empirically, we evaluate the performance of
optimized crystals obtained by running 50 gradient steps of optimization on
the learned surrogate models starting from an initial structure. Specifically,
we compute the relative improvement in energy values after optimization,
formally calculated as
$\left(E(\mathbf{x},c_{0})-E(\mathbf{x},\hat{c})\right)/|E(\mathbf{x},c_{0})|$,
where $c_{0}$ denotes the initialization and $\hat{c}$ denotes the optimized
structure in Figure 3.

Observe that while LCOMs generally produces positive improvement, the non-
conservative model leads to _negative_ improvement: the optimized structures
are generally worse than the random structure at initialization. This
indicates that conservative training is critical for latent space optimization
to work. We also perform a more fine-grained analysis, where we plot the
trajectory of evolution of the energy values over each round of optimization
in Figure 2. Observe that for the non-conservative model (orange), the energy
increases over the course of optimization indicating that the optimizer is
exploiting errors in the learned model. This exploitation is absent for LCOMs
(in blue), indicating that conservatism is crucial for attaining good
performance.

Figure 3: Comparison of energy improvements produced by LCOMs and the non-
conservative supervised learning (SL) baseline. Top: Energy improvement for 25
compositions, when training on the OQMD dataset; Bottom: energy improvement
over a set of 83 chemical compositions, when training on the MatBench dataset.
Note that structures found by LCOMs achieve better formation energy after
optimization, while the non-conservative supervised learning model actually
leads to structures with worse energy values (negative improvement). The
quantity on the y-axis represents the percentage of improvement (reduction) in
the energy value. These results indicate that conservative training is
essential for successfully instantiating a method with gradient-based latent
space optimization for crystal structure prediction.

###  7 Discussion and Future Directions

We presented a method for offline optimization that uses the latent space of a
CD-VAE to perform smooth gradient-based optimization of complex structures,
with application to crystal structure prediction. Our method combines concepts
from conservative objective models that robustify predictive models to make
them amenable to gradient-based optimization, with a latent-space generative
models of graphs, enabling us to use simple gradient-based optimization
methods. Our experiments show that our method can successfully optimize the
formation energy and recover the optimal structure of a chemical compound with
a good level of accuracy, comparing favorably with existing approaches, while
tremendously reducing computation time. An interesting avenue for future work
is to study the efficacy of LCOMs in more problems in computational chemistry.
Another direction for future work is to use the best performing models to
predict the optimized structure for novel chemicals and validate the
predictions experimentally.

### References

  * Brookes et al. (2019) D. Brookes, H. Park, and J. Listgarten.  Conditioning by adaptive sampling for robust design.  In _Proceedings of the 36th International Conference on Machine Learning_. PMLR, 2019.  URL http://proceedings.mlr.press/v97/brookes19a.html.
  * Chanussot et al. (2021) L. Chanussot, A. Das, S. Goyal, T. Lavril, M. Shuaibi, M. Riviere, K. Tran, J. Heras-Domingo, C. Ho, W. Hu, et al.  Open catalyst 2020 (oc20) dataset and community challenges.  _Acs Catalysis_ , 11(10):6059–6072, 2021.
  * Cheng et al. (2022) G. Cheng, X.-G. Gong, and W.-J. Yin.  Crystal structure prediction by combining graph network and optimization algorithm.  _Nature Communications_ , 13(1):1492, 2022.
  * Chermette (1998) H. Chermette.  Density functional theory: a powerful tool for theoretical studies in coordination chemistry.  _Coordination chemistry reviews_ , 178:699–721, 1998.
  * Clerc (2010) M. Clerc.  _Particle swarm optimization_ , volume 93.  John Wiley & Sons, 2010.
  * Dunn et al. (2020) A. Dunn, Q. Wang, A. Ganose, D. Dopp, and A. Jain.  Benchmarking materials property prediction methods: the matbench test set and automatminer reference algorithm.  _npj Computational Materials_ , 6(1):138, 2020\.
  * Enkovaara et al. (2010) J. Enkovaara, C. Rostgaard, J. J. Mortensen, J. Chen, M. Dułak, L. Ferrighi, J. Gavnholt, C. Glinsvad, V. Haikola, H. A. Hansen, H. H. Kristoffersen, M. Kuisma, A. H. Larsen, L. Lehtovaara, M. Ljungberg, O. Lopez-Acevedo, P. G. Moses, J. Ojanen, T. Olsen, V. Petzold, N. A. Romero, J. Stausholm-Møller, M. Strange, G. A. Tritsaris, M. Vanin, M. Walter, B. Hammer, H. Häkkinen, G. K. H. Madsen, R. M. Nieminen, J. K. Nørskov, M. Puska, T. T. Rantala, J. Schiøtz, K. S. Thygesen, and K. W. Jacobsen.  Electronic structure calculations with gpaw: a real-space implementation of the projector augmented-wave method.  _Journal of Physics: Condensed Matter_ , 22(25):253202, jun 2010.  doi: 10.1088/0953-8984/22/25/253202.  URL https://dx.doi.org/10.1088/0953-8984/22/25/253202.
  * Enkovaara et al. (2011) J. Enkovaara, N. A. Romero, S. Shende, and J. J. Mortensen.  Gpaw-massively parallel electronic structure calculations with python-based software.  _Procedia Computer Science_ , 4:17–25, 2011.
  * Gasteiger et al. (2021) J. Gasteiger, F. Becker, and S. Günnemann.  Gemnet: Universal directional graph neural networks for molecules.  _Advances in Neural Information Processing Systems_ , 34:6790–6802, 2021.
  * Ghavamzadeh et al. (2015) M. Ghavamzadeh, S. Mannor, J. Pineau, and A. Tamar.  Bayesian reinforcement learning: A survey.  _Found. Trends Mach. Learn._ , 8(5-6):359–483, Nov. 2015.  ISSN 1935-8237.  doi: 10.1561/2200000049.  URL http://dx.doi.org/10.1561/2200000049.
  * Glass et al. (2006) C. W. Glass, A. R. Oganov, and N. Hansen.  Uspex—evolutionary crystal structure prediction.  _Computer physics communications_ , 175(11-12):713–720, 2006.
  * Grisoni et al. (2020) F. Grisoni, M. Moret, R. Lingwood, and G. Schneider.  Bidirectional molecule generation with recurrent neural networks.  _Journal of chemical information and modeling_ , 2020.
  * Hafner (2008) J. Hafner.  Ab-initio simulations of materials using vasp: Density-functional theory and beyond.  _Journal of computational chemistry_ , 29(13):2044–2078, 2008.
  * Ho et al. (2020) J. Ho, A. Jain, and P. Abbeel.  Denoising diffusion probabilistic models.  _Advances in Neural Information Processing Systems_ , 33:6840–6851, 2020.
  * Kim et al. (2020) S. Kim, J. Noh, G. H. Gu, A. Aspuru-Guzik, and Y. Jung.  Generative adversarial networks for crystal structure prediction.  _ACS central science_ , 6(8):1412–1420, 2020\.
  * Kingma and Welling (2013) D. P. Kingma and M. Welling.  Auto-encoding variational bayes.  _arXiv preprint arXiv:1312.6114_ , 2013.
  * Klicpera et al. (2020) J. Klicpera, S. Giri, J. T. Margraf, and S. Günnemann.  Fast and uncertainty-aware directional message passing for non-equilibrium molecules.  _arXiv preprint arXiv:2011.14115_ , 2020.
  * Kumar and Levine (2019) A. Kumar and S. Levine.  Model inversion networks for model-based optimization.  _NeurIPS_ , 2019.
  * Kumar et al. (2020) A. Kumar, A. Zhou, G. Tucker, and S. Levine.  Conservative q-learning for offline reinforcement learning.  _arXiv preprint arXiv:2006.04779_ , 2020.
  * Kumar et al. (2021) A. Kumar, A. Yazdanbakhsh, M. Hashemi, K. Swersky, and S. Levine.  Data-driven offline optimization for architecting hardware accelerators.  _arXiv preprint arXiv:2110.11346_ , 2021.
  * Lonie and Zurek (2011) D. C. Lonie and E. Zurek.  Xtalopt: An open-source evolutionary algorithm for crystal structure prediction.  _Computer Physics Communications_ , 182(2):372–387, 2011.
  * Mortensen et al. (2005) J. J. Mortensen, L. B. Hansen, and K. W. Jacobsen.  Real-space grid implementation of the projector augmented wave method.  _Phys. Rev. B_ , 71:035109, Jan 2005.  doi: 10.1103/PhysRevB.71.035109.  URL https://link.aps.org/doi/10.1103/PhysRevB.71.035109.
  * Oganov et al. (2011) A. R. Oganov, A. O. Lyakhov, and M. Valle.  How evolutionary crystal structure prediction works and why.  _Accounts of chemical research_ , 44(3):227–237, 2011.
  * Parr (1989) R. G. Parr.  Density-functional theory of atoms and molecules.  1989\.
  * Pelikan et al. (1999) M. Pelikan, D. E. Goldberg, E. Cantú-Paz, et al.  Boa: The bayesian optimization algorithm.  In _Proceedings of the genetic and evolutionary computation conference GECCO-99_ , volume 1, pages 525–532. Citeseer, 1999.
  * Pereira et al. (2020) T. Pereira, M. Abbasi, B. Ribeiro, and J. P. Arrais.  Diversity oriented deep reinforcement learning for targeted molecule generation.  _Journal of Cheminformatics_ , 13, 2020.
  * Qi et al. (2022) H. Qi, Y. Su, A. Kumar, and S. Levine.  Data-driven offline decision-making via invariant representation learning.  _arXiv preprint arXiv:2211.11349_ , 2022.
  * Saal et al. (2013) J. E. Saal, S. Kirklin, M. Aykol, B. Meredig, and C. Wolverton.  Materials design and discovery with high-throughput density functional theory: the open quantum materials database (oqmd).  _Jom_ , 65:1501–1509, 2013.
  * Sarkisyan et al. (2016) K. S. Sarkisyan, D. A. Bolotin, M. V. Meer, D. R. Usmanova, A. S. Mishin, G. V. Sharonov, D. N. Ivankov, N. G. Bozhanova, M. S. Baranov, O. Soylemez, N. S. Bogatyreva, P. K. Vlasov, E. S. Egorov, M. D. Logacheva, A. S. Kondrashov, D. M. Chudakov, E. V. Putintseva, I. Z. Mamedov, D. S. Tawfik, K. A. Lukyanov, and F. A. Kondrashov.  Local fitness landscape of the green fluorescent protein.  _Nature_ , 533(7603):397–401, May 2016.  ISSN 1476-4687.  doi: 10.1038/nature17995.  URL https://doi.org/10.1038/nature17995.
  * Satorras et al. (2021) V. G. Satorras, E. Hoogeboom, F. B. Fuchs, I. Posner, and M. Welling.  E(n) equivariant normalizing flows for molecule generation in 3d.  _ArXiv_ , abs/2105.09016, 2021.
  * Simonovsky and Komodakis (2018) M. Simonovsky and N. Komodakis.  Graphvae: Towards generation of small graphs using variational autoencoders.  _ArXiv_ , abs/1802.03480, 2018.
  * Snoek et al. (2012) J. Snoek, H. Larochelle, and R. P. Adams.  Practical bayesian optimization of machine learning algorithms.  In _Proceedings of the 25th International Conference on Neural Information Processing Systems - Volume 2_ , NIPS’12, 2012.  URL http://dl.acm.org/citation.cfm?id=2999325.2999464.
  * Snoek et al. (2015) J. Snoek, O. Rippel, K. Swersky, R. Kiros, N. Satish, N. Sundaram, M. Patwary, M. Prabhat, and R. Adams.  Scalable bayesian optimization using deep neural networks.  In _Proceedings of the 32nd International Conference on Machine Learning_. PMLR, 2015.
  * Song and Ermon (2019) Y. Song and S. Ermon.  Generative modeling by estimating gradients of the data distribution.  In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett, editors, _Advances in Neural Information Processing Systems_ , volume 32. Curran Associates, Inc., 2019.  URL https://proceedings.neurips.cc/paper/2019/file/3001ef257407d5a371a96dcd947c7d93-Paper.pdf.
  * Trabucco et al. (2021a) B. Trabucco, X. Geng, A. Kumar, and S. Levine.  Design-bench: Benchmarks for data-driven offline model-based optimization, 2021a.  URL https://github.com/brandontrabucco/design-bench.
  * Trabucco et al. (2021b) B. Trabucco, A. Kumar, X. Geng, and S. Levine.  Conservative objective models for effective offline model-based optimization.  In _International Conference on Machine Learning_ , pages 10358–10368. PMLR, 2021b.
  * Walsh et al. (2012) A. Walsh, S. Chen, S.-H. Wei, and X.-G. Gong.  Kesterite thin-film solar cells: Advances in materials modelling of cu2znsns4.  _Advanced Energy Materials_ , 2(4):400–409, 2012\.
  * Woodley and Catlow (2008) S. M. Woodley and R. Catlow.  Crystal structure prediction from first principles.  _Nature materials_ , 7(12):937–946, 2008.
  * Xie et al. (2021) T. Xie, X. Fu, O.-E. Ganea, R. Barzilay, and T. Jaakkola.  Crystal diffusion variational autoencoder for periodic material generation.  _arXiv preprint arXiv:2110.06197_ , 2021.
  * Yamashita et al. (2016) T. Yamashita, H. Momida, and T. Oguchi.  Crystal structure predictions of naxc6o6 for sodium-ion batteries: First-principles calculations with an evolutionary algorithm.  _Electrochimica Acta_ , 195:1–8, 2016.
  * Yu et al. (2021) S. Yu, S. Ahn, L. Song, and J. Shin.  Roma: Robust model adaptation for offline model-based optimization.  _Advances in Neural Information Processing Systems_ , 34, 2021.

## Appendices

###  Appendix A Additional Ablation Study

Figure 4: Ablation study. The overestimation threshold, $\tau$, is the factor
controlling the level of conservatism imposed by LCOMs. The above plot shows
the performance of the crystal structures found by LCOMs by varying this
threshold for two sample compounds: SrS and MgO.

Since our method builds on existing conservative optimization algorithms, one
of the main hyper parameters of our method is the coefficient $\alpha$
controlling the strength of the conservatism regularizer. In the practical
instantiation of COMs (Trabucco et al., 2021b), this hyperparameter is
replaced by its Lagrangian dual Equation 6 in Trabucco et al. (2021b)), and
the corresponding hyperparameter in the practical algorithm is $\tau$, the
threshold of allowed over-estimation on adversarial examples (in our case,
adversarial latent vectors). A smaller $\tau$ enforces a stricter upper bound
on the allowed amount of distribution shift, whereas a larger $\tau$ does not
penalize distributional shift. As a result, energies of produced designs would
be close to the energy in the dataset when the coefficient $\tau$ is small,
but also get exploited when $\tau$ is too large. An intermediate value of
$\tau$ is expected to likely lead to the most favorable results.

As shown in Figure 4, an intermediate value of $\tau$ (e.g., 1.0 in this case)
leads to the best results as more gradient steps are performed to optimize the
crystal structure. As expected, while a very small $\tau=0.5$ plateaus in the
case of MgO, a very large value $\tau=5.0$ starts to get exploited for both
the sample compounds, MgO and SrS. These results align with our hypothesis.

###  Appendix B Details of Our Simulator

DFT simulators, underpinned by Density Functional Theory (DFT) Parr (1989),
serve as crucial computational tools for approximating the solution of the
Schrödinger equation for a given system of particles. Particularly, in our
study, these particles constitute the chemical structure of a crystal. By
providing an approximate solution of the underlying differential equation, DFT
simulators enable the calculation of system dynamics, including critical
properties such as total energy, and facilitate the simulation of system
relaxation to a stable, energy-minimal configuration.

DFT represents a class of computational algorithms rather than a single
operation method, which justifies the availability of multiple DFT simulators.
Examples of these simulators include licensed platforms like VASP, and open-
source ones like GPAW Mortensen et al. (2005), Enkovaara et al. (2010). We
leveraged the operational flexibility inherent to DFT in this work by using
GPAW to create random stable structures as initial points for the optimization
process. Its accessibility as an open-source tool, and ease of integration
with Python, made it the preferred choice.

However, GPAW does have limitations, most notably the absence of pseudo-
potentials for all chemical elements, essential for approximating the
potential experienced by valence electrons in atoms. This limitation hindered
the simulation of some structures used for evaluation, as highlighted by Cheng
et al. (2022). Consequently, our evaluation was limited to 25 out of the
original 29 compounds discussed in this prior work that informed our
evaluation procedure.

###  Appendix C Experiment Details

In this section, we detail the hyperparameters and configurations employed in
our experiments to facilitate reproducibility of the results. Please note that
for competing models, we rely on results reported in the original work instead
of replicating the experiments. For comprehensive information regarding these
models, we refer the reader to the work of Cheng et al. (2022).

Encoding and Decoding Following the method in Xie et al. (2021), we firstly
train an variational encoder to transform crystal structure to CD-VAE latent
space, with the same training protocol in Xie et al. (2021). We use a batch
size of 256 here when training the encoder.

Hyperparameters In LCOMs, we follow most of the hyper-parameters in the
implementation of COMs method Trabucco et al. (2021b). The number of epochs to
train the model $\widehat{E}_{\theta}(\phi(\mathbf{x},c),c)$ is 50 and the
number of gradient descent steps used in Equation 4.2 is 50. The number of
steps used to generate optimized results in latent space is $10$ for model
trained with OQMD dataset and $40$ for model trained with MatBench dataset.
Please note that this number is picked by evaluating the distance between
optimized groups and training dataset. The training batch size is 128 and the
learning rate for model training is $0.00003$. The model structure is followed
by the one in Trabucco et al. (2021b). The overestimation limit $\tau$ in
Equation 6 of Trabucco et al. (2021b) is picked as 1.0.