aainstitutetext: Sino-French Institute of Nuclear Engineering and Technology,
Sun Yat-sen University, Zhuhai 519082, Chinabbinstitutetext: INPAC; Shanghai
Laboratory for Particle Physics and Cosmology; Key Laboratory for Particle
Astrophysics and Cosmology (MOE),
School of Physics and Astronomy, Shanghai Jiao Tong University, 800 Dongchuan
Road, Shanghai 200240, Chinaccinstitutetext: School of Physics, Sun Yat-Sen
University, 135 Xingang Xi Road, Guangzhou, 510275, Chinaddinstitutetext:
SPEIT (SJTU-Paris Elite Institute of Technology), Shanghai Jiao Tong
University, 800 Dongchuan Road, Shanghai, 200240, China

# Reconstruction of the event vertex in the PandaX-III experiment with
convolution neural network

Tao Li c,1 Yu Chen b,d,1 Shaobo Wang,111Corresponding author. b Ke Han b Heng
Lin b Kaixiang Ni a,c Wei Wang<EMAIL_ADDRESS><EMAIL_ADDRESS><EMAIL_ADDRESS><EMAIL_ADDRESS><EMAIL_ADDRESS><EMAIL_ADDRESS><EMAIL_ADDRESS>
###### Abstract

The tracks left by charged particles in a gaseous time projection chamber
(TPC) incorporate important information about the interaction process and
drift of electrons in gas. The electron diffusion information carried by the
tracks is an effective signature to reconstruct $z_{0}$, the vertex position
in drift direction at which the event takes place. In this paper, we propose
to reconstruct $z_{0}$ with convolution neural network (CNN) in the PandaX-III
experiment. A CNN model VGGZ0net is built and validated with Monte Carlo
simulation data. It gives $z_{0}$ with a 11 cm precision for the events above
2 MeV uniformly distributed along a drift distance of 120 cm, and then the
electron lifetime can be deduced. The energy resolution of detector is
significantly improved after the electron lifetime correction, i.e., from
10.1% to 4.0% FWHM at the Q-value of double beta decay of 136Xe for the
scenario with an electron lifetime of 6.5 ms. The CNN model is also
successfully applied to the experimental data of the PandaX-III prototype
detector for $z_{0}$ reconstruction.

##  1 Introduction

Neutrinoless double beta decay ($0\nu\beta\beta$) processes are sensitive
probes of physics beyond the standard model of particle physics. A discovery
of such a process in experiments would directly indicate the Majorana nature
of neutrinos, and the violation of lepton number conservation Avignone:2007fu
. Many collaborations are searching for the $0\nu\beta\beta$ with different
isotopes, such as 76Ge GERDA:2020xhi , 130Te CUORE:2022jto , and 136Xe
KamLAND-Zen:2022tow ; EXO-200:2019rkq  through analyzing the energy spectrum
around the Q-value ($Q_{\beta\beta}$), which is the total energy released by
the decay and carried by the two emitted electrons. In recent years, gaseous
Time Projection Chambers (TPCs) emerged as powerful tools for the study of
rare events, in particular concerning $0\nu\beta\beta$ searches Gomez-
Cadenas:2019ges . The track, representing the energy depositions along the
particle movement trajectory, can be recorded by TPCs, which is expected to
reveal a comprehensive detector response and enhance the sensitivity of
$0\nu\beta\beta$ detection through the topological analysis NEXT:2019gtz ;
galan2020topological ; li2021signal .

The PandaX-III collaboration Chen:2016qcd  has proposed the construction of a
high-pressure gaseous xenon TPC for $0\nu\beta\beta$ of 136Xe. As described in
Wang:2020owr , the Micromegas-based TPC will hold about 140 kg of xenon (90%
enriched 136Xe) at 10 bar, with an expected energy resolution of 3% Full Width
at Half Maximum (FWHM) at the $Q_{\beta\beta}$. Xenon is mixed with 1%
trimethylamine (TMA) to improve signal quality in Micromegas
gonzalez2015accurate . The use of TMA provides a stable operation of the
Micromegas at high pressure, as well as a reduction of the electron diffusion
and a quenching of the primary xenon scintillation light while simultaneously
displaying Penning effect penning1928einfluss . However, the event vertex
$z_{0}$ in the drift direction is not available gonzalez2015accurate  due to
the loss of scintillation light signal in such a xenon-TMA mixture. Without
$z_{0}$, the correction of electron lifetime is unavailable, which limits the
energy resolution of the detector in actual operations due to gas purity. In
addition, the radioactive background level from materials and radon, which are
almost near the readout plane and the cathode, cannot be identified as
background events by an effective fiducial volume selection, which further
limits the detector sensitivity.

Therefore, we focus on the $z_{0}$ reconstruction of events using a method of
the Convolution Neural Network (CNN), which has been widely used in particle
physics experiments for signal recognition and physical feature extraction
Aurisano:2016jvx ; MicroBooNE:2020hho ; Qiao:2018edn ; Kekic:2020cne . In the
PandaX-III TPC, as the ionization electrons drift toward the anode along the
drift field, the diffusion effect carries the $z_{0}$ information of events:
the longer drift distance a specific particle goes through, the more diffused
its track shows. Thus, the track features of different diffusion are perceived
by the CNN to realize the $z_{0}$ reconstruction.

The paper will be organized as follows. The geometry of the PandaX-III
detector as well as the correlation between $z_{0}$ and the electron diffusion
effect are presented in Section 2. The realization of CNN together with the
data set preparation is described in Section 3. The performance of $z_{0}$
reconstruction is evaluated through the simulated $0\nu\beta\beta$ and
$\gamma$ events, and an accuracy of about 11 cm is obtained. It is also
successfully used to correct the energy spectrum distortion due to the
electron loss during drift process. In Section 4, we verified the
effectiveness of the model using the experimental data of the PandaX-III
prototype detector Lin:2018mpd . A summary and outlook are given in Section 5.

##  2 The PandaX-III detector and event vertex $z_{0}$

The specific design of the PandaX-III detector is presented in Wang:2020owr .
As shown in Figure 1, the TPC will be a cylindrical active volume with a
height of 120 cm and a diameter of 160 cm, containing 140 kg enriched 136Xe
mixed with 1% TMA mixture at 10 bar operating pressure. The charge readout
plane will consist of a tessellation of fifty-two 20$\times$20 cm2 Micromegas
modules with a series of 3 mm pitch strips readout in $x$ and $y$ directions.
When a charged particle propagates through the xenon gas, it would deposit
energy along its trajectory by ionizing the xenon atoms. The primary ionized
electrons can drift to the readout plane along with the drift field and then
get amplified and collected by the Micromegas strips on $x$ and $y$
directions. The drift time of the electrons will provide a position
measurement in the $z$ direction. Therefore, both the energy deposition and
three-dimensional track could be recorded.

Figure 1: Illustration of the PandaX-III TPC geometry. The TPC has a
cylindrical active volume with a height of 120 cm and a diameter of 160 cm. An
example $0\nu\beta\beta$ event track is shown. The pink star is the vertex
$z_{0}$ at which the decay takes place and the green star $z_{c}$ is the
charge center position of the track. They are expressed as a distance referred
to the readout plane, and $z_{0}$ is replaced by $z_{c}$ as the reconstruction
target.

In the PandaX-III detector, the $z_{0}$ of events cannot be measured directly,
which is a crucial issue of the energy resolution considering a relatively
long drift distance of 120 cm. A portion of the ionized electrons will be
absorbed due to the gas impurities during the process of drifting, resulting
in a so-called “electron attachment effect” and characterized by the electron
lifetime huxley1974diffusion . The energy resolution can be corrected event by
event only if the electron lifetime and $z_{0}$ are both known. Thus, a
precise $z_{0}$ reconstruction is necessary for a high-pressure TPC aiming for
an ideal energy resolution. In this work, the position of event charge center
$z_{c}$ is selected as a reference to $z_{0}$ as shown in Figure 1, since
$z_{c}$ is a more general observable of the track than $z_{0}$ for different
kinds of events, for example, $z_{0}$ of $0\nu\beta\beta$ is the occurrence
point of decay while that of background $\gamma$ events is one of the two
endpoints of track.

Figure 2: The simulated tracks (projection in $xz$ plane) for the same
$0\nu\beta\beta$ event with different $z_{c}$ (shown in red triangles) of 16
cm, 66 cm, and 106 cm.

While drifting to the readout plane, ionization electrons diffuse transversely
and longitudinally which results in the blurring of the tracks. In the PandaX-
III TPC, the corresponding drift velocity $v_{drift}$ is 1.86 mm/$\mu$s in 10
bar xenon-(1%)TMA mixture by applying a drift field of 1 kV/cm, and the
transverse (longitudinal) diffusion coefficient is 1.0 (1.5)$\times$10-2 cm1/2
li2021signal . The diffusion could affect the energy and track reconstruction.
However, the degree of such dispersion could reveal its drift distance. As
shown in Figure 2, the same event with different $z_{c}$ within the detector
is simulated and a coarser track could be observed after a longer drift
distance. Hence, we propose an CNN-based method to reconstruct the $z_{c}$ and
estimate electron lifetime by extracting the electron diffusion information
from tracks.

##  3 CNN and network architectures

A CNN regression model, named _VGGZ0net_ , is built to predict the $z_{c}$ of
events. VGGZ0net is based on the VGG16 classification model proposed by the
Visual Geometry Group at Oxford University in 2014 simonyan2014very . In the
PandaX-III scenario, VGGZ0net is designed to solve a continuous regression
problem considering the continuous distribution of event location along the
drift distance of 120 cm.

###  3.1 Simulation and data production

Monte Carlo (MC) simulation datasets are generated using the REST framework
altenmuller2022rest  to optimize and validate the model. Both $0\nu\beta\beta$
events of 136Xe and $\gamma$ events of 2.615 MeV from 232Th decay chain are
simulated. The events are firstly produced by $restG4$ of REST, based on
Geant4 Agostinelli:2002hh . Then the detector responses, including electron
diffusion, electron attachment, energy smearing, and Micromegas readout plane
scheme, are simulated by $restManager$. More details of the simulation can be
found in li2021signal . Compared with the previous works, the simulation is
further optimized including the actual detector and electronics response of
the experimental configurations, including 5 MHz sampling rate, 5 $\mu$s
waveform shaping time, and 0.3 fC electronic noise level. The spatial
positions of the triggered channels provide the $x$ and $y$ positions, and the
sample points of the each waveform provides a relative $z$ position once given
drift velocity and relative drift time. Figure 3(a) shows the tracks of such
events on the $xz$ and $yz$ planes respectively. The energy of tracks on the
$xz$ and $yz$ plane is normalized to the total energy of obtained with the
plane respectively, and the tracks are then filled into red ($xz$) and green
($yz$) channels of the input RGB image to VGGZ0net, as shown in Figure 3(b).
The image format is 64$\times$64 pixels with a size of 3$\times$3 mm2 each.
The image size is 19.2 cm in each direction can ensure keeping the whole track
of $0\nu\beta\beta$ inside the image area.

(a)

(b)

