# Vehicle lateral control using Machine Learning for automated vehicle
guidance

Akash Fogla*, Kanish Kumar, Sunnay Saurav, Bishnu ramanujan
*Department of Mathematics, Georgia State University

##  I ABSTRACT

Uncertainty in decision-making is crucial in the machine learning model used
for a safety-critical system that operates in the real world. Therefore, it is
important to handle uncertainty in a graceful manner for the safe operation of
the CPS. In this work, we design a vehicle’s lateral controller using a
machine-learning model. To this end, we train a random forest model that is an
ensemble model and a deep neural network model. Due to the ensemble in the
random forest model, we can predict the confidence/uncertainty in the
prediction. We train our controller on data generated from running the car on
one track in the simulator and tested it on other tracks. Due to prediction in
confidence, we could decide when the controller is less confident in
prediction and takes control if needed. We have two results to share: first,
even on a very small number of labeled data, a very good generalization
capability of the random forest-based regressor in comparison with a deep
neural network and accordingly random forest controller can drive on another
similar track, where the deep neural network-based model fails to drive, and
second confidence in predictions in random forest controller makes it possible
to let us know when the controller is not confident in prediction and likely
to fail. By creating a threshold, it was possible to take control when the
controller is not safe and that is missing in a deep neural network-based
controller.

### Keywords

Random forest, Deep neural networks, Torcs, surrogate modeling, Electronics
Control Unit

##  II Introduction

A self-driving car has many machine learning and AI-based components involved
that act as subsystems in a complex driving system like a perception system,
obstacle detection system, etc. However, there are very few examples of an
end-to-end driving system[1] by using ML and AI models. ML and AI have shown
extraordinary success in the field of controls[2, 3], prediction or modeling
of complex behavior like cancer detection, stock market prediction, etc [4,
5], to design automation [6]. In this work , we take two famous and
extensively used machine learning model and train it to control the vehicle
lateral control. The goal of this study is to understand the comparison
between different trained model when used in controlling a system when an
adequate amount of data is available. This comparison study can be used as a
baseline and can provide guidelines for AI based control designer an idea
about selection of AI-ML model and their respective strength and weakness. For
vehicle lateral control, we utilize the open-source car racing simulator
called TORCS[7]. On a given set speed , the data is generated using an
traditional PID controller designed by experimenter. The PID controller
produces brake value, acceleration and steering to control and drive the car
on a given track. During training , for lateral control, we are only
interested in steering value and brake and acceleration value is taken using
the PID controller. The collected data is distance measured by LIDAR sensor
suite and the steering value. For training the controller ,we select decision
tree based ML model called Random forest[8] and Deep neural network model[9,
10]. Both are trained on same data set and tested in same scenario (track and
velocity). Our experimentation shows following :

  1. 1.

Random forest-based controller provides better generalization capability than
Deep neural network when we done have access of very large data set.

  2. 2.

When both trained model is deployed on another track, the random forest-based
controller was able to complete the track without crashing, while the deep
neural network-based controller failed to complete the track.

  3. 3.

The random forest controller can quantify uncertainty in prediction that can
help to decide the failure probability of the controller and override the AI
control command if needed.

The rest of the paper is organised as follow. Section III formulate the
problem and provide background, approach and training details of controller.
Section IV provides the experimental results of the experimentation. The
related work is discussed in section V and at the end we produce our
conclusion and future direction of researchVI.

##  III Problem Formulation and Approach

###  III-A Problem Formulation

For formalization of the problem, we first introduce some notations. Let $x$
is the state of the vehicle that is estimated by the sensors attached to the
car and $x$ is sampled from an unknown distribution $P_{x}$ which we neither
control nor know but can sample from it. The goal of controller design is to
learn a nonlinear function $f$ such that it maps each $x$ to the control
action $c$.

$f:x\mapsto c\;where\;x\sim P_{x}$

