# StofNet: Super-resolution Time of Flight Network

###### Abstract

Time of Flight (ToF) is a prevalent depth sensing technology in the fields of
robotics, medical imaging, and non-destructive testing. Yet, ToF sensing faces
challenges from complex ambient conditions making an inverse modelling from
the sparse temporal information intractable. This paper highlights the
potential of modern super-resolution techniques to learn varying surroundings
for a reliable and accurate ToF detection. Unlike existing models, we tailor
an architecture for sub-sample precise semi-global signal localization by
combining super-resolution with an efficient residual contraction block to
balance between fine signal details and large scale contextual information. We
consolidate research on ToF by conducting a benchmark comparison against six
state-of-the-art methods for which we employ two publicly available datasets.
This includes the release of our SToF-Chirp dataset captured by an airborne
ultrasound transducer. Results showcase the superior performance of our
proposed StofNet in terms of precision, reliability and model complexity. Our
code is available at {https://github.com/hahnec/stofnet}.

Index Terms—  Super-resolution, Deep Learning, Neural, Audio, Time of Flight,
Time of Arrival, Acoustic, Ultrasound, Localization, Trilateration,
Multilateration

##  1 Introduction

Time of Flight (ToF) is a widely used non-contact depth sensing technology in
various fields, including geometric scene reconstruction [1, 2, 3, 4, 5, 6],
robotics [7, 8], autonomous vehicles [9], non-destructive testing (e.g.,
structural health monitoring) [10, 11, 12] and medical imaging [13]. For
example, the precise detection of wavefront arrival times enables geometric
localization of targets, which is also known as multilateration [14, 6, 8,
13]. The core concept of ToF relies on measuring the time it takes for a
signal carrier, such as light or sound, to travel from a source to a target
and back to the receiver, which is also known as Time of Arrival (ToA).

Despite its wide adoption, ToF sensing still faces challenges such as
mitigating the impact of ambient influences (e.g., target shape, temperature,
humidity), improving the accuracy and resolution of distance measurements and
handling multipath interference, which occurs when the signal reflects off
multiple surfaces or objects before reaching the sensor. Such ambient
conditions can be attributed to the physical emission and reflection
properties as well as the surrounding environment (target position, material
stiffness, signal penetration, temperature, etc.). Modeling these influences
as inverse problems rapidly turns out to be cumbersome. Instead, the rise of
deep neural networks offers new opportunities to address these demands and
further enhance the capabilities and robustness of ToF sensing systems. Can
recent advances in super-resolution imaging be transferred and equally applied
to localization in the temporal domain? What kind of modifications are
necessary to ensure enhanced operation considering established methods?

ToF determination breaks down to a localization problem across the temporal
domain of a captured one-dimensional (1-D) signal and has been thoroughly
studied in the past. Traditional ToA detection includes thresholding [10, 11,
15] and cross-correlation methods [16, 6]. The vastly growing field of deep
learning offers many considerable architectures for that task [17, 18, 19,
12]. Often, these models contain a contracting and expanding path following
the style of a U-Net [17]. This trend can also be seen in the acoustic signal
processing domain [19, 12] and particularly in the audio super-resolution
field [18]. The general motivation for such U-Net shape is to establish global
context among several local samples that do not expose local adjacency. While
this comes in handy for semantic segmentation, spatial or temporal downscaling
carries the risk of losing the capability to recover fine signal details.
Instead of a bottleneck contraction, advances in efficient image super-
resolution networks omit spatial downscaling while expanding the feature
channels to learn upsampling from a trailing feature channel shuffle operation
[20, 21]. On the contrary, these super-resolution networks lack in building
large-scale contextual information, which becomes essential among extended
signal regions, e.g. to avoid multiple detections for a single target.

In this paper, we customize a network based on [20, 21] for temporal 1-D
localization while balancing the above requirements to improve distance
accuracy and the reliability of ToF measurements. We compare our architecture
against traditional ToF detection techniques [10, 15], audio super-resolution
networks [18, 19, 12] and 1-D equivalents of models used in computer vision
[20, 21].

$C$$F$L1$F\times S$L2$F\times S$L3+$F$L4$F$L5+$F$L12$F$L13+$F$L14+$F$L15$R$

Fig. 1: The StofNet architecture consists of multiple convolutional layers and
employs residual skip connections for improved feature representation. The
model takes as input a 1-D signal with $C$ channels for optional feature
concatenation. The initial layer applies a 1-D convolution with $F$ filters
and a kernel size of 9 followed by a ReLu. Layer 2 and 3 represent the
bottleneck block consisting of 1-D convolutions, ReLus, and up- or
downsampling blocks, respectively. The subsequent layers (4 to 14) consist of
1-D convolutions with $F$ filters and a kernel size of 7. Residual connections
are added after the 3rd, 5th, 7th, 9th, 11th, 12th, and 13th layers, whereas a
ReLu follows convolutions without residual connections. The second last layer
uses a 1-D convolution with $F$ filters and a kernel size of 3, followed by an
element-wise addition with the third layer residual output. Finally, the
output is obtained by applying a 1-D convolution with the specified upsampling
factor $R$ and a kernel size of 3, followed by a sample shuffling operation.

##  2 Method

For an input signal $\mathbf{X}\in\mathbb{R}^{N\times C}$ of $N$ samples and
$C$ features, we wish our network $f(\mathbf{X}):~{}\mathbb{R}^{N\times
C}\mapsto\mathbb{R}^{NR}$ to learn how to discover local peaks using trained
weights. We develop our deep neural network along the lines of super-
resolution techniques in the computer vision field by using stacked residual
blocks followed by upscaling through learned feature channel shuffling [20,
21]. As we seek a localization method with refinement ability even in large
contextual regions, we introduce a single bottleneck block for semi-global
context recognition at an early stage of our proposed architecture. Contrary
to [12, 18, 19], our model omits a global bottleneck contraction to retain
crucial resolution information in the temporal domain. A schematic
visualization of our network is depicted in Fig. 1. While our proposed model
accepts varying input lengths with parameterizable block dimensions, we follow
the image super-resolution field with $F=64$ feature channels and an
upsampling factor of $R=4$.

Loss. Similar to previous works [22, 23], we construct a binary mask
$\mathbf{y}\in\\{0,1\\}^{NR}$ where each 1 represents a ground truth sample
position. We convolve this binary mask using a 1-D Gaussian kernel
$\mathbf{g}_{\sigma}\in\left[0,1\right)^{7}$ with $\sigma=1$ to learn the
prediction of smooth local maxima distributions. The loss function reads,

$\displaystyle\mathcal{L}(\mathbf{X},\mathbf{y})=\lVert
f(\mathbf{X})-\lambda_{0}(\mathbf{g}_{\sigma}\circledast\mathbf{y})\rVert_{2}^{2}+\lambda_{1}\lVert
f(\mathbf{X})\rVert_{1}$ (1)

where $\circledast$ denotes the convolution operator, $\lambda_{0}$ amplifies
the labels and $\lambda_{1}$ scales an $L_{1}$ regularization term that
prevents $f(\cdot)$ from predicting an excessive amount of false positives.

Training. Each model is trained using an AdamW optimizer with batch size 4,
weight decay of $1\mathrm{e}{-8}$ and a start learning rate of
$5\mathrm{e}{-4}$ using cosine annealing as a scheduler. The regularization
scales are set to
$\lambda_{0}=(\max(\mathbf{g}_{\sigma}\circledast\mathbf{y})/20)^{-1}$ and
$\lambda_{1}=1\mathrm{e}{-2}$. We train a network for a maximum of 80 epochs
with early stopping using a delta of $1\mathrm{e}{-6}$ and a patience of 5.

For better generalization, we prevent over-fitting through augmentation of the
input signals. This involves an amplitude normalization such that
$\mathbf{X}\in\left[-1,1\right]^{N\times C}$ and the addition of noise for the
training frames with a signal-to-noise ratio of $30~{}\text{dB}$. For learning
single echo detection, we apply a randomized positional frame cropping to
$\nicefrac{{3}}{{4}}\,N$ samples in the temporal domain with subsequent
padding for consistent frame length $N$ as required by baseline models [18,
19].

Inference. After training, a single target localization is accomplished via
${\operatorname{arg\,max}}\\{f(\mathbf{X})\\}$. For multiple localizations, we
employ thresholding in conjunction with non-maximum suppression. The ideal
threshold largely depends on the network and dataset and is determined using
G-means analysis of the ROC curve.

##  3 Experiments

For the evaluation of state-of-the-art ToF detection methods, we conduct a
benchmark comparison on two different datasets with suitable metrics that are
explained hereafter. The code of the following analyses, including backbone
models, can be found online111Access to our code repository at
https://github.com/hahnec/stofnet.

Datasets. For numerical evaluation, we employ the PALA dataset222Access to
PALA dataset at doi.org/10.5281/zenodo.4343435 published by Heiles et al. [24]
as it provides in silico A-scans coming with ground truth data from third-
party software that meets industrial standards. We customize settings for PALA
by interpolating In-phase and Quadrature (IQ) signals by factor 20 and use
$\lambda_{1}=1$ for training. Besides, we employ and release the SToF-Chirp
dataset333Access to our SToF-Chirp dataset at doi.org/10.5281/zenodo.8255670.
The dataset was collected using our prototype depicted in Fig. 2, which
carries a TDK DK-CH101 Smartsonic airborne ultrasound transducer operating at
175 kHz. We demodulate the IQ samples with an interpolation factor of 10. The
dataset comprises three target classes, namely: wooden ruler, metal stick, and
paper chart. The targets were moved along the $z$ direction with a total of
nine distance measurements captured for each class. We collected 100 frames
per class and distance, with 20 measurements designated for testing purposes
and 80 measurements for training. The determination of the true distance for
each measurement is based on the average of the manufacturer’s estimates from
200 separate frames (distinct from test and train sets), providing a reliable
reference for evaluating localization accuracy.

Fig. 2: Our prototype uses the DK-CH101 airborne ultrasound sensor (blue
circle) to acquire the SToF-Chirp dataset3.

Metrics. For localization assessment, the minimum Root Mean Squared Error
(RMSE) is calculated from the estimated position and the nearest ground truth
position. Here, only RMSEs less than a tolerance $\tau$ are considered true
positives and contribute to the total RMSE of all frames. In cases where the
RMSE distance exceeds the tolerance, an estimated position is considered a
false positive. As a result, ground truth locations without an estimate within
the $\tau$ neighborhood are treated as false negatives. The Jaccard Index is
employed to evaluate the ToA detection reliability, which takes into account
both true positives and false negatives, offering a robust measure of each
algorithm’s performance. In addition, we provide the single batch inference
time and parameter number for each model.

Baselines. Our benchmark comparison comprises a gradient-based peak detection
[15], the audio super-resolution network by Kuleshov et al. [18], the ToA-
tailored architectures by Zonzini et al. [12], the SincNet [19] and a 1-D
equivalent of the much celebrated EDSR net [21]. Zonzini et al. [12] propose a
large and a small architecture for embedded devices. While we employ the large
model for the PALA dataset, the number of its contracting downsample layers
cannot handle the frame lengths of the SToF-Chirp dataset, for which we
utilize the small Zonzini model in this study. All models are trained and
executed on an Nvidia GeForce RTX 3090 Ti.

Fig. 3: Single echo detection plot from the SToF-Chirp test set with a
magnified view of the onset target position on the right. Fig. 4: Multiple
echo detection plot from the PALA test set with a magnified view of the peak
target position on the right.

Results. For visual inspection, Figs. 3 and 4 depict exemplary channel data
frames with localization results. The magnified portion in Fig. 4 reveals that
ESPCN [20], as a standalone super-resolution network, tends to detect multiple
ToAs for a single target reflection. This observation serves as the motivation
for introducing the semi-global context block in our StofNet architecture,
facilitating our network to effectively investigate broader surroundings from
condensed temporal information.

Table 1: Single echo detection results from SToF-Chirp data3. Units are provided in brackets. The RMSE tolerance is $\tau=1$. Model |  RMSE [Sample] | Jaccard [%] |  Weights [k#] |  Time [ms]
---|---|---|---|---
Gradient [15] |  0.607 $\pm$ 0.000 | $0.185$ |  0 |  3.590
Zonzini [12] |  0.495 $\pm$ 0.161 | $0.370$ |  134 |  2.159
ESPCN [20] |  0.494 $\pm$ 0.295 | $21.111$ |  6 |  1.671
SincNet [19] |  0.359 $\pm$ 0.282 | $70.741$ |  329 |  2.879
EDSR [21] |  0.329 $\pm$ 0.297 | $51.481$ |  210 |  2.811
Kuleshov [18] |  0.229 $\pm$ 0.139 | $32.037$ |  57416 |  3.267
Ours |  0.145 $\pm$ 0.212 | $83.519$ |  645 |  3.273

Tables 1 and 2 list numerical results of the selected metrics for the SToF-
Chirp and PALA dataset, respectively.

In relation to RMSE, our StofNet model achieves the overall most precise
localization in comparison to the other models. As for the Jaccard index, our
model attains the highest score, indicating strong alignment in accurately
detecting valid echoes within the defined accuracy threshold $\tau$. All
models exhibited relatively similar computational performance with inference
time in the milliseconds range.

Table 2: Multiple echo detection results from PALA data2. Units are provided in brackets. The RMSE tolerance is $\tau=1$. Model |  RMSE [Sample] | Jaccard [%] |  Weights [k#] |  Time [ms]
---|---|---|---|---
Gradient [15] |  0.532 $\pm$ 0.242 | $2.637$ |  0 |  12.45
Zonzini [12] |  0.505 $\pm$ 0.290 | $0.151$ |  1259 |  2.236
ESPCN [20] |  0.405 $\pm$ 0.082 | $38.161$ |  6 |  0.914
SincNet [19] |  0.389 $\pm$ 0.074 | $75.678$ |  329 |  3.814
EDSR [21] |  0.287 $\pm$ 0.090 | $72.176$ |  210 |  3.695
Kuleshov [18] |  0.462 $\pm$ 0.145 | $34.858$ |  516631 |  4.345
Ours |  0.253 $\pm$ 0.085 | $79.811$ |  645 |  3.430

The super-resolution network proposed by Kuleshov et al. [18] is a heavyweight
with respect to its 517M trainable weights. This network may serve well for
spectral extrapolation as required by processing of audible signals, however,
the imbalanced ratio of RMSE and Jaccard versus parameter number in Table 2
makes it an unsuitable choice for ToF detection especially on embedded
devices. One may note the difference in model weights for Kuleshov [18] and
Zonzini [12] in Tables 2 and 1, which are due to varying input lengths of the
different datasets.

The comparison demonstrates the superiority of our proposed StofNet model in
terms of precision (RMSE), reliability (Jaccard) and complexity (time and
weights) underpinning the effectiveness of our approach and its potential for
practical metrology applications.

##  4 Conclusion

This study explored the potentials and limitations of modern super-resolution
techniques for accurate, fast and reliable Time of Flight (ToF) sensing
through a benchmark comparison against state-of-the-art networks using
publicly available datasets. Overall, our proposed 1-D localization network
exhibits best performance compared to state-of-the-art attempts owing to its
ability to build sufficient global context while maintaining high temporal
fidelity. Moving forward, future research can focus on expanding the
application of super-resolution networks to address specific challenges in
deep learning such as domain gap across sensors and environmental conditions.
In summary, our study has shed light on the potential of modern super-
resolution networks and their significant role in improving ToF sensing. The
findings presented here provide valuable insights for researchers and
practitioners in the field, paving the way for enhanced depth sensing
capabilities and applications in various domains.

## References

  * [1] Fengqiang Li, Pablo Ruiz, Oliver Cossairt, and Aggelos K Katsaggelos,  “Multi-frame super-resolution for time-of-flight imaging,”  in ICASSP 2019-2019 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP). IEEE, 2019, pp. 2327–2331.
  * [2] Jesper Haahr Christensen, Sascha Hornauer, and Stella X. Yu,  “Batvision: Learning to see 3d spatial layout with two ears,”  in 2020 IEEE International Conference on Robotics and Automation (ICRA), 2020, pp. 1581–1587.
  * [3] Miguel Heredia Conde,  “Raw data processing for practical time-of-flight super-resolution,”  in ICASSP 2021 - 2021 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), 2021, pp. 1650–1654.
  * [4] Kranti Kumar Parida, Siddharth Srivastava, and Gaurav Sharma,  “Beyond image to depth: Improving depth prediction using echoes,”  in Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2021, pp. 8268–8277.
  * [5] Go Irie, Takashi Shibata, and Akisato Kimura,  “Co-attention-guided bilinear model for echo-based depth estimation,”  in ICASSP 2022 - 2022 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), 2022, pp. 4648–4652.
  * [6] Axel Berg, Mark O’Connor, Kalle Åström, and Magnus Oskarsson,  “Extending GCC-PHAT using Shift Equivariant Neural Networks,”  in Proc. Interspeech 2022, 2022, pp. 1791–1795.
  * [7] Ethan Tracy and Navinda Kottege,  “CatChatter: Acoustic perception for mobile robots,”  IEEE Robotics and Automation Letters, vol. 6, no. 4, pp. 7209–7216, 2021.
  * [8] Christopher Hahne,  “3-dimensional sonic phase-invariant echo localization,”  in 2023 IEEE International Conference on Robotics and Automation (ICRA), 2023, pp. 4121–4127.
  * [9] Liangtian Wan, Yuchen Sun, Lu Sun, Zhaolong Ning, and Joel JPC Rodrigues,  “Deep learning based autonomous vehicle super resolution doa estimation for safety driving,”  IEEE Transactions on Intelligent Transportation Systems, vol. 22, no. 7, pp. 4301–4315, 2020.
  * [10] Jochen H Kurz, Christian U Grosse, and Hans-Wolf Reinhardt,  “Strategies for reliable automatic onset time picking of acoustic emissions and of ultrasound signals in concrete,”  Ultrasonics, vol. 43, no. 7, pp. 538–546, 2005.
  * [11] Matthew R Pearson, Mark Eaton, Carol Featherston, Rhys Pullin, and Karen Holford,  “Improved acoustic emission source location during fatigue and impact events in metallic and composite structures,”  Structural Health Monitoring, vol. 16, no. 4, pp. 382–399, 2017\.
  * [12] Federica Zonzini, Denis Bogomolov, Tanush Dhamija, Nicola Testoni, Luca De Marchi, and Alessandro Marzani,  “Deep learning approaches for robust time of arrival estimation in acoustic emission monitoring,”  Sensors, vol. 22, no. 3, 2022.
  * [13] Christopher Hahne and Raphael Sznitman,  “Geometric ultrasound localization microscopy,”  in Medical Image Computing and Computer-Assisted Intervention–MICCAI 2023: 26th International Conference, Vancouver, Canada, October 8-12, 2023, Proceedings. Springer, 2023.
  * [14] Kalle Åström, Martin Larsson, Gabrielle Flood, and Magnus Oskarsson,  “Extension of time-difference-of-arrival self calibration solutions using robust multilateration,”  in 2021 29th European Signal Processing Conference (EUSIPCO), 2021, pp. 870–874.
  * [15] Christopher Hahne,  “Multimodal exponentially modified gaussian oscillators,”  in 2022 IEEE International Ultrasonics Symposium (IUS), 2022, pp. 1–4.
  * [16] Silvia Adrián-Martínez, Manuel Bou-Cabo, Ivan Felis, Carlos D Llorens, Juan Antonio Martínez-Mora, María Saldaña, and Miguel Ardid,  “Acoustic signal detection through the cross-correlation method in experiments with different signal to noise ratio and reverberation conditions,”  in Ad-hoc Networks and Wireless: ADHOC-NOW 2014 International Workshops, ETSD, MARSS, MWaoN, SecAN, SSPA, and WiSARN, Benidorm, Spain, June 22–27, 2014, Revised Selected Papers 13. Springer, 2015, pp. 66–79.
  * [17] Olaf Ronneberger, Philipp Fischer, and Thomas Brox,  “U-net: Convolutional networks for biomedical image segmentation,”  in Medical Image Computing and Computer-Assisted Intervention–MICCAI 2015: 18th International Conference, Munich, Germany, October 5-9, 2015, Proceedings, Part III 18. Springer, 2015, pp. 234–241.
  * [18] Volodymyr Kuleshov, S Zayd Enam, and Stefano Ermon,  “Audio super-resolution using neural nets,”  in ICLR (Workshop Track), 2017.
  * [19] Mirco Ravanelli and Yoshua Bengio,  “Speaker recognition from raw waveform with sincnet,”  in 2018 IEEE spoken language technology workshop (SLT). IEEE, 2018, pp. 1021–1028.
  * [20] Wenzhe Shi, Jose Caballero, Ferenc Huszár, Johannes Totz, Andrew P Aitken, Rob Bishop, Daniel Rueckert, and Zehan Wang,  “Real-time single image and video super-resolution using an efficient sub-pixel convolutional neural network,”  in Proceedings of the IEEE conference on computer vision and pattern recognition, 2016, pp. 1874–1883.
  * [21] Bee Lim, Sanghyun Son, Heewon Kim, Seungjun Nah, and Kyoung Mu Lee,  “Enhanced deep residual networks for single image super-resolution,”  in Proceedings of the IEEE conference on computer vision and pattern recognition workshops, 2017, pp. 136–144.
  * [22] Ruud JG van Sloun, Oren Solomon, Matthew Bruce, Zin Z Khaing, Hessel Wijkstra, Yonina C Eldar, and Massimo Mischi,  “Super-resolution ultrasound localization microscopy through deep learning,”  IEEE transactions on medical imaging, vol. 40, no. 3, pp. 829–839, 2020.
  * [23] Xin Liu, Tianyang Zhou, Mengyang Lu, Yi Yang, Qiong He, and Jianwen Luo,  “Deep learning for ultrasound localization microscopy,”  IEEE transactions on medical imaging, vol. 39, no. 10, pp. 3064–3078, 2020.
  * [24] Baptiste Heiles, Arthur Chavignon, Vincent Hingot, Pauline Lopez, Eliott Teston, and Olivier Couture,  “Performance benchmarking of microbubble-localization algorithms for ultrasound localization microscopy,”  Nature Biomedical Engineering, vol. 6, no. 5, pp. 605–616, 2022\.