Figure 3: Demonstration of the dataset production for the simulated
$0\nu\beta\beta$ event: (a) two-dimensional tracks on $xz$ (left) and $yz$
(right) planes; (b) the converted RGB image as the input of CNN: all the
pixels shown here are set to the maximum values just for better visualization.

The events with a total energy more than 2.0 MeV are selected as the dataset.
The images are labeled as $z_{c}$ of the principle track, which is a
continuous track with significant energy deposition. For both the
$0\nu\beta\beta$ and $\gamma$ events, some subordinate tracks may be
generated. In particular, the $\gamma$ events will generate multiple tracks
through Compton scattering. To simplify the data processing, the principle
track within each event is selected and its $z_{c}$ will be predicted by
VGGZ0net. Once the position of the principle tracks is derived, all the
subordinate tracks can be located through their relevant distance. A DBSCAN
clustering algorithm ester1996density  is used to pick out the principle track
of events as a preprocessing step. An example of $\gamma$ dataset is shown in
Figure. 4. We also produced different scenarios to study electron lifetime
reconstruction and energy correction. For the ideal scenario of zero-
attachment, the energy spectrum is smeared by a Gaussian function which
represents the energy resolution of 3.3% FWHM (Full-Width Half Maximum) at
$Q_{\beta\beta}$ of 136Xe (2.458 MeV). Electron lifetime is endowed with
introducing a coefficient of electron adsorption in gas based on the ideal
scenario.

Figure 4: An example of simulated $\gamma$ event 232Th, the tracks on $xz$
(left) and $yz$(right) plane are shown. The DBSCAN algorithm ester1996density
extracts the principle track in red.

###  3.2 Network architectures and performance

The structure of VGGZ0net is represented in Figure 5, consisting of 5 sets of
convolution blocks as the feature extraction part. The size of each layer is
shown in Figure 5 and LeakyReLU maas2013rectifier  is chosen as the activation
function of the convolution layers. Thereinto, batch normalization layer is
utilized to improve accuracy and speed up trainingbjorck2018understanding .
After the feature extraction part, a fully connected part is designed with 2
fully connected layers with LeakyReLU as the activation function, as well as a
final connected layer and a linear activation function. A mean square error
(MSE) loss function is adopted for the regression requirements. A stochastic
gradient descent method, Adam kingma2014adam , is chosen as the optimizer.

Figure 5: The Structure of VGGZ0net. Table 1: The tuning results of the optimal hyper-parameters in VGGZ0net. Hyper-parameters | tuning result
---|---
Label range of dataset | (-1, 1)
Batch size per epoch | 128
Negative slope coefficient of leakyReLU | 0.6
Learning rate of Adam optimizer | 7$\times$10-5

The model is first trained through a supervised learning process based on the
MC $0\nu\beta\beta$ events of zero-attachment scenario. The size of both the
training and validation datasets is 120,000 events. A grid search method was
conducted within four main hyper-parameters in VGGZ0net to optimize the model
performance. The optimized values of hyper-parameters used are shown in Table
1.

The results of VGGZ0net on $0\nu\beta\beta$ training and validation datasets
are shown in Figure 6, where the MSE loss on the training and validation
datasets decreases simultaneously along with the epochs. We defined the
regression error:

$\Delta z=\hat{z}_{c}-z_{c},$ (1)

where $z_{c}$ is from the label of events, and $\hat{z}_{c}$ is that
reconstructed by VGGZ0net. $\sigma(\Delta z)$ is the standard deviation of
$\Delta z$. As shown in Figure 6, both the training and the validation MSE
loss of VGGZ0net are continuously reduced during the model training process,
resulting in $\sigma(\Delta z)$ decreases gradually. The optimal model is
chosen conservatively around the $50_{th}$ epoch without overfitting.

Figure 6: Top: The evolution of the MSE loss function. Bottom: $\sigma(\Delta
z)$ evolution on validation dataset. The values of hyper-parameters in Table 1
is used, and the optimal model at around $50_{th}$ epoch is chosen.

Figure 7: Top: Using VGGZ0net model in a dataset of $0\nu\beta\beta$ with
$l_{e}$ of 1200 cm, the distribution of $\Delta z$ is approximately a Gaussian
distribution with zero mean and a $\sigma(\Delta z)$ of about 11 cm. Bottom:
The distribution of raw energy $E_{r}$ over $\hat{z}_{c}$ shows the energy
loss along the drift distance. $\hat{l}_{e}$ of 1217$\pm$30 cm is obtained
through fitting the energy band of $Q_{\beta\beta}$ within the red dotted box
region. The corresponding $\hat{\tau}_{e}$ is 6.54$\pm$0.16 ms.

The trained model is applied to the dataset of $0\nu\beta\beta$ with a defined
electron absorption distance $l_{e}$ of 1200 cm, which is equivalent to the
electron lifetime $\tau_{e}$ of 6.45 ms ($\tau_{e}=l_{e}/v_{drift}$).
$\sigma(\Delta z)$ obtained by VGGZ0net is about 11 cm of the events uniformly
distributed within 120 cm distance, as shown in Figure 7 Top. The event energy
distribution is plotted over the reconstructed $\hat{z}_{c}$ shown in Figure 7
Bottom. An exponential function is introduced to obtain the electron
absorption distance and electron lifetime:

$E_{Q}=A_{0}\cdot e^{-\hat{l}_{e}/\hat{z}_{c}}=A_{0}\cdot
e^{-\hat{\tau}_{e}/\hat{t}_{c}},$ (2)

where the $E_{Q}$ is the fitted mean of detected $Q_{\beta\beta}$ peak in the
TPC. $A_{0}$ is an amplitude coefficient. The fitted electron absorption
distance $\hat{l}_{e}$ of 1217$\pm$30 cm conforms to the defined $l_{e}$
within one standard deviation. $\hat{t}_{c}$ is the drift time related to
$\hat{z}_{c}$, the corresponding electron lifetime $\hat{\tau}_{e}$ is 6.54
$\pm$ 0.16 ms. Given the charge normalized images, our model is insensitive to
the absolute energy of the event. Therefore, the trained model is also
applicable to the dataset of $\gamma$, the $\sigma(\Delta z)$ by the same
trained model is still about 11 cm. The reconstructed $\hat{l}_{e}$ is
1194$\pm$27 cm for the dataset of $\gamma$ with $l_{e}$ of 1200 cm. Thus, it
demonstrates the effectiveness and applicability of this method for the
$0\nu\beta\beta$ and $\gamma$ events in the same energy range.

###  3.3 Energy correction

The reconstructed $\hat{z}_{c}$ and $\hat{l}_{e}$ are applied to correct the
energy spectrum event by event for both the $0\nu\beta\beta$ and $\gamma$
datasets, and the performance is shown in Figure 8. $E_{r}$ is the raw energy
spectrum of the dataset with $l_{e}$ of 1200 cm, $E_{0}$ is that of the ideal
scenario. The difference between $E_{r}$ and $E_{0}$ represents the distortion
effect due to the limited electron lifetime. The shape of the $E_{r}$ spectrum
is distorted and the energy peak becomes blurred for both two datasets. The
energy spectrum $E_{r}$ could be corrected to $E_{c}$ by VGGZ0net, of which
the spectrum shape is comparable with that of $E_{0}$. In this way, by just
using the background events or calibration sources, the $z_{c}$ reconstruction
accuracy of 11 cm is enough to determine the electron lifetime and realize the
energy correction. It is beneficial to calibrate and monitor the detector
performance, especially during a long-term operation.

Figure 8: The energy spectra of simulated $0\nu\beta\beta$ (Top) and $\gamma$
from 232Th (Bottom) datasets. $E_{0}$ represents the energy spectrum of the
ideal scenario, $E_{r}$ is the raw energy spectrum affected by the introduced
electron lifetime, and $E_{c}$ is the energy spectrum corrected by VGGZ0net.

For different electron lifetime scenarios, the energy correction with the
trained VGGZ0net is studied. Seven datasets of $0\nu\beta\beta$ events under
different gas attachment coefficients are simulated and then undergoing the
above processes. The results are shown in Table. 2. $\sigma(\Delta z)$ of
VGGZ0net is found to be independent of the electron lifetime, indicating that
the model is insensitive to the carried energy of each event and only makes
predictions based on the diffusion of tracks. Both $\hat{l}_{e}$ and
$\hat{\tau}_{e}$ match the defined value well, and the energy resolution is
significantly improved after the correction based on VGGZ0net: the energy
resolution at $Q_{\beta\beta}$ is improved from 10.1% to 4.0% FHWM for the
scenario shown in Figure 8. The energy resolution of different scenarios
before the correction is not given in Table. 2, because the energy peak has
been very blurred under a lower electron lifetime. Effectively, the improved
energy resolution is a more valuable index to evaluate the performance of
VGGZ0net. As listed in Table. 2, the improved energy resolution of the 7
datasets is comparable to that of the ideal scenario. However, the energy
correction becomes more difficult to recover the energy resolution to the
ideal level when the electron lifetime is too small.

Table 2: The performance of electron lifetime reconstruction and energy correction based on VGGZ0net in different electron lifetime scenarios. The corrected energy resolution at $Q_{\beta\beta}$ is presented.  $l_{e}$ (cm) |  $\sigma(\Delta z)$ (cm) |  $\hat{l}_{e}$ (cm) |  $\hat{\tau}_{e}$ (ms) |  Corrected energy resolution at $Q_{\beta\beta}$ (%) FWHM
---|---|---|---|---
Infinity | 11 | - | - | 3.3
2000 | 11 | 2015 $\pm$ 55 | 10.83 $\pm$ 0.30 | 3.4
1800 | 11 | 1815 $\pm$ 53 | 9.76 $\pm$ 0.28 | 3.5
1600 | 11 | 1614 $\pm$ 42 | 8.68 $\pm$ 0.23 | 3.6
1400 | 11 | 1408 $\pm$ 33 | 7.57 $\pm$ 0.18 | 3.7
1200 | 11 | 1217 $\pm$ 30 | 6.54 $\pm$ 0.16 | 4.0
1000 | 11 | 1008 $\pm$ 25 | 5.42 $\pm$ 0.13 | 4.2
800 | 11 | 809 $\pm$ 20 | 4.35 $\pm$ 0.11 | 4.6

##  4 Application and performance test in experiment

Figure 9: Left: Structure of the PandaX-III prototype detector Lin:2018mpd ,
the maximun drift distance of detector is 78 cm, the 137Cs source (Green
point) is placed on the cathode plane. Right: Structure of readout plane, the
location of the Micromegas for data taking is in blue, the relative position
between the Micromegas and the projection of 137Cs source is shown.