The choice of $f$ that is an instance of a class of architecture
($\mathcal{A}$) affects the quality of the control action. In this context, we
chose two famous or extensively used architectures called random forest and
Deep neural networks. The choice of model architecture also affects the
training process and the parameters to select. In this context, we frame this
modeling problem as a supervised learning problem where it is possible to get
a labeled dataset that can be used for training. The collection of labeled
sets of $<x_{i},c_{i}>$ represents our data $D$. This data set is further
split into training and test dataset called $D_{train}$ and $D_{test}$. Then
the learning goal is to find the parameters (in the case of neural networks or
feature decision graphs (in the case of the random forest) that can represent
the training data the most and work well on test data as well.

###  III-B Approach

In this section, we discuss a little background literature on our
experimentation components, the approach of the experimentation, model
training, and related details.

Torcs simulator TORCS (The Open Racing Car Simulator) is an open-source 3D car
racing simulator that has the ability to drive manually using peripheral
devices or by design controller- AI-based or traditional controller. The key
components of vehicle dynamics, including mass, rotational inertia, collision,
the mechanics of suspensions, links, differentials, friction, and
aerodynamics, may all be accurately simulated by it. Using a temporal
discretization level of $0.002$ seconds, Euler integration of differential
equations is used to simplify and carry out physics simulation. TORCS offers a
large variety of tracks and cars as assets.

Data generation For data generation purposes, we used the TORCS’ asset car on
the road track called ’Forza’ (refer to fig 1). For sensing the environment,
we deployed multiple LIDAR sensors attached to the car that can sense the
track and the wall. For data generation purposes, we designed a tuned PID
controller that can drive the car successfully on a given speed on the track
Forza. Our data set consists of input as a set of distances measured by
various LIDAR sensors and output as the steering value. We ran multiple rounds
on this track and collected this data. This labeled data set is used for
training our models. Both models (random forest and deep neural network) can
be used for both classification and regression tasks. In our case, we want to
use it as a regressor that can predict the steering value on the given LIDAR
sensor data.

Figure 1: track: Forza

Random forest model, parameters and training  Random forest is a learning
model that is an ensemble model of various decision trees. For regression
tasks, the mean prediction of the individual decision trees is the estimated
prediction. Random forests do various other tweaks in the vanilla decision
tree for taking care of the habit of over-fitting of decision trees’ on their
training set.

In particular, decision trees are grown very deep when trained on data and
consequently tend to learn highly nonlinear patterns but over-fit, i.e. have
low bias, but very high variance. A random forest is an approach of averaging
multiple decision trees, where each decision tree is trained on a different
sub-sample of the training data set, with the goal of reducing the variance in
prediction.

For training the decision tree in a random forest we apply the technique of
bootstrap aggregating, or bagging, to each decision tree learner. On a given
input data set $D=\\{X,Y\\}$ where $X$ is LIDAR’s measurement of obstacle
distance; $X=\\{x_{1},x_{2},...,x_{n}\\}$ and steering output
$Y=\\{y_{1},y_{2},...,y_{n}\\}$, the bagging process involves repeatedly
sampling a subset of data with $B$ data points from the data set $D$ with
replacement and train each decision tree on this selected samples. If $n$ is
the number of trees in the random forest the process involves as below: for
i=1,…,n:

  1. 1.

Sample with replacement $B$ training examples from $D$; call it $D_{i}$.

  2. 2.

train the $i^{th}$ decision tree $T_{i}$ on data set $D_{i}$.

For prediction on unseen data $x^{{}^{\prime}}$, the steering output is
estimated by estimating the mean of prediction from each decision tree.

$\hat{s}=\frac{1}{n}\sum_{i=1}^{i=n}T_{i}(x^{{}^{\prime}})$

Additionally, the uncertainty in the prediction can be estimated by
calculating the standard deviation of the predictions from all the individual
regression trees on $x^{{}^{\prime}}$:

$\sigma=\sqrt{\frac{\sum_{i=1}^{i=n}(T_{i}(x^{{}^{\prime}})-\hat{s})^{2}}{n-1}}$

The prediction value of steering $\hat{s}$ is used for steering the vehicle
and the standard deviation and other related statistics -like Coefficient of
Variance (CoV). We select $n=100$ i.e. total 100 number of decision trees. The
other hyper-parameters are ’splitting criteria= squared error’, minimum
samples for split=2, the maximum number of features=1.0, and minimum impurity
decrease=$0.001$.
Deep Neural Network model, parameters and training DNN are computing functions
inspired by the biological neural networks of animal brains. These systems
learn to do a task by observing examples that are already done/labeled. It has
wide ranging application from image recognition, engineering design[11, 6],
control design, anomaly detection[vardhan],etc. A DNN is based on a layered
network of smaller computational units called artificial neurons. These
neurons are in multiple layers and fully connected between the input and
output layers. A fully connected feed-forward neural network is used in this
experiment that is defined by its architecture $a=(L,N,\delta)$, where $L$ is
the number of layers in the networks $L\in N$, $N$ is the set that represents
the number of neurons in each layer represented by $N_{l}$ where $l\in[L-1]$,
$\delta$ is the activation function $\delta:R\mapsto R$. For this experiment,
we used a fully connected feed-forward 6 layers neural network with
$\\{256,128,64,32,16\\}$ neurons in hidden layers and $1$ neuron in the output
layer. All hidden neurons have ReLU activation units and the output neuron is
the linear activation unit. The loss function, in this case, is the mean
square error and the weight initialization is Xavier normal. The training is
done for $500$ epochs by splitting the training and validation data into
$90:10$ and early stopping is used to stop over-fitting.

Figure 2: Training and validation loss for deep neural network training

##  IV Experimental Results

Once both models are trained, these models are tested on a similar but another
track called ’E-Track 4’ (refer to figure 3). The selection of the track is
based on twists and turns and complexity in the track profile. We set the
target speed of the car to 60 miles/hrs. For the first experiment, we used
both controllers to complete the track without any manual interventions. The
task was to finish one complete lap of this track. One complete lap of this
track is 7.041 Km long with a track width of 15 meters. The random forest
controller was able to complete the task without any crash or off-the-track
navigation, while the deep neural network-based controller failed to complete
the track and crashed multiple times and could not complete even $10\%$ of the
track without a crash.

Figure 3: track: E-track 4

In the second experiment, we let the experimenter intervene whenever the
random forest controller gives a prediction with a high coefficient of
variance. A high coefficient of variance (CoV) reflects less confidence in
prediction and consequently, its prediction cannot rely upon. In such cases,
the control action is shifted to manual control. We observed that the trained
random forest controller produced high CoV in scenarios which was very complex
and never encountered during training and also when we tried to run it on a
significantly different track. This feature is missing with the Deep Neural
network controller. A snapshot of the random forest controller driving the car
on the test track is shown in figure 4.

Figure 4: RF controller in Torcs and related statistics, the statistics of our
interest are: odd counts -number of decision trees out of 2 standard
deviations from the mean, total\- total number of odd counts during whole
simulation, steer\- steering value predicted by regressor, std\- standard
deviation of prediction, CoV\- coefficient of variance, that explain the
spread of the distribution.

##  V Related Work

The dominant approach to solving engineering problems are model-driven, like
model-based design[12], model-based control[13], and model-based
optimization[14, 15] etc. With the discovery of data-driven models especially
deep neural networks the domain of model-based engineering is revolutionizing.
In recent times these data-driven models have been used and have worked well
in image classification (AlexNet[16] ), engineering design[17, 18, 19, 20],
autonomous driving car[1, 21, 22, 23], radiology [24], human genome[25], and
many more for developing state of the art control[26, 27, 28] and prediction
systems[29].

According to research[30], self-driving cars could potentially save up to
$35,000$ lives annually in the United States alone by reducing traffic deaths
by up to $99$ percent. $93$ percent of car accidents are the result of driver
mistake, according to information from 2007 US report[31], and rookie
drivers—who seldom have more than $30$ hours of driving experience before
receiving a license—are over-represented in fatal car accidents. Moreover,
crashes among senior drivers rise at the same time because they have a
diminished capacity to assess their surroundings for unforeseen dangers,
according to [30]. Self-driving cars will have been trained using machine
learning data-sets made up of a wide range of driving circumstances and driver
behaviors for the equivalent of hundreds of hours behind the wheel before they
hit the roads[1]. The practical utility of AI-ML model is enormous [32, 33,
34, 35] and a better understanding and insight about the model’s performance
would be useful for practical control designer. Since Deep learning is a data
hungry model, it is not possible to generate an enormous amount of data for
each application. In such cases, this DNN models fail to perform and
alternative learning model can be useful.

##  VI Conclusion and Future Work

###  VI-A Conclusion

In this work, we attempted to design a controller for the lateral control of a
vehicle using AI-ML models. Experiments’ results show that:

  1. 1.

Random forest-based controller has better generalization capability than Deep
neural network when limited data is available.

  2. 2.

Trained on one track and deployed on another track, the RF-based controller
was able to complete the track without crashing, while the deep neural
network-based controller could not complete even 10% of the track before
crashing.

  3. 3.

Due to an ensemble of trees in RF, it is possible to quantify uncertainty and
can help to decide the failure probability of the controller and take command
of control if needed.

###  VI-B Future Work

Our experiments are use cases to understand deeper about two prominently used
AI models when used in the context of control. For future work, we want to
extend this work by replacing the deep neural network with a Bayesian neural
network that has better generalization capability and can predict confidence
in prediction and evaluate its performance.

## References

  * [1] M. Bojarski, D. Del Testa, D. Dworakowski, B. Firner, B. Flepp, P. Goyal, L. D. Jackel, M. Monfort, U. Muller, J. Zhang _et al._ , “End to end learning for self-driving cars,” _arXiv preprint arXiv:1604.07316_ , 2016.
  * [2] P. Abbeel, A. Coates, and A. Y. Ng, “Autonomous helicopter aerobatics through apprenticeship learning,” _The International Journal of Robotics Research_ , vol. 29, no. 13, pp. 1608–1639, 2010.
  * [3] H. Vardhan and J. Sztipanovits, “Rare event failure test case generation in learning-enabled-controllers,” in _2021 6th International Conference on Machine Learning Technologies_ , 2021, pp. 34–40.
  * [4] A. Al Bataineh, “A comparative analysis of nonlinear machine learning algorithms for breast cancer detection,” _International Journal of Machine Learning and Computing_ , vol. 9, no. 3, pp. 248–254, 2019.
  * [5] M. A. Ghazanfar, S. A. Alahmari, Y. F. Aldhafiri, A. Mustaqeem, M. Maqsood, and M. A. Azam, “Using machine learning classifiers to predict stock exchange index,” _International Journal of Machine Learning and Computing_ , vol. 7, no. 2, pp. 24–29, 2017.
  * [6] H. Vardhan and J. Sztipanovits, “Deepal for regression using epsilon-weighted hybrid query strategy,” _arXiv preprint arXiv:2206.13298_ , 2022.
  * [7] B. Wymann, E. Espié, C. Guionneau, C. Dimitrakakis, R. Coulom, and A. Sumner, “Torcs, the open racing car simulator,” _Software available at http://torcs. sourceforge. net_ , vol. 4, no. 6, p. 2, 2000.
  * [8] L. Breiman, “Random forests,” _Machine learning_ , vol. 45, pp. 5–32, 2001\.
  * [9] A. G. Ivakhnenko, “The group method of data handling a rival of stochastic approximation,” _Soviet Automatic Control_ , vol. 13, pp. 43–55, 1968.
  * [10] K. Fukushima, “Neocognitron: A hierarchical neural network capable of visual pattern recognition,” _Neural networks_ , vol. 1, no. 2, pp. 119–130, 1988\.
  * [11] H. Vardhan, P. Volgyesi, and J. Sztipanovits, “Machine learning assisted propeller design,” in _Proceedings of the ACM/IEEE 12th International Conference on Cyber-Physical Systems_ , 2021, pp. 227–228.
  * [12] H. Neema, H. Vardhan, C. Barreto, and X. Koutsoukos, “Web-Based Platform for Evaluation of Resilient and Transactive Smart-Grids,” in _2019 7th Workshop on Modeling and Simulation of Cyber-Physical Energy Systems (MSCPES)_. IEEE, 2019, pp. 1–6.
  * [13] C. Brosilow and B. Joseph, _Techniques of model-based control_. Prentice Hall Professional, 2002.
  * [14] H. Vardhan, N. M. Sarkar, and H. Neema, “Modeling and optimization of a longitudinally-distributed global solar grid,” in _2019 8th International Conference on Power Systems (ICPS)_. IEEE, 2019, pp. 1–6.
  * [15] H. Neema _et al._ , “Design and Simulation Platform for Evaluation of Grid Distribution System and Transactive Energy,” in _Proceedings of the 6th Annual Symposium on Hot Topics in the Science of Security_. ACM, 2019, p. 15.
  * [16] A. Krizhevsky, I. Sutskever, and G. E. Hinton, “Imagenet classification with deep convolutional neural networks,” _Communications of the ACM_ , vol. 60, no. 6, pp. 84–90, 2017.
  * [17] H. Vardhan, U. Timalsina, P. Volgyesi, and J. Sztipanovits, “Data efficient surrogate modeling for engineering design: Ensemble-free batch mode deep active learning for regression,” _arXiv preprint arXiv:2211.10360_ , 2022\.
  * [18] S. M. Moosavi, K. M. Jablonka, and B. Smit, “The role of machine learning in the understanding and design of materials,” _Journal of the American Chemical Society_ , vol. 142, no. 48, pp. 20 273–20 287, 2020.
  * [19] H. Vardhan and J. Sztipanovits, “Search for universal minimum drag resistance underwater vehicle hull using cfd,” _arXiv preprint arXiv:2302.09441_ , 2023\.
  * [20] M. J. Volk, I. Lourentzou, S. Mishra, L. T. Vo, C. Zhai, and H. Zhao, “Biosystems design by machine learning,” _ACS synthetic biology_ , vol. 9, no. 7, pp. 1514–1533, 2020.
  * [21] C. Badue, R. Guidolini, R. V. Carneiro, P. Azevedo, V. B. Cardoso, A. Forechi, L. Jesus, R. Berriel, T. M. Paixao, F. Mutz _et al._ , “Self-driving cars: A survey,” _Expert Systems with Applications_ , vol. 165, p. 113816, 2021.
  * [22] J. Ni, Y. Chen, Y. Chen, J. Zhu, D. Ali, and W. Cao, “A survey on theories and applications for self-driving cars based on deep learning methods,” _Applied Sciences_ , vol. 10, no. 8, p. 2749, 2020.
  * [23] S. Shalev-Shwartz, S. Shammah, and A. Shashua, “On a formal model of safe and scalable self-driving cars,” _arXiv preprint arXiv:1708.06374_ , 2017.
  * [24] G. W. Gross, J. M. Boone, V. Greco-Hunt, and B. Greenberg, “Neural networks in radiologic diagnosis. ii. interpretation of neonatal chest radiographs.” _Investigative Radiology_ , vol. 25, no. 9, pp. 1017–1023, 1990.
  * [25] L. Sundaram, H. Gao, S. R. Padigepati, J. F. McRae, Y. Li, J. A. Kosmicki, N. Fritzilas, J. Hakenberg, A. Dutta, J. Shon _et al._ , “Predicting the clinical impact of human mutation with deep neural networks,” _Nature genetics_ , vol. 50, no. 8, pp. 1161–1170, 2018.
  * [26] H. Vardhan and J. Sztipanovits, “Reduced robust random cut forest for out-of-distribution detection in machine learning models,” _arXiv preprint arXiv:2206.09247_ , 2022.
  * [27] Z. Wu, A. Tran, D. Rincon, and P. D. Christofides, “Machine learning-based predictive control of nonlinear processes. part i: theory,” _AIChE Journal_ , vol. 65, no. 11, p. e16729, 2019.
  * [28] S. M. Tahsien, H. Karimipour, and P. Spachos, “Machine learning based solutions for security of internet of things (iot): A survey,” _Journal of Network and Computer Applications_ , vol. 161, p. 102630, 2020.
  * [29] H. Vardhan and J. Sztipanovits, “Deep learning based fea surrogate for sub-sea pressure vessel,” in _2022 6th International Conference on Computer, Software and Modeling (ICCSM)_. IEEE, 2022, pp. 36–39.
  * [30] Y. A. Momtaz, R. Kargar, R. Hosseiny, and R. Sahaf, “Rate and pattern of road traffic accidents among older and younger drivers,” _Healthy aging research_ , vol. 7, no. 2, p. e18, 2018.
  * [31] N. T. S. Facts, “Us department of transportation, national highway traffic safety administration, 2008,” 2007.
  * [32] H. Vardhan, P. Volgyesi, and J. Sztipanovits, “Fusion of ml with numerical simulation for optimized propeller design,” _arXiv preprint arXiv:2302.14740_ , 2023.
  * [33] T. Dreossi, D. J. Fremont, S. Ghosh, E. Kim, H. Ravanbakhsh, M. Vazquez-Chanlatte, and S. A. Seshia, “Verifai: A toolkit for the formal design and analysis of artificial intelligence-based systems,” in _Computer Aided Verification: 31st International Conference, CAV 2019, New York City, NY, USA, July 15-18, 2019, Proceedings, Part I 31_. Springer, 2019, pp. 432–442.
  * [34] H. Vardhan, P. Volgyesi, and J. Sztipanovits, “Constrained bayesian optimization for automatic underwater vehicle hull design,” _arXiv preprint arXiv:2302.14732_ , 2023.
  * [35] R. Vaishya, M. Javaid, I. H. Khan, and A. Haleem, “Artificial intelligence (ai) applications for covid-19 pandemic,” _Diabetes & Metabolic Syndrome: Clinical Research & Reviews_, vol. 14, no. 4, pp. 337–339, 2020.


# Navigating Unknowns: Deep learning robustness for gravitational wave signal
reconstruction

Chayan Chatterjee Department of Physics and Astronomy, Vanderbilt University
2201 West End Avenue, Nashville, Tennessee - 37235,  Data Science Institute,
Vanderbilt University
1400 18th Avenue South Building, Suite 2000, Nashville, Tennessee - 37212,
Karan Jani Department of Physics and Astronomy, Vanderbilt University
2201 West End Avenue, Nashville, Tennessee - 37235,

###### Abstract

We present a rapid and reliable deep learning-based method for gravitational
wave signal reconstruction from elusive, generic binary black hole mergers in
LIGO data. We demonstrate that our model, AWaRe, effectively recovers
gravitational waves with parameters it has not encountered during training.
This includes features like higher black hole masses, additional harmonics,
eccentricity, and varied waveform systematics, which introduce complex
modulations in the waveform’s amplitudes and phases. The accurate
reconstructions of these unseen signal characteristics demonstrates AWaRe’s
ability to handle complex features in the waveforms. By directly incorporating
waveform reconstruction uncertainty estimation into the AWaRe framework, we
show that for real gravitational wave events, the uncertainties in AWaRe’s
reconstructions align closely with those achieved by benchmark algorithms like
BayesWave and coherent WaveBurst. The robustness of our model to real
gravitational wave events and its ability to extrapolate to unseen data open
new avenues for investigations in various aspects of gravitational wave
astrophysics and data analysis, including tests of General Relativity and the
enhancement of current gravitational wave search methodologies.

##  1 Introduction

Gravitational wave (GW) astronomy has undergone a remarkable evolution since
the first detection in 2015, profoundly impacting our understanding of the
universe’s most violent and energetic phenomena (Abbott et al., 2016). The
field now frequently reports new discoveries, signifying a major leap in our
observational capabilities. The fourth observing run (O4) of the LIGO-Virgo-
KAGRA Collaboration (Aasi et al., 2015; Acernese et al., 2014; Akutsu et al.,
2019) is currently in progress and plans are set for the next decade to
further expand this network with both ground and space-based detectors (Babak
et al., 2021; Amaro-Seoane et al., 2017; Jani et al., 2020). These
advancements promise to increase the frequency of GW detections significantly,
thereby enhancing our ability to conduct follow-up observations across various
cosmic messengers, such as electromagnetic radiation and astrophysical
neutrinos. This evolving landscape of GW astronomy sets the stage for
significant advancements in understanding the universe through the lens of GW
and their accompanying signals.

GW signals from merging compact binaries detected by the LIGO and Virgo
interferometers are primarily identified using the template-based matched
filtering technique (Finn, 1992; Allen et al., 2012), which is most effective
with stationary and Gaussian noise, but can be compromised by non-stationary
noise transients in the detectors. Additionally, current low-latency matched
filtering-based detection approaches often overlook the impact of precession
(Ajith et al., 2011; Schmidt et al., 2015) in black hole spins and higher-
order waveform harmonics, instead focusing only on the primary quadrupolar
modes. This restriction limits our ability to explore significant
astrophysical phenomena, such as the presence of intermediate-mass black hole
(IMBH) binaries (Calderón Bustillo et al., 2018) and hierarchical mergers,
although recent work in these directions present significant advancements in
the search for these elusive systems (McIsaac et al., 2023; Schmidt et al.,
2024; Harry et al., 2018; Chandra et al., 2022; Wadekar et al., 2024). Machine
learning offers a transformative potential for GW data analysis, bringing
sophisticated tools that can enhance detection capabilities and reduce
computational burdens (Schäfer et al., 2023; Cuoco et al., 2020). In
particular, deep learning models (Goodfellow et al., 2016) have been
successfully applied to various aspects of GW astronomy, including rapid
signal detection, noise characterization, and parameter estimation (George &
Huerta, 2018a; Kim et al., 2015; George & Huerta, 2018b; Schäfer et al., 2020;
Astone et al., 2018; Skliris et al., 2022; Krastev, 2020; Qiu et al., 2023;
Dax et al., 2023; Gabbard et al., 2022).

In Chatterjee & Jani (2024), we demonstrated for the first time that deep
learning can effectively reconstruct waveforms from precessing binary black
hole (BBH) signals that include higher-order modes. This means given a segment
of noisy detector strain data for such GW events, our model outputs the
underlying astrophysical whitened waveform, removing it completely from
background noise. Alternately, if the model is presented with a segment of
pure noise data, with no underlying astrophysical signal, the model outputs a
vector of tiny random fluctuations around zero amplitude, indicating that this
method can be readily adapted into a rapid GW search algorithm. We will
explore this application in a future work. Our model, AWaRe, or Attention-
boosted Waveform Reconstruction network, evaluated using simulated signals
embedded in actual LIGO Livingston noise from the third observation run
(2019-2020) (Abbott et al., 2023), consistently achieved a high overlap with
the injected signals. This level of accuracy persisted across a broad spectrum
of BH mass and spin configurations considered in our study. When applied to
real GW events, our model’s reconstructions achieved overlaps ranging from 85%
to 98% with those obtained by Coherent WaveBurst (unmodeled) (Klimenko et al.,
2016; cWB, 2021) and LALInference (modeled) (Veitch et al., 2015) methods,
indicating that AWaRe is an effective and versatile tool for signal
reconstrcution from a diverse range of compact binary systems.

One of the limitations of the approach presented in Chatterjee & Jani (2024)
is that AWaRe was designed to only produce point-estimate predictions of the
reconstructed BBH waveform, given a noisy strain time-series, without an
estimate of the corresponding reconstruction uncertainty. This leads to biased
estimates of the model’s learning ability, since multiple waveforms can fit
the same data, especially in the presence of noise and model limitations.
Furthermore, having robust uncertainty estimates would enable a more rigorous
validation of the model’s generalization capabilities. Specifically,
uncertainty quantification is critical for assessing whether AWaRe can
accurately reconstruct waveforms it has not previously encountered, such as
those from different waveform approximants, higher-order modes (HM), or
scenarios with eccentric orbits. By providing confidence intervals or
probability distributions around the predicted waveforms, we can determine the
reliability of the model’s predictions across these varied scenarios. This is
essential for ensuring that the model not only performs well on familiar data
but also adaptably handles the complex and diverse nature of real-world GW
signals.

Figure 1: Top: Reconstructed signal (blue) corresponding to the injected
waveform (red) plotted against the whitened strain (grey), obtained by
injecting the simulated waveform in LIGO O3 noise. The dark and light blue
regions show 50% and 90% confidence intervals of the reconstructed waveform.
Bottom: The frequency evolution over time of the waveform and its
reconstruction shown in the figure on top.

In this article, we present an enhanced version of AWaRe, introduced in
Chatterjee & Jani (2024), in which we incorporate uncertainty estimation into
the GW waveform reconstruction. We also perform tests of the generalizability
of the model for different types of BBH sources. The remainder of this paper
is organized as follows: Section 2 describes the modifications made to our
model from Chatterjee & Jani (2024) for uncertainty estimation. Section 3.1
demonstrates AWaRe’s performance on real GW events, along with a comparison
against BayesWave (BW) (Cornish & Littenberg, 2015; Cornish et al., 2021) and
cWB algorithms. In Sections 3.2 – 3.3, we discuss our model’s robustness on
injections in real LIGO Livingston noise for signals whose parameters lie
outside our training set distributions, including different waveform families,
higher primary masses and non-zero eccentricities. Finally, in Section 4, we
discuss the implications of this work and the applicability of this model for
GW detection and detector characterization.

Figure 2:  Clockwise from top left: AWaRe mean reconstruction and 90% C.I.
(blue), cWB 90% C.I. (yellow) and BW 90% C.I. (orange) for GW190521
($\mathcal{M_{\text{det}}}$ = 113.48 M⊙), GW190602_175927
($\mathcal{M_{\text{det}}}$ = 72.18 M⊙), GW190412 ($\mathcal{M_{\text{det}}}$
= 15.30 M⊙) and (d) GW190517_055101 ($\mathcal{M_{\text{det}}}$ = 35.68 M⊙).
The numbers in the brackets refer to the median detector frame chirp mass of
these events (Abbott et al., 2021). The grey lines show the whitened strain
data provided as input to the model.

##  2 Method

Our model, AWaRe is a deep learning model consisting of two neural networks –
an encoder, consisting of one-dimensional convolutional layers (Lecun et al.,
1998; Krizhevsky et al., 2012) that extract features from noisy GW strain
data, and a decoder, consisting of Long Short-Term Memory (LSTM) (Hochreiter &
Schmidhuber, 1997) layers that reconstruct the pure waveforms from the
embeddings learned by the encoder. In addition to the encoder and decoder
layers, we apply a multi-headed attention (MHA) mechanism (Vaswani et al.,
2017) on top of the encoder stack which generates a richer feature
representation from the input, allowing the model to better focus on relevant
patterns in the data, thereby enhancing the accuracy and efficiency of the
waveform reconstruction. This integration of MHA provides significant
improvements in handling complex waveform features and also makes AWaRe
particularly effective in distinguishing genuine signals from noise. Section
2.2 in Chatterjee & Jani (2024) provides a detailed description of the AWaRe
model architecture and the MHA algorithm.

In the present work, we enhance the predictions of AWaRe by introducing the
estimation of waveform reconstruction uncertainty. To achieve this, we assume
a Gaussian distribution for every sample of our model’s prediction and train
our network to output the mean and standard deviation for these Gaussian
distributions. In particular, for a 1-second long input strain data sampled at
2048 Hz, our model predicts 2048 means and standard deviations, corresponding
to the 2048 independent Gaussian random variables used to parameterize the
reconstructed waveform. These mean and standard deviation values are then used
to compute the 90% confidence intervals of the model’s prediction of the
recovered waveforms. Fig. 1 shows the amplitude and frequency evolution of a
30 + 30 M⊙ GW signal (red) and its reconstruction with associated uncertainty
(blue) obtained using AWaRe. To obtain the frequency evolution plot, the
method described in Sec. V B in (Cornish et al., 2021) was used.

##  3 Results

We present results obtained by AWaRe on real GW events detected during the
third observation run (O3) (Abbott et al., 2023), as well as on simulated
injections with source parameters beyond the chosen prior distributions of our
training set. To characterize the quality of waveform reconstruction, we
compute the mismatch ($\mathcal{M}$) between AWaRe’s mean reconstruction and
the maximum likelihood waveform reconstruction from cWB and BW, for real
events, and the injected waveforms for simulated injections.

The mismatch quantifies the dissimilarity between two waveforms and is defined
as,

$\mathcal{M}(h,s)=1-\max_{t_{c},\phi_{c}}\mathcal{O}(h,s),$ (1)

where $\max_{t_{c},\phi_{c}}\mathcal{O}(h,s)$ represents the maximum overlap,
as defined in Eq. 7 in (Chatterjee & Jani, 2024), maximized over the time
($t_{c}$) and phase ($\phi_{c}$) of coalescence. Here, $h$ may refer to the
reconstructed waveform produced by AWaRe, and $s$ refers either to the
injection waveform or to the maximum likelihood reconstruction obtained by cWB
or BW for real events. A higher value of $\mathcal{M}$ denotes higher
dissimilarity between two waveforms. The mismatch values quoted throughout the
paper refers to the mismatch computed between the mean reconstruction obtained
from AWaRe (waveform obtained by taking the mean value of the 2048 independent
Gaussians) and the injection sample.

###  3.1 Tests on real O3 events

We tested AWaRe’s performance on real O3 events and compared the
reconstruction with results from cWB and BW. Figures 2 shows the mean AWaRe
reconstruction (in blue) and 90% confidence interval (in light blue) against
the 90% confidence intervals from cWB (in yellow) and BW (in orange). The grey
curves across all the plots show the whitened strain data provided as input to
the model.

The top left plot shows our result for GW190521, an interesting candidate due
to its high total mass and support for misaligned spins (Bayes factor = 11.5)
(Abbott et al., 2020a). There is also strong observational evidence for a
multimode BH ringdown spectrum (Capano et al., 2023). Our findings indicate
that the mean AWaRe reconstruction lies well within the 90% C.I. of both cWB
and BW, and the 90% C.I. also overlaps well with that of the cWB and BW 90%
C.I. However, AWaRe predicts an additional cycle at the end of the merger
region, which is not present in the other two analyses.

GW190602_175927 (top right, Fig. 2) is another interesting GW candidate due to
its high total mass. Our results again demonstrate very high reconstruction
accuracy relative to cWB and BW around the merger, with small additional
fluctuations in the inspiral phase that are not observed in the other
analyses. The high sensitivity of AWaRe towards heavy binaries and burst-like
real sources motivates applying this method for the analysis of IMBH binaries.
We will investigate this in our future work.

Figure 3: AWaRe mean reconstruction with 50% and 90% C.I. (blue) for
GW190412-like injection (red) generated using (a) IMRPhenomXPHM (b)
SEOBNRv4PHM and (c) NRHybSur3dq8 waveform approximants. The grey lines show
the whitened strain data.

The bottom left figure in Fig. 2 shows our reconstruction for GW190412, a
notable case due to its higher order mode contributions (Abbott et al.,
2020b). The event’s unequal masses ($q=m_{2}/m_{1}=0.28^{+0.12}_{-0.07}$),
non-zero $\chi_{\text{{eff}}}$ ($0.25^{+0.08}_{-0.11}$), and $\chi_{p}$
($0.31^{+0.19}_{-0.16}$) make it an ideal test for our model. AWaRe’s mean
reconstruction and 90% confidence interval (C.I.) align in phase with the cWB
90% C.I., though AWaRe shows greater uncertainty during the inspiral. The BW
90% C.I., however, exhibits significant mismatches in amplitude and phase
compared to both AWaRe and cWB, predicting more oscillations. This discrepancy
may arise because both cWB and BW use wavelets optimized for burst-like
sources, not low chirp-mass events like GW190412. BW uses wavelets within a
Bayesian framework for precise reconstruction (Cornish & Littenberg, 2015;
Cornish et al., 2021; Millhouse et al., 2018), while cWB employs wavelets for
time-frequency analysis and coherent detection across detectors, focusing on
unmodeled burst signals (Klimenko et al., 2016). Since these methods do not
rely on modelled template banks for searches and reconstruction, they are most
effective for analysis of burst-like GW sources, corresponding to events with
high total masses ($>$ 60 M⊙).

Figure 4: (Left) Mismatch between AWaRe mean reconstructions and pure
waveforms as a function of SNR for injections in LIGO O3b noise. (Right)
Mismatch between AWaRe mean reconstructions and pure waveforms as a function
of eccentricity for injections in LIGO O3b noise. The red dots (left) and blue
dots (right) and associated error bars show the mean mismatches and the
1-$\sigma$ uncertainties obtained by running AWaRe on injections in 10
different realizations of O3b noise. The dark blue lines shows the best fit
curves. The inset plots shows pure waveforms (red) and 90% C.I. of AWaRe
reconstruction (light blue) for a SNR-12 and SNR-22 sample respectively (left)
and waveforms with eccentricities of 0.11 and 0.37 respectively (right).

The bottom right figure in Fig. 2 shows result for GW190517_055101, which is a
more generic BBH signal with primary and secondary masses,
$37.4^{+11.7}_{-7.6}$ and $25.3^{+7.0}_{-7.3}$ respectively (Abbott et al.,
2021) in the source frame. The result shows better match with cWB and BW than
GW190412, because of its relatively higher masses for the reason described
earlier. Our testing of AWaRe on real O3 events, compared with cWB and BW,
demonstrates its robust performance. The framework’s ability to handle a
diverse range of events underscores its potential for broader application in
GW detection and detector characterization.

###  3.2 Tests on different waveform approximants

In this study, we demonstrate the performance of AWaRe, trained exclusively on
IMRPhenomXPHM waveforms, in reconstructing GW signals from three different
waveform families: IMRPhenomXPHM (Pratten et al., 2021), SEOBNRv4PHM (Ossokine
et al., 2020), and NRHybSur3dq8 (Varma et al., 2019). Since these waveform
families significantly differ in their higher order mode content, this test
serves as a strong indicator of the generalizability of our model to complex
waveform systematics outside our training set.

The top, middle, and bottom panels in Fig. 3 show the whitened strain data,
the respective injected waveforms from each family, and the mean
reconstructions by AWaRe for a GW190412-like injection sample (Abbott et al.,
2020b). The AWaRe reconstructions are shown in blue and the injected waveforms
are shown in red across all three waveform families. GW190412 was chosen in
particular because of the statistical evidence for the presence of higher
order modes, as previously discussed (Abbott et al., 2020b). We performed the
injections with SNR of 22, which is higher than the observed network SNR of
$19.0^{+0.2}_{-0.3}$ (Abbott et al., 2020b) in order to enhance the effect of
the available higher harmonics in the three waveform families. The shaded
regions represent the 50% and 90% confidence intervals (C.I.) of the
reconstructions.

For IMRPhenomXPHM, the model displays high reconstruction accuracy
($\mathcal{O}$ = 0.88 with injected waveform), which is expected, by virtue of
the fact that AWaRe is trained on IMRPhenomXPHM injections. For SEOBNRv4PHM
($\mathcal{O}$ = 0.71 with injected waveform), the amplitude near the merger
is underestimated, which could be attributed to the discrepancy in the higher
order mode content, which has its strongest influence where the signal is
loudest, i.e., near merger. The reconstruction of NRHybSur3dq8 ($\mathcal{O}$
= 0.87 with injected waveform) shows the full injection waveform is captured
within the 90% AWaRe bounds, and the mean reconstruction also closely follow
the injected waveforms indicating high reconstruction accuracy. These results
imply that the model’s training on a single waveform family (IMRPhenomXPHM)
generalizes well to other families, enhancing its robustness and
generalizability, which is crucial for improving GW detection in various
astrophysical scenarios.

###  3.3 Tests on masses beyond the training set

In order to further test the generalizability of our model, we tested AWaRe’s
performance on injections with primary mass beyond the training set. AWaRe was
initially trained using IMRPhenomXPHM waveform injections with component
masses $m_{1}$ and $m_{2}$ uniformly distributed in the range [5, 100] M⊙ in
the detector frame. To evaluate the model’s robustness, we tested it on
injections with primary mass, m1 uniformly distributed in the range 200 and
250 M⊙. In addition, we chose the mass ratio to be strictly $>$ 4, and the
inclination angle was uniformly sampled between 75 and 105 degrees. These
parameters were chosen to enhance the effect of higher-order modes in the
signals, providing a rigorous test of the model’s capabilities to reconstruct
complex waveform systematics in a parameter space outside the scope of our
model’s training.

The left plot in Fig. 4 presents the mismatch between the reconstructed
waveforms and the injected waveforms as a function of the injected SNR. The
primary y-axis on the left shows the mismatch values, while the x-axis
represents the injected SNR. The red dots correspond to the mean mismatch
values, with error bars representing the 1-$\sigma$ uncertainty. The blue line
indicates a fitted curve to the data points. To obtain the points and the
error bars in the mismatch plot, we injected the same waveform in 10 different
realizations of real LIGO O3b noise and then computed the mean and standard
deviation of the mismatches between the mean reconstructions from AWaRe and
the injected waveform. The main plot shows that the mismatch decreases as the
injected SNR increases, demonstrating that AWaRe performs better with higher
SNR signals. At lower SNRs, the mismatch values are higher and exhibit larger
error bars, indicating greater variability and less accurate reconstructions.
As the SNR increases beyond approximately 15, the mismatch values stabilize
and the error bars become smaller, suggesting more consistent and reliable
reconstructions.

The insets provide a detailed view of the time-domain reconstructions. The top
inset corresponds to an injected waveform with SNR - 12, while the bottom
inset shows a waveform with SNR - 22. In both cases, the gray lines represent
the whitened strain data, the red lines represent the whitened injected
waveforms, and the light blue bands show the 90% C.I. of the reconstructions.
These insets illustrate that the AWaRe can closely follow the injected
waveforms with the confidence intervals adequately capturing the uncertainty
in the reconstructions.

###  3.4 Tests on eccentric waveforms

In this study, we extend the evaluation of AWaRe to reconstruct eccentric
waveform injections generated with the EccentricTD waveform approximant (Tanay
et al., 2016). The injected signals have component masses of 30 and 30 M⊙, are
spinless, and have a SNR of 12. Notably, all waveforms in the training set had
zero eccentricity, providing a stringent test of the model’s ability to
generalize to waveforms with non-zero eccentricity, which introduce
modulations in amplitude and phase. The reconstruction of signals with non-
zero eccentricity using AWaRe has important implications in the future
application of our method for detecting eccentric BBH merger signals. This
work would add to recent works in the reconstruction and detection of
eccentric GW signals using modelled, unmodelled and semi-modelled methods
(Gadre et al., 2024; Roy, 2022).

The right plot of Fig. 4 illustrates the mismatch between the reconstructed
waveforms and the injected waveforms as a function of the initial eccentricity
(i.e., eccentricity at 20 Hz). The y-axis shows the mismatch values, while the
x-axis represents the initial eccentricity. Blue dots indicate the mean
mismatch values, with error bars representing the 1-$\sigma$ errors. The dark
blue line shows a fitted curve to the data points. To obtain the data points
and error bars, we injected the same waveform in 10 different realizations of
real LIGO O3b noise and then computed the mean and standard deviation as
described in the previous section. The plot reveals that while the mismatch
increases with higher eccentricities, indicating the greater challenge in
accurately reconstructing waveforms with significant eccentricity, AWaRe
nonetheless shows remarkable generalization capabilities. At lower
eccentricities, the mismatch values are lower with smaller error bars,
indicating more accurate and consistent reconstructions. Even at higher
eccentricities, although the mismatch values increase, the model still manages
to produce reliable reconstructions.

The insets provide a closer look at the time-domain reconstructions for two
specific eccentricities. The top inset, with an eccentricity of 0.11, shows
that AWaRe can closely follow the injected waveform, with the reconstruction
nearly indistinguishable from the original signal. In contrast, the bottom
inset, with an eccentricity of 0.31, while exhibiting more pronounced
deviations, still demonstrates the model’s ability to capture the essential
features of the waveform. The gray lines represent the whitened strain data,
the red lines indicate the whitened injected waveforms, and the blue bands
depict the 90% C.I. of the reconstructions. These insets highlight the model’s
capacity to handle the modulation in amplitude and phase introduced by non-
zero eccentricities.

The results presented here underscore AWaRe’s robust generalization
capabilities beyond its training set. The model’s ability to accurately
reconstruct waveforms with low to moderate eccentricities, and still produce
reliable results at higher eccentricities, demonstrates its potential for
broader applicability in gravitational wave data analysis. This robustness is
particularly important for detecting and analyzing signals from sources with
significant eccentricities, which provide valuable insights into the formation
and evolution of compact binary systems.

##  4 Conclusion

In this paper, we demonstrated the robustness and generalizability of AWaRe
for reconstructing GW signals from BBH mergers, including signals with higher-
order modes and eccentricities. Building on our previous study, we have
incorporated waveform reconstruction uncertainty in the our deep learning
framework, enhancing the robustness of its predictions. Our tests, conducted
on a diverse set of injections with parameters beyond the training set,
illustrate AWaRe’s capacity to handle a variety of complex waveform features.
Specifically, the model was able to reconstruct waveforms with higher masses,
significant mass ratios, varied waveform systematics and non-zero
eccentricities, despite being trained exclusively on non-eccentric waveforms
with much smaller primary masses and a fixed waveform approximant.

The ability of AWaRe to generalize beyond its training set has significant
implications for GW data analysis. This robustness enhances its utility in
real-world scenarios, where the variety of source parameters can be vast and
unpredictable. By accurately reconstructing waveforms with complex features
not present in the training set, AWaRe opens new avenues for the detection and
analysis of GW signals from diverse astrophysical sources. This generalization
capability is crucial for advancing our understanding of the universe’s most
extreme events and improving current GW search methodologies.

We plan to extend this work by adapting AWaRe into a full GW search and
reconstruction pipeline, operating at low latency. This is motivated by
AWaRe’s demonstrated ability to distinguish between actual GW signals and pure
noise, as shown in our previous work (Chatterjee & Jani, 2024). The robustness
of our model to outliers from the source parameters’ prior distributions
further builds confidence in its applicability for GW detection in future
observation runs. Additionally, AWaRe is uniquely suited to produce rapid
reconstructions of a wide variety of GW transients. While cWB and BayesWave
reconstructions are mostly optimized for burst-like sources, AWaRe can
generate reconstructions at millisecond latency for signals originating from a
much wider range of BBH masses. We also plan to perform robust tests of AWaRe
against glitches in our future work. This study is especially important to
characterize AWaRe’s ability to accurately distinguish real GW signals from
noise transients and perform high-fidelity reconstructions of astrophysical
waveforms. Furthermore, training AWaRe on synthetic glitches sampled from a
wide range of glitch classes will make it suitable for application in any
detector characterization framework and complement online GW search methods.

The authors would like to thank Thomas Dent and Alistair McLeod for their
helpful comments and suggestions. This research was undertaken with the
support of compute grant and resources, particularly the DGX A100 AI Computing
Server, offered by the Vanderbilt Data Science Institute (DSI) located at
Vanderbilt University, USA. This research used open public data obtained from
the Gravitational Wave Open Science Center (https://www.gw-openscience.org), a
service of LIGO Laboratory, the LIGO Scientific Collaboration and the Virgo
Collaboration. LIGO is funded by the U.S. National Science Foundation. Virgo
is funded by the French Centre National de Recherche Scientifique (CNRS), the
Italian Istituto Nazionale della Fisica Nucleare (INFN) and the Dutch Nikhef,
with contributions by Polish and Hungarian institutes. This material is based
upon work supported by NSF’s LIGO Laboratory which is a major facility fully
funded by the National Science Foundation.

## References

  * cWB (2021) 2021, Coherent Wave Burst public data release, https://gitlab.com/gwburst/public/analysis
  * Aasi et al. (2015) Aasi, J., Abbott, B. P., Abbott, R., et al. 2015, Classical and Quantum Gravity, 32, 074001, doi: 10.1088/0264-9381/32/7/074001
  * Abbott et al. (2016) Abbott, B. P., Abbott, R., Abbott, T. D., et al. 2016, Phys. Rev. Lett., 116, 061102, doi: 10.1103/PhysRevLett.116.061102
  * Abbott et al. (2020a) Abbott, R., Abbott, T. D., Abraham, S., et al. 2020a, Phys. Rev. Lett., 125, 101102, doi: 10.1103/PhysRevLett.125.101102
  * Abbott et al. (2020b) —. 2020b, Phys. Rev. D, 102, 043015, doi: 10.1103/PhysRevD.102.043015
  * Abbott et al. (2021) —. 2021, Phys. Rev. X, 11, 021053, doi: 10.1103/PhysRevX.11.021053
  * Abbott et al. (2023) Abbott, R., Abbott, T. D., Acernese, F., et al. 2023, Phys. Rev. X, 13, 041039, doi: 10.1103/PhysRevX.13.041039
  * Acernese et al. (2014) Acernese, F., Agathos, M., Agatsuma, K., et al. 2014, Classical and Quantum Gravity, 32, 024001, doi: 10.1088/0264-9381/32/2/024001
  * Ajith et al. (2011) Ajith, P., Hannam, M., Husa, S., et al. 2011, Phys. Rev. Lett., 106, 241101, doi: 10.1103/PhysRevLett.106.241101
  * Akutsu et al. (2019) Akutsu, T., Ando, M., Arai, K., et al. 2019, Nature Astronomy, 3, 35, doi: 10.1038/s41550-018-0658-y
  * Allen et al. (2012) Allen, B., Anderson, W. G., Brady, P. R., Brown, D. A., & Creighton, J. D. E. 2012, Phys. Rev. D, 85, 122006, doi: 10.1103/PhysRevD.85.122006
  * Amaro-Seoane et al. (2017) Amaro-Seoane, P., Audley, H., Babak, S., et al. 2017, Laser Interferometer Space Antenna.  https://arxiv.org/abs/1702.00786
  * Astone et al. (2018) Astone, P., Cerdá-Durán, P., Di Palma, I., et al. 2018, Phys. Rev. D, 98, 122002, doi: 10.1103/PhysRevD.98.122002
  * Babak et al. (2021) Babak, S., Hewitson, M., & Petiteau, A. 2021, LISA Sensitivity and SNR Calculations.  https://arxiv.org/abs/2108.01167
  * Calderón Bustillo et al. (2018) Calderón Bustillo, J., Salemi, F., Dal Canton, T., & Jani, K. P. 2018, Phys. Rev. D, 97, 024016, doi: 10.1103/PhysRevD.97.024016
  * Capano et al. (2023) Capano, C. D., Cabero, M., Westerweck, J., et al. 2023, Phys. Rev. Lett., 131, 221402, doi: 10.1103/PhysRevLett.131.221402
  * Chandra et al. (2022) Chandra, K., Bustillo, J. C., Pai, A., & Harry, I. W. 2022, Phys. Rev. D, 106, 123003, doi: 10.1103/PhysRevD.106.123003
  * Chatterjee & Jani (2024) Chatterjee, C., & Jani, K. 2024, Reconstruction of binary black hole harmonics in LIGO using deep learning.  https://arxiv.org/abs/2403.01559
  * Cornish & Littenberg (2015) Cornish, N. J., & Littenberg, T. B. 2015, Classical and Quantum Gravity, 32, 135012, doi: 10.1088/0264-9381/32/13/135012
  * Cornish et al. (2021) Cornish, N. J., Littenberg, T. B., Bécsy, B., et al. 2021, Phys. Rev. D, 103, 044006, doi: 10.1103/PhysRevD.103.044006
  * Cuoco et al. (2020) Cuoco, E., Powell, J., Cavaglià, M., et al. 2020, Machine Learning: Science and Technology, 2, 011002, doi: 10.1088/2632-2153/abb93a
  * Dax et al. (2023) Dax, M., Green, S. R., Gair, J., et al. 2023, Phys. Rev. Lett., 130, 171403, doi: 10.1103/PhysRevLett.130.171403
  * Finn (1992) Finn, L. S. 1992, Phys. Rev. D, 46, 5236, doi: 10.1103/PhysRevD.46.5236
  * Gabbard et al. (2022) Gabbard, H., Messenger, C., Heng, I. S., Tonolini, F., & Murray-Smith, R. 2022, Nature Physics, 18, 112, doi: 10.1038/s41567-021-01425-7
  * Gadre et al. (2024) Gadre, B., Soni, K., Tiwari, S., et al. 2024, Detectability of eccentric binary black holes with PyCBC and cWB pipelines during the third observing run of LIGO-Virgo-KAGRA.  https://arxiv.org/abs/2405.04186
  * George & Huerta (2018a) George, D., & Huerta, E. 2018a, Physics Letters B, 778, 64, doi: https://doi.org/10.1016/j.physletb.2017.12.053
  * George & Huerta (2018b) George, D., & Huerta, E. A. 2018b, Phys. Rev. D, 97, 044039, doi: 10.1103/PhysRevD.97.044039
  * Goodfellow et al. (2016) Goodfellow, I., Bengio, Y., & Courville, A. 2016, Deep Learning (MIT Press)
  * Harry et al. (2018) Harry, I., Bustillo, J. C., & Nitz, A. 2018, Phys. Rev. D, 97, 023004, doi: 10.1103/PhysRevD.97.023004
  * Hochreiter & Schmidhuber (1997) Hochreiter, S., & Schmidhuber, J. 1997, Neural Computation, 9, 1735
  * Jani et al. (2020) Jani, K., Shoemaker, D., & Cutler, C. 2020, Nature Astronomy, 4, 260, doi: 10.1038/s41550-019-0932-7
  * Kim et al. (2015) Kim, K., Harry, I. W., Hodge, K. A., et al. 2015, Classical and Quantum Gravity, 32, 245002, doi: 10.1088/0264-9381/32/24/245002
  * Klimenko et al. (2016) Klimenko, S., Vedovato, G., Drago, M., et al. 2016, Phys. Rev. D, 93, 042004, doi: 10.1103/PhysRevD.93.042004
  * Krastev (2020) Krastev, P. G. 2020, Physics Letters B, 803, 135330, doi: https://doi.org/10.1016/j.physletb.2020.135330
  * Krizhevsky et al. (2012) Krizhevsky, A., Sutskever, I., & Hinton, G. E. 2012, in Advances in Neural Information Processing Systems, ed. F. Pereira, C. Burges, L. Bottou, & K. Weinberger, Vol. 25 (Curran Associates, Inc.)
  * Lecun et al. (1998) Lecun, Y., Bottou, L., Bengio, Y., & Haffner, P. 1998, Proceedings of the IEEE, 86, 2278, doi: 10.1109/5.726791
  * McIsaac et al. (2023) McIsaac, C., Hoy, C., & Harry, I. 2023, Phys. Rev. D, 108, 123016, doi: 10.1103/PhysRevD.108.123016
  * Millhouse et al. (2018) Millhouse, M., Cornish, N. J., & Littenberg, T. 2018, Phys. Rev. D, 97, 104057, doi: 10.1103/PhysRevD.97.104057
  * Ossokine et al. (2020) Ossokine, S., Buonanno, A., Marsat, S., et al. 2020, Phys. Rev. D, 102, 044055, doi: 10.1103/PhysRevD.102.044055
  * Pratten et al. (2021) Pratten, G., García-Quirós, C., Colleoni, M., et al. 2021, Phys. Rev. D, 103, 104056, doi: 10.1103/PhysRevD.103.104056
  * Qiu et al. (2023) Qiu, R., Krastev, P. G., Gill, K., & Berger, E. 2023, Physics Letters B, 840, 137850, doi: https://doi.org/10.1016/j.physletb.2023.137850
  * Roy (2022) Roy, S. 2022, Phys. Rev. Res., 4, 033078, doi: 10.1103/PhysRevResearch.4.033078
  * Schäfer et al. (2020) Schäfer, M. B., Ohme, F., & Nitz, A. H. 2020, Phys. Rev. D, 102, 063015, doi: 10.1103/PhysRevD.102.063015
  * Schäfer et al. (2023) Schäfer, M. B., Zelenka, O. c. v., Nitz, A. H., et al. 2023, Phys. Rev. D, 107, 023021, doi: 10.1103/PhysRevD.107.023021
  * Schmidt et al. (2015) Schmidt, P., Ohme, F., & Hannam, M. 2015, Phys. Rev. D, 91, 024043, doi: 10.1103/PhysRevD.91.024043
  * Schmidt et al. (2024) Schmidt, S., Caudill, S., Creighton, J. D. E., et al. 2024, Searching for gravitational-wave signals from precessing black hole binaries with the GstLAL pipeline.  https://arxiv.org/abs/2403.17186
  * Skliris et al. (2022) Skliris, V., Norman, M. R. K., & Sutton, P. J. 2022, Real-Time Detection of Unmodelled Gravitational-Wave Transients Using Convolutional Neural Networks.  https://arxiv.org/abs/2009.14611
  * Tanay et al. (2016) Tanay, S., Haney, M., & Gopakumar, A. 2016, Phys. Rev. D, 93, 064031, doi: 10.1103/PhysRevD.93.064031
  * Varma et al. (2019) Varma, V., Field, S. E., Scheel, M. A., et al. 2019, Phys. Rev. D, 99, 064045, doi: 10.1103/PhysRevD.99.064045
  * Vaswani et al. (2017) Vaswani, A., Shazeer, N., Parmar, N., et al. 2017, CoRR, abs/1706.03762
  * Veitch et al. (2015) Veitch, J., Raymond, V., Farr, B., et al. 2015, Phys. Rev. D, 91, 042003, doi: 10.1103/PhysRevD.91.042003
  * Wadekar et al. (2024) Wadekar, D., Venumadhav, T., Roulet, J., et al. 2024, A new search pipeline for gravitational waves with higher-order modes using mode-by-mode filtering.  https://arxiv.org/abs/2405.17400


# Spacerini: Plug-and-play Search Engines with Pyserini and Hugging Face

Christopher Akiki  0000-0002-1634-5068 Leipzig University and ScaDS.AI ,
Odunayo Ogundepo  University of Waterloo ,  Aleksandra Piktus  Hugging Face
and Sapienza University ,  Xinyu Zhang  University of Waterloo ,  Akintunde
Oladipo  University of Waterloo ,  Jimmy Lin  University of Waterloo and
Martin Potthast  Leipzig University and ScaDS.AI

(2023)

###### Abstract.

We present Spacerini, a modular framework for seamless building and deployment
of interactive search applications, designed to facilitate the qualitative
analysis of large scale research datasets. Spacerini integrates features from
both the Pyserini toolkit and the Hugging Face ecosystem to ease the indexing
text collections and deploy them as search engines for ad-hoc exploration and
to make the retrieval of relevant data points quick and efficient. The user-
friendly interface enables searching through massive datasets in a no-code
fashion, making Spacerini broadly accessible to anyone looking to
qualitatively audit their text collections. This is useful both to IR
researchers aiming to demonstrate the capabilities of their indexes in a
simple and interactive way, and to NLP researchers looking to better
understand and audit the failure modes of large language models. The framework
is open source and available on GitHub: https://github.com/castorini/hf-
spacerini, and includes utilities to load, pre-process, index, and deploy
local and web search applications. A portfolio of applications created with
Spacerini for a multitude of use cases can be found by visiting
https://hf.co/spacerini.

information retrieval, search interface, text auditing

††ccs: Information systems Search interfaces††copyright:
acmcopyright††journalyear: 2023††doi: XX.XXXX/XXXXXX.XXXXXX††conference:
Proceedings of the 46th International ACM SIGIR Conference on Research and
Development in Information Retrieval; July 23–27, 2023; Taipei,
Taiwan††booktitle: Proceedings of the 46th International ACM SIGIR Conference
on Research and Development in Information Retrieval (SIGIR ’23), July 23–27,
2023, Taipei, Taiwan††price: 15.00††isbn: XXX-X-XXX-XXXX-X/XX/XX

##  1\. Introduction

The commodification of data has had a transformative effect on science in
general (Hey et al., 2009) and machine learning in particular (Mitchell et
al., 2022). The race to train ever-larger language models is so contingent
upon having access to immeasurable quantities of text (Hoffmann et al., 2022)
that datasets have become—as Bender et al. (2021) contend—“too large to
document” both technically (Bender et al., 2021) and methodologically (Jo and
Gebru, 2020). This often leads to a “train first, ask questions later”
approach to model training (Akiki et al., 2022), itself a concrete instance of
convenience experimentation (Krohs, 2012), an approach to research that is
modulated by the availability of a resource and ease of a method, rather than
the suitability thereof to the problem at hand. In this sense, Beaulieu and
Leonelli (2021) deem it important to differentiate between the availability of
data and its suitability, especially in light of the misconception that Web
data is representative of all human experiences and immune to the ever-
widening _digital divide_ (Leonelli, 2020). They point out that the divide not
only exists but actively limits this representativeness—both socio-economic
and socio-cultural—of the Web, which in turn reinforces biases of the
artifacts that leverage it (Bender et al., 2021).

Being unable to easily audit large datasets incentivizes researchers to
release models trained on data they do not truly understand (Mitchell et al.,
2022) leading to model behaviors that are hard to study, predict or trace
(Mitchell et al., 2022; Siddiqui et al., 2022; Akyurek et al., 2022). This is
especially problematic in light of the potential real-world harms that ensue
(Weidinger et al., 2021; Hutchinson et al., 2020; Founta et al., 2018; Fast et
al., 2016). Being able to properly understand the limitations of our datasets
is a necessary first step toward understanding the harmful biases and failure
modes of the artifacts that build upon them. Understanding the training data
is therefore a critical step in the process of auditing large language models
(Mökander et al., 2023).

It is from this vantage point that we initially developed Spacerini as an
open-source tool for the quick indexing and deployment of shareable search
engines, but have also since come to realize its potential in being useful for
an even wider audience interested in making their text artifacts searchable.
This includes IR students, Digital Humanists, Shared Tasks organizers, and
digital investigative journalists. We cover these use cases in more detail in
Section 4. Spacerini was designed as a modular framework that facilitates the
indexing, creation, and free hosting of graphical search interfaces for text
datasets. It helps streamline the process of auditing large datasets by
allowing users to effortlessly index their text collections and deploy them as
as search systems that can be accessed by anyone. It does so by standing on
the shoulders of battle-tested open-source libraries from the Castorini (Lin
et al., 2021) and Hugging Face (Abid et al., 2019; Lhoest et al., 2021)
ecosystems and enhancing the interoperability between them to enable quick
indexing and free deployment of search interfaces.

The key advantage of Spacerini is its ability to simplify the search process
and retrieve relevant data points from massive datasets with ease, allowing
researchers to conduct quick and efficient audits, while abstracting away all
the minutiae of indexing data or hosting services. We believe that this
provides an opportunity for collaboration and transparency in IR and NLP
research. With the creation and sharing of search indexes publicly,
practitioners, researchers and the general public can work together to
pinpoint problematic content, find duplicates, and identify biases in
datasets. Specific to the IR community, Spacerini allows a convenient and
universal approach to index sharing, which greatly facilitates the
demonstration of retrieval systems and the reproduction from other researcher
teams.

##  2\. Background

Large scale, predominantly web-mined text datasets have been proliferating in
NLP recently, giving rise to publications (Laurençon et al., 2022; Gao et al.,
2021; Ortiz Suárez et al., 2019; Raffel et al., 2022) which often contain
interesting analyses of the specific datasets being presented, however,
usually lack any comparison to existing resources beyond basic metrics such as
sizes of the datasets or languages they contain.

As discussed in Section 1, in the face of an increased scrutiny of the models
trained on datasets in question, the topic of data understanding and
governance has been gaining more traction, being accepted as an important part
of research. Efforts such as those of Mitchell et al. (2022) contribute
frameworks for more standardised and reproducible metrics and measurements of
datasets, and we position ourselves as a complementary continuation of their
work, focusing on a more curatorial and qualitative assessment that might not
readily fit under the umbrella of “measurements”. We therefore aim to fill the
gap in the evaluation landscape by facilitating qualitative, rather than
quantitative analysis of large scale data collections.

Similarly to the authors of Gradio (Abid et al., 2019), a Python package for
fast development of Machine Learning demos, we believe that the accessibility
of data and model analysis tools is crucial to building both the understanding
of and the trust in the underlying resources. The potential of relevance-based
interfaces to massive textual corpora, the creation of which can be
facilitated by leveraging toolkits such as Pyserini (Lin et al., 2021), has
previously been tapped into by the researchers at the Allen Institute of AI
who propose a C4 (Raffel et al., 2022) search
engine111https://c4-search.apps.allenai.org/. Similar interfaces have also
been found useful in more specialised domains, e.g. in COVID-related datasets
(Zhang et al., 2020), news quotes (Vuković et al., 2022), or medical
literature (Niezni et al., 2022). However, while these solutions are
undeniably useful, they remain very contextual: dataset-specific, and project-
specific. We believe Spacerini to be the first generalizable tool which
proposes an end-to-end pipeline automating the route from raw text to
qualitative analysis.

##  3\. Spacerini

Spacerini is a modular framework that integrates Pyserini with the Hugging
Face ecosystem to streamline the process of going from any Hugging Face text
dataset—either local or hosted on the Hugging Face Hub—to a search interface
driven by a Pyserini index that can be deployed for free on the Hugging Face
Hub.

In what follows, we deconstruct an example script222gradio-demo.py to showcase
the different features enabled by Spacerini. When run end-to-end, the script
pulls a dataset from the Hugging Face hub, pre-processes it, indexes it,
creates a gradio-based search interface and pushes that to the Web as a
Hugging Face Spaces demo. This is only meant as a feature-complete demo, and
we don’t expect most people to want to integrate every step into their
workflows, but rather to cherry-pick and decide what best to use depending on
context.

Figure 1.  Search interface for the Extreme Summarization (xsum) dataset
deployed on Hugging Face Spaces using Spacerini.

###  3.1. Loading Data

All our workflows are backed by the Hugging Face datasets library (Lhoest et
al., 2021), itself based on the extremely efficient Apache Arrow format.
Datasets is a mature library which provides a standardized interface to any
tabular dataset, in particular, to tens of thousands of community datasets
hosted on the Hugging Face Hub333https://hf.co/datasets. The datasets library
gives fine-grained control over the lifecycle of tabular datasets, which we
choose to abstract away through a set of opinionated data loading functions
that cover the use cases we deem relevant to information retrieval. We also
add new functionality, such as the ability to load any document dataset from
the ir_datasets library using for example the following one-liner to load MS
MARCO as a Hugging Face datasets.Dataset object using a function from the data
subpackage:

⬇

from spacerini.data import load_ir_dataset

hf_dset = load_ir_dataset("msmarco-passage")

We include wrappers to load database tables, pandas DataFrames (pandas
development team, 2020; Wes McKinney, 2010), and text datasets on disk, as
well as the ability to load any dataset either in memory-mapped mode or in
streaming mode: the former makes it possible to handle larger-than-memory
datasets, and the latter larger-than-disk datasets that can be streamed from a
remote location such as the Hugging Face Hub.

###  3.2. Pre-processing

Spacerini also provices a preprocess subpackage which offers a range of
customizable pre-processing options for preparing datasets. This module
includes a sharding utility that enables the partitioning of large datasets
into smaller, more manageable chunks for efficient parallel processing.

⬇

from spacerini.preprocesss import shard_dataset

shard_dataset(

hf_dataset=hf_dset,

shard_size="1GB",

column_to_index="text",

shards_paths="msmarco-shards",

)

###  3.3. Indexing

Spacerini’s index subpackage leverages Pyserini to provide very efficient
Lucene indexing and allow users to easily and quickly index large datasets,
either sharded in the pre-processing step, or any text format accepted by
Pyserini, and streaming text datasets, such as those returned by Spacerini’s
data subpackage. This subpackage also exposes several tokenization options
using existing language-specific analyzers
444https://lucene.apache.org/core/9_5_0/analysis/common/index.html as well as
Hugging Face subword tokenizers (MOI et al., 2022).

⬇

from spacerini.index import index_json_shards

index_json_shards(

shards_path="msmarco-shards",

index_path="app/index",

)

###  3.4. Template-based Search Interfaces

Having indexed a collection, one can easily spin up a frontend using the
frontend subpackage and one of many provided
templates555https://github.com/castorini/hf-spacerini/tree/main/templates.
These are built using
cookiecutter,666https://github.com/cookiecutter/cookiecutter a Python
templating library for software projects. We provide a few batteries-included
frontend templates based both the Gradio777https://gradio.app/ and Streamlit
888https://streamlit.io/ demo app frameworks, both of which are natively
supported by Hugging Face Spaces. Figure 1 showcases a search engine built
using one of our Streamlit templates.

⬇

from spacerini.frontend import create_app

cookiecutter_vars = {"dset_text_field": "text,

"metadata_field": "docid",

"space_title": "MS MARCO Search",

"local_app": "app"}

create_app(

template="gradio-vanilla",

extra_context_dict=cookiecutter_vars,

output_dir=".",

)

###  3.5. Deployment to Hugging Face Spaces

The local apps developed in the previous subsection can then be pushed to
Hugging Face Spaces and hosted there for free. One can then further customize
the running app from the browser, for example to add functionality not
provided by the chosen template.

⬇

from spacerini.frontend import create_space_from_local

create_space_from_local(

space_slug="msmarco-passage-search,

organization="spacerini,

space_sdk="gradio",

local_dir=LOCAL_APP,

delete_after_push=True,

)

###  3.6. Sharing Indexes as Hugging Face Datasets

Orthogonal to the workflow presented so far, is the ability to upload Lucene
indexes to the Hugging Face Hub using shareable dataset repositories and
enabling reproducible retrieval experiments.

⬇

from spacerini.index import push_index_to_hub

\parpush_index_to_hub(

dataset_slug="lucene-english-analyzer-msmarco",

index_path="index",

)

Any index can then just as easily be downloaded for local use:

⬇

from spacerini.index import load_index_from_hub

\parindex_path = load_index_from_hub("lucene-fr-analyzer-")

###  3.7. Search and Pagination

Search features are provided by the search subpackage and leverage the memory-
mapping feature of Arrow tables to load the entire table of results—no matter
how big—only materializing the specific shard that corresponds to the
requested result page.

⬇

from spacerini.search import result_indices, result_page

\parix = result_indices(

"Lorem Ipsum",

num_results=1_000,

INDEX_PATH,

)

\parlast_results_page = result_page(

hf_dset,

ix,

page=-1,

results_per_page=20,

).to_pandas()

##  4\. Use Cases and Examples

Spacerini is designed to enable qualitative analysis of large-scale textual
corpora without the need for extensive engineering work. The tool can be used
in dataset auditing campaigns, such as those carried out by Kreutzer et al.
(2022) or in data annotation efforts. It may be applied to inspect failures of
large scale language model predictions and find potential sources of memorized
generations. AI Ethics researchers can employ the tool to find evidence
supporting their hypotheses about the content of the models’ training
datasets.

Given its tight integration with Pyserini, Spacerini can also be leveraged by
IR researchers to experiment with modifications of their retrieval pipelines
in user studies (e.g. by exposing retrieval options such as BM25 and RM3
parameters in the the user interface) or to deploy demos of their working
prototypes. Reproducibility for IR experiments is further enhanced thanks to
the index sharing abilities introduced in Section 3.6.

Spacerini can also be leveraged by Digital Humanists, Archivists and
Librarians looking to index their collections. Indeed, GLAM (Galleries,
libraries, archives, and museums) collections are increasingly being made
available as datasets. Furthermore, there is a growing interest in the digital
humanities in training and using languages models, as demonstrated by the
success of projects such as the _BERT for Humanists Project_
999https://www.bertforhumanists.org/. This makes it especially pressing to
have easy access ways to critically examine the data that goes into such
models. In the context of the digital and computational humanities, indexing
data relevant to these efforts is often not an easy task, and often project-
based and contingent upon precarious funding arrangements. Having a project-
agnostic tool like Spacerini could prove valuable to this community and a
useful addition to toolkits such as the GLAM Workbench (Sherratt, 2021).

Given its low engineering barrier of entry, Spacerini can be a good addition
to IR courses with a practical component, where students are tasked with
developing search engines, by providing an easy-to-deploy interface for their
developed retrieval systems.

Spacerini can also be leveraged by organizers of shared tasks such as MIRACL
(Zhang et al., 2022) and Touché (Bondarenko et al., 2022), who want to help
participants explore the datasets without forcing them to download large
volumes of data. It can also be used as a platform for participants to deploy
working prototypes with a unified interface.

Spacerini can help data journalists and digital investigative journalists
index, explore, and understand open data, in a similar vein to the
functionality provided by the Aleph suite.101010https://docs.alephdata.org/
Providing technical tools to data journalists is a crucial in uncovering
matters of public interest, as was evident by role played by the collaborative
use of the Neo4j graph database in unraveling the corrupt network surrounding
tax havens (Díaz-Struck and Cabra, 2018).

Finally, three features of Hugging Face Spaces make them especially attractive
for users: (1) they can leverage private datasets, meaning that one can
provide search access to a dataset without sharing the underlying data,(2)
they can be seamlessly embedded into HTML, specifically Gradio-based Spaces
which can be embedded as _Web Components_
111111https://developer.mozilla.org/en-US/docs/Web/Web_Components so that
users can easily integrate a Spacerini-based search feature into their own
sites121212For example: https://cakiki.github.io/search-engine/, and (3)
Gradio-based Spaces expose a FastAPI131313https://github.com/tiangolo/fastapi
endpoint that can be queried to access the functionality of the space, making
deployed search engines accessible through HTTP calls.

##  5\. Limitations and Future Plans

The main limitation of the off-the-shelf variant of Spacerini is the disk
space limit imposed by Hugging Face Spaces, which is currently set to 50
GB.141414https://huggingface.co/docs/hub/spaces-overview#hardware-resources
While not enough to accommodate entire corpora such as ROOTS or The Pile, such
large datasets are most often amalgamations of constitutents datasets which
can each be studied independently. This limit has no bearing on Spacerini
search apps deployed locally. Should users still want to get more disk space
for their Spaces-hosted indexes, they are welcome to apply for community
grants offered by Hugging Face.

Planned improvements to the library include automating the creation of dataset
cards (Or rather “index cards”) when pushing an index to the Hugging Face Hub,
more support for dense indexing features provided by Pyserini, as well as more
fine-grained support of document tokenization. We also look forward to
community contributions both to the codebase and to the frontend templates.

##  6\. Conclusion

We presented Spacerini, a modular framework that enables the quick and free
deployment and serving of template-based search indexes as interactive
applications. The need for such a tool is especially pressing as large
language models have come to consume inordinate amounts of text data,
reinforcing the need for a qualitative exploration and understanding of
datasets to assess them in a way that is impenetrable to quantitative analyses
alone.

Spacerini leverages features from both the Pyserini toolkit and the Hugging
Face ecosystem to facilitate the creation and hosting of user-friendly search
systems for text datasets. Users can easily index their collections and deploy
them as ad-hoc search interfaces, making the retrieval of relevant data points
a quick and efficient process. The user-friendly interface enables non-
technical users to effectively search massive datasets, making Spacerini a
valuable tool for anyone looking to audit their text collections
qualitatively. The framework is open-source and available on GitHub:
https://github.com/castorini/hf-spacerini.

Finally, we emphasize that Spacerini is a first step in the direction of
systematic dataset auditing, and more work is still needed to create
standardized structures that leverage tools such as Spacerini to properly
document the different axes of interest that are appropriate for a given usage
context.

## Acknowledgments

We are grateful to Daniel van Strien, Yuvraj Sharma, Omar Sanseviero, Julien
Chaumond, Lucain Pouget, Abubakar Abid, and Leandro von Werra for their
tireless support and invaluable advice throughout this project.

## References

  * (1)
  * Abid et al. (2019) Abubakar Abid, Ali Abdalla, Ali Abid, Dawood Khan, Abdulrahman Alfozan, and James Zou. 2019\.  Gradio: Hassle-Free Sharing and Testing of ML Models in the Wild.  _arXiv preprint arXiv:1906.02569_ (2019).
  * Akiki et al. (2022) Christopher Akiki, Giada Pistilli, Margot Mieskes, Matthias Gallé, Thomas Wolf, Suzana Ilic, and Yacine Jernite. 2022.  BigScience: A Case Study in the Social Construction of a Multilingual Large Language Model.  _CoRR_ abs/2212.04960 (2022).  https://doi.org/10.48550/arXiv.2212.04960 arXiv:2212.04960
  * Akyurek et al. (2022) Ekin Akyurek, Tolga Bolukbasi, Frederick Liu, Binbin Xiong, Ian Tenney, Jacob Andreas, and Kelvin Guu. 2022.  Towards Tracing Knowledge in Language Models Back to the Training Data. In _Findings of the Association for Computational Linguistics: EMNLP 2022_. Association for Computational Linguistics, Abu Dhabi, United Arab Emirates, 2429–2446.  https://aclanthology.org/2022.findings-emnlp.180
  * Beaulieu and Leonelli (2021) Anne Beaulieu and Sabina Leonelli. 2021.  _Data and Society: A Critical Introduction_.  Sage.
  * Bender et al. (2021) Emily M. Bender, Timnit Gebru, Angelina McMillan-Major, and Shmargaret Shmitchell. 2021.  On the Dangers of Stochastic Parrots: Can Language Models Be Too Big?. In _Proceedings of the 2021 ACM Conference on Fairness, Accountability, and Transparency_ (Virtual Event, Canada) _(FAccT ’21)_. Association for Computing Machinery, New York, NY, USA, 610–623.  https://doi.org/10.1145/3442188.3445922
  * Bondarenko et al. (2022) Alexander Bondarenko, Maik Fröbe, Johannes Kiesel, Shahbaz Syed, Timon Gurcke, Meriem Beloucif, Alexander Panchenko, Chris Biemann, Benno Stein, Henning Wachsmuth, Martin Potthast, and Matthias Hagen. 2022\.  Overview of Touché 2022: Argument Retrieval. In _Experimental IR Meets Multilinguality, Multimodality, and Interaction - 13th International Conference of the CLEF Association, CLEF 2022, Bologna, Italy, September 5-8, 2022, Proceedings_ _(Lecture Notes in Computer Science, Vol. 13390)_ , Alberto Barrón-Cedeño, Giovanni Da San Martino, Mirko Degli Esposti, Fabrizio Sebastiani, Craig Macdonald, Gabriella Pasi, Allan Hanbury, Martin Potthast, Guglielmo Faggioli, and Nicola Ferro (Eds.). Springer, 311–336.  https://doi.org/10.1007/978-3-031-13643-6_21
  * Díaz-Struck and Cabra (2018) Emilia Díaz-Struck and Mar Cabra. 2018.  _Uncovering International Stories with Data and Collaboration_.  Springer International Publishing, Cham, 55–65.  https://doi.org/10.1007/978-3-319-97283-1_6
  * Fast et al. (2016) Ethan Fast, Tina Vachovsky, and Michael S. Bernstein. 2016\.  Shirtless and Dangerous: Quantifying Linguistic Signals of Gender Bias in an Online Fiction Writing Community.  _CoRR_ abs/1603.08832 (2016).  arXiv:1603.08832 http://arxiv.org/abs/1603.08832
  * Founta et al. (2018) Antigoni-Maria Founta, Constantinos Djouvas, Despoina Chatzakou, Ilias Leontiadis, Jeremy Blackburn, Gianluca Stringhini, Athena Vakali, Michael Sirivianos, and Nicolas Kourtellis. 2018.  Large Scale Crowdsourcing and Characterization of Twitter Abusive Behavior.  _CoRR_ abs/1802.00393 (2018).  arXiv:1802.00393 http://arxiv.org/abs/1802.00393
  * Gao et al. (2021) Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, Shawn Presser, and Connor Leahy. 2021\.  The Pile: An 800GB Dataset of Diverse Text for Language Modeling.  https://doi.org/10.48550/ARXIV.2101.00027
  * Hey et al. (2009) Tony Hey, Stewart Tansley, Kristin Tolle, and Jim Gray. 2009\.  _The Fourth Paradigm: Data-Intensive Scientific Discovery_.  Microsoft Research.  https://www.microsoft.com/en-us/research/publication/fourth-paradigm-data-intensive-scientific-discovery/
  * Hoffmann et al. (2022) Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, Tom Hennigan, Eric Noland, Katie Millican, George van den Driessche, Bogdan Damoc, Aurelia Guy, Simon Osindero, Karen Simonyan, Erich Elsen, Jack W. Rae, Oriol Vinyals, and Laurent Sifre. 2022\.  Training Compute-Optimal Large Language Models.  _CoRR_ abs/2203.15556 (2022).  https://doi.org/10.48550/arXiv.2203.15556 arXiv:2203.15556
  * Hutchinson et al. (2020) Ben Hutchinson, Vinodkumar Prabhakaran, Emily Denton, Kellie Webster, Yu Zhong, and Stephen Denuyl. 2020.  Social Biases in NLP Models as Barriers for Persons with Disabilities. In _Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics_. Association for Computational Linguistics, Online, 5491–5501.  https://doi.org/10.18653/v1/2020.acl-main.487
  * Jo and Gebru (2020) Eun Seo Jo and Timnit Gebru. 2020.  Lessons from Archives: Strategies for Collecting Sociocultural Data in Machine Learning. In _Proceedings of the 2020 Conference on Fairness, Accountability, and Transparency_ (Barcelona, Spain) _(FAT* ’20)_. Association for Computing Machinery, New York, NY, USA, 306–316.  https://doi.org/10.1145/3351095.3372829
  * Kreutzer et al. (2022) Julia Kreutzer, Isaac Caswell, Lisa Wang, Ahsan Wahab, Daan van Esch, Nasanbayar Ulzii-Orshikh, Allahsera Tapo, Nishant Subramani, Artem Sokolov, Claytone Sikasote, Monang Setyawan, Supheakmungkol Sarin, Sokhar Samb, Benoît Sagot, Clara Rivera, Annette Rios, Isabel Papadimitriou, Salomey Osei, Pedro Ortiz Suarez, Iroro Orife, Kelechi Ogueji, Andre Niyongabo Rubungo, Toan Q. Nguyen, Mathias Müller, André Müller, Shamsuddeen Hassan Muhammad, Nanda Muhammad, Ayanda Mnyakeni, Jamshidbek Mirzakhalov, Tapiwanashe Matangira, Colin Leong, Nze Lawson, Sneha Kudugunta, Yacine Jernite, Mathias Jenny, Orhan Firat, Bonaventure F. P. Dossou, Sakhile Dlamini, Nisansa de Silva, Sakine Çabuk Ballı, Stella Biderman, Alessia Battisti, Ahmed Baruwa, Ankur Bapna, Pallavi Baljekar, Israel Abebe Azime, Ayodele Awokoya, Duygu Ataman, Orevaoghene Ahia, Oghenefego Ahia, Sweta Agrawal, and Mofetoluwa Adeyemi. 2022.  Quality at a Glance: An Audit of Web-Crawled Multilingual Datasets.  _Transactions of the Association for Computational Linguistics_ 10 (2022), 50–72.  https://doi.org/10.1162/tacl_a_00447
  * Krohs (2012) Ulrich Krohs. 2012\.  Convenience experimentation.  _Studies in History and Philosophy of Science Part C: Studies in History and Philosophy of Biological and Biomedical Sciences_ 43, 1 (2012), 52–57.  https://doi.org/10.1016/j.shpsc.2011.10.005 Data-Driven Research in the Biological and Biomedical Sciences On Nature and Normativity: Normativity, Teleology, and Mechanism in Biological Explanation.
  * Laurençon et al. (2022) Hugo Laurençon, Lucile Saulnier, Thomas Wang, Christopher Akiki, Albert Villanova del Moral, Teven Le Scao, Leandro Von Werra, Chenghao Mou, Eduardo González Ponferrada, Huu Nguyen, Jörg Frohberg, Mario Šaško, Quentin Lhoest, Angelina McMillan-Major, Gérard Dupont, Stella Biderman, Anna Rogers, Loubna Ben allal, Francesco De Toni, Giada Pistilli, Olivier Nguyen, Somaieh Nikpoor, Maraim Masoud, Pierre Colombo, Javier de la Rosa, Paulo Villegas, Tristan Thrush, Shayne Longpre, Sebastian Nagel, Leon Weber, Manuel Romero Muñoz, Jian Zhu, Daniel Van Strien, Zaid Alyafeai, Khalid Almubarak, Vu Minh Chien, Itziar Gonzalez-Dios, Aitor Soroa, Kyle Lo, Manan Dey, Pedro Ortiz Suarez, Aaron Gokaslan, Shamik Bose, David Ifeoluwa Adelani, Long Phan, Hieu Tran, Ian Yu, Suhas Pai, Jenny Chim, Violette Lepercq, Suzana Ilic, Margaret Mitchell, Sasha Luccioni, and Yacine Jernite. 2022.  The BigScience ROOTS Corpus: A 1.6TB Composite Multilingual Dataset. In _Thirty-sixth Conference on Neural Information Processing Systems Datasets and Benchmarks Track_.  https://openreview.net/forum?id=UoEw6KigkUn
  * Leonelli (2020) Sabina Leonelli. 2020\.  Scientific Research and Big Data.  In _The Stanford Encyclopedia of Philosophy_ (Summer 2020 ed.), Edward N. Zalta (Ed.). Metaphysics Research Lab, Stanford University.
  * Lhoest et al. (2021) Quentin Lhoest, Albert Villanova del Moral, Yacine Jernite, Abhishek Thakur, Patrick von Platen, Suraj Patil, Julien Chaumond, Mariama Drame, Julien Plu, Lewis Tunstall, Joe Davison, Mario Sasko, Gunjan Chhablani, Bhavitvya Malik, Simon Brandeis, Teven Le Scao, Victor Sanh, Canwen Xu, Nicolas Patry, Angelina McMillan-Major, Philipp Schmid, Sylvain Gugger, Clément Delangue, Théo Matussière, Lysandre Debut, Stas Bekman, Pierric Cistac, Thibault Goehringer, Victor Mustar, François Lagunas, Alexander M. Rush, and Thomas Wolf. 2021\.  Datasets: A Community Library for Natural Language Processing. In _Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing: System Demonstrations, EMNLP 2021, Online and Punta Cana, Dominican Republic, 7-11 November, 2021_ , Heike Adel and Shuming Shi (Eds.). Association for Computational Linguistics, 175–184.  https://doi.org/10.18653/v1/2021.emnlp-demo.21
  * Lin et al. (2021) Jimmy Lin, Xueguang Ma, Sheng-Chieh Lin, Jheng-Hong Yang, Ronak Pradeep, and Rodrigo Nogueira. 2021\.  Pyserini: A Python Toolkit for Reproducible Information Retrieval Research with Sparse and Dense Representations. In _Proceedings of the 44th Annual International ACM SIGIR Conference on Research and Development in Information Retrieval (SIGIR 2021)_. 2356–2362.
  * Mitchell et al. (2022) Margaret Mitchell, Alexandra Sasha Luccioni, Nathan Lambert, Marissa Gerchick, Angelina McMillan-Major, Ezinwanne Ozoani, Nazneen Rajani, Tristan Thrush, Yacine Jernite, and Douwe Kiela. 2022\.  Measuring Data.  _CoRR_ abs/2212.05129 (2022).  https://doi.org/10.48550/arXiv.2212.05129 arXiv:2212.05129
  * MOI et al. (2022) Anthony MOI, Nicolas Patry, Pierric Cistac, Pete, Funtowicz Morgan, Sebastian Pütz, Mishig, Bjarte Johansen, Thomas Wolf, Sylvain Gugger, Clement, Julien Chaumond, Lysandre Debut, François Garillot, Luc Georges, dctelus, JC Louis, MarcusGrass, Taufiquzzaman Peyash, 0xflotus, Alan deLevie, Alexander Mamaev, Arthur, Cameron, Colin Clement, Dagmawi Moges, David Hewitt, Denis Zolotukhin, and Geoffrey Thomas. 2022\.  _huggingface/tokenizers: Rust 0.13.2_.  https://doi.org/10.5281/zenodo.7298413
  * Mökander et al. (2023) Jakob Mökander, Jonas Schuett, Hannah Rose Kirk, and Luciano Floridi. 2023.  Auditing large language models: a three-layered approach.  https://doi.org/10.48550/ARXIV.2302.08500
  * Niezni et al. (2022) Danna Niezni, Hillel Taub-Tabib, Yuval Harris, Hagit Sason-Bauer, Yakir Amrusi, Dana Azagury, Maytal Avrashami, Shaked Launer-Wachs, Jon Borchardt, M Kusold, Aryeh Tiktinsky, Tom Hope, Yoav Goldberg, and Yosi Shamay. 2022\.  Extending the Boundaries of Cancer Therapeutic Complexity with Literature Data Mining.  _bioRxiv_ (2022).  https://doi.org/10.1101/2022.05.03.490286 arXiv:https://www.biorxiv.org/content/early/2022/10/30/2022.05.03.490286.full.pdf
  * Ortiz Suárez et al. (2019) Pedro Javier Ortiz Suárez, Benoît Sagot, and Laurent Romary. 2019\.  Asynchronous Pipeline for Processing Huge Corpora on Medium to Low Resource Infrastructures. In _7th Workshop on the Challenges in the Management of Large Corpora (CMLC-7)_ , Piotr Bański, Adrien Barbaresi, Hanno Biber, Evelyn Breiteneder, Simon Clematide, Marc Kupietz, Harald Lüngen, and Caroline Iliadi (Eds.). Leibniz-Institut für Deutsche Sprache, Cardiff, United Kingdom.  https://doi.org/10.14618/IDS-PUB-9021
  * pandas development team (2020) The pandas development team. 2020\.  _pandas-dev/pandas: Pandas_.  https://doi.org/10.5281/zenodo.3509134
  * Raffel et al. (2022) Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. 2022.  Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer.  _J. Mach. Learn. Res._ 21, 1, Article 140 (jun 2022), 67 pages.
  * Sherratt (2021) Tim Sherratt. 2021\.  _GLAM Workbench_.  https://doi.org/10.5281/zenodo.5603060
  * Siddiqui et al. (2022) Shoaib Ahmed Siddiqui, Nitarshan Rajkumar, Tegan Maharaj, David Krueger, and Sara Hooker. 2022\.  Metadata Archaeology: Unearthing Data Subsets by Leveraging Training Dynamics.  _CoRR_ abs/2209.10015 (2022).  https://doi.org/10.48550/arXiv.2209.10015 arXiv:2209.10015
  * Vuković et al. (2022) Vuk Vuković, Akhil Arora, Huan-Cheng Chang, Andreas Spitz, and Robert West. 2022.  Quote Erat Demonstrandum: A Web Interface for Exploring the Quotebank Corpus. In _Proceedings of the 45th International ACM SIGIR Conference on Research and Development in Information Retrieval_. ACM.  https://doi.org/10.1145/3477495.3531696
  * Weidinger et al. (2021) Laura Weidinger, John Mellor, Maribeth Rauh, Conor Griffin, Jonathan Uesato, Po-Sen Huang, Myra Cheng, Mia Glaese, Borja Balle, Atoosa Kasirzadeh, Zac Kenton, Sasha Brown, Will Hawkins, Tom Stepleton, Courtney Biles, Abeba Birhane, Julia Haas, Laura Rimell, Lisa Anne Hendricks, William S. Isaac, Sean Legassick, Geoffrey Irving, and Iason Gabriel. 2021.  Ethical and social risks of harm from Language Models.  _CoRR_ abs/2112.04359 (2021).  arXiv:2112.04359 https://arxiv.org/abs/2112.04359
  * Wes McKinney (2010) Wes McKinney. 2010\.  Data Structures for Statistical Computing in Python. In _Proceedings of the 9th Python in Science Conference_ , Stéfan van der Walt and Jarrod Millman (Eds.). 56 – 61.  https://doi.org/10.25080/Majora-92bf1922-00a
  * Zhang et al. (2020) Edwin Zhang, Nikhil Gupta, Raphael Tang, Xiao Han, Ronak Pradeep, Kuang Lu, Yue Zhang, Rodrigo Nogueira, Kyunghyun Cho, Hui Fang, and Jimmy Lin. 2020.  Covidex: Neural Ranking Models and Keyword Search Infrastructure for the COVID-19 Open Research Dataset. In _Proceedings of the First Workshop on Scholarly Document Processing_. Association for Computational Linguistics, Online, 31–41.  https://doi.org/10.18653/v1/2020.sdp-1.5
  * Zhang et al. (2022) Xinyu Zhang, Nandan Thakur, Odunayo Ogundepo, Ehsan Kamalloo, David Alfonso-Hermelo, Xiaoguang Li, Qun Liu, Mehdi Rezagholizadeh, and Jimmy Lin. 2022\.  Making a MIRACL: Multilingual Information Retrieval Across a Continuum of Languages.  https://doi.org/10.48550/ARXIV.2210.09984