11institutetext: L3S Research Center - Leibniz University Hannover, Germany
11email: {tai, szerr, ntoutsi<EMAIL_ADDRESS>
# Data augmentation for dealing with low sampling rates in NILM

Tai Le Quy  Sergej Zerr  Eirini Ntoutsi  Wolfgang Nejdl

###### Abstract

Data have an important role in evaluating the performance of NILM algorithms.
The best performance of NILM algorithms is achieved with high-quality
evaluation data. However, many existing real-world data sets come with a low
sampling quality, and often with gaps, lacking data for some recording
periods. As a result, in such data, NILM algorithms can hardly recognize
devices and estimate their power consumption properly. An important step
towards improving the performance of these energy disaggregation methods is to
improve the quality of the data sets. In this paper, we carry out experiments
using several methods to increase the sampling rate of low sampling rate data.
Our results show that augmentation of low-frequency data can support the
considered NILM algorithms in estimating appliances’ consumption with a higher
F-score measurement.

###### Keywords:

data augmentation energy disaggregation NILM temporal disaggregation

##  1 Introduction

Non-Intrusive Load Monitoring (NILM) or energy disaggregation is a state-of-
the-art technology to disaggregate and estimate the power consumption of
individual appliances from the aggregated signal in households or companies.
It is preferred over intrusive approaches due to bounded costs compared to the
monitoring of each device separately. From the first work of Hart [3], a
number of NILM techniques have been proposed [15]. NILM research requires
large amounts of high-quality aggregated data [2, 6]. It has been shown that
NILM algorithms work efficiently with a sampling rate at 1Hz because at this
granularity the data can provide several electricity measurements such as
active and reactive power [2, 6]. However, many data sets exist with lower
sampling rates, like Pecan Street Sample [4] at 1/60 Hz or HES [14] at 1/120
and 1/600 Hz. Recognizing devices in such data sets is extremely difficult for
the existing NILM algorithms. Besides the low sampling rates, the algorithms
have often to deal with data incompleteness like missing aggregated signal for
some time points or even, for certain time periods. Furthermore, recently,
deep learning techniques have shown their potential for NILM, e.g., [5, 9],
however, they require huge amounts of training data. As a result, constructing
higher sampling rate data even from a small amount of the lower samples is an
important direction in order to improve the performance of NILM algorithms.

In this research, we investigate and propose data augmentation methods in
order to construct high sampling rate data from the low sampling rate signal
that can be used for NILM. To this end, we apply interpolation techniques such
as Denton-Cholette method for temporal disaggregation, stepwise method and
Cubic spline interpolation on power consumption data of two selected data sets
(ECO [2] and iAWE [1]) in order to generate a higher sampling rate data. We
then report the results of the performance of two NILM methods (Weiss’s
algorithm [13] and Parson’s algorithm [10]).

The rest of the paper is structured as follows: In Section 2, we overview the
related work. In Section 3, we describe the augmentation techniques for
generating high sampling rate data from the lower samples. In Section 4, we
present our experimental results as the effect of augmentation on the
performance of two well-known NILM algorithms. Conclusions and outlook are
finally presented in Section 5.

##  2 Related work

The first NILM method has been introduced by Hart [3] to extract device
consumption profiles called signatures. Following this work, different methods
have been proposed relying on state analysis (e.g., tracking on/off operation
by using real power and reactive power), utilize learning methods or different
data granularities [6, 15]. Parson et al. [10] proposed a semi-supervised
approach using factorial HMMs which was evaluated on data sets at a sampling
rate 1/60 Hz. Weiss et al. [13] proposed a supervised approach to extract
switching events and find the best match in a signature database by using real
and reactive power information with granularity 1 Hz.

Unfortunately, the publicly available data sets do not always come with the
desired granularity. In fact, data signal collection is a very costly process,
in terms of the time required to collect reasonably large data sets but also
due to other reasons like privacy. In Table 1, we survey the data sets that
are often used for NILM evaluation. As depicted, different datasets come with
different granularities and for some cases, the sampling rate is too low.

Table 1: Datasets for NILM Dataset |  Institution |  Sampling rate
---|---|---
REDD (2011) |  MIT |  1 Hz and 15 kHz
BLUED (2012) |  CMU |  12 kHz
Smart* (2012) |  UMass |  1 Hz
Sample (2013) |  Pecan Street |  1/60 Hz
HES (2013) |  DECC DEFRA |  1/120 Hz and 1/600 Hz
iAWE (2013) |  III Delhi |  1 Hz
ECO (2013) |  ETH |  1 Hz
UK-DALE (2014) |  Imperial College |  1/6 Hz and 16 kHz

In contrast to existing works focussing on better NILM methods to cope with
real-world electricity data, we follow an approach by proposing data
augmentation for NILM in order to generate high granularity samples from low
granularity ones. Except for the low sample case, such an augmentation can
also help with data incompleteness. Our method independent approach can work
with a variety of different NILM algorithms because it is applied at the data
level.

Data augmentation is summarizing techniques for dealing with data sparsity by
deducing missing values using historical information or third-party
information sources. Recently, it has gained a lot of attention as a way to
cope with the huge data demands of complex learning models such as deep neural
networks, and it has been proven that for many architectures and different
applications, it improves the robustness and the generalization power of the
models. Data augmentation techniques increase the volume of the data by
generating new instances from the existing instances. In the image domain, for
example, this is achieved by applying label-preserving transformations like
rotation and illumination [7] to teach a machine model and achieve higher
accuracy. Similarly, in the audio domain augmentation is achieved by applying
transformations like adding background noise or changing the pitching level
[11]. In the case of signal data being undersampled, augmentation is related
to interpolation, which is used traditionally for e.g., time series and
trajectory data [8]. In trajectory data, for example, interpolation between
successive moving object positions is used in order to simulate the continuous
nature of the actual movement.

To our best knowledge, data augmentation for NILM has not been investigated
thus far, however, our results show that it comprises a promising approach.

##  3 Generating high sampling rate data from low rate samples

In our work, the augmentation methods are used to generate augmented time
series data in between every timestamp $t_{i}$ and $t_{i+1}$.

###  3.1 Stepwise interpolation

Augmented data in between every time-stamp ti and $t_{i+1}$ is generated by
dividing the time-gap of $t_{i}$ and $t_{i+1}$ into $k$ parts, then we create
data for each part by the following formula:

$data(part_{j})=data(t_{i})+\frac{data(t_{i+1})-data(t_{i})}{k}*j,j=1..k$ (1)

###  3.2 Cubic Spline interpolation

This is a form of interpolation using a special type of piece-wise polynomial
called a spline. In our work, we use a spline function in Numerical Recipes in
C111http://www.aip.de/groups/soe/local/numres/bookcpdf/c3-3.pdf, that returns
interpolated values between data at time-stamp $t_{i}$ and $t_{i+1}$. The
distance to the next interpolated point is calculated by a parameter $t$ with
range from 0 for time-stamp $t_{i}$ to 1 for timestamp $t_{i+1}$.

###  3.3 Denton-Cholette interpolation

This is a method for temporal disaggregation, it can disaggregate low-
frequency time series data with or without high-frequency indicator series.
This method is primarily concerned with movement preservation. Augmented data
that is similar to the indicator is generated by considering the most common
case in the indicator. Mathematical techniques are used to distribute low-
frequency data in high-frequency series when the indicator has a poor quality.
In this paper, we use an implementation of this method in a R package by
Christoph et al. [12].

###  3.4 Device interpolation

In this method, the changes in the power consumption of devices with high
sampling rates are used as indicators to estimate the augmented values for
aggregated signals. For every time points in between two time-stamps $t_{i}$
and $t_{i+1}$, if the power consumption of any appliance changes beyond a
threshold value. In our experiment, we set the threshold value from 5 to 10W.
We will increase or decrease the aggregated value to the power level of time-
stamp $t_{i+1}$. Otherwise we set the same value for every time series data in
between $t_{i}$ and $t_{i+1}$.

Fig. 1 visualizes 10 minutes of aggregated data on 01/06/2012 in house 2 of
ECO data set with different augmentation methods.

Figure 1: A sample augmented data with different methods

##  4 Experiments

We evaluate the effect of data augmentation for NILM on two real data sets
(Section 4.1), the evaluation is presented in Section 4.3.

###  4.1 Datasets

The iAWE dataset [1] contains ambient, water and electricity data from a
single house in Delhi over a period of 73 days (May-August 2013), 10 jPlugs
are used to measure 10 appliances with multiple parameters including voltage,
current, phase and frequency. It consists of almost 15M measurements at a
sampling rate 1 Hz. The average of missing values in this data set is 31%.

The ECO dataset [2] contains more than 650M measurements from 45 smart plugs
in 6 households in Switzerland over a period of 8 months (June 2012-January
2013). The data set is collected at a sampling rate 1 Hz with the aggregated
consumption data and appliances’ consumption. Each of the measurements
contains information on power consumption, voltage, current and the phase
shift between voltage and current. For each household, there are 6-10 devices
connected to smart plugs to measure power consumption in order to obtain
ground truth data for analysis. The average of missing values in smart meter
data is 0.8%. The proportion of appliance consumption measured by plug meters
in households varies from 10% to 80%.

These datasets contain both real and reactive power measurements which are a
prerequisite for several NILM algorithms, e.g., Weiss et al. [13]. For iAWE
dataset, we use the whole observation period (73 days). We evaluate the
performance of NILM algorithms on five appliances: two air conditionals,
fridge, laptop-PC and TV. In the ECO dataset, each house has a different
observation period. Therefore, we select 30 days for each house as follows:
House 1, 2: from 01/06/2012 to 30/06/2012, House 3: from 06/12/2012 to
04/01/2013, House 4, 5, 6: from 01/07/2012 to 30/07/2012. We evaluate the
performance on 15 devices: dryer, freezer, fridge, water kettle, PC, laptop,
dishwasher, lamp, microwave, stove, stereo system, TV, coffee machine,
entertainment system and fountain. For evaluation, we took the first 2/3 of
the recording period from each data set for training and used the remaining
segment for testing. Resulting in 20 and 50 days training data for ECO and
iAWE datasets respectively.

###  4.2 Experimental setup

We evaluate the performance, in terms of the F-score on the estimation of
power consumption, of two well-known NILM algorithms, namely Parson [10] and
Weiss [13]. For the experiments, we used the NILM-Eval framework [2].

To evaluate the effect of data augmentation for creating high sampling data
from low sampling ones, we first down sample the original data sets to the
1/60 Hz granularity by keeping the first second data for each minute. For
Parson algorithm, because this algorithm is designed to work with data at a
sampling rate 1/60 Hz, we down-sample the data to the granularity of 1/600 Hz
(10 minutes) before evaluating the performance of this algorithm. We refer to
the undersampled data set as “undersampled” and to the original data set as
“original”. For the experiment of Weiss algorithm, we generate the augmented
data at granularity 1Hz from the “undersampled” data at 1/60Hz. Parson
algorithms use the data at the sampling rate 1/60Hz which is reconstructed
from the “undersampled” data at granularity 1/600Hz. We also do several
experiments to find the best parameter for our augmentation methods. In the
stepwise method, we carried out the experiments with different values of k (k
= 2; 3; 4; 6; 10) and we found that with k = 4 our stepwise interpolation
shows the best results across the available datasets. We then compare the
performance of the NILM algorithms for the different data sets: “original”,
“undersampled” and several proposed augmented methods. Our goal is to
investigate, how augmentation helps the undersampled data set to reach a
performance close to the original high sampling data set.

###  4.3 Performance evaluation

Table 2 describes the Mean Squared Error (MSE) and Root Mean Squared Error
(RMSE) that shows the difference between the augmented data generated by our
augmentation methods and the original data on two data sets. We calculate the
average values of MSE and RMSE of 6 households in the ECO data set. These
measurements are calculated on the data set at granularity 1Hz.

Table 2: MSE and RMSE of augmentation methods Augmentation methods | ECO dataset | iAWE dataset
---|---|---
MSE | RMSE | MSE | RMSE
Stepwise |  85,374.7 |  286.7 |  137,600.8 |  370.9
Cubic |  247,425.5 |  473.7 |  100,464.6 |  316.9
Denton-Cholette |  34,921.3 |  181.9 |  125,423.4 |  354.1
Device |  61,021.8 |  238.6 |  330,394.8 |  574.8

Comparison of the NILM Algorithm performance on original data: During our
experiments, we noticed the difference in performance and properties of the
used NILM techniques. Whilst Parson algorithm was showing good performance for
smaller period data sets, it was lagging behind for data sets comprising large
time periods. The possible explanation is that Parson uses pre-trained models,
whilst Weiss needs more time to train, but in the long run, it is able to
identify more devices, especially outperforming at the “outlier-devices” such
as Fountain and Dryer. For similar reasons, the smaller iAWE dataset was
dominated by Parson and the larger ECO dataset by Weiss. In this context,
while Parson showed high precision for a small number of devices, Weiss was
able to identify more devices with moderate precision.

Down-sampling effect: Separately for each method, we measured the effect of
the down-sampling. Down-sampling has shown a strong effect on the Weiss
method, while it had a light effect on Parson algorithm. This can be explained
by the fact that the Weiss method requires information about active power and
reactive power as shown for the ECO data set in Figure 2

Figure 2: Performance of NILM algorithms on ECO dataset

Comparison of the augmentation methods: The performance of the augmentation
methods (measured by F-score value) on iAWE data set are shown in Table 3 and
Figure 3 3, the results on ECO data set are presented in Table 4 and Figure 2.
Comparing the different augmentation techniques, device interpolation shows
the best performance for most of the devices in both the Weiss and Parson
algorithm. The stepwise method follows, whilst the cubic interpolation shows
the worse performance among the remaining augmentation techniques. One of the
reasons, as we noticed from the augmented signal results, is the introduced
smoothness of the produced time series, which can hinder event detection and
disturb the inferring power consumption of the appliances.

Figure 3: Performance of NILM algorithms on iAWE dataset Table 3: Performance of NILM on augmentation methods on iAWE dataset Augmentation methods | Weiss’s algorithm | Parson’s algorithm
---|---|---
Avg. | Dev. | Avg. | Dev.
Original |  0.736 |  0.136 |  0.866 |  0.007
Undersampled |  0.206 |  0.178 |  0.877 |  0.048
Stepwise |  0.416 |  0.217 |  0.793 |  0.099
Cubic |  0.145 |  0.136 |  0.775 |  0.085
Denton-Cholette |  0.305 |  0.227 |  0.752 |  0.136
Device |  0.275 |  0.175 |  0.718 |  0.114
Table 4: Performance of NILM on augmentation methods on ECO dataset Augmentation methods | Weiss’s algorithm | Parson’s algorithm
---|---|---
Avg. | Dev. | Avg. | Dev.
Original |  0.604 |  0.257 |  0.359 |  0.331
Undersampled |  0.269 |  0.240 |  0.27 |  0.274
Stepwise |  0.332 |  0.278 |  0.328 |  0.307
Cubic |  0.176 |  0.236 |  0.277 |  0.323
Denton-Cholette |  0.128 |  0.234 |  0.325 |  0.305
Device |  0.467 |  0.331 |  0.331 |  0.310

Performance with respect to device type: We categorize devices into five
groups based on their function and characteristics: Cooling devices (Freezer,
Fridge), Cooking devices (Coffee Machine, Microwave, Water kettle, Stove),
Entertainment (TV, Stereo), Computer (PC, Laptop), Lighting device (Lamp). A
summary of results for groups of appliances in ECO data set is presented in
Table 5 and Table 6. The device interpolation method showed the best
performance among the four methods, although in some cases it failed for
Lighting devices, as such appliances consume a low amount of electric power.
Another observation is that Parson’s algorithm can work well with data
generated by Device interpolation and Stepwise methods for cooling devices,
entertainment and computer, as these appliances do not have significant
changes in power consumption over time.

Table 5: Performance of Weiss’s alg. with appliance’s groups ECO dataset: Augmentation methods |  Cooling |  Cooking |  Ent. |  Computer |  Lighting
---|---|---|---|---|---
Original |  0.711 |  0.613 |  0.585 |  0.357 |  0.281
Undersampled |  0.318 |  0.101 |  0.396 |  0.208 |  0.174
Stepwise |  0.482 |  0.231 |  0.312 |  0.045 |  0.157
Cubic |  0.273 |  0.124 |  0.246 |  0.021 |  0.078
Denton-Cholette |  0.239 |  0.037 |  0.181 |  0.014 |  0.060
Device |  0.744 |  0.300 |  0.416 |  0.043 |  0.193
Table 6: Performance of Parson’s alg. with appliance’s groups on ECO dataset Augmentation methods |  Cooling |  Cooking |  Ent. |  Computer |  Lighting
---|---|---|---|---|---
Original |  0.667 |  0.042 |  0.410 |  0.454 |  0.03
Undersampled |  0.504 |  0.021 |  0.326 |  0.306 |  0.028
Stepwise |  0.596 |  0.048 |  0.306 |  0.485 |  0.037
Cubic |  0.482 |  0.038 |  0.319 |  0.373 |  0.023
Denton-Cholette |  0.574 |  0.047 |  0.319 |  0.472 |  0.036
Device |  0.620 |  0.049 |  0.309 |  0.456 |  0.037

##  5 Conlusion and outlook

In this work, we presented an attempt to assist NILM by improving the quality
of low sampling rate energy consumption data sets through data augmentation by
using several interpolation techniques. Our approach works at the data level
and therefore it is method-independent and applicable for a variety of
different NILM algorithms. Augmentation was also shown to be helpful for data-
intensive machine learning models like deep neural networks which recently
have been successfully used also for NILM [5, 9]. Our results show that data
augmentation is applicable for increasing the sampling rates of a data set. We
believe that this is a promising direction for NILM and further research
should be carried, in parallel to the development of new, more sophisticated
methods.

In this preliminary work, we adapted simple augmentation techniques, which
nevertheless yield improvements over the non-augmented low-sample data sets.
As part of our ongoing work, we are investigating more sophisticated
augmentation approaches which take into account the consumption profile of the
household as well as the profiles of individual devices and do not require the
high sampling rate data of appliances. Deep learning technology is also a
potential approach that can learn from devices’ using patterns in order to
construct the aggregated data. Such “informed”-augmentation approaches are
expected to yield better augmented data and therefore, better predictions.

## Acknowledgements

The work of the first author is financially supported by the Ministry of
Education and Training, Vietnam, within the program “Training doctoral degrees
for lecturers at universities and colleges from 2010 to 2020” (Project 911).

## References

  * [1] Batra, N., Gulati, M., Singh, A., Srivastava, M.B.: It’s different: Insights into home energy consumption in india. In: Proceedings of the 5th ACM Workshop on Embedded Systems For Energy-Efficient Buildings. pp. 1–8 (2013)
  * [2] Beckel, C., Kleiminger, W., Cicchetti, R., Staake, T., Santini, S.: The eco data set and the performance of non-intrusive load monitoring algorithms. In: Proceedings of the 1st ACM conference on embedded systems for energy-efficient buildings. pp. 80–89 (2014)
  * [3] Hart, G.W.: Nonintrusive appliance load monitoring. Proceedings of the IEEE 80(12), 1870–1891 (1992)
  * [4] Holcomb, C.: Pecan street inc.: A test-bed for nilm. In: International Workshop on Non-Intrusive Load Monitoring, Pittsburgh, PA, USA (2012)
  * [5] Kelly, J., Knottenbelt, W.: Neural nilm: Deep neural networks applied to energy disaggregation. In: Proceedings of the 2nd ACM international conference on embedded systems for energy-efficient built environments. pp. 55–64 (2015)
  * [6] Klemenjak, C., Goldsborough, P.: Non-intrusive load monitoring: A review and outlook. Informatik 2016 (2016)
  * [7] Krizhevsky, A., Sutskever, I., Hinton, G.E.: Imagenet classification with deep convolutional neural networks. Advances in neural information processing systems 25, 1097–1105 (2012)
  * [8] Macedo, J., Vangenot, C., Othman, W., Pelekis, N., Frentzos, E., Kuijpers, B., Ntoutsi, I., Spaccapietra, S., Theodoridis, Y.: Trajectory data models. In: Mobility, Data Mining and Privacy, pp. 123–150. Springer (2008)
  * [9] Mauch, L., Yang, B.: A new approach for supervised power disaggregation by using a deep recurrent lstm network. In: 2015 IEEE Global Conference on Signal and Information Processing (GlobalSIP). pp. 63–67. IEEE (2015)
  * [10] Parson, O., Ghosh, S., Weal, M., Rogers, A.: Non-intrusive load monitoring using prior models of general appliance types. In: Proceedings of the AAAI Conference on Artificial Intelligence. vol. 26 (2012)
  * [11] Salamon, J., Bello, J.P.: Deep convolutional neural networks and data augmentation for environmental sound classification. IEEE Signal Processing Letters 24(3), 279–283 (2017)
  * [12] Sax, C., Steiner, P.: Temporal disaggregation of time series. A peer-reviewed, open-access publication of the R Foundation for Statistical Computing p. 80
  * [13] Weiss, M., Helfenstein, A., Mattern, F., Staake, T.: Leveraging smart meter data to recognize home appliances. In: 2012 IEEE International Conference on Pervasive Computing and Communications. pp. 190–197. IEEE (2012)
  * [14] Zimmermann, J.P., Evans, M., Griggs, J., King, N., Harding, L., Roberts, P., Evans, C.: Household electricity survey: A study of domestic electrical product usage. Intertek Testing & Certification Ltd pp. 213–214 (2012)
  * [15] Zoha, A., Gluhak, A., Imran, M.A., Rajasegarar, S.: Non-intrusive load monitoring approaches for disaggregated energy sensing: A survey. Sensors 12(12), 16838–16866 (2012)


# Rotationally Invariant Latent Distances for Uncertainty Estimation of
Relaxed Energy Predictions by Graph Neural Network Potentials

Joseph Musielewicz  [  Janice Lan  Matt Uyttendaele  [<EMAIL_ADDRESS>John R.
Kitchin  [<EMAIL_ADDRESS>
###### Abstract

Graph neural networks (GNNs) have been shown to be astonishingly capable
models for molecular property prediction, particularly as surrogates for
expensive density functional theory calculations of relaxed energy for novel
material discovery. However, one limitation of GNNs in this context is the
lack of useful uncertainty prediction methods, as this is critical to the
material discovery pipeline. In this work, we show that uncertainty
quantification for relaxed energy calculations is more complex than
uncertainty quantification for other kinds of molecular property prediction,
due to the effect that structure optimizations have on the error distribution.
We propose that distribution-free techniques are more useful tools for
assessing calibration, recalibrating, and developing uncertainty prediction
methods for GNNs performing relaxed energy calculations. We also develop a
relaxed energy task for evaluating uncertainty methods for equivariant GNNs,
based on distribution-free recalibration and using the Open Catalyst Project
dataset. We benchmark a set of popular uncertainty prediction methods on this
task, and show that latent distance methods, with our novel improvements, are
the most well-calibrated and economical approach for relaxed energy
calculations. Finally, we demonstrate that our latent space distance method
produces results which align with our expectations on a clustering example,
and on specific equation of state and adsorbate coverage examples from outside
the training dataset.

##  1 Introduction

To keep up with growing energy demands, it is necessary to search for novel
catalyst materials to enable more efficient storage of renewable sources of
energy 1, 2, 3, 4. Computational material discovery is crucial to this
process, as it enables less expensive screening of an enormous space of
possible catalyst materials than physical experiments. Faster and more
accurate computational material discovery methods will be required to meet our
society’s renewable energy needs in the face of a rapidly changing climate.

Graph neural networks are state of the art in accelerating computational
material discovery pipelines with machine learning potentials. Machine
learning potentials work as surrogate models trained to approximate
computationally expensive density functional theory (DFT) calculations of
energy and forces on atomistic structures. This task is referred to as
structure to energy and forces (S2EF). These energy and force calculations are
used to iteratively perform geometric optimizations of atomic positions
(referred to in this work as "relaxations"), to minimize their energy. These
relaxed structure and relaxed energy calculations are what enable high-
throughput predictions of catalyst performance in the real world. For a given
catalyst-adsorbate system, the global minimum relaxed energy (adsorption
energy) directly correlates with the reactivity and selectivity of reaction
pathways on that catalyst surface 5, 6, 7, 8, 9, 10.

In recent years, graph neural networks (GNN) s have made tremendous strides in
replacing DFT codes with an inexpensive, accurate alternative 11, 12, 13, 14.
Thanks to methods like AdsorbML, GNN s can speed up adsorption energy
calculations alone at the cost of accuracy, or in tandem with DFT at the cost
of speed 15. The recent OCP Demo (https://open-catalyst.metademolab.com) is a
publicly available tool where GNN s are used to calculate these adsorption
energies without any expensive DFT calculations. However, a major limitation
of current GNN s is their lack of uncertainty estimates for relaxed energy
predictions. Ideally, users of these methods would know when it is safe to
trust the GNN predictions, and when additional DFT calculations are warranted.
In this work, we specifically examine methods of uncertainty quantification
(UQ) estimates for GNN predictions on this relaxed structure to relaxed energy
(RS2RE) task.

##  2 Background

###  2.1 AdsorbML

AdsorbML 15 is a method to calculate adsorption energy using machine learning
potentials. In order to find the global minimum relaxed energy for a specified
surface and adsorbate, this method places the adsorbate in many different
starting configurations, relaxes each configuration, and returns the minimum
of all the relaxed energies. Traditionally, this would be done using an ab
initio method such as DFT, but DFT is very costly and this approach is
infeasibly expensive. AdsorbML uses GNN s as a surrogate for DFT to perform
the relaxations instead, requiring only a solitary DFT single point
calculation for the relaxed structure to verify the relaxed energy. The
required number of expensive DFT calculations is further reduced by using GNN
s to filter out all but the few most promising candidates. This method
provides an adjustable spectrum of trade-offs between accuracy and efficiency,
with one balanced option finding an equivalent or better adsorption energy
87.36% of the time while reducing DFT compute by more than a factor of 2000.
Ideally, no DFT would be required, but even state of the art GNN s are
unreliable energy predictors when out of domain, and using them alone drops
the success rate to 56%. Without DFT, such as in the OCP demo, we need
uncertainty metrics, so users know when to trust the results of these models.

###  2.2 Graph Neural Networks

This work focuses on quantifying uncertainty prediction methods for
EquiformerV2 14, a GNN model architecture for molecular property prediction.
We choose EquiformerV2 because it is the current state of the art in molecular
property prediction for catalyst materials, according to the Open Catalyst
Project (OCP) leaderboard2. We also compare it to Gemnet-OC 11, another high
performing GNN on the leaderboard. Both of these models are used in the OCP
Demo, to run the AdsorbML algorithm and predict minimum relaxed energies
without the use of expensive DFT calculations15.

###  2.3 Uncertainty Quantification

Many prior studies have examined the application of UQ techniques to machine
learning potentials and molecular property prediction. Most UQ validation
metrics seek to measure some description of the calibration of an uncertainty
prediction method. The most popular UQ validation metrics are miscalibration
area, Spearman’s rank correlation coefficient, and the negative log likelihood
of the errors given the uncertainties16, 17, 18, 19, 20, 21, 22. These metrics
each face drawbacks, especially when used for overall validation of
uncertainty estimates. Specifically, the miscalibration area and negative log
likelihood (NLL) rely on an assumption of normally distributed errors, which
adds fragility to the metrics 23. Additionally, the NLL and Spearman’s rank
correlation coefficient validation metrics are not well suited to the absolute
validation of uncertainty estimates, which limits them to comparing models and
does not allow for the overall validation of the uncertainty estimates 24.
Notably these metrics each target different properties of uncertainty
estimates, and are not consistently in agreement about which uncertainty
estimator performs best, even within a single study 25.

Calibration is the primary UQ validation metric for making direct comparisons
between uncertainty prediction methods. Prior work by Pernot, and Levi et al.
show distribution free methods of measuring local and global miscalibration:
the CI(Var(Z)) test and the error-based calibration plot 23, 26. These
approaches can be more effective in describing the performance of predicted
uncertainties of surrogate models when the expected error distribution cannot
be assumed to be Gaussian. Rasmussen et al. finds that the error-based
calibration plot generally serves as the best overall method of validating of
uncertainty estimates for machine learning potentials for chemistry 24. Error-
based calibration measures can also be used to recalibrate uncertainty
predictions, allowing a variety of uncertainty quantification methods to be
recalibrated and compared.

##  3 Methods

###  3.1 Uncertainty Quantification Methods for GNNs

Figure 1: Diagram of the components of the latent distance uncertainty
prediction method. It details the setup procedure and the inference procedure.
The setup procedure must be performed once per model training set. The
inference procedure is first used on the calibration set during the setup
procedure, and then every time inference is performed afterwards to predict
the model uncertainty. Note that the calibration steps at the end of the setup
procedure and inference procedure can be applied to any of the uncertainty
prediction methods.

In this work we examine four common methods of uncertainty prediction on a
pre-trained GNN: ensembles, latent space distances, mean variance estimation,
and sequence regression models. The most widely used of these are ensemble
methods, where a set of similar surrogate models are trained on similar sets
of data to perform the same task, and the variance between their predictions
is used to calculate the uncertainty of the model. We train ensembles of GNN s
on the S2EF task to calculate energies of adsorbate-catalyst structures, but
we are interested in the performance of predicting the uncertainty of
EquiformerV2 on the relaxed structures. We construct three different ensembles
for testing this approach, which we refer to as architecture (11 members),
bootstrap (10 members), and parameter ensembles (6 members). Prior work has
shown that diversity between members is typically the most important factor in
developing a well-calibrated and expressive ensemble for uncertainty
prediction 27, 28, 29, 30. The architecture ensemble contains a variety of GNN
model architectures, while the parameter and bootstrap ensembles contain only
EquiformerV2 models, but vary the number of parameters, and the composition of
training data respectively. More details on ensemble construction can be found
in the supporting information. Because we are specifically interested in the
uncertainty of relaxed energy predictions, which requires a sequence of prior
energy/force predictions during the relaxation, we hypothesize that taking the
mean or the maximum of the predicted variances over each step (here referred
to as a frame) of the trajectory might contain additional information which
better models the uncertainty. For each ensemble composition, we test this
theory by computing the uncertainty using the variance at the first frame,
last frame, mean over all the frames, and max over over all the frames.

Another proven uncertainty method is the use of latent space distances,
commonly referred to as conformal prediction methods 16, 17, 24. In this
approach, we extract some latent representation of each training point from
the GNN, and create an index of these points to compute the L2-norm of the
distance from any new test point to the training points. In practice, using
libraries such as FAISS, the computational cost of this method at inference
time is the lowest of any of the uncertainty methods we test 31. This method
also comes with an upfront cost of creating the index of training point
representations, but this is generally significantly less expensive than the
upfront costs associated with training an additional model or models. Because
this approach produces a distance of arbitrary scale, it is necessary to
recalibrate on some calibration set to produce a meaningful uncertainty
estimate. Figure 1 shows the pipeline we use to setup, predict uncertainties,
and calibrate the uncertainty predictions for this latent distance method.
Prior work has shown the latent distance method to be effective for models
which preserve some rotationally invariant information through the model
architecture, such as GemNet-OC. However EquiformerV2 improves upon the
accuracy of GemNet-OC by preserving rotational equivariance all the way
through the model 14, 11. We expect this rotational equivariance will
contribute undesirable noise to the L2 distance between latent
representations, so we test this hypothesis by comparing the performance of
the full EquiformerV2 latent representation, and the latent representation of
its single rotationally invariant channel. We also compare these methods to
the performance of using the latent space representation of other GNN models,
such as GemNet-OC, trained on the same data, but used to predict the
uncertainty of the same EquiformerV2 model. Additionally, for each latent
distance approach, we test a novel strategy of computing the latent distance
on a per-atom basis, and then taking the mean, max, or sum over these
distances. We compare this method to the more common approach of computing the
latent distance for an entire frame by taking the mean of latent
representations over all the atoms. More details on how all of the latent
representations were extracted, and how the distances were computed can be
found in the supporting information.

The final categories of methods we test are mean variance estimation (MVE) and
sequence regression models 22, 20, 25, 32. We implement these methods in
similar ways. For MVE methods we append an output-head, or an ensemble of
output-heads to the EquiformerV2 architecture, and fine-tune this new output
head on the calibration set to predict the residual of the energy prediction
of the larger model. For the sequence regression model, we extract the same
latent space representations used in the latent distance method, and train a
sequence regression transformer architecture to predict the residual of the
energy prediction on the calibration set 33. A significant distinction between
these approaches is that the sequence regressor takes the latent
representations of each frame of the trajectory as input, in sequence, with
the hypothesis that some additional information about the uncertainty of the
GNN model might be contained within its latent representations along the
trajectory. More details on the implementation of both of these approaches can
be found in the supporting information.

###  3.2 Uncertainty Validation Metrics

We suspect that the NLL, Spearman’s rank correlation coefficient, and
miscalibration area, will not serve as ideal metrics for UQ validation on this
task, due to the non Gaussian distribution of the RS2RE task. We test this by
evaluating each of the uncertainty estimation methods using each of the
commonly used UQ validation metrics. We anticipate significant disagreement
among many of the UQ validation metrics, especially those not based on
distribution-free methods. Further examination of the error distribution of
the RS2RE task can be found in the supporting information.

Prior work by Pernot suggests the use of a distribution-free method to test
whether an uncertainty method is calibrated 23. The CI(Var(Z)) test uses the
BCa boostrap method to compute a confidence interval of Var(Z) on a set of
errors and uncertainties, without making any assumptions about the
distribution 34, 35. If 1 lies within the confidence interval, the uncertainty
method is considered calibrated, because the Z values for a calibrated
uncertainty metric are expected to have a variance of 1. When applied to the
calibration set, this metric is useful for characterizing whether error-based
recalibration is effective in achieving globally well calibrated
uncertainties. On the test set (after recalibration on the calibration set)
estimated uncertainties are typically no longer perfectly calibrated. However
the closer the confidence interval of the variance is to the ideal (Var(Z)=1),
the better the calibrated uncertainty estimates remain globally calibrated
under distribution shift. Unless otherwise stated, we report the CI(Var(Z))
metric on the test set.

Rasmussen et al. expands on this to suggest using an error-based calibration
plot to quantify local calibration, distribution-free 26, 24. The error-based
calibration plot is predicated on the expected relationship between the root
mean variance (RMV) and root mean square error (RMSE) being one-to-one.

$RMSE=\sqrt{1/N_{bin}\sum_{i}\epsilon^{2}_{i}}\quad
RMV=\sqrt{1/N_{bin}\sum_{i}\sigma^{2}_{i}}\quad\frac{RMSE}{RMV}\approx 1$ (1)

We sort the test points by their predicted uncertainty, and then bin them into
20 bins. We compute the RMV of each bin, and the RMSE of each bin, using the
BCa bootstrap method to compute a 95% confidence interval for the RMSE. Then
we fit a line through the points to test for calibration, ideally the fitted
line should have a high $R^{2}$ correlation with the points, and be as close
as possible to the parity line. We can identify problems with local
miscalibration where the parity line does not lie within the binned RMSE
confidence intervals.

###  3.3 Recalibration and Evaluation

We compare all of the uncertainty methods benchmarked in this work after
recalibration on a calibration set. In this case we use Open Catalyst 2020
dataset (OC20) in-domain validation set, and relax each structure
(approximately 25,000 structures) to the same relaxation criterion as OC20,
with the publicly available EquiformerV2 31M parameter checkpoint. We then
compute a single point calculation using the Vienna Ab initio Simulation
Package (VASP) on the final frame to serve as the ground truth energy value
for the RS2RE task 36, 37, 38, 39. This set of relaxed energy predictions
serves as the calibration set, and we repeat this process with the out-of-
domain-both validation set to serve as the test set for this task. To
recalibrate each uncertainty method, we use same the approach as the error-
based calibration plot, to find the line of best fit through the binned
RMSE/RMV points of the calibration set. We then simply recalibrate the
uncertainty using the formula for the line of best fit:

$\sigma_{recalibrated}=\text{slope}_{fit}*\sigma_{uncalibrated}+\text{intercept}_{fit}$
(2)

We propose a small modification to the characterization of global calibration
of uncertainties recalibrated with the error-based method. For this task, in
addition to checking the global calibration CI(Var(Z)) test, we also select
the most effective uncertainty metric by measuring the $R^{2}$ correlation of
the binned RMSE/RMV ratio with the parity line. By recalibrating the
uncertainty method on the calibration set and then testing on the test set, we
can make direct comparisons between different uncertainty methods on the basis
of their $R^{2}$ correlation with the parity line.

##  4 Results

###  4.1 Error Distribution and Uncertainty Validation Metrics

(a)

(b)

Figure 2:  Violin and parity plot for a latent distance uncertainty
prediction. After calibrating uncertainty predictions with error based
calibration, they can be interpreted as predictions of the dispersion of the
residual for that inference calculation. In 2(a) we see the measured error
distributions for five uncertainty prediction bins. In 2(b) we see contour
lines tracing five different quantiles of measured error, rolling along the
predicted uncertainty.

To interpret the meaning of these recalibrated uncertainties, we look at the
distribution of errors compared to the corresponding uncertainties. In Figure
2(a) we plot the error distributions of points in the test set binned by their
estimated uncertainties, after error-based recalibration. When we recalibrate
uncertainty predictions, we fit the ratio of the RMSE to the RMV of these bins
to the parity line. We see that the distribution of errors is generally
skewed, with the bulk of the distribution falling below the predicted RMV
(uncertainty) of the bin. This is what we expect to see for uncertainty
predictions, since calibrated uncertainties should describe the dispersion of
the distribution that the error is sampled from, most of the sampled errors
should lie below the predicted uncertainty. We also see that the measured RMSE
generally corresponds to the predicted RMV of each bin. However the lowest
uncertainty bin is underconfident, and the highest uncertainty bin is
overconfident, which is a common trend across many uncertainty estimators we
examine in this work.

In the parity plot in Figure 2(b) we see the effect of the error-based
recalibration on the distribution of RS2RE errors. The contour lines trace the
n-th quantile of the measured errors. The contour lines curve somewhat
relative to the parity line, this is partly explained by the underconfidence
towards the lower end of the plot, and overconfidence towards the highest
uncertainties. However this is also an effect of the error distribution being
non-Gaussian. We do not necessarily expect to see a perfect linear
relationship between the contour lines and the parity line, unless the errors
follow a Gaussian distribution, since calibration is only based on a linear
relationship between RMSE/RMV. Even with a non-Gaussian distribution, the
uncertainties can still be well-calibrated with respect to the errors, as long
as the RMSE/RMV relationship follows the parity line. The goal of a well-
calibrated uncertainty estimate is to predict the dispersion characteristic of
the error distribution at that point, regardless of the distribution from
which it is sampled. If, for instance, the errors more closely follow a
Laplace distribution, we would expect to see contour lines similar to the ones
we observe here.

Table 1: Best performers for each of the UQ validation metrics. Many of the popular UQ validation metrics assume normally distributed errors, such as miscalibration area ($A_{mis}$), before and after error-based recalibration, and the negative log likelihood (NLL). Other metrics target properties other than calibration, such as Spearman’s rank correlation coefficient ($\rho$) and the area under the receiver operating characteristic curve (AUROC). These methods tend to be inconsistent with one another. The error based calibration (Parity $R^{2}$) measure and the CI(Var(Z)) test both measure calibration without a distribution assumption, and agree that the GNOC distance method performs best. Method | Parity $R^{2}$ | $A_{mis,u}$ | $A_{mis,r}$ | $NLL$ | $\rho$ | AUROC | CI(Var(Z))
---|---|---|---|---|---|---|---
GNOC distance | 0.952 | 0.365 | 0.178 | 0.279 | 0.468 | 0.755 | [1.13, 1.79]
Architecture ens. | 0.830 | 0.050 | 0.156 | 0.411 | 0.453 | 0.746 | [1.58, 2.19]
Residual model | 0.331 | 0.174 | 0.102 | 1.170 | 0.360 | 0.695 | [3.38, 4.47]
MVE head | 0.831 | 0.340 | 0.114 | 0.251 | 0.526 | 0.785 | [1.59, 2.40]
Bootstrap ens. | 0.870 | 0.248 | 0.121 | 0.303 | 0.585 | 0.817 | [1.58, 2.48]

In Table 1 we compare the best performing UQ methods according to several
commonly used UQ validation metrics. We observe that top performance in one
metric does not correspond to top performance in another, except with respect
to the distribution-free validation metrics, Parity $R^{2}$ and the
CI(Var(Z)), which tend to agree. Unsurprisingly, Spearman’s rank correlation
coefficient and AUROC also agree, as both of these metrics are aimed at
classification performance, but do not provide much insight into overall
calibration. These results indicate that the distribution-free validation
metrics output significantly different results from other common UQ validation
metrics, and that other such metrics cannot be relied upon to agree with one
another. Due to the fragility of assumptions of Gaussian distributed errors
for the RS2RE task, and the recommendations from previous work on this topic,
we are inclined to believe that these distribution-free metrics are more
reliable than other options for generally assessing and comparing calibration
on the RS2RE task 23, 26, 24.

###  4.2 Benchmarking Uncertainty Quantification

Table 2: Distribution free calibration metrics for the best performing candidate uncertainty estimation method for each of the categories we benchmark. Note that all of the uncertainty estimates are able to be globally calibrated on the calibration set. The latent space distance method achieves highest local calibration, and is closest to well calibrated globally on the test set, according to the CI(Var(Z)) test. Method | Parity $R^{2}$ | Fit $R^{2}$ | Slope | Int. | Cal. set | Test set
---|---|---|---|---|---|---
|  |  |  |  | CI(Var(Z)) | CI(Var(Z))
Latent distance | 0.952 | 0.967 | 1.022 | 0.026 | [0.84, 1.35]* | [1.13, 1.79]
Ensemble | 0.905 | 0.940 | 0.946 | 0.070 | [0.83, 1.57]* | [1.42, 2.57]
Residual model | 0.331 | 0.955 | 1.200 | 0.127 | [0.89, 1.50]* | [3.38, 4.47]
MVE head | 0.849 | 0.964 | 1.187 | 0.041 | [0.86, 2.09]* | [1.66, 2.46]

Of the four methods we benchmark on the RS2RE task, the latent distance method
is the best performer according to the distribution-free uncertainty
quantification techniques, as seen in Table 2. The CI(Var(Z)) global
calibration test shows that all four UQ methods seem to be well calibrated by
error based recalibration on the calibration set. However on the test set, all
of the UQ some degree of miscalibration due to distribution shift. The latent
distance approach demonstrates the best test set performance, according to the
CI(Var(Z)), since the confidence interval is closest to 1. It also achieves
the best calibration according to the parity $R^{2}$ metric on the out-of-
domain test set. Note that all four models are capable of a comparable fit (if
this test set were used for recalibration) as evidenced by the similar scores
on the fit $R^{2}$.

Figure 3: Calibration plots for the best performing iterations of the four
categories of methods discussed in this work. The diagram above each plot
shows how each method estimates the variance of the relaxed energy trajectory
predicted by the graph neural network.

We observe superior local calibration for the top performing latent distance
method, as compared to the other UQ methods in Figure 3. We see that the
latent distance method generally retains good local calibration throughout,
which we can see by the error bars of nearly every confidence interval
intersect the parity line. The latent distance method also maintains the best
overall calibration, with the line of best fit on the the test set only being
offset slightly from the parity line. The ensemble method, which is the next
best method class, sees more significant local deviations, displaying
intermittent overconfidence below 0.3 eV, then systematic overconfidence
between 0.3 eV and 1 eV. The ensemble method does maintain a good overall
calibration fit on the test set. The remaining two methods both result in more
significant deviations on the test set. The residual model is systematically
over-confident on the entire test set, and this shows in its overall parity
$R^{2}$. The MVE method only achieves good local calibration at the lowest
uncertainty estimates, before becoming systematically overconfident above 0.2
eV.

(a) Best ensemble contour plot
(Architecture ensemble, mean over trajectory)

(b) Best latent distance contour plot
(GemNet-OC, sum of per-atom distances)

Figure 4: Contour parity plots for the best two method classes, predicting on
the test set after recalibration on the calibration set. Uncertainties
predicted by the best performing ensemble method (4(a)) and the best
performing latent distance method (4(b)). The quantile curves of the latent
distance method are generally montonic, which aligns with the better local
calibration performance. The ensemble method uncertainties are shifted right
by recalibration, and shows a significant drop in the uncertainty estimates
after 0.2 eV, and an abrupt increase after 0.3 eV.

We examine the contour parity plots of the two top performing UQ methods in
Figure 4. The ensemble method shows a shift which lower bounds the predicted
uncertainties resulting from recalibration, while the recalibrated
uncertainties of the latent distance method result in no such lower bound.
This lower bounding limits the effectiveness of this approach at low
uncertainty estimates, which may be a critical region of interest depending on
the use case. We also observe a sharp drop and increase in the measured error
of the ensemble method from 0.2 eV to 0.3 eV, which seems to be responsible
for the intermittent local miscalibration in that range. Overall, we see the
error distribution predicted by the latent distance method fans out in a more
linear fashion. The parity line appears to be a rough upper bound for the
errors for the latent distance method, which is desirable in a well-calibrated
uncertainty method.

###  4.3 Comparing Distance Methods

The latent space representation sampled from EquiformerV2 after all graph
convolutional interactions have been performed is inherently equivariant with
respect to rotations. This is a valuable property for training GNN machine
learning potentials to predict properties of molecular systems, but it may
contribute noise to distance measures in the latent space. Since the latent
space distance should measure the similarity of atomic structure
representations, it should ideally be invariant to the rotations of input
structures. Since the same structure can be rotated to a numerically different
latent space representation, many more samples must be present in training set
for the latent space distance to be meaningful. Therefore we compare distances
measured using the full latent space representation, which includes all
spherical harmonic channels, to distances measured using only the rotationally
invariant (degree 0 spherical harmonic channels) latent space representations.

(a) Equivariant Latent Space

(b) Invariant Latent Space

Figure 5: Plots of UMAP dimensionality reduction performed on the equivariant
(all channels) and invariant (m=0, l=0) latent space representations sampled
from Equiformer V2 without edge alignment. Latent space representations were
sampled for a subset of the training set containing specific elements. We see
that the different elements represented are clearly clustered in both plots,
but that there is significantly more noise in the clustering of similar
elements in the equivariant latent space, while the invariant latent space
clusters are much denser and less noisy.

In Figure 5 we observe the effects of these two different latent space
sampling methods on Equiformer V2 using UMAP, a dimensionality reduction
package 40. We plot the first two dimensions of the UMAP-reduced latent space,
sampled from a subset of systems containing only three metal species, and four
adsorbate species. We observe clear clustering of the latent space
representations based on the atomic species, but with significantly more noise
present in the fully equivariant latent representation. By contrast, the
invariant latent representation produces much denser clusters, which are more
clearly differentiated by atomic species. The key features we would expect to
see in a good atomic representation are: clear distinctions between metals and
adsorbate atoms, and more similarity between group 10 metals like Pd and Pt
than between those and group 11 metals like Cu. Both latent representations
produce definitive segregation of adsorbate species from metal species,
however the equivariant representation generates a cloud of overlapping metal
species which includes Cu, Pd, Pt. The invariant representation creates
distinct clusters for Cu, and clusters with significant overlap between Pt and
Pd. This seems to indicate that the rotationally equivariant information,
which is preserved by the structure of Equiformer V2, will appear as noise to
simple low dimensional comparisons of latent representations. This explanation
also aligns with the superior performance of distance measures within
rotationally invariant latent space.

Table 3:  Distribution free calibration metrics for each latent distance method tested. The latent representations correspond to model and method that was used to extract the latent representation for each atom. The distance column refers to the approach used to compute the distance between systems, which are described in the methods section. A * indicates the method is globally calibrated according the CI(Var(Z)) test.  Latent Rep. | Distance | Parity $R^{2}$ | Slope | Int. | Cal. set | Test set
---|---|---|---|---|---|---
|  |  |  |  | CI(Var(Z)) | CI(Var(Z))
EqV2 equiv. | atom max | -0.078 | 0.761 | 0.182 | [0.90, 1.16]* | [1.61, 2.02]
EqV2 equiv. | atom mean | 0.813 | 0.974 | 0.060 | [1.03, 1.74] | [1.51, 5.45]
EqV2 equiv. | sys. mean | -14.709 | -0.236 | 0.547 | [1.24, 2.74] | [1.90, 3.93]
EqV2 equiv. | atom sum | 0.672 | 0.711 | 0.159 | [0.89, 1.27]* | [1.52, 2.78]
EqV2 inv. l=0 (nodes) | atom max | 0.842 | 0.983 | 0.054 | [0.86, 1.24]* | [1.22, 1.70]
EqV2 inv. l=0 (nodes) | atom mean | 0.866 | 0.974 | -0.037 | [1.16, 3.10] | [1.04, 5.44]
EqV2 inv. l=0 (nodes) | sys. mean | -0.628 | 1.540 | -0.061 | [0.91, 1.14]* | [1.69, 2.05]
EqV2 inv. l=0 (nodes) | atom sum | 0.826 | 0.750 | 0.081 | [0.90, 1.47]* | [0.93, 1.59]*
EqV2 inv. l=0 (edges) | atom max | 0.548 | 1.124 | 0.057 | [0.94, 1.23]* | [1.61, 2.00]
EqV2 inv. l=0 (edges) | atom mean | 0.910 | 1.014 | -0.002 | [0.94, 1.65]* | [1.34, 6.37]
EqV2 inv. l=0 (edges) | sys. mean | 0.808 | 1.107 | 0.021 | [0.91, 1.14]* | [1.25, 1.55]
EqV2 inv. l=0 (edges) | atom sum | 0.882 | 0.814 | 0.094 | [0.88, 1.25]* | [1.14, 1.59]
EqV2 inv. l=4 (edges) | atom max | 0.480 | 1.011 | 0.082 | [0.95, 1.22]* | [1.54, 1.88]
EqV2 inv. l=4 (edges) | atom mean | 0.906 | 0.998 | -0.020 | [1.00, 2.33] | [1.10, 5.16]
EqV2 inv. l=4 (edges) | sys. mean | 0.824 | 1.121 | 0.006 | [0.90, 1.15]* | [1.19, 1.48]
EqV2 inv. l=4 (edges) | atom sum | 0.828 | 0.740 | 0.097 | [0.90, 1.38]* | [1.00, 1.63]*
GNOC | atom max | 0.924 | 0.951 | 0.064 | [0.81, 1.17]* | [1.23, 1.73]
GNOC | atom mean | 0.932 | 1.108 | 0.006 | [0.84, 1.55]* | [1.17, 2.18]
GNOC | sys. mean | 0.818 | 1.126 | 0.050 | [0.87, 1.16]* | [1.56, 2.04]
GNOC | atom sum | 0.952 | 1.022 | 0.026 | [0.84, 1.35]* | [1.13, 1.79]

Table 3 shows that the choice of latent space representation seems to be the
most important factor in choosing a good distance metric for predicting
uncertainty. Depending on the approach to sampling an invariant latent space,
the performance of using the latent space distance can vary wildly. These
results suggest the second most important factor is choosing a per-atom
distance metric, instead of taking the mean of the representations over all
the atoms. Using a distance approach for latent representations pooled over
the entire can work passably well in some cases, such as for GemNet-OC, but
generally a per-atom approach performs better. Finally, selecting the right
type of per atom distance measure appears to make some difference, but seems
to be less significant if a good latent representation is chosen. Taking the
sum or the mean of the per-atom distances tends to work well for most methods
shown here.

We also consider how edge alignment impacts the latent space representation
during graph convolutions within Equiformer V2. Atoms in Equiformer V2 are
embedded as nodes, with edge-degree embeddings, during equivariant graph
attention blocks, the edge embeddings are rotated to align with the y axis and
processed by the Transformer block. The edge aligned embeddings should be
invariant with respect to rotation along all of the m=0 spherical harmonic
channels, so we test this by sampling these edge embeddings and pooling them
to obtain rotationally invariant per-atom edge embeddings for all of the
channels through [l=4,m=0]. We compare this approach to the previously
mentioned approach of sampling the rotationally invariant [l=0, m=0] node
embedding as a latent representation. We also compare to using the [l=0, m=0]
edge aligned embedding as an invariant latent representation, and to using the
GemNet-OC node embeddings as a latent representation, since GemNet-OC
inherently maintains rotational invariance through its latent space. We find
that GemNet-OC performs best of these of these approaches, with both of the
edge aligned embedding approaches coming in close second, ahead of simply
taking the Equiformer V2 [l=0, m=0] node embeddings. There is little
discernible difference between the two edge aligned embedding approaches for
Equiformer V2, perhaps indicating that a significant portion of the relevant
information is contained in the [l=0, m=0] channel. The fact that both edge
embedding approaches outperform the node embedding approach also may suggest
that there is some advantage to using the edge embeddings for distance
metrics, that goes beyond eliminating noise from rotational equivariance. All
of these approaches far exceed the performance of using all the equivariant
spherical harmonic node embeddings.

###  4.4 Interpretable Examples

(a) Cu bulk

(b) TiO2 bulk

(c) Al slab

Figure 6:  Equation of state plots created using GemNet-OC, which describe the
energy (relative to the plotted minimum) as a function of interatomic distance
for three systems. 6(a) and 6(b) are both bulk systems, which are not
represented in the OC20 training data set. The predicted uncertainties for
these systems are quite large, as expected. Note that the uncertainty scale is
much larger 6(b) which is an oxide bulk, and no oxide compositions of any kind
are represented in the training set. 6(c) is an aluminum slab, which is
similar to the slabs used to create the systems in the OC20 training data set,
although this type of change to the interatomic distance of this slab is not
represented in the training data. The predicted uncertainty for this system is
significantly smaller, but grows towards the extremes, which also aligns with
our expectations.

For a more interpretable test case, we consider several specific examples of
adsorbate and materials systems. The goal here is to check whether the latent
space distance behaves as we would expect an uncertainty prediction method to
behave. In Figure 6 we look at how the predicted uncertainty (calibrated on
relaxed adsorbate-slab energies) scales for predicting equation of states of
two bulk materials, and an aluminum slab system. The predicted uncertainty is
compared to the relative change in energy from the local minimum as the
interatomic distance widens and contracts between the atoms in the metal
structure. The underlying GemNet-OC model was trained exclusively on
combinations of adsorbates and metal slabs present in the OC20 training
dataset, and the latent distance uncertainty predictions were calibrated on
similar systems. Since bulk materials, and oxide materials, are not present in
the training data, we expect to see high uncertainty predictions for the Cu
bulk example, and even higher for the Ti oxide bulk. This behavior is exactly
what we observe: the fluctuations in energy due to changes in the lattice
constant are completely dwarfed by the uncertainty predictions for both bulk
systems. We also expect the uncertainty predictions to be comparatively more
certain about the slab system, which is what we observe. The slab should be
more similar than the bulk systems to the training data it has seen, despite
this Al slab not being present in the training data. We would expect the
training data to contain many examples of stable slab structures, with near
optimum interatomic distances. Therefore we expect the GNN to be fairly
certain about its energy predictions near the stable interatomic distance for
a given slab, relative to the rate of change in energy with respect to
distance it predicts near the local minimum. We also observe increased
uncertainty further from the local minimum energy for interatomic distance,
which is another promising sign.

Figure 7:  Example using GemNet-OC with the latent distance uncertainty to
predict platinum oxygen coverage. One oxygen adsorbate resting on the platinum
slab should be very similar to examples seen in the training set, but with
increasing coverage of oxygen atoms on the surface the system should appear to
be further and further out of domain. The uncertainty boundary increases
smoothly as expected as more oxygen adsorbates are added, and the error
appears to increase in a similar manner.

In Figure 7 we look at another interpretable example, this time of oxygen
coverage on a platinum surface. Each system within the training and
calibration data sets contains only a single adsorbate placement per unit
cell, and therefore multiple adsorbate coverage on a single surface should
appear to be increasingly out of domain for a useful uncertainty prediction.
The majority of unit cells in the training data set contain surfaces with
multiple adsorption sites, but with adsorbate coverage on only one of them, so
increasing the coverage on the surface should be novel from the perspective of
the machine learning potential. For this example we measure the actual error
in the total energy predicted betweenVASP and the machine learning potential,
and compare it to the uncertainty prediction. The error is measured after
relaxation with the machine learning potential, since GNN relaxed systems with
adsorbates on slab surfaces is exactly what this RS2RE uncertainty method was
calibrated for. We expect to see the error and the uncertainty predictions
increase commensurately with increasing coverage. One key feature to note is
that we expect the uncertainty prediction to increase smoothly and
monotonically, which is exactly what we observe in Figure 7. We also see that
the magnitude of the errors does not increase perfectly monotonically, but
does follow the general trend of increasing uncertainty, and this aligns with
our expectation that the uncertainty predictions represent the dispersion of
the distribution that the measured errors are sampled from.

##  5 Conclusion

Effective uncertainty prediction methods for GNN relaxed energies are key to
the development of faster and more accurate screening techniques for novel
material discovery. Quantifying the performance of uncertainty methods on
relaxed energy predictions is especially complex, due to distribution
assumptions built into most commonly employed UQ techniques. Distribution-free
techniques which employ bootstrapped confidence intervals, such as the
CI(Var(Z)) test and error-based calibration plots, have been shown to be
better metrics for analyzing the calibration of a UQ method in similar
contexts, and we employ them here to great effect. We show that latent
distance methods outperform ensembles and other uncertainty methods on the
RS2RE task, which is of practical relevance to workflows such as AdsorbML. We
also show that the choice of latent representation is very important to the
calibration of the latent distance as an uncertainty metric. In the GNN latent
space, atom-wise distances produce better calibrated than system-wise
distances. Using rotationally invariant latent representations is crucial to
producing calibrated distance measures, and the rotationally invariant latent
space of GemNet-OC, a less accurate model, serves to compute a more well
calibrated measure of uncertainty for EquiformerV2 than its own rotationally
invariant latent space. We compare the latent space uncertainty predictions of
GemNet-OC to equation of state and coverage examples with expected outcomes,
and find that its behavior aligns with our expectations of a useful
uncertainty estimation method. Finally, we challenge the community to improve
on this RS2RE task for predicting uncertainties, using our proposed
recalibration framework as a measure. Future work in this area should also
explore the prediction of global minimum energy uncertainties directly, and
the development of model architectures training methods or distance measures
which preserve rotational equivariance while producing meaningful latent space
distances.

##  6 Data and Software Availability

Details of the implementations, and additional tables of results, for the
ensemble, distance, MVE, and sequence regression methods, as well as an
examination of the error distributions, can be found in the supporting
information.

## References

  * Zitnick et al. 2020 Zitnick, C. L. et al. An Introduction to Electrocatalyst Design using Machine Learning for Renewable Energy Storage. 2020, abs/2010.09435
  * Chanussot et al. 2021 Chanussot, L. et al. Open Catalyst 2020 (OC20) Dataset and Community Challenges. _ACS Catalysis_ 2021, _11_ , 6059–6072
  * Nørskov et al. 2014 Nørskov, J. K.; Studt, F.; Abild-Pedersen, F.; Bligaard, T. _Fundamental concepts in heterogeneous catalysis_ ; John Wiley & Sons, 2014
  * Dumesic et al. 2008 Dumesic, J. A.; Huber, G. W.; Boudart, M. _Principles of heterogeneous catalysis_ ; Wiley Online Library, 2008
  * Tran and Ulissi 2018 Tran, K.; Ulissi, Z. W. Active learning across intermetallics to guide discovery of electrocatalysts for CO2 reduction and H2 evolution. _Nat. Catal._ 2018, _1_ , 696–703
  * Zhong et al. 2020 Zhong, M. et al. Accelerated discovery of CO2 electrocatalysts using active machine learning. _Nature_ 2020, _581_ , 178–183
  * Liu et al. 2017 Liu, X.; Xiao, J.; Peng, H.; Hong, X.; Chan, K.; Nørskov, J. K. Understanding trends in electrochemical carbon dioxide reduction rates. _Nat. Commun_ 2017, _8_ , 1–7
  * Nørskov et al. 2005 Nørskov, J. K.; Bligaard, T.; Logadottir, A.; Kitchin, J. R.; Chen, J. G.; Pandelov, S.; Stimming, U. Trends in the Exchange Current for Hydrogen Evolution. _Journal of The Electrochemical Society_ 2005, _152_ , J23
  * Wan et al. 2022 Wan, X.; Zhang, Z.; Yu, W.; Niu, H.; Wang, X.; Guo, Y. Machine-learning-assisted discovery of highly efficient high-entropy alloy catalysts for the oxygen reduction reaction. _Patterns_ 2022, _3_ , 100553
  * Ulissi et al. 2017 Ulissi, Z. W.; Medford, A. J.; Bligaard, T.; Nørskov, J. K. To address surface reaction network complexity using scaling relations machine learning and DFT calculations. _Nat. Commun._ 2017, _8_ , 1–7
  * Gasteiger et al. 2022 Gasteiger, J.; Shuaibi, M.; Sriram, A.; Günnemann, S.; Ulissi, Z.; Zitnick, C. L.; Das, A. GemNet-OC: Developing Graph Neural Networks for Large and Diverse Molecular Simulation Datasets. _arXiv_ 2022, abs/2204.02782
  * Zitnick et al. 2022 Zitnick, C. L.; Das, A.; Kolluru, A.; Lan, J.; Shuaibi, M.; Sriram, A.; Ulissi, Z.; Wood, B. Spherical Channels for Modeling Atomic Interactions. _arXiv_ 2022, abs/2206.14331 [physics]
  * Passaro and Zitnick 2023 Passaro, S.; Zitnick, C. L. Reducing SO(3) Convolutions to SO(2) for Efficient Equivariant GNNs. Proceedings of the 40th International Conference on Machine Learning. 2023; pp 27420–27438
  * Liao et al. 2023 Liao, Y.-L.; Wood, B.; Das, A.; Smidt, T. EquiformerV2: Improved Equivariant Transformer for Scaling to Higher-Degree Representations. _arXiv_ 2023, abs/2306.12059
  * Lan et al. 2023 Lan, J.; Palizhati, A.; Shuaibi, M.; Wood, B. M.; Wander, B.; Das, A.; Uyttendaele, M.; Zitnick, C. L.; Ulissi, Z. W. AdsorbML: a leap in efficiency for adsorption energy calculations using generalizable machine learning potentials. _npj Computational Materials_ 2023, _9_ , 172
  * Hu et al. 2022 Hu, Y.; Musielewicz, J.; Ulissi, Z. W.; Medford, A. J. Robust and scalable uncertainty estimation with conformal prediction for machine-learned interatomic potentials. _Machine Learning: Science and Technology_ 2022, _3_ , 045028
  * Janet et al. 2019 Janet, J. P.; Duan, C.; Yang, T.; Nandy, A.; Kulik, H. J. A quantitative uncertainty metric controls error in neural network-driven chemical discovery. _Chemical Science_ 2019, _10_ , 7913–7922
  * Tran et al. 2020 Tran, K.; Neiswanger, W.; Yoon, J.; Zhang, Q.; Xing, E.; Ulissi, Z. W. Methods for comparing uncertainty quantifications for material property predictions. _Machine Learning: Science and Technology_ 2020, _1_ , 025006
  * Wollschläger et al. 2023 Wollschläger, T.; Gao, N.; Charpentier, B.; Ketata, M. A.; Günnemann, S. Uncertainty Estimation for Molecules: Desiderata and Methods. _arXiv_ 2023, abs/2306.14916 [physics, stat]
  * Carrete et al. 2023 Carrete, J.; Montes-Campos, H.; Wanzenböck, R.; Heid, E.; Madsen, G. K. H. Deep Ensembles vs. Committees for Uncertainty Estimation in Neural-Network Force Fields: Comparison and Application to Active Learning. _The Journal of Chemical Physics_ 2023, _158_ , 204801, abs/2302.08805 [physics]
  * Varivoda et al. 2023 Varivoda, D.; Dong, R.; Omee, S. S.; Hu, J. Materials property prediction with uncertainty quantification: A benchmark study. _Applied Physics Reviews_ 2023, _10_ , 021409
  * Tan et al. 2023 Tan, A. R.; Urata, S.; Goldman, S.; Dietschreit, J. C. B.; Gómez-Bombarelli, R. Single-model uncertainty quantification in neural network potentials does not consistently outperform model ensembles. _arXiv_ 2023, abs/2305.01754 [physics]
  * Pernot 2022 Pernot, P. Prediction uncertainty validation for computational chemists. _The Journal of Chemical Physics_ 2022, _157_ , 144103, abs/2204.13477 [physics]
  * Rasmussen et al. 2023 Rasmussen, M. H.; Duan, C.; Kulik, H. J.; Jensen, J. H. Uncertain of uncertainties? A comparison of uncertainty quantification metrics for chemical data sets. _ChemRxiv_ 2023,
  * Hirschfeld et al. 2020 Hirschfeld, L.; Swanson, K.; Yang, K.; Barzilay, R.; Coley, C. W. Uncertainty Quantification Using Neural Networks for Molecular Property Prediction. _arXiv_ 2020, abs/2005.10036 [cs, q-bio, stat]
  * Levi et al. 2022 Levi, D.; Gispan, L.; Giladi, N.; Fetaya, E. Evaluating and Calibrating Uncertainty Prediction in Regression Tasks. _Sensors_ 2022, _22_ , 5540
  * Gawlikowski et al. 2023 Gawlikowski, J.; Tassi, C. R. N.; Ali, M.; Lee, J.; Humt, M.; Feng, J.; Kruspe, A.; Triebel, R.; Jung, P.; Roscher, R.; Shahzad, M.; Yang, W.; Bamler, R.; Zhu, X. X. A survey of uncertainty in deep neural networks. _Artificial Intelligence Review_ 2023, _56_
  * Zaidi et al. 2021 Zaidi, S.; Zela, A.; Elsken, T.; Holmes, C. C.; Hutter, F.; Teh, Y. Neural Ensemble Search for Uncertainty Estimation and Dataset Shift. Advances in Neural Information Processing Systems. 2021; pp 7898–7911
  * Ovadia et al. 2019 Ovadia, Y.; Fertig, E.; Ren, J.; Nado, Z.; Sculley, D.; Nowozin, S.; Dillon, J. V.; Lakshminarayanan, B.; Snoek, J. Can You Trust Your Model’s Uncertainty? Evaluating Predictive Uncertainty under Dataset Shift. Proceedings of the 33rd International Conference on Neural Information Processing Systems. 2019
  * Egele et al. 2022 Egele, R.; Maulik, R.; Raghavan, K.; Lusch, B.; Guyon, I.; Balaprakash, P. AutoDEUQ: Automated Deep Ensemble with Uncertainty Quantification. 2022 26th International Conference on Pattern Recognition (ICPR). 2022; pp 1908–1914
  * Johnson et al. 2019 Johnson, J.; Douze, M.; Jégou, H. Billion-scale similarity search with GPUs. _IEEE Transactions on Big Data_ 2019, _7_ , 535–547
  * Navrátil et al. 2020 Navrátil, J.; Arnold, M.; Elder, B. Uncertainty Prediction for Deep Sequential Regression Using Meta Models. _arXiv_ 2020, abs/2007.01350
  * Sanh et al. 2019 Sanh, V.; Debut, L.; Chaumond, J.; Wolf, T. DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter. _arXiv_ 2019, abs/1910.01108
  * DiCiccio and Efron 1996 DiCiccio, T. J.; Efron, B. Bootstrap confidence intervals. _Statistical Science_ 1996, _11_ , 189–228
  * Virtanen et al. 2020 Virtanen, P. et al. SciPy 1.0: fundamental algorithms for scientific computing in Python. _Nature Methods_ 2020, _17_ , 261–272
  * Kresse and Furthmüller 1996 Kresse, G.; Furthmüller, J. Efficiency of ab-initio total energy calculations for metals and semiconductors using a plane-wave basis set. _Computational Materials Science_ 1996, _6_ , 15–50
  * Kresse 1995 Kresse, G. Ab initio molecular dynamics for liquid metals. _Journal of Non-Crystalline Solids_ 1995, _192-193_ , 222–229
  * Kresse and Hafner 1994 Kresse, G.; Hafner, J. Ab initio molecular-dynamics simulation of the liquid-metal amorphous-semiconductor transition in germanium. _Physical Review B_ 1994, _49_ , 14251–14269
  * Kresse and Furthmüller 1996 Kresse, G.; Furthmüller, J. Efficient iterative schemes for ab initio total-energy calculations using a plane-wave basis set. _Physical Review B - Condensed Matter and Materials Physics_ 1996, _54_ , 11169–11186
  * McInnes et al. 2018 McInnes, L.; Healy, J.; Saul, N.; Großberger, L. UMAP: Uniform Manifold Approximation and Projection. _Journal of Open Source Software_ 2018, _3_ , 861

TOC graphic