The PandaX-III prototype detector is built to study the performance of high-
pressure xenon TPC. A detailed description of the detector components and
subsystems is given in ref Lin:2018mpd . As shown in Figure 9 Left, the
detector vessel has an inner volume of about 600 L and the active volume
inside the TPC is about 270 L, with a 78 cm drift distance and a readout plane
composed of seven 20$\times$20 cm2 Micromegas modules. The detector has
operated in different working gases and detected several calibration sources
Lin:2018mpd . For the validation of VGGZ0net, one of the seven Micromegas at
the bottom right of the readout plane was selected, as shown in Figure 9
Right. The detector was filled with 2 bar Ar-2.5% isobutane. We applied -20 kV
high voltage on the cathode to form a drift electric field of 256.4 V/cm.
Under this condition, the transverse (longitudinal) diffusion coefficient is
4.5 (2.1)$\times$10-2 cm1/2, and $v_{drift}$ is 30.3 mm/$\mu$s. A 137Cs source
was placed on the cathode and below the center Micromegas. At the same time,
the internal calibration gaseous source 83mKr was injected into the detector
with the gas circulation through the gas handling system, generating the
uniformly distributed, 41.5 keV monoenergetic events with a rate of about 1
Hz. The strip signals of Micromegas are collected and digitized by a
commercial waveform-sampling electronic system, AsAd/CoBo pollacco2018get ,
with a sampling rate of 5 MHz and a record length of 512 samples,
corresponding to a time window of 102.4 $\mu$s and a time bin of 0.2 $\mu$s.

Figure 10:  Examples of RGB images of 83mKr, the left one is from MC and the
right one is from the experimental data. All the pixels take the maximum
values and the two channels (R and G) are presented separately for better
visualization.

Specific adjustments of the input image size and the VGGZ0net model are
performed in the dozens keV energy range. MC datasets of 83mKr are simulated
according to the configurations of data taking. The 83mKr event tracks of MC
and data are shown in Figure 10, and the tracks in R and G channels are
presented separately for better visualization. The 41.5 keV events of 83mKr
triggered about 5 channels in $x$ or $y$ direction and 20 time bins in $z$
direction, leading to a cluster-like electron track. Therefore, the size of
the RGB image is reduced to 32$\times$32 to increase the proportion of pixels
occupied by the short tracks, and we directly reserved the $x$($y$) strips and
sampling points as the pixels. Meanwhile, a preprocessing step is added to
convert the image size back to 64$\times$64 for matching the input size of
VGGZ0net.

Figure 11: The distribution of $\Delta z$ is approximately a Gaussian
distribution with $\sigma(\Delta z)$ of about 4 cm.

The model is first trained through MC data of 83mKr with a uniform
distribution along the drift distance. The distribution of $\Delta z$ is shown
in Figure 11 with a standard deviation $\sigma(\Delta z)$ is 4 cm. It is
better than that of MeV-scale events presented in Section 3.1 and could be
explained that the cluster-like tracks can better characterize the electron
diffusion, while the features will be affected by the zigzag of long tracks in
a higher energy range.

Figure 12: The distribution of detected charge $E_{r}$ along $\hat{z}_{c}$ of
experimental data. The 83mKr band of 41.5 keV is in the dashed box in red, and
the 137Cs of 32.2 keV is in the blue one.

The trained model is applied to the experimental data. The distribution of the
collected charge by detector along $\hat{z}_{c}$ is shown in Figure 12. The
events within the 83mKr band generally distribute uniformly in the whole
interval while a large number of the events within the 137Cs band of 32.2 keV
is near the cathode, which is consistent with the physical expectation. The
reconstructed $\hat{l}_{e}$ is 1169$\pm$117 cm through fitting the 83mKr band,
the fitting error is generally introduced by statistics of a limited number of
events. The equivalent $\hat{\tau}_{e}$ is 0.39$\pm$0.04 ms. Thus, the trained
VGGZ0net based on the MC dataset is effective to determine $z_{c}$ and
electron lifetime in the experimental data. As a note, the reconstruction may
have a possible bias due to the deviation between the data and MC, such as the
small differences in pressure and gas ratio. A more quantitative optimization
of the conformity can be done by tuning the detector response model in MC
based on the calibration data, which is left for future work.

##  5 Summary

In this paper, we report the event vertex $z_{0}$ reconstruction using CNN in
the high-pressure TPC of the PandaX-III experiment. The lack of $z_{0}$ will
directly affect the detector performance and the sensitivity for searching
$0\nu\beta\beta$. As the scintillation signal cannot be measured in the xenon-
TMA gas mixture, the diffusion information carried by the tracks is an
effective and unique signature to reconstruct $z_{0}$. VGGZ0net is built to
predict the $z_{0}$ of the events distributed continuously between the cathode
and the readout plane within the TPC. MC simulation datasets of
$0\nu\beta\beta$ events are prepared to train the CNN model and tune its
hyperparameters. The reconstruction accuracy of VGGZ0net can reach a standard
deviation of about 11 cm. On this basis, the electron lifetime can be derived
and the energy spectrum can be corrected for both the $0\nu\beta\beta$ and
$\gamma$ datasets. VGGZ0net is proven to be effective and the improvement of
energy resolution is significant. This method has been validated by the
experimental data of the PandaX-III prototype detector using the 83mKr
calibration run. In the future, we will use it to study the detector
performance and suppress the location-dependent background in the PandaX-III
experiment.

###### Acknowledgements.

This work is supported by the grants from National Natural Sciences Foundation
of China (No.11775142, and No.11905127). We thank the support from the Key
Laboratory for Particle Physics, Astrophysics and Cosmology, Ministry of
Education. This work is supported in part by the Chinese Academy of Sciences
Center for Excellence in Particle Physics (CCEPP).

## References

  * (1) F.T. Avignone, III, S.R. Elliott and J. Engel, _Double Beta Decay, Majorana Neutrinos, and Neutrino Mass_ , _Rev. Mod. Phys._ 80 (2008) 481 [0708.1033].
  * (2) GERDA collaboration, _Final Results of GERDA on the Search for Neutrinoless Double- $\beta$ Decay_, _Phys. Rev. Lett._ 125 (2020) 252502 [2009.06079].
  * (3) CUORE collaboration, _Search for Majorana neutrinos exploiting millikelvin cryogenics with CUORE_ , _Nature_ 604 (2022) 53.
  * (4) KamLAND-Zen collaboration, _First Search for the Majorana Nature of Neutrinos in the Inverted Mass Ordering Region with KamLAND-Zen_ , 2203.02139.
  * (5) EXO-200 collaboration, _Search for Neutrinoless Double- $\beta$ Decay with the Complete EXO-200 Dataset_, _Phys. Rev. Lett._ 123 (2019) 161802 [1906.02723].
  * (6) J.J. Gomez-Cadenas, F. Monrabal Capilla and P. Ferrario, _High Pressure Gas Xenon TPCs for Double Beta Decay Searches_ , _Front. in Phys._ 7 (2019) 51 [1903.02435].
  * (7) NEXT collaboration, _Demonstration of the event identification capabilities of the NEXT-White detector_ , _JHEP_ 10 (2019) 052 [1905.13141].
  * (8) J. Galan, X. Chen, H. Du, C. Fu, K. Giboni, F. Giuliani et al., _Topological background discrimination in the pandax-iii neutrinoless double beta decay experiment_ , _Journal of Physics G: Nuclear and Particle Physics_ 47 (2020) 045108.
  * (9) T. Li, S. Wang, Y. Chen, K. Han, H. Lin, K. Ni et al., _Signal identification with kalman filter towards background-free neutrinoless double beta decay searches in gaseous detectors_ , _Journal of High Energy Physics_ 2021 (2021) 1.
  * (10) X. Chen et al., _PandaX-III: Searching for neutrinoless double beta decay with high pressure 136Xe gas time projection chambers_, _Sci. China Phys. Mech. Astron._ 60 (2017) 061011 [1610.08883].
  * (11) S. Wang, _The TPC detector of PandaX-III Neutrinoless Double Beta Decay experiment_ , _JINST_ 15 (2020) C03052 [2001.01356].
  * (12) D. González-Díaz, V. Álvarez, F. Borges, M. Camargo, S. Cárcel, S. Cebrián et al., _Accurate $\gamma$ and mev-electron track reconstruction with an ultra-low diffusion xenon/tma tpc at 10 atm_, _Nuclear Instruments and Methods in Physics Research Section A: Accelerators, Spectrometers, Detectors and Associated Equipment_ 804 (2015) 8.
  * (13) F.M. Penning, _Über den einfluß sehr geringer beimischungen auf die zündspannung der edelgase_ , _Zeitschrift für Physik_ 46 (1928) 335.
  * (14) A. Aurisano, A. Radovic, D. Rocco, A. Himmel, M.D. Messier, E. Niner et al., _A Convolutional Neural Network Neutrino Event Classifier_ , _JINST_ 11 (2016) P09001 [1604.01444].
  * (15) MicroBooNE collaboration, _Convolutional neural network for multiple particle identification in the MicroBooNE liquid argon time projection chamber_ , _Phys. Rev. D_ 103 (2021) 092003 [2010.08653].
  * (16) H. Qiao, C. Lu, X. Chen, K. Han, X. Ji and S. Wang, _Signal-background discrimination with convolutional neural networks in the PandaX-III experiment using MC simulation_ , _Sci. China Phys. Mech. Astron._ 61 (2018) 101007 [1802.03489].
  * (17) NEXT collaboration, _Demonstration of background rejection using deep convolutional neural networks in the NEXT experiment_ , 2009.10783.
  * (18) H. Lin et al., _Design and commissioning of a 600 L Time Projection Chamber with Microbulk Micromegas_ , _JINST_ 13 (2018) P06012 [1804.02863].
  * (19) L.G.H. Huxley and R.W. Crompton, _Diffusion and drift of electrons in gases_ (1974).
  * (20) K. Simonyan and A. Zisserman, _Very deep convolutional networks for large-scale image recognition_ , _arXiv preprint arXiv:1409.1556_ (2014) .
  * (21) K. Altenmüller, S. Cebrián, T. Dafni, D. Díez-Ibáñez, J. Galán, J. Galindo et al., _Rest-for-physics, a root-based framework for event oriented data analysis and combined monte carlo response_ , _Computer Physics Communications_ (2022) 108281.
  * (22) GEANT4 collaboration, _GEANT4–a simulation toolkit_ , _Nucl. Instrum. Meth. A_ 506 (2003) 250.
  * (23) M. Ester, H.-P. Kriegel, J. Sander, X. Xu et al., _A density-based algorithm for discovering clusters in large spatial databases with noise._ , in _kdd_ , vol. 96, pp. 226–231, 1996.
  * (24) A.L. Maas, A.Y. Hannun, A.Y. Ng et al., _Rectifier nonlinearities improve neural network acoustic models_ , in _Proc. icml_ , vol. 30, p. 3, Citeseer, 2013.
  * (25) N. Bjorck, C.P. Gomes, B. Selman and K.Q. Weinberger, _Understanding batch normalization_ , _Advances in neural information processing systems_ 31 (2018) .
  * (26) D.P. Kingma and J. Ba, _Adam: A method for stochastic optimization_ , _arXiv preprint arXiv:1412.6980_ (2014) .
  * (27) E. Pollacco, G. Grinyer, F. Abu-Nimeh, T. Ahn, S. Anvar, A. Arokiaraj et al., _Get: A generic electronics system for tpcs and nuclear physics instrumentation_ , _Nuclear Instruments and Methods in Physics Research Section A: Accelerators, Spectrometers, Detectors and Associated Equipment_ 887 (2018) 81.


# Modelling Technical and Biological Effects in single-cell RNA-seq data with
Scalable Gaussian Process Latent Variable Models (GPLVMs)

Vidhi Lalchand  Aditya Ravuri  Emma Dann  Natsuhiko Kumasaka  Dinithi
Sumanaweera  Rik G.H. Lindeboom  Shaista Madad  Sarah A. Teichmann  Neil D.
Lawrence

###### Abstract

Single-cell RNA-seq datasets are growing in size and complexity, enabling the
study of cellular composition changes in various biological/clinical contexts.
Scalable dimensionality reduction techniques are in need to disentangle
biological variation in them, while accounting for technical and biological
confounders. In this work, we extend a popular approach for probabilistic non-
linear dimensionality reduction, the Gaussian process latent variable model
(Lawrence, 2004), to scale to massive single-cell datasets while explicitly
accounting for technical and biological confounders. The key idea is to use an
augmented kernel which preserves the factorisability of the lower bound
allowing for fast stochastic variational inference. We demonstrate its ability
to reconstruct latent signatures of innate immunity recovered in Kumasaka et
al. (2021) with 9x speed-up on training time. We further analyse a dataset of
blood cells from COVID-19 patients and demonstrate that this framework enables
to capture interpretable signatures of infection, while integrating data
across individuals and technical batches. Specifically, we explore COVID-19
severity as a latent dimension to refine patient stratification and capture
disease-specific gene expression signatures.

##  1 Introduction

The development of single-cell transcriptomics technologies (scRNA-seq) has
led to the discovery of new cell types in tissues and of cellular composition
changes associated with disease, aging or experimental perturbations.
Dimensionality reduction and unsupervised learning techniques are widely
adopted to identify patterns in scRNA-seq data, capturing complex biological
phenomena (e.g. cell differentiation). With datasets growing in size and
complexity (often including many samples generated across multiple
laboratories, conditions and individuals), models that resolve both technical
and biological confounders (e.g. batch effect, inter-individual variation,
proliferation signatures) are needed to learn interpretable representations of
cell states.

The Gaussian Process Latent Variable Model (GPLVM) has been previously used as
a non-linear probabilistic method for dimensionality reduction, trajectory
inference and noise reduction in single-cell data (Campbell and Yau (2015);
Buettner et al. (2015); Ahmed et al. (2019); Verma and Engelhardt (2020)).
This model allows for uncertainty in inferences and incorporation of
additional information on biological processes through priors, for example to
model spatio-temporal dependencies between cells (Svensson et al. (2018);
Velten et al. (2022)). However, existing implementations either lack the
ability to account for confounding covariates or are too inefficient for
application to large datasets.

In this work, we extend the GPLVM to attain scalability across large scRNA-seq
datasets through an augmented kernel function which jointly accounts for both
technical and biological covariates, and successfully demonstrate its ability
to capture biologically interpretable signatures from an innate immunity
dataset (Kumasaka et al., 2021) and a COVID infection dataset (Stephenson et
al., 2021).

##  2 Methods

###  2.1 Model

The GPLVM uses a Gaussian process (GP) to define a non-parametric mapping from
a set of latent variables to gene expression. The covariance function of the
GP controls the inductive biases of the mapping, controlling properties like
smoothness and periodicity. (See Rasmussen and Williams (2006) for an
extensive review on GPs). Let the gene expression data be represented by
$\\{\bm{y}_{n}\\}_{n=1}^{N}\equiv Y\in\mathbb{R}^{N\times D}$ where $N$
denotes the number of cells and $D$ the number of genes. We wish to capture
the most significant factors of variation in $Y$ by learning a $Q$ dimensional
latent encoding per cell, $\\{\bm{x}_{n}\\}_{n=1}^{N}\equiv
X\in\mathbb{R}^{N\times Q}$ where $Q<D$ provides dimensionality reduction.

The core component of this model is a sparse GP (Hensman et al., 2013) that
models latent effects, also allowing to include additional terms for random
effects associated with known technical or biological confounding covariates
(via a design matrix $\Phi$) as in Kumasaka et al. (2021),

$Y=\underbrace{F}_{\textrm{\tiny{{Sparse
GP}}}}+\underbrace{\Phi}_{\textrm{\tiny{design
matrix}}}\times\underbrace{B}_{\textrm{\tiny{random
effects}}}+\underbrace{\bm{\epsilon}}_{\textrm{\tiny{noise model}}}.\hskip
5.69054pt$ (1)

The model above renders the likelihood non-factorisable (see section B.1),
preventing truly scalable inference through mini-batching. In the sections
below, we interpret the additive model above as a GP with an augmented kernel
function that jointly models the process $F+\Phi B$. Then, replacing this
joint GP with a sparse GP makes the formulation amenable to stochastic
variational inference (SVI),

$Y=\underbrace{F+\Phi\times B}_{\textrm{\tiny{GP with augmented
kernel}}}+\underbrace{\bm{\epsilon}}_{\textrm{\tiny{noise model}}}.\hskip
5.69054pt$ (2)

Below sections will unpack each component of the generative model and present
the formulation amenable to SVI.

####  2.1.1 Sparse GP

The sparse GP component is a stack of $D$ independent Gaussian processes
$F\equiv\\{\boldsymbol{\bm{f}}_{d}\\}_{d=1}^{D}$, with each column modelling a
particular gene expression profile,

$\boldsymbol{\bm{f}}_{d}\sim\mathcal{N}(\mu_{f}\mathbb{I}_{N},\sigma^{2}_{d}K_{nn}),$
(3)

where $\sigma^{2}_{d}$ controls the scaling of the data in column $D$. The
matrix $K_{nn}$ is the usual GP covariance matrix computed with a kernel
function of choice on latent inputs $X$.

To avoid the $\mathcal{O}(N^{3})$ scaling we introduce inducing variables per
gene dimension $\\{\bm{u}_{d}\\}_{d=1}^{D}\equiv
U,\boldsymbol{\bm{u}}_{d}\in\mathbb{R}^{M}$. The $\bm{u}_{d}$’s are
evaluations of $\bm{f}_{d}$ at inducing inputs $Z\in\mathbb{R}^{M\times Q}$
which lie in latent space (Hensman et al., 2013). Similar to $F$,

$\displaystyle\boldsymbol{\bm{u}}_{d}$
$\displaystyle\sim\mathcal{N}(\boldsymbol{\bm{0}},\sigma^{2}_{d}K_{mm})$ (4)
$\displaystyle\boldsymbol{\bm{f}}_{d}|\boldsymbol{\bm{u}}_{d}$
$\displaystyle\sim\mathcal{N}({K}_{nm}K_{mm}^{-1}\bm{u}_{d},\sigma^{2}_{d}(K_{nn}-K_{nm}K_{mm}^{-1}K_{mn})).$

Further, $K_{nn},K_{mm}$ and $K_{nm}$ represent covariance matrices computed
on latent inputs $(X)$, inducing inputs $(Z)$ and $X,Z$ respectively. The data
noise model is given by,
$\\{\epsilon_{n}\\}_{n=1}^{N}\equiv\bm{\epsilon}\in\mathbb{R}^{N\times D}$,
where $\epsilon_{nd}\sim\mathcal{N}(0,\sigma^{2}_{y})$.

####  2.1.2 Augmenting the Kernel Function

This section describes a formulation where we interpret $F+\Phi B$ as a joint
process (a sum of Gaussian processes). Let, $\tilde{F}=F+\Phi B$ where the
columns
$\tilde{\bm{f}_{d}}\sim\mathcal{N}(\mu_{f}\mathbb{I}_{N}+\Phi\zeta_{d},K_{nn}+\nu\Phi\Phi^{T})$
are distributed identically ($\zeta_{d}$ is identical across columns $D$). We
use straight-forward identities to derive the mean and covariance for the
joint process,

$\displaystyle\mathbb{E}(\tilde{\bm{f}_{d}})$
$\displaystyle=\mathbb{E}(\bm{f}_{d})+\mathbb{E}(\Phi
B_{d})=\mu_{f}\mathbb{I}_{N}+\Phi\zeta_{d}$ (5)
$\displaystyle\textrm{Cov}(\tilde{\bm{f}_{d}})$
$\displaystyle=\textrm{Cov}(\bm{f}_{d})+\textrm{Cov}(\Phi
B)=K_{nn}+\nu\Phi\Phi^{T},$ (6)

where we assume a constant mean $\mu_{f}\in\mathbb{R}$ for the $\bm{f}$
process, the design matrix $\Phi$ with covariates is specified and $\zeta_{d}$
encapsulates the mean of random effects $B$. The mean parameters $\mu_{f}$ and
$\zeta_{d}$ are model hyperparameters learnt during training. The expression
matrix $Y$ is driven by this joint process $\tilde{F}$ with columns
$\bm{\tilde{f}_{d}}$ distributed as individual Gaussian processes,

$\displaystyle
p(\tilde{F})=\prod_{d=1}^{D}p(\tilde{\bm{f}}_{d})=\prod_{d=1}^{D}\mathcal{N}(\Phi\bm{\zeta}_{d},$
$\displaystyle K_{nn}+\nu\Phi\Phi^{T}),$ (7)

where the term $\nu\Phi\Phi^{T}$ contributes a linear kernel term added to the
canonical kernel $K_{nn}$ driven by the kernel function described below. The
parameter $\nu\in\mathbb{R}$ controls the scaling over the linear kernel. We
assume that the gene expression is affected by cell cycle and other biological
effects. These effects are not known a priori, and have to be inferred from
the data. For $Q$-dimensional latent vectors
$\bm{x}_{1},\bm{x}_{2}\in\mathbb{R}^{Q}$, the kernel function designed to
capture these effects is given by the product of a periodic and SE-ARD
($k_{rbf}$) kernel (each acting on non-overlapping dimensions of the $Q$
dimensional latent variable), where the periodic term models the cell cycle
effect and the $k_{rbf}$ captures the other effects,

$\displaystyle k_{\tilde{f}}(\bm{x},\bm{x}^{\prime})$
$\displaystyle=\sigma^{2}_{f}\exp\left\\{\frac{-2\sin^{2}(|\bm{x}_{1}-\bm{x}_{1}^{\prime}|/2)}{l_{1}^{2}}\right\\}\times\exp\left\\{-\sum_{q=2}^{Q}\frac{(\bm{x}_{q}-\bm{x}^{\prime}_{q})^{2}}{2l_{q}^{2}}\right\\}+\nu\Phi\Phi^{T}$
(8) $\displaystyle=k_{per}\times k_{rbf}+k_{lin}.$ (9)

The hyperparameters for the mean and covariance are collected in
$\bm{\theta}=\\{\sigma^{2}_{f},\\{l_{q}\\}_{q=1}^{Q}\nu,\mu_{f},\zeta_{d}\\}.$

Let $\tilde{K}_{nn}=K_{nn}+\nu\Phi\Phi^{T}$; it is important to clarify how
the cross-covariance terms $\tilde{K}_{nm}$ and $\tilde{K}_{mm}$ are computed
in the augmented formulation. The inducing inputs/locations $Z$ now have to
incorporate additional dimensions to account for linear kernel term. Hence,
they have dimension $M\times(Q+\textrm{col. dim}(\Phi))$. One can think of
this as a partitioning of $Z=[\underbrace{Z_{per}|Z_{rbf}}_{Q-dims}|Z_{lin}]$
where the first $Q$ columns feed into $k_{per}\times k_{rbf}$ and the columns
corresponding to the col. dim($\Phi$) are used to compute the linear kernel
term,

$\displaystyle\tilde{K}_{nm}$ $\displaystyle=K_{nm}+\nu\Phi Z_{lin}^{T}$ (10)
$\displaystyle\tilde{K}_{mm}$ $\displaystyle=K_{mm}+\nu Z_{lin}Z_{lin}^{T}$
(11)

where $K_{nm}$ and $K_{mm}$ have been computed on the first $Q$ columns
$Z_{per}|Z_{rbf}$.

####  2.1.3 Variational Formulation

The joint probabilistic model is given by,

$\displaystyle p(Y,\tilde{F},\tilde{U})$
$\displaystyle=p(Y|\tilde{F},\Sigma)p(\tilde{F}|\tilde{U})p(\tilde{U}|Z)$ (12)
$\displaystyle=\mathcal{N}(Y|F+\Phi\bm{\zeta},\Sigma)\mathcal{N}(\tilde{F}|\tilde{K}_{nm}\tilde{K}_{mm}^{-1}\tilde{U},\tilde{Q}_{nn})\mathcal{N}(\tilde{U}|0,\tilde{K}_{mm}),$

where
$\tilde{Q}_{nn}=\tilde{K}_{nn}-\tilde{K}_{nm}\tilde{K}_{mm}^{-1}\tilde{K}_{mn}$.
The marginal likelihood $p(Y)$ is intractable and we resort to variational
inference over $\tilde{U}$ and point estimation for latents $X$, model
hyperparameters $\bm{\theta}$ and variational parameters $Z$. The posterior
over unknowns is approximated by the variational distribution,

$\displaystyle p(\tilde{F},\tilde{U}|Y)$ $\displaystyle\approx
p(\tilde{F}|\tilde{U},X)q(\tilde{U}|Z)$ (13)
$\displaystyle=\prod_{d=1}^{D}[p(\tilde{\bm{f}_{d}}|\tilde{\bm{u}}_{d},X)q(\tilde{\bm{u}}_{d}|Z)],$
(14)

where $p(\tilde{\bm{f}_{d}}|\tilde{\bm{u}}_{d},X)$ is as in eq. 2.1.1 and
$q(\tilde{\bm{u}}_{d}|Z)=\mathcal{N}(\bm{m}_{d},S_{d})$ where $S_{d}$ is a
dense $M\times M$ covariance matrix. $\\{\bm{m}_{d},S_{d}\\}_{d=1}^{D}$ are
global variational parameters.

We use the standard stochastic variational evidence lower bound (ELBO) which
preserves factorisation across the data-points $N$, and the final bound used
to generate the results takes the form,

$\displaystyle\log p(Y)$
$\displaystyle\geq\sum_{n=1}^{N}\sum_{d=1}^{D}\Big{\\{}\log\mathcal{N}(y_{n,d}|\tilde{k}_{n}\tilde{K}^{-1}_{mm}\bm{m}_{d},\sigma^{2}_{y})-\dfrac{q_{n,n}}{2\sigma^{2}_{y}}-\dfrac{\textrm{Tr}(S_{d}\Lambda_{n})}{2\sigma^{2}_{y}}\Big{\\}}-\textrm{KL}(q(\tilde{U}|Z)||p(\tilde{U}|Z)),$
(15)

where
$\Lambda_{n}=\tilde{K}_{mm}^{-1}\tilde{k}_{mn}\tilde{k}_{nm}\tilde{K}^{-1}_{mm}$,
$\tilde{k}_{n}$ is the $n^{th}$ column of $\tilde{K}_{mn}$, and $q_{n,n}$ is
the $n^{th}$ diagonal entry of the matrix $Q_{nn}$ defined earlier. Note that
this is the multi-output lower bound used by Hensman et al. (2013) in a
regression context and by Lalchand et al. (2022) in an unsupervised context
with variational inference over $X$, whereas in here we learn a point-estimate
for $X$. We include our algorithm in Appendix C.

###  2.2 Data

The innate immunity dataset (Kumasaka et al. (2021)) comprises gene expression
profiles of 22,188 primary dermal fibroblasts from 68 donors exposed to two
stimulants to mimic innate immune response: (1) the synthetic dsRNA Poly(I:C)
for primary antiviral and inflammatory responses, and (2) Interferon-beta
(IFN-beta) for secondary antiviral response. Cells were processed in 128
plates representing the technical batch. The COVID dataset (Stephenson et al.
(2021)) consists of gene expression profiles of peripheral blood mononuclear
cells from a cohort of 107 patients with varying severities of COVID-19 and 23
healthy individuals. Samples have been collected and processed at three
different sites (Sanger, Cambridge, Newcastle). We sampled 5% (32,368 cells)
and 10% (54,941 cells) of the full dataset for dimensionality reduction
comparison analysis (Fig. 2\- 3) and latent severity analysis, respectively.
For the encoder dimensionality reduction experiment, we used the full dataset
(647,366 cells).

For all datasets, their gene expression counts were normalized to sum up to
10,000 counts per cell and $\log(x+1)$ transformed as per the standard
practice of preprocessing. For dimensionality reduction analysis, we selected
the 5000 most highly variable genes, as in Scanpy (Wolf et al., 2018).

###  2.3 Model Inference

Across all datasets, our model inference was done by initializing latent
variables by the $Q$-largest principal components of the data, and utilizing
expression of S, G2 and M state gene markers to initialise the cell cycle
latent variable. (See Appendix D for a detailed discussion on initialization).
The models were trained using learning rates of $5*10^{-3}\text{ to }10^{-2}$
for $\sim$50 epochs using batch sizes of $\sim$200-700 data points for SVI. We
observe that, training the model excluding the latent variables (i.e. the
sparse GP) for few epochs with high learning rate before jointly optimising
all parameters with a lower learning rate leads to good model performance.

##  3 Results

###  3.1 Reproducing Innate Immunity Analysis

We first tested our framework on a dataset previously analysed with the
additive GPLVM framework (eq. 1) used in Kumasaka et al. (2021), to
demonstrate the ability to capture validated latent variables at a fraction of
the speed. Our run was about nine times faster (30 min compared to a baseline
of 4.5 hours) and is able to reproduce key results described in Kumasaka et
al. (2021). Our GPLVM model can disentangle batch effects and the gene
expression signature related to the cell cycle (Fig. S3), which is captured
with the use of the periodic kernel (Figure 1A). In addition, latent variables
recover the two main axes of biological variation in this data: (1) response
to Poly(I:C) treatment, and (2) response to IFN-beta treatment from the naive
state, mimicking primary and secondary immune response (as known as response
pseudotime) (Fig. 1B).

Figure 1:  (A) Scatterplots of scaled expression of known cell cycle genes
(UBE2C, CDC6, FN1). The dotted curves show the posterior mean expression
estimates by GPLVM. Points are colored by predicted cell cycle phase. (B) UMAP
embedding of innate immunity dataset latent space learnt by GPLVM. Cells are
colored by treatment condition (pIC: Poly(I:C), IFN: interferon-beta) (left),
primary immune response pseudotime latent variable (center) and secondary
immune response pseudotime latent variable (right), as described by Kumasaka
et al. (2021). Latent dimensions for hidden technical effects (batch and plate
border effects) were excluded for UMAP embedding.

###  3.2 Modelling COVID Severity with GPLVM

The ability to correct for confounders with random effects makes the GPLVM an
amenable model for analysis of cohort-level scRNA-seq datasets.

We first assessed the ability to identify cell type clusters while accounting
for batch effects. We trained our model by encoding the site of collection and
sample ID in the design matrix for random effects. We compared our model to
PCA (dimensionality reduction without batch correction) and a conditional VAE
(cVAE) for dimensionality reduction (scVI Lopez et al. (2018)) trained by
encoding the same technical conditions as categorical covariates and the cell
cycle initialization as a continuous covariate. We visualize the results by
UMAP embedding of latent spaces and we quantify mixing between cell labels by
measuring the fraction of k-nearest neighbors of a cell (k=100) harboring the
same label (KNN purity). While the scVI model performs better on the cell type
clustering task, our model was able to recover broad cell type clusters and
minimize differences driven by batch effects (Fig. 2). In addition, the GPLVM
latent dimensions capture interpretable biological differences between cells,
which are not discernible using the latent dimensions from cVAE based models.
For example, we find latent variables that capture the cell type-specific gene
expression signatures (e.g. Platelets in LV1 and LV2, separation between B
cell and NK/T cell lineage in LV3) (Fig. 3, top). Moreover, Stephenson et al.
(2021) found increased differentiation of megakaryocyte (MK) cells
(megakaryopoiesis) and platelet activation associated to COVID. One of the
GPLVM latent variables was strongly associated with the expression of
megakaryopoiesis genes in platelets (Fig. 3, middle) and this variable
separates platelets from healthy and asymptomatic patients from those of
patients with severe disease (Fig. 3, bottom). In contrast, none of the scVI
latent variables were significantly associated with this signature.

Figure 2:  Cell type clustering and batch correction task on COVID dataset (a)
UMAP embedding of cells generated from latent dimensions (excluding batch and
cell cycle variables) learnt by PCA (left), scVI (middle), and scalable GPLVM
with random effects (right). Cells are colored by annotated cell type (top)
and batch (bottom) (b) Cumulative Distribution of k-nearest neighbor (KNN)
purity of cell type labels (top) and batch labels (bottom) in each latent
space.  Figure 3:  Scatterplots of four latent dimensions learnt with GPLVM
(left) and scVI (right). Points are colored by cell type.

Figure 4:  (middle) Barplot of Pearson correlation between latent variable
values and megakaryopoiesis gene expression signature. The expression of 7
marker genes (TUBB1, PPBP, PF4, TREML1, GP9, MYL9, F13A1) was aggregated by
taking the average expression of these genes subtracted with the average
expression of a background set of genes with matched expression across cells.
(bottom) Distribution of values for LVs (y-axis) in platelet cells grouped by
disease severity (x-axis). LVs most associated with megakaryopoiesis signature
with GPLVM (left) and scVI (right) are shown.

One powerful feature of the GPLVM is its ability to incorporate covariates of
interest in the latent dimensions to generate a meaningful dimensionality
reduction. In cohort studies, patients are often grouped by disease severity
using qualitative assessments, which can be inaccurate and do not fully
reflect the true spectrum of observed severities. Incorporating a latent
variable in the GPLVM to capture disease severity supports better patient
stratification and detection of gene expression signatures predictive of
clinical outcome. We demonstrate this feature by training the GPLVM on 54,941
cells from the COVID dataset with an additional continuous latent variable
initialized on disease severity (converted to an integer encoding). After
training, the GPLVM refines disease severity assignment in a meaningful and
cell-specific manner: cells from healthy individuals and asymptomatic infected
patients are placed closer together in latent severity than cells from
symptomatic infections (Fig. 5A). Symptomatic cases are placed along a more
fine grained spectrum of disease severity, which could be helpful to better
classify cases. For example, in the moderate cases, we find that the predicted
severity is significantly correlated with the days since onset of symptoms
($R^{2}=0.25,\text{p-val}<2.2e-16$, Fig. S4). In addition, we can use the
generative model to predict genes that are most affected by changes in disease
severity, while minimizing the effect of confounders. This is done by
inputting a uniformly varying disease severity (with all other inputs set to
be constant) to the sparse GP generative model, which outputs predicted gene
expression for different levels of severity. Amongst the top 20 genes most
affected by changes in severity, we find genes associated with interferon
signalling (IRF7, ISG15), which is a recognized signature of COVID disease
(Stephenson et al. (2021)), markers of viral entry (HSPA5), as well as several
cell type markers (e.g. KLRK1, TRBC1, CD7, KLRD1), potentially reflecting
shifts in cell abundances in response to infection (Fig. 5B).

Figure 5:  (A) Violin plots of latent COVID severity inferred by GPLVM for
each patient in COVID study. Patients are ordered by the median latent
severity (point) and colored by reported disease severity (worst clinical
outcome). (B) Average expression of genes identified as variable across
severities in perturbation experiment.

##  4 Scaling up to Large Scale Dataset with Amortized Inference

With datasets now collecting up to millions of cells, our SVI formulation can
be combined with strategies such as amortized inference for application to
massive scale datasets. We demonstrate scalability to very large datasets on
the full COVID infection dataset ($>$ 600k cells) using neural network based
encoders as variational approximations for GPLVM latent variables.

In this version of the GPLVM, we variationally integrate out the latent
variables $X$, with a standard normal prior $p(X)$, the mean and variance of
the variational Gaussian distribution are parameterised as outputs of
individual neural networks $G_{\phi_{1}}$ and $H_{\phi_{2}}$ with network
weights $\phi_{1}$ and $\phi_{2}$. The network weights are shared across all
the cells ($N$) enabling amortised learning (Bui and Turner, 2015).

$q(X)=\prod_{n=1}^{N}\mathcal{N}(\bm{x}_{n};G_{\phi_{1}}(\bm{y}_{n}),\text{diag}(H_{\phi_{2}}(\bm{y_{n}})))$
(16)

This function is usually referred to as the back-constraint or an encoder and
its parameters are global, i.e. shared between all the data points and updated
in each iteration. This allows for fast amortised inference as it precludes
from learning free-form mean and standard deviations for each individual data
point (cell). In other words, the number of variational (or model) parameters
does not grow with the number of data points. We choose to model the
covariance of each $\bm{x}_{n}$ a diagonal covariance as this leads to better
behaviour while training (as opposed to when a dense covariance is allowed).

Figure 6 shows results of a run using this variational set-up on the full
COVID dataset comprising of over 600K cells, demonstrating that our model can
be scaled up to very large datasets while learning meaningful biological
representations such as broad cell types.

Figure 6:  Top two latent dimensions (selected by inverse lengthscales)
learned using a GPLVM fit on the full COVID dataset colored by cell types and
site information.

##  5 Summary

This work introduces an augmented kernel function to jointly model known and
unknown technical and biological covariates in scRNA-seq datasets, with a
formulation amenable for scale-up using stochastic variational inference. With
datasets now collecting up to millions of cells, our SVI formulation can be
combined with strategies such as amortized inference for application to
massive scale datasets. We demonstrate scalability to very large datasets on
the full COVID infection dataset ($>$ 600k cells) using neural network based
encoders as variational approximations for GPLVM latent variables. We envision
this model will have broad application to population-scale scRNA-seq studies,
where complex confounders and inter-individual variability need to be
accounted for, as shown in our analysis on COVID severity.

##  6 Code availability

The model and code to reproduce the analyses presented here is available at
https://github.com/InfProbSciX/gplvm_scrna.Our software is implemented in
GPyTorch and can be integrated into existing libraries of probabilistic
single-cell data analysis (https://scvi-tools.org/) for accessibility to the
wider community.

##  7 Acknowledgements

This work was supported by the Wellcome Trust grant reference: 221052/Z/20/Z,
221052/A/20/Z, 221052/B/20/Z, 221052/C/20/Z, and 221052/E/20/Z.

## References

  * Ahmed et al. (2019) Sumon Ahmed, Magnus Rattray, and Alexis Boukouvalas.  Grandprix: scaling up the Bayesian GPLVM for single-cell data.  _Bioinformatics_ , 35(1):47–54, 2019.
  * Buettner et al. (2015) Florian Buettner, Kedar N. Natarajan, F. Paolo Casale, Valentina Proserpio, Antonio Scialdone, Fabian J. Theis, Sarah A. Teichmann, John C. Marioni, and Oliver Stegle.  Computational analysis of cell-to-cell heterogeneity in single-cell RNA-sequencing data reveals hidden subpopulations of cells.  _Nature Biotechnology_ , 33(2):155–160, February 2015.  ISSN 1546-1696.  doi: 10.1038/nbt.3102.
  * Bui and Turner (2015) Thang D. Bui and Richard E. Turner.  Stochastic variational inference for Gaussian process latent variable models using back constraints.  In _Black Box Learning and Inference NIPS workshop_ , 2015.
  * Campbell and Yau (2015) Kieran Campbell and Christopher Yau.  Bayesian Gaussian process latent variable models for pseudotime inference in single-cell rna-seq data.  _bioRxiv_ , page 026872, 2015.
  * Hensman et al. (2013) James Hensman, Nicoló Fusi, and Neil D. Lawrence.  Gaussian processes for big data.  _Proceedings of the Twenty-Ninth Conference on Uncertainty in Artificial Intelligence (UAI2013)_ , 2013.
  * Hensman et al. (2014) James Hensman, Alex Matthews, and Zoubin Ghahramani.  Scalable variational gaussian process classification, 2014.  URL https://arxiv.org/abs/1411.2005.
  * Kumasaka et al. (2021) Natsuhiko Kumasaka, Raghd Rostom, Ni Huang, Krzysztof Polanski, Kerstin Meyer, Sharad Patel, Rachel Boyd, Celine Gomez, Sam Barnett, Nikolaos Panousis, et al.  Mapping interindividual dynamics of innate immune response at single-cell resolution.  _bioRxiv_ , 2021.
  * Lalchand et al. (2022) Vidhi Lalchand, Aditya Ravuri, and Neil D. Lawrence.  Generalised GPLVM with Stochastic Variational Inference.  In Gustau Camps-Valls, Francisco J. R. Ruiz, and Isabel Valera, editors, _Proceedings of The 25th International Conference on Artificial Intelligence and Statistics_ , volume 151 of _Proceedings of Machine Learning Research_ , pages 7841–7864. PMLR, 28–30 Mar 2022.  URL https://proceedings.mlr.press/v151/lalchand22a.html.
  * Lawrence (2004) Neil D. Lawrence.  Gaussian process latent variable models for visualisation of high dimensional data.  In _Advances in neural information processing systems_ , pages 329–336, 2004.
  * Lopez et al. (2018) Romain Lopez, Jeffrey Regier, Michael B. Cole, Michael I. Jordan, and Nir Yosef.  Deep generative modeling for single-cell transcriptomics.  _Nature Methods_ , 15(12):1053–1058, December 2018.  ISSN 1548-7105.
  * Rasmussen and Williams (2006) Carl Edward Rasmussen and Christopher K. I. Williams.  _Gaussian processes in machine learning_.  Springer, 2006.
  * Stephenson et al. (2021) Emily Stephenson, Gary Reynolds, Rachel A. Botting, Fernando J. Calero-Nieto, Michael D. Morgan, Zewen Kelvin Tuong, Karsten Bach, Waradon Sungnak, Kaylee B. Worlock, Masahiro Yoshida, Natsuhiko Kumasaka, Katarzyna Kania, Justin Engelbert, Bayanne Olabi, Jarmila Stremenova Spegarova, Nicola K. Wilson, Nicole Mende, Laura Jardine, Louis C. S. Gardner, Issac Goh, Dave Horsfall, Jim McGrath, Simone Webb, Michael W. Mather, Rik G. H. Lindeboom, Emma Dann, Ni Huang, Krzysztof Polanski, Elena Prigmore, Florian Gothe, Jonathan Scott, Rebecca P. Payne, Kenneth F. Baker, Aidan T. Hanrath, Ina C. D. Schim van der Loeff, Andrew S. Barr, Amada Sanchez-Gonzalez, Laura Bergamaschi, Federica Mescia, Josephine L. Barnes, Eliz Kilich, Angus de Wilton, Anita Saigal, Aarash Saleh, Sam M. Janes, Claire M. Smith, Nusayhah Gopee, Caroline Wilson, Paul Coupland, Jonathan M. Coxhead, Vladimir Yu Kiselev, Stijn van Dongen, Jaume Bacardit, Hamish W. King, Anthony J. Rostron, A. John Simpson, Sophie Hambleton, Elisa Laurenti, Paul A. Lyons, Kerstin B. Meyer, Marko Z. Nikolić, Christopher J. A. Duncan, Kenneth G. C. Smith, Sarah A. Teichmann, Menna R. Clatworthy, John C. Marioni, Berthold Göttgens, and Muzlifah Haniffa.  Single-cell multi-omics analysis of the immune response in COVID-19.  _Nature Medicine_ , 27(5):904–916, May 2021.  ISSN 1546-170X.  doi: 10.1038/s41591-021-01329-2.
  * Svensson et al. (2018) Valentine Svensson, Sarah A. Teichmann, and Oliver Stegle.  SpatialDE: Identification of spatially variable genes.  _Nature Methods_ , 15(5):343–346, May 2018.  ISSN 1548-7105.  doi: 10.1038/nmeth.4636.
  * van der Wilk (2019) Mark van der Wilk.  _Sparse Gaussian Process Approximations and Applications (Doctoral thesis)_.  PhD thesis, University of Cambridge, 2019.
  * Velten et al. (2022) Britta Velten, Jana M. Braunger, Ricard Argelaguet, Damien Arnol, Jakob Wirbel, Danila Bredikhin, Georg Zeller, and Oliver Stegle.  Identifying temporal and spatial patterns of variation from multimodal data using MEFISTO.  _Nature Methods_ , 19(2):179–186, February 2022\.  ISSN 1548-7105.  doi: 10.1038/s41592-021-01343-9.
  * Verma and Engelhardt (2020) Archit Verma and Barbara E Engelhardt.  A robust nonlinear low-dimensional manifold for single cell rna-seq data.  _BMC bioinformatics_ , 21(1):1–15, 2020.
  * Wolf et al. (2018) F. Alexander Wolf, Philipp Angerer, and Fabian J. Theis.  SCANPY: Large-scale single-cell gene expression data analysis.  _Genome Biology_ , 19(1):15, February 2018.  ISSN 1474-760X.  doi: 10.1186/s13059-017-1382-0.

## Supplementary Material

##  Appendix A Graphical Model

Figure S1 depicts a graphical model of the sparse Gaussian process used in
this paper.

$Z$$\bm{u}_{d}$$\bm{f}_{d}$$x_{n}$$y_{n}$$\sigma^{2}_{y}$$\\{\theta\\}$$\bm{m}_{d}$$S_{d}$$\phi$$N$$D$
Figure S1: The probabilistic graphical model presented in this work. The
unshaded circles denote latent variables to be inferred. The point estimate
variational and model parameters are denoted by just their symbol. The dataset
$Y$ is informed by the sparse GPs, random effects and noise term. The latent
variables $x_{n}$ are inferred through a neural network encoder with weights
$\phi$, alternatively, just learnt as a point estimate.

##  Appendix B Random Effects

The matrix $B$ models random effects across $P$ categorical and $J$ continuous
variables. Let $c_{p}$ denote the number of levels of the $p^{th}$ categorical
variable. The design matrix $\Phi$ has $N$ rows and column dimension
$\sum_{p=1}^{P}c_{p}+J$ (there is a column per level per categorical variable
and one for each continuous variable). Each cell (row of $Y$, $\bm{y}_{n}$) is
associated with a particular level within each of the $P$ categories 111For
example, if the two categories are gender (with 2 levels) and ethnicity (with
8 levels) then each cell $\bm{y}_{n}$ is associated with one of two levels in
the gender category and one of eight levels in ethnicity category.. The
$n^{th}$ row of the design matrix $\Phi$ would have a 1 in the columns to
reflect membership to a particular category and reals in the columns which
reflect co-efficients for the continuous covariates. The row dimension of $B$
is $\sum_{p}c_{p}+J$. The dimension of $\Phi B$ is thus, $N\times D$.

###  B.1 Additive Framework: Non-factorisability

In sparse GP additive model the final likelihood of the data $Y$ can be
expressed as,

$\displaystyle p(Y|F,B,\sigma^{2}_{y})$
$\displaystyle=\prod_{n=1}^{N}\prod_{d=1}^{D}\mathcal{N}(y_{n,d}|\bm{f}_{d}(\bm{x}_{n})+(\Phi
B)_{n,d},\sigma^{2}_{y})$ (17)
$\displaystyle=\prod_{d=1}^{D}\mathcal{N}(\bm{y}_{d}|\bm{f}_{d}(X)+(\Phi
B)_{:,d},\sigma^{2}_{y}\mathbb{I}_{N})$ (18)

A common approach is to integrate out the random effects term $B$ which we can
do in closed form due to the linear dependence. We assume independence between
the columns of $B$, with each $B_{d}\sim\mathcal{N}(\zeta_{d},\Delta_{d})$.
Let $\Sigma$ denote a diagonal matrix with entries $\sigma^{2}_{y}$ on the
diagonal.

$\displaystyle p(Y|F,\sigma^{2}_{y})$
$\displaystyle=\int\prod_{d=1}^{D}p(\bm{y}_{d}|\bm{f}_{d},B_{d})p(B_{d})dB_{d}$
(19)
$\displaystyle=\int\prod_{d=1}^{D}\mathcal{N}(\bm{y}_{d}|\bm{f}_{d}(X)+(\Phi
B)_{:,d},\sigma^{2}_{y}\mathbb{I}_{N})p(B_{d})dB_{d}$
$\displaystyle=\int\prod_{d=1}^{D}\mathcal{N}(\bm{y}_{d}|\bm{f}_{d}(X)+(\Phi
B)_{:,d},\Sigma)\mathcal{N}(\zeta_{d},\Delta_{d})dB_{d}$
$\displaystyle=\prod_{d=1}^{D}\mathcal{N}(\bm{y}_{d}|\bm{f}_{d}(X)+\Phi\zeta_{d},\Sigma+\Phi\Delta_{d}\Phi^{T})$

where $\zeta_{d}$ is a vector of size $\sum_{p=1}^{P}c_{p}+J$ and $\Delta_{d}$
is a matrix of size $(\sum_{p=1}^{P}c_{p}+J)\times(\sum_{p=1}^{P}c_{p}+J)$.
Integrating out the random effects $B$ from the likelihood introduces a dense
covariance matrix $\Phi\Delta_{d}\Phi^{T}$ in the likelihood. The non-
factorisable likelihood (in each of the $D$ output dimensions) prevents the
application of stochastic mini-batching as we need the whole dataset in memory
to compute a single gradient step. Instead, we present an alternative
formulation in section 2.1.2 which subsumes the random effects within the
kernel function and preserves factorisability.

##  Appendix C Inference Algorithm

Algorithm 1 Stochastic Variational Inference for Scalable GPLVM with random
effects

Input: ELBO objective $\mathcal{L}$, gradient based optimiser optim(),
training data $Y=\\{\bm{y}_{n}\\}_{n=1}^{N}$
Initial model params:
$\bm{\theta}$ (covariance hyperparameters for GP mappings $\bm{f}_{d}$),
$\sigma^{2}_{y}$ (variance of the noise model),
$X\equiv\\{\bm{x}_{n}\\}_{n=1}^{N}$ (point estimates for latent covariates)
Initial variational params:
$Z\in\mathbb{R}^{M\times(Q+col.dim(\Phi))}$ (inducing locations),
$\lambda=\\{m_{d},S_{d}\\}_{d=1}^{D}$ (global variational params for inducing
variables per dimension $\bm{u}_{d}$),
while  _not converged_ do

  * •

Choose a random mini-batch ${Y}_{B}\subset Y$.

  * •

Form a mini-batch estimate of the ELBO:
$\mathcal{L}(Y_{B})=\dfrac{N}{B}\left(\sum_{b}\sum_{d}\mathcal{L}_{b,d}\right)-\sum_{d}\textrm{KL}(q(\bm{u}_{d})||p(\bm{u}_{d}|Z))$

  * •

Gradient step:
$Z,\bm{\theta},\sigma^{2}_{y},\lambda,X_{B}\equiv\\{\bm{x}_{n}\\}_{n=1}^{B}\longleftarrow$
optim$(\mathcal{L}(Y_{B}))$

end while

return _$\bm{\theta},\sigma^{2}_{y},Z,\lambda,X$_

##  Appendix D Initialization and Training Configurations

In the experiments presented here, we use informative initializations for the
GPLVM latent variables.

We use the expression of S, G2 and M state (cell cycle) gene markers to
initialize the cell cycle effect latent variable modelled with periodic
covariance (Note: S,G2, M denote the synthesis phase, cell growth phase and
cell division phase in the cell cycle, respectively). While alternative
analysis workflows suggest to regress out this signal before dimensionality
reduction, the GPLVM provides a way to inject prior knowledge regarding this
(through the periodic kernel and latent dimension initialisation). Overall, we
can consider the GPLVM as a unified probabilistic framework that gives
flexibility and consistency to analyse multiple covariates in the data. We
initialize the other latent variables by top principal components of the data,
which turn out to be interpretable in some cases, for example, by
corresponding to known batch and technical effects in the data.

We find that informative initializations help model performance. Our software
allows flexibility for incorporating various combinations of known and latent
variables to be used with the periodic, linear and smooth components of our
covariance function. For example, spatial coordinates corresponding to well
location of sequencing plates can be input to the smooth component of the
model as a known covariate if it is expected to reduce confounding effect.
Additionally, when neural-network based encoders are used (as described in
Section 4), a cell cycle initialization vector can be passed as a known
covariate to the periodic component of our kernel, somewhat mirroring to the
process used in conditional VAEs.

We use a number of inducing variables strictly greater than the number of
covariates (Col. dim ($\Phi$)) used for the linear kernel as this is the
minimum number of data points needed to define a hyperplane in Col.
dim($\Phi$) dimensions.

We show an example of how the GPLVM diverges from initialization in a
meaningful way on the COVID dataset in Fig. S6. Here we compare one GPLVM
latent variable with the principal component provided during initialization.
The two are correlated and both capture cell-type specific gene expression
signatures (Fig. S6B). However, in PCA the cell type specific signal is
confounded with the batch effect (Fig. S6C) while with GPLVM we can
intuitively use this variable to distinguish NK cells (NK_16hi, NK_56hi) from
T cell clusters (T CD4, T CD8, T gd).

We further note how the latent variable initialization varies as a function of
the number of training iterations/epochs. Considering the first latent
dimension of the GPLVM as the cell cycle (CC) effect along with 10 further
dimensions initialised by the 10 PCs of the COVID dataset, Fig. S5 shows how
they start to deviate significantly from the initial point as the number of
iterations across the epochs increases. (The distance between an initial
dimension and the optimised dimension is measured by their Root Mean Squared
Deviation – RMSD). We observe that the CC effect is the least changed overall,
implying that our initialization of the average CC gene expression well
defines the latent dimension that captures CC effect. PC deviations are
distinctly visible, specifically as the degree of deviation starts to spread
out across a range of RMSD values.

Table S1: Training configurations for the experimental results in the
manuscript. The Innate Immunity experiment used a learning rate of 0.05 for
the first 3 epochs and 0.005 thereafter.

Experiment | Learning rate | Batch size | Epochs | Latent dim $(Q)$ | Col. dim($\Phi$) | Num inducing $(M)$
---|---|---|---|---|---|---
Innate Immunity | 0.05 / 0.005 | 220 | 50 | 7 = 6 (rbf) + 1 (periodic) | 200 | 201
COVID w/o severity | 0.01 | 200 | 100 | 11 = 10 (rbf) + 1 (periodic) | 146 | 147
COVID w. severity | 0.005 | 300 | 50 | 12 = 10 (rbf) + 1 (periodic) + 1 (sev.) | 146 | 147

An interesting effect observed while training the model on the Innate Immunity
experiment is that overtraining can occur, albeit the effect is very subtle.
Figure S2 shows the effect of running the model well past our stopping point
(50 epochs) on the quality of the first latent variable in classifying batch
effects. Quality of classification is measured as the percentage accuracy
obtained using a standard classifier (a logistic regression) in being able to
separate the two batches given just the first latent variable. The accuracy
score drops slightly if training is continued well past 100 epochs in this
experiment.

Figure S2:  Effect of overtraining on the ability of the first latent variable
(that captures batch effect) to classify batch one and two. After about a
hundred epochs, the quality of the batch effect latent dimensions starts to
diminish slightly.

##  Appendix E Simplifying the Model to Independence across Linear and Non-
linear Terms

It’s possible to simplify sparse GPs with two additive kernel components which
can each be made zero by a choice of inputs (and thus, our model) so that the
linear and non-linear parts of each of the $\boldsymbol{\bm{f}}_{d}$’s are
independent. We drop the column index $d$ for clarity. For certain downstream
tasks, we may like to interpret
$\bm{f}=\boldsymbol{\bm{f}}_{1}+\boldsymbol{\bm{f}}_{2}$ as a sum of GPs,
$\boldsymbol{\bm{f}}_{1}$ corresponding to the kernel $k_{per}\times k_{rbf}$
and $\boldsymbol{\bm{f}}_{2}$ corresponding to $k_{lin}$. We’d like to have
the prior conditioned on inducing variables $\bm{f}|\bm{u}$ look like a sum of
sparse GPs as in van der Wilk [2019],
$\boldsymbol{\bm{f}}|\boldsymbol{\bm{u}}=\boldsymbol{\bm{f}}_{\not{\text{lin}}}|\boldsymbol{\bm{u}}_{\not{\text{lin}}}+\boldsymbol{\bm{f}}_{\text{lin}}|\boldsymbol{\bm{u}}_{\text{lin}}$
given by,

$\displaystyle\bm{f}|\bm{u}\sim\mathcal{N}\Big{(}\sum_{i\in\not{\text{lin}},\text{lin}}K_{nm}^{(i)}K_{mm}^{{-1}^{(i)}}\bm{u}_{i},\sum_{i\in\not{\text{lin}},\text{lin}}K_{nn}^{(i)}-K_{nm}^{(i)}K_{mm}^{{-1}^{(i)}}K_{mn}^{(i)}\Big{)}$
(20)

where $K_{nm}^{(i)}$ is the cross-covariance matrix between latent covariates
$(X_{per}|X_{rbf}|\Phi)$ and $(Z_{per}|Z_{rbf}|Z_{lin})$ and with
$\boldsymbol{\bm{u}}_{\text{lin}}\perp\\!\\!\\!\perp\boldsymbol{\bm{u}}_{\not{\text{lin}}}$.
The covariance matrices are given by,

$\displaystyle K_{nm}^{\not{\text{lin}}}$
$\displaystyle=k_{per}(X_{per},Z_{per})\times k_{rbf}(X_{rbf},Z_{rbf})$ (21)
$\displaystyle K_{nm}^{\text{lin}}$ $\displaystyle=k_{lin}(\Phi,Z_{lin}),$
(22) $\displaystyle K_{mm}^{\not{\text{lin}}}$
$\displaystyle=k_{per}(Z_{per},Z_{per})\times k_{rbf}(Z_{rbf},Z_{rbf}),$ (23)
$\displaystyle K_{mm}^{\text{lin}}$ $\displaystyle=k_{lin}(Z_{lin},Z_{lin}).$
(24)

This simplification can be achieved by, as an example pertaining to our
kernel, partitioning the inducing points $Z$ as follows,

$Z\equiv(Z_{\not{\text{lin}}}|Z_{\text{lin}})=\begin{bmatrix}Z_{1}&\boldsymbol{\bm{0}}\\\
\pm\boldsymbol{\bm{\infty}}&Z_{2}\end{bmatrix}.$ (25)

The zeros and infinities in this inducing inputs matrix are selected by
considering inputs to a kernel that can set it to zero - for example,
$k_{\text{rbf}}(.,\pm\infty)=0$ and $k_{\text{lin}}(.,0)=0$. Due to this
partitioning, the matrices of the prior,

$\boldsymbol{\bm{f}}|\boldsymbol{\bm{u}}\sim\mathcal{N}(K_{nm}K_{mm}^{-1}\boldsymbol{\bm{u}},K_{nn}-K_{nm}K_{mm}^{-1}K_{mn})$
(26)

simplifies as follows,

$\displaystyle K_{mm}$
$\displaystyle=\begin{bmatrix}k_{\not{\text{lin}}}(Z_{1},Z_{1})+k_{\text{lin}}(\boldsymbol{\bm{0}},\boldsymbol{\bm{0}})&k_{\not{\text{lin}}}(Z_{1},\pm\boldsymbol{\bm{\infty}})+k_{\text{lin}}(\boldsymbol{\bm{0}},Z_{2})\\\
k_{\not{\text{lin}}}(\pm\boldsymbol{\bm{\infty}},Z_{1})+k_{\text{lin}}(Z_{2},\boldsymbol{\bm{0}})&k_{\not{\text{lin}}}(\pm\boldsymbol{\bm{\infty}},\pm\boldsymbol{\bm{\infty}})+k_{\text{lin}}(Z_{2},Z_{2})\end{bmatrix}=\begin{bmatrix}k_{\not{\text{lin}}}(Z_{1},Z_{1})&\boldsymbol{\bm{0}}\\\
\boldsymbol{\bm{0}}&k_{\text{lin}}(Z_{2},Z_{2})\end{bmatrix},$ (27)
$\displaystyle K_{nm}$
$\displaystyle=\begin{bmatrix}k_{\not{\text{lin}}}(X_{\not{\text{lin}}},Z_{1})+k_{\text{lin}}(\Phi,\boldsymbol{\bm{0}})&k_{\not{\text{lin}}}(X_{\not{\text{lin}}},\pm\boldsymbol{\bm{\infty}})+k_{\text{lin}}(\Phi,Z_{2})\end{bmatrix}=\begin{bmatrix}k_{\not{\text{lin}}}(X_{\not{\text{lin}}},Z_{1})&k_{\text{lin}}(\Phi,Z_{2})\end{bmatrix}.$
(28)

Cleaning up the notation,

$\displaystyle
K_{mm}=\begin{bmatrix}K^{\not{\text{lin}}}_{mm}&\boldsymbol{\bm{0}}\\\
\boldsymbol{\bm{0}}&K^{\text{lin}}_{mm}\end{bmatrix},K_{nm}=\begin{bmatrix}K^{\not{\text{lin}}}_{nm}&K^{\text{lin}}_{nm}\end{bmatrix}.$
(29)

By using block matrix multiplication rules, the prior simplifies to Equation
20.

$\displaystyle\mathbb{E}(\boldsymbol{\bm{f}}|\boldsymbol{\bm{u}})$
$\displaystyle=K_{nm}K_{mm}^{-1}\boldsymbol{\bm{u}}=\begin{bmatrix}K^{\not{\text{lin}}}_{nm}&K^{\text{lin}}_{nm}\end{bmatrix}\begin{bmatrix}K^{\not{\text{lin}}}_{mm}&\boldsymbol{\bm{0}}\\\
\boldsymbol{\bm{0}}&K^{\text{lin}}_{mm}\end{bmatrix}^{-1}\begin{bmatrix}\boldsymbol{\bm{u}}_{\not{\text{lin}}}\\\
\boldsymbol{\bm{u}}_{\text{lin}}\end{bmatrix}$ (30)
$\displaystyle=K^{\not{\text{lin}}}_{nm}{K^{\not{\text{lin}}}}^{-1}_{mm}\boldsymbol{\bm{u}}_{\not{\text{lin}}}+K^{\text{lin}}_{nm}{K^{\text{lin}}}^{-1}_{mm}\boldsymbol{\bm{u}}_{\text{lin}},$
(31) $\displaystyle\mathbb{C}(\boldsymbol{\bm{f}}|\boldsymbol{\bm{u}})$
$\displaystyle=K_{nn}+K_{nm}K_{mm}^{-1}K_{mn}$ (32)
$\displaystyle=K^{\text{lin}}_{nn}+K^{\not{\text{lin}}}_{nn}+\begin{bmatrix}K^{\not{\text{lin}}}_{nm}&K^{\text{lin}}_{nm}\end{bmatrix}\begin{bmatrix}K^{\not{\text{lin}}}_{mm}&\boldsymbol{\bm{0}}\\\
\boldsymbol{\bm{0}}&K^{\text{lin}}_{mm}\end{bmatrix}^{-1}\begin{bmatrix}K^{\not{\text{lin}}}_{nm}\\\
K^{\text{lin}}_{nm}\end{bmatrix}$ (33)
$\displaystyle=(K^{\not{\text{lin}}}_{nn}+K^{\not{\text{lin}}}_{nm}{K^{\not{\text{lin}}}}^{-1}_{mm}K^{\not{\text{lin}}}_{mn})+(K^{\text{lin}}_{nn}+K^{\text{lin}}_{nm}{K^{\text{lin}}}^{-1}_{mm}K^{\text{lin}}_{mn}).$
(34)

Note that $k_{\text{per}}$, which is part of our kernel, doesn’t matter in
this simplification as if $k_{\text{rbf}}(...)=0$, then
$k_{\not{\text{lin}}}(...)=k_{\text{rbf}}(...)*k_{\text{per}}(...)=0$.
Further, the number of inducing variables across $i$, $\bm{u}_{i}$’s can be
different, and have shapes $(M1\times 1)$ and $(M2\times 1)$ without loss of
generality. This formulation allows one to extract the terms in the posterior
process corresponding to just the linear (random-effects) term,

$\displaystyle q\left(\begin{bmatrix}\boldsymbol{\bm{u}}_{\not\text{lin}}\\\
\boldsymbol{\bm{u}}_{\text{lin}}\end{bmatrix}\right)$
$\displaystyle=\mathcal{N}\left(\begin{bmatrix}\boldsymbol{\bm{u}}_{\not\text{lin}}\\\
\boldsymbol{\bm{u}}_{\text{lin}}\end{bmatrix}|\begin{bmatrix}\boldsymbol{\bm{m}}_{\not\text{lin}}\\\
\boldsymbol{\bm{m}}_{\text{lin}}\end{bmatrix},\begin{bmatrix}S_{\not\text{lin}}&S_{c}\\\
S_{c}^{T}&S_{\text{lin}}\end{bmatrix}\right),$ (35) $\displaystyle
q(\boldsymbol{\bm{f}}_{\text{lin}})$ $\displaystyle=\int
p(\boldsymbol{\bm{f}}_{\text{lin}}|\boldsymbol{\bm{u}}_{\text{lin}})q(\boldsymbol{\bm{u}}_{\text{lin}})d\boldsymbol{\bm{u}}_{\text{lin}}=\mathcal{N}(\Phi
Z_{\text{lin}}^{T}(Z_{\text{lin}}Z_{\text{lin}}^{T})^{-1}\boldsymbol{\bm{m}}_{\text{lin}},\Phi(\mathbb{I}+A(S_{\text{lin}}-Z_{\text{lin}}Z_{\text{lin}}^{T})A^{T})\Phi^{T}),$
(36)

where $A=Z_{\text{lin}}^{T}(Z_{\text{lin}}Z_{\text{lin}}^{T})^{-1}$. We
compute the variational posterior $q(\boldsymbol{\bm{f}}_{\text{lin}})$ as in
Hensman et al. [2014]. A small jitter $\delta\mathbb{I}$ is added to make
$Z_{\text{lin}}Z_{\text{lin}}^{T}$ invertible in practice.

##  Appendix F Supplementary Figures

Figure S3:  UMAP embedding of innate immunity dataset 6 dimensional latent
space learnt by PCA (top) and BGPLVM (bottom). Cells are colored by treatment
condition (pIC: Poly(I:C), IFN: interferon-beta) (left), experimental batch
(center) and cell cycle signature (right).  Figure S4:  Latent COVID severity
in moderate COVID patients, ordered and colored by number of days since onset
of symptoms.  Figure S5: Latent variable deviations from initialization as the
total number of epochs increases during training. CC corresponds to the cell
cycle effect, while the other 10 lines correspond to the top 10 principal
components of the COVID dataset.

Figure S6:  (top) Scatterplot of GPLVM latent variable 3 (GPLVM3) and the
principal component used for initialization (PC3). Cells are colored by cell
type. (middle) Distribution of latent variable values for each cell type for
PC3 (left) and GPLVM3 (right) (bottom) Distribution of latent variable values
for each experimental batch (site of sample collection) for PC3 (left) and
GPLVM3 (right)