# Data-Driven Time Propagation of Quantum Systems with Neural Networks

James Nelson<EMAIL_ADDRESS>School of Physics, AMBER and CRANN Institute,
Trinity College, Dublin 2, Ireland  Luuk Coopmans<EMAIL_ADDRESS>School of
Physics, AMBER and CRANN Institute, Trinity College, Dublin 2, Ireland  Dublin
Institute for Advanced Studies, School of Theoretical Physics, 10 Burlington
Rd, Dublin, Ireland  Graham Kells<EMAIL_ADDRESS>Dublin Institute for
Advanced Studies, School of Theoretical Physics, 10 Burlington Rd, Dublin,
Ireland  Dublin City University, School of Physical Sciences, Glasnevin,
Dublin 9, Ireland  Stefano Sanvito<EMAIL_ADDRESS>School of Physics, AMBER
and CRANN Institute, Trinity College, Dublin 2, Ireland

###### Abstract

We investigate the potential of supervised machine learning to propagate a
quantum system in time. While Markovian dynamics can be learned easily, given
a sufficient amount of data, non-Markovian systems are non-trivial and their
description requires the memory knowledge of past states. Here we analyse the
feature of such memory by taking a simple 1D Heisenberg model as many-body
Hamiltonian, and construct a non-Markovian description by representing the
system over the single-particle reduced density matrix. The number of past
states required for this representation to reproduce the time-dependent
dynamics is found to grow exponentially with the number of spins and with the
density of the system spectrum. Most importantly, we demonstrate that neural
networks can work as time propagators at any time in the future and that they
can be concatenated in time forming an autoregression. Such neural-network
autoregression can be used to generate long-time and arbitrary dense time
trajectories. Finally, we investigate the time resolution needed to represent
the system memory. We find two regimes: for fine memory samplings the memory
needed remains constant, while longer memories are required for coarse
samplings, although the total number of time steps remains constant. The
boundary between these two regimes is set by the period corresponding to the
highest frequency in the system spectrum, demonstrating that neural network
can overcome the limitation set by the Shannon-Nyquist sampling theorem.

##  I Introduction

The notion of a memory kernel arises frequently in the study of dynamical
models, where one attempts to understand or recreate dynamical behaviours
using only a subset of the total available degrees of freedom. In particular,
a kernel becomes necessary in situations where the subset of information
available at a given instant in time is not enough to predict the subsequent
behaviour. In open quantum systems, where a restricted number of degrees of
freedom is in contact with a bath, one formally understands this subset
restriction as the tracing out of the bath. Here, the need of a memory kernel
to obtain accurate long-time trajectories is a defining criteria that
distinguishes Markovian from non-Markovian dynamics.

A number of key techniques and approaches have emerged enabling a systematic
study of this effect. One such family are the so-called projection techniques
(e.g the Nakajima-Zwanzig equation Nakajima1958 ; Zwanzig1960 ; Prigogine1962
and their time-convolutionless variants Shibata1977 ; Shibata1980 ) that
render the history (memory) within an integro-differential equation and can be
used to provide a perturbative justification for the Markovian behaviour.
Another approach for exploring memory effects is to embed the system into an
enlarged Hilbert space such that the non-Markovian dynamics of the reduced
system can be represented via the Markovian evolution of an enlarged whole
Ciccarello2013 ; McCloskey2014 ; Kretschmer2016 ; Lorenzo2017 ; Cakmak2017 ;
Campbell2018 . In this respect, a useful class of embeddings are the so-called
collision models, where ancillas representing the bath interact in a time-
dependent fashion. Here the concept of memory depth naturally arises, when one
allows these ancillas to interact between themselves.

Both the projection and embedding schemes represent intuitive and inherently
human-learned/reasoned approaches for investigating the interplay between a
quantum system and its associated bath. Given the now widespread application
of machine learning (ML) in the domain of classical and quantum dynamics, it
is natural that one might apply machine learning to the same domain. For a
selective snapshot of this field we refer the reader to some representative
published literature mardt2018vampnets ; breen2020newton ; carleo2017solving ;
hartmann2019neural ; lin2021scaling ; banchi2018modelling ; Luchnikov2020  and
references therein.

When looking at the non-Markovianity of a time evolution, in particular with
embeddings and memory kernels, an approach that has recently emerged is to
construct a fixed neural network (NN) that learns the behaviours of local
observables and then attempts to match the time dependence of the observables
beyond the training time. For instance, in reference Mazza2021  this data-
driven method was applied to learn time-local generators of the dynamics of
open quantum systems. Similarly, in Luchnikov2021  a related approach was used
to find an effective Markovian embedding from which several important
properties and the exact dynamics of the reduced system could be extracted.

One of the key advantages of these data-driven methods is their universality.
In fact, machine-learning time propagators can be used with any method to
generate quantum evolution and show promise for use in investigations into the
nature of the non-Markovian evolution. In this respect, the methodology blends
features of Markovian embeddings and integro-differential memory kernels, but
where the functional that encodes the dynamics is fixed so that the embeddings
and the kernel/memory encoding are time or state independent.

In this paper we take a broadly similar approach in that we train a NN on a
dataset of time-dependent quantum trajectories and then ask it to evolve an
arbitrary starting state up to times longer than the final time of the learned
trajectories. Distinguishing our work from previous studies, we look at the
complete collection of reduced single-particle density matrices that together
make up the full (strongly interacting) system and focus, in particular, on
the accuracy of the method with respect to the depth of the sampling history
that is used as input for the NN. While in general, as we show below, the
Markovian dynamics of the full system can be learned by only one history
sample, for the representation in terms of reduced density matrices this is no
longer the case. In fact, we find that the sampling history needs to be
increased exponentially with system size to keep the error rates on the
predicted dynamics below a fixed threshold. Aside from studying such memory
effects, an advantage of this approach is that its malleability allows for
non-fixed time sampling and naturally protects against oversampling error.
Moreover, as the error growth on a given NN is random we can use the idea of
neural-network ensembling to monitor and reduce errors.

We structure our paper as follows: in section II we give an overview of
various methods to study memory effects in quantum dynamics and introduce our
specific setup and method to generate the data of the quantum evolution. In
section III and IV we present our results for Markovian and non-Markovian
dynamics, respectively, discussing the interplay between the non-locality in
time and space. Then we conclude.

##  II Methodology and background

###  II.1 Brief review of existing methods

We consider a quantum system described by its density matrix, $\hat{\rho}$,
that consists of a total of $N$ qubits. The reduced density matrices of such
system can be obtained by tracing out selected degrees of freedom, here
collectively called the environment,
$\hat{\rho}_{\text{red}}=\text{Tr}_{\text{E}}\hat{\rho}$, where ‘E’ stands for
some environment (see Figure 1). When the traced-out region is all but one
qubit we use the parameterisation
$\hat{\rho}_{\text{red}}=\frac{1}{2}(\mathbbm{1}+\textbf{r}\cdot\boldsymbol{\sigma})$,
where r is the Bloch vector. Since the reduced density matrices leave out much
information they alone cannot determine the future evolution of the entire
system. Thus, for a typical dynamical reconstruction of the reduced quantities
over time one may also need some historical knowledge of past states. Several
analytical and numerical techniques have been developed to study the memory
effects in this type of setup.

Projection methods: A technique often mentioned in literature to render the
past knowledge explicit in the time-evolution equations is the Nakajima-
Zwanzig projector method Nakajima1958 ; Zwanzig1960 ; Prigogine1962  (see
appendix A for a quick derivation). When the time-evolution starts ($t=0$)
from an initially unentangled system and bath,
$\hat{\rho}(t=0)=\hat{\rho}_{\text{red}}\otimes\hat{\rho}_{\text{E}}$, this
approach allows one to reduce the Liouville-Von Neumann equation for the
entire system,

$\partial_{t}\hat{\rho}=\frac{i}{\hbar}\left[\hat{\rho},\hat{H}\right]=\hat{L}\hat{\rho}\>,$
(1)

to an equation containing just the relevant reduced density operator,
$\hat{\rho}_{\text{red}}$,

$\partial_{t}[\hat{\rho}_{\text{red}}]=\mathcal{\hat{P}}\hat{L}[\hat{\rho}_{\text{rel}}]+\int_{0}^{t}dt^{\prime}\mathcal{\hat{K}}(t^{\prime})[\mathcal{\hat{P}}\hat{\rho}_{\text{red}}(t-t^{\prime})]\>,$
(2)

where

$\mathcal{\hat{K}}(t)=\mathcal{\hat{P}}\hat{L}e^{\mathcal{\hat{Q}}\hat{{L}t}}\mathcal{\hat{Q}}\hat{L}\mathcal{\hat{P}}\>.$
(3)

Here we have introduced the Liouvillian operator, $\hat{L}$, and the
projectors $\mathcal{\hat{P}}$ and $\mathcal{\hat{Q}}=\mathcal{I}-\mathcal{P}$
($\mathcal{I}$ is the identity), which are defined in terms of the reduced
density matrix $\hat{\rho}_{\text{red}}$ and a fixed reference state
$\hat{\rho}_{\mathrm{B}}$ of the environment degrees of freedom,

$\hat{\rho}_{\text{red}}=\hat{\mathcal{P}}\hat{\rho}=\text{Tr}_{\mathrm{B}}(\hat{\rho})\otimes\hat{\rho}_{\mathrm{B}}\>.$
(4)

The memory kernel $\mathcal{\hat{K}}(t)$ includes the effects of the history
of the reduced state. Solving the Nakajima-Zwanzig equation and finding the
memory kernel is in general a computationally hard task. Transfer-tensor
methods Cerrillo2014 ; Gelzinis2017 ; Gherardini2021  have been recently
proposed as efficient techniques to reconstruct the memory kernel. In these
the individual quantum trajectories are used to learn dynamical maps that can
be converted into the so-called transfer tensors. The transfer tensors are
then directly related to the memory kernel of the Nakajima-Zwanzig equation
Cerrillo2014 .

Figure 1: (Color on line) Here the two types of non-Markovian mappings are
shown. The small blue arrows symbolise the Block vectors of the individual
qubits, while the dashed lines are the time steps. $N$ is the number of
qubits, $h$ is the history and $m$ is the future time where we map the
evolution to. In the ‘single’ case, the history of a single qubit (evolving in
the presence of the others, which then define the environment) is used to
predict its future dynamics, while in the ‘all’ case, the history of all the
qubits, is used to predict the future dynamics of the entire system.

Embeddings and collision models: An alternative general approach to the study
of memory consists in collision models and the so-called Markovian embeddings
[see, for instance, Ciccarello2013 ; McCloskey2014 ; Kretschmer2016 ;
Lorenzo2017 ; Cakmak2017 ; Campbell2018 ]. The general idea is that the system
interacts in a controlled way with an environment comprising of ancillas that,
in turn, can interact amongst themselves in a variety of ways. In this
approach, which has been very fruitful for understanding the boundary between
Markovian and non-Markovian dynamics, one typically employs some spin
Hamiltonian, where the interaction is turned on and off sequentially in some
predetermined fashion. For the purposes of our study we similarly employ a
one-dimensional spin description of our system, specifically using the
Heisenberg model as generator of the dynamics. Thus the Hamiltonian reads,

$H=-\frac{J}{2}\sum_{j=1}^{N}\boldsymbol{\sigma}^{(j)}\cdot\boldsymbol{\sigma}^{(j+1)}\>,$
(5)

where $\boldsymbol{\sigma}$ is the vector of Pauli matrices for spin 1/2, $N$
is the total number of qubits and $J$ defines the time scale of the problem.
In particular here we consider finite periodic spin arrangements, namely rings
of spins. In our case, since our motivation is to understand how well the full
dynamics can be predicted by using a machine-learning (ML) data-driven
technique, no piecewise temporal distinctions are made.

###  II.2 Data-driven methods

The data driven approach for predictive dynamics used here and, for instance,
in references Mazza2021 ; Luchnikov2021  has some clear similarities to both
the projection and the collision methods. One defines $\textbf{r}_{i}^{t}$ as
the Bloch vector for the $i$-th qubit at the time $t$, which is then
discretised into $t_{n}=n\Delta$ with $n$ an integer and some fundamental time
$\Delta$, and sets

$R^{t_{n}}=R^{n}=(\textbf{r}_{1}^{n},\textbf{r}_{2}^{n},...,\textbf{r}_{N}^{n})\>,$
(6)

namely, $R^{t_{n}}$ is the collection of all Bloch vectors at a given time. We
consider two types of mappings (in the underbraces we show the dimensions): i)
‘all’, where the entire collection of Bloch vectors is propagated in time

$\underbrace{(R^{n},R^{n-1},...,R^{n-h+1})}_{3\times N\times
h}\rightarrow\underbrace{R^{n+m}}_{3\times N}\>,$ (7)

and ii) ‘single’, where the time-evolution of only one Bloch vector is
computed,

$\underbrace{(\textbf{r}^{n},\textbf{r}^{n-1},...,\textbf{r}^{n-h+1})}_{3\times
h}\rightarrow\underbrace{\textbf{r}^{n+m}}_{3}\>.$ (8)

Here, the history $h$ is the Markov order, which determines the depth of the
time memory, and $m$ is the distance into the future where the prediction will
take place. Unlike when propagating the full many-body wave-function, which
grows as $2^{N}$, both mappings scale linearly with the number of qubits $N$.
A schematic of both these mappings is shown in Figure 1.

In this set-up a Markovian evolution corresponds to the $h=1$ case, while the
accurate prediction of an increasingly non-Markovian dynamics Campbell2018
can be obtained by enlarging the memory, $h$. In lay terms, the ability of a
NN to predict the dynamics from only a partial knowledge of the full density
matrix (the reduced density matrix) is traded off with a non-locality in time.
An important feature of this scheme is that we can perform autoregression,
namely we can feed the predictions at some given time back into the model to
make predictions at subsequent times. Such ability makes our mappings to be
universal time propagators, regardless of whether the dynamics is Markovian or
not. Notably, as one can train the model to predict at arbitrary times in the
future (one can train different networks for different $m$’s) both short- and
long-time trajectories are accessible, with the error being dictated by the
accuracy of the specific network.

Although recurrent NNs have been used previously in the literature to model
quantum evolution banchi2018modelling ; zhang2019predicting , here we opt for
a fully connected NN to learn the propagator of the dynamics, since we only
deal with inputs of a fixed size. In particular, we use a two-layer network
with 64 nodes in each hidden layer, and the Exponential Linear Unit (ELU)
activation function, which was found to perform best. In order to learn the NN
parameters we minimize the trace distance between the predicted and true Bloch
vectors NoteD . The networks have been implemented using PyTorch
paszke2019pytorch .

Figure 2: (Color on line) The mean single-particle entropy,
$S_{\mathrm{sp}}=\mbox{Tr}({\hat{\rho}_{\text{red}}\log{\hat{\rho}_{\text{red}}}})$,
as a function of time for differently initialized states. The dashed lines
correspond to randomly generated initial states, while the continuous lines
are associated to product initial states. The colour encodes the total number
of qubits in the system. Note that product states have vanishing initial
entropy, while the entropy of random states grows with the number of qubits
considered.

###  II.3 Data generation

The time-dependent quantum trajectories used for the training of the NN have
been generated by taking two different types of pure initial states,
$\hat{\rho}(t=0)$, namely random and product states. These are, respectively,
constructed from initial wavefunctions of the form

$|\psi(0)\rangle_{\mathrm{rand}}=\sum_{i_{1}=0,1}\sum_{i_{2}=0,1}...\sum_{i_{N}=0,1}a_{i_{1}i_{2}...i_{N}}|i_{1}i_{2}...i_{N}\rangle\>,$
(9)
$|\psi(0)\rangle_{\mathrm{prod}}=\sum_{i_{1}=0,1}a_{i_{1}}|i_{1}\rangle\sum_{i_{2}=0,1}a_{i_{2}}|i_{2}\rangle...\sum_{i_{N}=0,1}a_{i_{N}}|i_{N}\rangle\>,$
(10)

where the vector $|i_{n}\rangle$ spans the Fock space of the $n$-th qubit and
the coefficients $a_{i_{n}}$ and $a_{i_{1}i_{2}...i_{N}}$ are selected
randomly in both cases. Note that, on the one hand, the single-particle
reduced density matrices associated to product states have all zero Von
Neumann entanglement entropy,
$\mbox{Tr}({\hat{\rho}_{\text{red}}\log{\hat{\rho}_{\text{red}}}})$, and hence
there is no entanglement initially present in the system. On the other hand,
random initial states have, on average, a high entanglement entropy, as shown
in figure 2.

These initial states are then time evolved with the propagator of the total
system (the specific time-dynamic generators are discussed below). In
particular, at a practical level the equations of motion have been integrated
numerically by using the QuTiP Python package QuTiP . Then, we trace out the
environment degrees of freedom (some of the quibits) to generate the quantum
trajectories of the single qubits. For all our numerical experiments - unless
otherwise specified - we use a dataset consisting of 11,000 samples for each
system size, $N$. These are divided into 8,000 training samples, 2,000
validation ones, while 1,000 samples are used as test set. We employ the
validation set to determine when to halt the training. Throughout this work
all the results shown are calculated on the test set.

##  III Predicting the evolution of Markovian systems

For Markovian dynamics the state of the system at a given time contains
sufficient information to predict its future evolution. This is the case when
the state is described by the wavefunction or by the full density matrix.
Here, since either the Schrödinger or the Von Neumann equation imply a linear
mapping between states at different times, given a set of trajectories we can
use a linear regression to learn the propagator. Let the vector
$\textbf{x}_{n}$ represent the wavefunction or density matrix of a system at
time step $n$, then the propagator, $P$, is a matrix with
$\textbf{x}_{n+1}=P\textbf{x}_{n}$. Given a dataset with $p$ pairs of single
time step evolutions
$\\{(\textbf{x}_{n}^{(1)},\textbf{x}_{n+1}^{(1)}),...,(\textbf{x}_{n}^{(p)},\textbf{x}_{n+1}^{(p)})\\}$
we can learn the propagator matrix by writing

$P=(X_{n}^{T}X_{n})^{-1}X_{n}^{T}X_{n+1}\>,$ (11)

where $X_{n}$ is a matrix, whose $i$-th row is $\textbf{x}_{n}^{(i)}$. Thus,
we can exactly learn the propagator, if the matrix $X_{n}^{T}X_{n}$ is
invertible. This is the case when the number of samples, $p$, is greater than
the dimension of the vector. If $\Omega$ is the dimension of the Hilbert
space, then we require $2\Omega-1$ samples to learn the wavefunction
propagator and $\Omega^{2}-1$ samples to learn that of the density matrix Note
.

As a warm-up example we consider a single qubit evolving via the Lindblad
master equation breuer2002theory

$\frac{d\hat{\rho}}{dt}=-i[\hat{H},\hat{\rho}]+2\hat{L}\hat{\rho}\hat{L}^{\dagger}-\\{\hat{L}^{\dagger}\hat{L},\hat{\rho}\\},$
(12)

where we have chosen
$\hat{H}=0.5\hat{\sigma}_{x}+0.3\hat{\sigma}_{y}+0.2\hat{\sigma}_{z}$ and
$\hat{L}=\sqrt{0.1}\hat{\sigma}_{x}$. Note that this is non-unitary but
Markovian. Here the propagator is of the form
$\textbf{x}_{n+1}=A\textbf{x}_{n}+\textbf{a}$, where $A$ is a matrix and a is
a vector. However, by defining $\tilde{\textbf{x}}=(1,x_{1},x_{2},...)$ we can
write it in the form $\tilde{\textbf{x}}_{n+1}=P\tilde{\textbf{x}}_{n}$ and
thus apply Eq. (11). Figure 3 compares the exacts Bloch vector trajectories
with those generated by the learnt propagator. Here four samples were required
to learn the propagator and the two trajectories agree up to floating point
error.

Figure 3: (Color on line) Comparison between the exact time-dependent
trajectories (continuous line) and those predicted by the machine-learned
regression (points). Here we plot the time-dependent expectation values of the
three Pauli operators over the evolution determined by the Lindblad equation,
Eq. (12). The predictions agree with the exact values up to a floating point
error.

##  IV Predicting non-markovian dynamics

Let us now move to the non-Markovian dynamics. In this case our Hamiltonian is
the Heisenberg model of Eq. (5), calculated for a number of qubits going from
$N=2$ to $N=7$. Information about the time evolution can be extracted by
plotting the wave-function fidelity, $|\langle\psi(t=0)|\psi(t)\rangle|^{2}$,
as a function of time for a range of random initial states, as presented in
Figure 4. We note that for $N\leq 4$ the fidelity appears periodic in time
over the time interval considered, with periods $T_{2}=\pi/2$, $T_{3}=2\pi/3$
and $T_{4}=\pi$, respectively for $N=2,3$ and 4 (all times are in units of
$1/J$). We also note that after a time of approximately $1$ ($1/J$) the
fidelity approaches zero regardless of $N$. Such time can be defined as the
characteristic de-correlation time.

Figure 4: (Color on line) The fidelity of the wavefunction, $|\psi\rangle$,
with respect to the initial state, $|\psi(t=0)\rangle$, is computed as a
function of time for states initialized randomly. Here, the fidelity is
defined as $|\langle\psi(t=0)|\psi(t)\rangle|^{2}$. The blue lines are the
individual trajectories, while the black lines their mean. Note that for
$N=2,3,4$ the fidelity is periodic in time. The time evolution are for the
Heisenberg Hamiltonian of Eq. (5).

Next, we describe our results obtained for the non-Markovian dynamics of
single particle Bloch vectors.

###  IV.1 Memory

Let us now consider the dynamics of the single-particle reduced density matrix
and evaluate the memory needed to propagate a system initialized in a random
state. In particular, we set the time step at $\Delta=0.08\pi\approx 0.25$ and
predict several times into the future. We consider two different cases. In the
first, figure 5, the feature vector for the NN comprises the Bloch vectors of
all qubits, namely we simultaneously follow the dynamics of the single-
particle reduced density matrices of all the qubits. In the second case,
figure 6, the NN propagates the Bloch vector of a single qubit only (note that
the qubit choice here is irrelevant since the system is a ring, namely it has
translational symmetry). In the figures we plot the NN error, the mean trace
distance (MTD), as a function of the memory (in units of the time step,
$\Delta=0.08\pi$) for NNs that propagate at different times into the future
(color code).

Figure 5: (Color on line) Error of the NN (the mean trace distance - MTD) as a
function of the memory, $h$, (in units of $\Delta=0.08\pi\approx 0.25/J$), for
NNs that propagate at different points in the future (color code). Here the
results are for NNs that use as feature the reduced density matrix of all
qubits (‘all’ case).

In general, in both situations one can always find a sufficiently long memory
to converge the NNs to small errors (in the figures an error of 0.01 is
indicted as a dashed black line). For the ‘all’ case such memory seems to be
rather independent of the final time of propagation, namely it takes
approximately the same memory to propagate the entire Bloch-vector manifold to
either short or long times. This is more evident when the total number of
qubits is small, while some scattering in the data appears for $N\geq 5$. In
this case propagating at longer times in the future seems to require a deeper
memory. Similar results are found when constructing NNs using a single Bloch
vector as feature (see Fig. 6), in particular when the qubit count remains low
($N\leq 4$). Most importantly, in both situations the memory needed to
converge the NNs increases drastically with the number of qubits.

Figure 6: (Color on line) Error of the NN (the mean trace distance - MTD) as a
function of the memory, $h$, (in units of $\Delta=0.08\pi\approx 0.25/J$), for
NNs that propagate at different points in the future (color code). Here the
results are for NNs that use as feature the reduced density matrix of a single
qubit (‘single’ case).

The scaling of the memory with the number of qubits is presented in the left-
hand side panel of figure 7 for NNs propagating $0.48\pi\approx 1.5$ in the
future ($\Delta=0.08\pi$). This is a time significant longer than the de-
correlation time observed in Fig. 4. We take the operational definition of
‘necessary’ memory, $h_{\mathrm{nec}}$, as the memory needed to reduce the
error below 0.01, and this is plotted on a base-2 logarithmic scale against
$N$. We find that $h_{\mathrm{nec}}$ scales exponentially with $N$, namely
$h_{\mathrm{nec}}\propto 2^{\alpha N}$, with $\alpha\sim 1$ (this cannot be
determined with precision from our limited number of data points).
Interestingly, we find little difference in the $h_{\mathrm{nec}}$ scaling
between the ‘all’ and ‘single’ case, although for large $N$’s
$h_{\mathrm{nec}}$ is systematically lower when all Bloch vectors are used in
the NNs.

In Figure 7 we also plot $\log_{2}(h_{\mathrm{nec}}$), respectively as a
function of $-\log_{2}(f_{\mathrm{min}})$ (middle panel), where
$f_{\mathrm{min}}$ is the lowest frequency of a given system, and as a
function of the $\log_{2}$ of the number of unique frequencies associated to
the spectrum of the corresponding system (right-hand side panel). The
frequencies are computed as $\Delta\epsilon_{nm}=\epsilon_{n}-\epsilon_{m}$,
with $\epsilon_{n}$ being an eigenvalue of the Heisenberg Hamiltonian,
obtained by exact diagonalization. We find an extremely good correlation
between $h_{\mathrm{nec}}$ and the number of unique frequencies, suggesting
that the memory needed for an accurate propagation of the time propagation may
scale more favourably for systems presenting a limited number of frequencies
(e.g. in the case of high degeneracy). This correlation, however, may still be
accidental, so that it would be interesting to investigate a many-body
Hamiltonian where the multiplicity of the frequency spectrum presents a
scaling, as a function of the number of particles, different from that of the
Hilbert space.

Figure 7: (Color on line) Scaling of the memory needed to achieve accurate
propagation (see definition in the text), $h_{\mathrm{nec}}$, as a function
of: (left) the number of qubits, (middle) the minimum frequency,
$f_{\mathrm{min}}$, and (right) the number of unique frequencies of the
spectrum of the corresponding Heisenberg chain, $N_{\mathrm{uni}}^{f}$. Note
that all quantities are plot on a $\log_{2}$ scale. Light (dark) blue symbols
are for the ‘all’ (‘single’) case. Here data are presented for NNs predicting
$\Delta=0.48\pi\sim 1.5/J$ in the future.

###  IV.2 Propagation

In order to show that our NNs can act as true time propagators we need to
demonstrate that their time evolution can be concatenated, namely that one can
use the predictions made for time $t$ as an input for the subsequent
prediction at time $t+\Delta t$. Repeating such a process, an operation called
autoregression, allows one to follow the dynamics at, in principle, arbitrary
times. Such exercise is performed here for the case of six qubits, $N=6$, a
propagator with a time time step of $\Delta=0.16\pi\sim 0.5/J$ and a memory of
$h=45$. We consider NNs using the entire set of Bloch vectors as feature (the
‘all’ case) and we initiate the dynamics from a random state.

The results of this test are presented in Figure 8, where we plot the error
(the mean trace distance) as a function of time for propagation up to $100/J$,
namely for about 200 time steps. In order to minimize the error we train
multiple NNs along the same time trajectories and propagate the Bloch vectors
by using the average of the networks’ predictions. Thus, in Fig. 8 we show
results for a single NN (top panel), an average of three NNs (middle panel)
and one of ten (bottom panel). Two main conclusions can be drawn from the
figure. On the one hand, it is clear that our NNs are well capable of
performing autoregression, with the error growing relatively slowly in time.
On the other hand, the figure clearly shows that averaging over several NNs
significantly improves the predictions; the average is more accurate and the
fluctuations around the average are reduced. This effectively demonstrates
that our NNs can be used as universal time-evolution operators.

Figure 8: (Color on line) Demonstration of autoregression, where the output of
a NN at time $t$ is used as the input for the propagation at time $t+\Delta
t$. Results are presented for NNs with a time time step of $\Delta=0.16\pi\sim
0.5/J$ and a memory of $h=45$. The system comprises $N=6$ qubits, the initial
state is a random state and the NNs use the entire manifold of Bloch vectors
as feature. The graphs show the mean trace distance error as a function of
time for different trajectories (blue lines). The bold line corresponds to the
time-averaged mean error. In the top panel we use a single NN, in the middle
panel an ensemble of three NNs and in the bottom panel an ensemble of ten.
Note that both the time-averaged error and the fluctuations around the average
improve as one uses a large number of NNs.

When considering an ensemble of NNs one can then use the disagreement among
the NNs as a measure of the confidence over the prediction. This is
essentially the variance of the predicted quantity over the different NNs.
Such variance is plotted in the top panel of Figure 9 for the expectation
value of $\sigma_{z}$ along the time trajectory of a single qubit in a system
of $N=6$ qubits. Also in this case we simulate the time-evolution of $N=6$
qubits with a memory $h=45$. In this particular case the variance does not
significantly change over time, indicating that the accuracy is largely
preserved over the trajectory. However, one can also note that the variance is
larger along particular branches of the trajectory, where the predictions of
the NNs agree less well with the exact results. As such, the variance can be
used as a measure of the accuracy of the autoregression. Specifically, one can
monitor the increase in variance and use it as an indicator of the fidelity of
the prediction as a function of time.

Figure 9: (Color on line) Time trajectories generated via NN autoregression.
Here we consider $N=6$ qubits and monitor the time dependence of the
expectation value of $\sigma_{z}$ of one qubit. The memory in this case is
$h=45$. In the top panel we show the predicted trajectory (dots) against the
exact one (solid black line), together with the variance of the prediction
over ten NNs (blue shadow). The middle panel show the predictions obtained
with a time-propagator NN of step $\Delta=0.16\pi$, while the bottom one uses
the same predictions as starting point for time propagators with time steps
0.04$\pi$, 0.08$\pi$ and 0.12$\pi$. Note NNs trained to predict the dynamics
at different time scales can be effectively concatenated.

Finally, we show that NNs trained to propagate the dynamics at different time
steps can be concatenated, namely that the output of a network of time step
$\Delta t_{1}$ can be used as input for a NN with time step $\Delta t_{2}$.
This is shown in the two lower panels of Figure 9. Firstly, we present the
trajectory of the expectation value of $\sigma_{z}$ of one qubit, computed
with a NN using a time step of $\Delta=0.16\pi$. Then, we use the output of
such NN as input in NNs of steps $0.04\pi$, $0.08\pi$ and $0.12\pi$. This
effectively allows us to increase the density of the predicted points along
the trajectory. The figure shows clearly that such an operation is possible,
without any significant loss of accuracy. This is an important results, as it
allows one to construct a range of NNs, predicting at different times in the
future, and use appropriate combinations of them to reach any point in time
along the trajectory. Crucially, this means that the NNs autoregression can be
used to generate long-time and arbitrary dense time trajectories.

###  IV.3 Dynamics initiated from a product state

So far we have analyzed the dynamics initialized from a random state, namely
from a state with a high entanglement entropy. Here we repeat the analysis for
an initial state corresponding to a product state. Recalling figure 2, the
single-particle entropy of a product state grows with the time, meaning that a
NN trained from the early-time trajectory will not contain enough information
to reproduce the correct dynamics and cannot be used in an autoregression.
This is because the states encountered during the early dynamics of a product
state are qualitatively different from those encountered over long times. As
such, we have now trained NNs by using as an initial state the one evolved
from a product state at time $0.32\pi\approx 1/J$ (this time is sufficient to
reach equilibrium, see Fig. 2), and our results are summarised in Figure 10
for system with $N\leq 5$.

Figure 10: (Color on line) Error of the NN (the mean trace distance) as a
function of the memory, $h$, (in units of $\Delta=0.08\pi\approx 0.25/J$), for
NNs that propagate at different points in the future (color code). Here the
results are for NNs that use as feature either the reduced density matrix of
all qubits (‘all’ case - top panels) or of a single qubit (‘single’ case -
lower panels).

In the figure we show the NN error (the mean trace distance) as a function of
the memory for different time propagators computed for both the ‘all’ and
‘single’ case. The results are qualitatively similar to those encountered for
networks trained over dynamics initiated with random states (see Fig. 5 and
Fig. 6), except that the memory required to converge the propagator is now, in
general, significantly shorter. This follows from the fact that the dynamics
initiated from a product state spans only a subset of the entire Hilbert
space, an observation corroborated by the result that product states never
evolve to configurations with single-particle entropy close to that of random
states (see Fig. 2). Such behaviour is not surprising, given the high-symmetry
form of the Hamiltonian.

We investigate further this aspect by constructing an autoencoder
goodfellow2016deep  performing a non-linear compression of the wavefunctions
corresponding to both random and product states. To perform this task the
wavefunction is expanded over the full Fock space $\\{|i\rangle\\}$ and the
real and imaginary coefficients of expansion, $\psi_{i}=\langle
i|\psi\rangle$, are taken to form a 2$\Omega$-dimensional vector x with

$x_{j}=\begin{cases}\text{Re}(\psi_{j})\>\>\>\>\>\quad j\leq\Omega\\\
\text{Im}(\psi_{j-\Omega})\quad j>\Omega\\\ \end{cases}$ (13)

where $\sum_{i=1}^{\Omega}|\psi_{i}|^{2}=|\textbf{x}|^{2}=1$ and $\Omega$ is
the dimension of the Hilbert space. The encoder and decoder forming the
autoencoder are two fully connected NNs, with two 64-nodes hidden layers and
the ELU activation function. By minimizing the reconstruction error, as
measured by the Euclidean distance between the vector representations, we can
quantify the level of compression that a wavefunction can undergo at different
times. Figure 11 shows our results obtained with a dataset of 2,000 samples,
with 1,000 in the training set and 500 in the validation and test ones. For
this numerical experiment we consider $N=5$. Clearly, a product state can be
compressed much more efficiently than a random one (a 15-dimensional latent
space appears to be sufficient at all times), regardless of the time at which
the wavefunction is measured. Note here that for $N=5$ the dimension of the
Hilbert space is 32. Therefore, when the latent space is 64-dimensional there
is no compression and the autoencoder just learns the identify function. As
such, one expects the fidelity to approach unity regardless of the state. The
deviation observed for the random state is then attributed to relatively small
training set.

Figure 11: (Color on line) Wavefunction fidelity as a function of the
dimension of the latent space (reduced dimension) for wavefunction evolved to
times 0, 0.25, 0.5 and 1 (in units of $1/J$). Here we compare evolutions
initiated from either a random or a product state for $N=5$.

###  IV.4 Frequency of the memory sampling

We now investigate further the nature of the memory required for non-Markovian
dynamics. The first question we want to address is whether the convergence of
the memory is smooth or sharp, namely whether the error in the propagation
drops sharply when one reaches the required memory. This is explored in Figure
12, where we present the error of the NN (the mean trace distance of a single
NN for all the samples in the training set) as a function of the memory
duration for propagation to a time of $0.48\pi\approx 1.5/J$, where the memory
is sampled with the fine time resolution of $0.04\pi\approx 0.13/J$. The
exercise is performed for both $N=5$ (left-hand side panels) and $N=6$ (right-
hand side panels), in the ‘all’ (top panels) and ‘single’ (bottom panels)
cases. The graphs clearly show a very sharp transition, with the error
abruptly reducing below the 0.01 threshold as soon as the memory reaches a
critical duration. This means that, if the memory is accurately sampled, one
will just need to construct a NN time propagator with a memory longer than the
critical one. In fact, considering longer memories does not improve the
convergence. In other words, it appears that the range of the time-propagator
kernel is finite and well defined for any given system.

Figure 12: (Color on line) Error of the NN (the mean trace distance) as a
function of the memory, $h$, (in units of $\Delta=0.04\pi\sim 0.13/J$), for
$N=5$ and $N=6$ and both the ‘all’ (upper panels) and ‘single’ (lower panels)
case. In this case the NN propagates $0.48\pi\approx 1.5/J$ in the future. The
shaded region corresponds to the 25th and 75th error percentiles.

Having established that there is a critical memory for each system, one now
needs to find out how finely such memory should be sampled. This question is
answered in Figure 13. Here we have constructed a single NN to propagate a
$N=5$ system $0.48\pi\approx 1.5/J$ in the future (‘all’ case). The memory is
then sampled with different time steps, going from 0.04$\pi$ to 0.64$\pi$.
Surprisingly, we find that when the memory is finely sampled
($\Delta=0.04\pi$, $0.08\pi$ and $0.16\pi$) its duration does not change,
namely feature vectors of different length can equally well predict the
dynamics as long as enough time-evolution history is learned. In contrast,
more coarse samplings ($\Delta>0.16\pi$) require longer memories.
Intriguingly, in this case of coarse sampling, one approximately needs the
same number of time steps (about 14), although of different duration, to
converge the NN.

Notably, the highest frequency found in the spectrum of the $N=5$ spin chain
is $f_{\mathrm{max}}=0.99J$. According to the Shannon-Nyquist sampling theorem
allen2004signal  the largest possible period required to sample a time-
dependent dynamics is $\Delta_{\mathrm{max}}=1/2f_{\mathrm{max}}$, which in
our case is $\Delta_{\mathrm{max}}\approx 0.16\pi$. Intriguingly, this
corresponds to the critical sampling time above which the memory is no longer
time-step independent. We can then conclude that the two regimes found in
Figure 13 are simply separated by the Shannon-Nyquist limit. In general, for a
linear model the dynamics cannot be reproduced at all if sampled at a time
step larger than $\Delta_{\mathrm{max}}$. Here, however, our time-propagator
(the NN) is highly non-linear and the Shannon-Nyquist limit can be avoided by
sampling longer memories at a coarser resolution. Further tests using non-
linear memory samplings have not given conclusive results.

Figure 13: (Color on line) NN error (the mean trace distance) as a function of
the memory for the propagation to the time $0.48\pi\approx 1.5/J$ of a $N=5$
system. Different curves correspond to memories sampled with different time
steps. The numbers on each curve correspond to the number of time steps
included in the memory (the dimension of the time component of the feature
vector), while the $x$-axis scale is absolute ($\Delta=0.04\pi$). The black
diagonal line separates the results for which the memory duration is
independent from the sampling.

###  IV.5 Information locality

Finally we investigate how the information included in the feature vector
affects our ability to make predictions. For this experiment we take the $N=5$
spin chain with dynamics initiated at a random initial state. In this case we
compute the error (the mean trace distance) of a single NN as a function of
the memory for a generic case where the feature vector contains the Bloch
vectors of $m$ qubits. The error is then computed over both the $m$ qubits
used in the feature vector and the remaining $N-m$ ones. With this notation
$m=1$ and $m=5$ corresponds to the ‘single’ and ‘all’ case, respectively. Our
results are summarized in Figure 14.

In general we find that the models cannot satisfactory predict the evolution
of qubits, whose Bloch vector was not included in the training, although the
error reduces as $m$ gets larger. Interestingly, this is true even for the
$m=4$ case, where only the Bloch vector of one qubit is left outside the NN
feature vector. A perhaps more intriguing feature is found for the $m=3$ case.
In this situation the three qubits used for the training are inequivalent,
since only one neighbour other qubits included in training the model. In this
case we found that the memory required to bring the error below threshold is
different depending on the location of the qubit. This suggests that there may
be a trade off between time and space locality in the time propagator kernel.
Similar results are also found for the $N=6$ case (not presented here), where
several inequivalent configurations can be designed.

Figure 14: (Color on line) NN error (the mean trace distance) as a function of
the memory for the propagation to the time $0.48\pi\approx 1.5/J$ of a $N=5$
system. Here the NN is trained by including only the Bloch vector of $m$
qubits, while the remaining $N-m$ are not considered. In the graphs the error
over the qubits included in the training is in blue, while that of the qubits
outside the NN feature vector are in black.

##  V Conclusion

In this work we have explored the use of machine learning to propagate quantum
systems in time, using the Heisenberg Hamiltonian as many-body model. For
Markovian dynamics the time propagator can be learnt easily with a linear
regression, as long as the training dataset is sufficiently large. In
contrast, the propagation of non-Markovian systems requires a history, meaning
that the system state at a given time is determined by a number of states in
the past. The duration of such memory appears to scale exponentially with the
system size, regardless of whether one uses the density matrix of a single
qubit or of the entire ensemble as feature. However, shorter memories are
required when the system explores only a sub-set of the available spectrum
during the time evolution, as in the case of dynamics initiated from a product
state.

Crucially, we have shown that the machine-learning propagators can be
concatenated in an autoregression, meaning that the state evolved with one
neural network can be used as input for another propagation. This allows us to
propagate at arbitrary long times with any desired resolution. Our method can
then be applied to quantum dynamical data generated from any computational
scheme, whether it be propagation, tensor networks PAECKEL2019167998 ,
restricted Boltzmann machines carleo2017solving , or data obtained
experimentally. Furthermore, by using an ensemble of machine-learning
propagators we can maintain accuracy for a large number of iterations and
constantly estimate the error of our predictions.

Finally, we have investigate in detail the time resolution needed to represent
the system memory, and found two regimes separated by the Shannon-Nyquist
limit. Namely, when the memory is sampled at a time step shorter than the
period corresponding to the fastest frequency of the system, the memory
remains constant. This means that one has to sample a fixed time interval but
different time steps can be used. In contrast, if the time step is longer than
such period, the required memory is no longer constant, but the total number
of time steps is. This demonstrates that the non-linearity built in the neural
networks can overcome the limitation set by the Shannon-Nyquist sampling
theorem.

For future studies it will be useful to compare the method presented here with
the other approaches proposed for solving the reduced dynamics of the system,
such as the discussed Markovian embeddings, the Nakajima-Zwanzig technique and
also the transfer tensor methods. In particular, it will be interesting to see
if a direct connection between the memory kernel appearing in the Nakajima-
Zwanzig equation and the history depth $h$ that we observed for the trained
NNs can be made.

### Acknowledgment

This work has been supported by the Irish Research Council Postgraduate
Scholarship (J.N.) and Advance Laureate Award (S.S. - Award IRCLA/2019/127).
L.C. and G.K. acknowledge Science Foundation Ireland for financial support
through Career Development Award 15/CDA/3240. We acknowledge the
DJEI/DES/SFI/HEA Irish Centre for High-End Computing (ICHEC) and Trinity
Centre for High Performance Computing (TCHPC) for the provision of
computational resources. The authors thank Plamen Stamenov and John Goold for
interesting discussion.

##  Appendix A Nakajima-Zwanzig equation

Starting from the Liouville-Von Neumann equation,

$\partial_{t}\hat{\rho}=\frac{i}{\hbar}\left[\hat{\rho},\hat{H}\right]=\hat{L}\hat{\rho}\>,$
(14)

one splits the density operator up,
$\hat{\rho}=(\hat{\mathcal{P}}+\hat{\mathcal{Q}})\hat{\rho}$, by means of the
projection operators $\hat{\mathcal{P}}$ and $\hat{\mathcal{Q}}$, with
$\hat{\mathcal{Q}}=\hat{\mathcal{I}}-\hat{\mathcal{P}}$ ($\mathcal{I}$ is the
identity). The definition of $\hat{\mathcal{P}}$ and $\hat{\mathcal{Q}}$ is
somewhat unusual. Consider a reference state $\hat{\rho}_{\mathrm{B}}$ of the
environment and then define $\hat{\mathcal{P}}$ as the projector such that

$\hat{\rho}_{\text{red}}=\hat{\mathcal{P}}\hat{\rho}=\text{Tr}_{\mathrm{B}}(\hat{\rho})\otimes\hat{\rho}_{\mathrm{B}}\>,$
(15)

where $\rho_{\mathrm{red}}$ is the reduced density matrix of the degrees of
freedom of interest. Thus $\hat{\mathcal{P}}$ is defined as the partial
tracing out of the environment and then the product of the resulting state
with the predetermined environmental state, which is typically time
independent. The definition of $\hat{\mathcal{Q}}$ then follows as

$\hat{\mathcal{Q}}\hat{\rho}=\hat{\rho}-\text{Tr}_{\mathrm{B}}(\hat{\rho})\otimes\hat{\rho}_{\mathrm{B}}\>.$
(16)

Although unusual the projectors are perfectly well defined and it is easy to
show that

$\displaystyle\hat{\mathcal{P}}^{2}=\hat{\mathcal{P}}\>,$ (17)
$\displaystyle\hat{\mathcal{Q}}\hat{\mathcal{P}}=\hat{\mathcal{P}}\hat{\mathcal{Q}}=0\>.$
(18)

Starting from Eq. (14) one inserts $\mathcal{I}=\mathcal{P}+\mathcal{Q}$ to
obtain

$\displaystyle\partial_{t}\hat{\mathcal{P}}\hat{\rho}=\hat{\mathcal{P}}\hat{L}\hat{\mathcal{P}}\hat{\rho}+\hat{\mathcal{P}}\hat{L}\hat{\mathcal{Q}}\hat{\rho}\>,$
(19)
$\displaystyle\partial_{t}\hat{\mathcal{Q}}\hat{\rho}=\hat{\mathcal{Q}}\hat{L}\hat{\mathcal{Q}}\hat{\rho}+\hat{\mathcal{Q}}\hat{L}\hat{\mathcal{P}}\hat{\rho}\>,$
(20)

and the second identity can be formally solved as

$\displaystyle\hat{\mathcal{Q}}\hat{\rho}=e^{\hat{\mathcal{Q}}\hat{L}t}\hat{\mathcal{Q}}\hat{\rho}(0)+\int_{0}^{t}dt^{\prime}e^{\hat{\mathcal{Q}}\hat{L}t^{\prime}}\hat{\mathcal{Q}}\hat{L}\hat{\mathcal{P}}\hat{\rho}(t-t^{\prime})\>.$
(21)

This can then be inserted back into the first equation to give

$\displaystyle\partial_{t}\hat{\mathcal{P}}\hat{\rho}=\hat{\mathcal{P}}\hat{L}\hat{\mathcal{P}}\hat{\rho}+\hat{\mathcal{P}}\hat{L}e^{\hat{\mathcal{Q}}\hat{L}t}\hat{Q}\hat{\rho}(0)+$
(22)
$\displaystyle\hat{\mathcal{P}}\hat{L}\int_{0}^{t}dt^{\prime}e^{\hat{\mathcal{Q}}\hat{L}t^{\prime}}\hat{\mathcal{Q}}\hat{L}\hat{\mathcal{P}}\hat{\rho}(t-t^{\prime})\>.$
(23)

Finally, if we assume that at the time $t=0$ the system is a product state
$\hat{\rho}(0)=\hat{\rho}_{\mathrm{A}}\times\hat{\rho}_{\mathrm{B}}$, then
$\hat{\mathcal{P}}\hat{\rho}(0)=\hat{\rho}(0)$ and
$\hat{\mathcal{Q}}\hat{\rho}(0)=0$, so that we can drop the middle term to
obtain

$\partial_{t}[\hat{\mathcal{P}}\hat{\rho}]=\hat{\mathcal{P}}\hat{L}[\hat{\mathcal{P}}\hat{\rho}]+\int_{0}^{t}dt^{\prime}\hat{\mathcal{K}}(t^{\prime})[\hat{\mathcal{P}}\hat{\rho}(t-t^{\prime})]\>,$
(24)

where the time kernel writes

$\hat{\mathcal{K}}(t)=\hat{\mathcal{P}}\hat{L}e^{\hat{\mathcal{Q}}\hat{L}t}\hat{\mathcal{Q}}\hat{L}\hat{\mathcal{P}}\>.$
(25)

## References

  * [1] S. Nakajima, On Quantum Theory of Transport Phenomena: Steady Diffusion, Prog. Theor. Phys. 20, 948 (1958).
  * [2] R. Zwanzig, Ensemble method in the theory of irreversibility, J. Chem. Phys. 33 1338, (1960).
  * [3] I. Prigogine, Non-Equilibrium Statistical Mechanics, Interscience Publishers, New York, (1962).
  * [4] F. Shibata, Y. Takahashi and N. Hashitsume, A generalized stochastic Liouville equation. non-Markovian versus memoryless master equations, J. Stat. Phys. 17, 171 (1977).
  * [5] F. Shibata and T. Arimitsu, Expansion formulas in nonequilibrium statistical mechanics, J. Phys. Soc. Jpn. 49 891, (1980).
  * [6] F. Ciccarello and V. Giovannetti, A quantum non-Markovian collision model: incoherent swap case, Phys. Scr. 2013, 014010 (2013).
  * [7] R. McCloskey and M. Paternostro, Non-markovianity and system-environment correlations in a microscopic collision model, Phys. Rev. A 89, 052120 (2014).
  * [8] S. Kretschmer, K. Luoma and W.T. Strunz, Collision model for non-markovian quantum dynamics, Phys. Rev. A 94, 012106 (2016).
  * [9] S. Lorenzo, F. Ciccarello and G.M. Palma, Composite quantum collision models, Phys. Rev. A 96, 032107 (2017).
  * [10] B. Çakmak, M. Pezzutto, M. Paternostro and Ö.E. Müstecaplıoğlu, Non-markovianity, coherence, and system-environment correlations in a long-range collision model, Phys. Rev. A 96, 022109 (2017).
  * [11] S. Campbell, F. Ciccarello, G.M. Palma and B. Vacchini, System-environment correlations and markovian embedding of quantum non-markovian dynamics, Phys. Rev. A 98, 012142 (2018).
  * [12] A. Mardt, L. Pasquali, H. Wu and F. Noé, Vampnets for deep learning of molecular kinetics, Nature Commun. 9, 1 (2018).
  * [13] P.G. Breen, C.N. Foley, T. Boekholt and S.P. Zwart, Newton versus the machine: solving the chaotic three-body problem using deep neural networks, Mon. Not. R. Astron. Soc. 494, 2465 (2020).
  * [14] G. Carleo and M. Troyer, Solving the quantum many-body problem with artificial neural networks, Science 355, 602 (2017).
  * [15] M.J. Hartmann and G. Carleo, Neural-network approach to dissipative quantum many-body dynamics, Phys. Rev. Lett. 122, 250502 (2019).
  * [16] S.-H. Lin and F. Pollmann, Scaling of neural-network quantum states for time evolution, Phys. Status Solidi B 2100172 (2022).
  * [17] L. Banchi, E. Grant, A. Rocchetto and S. Severini, Modelling non-markovian quantum processes with recurrent neural networks, New J. Phys. 20, 123030 (2018).
  * [18] I.A. Luchnikov, S.V. Vintskevich, D.A. Grigoriev and S.N. Filippov, Machine learning non-markovian quantum dynamics, Phys. Rev. Lett. 124, 140502 (2020).
  * [19] P.P. Mazza, D. Zietlow, F. Carollo, S. Andergassen G. Martius and I. Lesanovsky, Machine learning time-local generators of open quantum dynamics, Phys. Rev. Research 3, 023084 (2021).
  * [20] I.A. Luchnikov, E.O. Kiktenko, M.A. Gavreev, H. Ouer- dane, S.N. Filippov and A.K. Fedorov. Probing non-markovian quantum dynamics with data-driven analysis: Beyond “black-box” machine learning models arXiv:2103.14490 (2021).
  * [21] J. Cerrillo and J. Cao, Non-markovian dynamical maps: Numerical processing of open quantum trajectories, Phys. Rev. Lett. 112, 110401 (2014).
  * [22] A. Gelzinis, E. Rybakovas and L. Valkunas, Applicability of transfer tensor method for open quantum system dynamics, J. Chem. Phys. 147, 234108 (2017).
  * [23] S. Gherardini, A. Smirne, S. Huelga and F. Caruso, Transfer-tensor description of memory effects in open-system dynamics and multi-time statistics, Quantum Sci. Technol. (2021).
  * [24] Z. Zhang, S. Yang, C. Liu, Y. Han, C.H. Lee, Z. Sun, G. Li and X. Zhang, Predicting quantum many-body dynamics with long short-term memory based neural networks, arXiv:1905.09168 (2019).
  * [25] Throughout this work we use the trace distance, $D$, as our metric for comparing density matrices with $D=\frac{1}{2}|\sigma-\rho|$. Note that $0\leq D\leq 1$ and for Bloch vectors x and y it reduces to $\frac{1}{2}|\textbf{x}-\textbf{y}|$.
  * [26] A. Paszke, et al., Pytorch: An imperative style, high-performance deep learning library, in “Advances in neural information processing systems”, pages 8026–8037, MIT Press (2019).
  * [27] J.R. Johansson, P.D. Nation and F. Nori, Qutip 2: A python framework for the dynamics of open quantum systems, Comp. Phys. Comm. 184, 1234 (2013).
  * [28] Since linear regressions require the weight matrix to be real we must write the wavefunction as a $2\Omega$ vector and, hence, require $2\Omega$ samples. For the density matrix one needs to evaluate $\Omega(1)+1/2(\Omega^{2}-\Omega)(2)=\Omega^{2}$ real numbers and hence we need $\Omega^{2}$ samples. We subtract 1 from both as we have the normalization constraint.
  * [29] H.-P. Breuer and F. Petruccione, The theory of open quantum systems, Oxford University Press (2002).
  * [30] I. Goodfellow, Y. Bengio and A. Courville, Deep learning, MIT press (2016).
  * [31] R.L. Allen and D. Mills, Signal analysis: time, frequency, scale, and structure, John Wiley & Sons (2004).
  * [32] S. Paeckel, T. Köhler, A. Swoboda, S.R. Manmana, U. Schollwöck and C. Hubig, Time-evolution methods for matrix-product states, Ann. Phys. 411, 167998 (2019).


# Mixture of Regression Experts in fMRI Encoding

Subba Reddy Oota*1, Adithya Avvaru1, Naresh Manwani1, Raju S. Bapi1,2
1 IIIT Hyderabad, 2 University of Hyderabad
<EMAIL_ADDRESS><EMAIL_ADDRESS>

###### Abstract

fMRI semantic category understanding using linguistic encoding models attempt
to learn a forward mapping that relates stimuli to the corresponding brain
activation. Classical encoding models use linear multi-variate methods to
predict the brain activation (all voxels) given the stimulus. However, these
methods essentially assume multiple regions as one large uniform region or
several independent regions, ignoring connections among them. In this paper,
we present a mixture of experts-based model where a group of experts captures
brain activity patterns related to particular regions of interest (ROI) and
also show the discrimination across different experts. The model is trained
word stimuli encoded as 25-dimensional feature vectors as input and the
corresponding brain responses as output. Given a new word (25-dimensional
feature vector), it predicts the entire brain activation as the linear
combination of multiple experts’ brain activations. We argue that each expert
learns a certain region of brain activations corresponding to its category of
words, which solves the problem of identifying the regions with a simple
encoding model. We showcase that proposed mixture of experts-based model
indeed learns region-based experts to predict the brain activations with high
spatial accuracy.

##  1 Introduction

In recent years, the use of both linear and non-linear multivariate
encoding/decoding approaches for analyzing fMRI brain activity have become
increasingly popular (Mitchell et al., 2008; Naselaris et al., 2011; Mesgarani
et al., 2014; Di Liberto et al., 2015; Pereira et al., 2018). An encoding
model that predicts brain activity in response to stimuli is important for
neuroscientists who can use the model predictions to investigate and test
hypotheses about the transformation from stimulus to brain response in
patients. Recent approaches for modeling fMRI data use word embedding
representations to build encoding systems (Oota et al., 2018; Abnar et al.,
2018). However, all these models do not have a principled way of learning
regions specializing in particular category of stimuli. Instead, they simply
predict the voxel intensity values either for the whole brain or for a pre-
selected set of voxels. In particular, interpreting the non-linear
multivariate models can be difficult due to capturing of unsuspected and
enigmatic representations (Benjamin et al., 2017; Kording et al., 2018).

Mixture of experts (MoE) models, proposed by Jordan’s group (Jordan & Xu,
1995) offer an interesting choice for the problem of learning distinct models
for different regions of the input space. MoE were applied to MRI data to
investigate the complex patterns of brain changes associated with non-
pathological and pathological processes, such as the effects of growth, aging,
injury or a disease (Kim et al., 2010; Eavani et al., 2016). MoEs have great
potential for use in medical diagnostics to diagnose a variety of clinical
conditions such as depression, Alzheimer’s dementia etc. (Yao et al., 2009)
proposed a Hidden Conditional Random Framework (HCRF) inspired from a mixture
of experts model to make predictions in all ROIs which are interconnected.

In this paper, we present a mixture of regression experts (MoRE) model that
predicts the brain activity associated with concrete nouns as stimuli. The
theory underlying this model is that each expert learns to cover a certain
brain region of interest (set of voxels that are significantly activated)
based on the category of words that are represented by the model. Finally,
fMRI signals are modeled as a linear combination of all expert predictions
corresponding to the gate probabilities of each expert. We present an
experimental evidence showing that the best encoding model is achieved with
the mixture of regression experts rather than using a simple linear/non-linear
model alone.

##  2 fMRI Encoding: Mixture of Regression Experts (MoRE) Approach

Figure 1: The sequence of slices, for the word “Apartment” after converting
voxel activation per subject into 23 slices, (i) in actual brain activation
(top row), (ii) for the model trained with mixture of experts model (second
row from top), (iii) for the model trained with multilayer perceptron model
(third row from top), (iv) for the model trained with ridge regression (bottom
row).

Current encoding methods attempt to learn either weights in case of linear
models (Mitchell et al., 2008; Pereira et al., 2018) or complex
representations in non-linear models (Oota et al., 2018). In the next
sections, we discuss the proposed MoRE approach and our enhancements.

###  2.1 Architecture:

We use a mixture of experts-based encoder model, whose architecture is
inspired from Jordan & Xu (1995). The mixture of experts architecture is
composed of a gating network and several expert networks, each of which solves
a function approximation problem over a local region of the input space.

Let $D=\\{X,Y\\}$ denote the data where $X=\\{x^{n}\\}_{n=1}^{N}$ is the
input, $Y=\\{y^{n}\\}_{n=1}^{N}$ is the target, $K$ is number of experts and
$N$ is the number of training points. Also, let
$\Theta=\\{\theta_{0},W_{e},\Sigma_{e}\\}$ denote the set of parameters where
$\theta_{0}$ is the set of gate parameters, $W_{e}$ is the set of expert
parameters, and $\Sigma_{e}$ is the set of expert co-variance matrices.

The mixtures of experts model is based on the following conditional mixture

$p(\mathbf{y}|\mathbf{x})=\sum_{j=1}^{K}P(j|\mathbf{x})p(\mathbf{y}|\mathbf{x},\theta_{j})=\sum_{j=1}^{K}g_{j}(\mathbf{x},\theta_{0})p(\mathbf{y}|\mathbf{x},\theta_{j})$
(1)

where $g_{j}(\mathbf{x},\theta_{0})$ is the gate’s rating, i.e., the
probability of the $j^{th}$ expert, given $\mathbf{x}$ and
$p(\mathbf{y}|\mathbf{x},\theta_{j})$ denotes the density function for the
output vector associated with the $j^{th}$ expert, i.e., the probability of
$j^{th}$ expert generating $\mathbf{y}$ given $\mathbf{x}$.

Here, we choose $p(\mathbf{y}|\mathbf{x},\theta_{j})$ as multivariate Gaussian
probability density for each of the experts, denoted by.

$p(\mathbf{y}|\mathbf{x},W_{j},\Sigma_{j})=\frac{1}{(2\pi)^{m/2}|\Sigma_{j}|^{1/2}}\exp\left(-\frac{1}{2}(\mathbf{y}-W_{j}\mathbf{x})^{T}\Sigma_{j}^{-1}(\mathbf{y}-W_{j}\mathbf{x})\right)$
(2)

where $W_{j}\in\mathbb{R}^{m\times n}$ is the weight matrix and
$\Sigma_{j}\in\mathbb{R}^{m\times m}$ is the covariance matrix associated with
the $j^{th}$ expert. In our problem, we consider two different cases of
$\Sigma_{j}$.

  1. 1.

$\Sigma_{j}=\sigma_{j}^{2}I,\;\forall j\in[K]$ and

  2. 2.

$\Sigma_{j}=diag(\sigma_{j,1}^{2},\sigma_{j,2}^{2},\ldots,\sigma_{j,m}^{2}),\;\forall
j\in[K]$.

Thus, we assume that the components of the output vector ${\bf
y}\in\mathbb{R}^{m}$ are statistically independent of each other. We use this
assumption to make the model very simple and reduce the number of overall
parameters. This assumption will also make the algorithm computationally less
expensive.

Assuming
$\Sigma_{j}=diag(\sigma_{j,1}^{2},\sigma_{j,2}^{2},\ldots,\sigma_{j,m}^{2})$,
we rewrite the conditional probability density model for $j^{th}$ expert as
follows:

$P({\bf y}|{\bf
x},W_{j},\Sigma_{j})=\frac{1}{(2\pi)^{m/2}(\sigma_{j,1}\sigma_{j,2}..\sigma_{j,m})}\exp\left(-\sum_{i=1}^{m}\frac{(y_{i}-{\bf
w}_{j,i}^{T}\mathbf{x})^{2}}{2\sigma_{j,i}^{2}}\right)$ (3)

where ${\bf w}_{j,i}$ is the $i^{th}$ row of $W_{j}$.

We use softmax function for $g(\mathbf{x},\theta_{0})$.

$g_{j}(\mathbf{x},\theta_{0})=\frac{\exp\left({\bf
v}_{j}^{T}\mathbf{x}\right)}{\sum_{i=1}^{K}\exp\left({\bf
v}_{i}^{T}\mathbf{x}\right)}$ (4)

where ${\bf v}_{j}\in\mathbb{R}^{n},\;\forall j\in[K]$. Thus,
$\theta_{0}=\\{{\bf v}_{1},\ldots,{\bf v}_{K}\\}$.

##  3 Experiments

#### Dataset

Figure 2: MoRE log-likelihood convergence

We used data from fMRI experiment 1 (Mitchell et al., 2008), where authors
conducted experiments with multiple subjects by showing different categories
of words as stimuli. Each category might correspond to activation of distinct
brain regions. In the experiment, the target word was presented with a picture
that depicted some aspect(s) of the relevant meaning. This fMRI dataset was
collected from a total of 9 participants. For each participant in the
experiment, a total set of 60 words (12 categories) were used as stimuli in
multi-modal form (word, picture). The fMRI dataset constitutes $51\times 61$
voxel windows arranged as 23 slices, per subject per stimulus. We use the
publicly available Mitchell’s 25-feature vector data (Mitchell et al., 2008)
of 60 words as input and the corresponding brain response of each participant
(containing 21000 voxels) as output to train the model.

Figure 3: Comparison of brain activity (ROIs) in two experts across slices.
The squared box indicates the ROIs’ activation differences: i) Expert-1 has
higher brain activity in “cingulate region" and “Superior and Middle frontal
gyrus regions”(top row), ii) ROIs captured in Expert2 are “occipital lobe"
from visual region”, “Inferior temporal gyrus” and “Middle temporal gyrus”
(bottom row)

#### Selection of Number of Experts:

In this paper, we selected the number of experts based on the number of
categories considered in the dataset. The dataset is divided into 12
categories in which each category has 5 concrete nouns. The idea is that each
expert should master a particular category and the corresponding activation in
relevant brain regions. The detailed selecting number expert models using
Bayesian Information Criteria provided in Appendix A.

#### Description of the Experiment

Here, we considered Mitchell’s feature vector, $\mathbf{x}$ as input to all
the experts and calculated $p(\mathbf{y}|\mathbf{x},\theta_{j})$ where
$\mathbf{y}$ is the corresponding brain responses for that $\mathbf{x}$.
Gating network trains to pick the expert which gives maximum
$g_{j}(\mathbf{x},\theta_{0})$. The detailed training methodology is provided
in Appendix B.

#### Results and Discussion

Using the approach discussed in Section 2 and Appendix B, we trained separate
mixture of regression experts model for each subject. The encoding performance
was evaluated by training and testing models using different subsets of the 60
concrete words in a 5-fold cross-validation scheme. The encoder models were
trained until the model reached convergence as shown in Figure 2. Figure 3
shows the region of activity in two different experts for the word “Church”.
Across all the slices, there is a considerable number of ROI differences among
all the experts. In particular, we depict regions of activation for two
experts in Figure 3. Figure 3 shows that Expert 1 has higher brain activity in
“cingulate region" and “Superior and Middle frontal gyrus regions”. In Expert
2, the ROIs captured are “Occipital lobe" from visual region”, “Inferior
temporal gyrus” and “Middle temporal gyrus”. It appears that while Expert 1
mediates activation in the frontal regions usually involved in decision
making, Expert 2 seems to activate regions responsible for visual processing
and object recogntion. In addition, we observed some brain activity being
common to ROIs across all the experts (results not shown here). In order to
assess the similarity between the actual and predicted activation in each
brain slice, we compared slice-wise voxel coordinates and their intensity
values. We measured the model scores using predictive $r^{2}$ score, for voxel
$v$ as follows:
$r_{v}^{2}=1-\frac{||\mathbf{y}^{v}_{test}-\mathbf{y}^{v}_{pred}||^{2}}{||\mathbf{y}^{v}_{test}-mean(\mathbf{y}^{v}_{test})||^{2}}$,
where $\mathbf{y}^{v}_{test}$ is the testing set brain activations,
$\mathbf{y}^{v}_{pred}$ is the predicted brain activations and
mean($\mathbf{y}^{v}_{test}$) is the mean activation on the test set. Table 1
depicts subject-wise (9 subjects) $r^{2}$ score comparisons among the mixture
of experts model, the model trained with the multilayer perceptron and with
the ridge regression model.From Table 1, we observe that $r^{2}$ score for
mixture of experts model is better than that of the other two models and the
score range indicates that MoRE model typically yields better predictions.

| Subjects ($r^{2}$ score)
---|---
Method | (1) | (2) | (3) | (4) | (5) | (6) | (7) | (8) | (9)
Mixture of Regression Experts | 0.811 | 0.813 | 0.806 | 0.811 | 0.814 | 0.815 | 0.813 | 0.808 | 0.813
Multilayer Perceptron | 0.755 | 0.745 | 0.745 | 0.772 | 0.75 | 0.765 | 0.756 | 0.761 | 0.751
Ridge Regression | 0.772 | 0.769 | 0.758 | 0.776 | 0.778 | 0.782 | 0.776 | 0.770 | 0.771
Table 1: Subject-wise $r^{2}$ scores for different models

##  4 Conclusion

In this paper, we present a mixture of experts-based model where group of
experts capture brain activity patterns related to particular regions of
interest (ROIs) and also show the discrimination across different experts.
Different from previous work, the underlying model depicts that each expert
trains on certain brain regions of interest (set of voxels that are
significantly activated) based on the category of words that are represented
by the model. The key distinction of our work is the utilization of experts
model to capture the ROIs based on various categories of input. In future, we
plan to experiment on all paradigms and conduct experiments on other fMRI
datasets and identify the number of experts using Bayesian Information
Criterion (BIC), with a primary focus on hierarchical mixture of experts at
slice-level instead of voxel-level predictions.

## References

  * Abnar et al. (2018) Samira Abnar, Rasyan Ahmed, Max Mijnheer, and Willem Zuidema.  Experiential, distributional and dependency-based word embeddings have complementary roles in decoding brain activity.  In _Proceedings of the 8th Workshop on Cognitive Modeling and Computational Linguistics (CMCL 2018)_ , pp. 57–66, 2018.
  * Benjamin et al. (2017) Ari S Benjamin, Hugo L Fernandes, Tucker Tomlinson, Pavan Ramkumar, Chris VerSteeg, Lee Miller, and Konrad Paul Kording.  Modern machine learning far outperforms glms at predicting spikes.  _BioRxiv_ , pp. 111450, 2017.
  * Di Liberto et al. (2015) Giovanni M Di Liberto, James A O’Sullivan, and Edmund C Lalor.  Low-frequency cortical entrainment to speech reflects phoneme-level processing.  _Current Biology_ , 25(19):2457–2465, 2015.
  * Eavani et al. (2016) Harini Eavani, Meng Kang Hsieh, Yang An, Guray Erus, Lori Beason-Held, Susan Resnick, and Christos Davatzikos.  Capturing heterogeneous group differences using mixture-of-experts: Application to a study of aging.  _Neuroimage_ , 125:498–514, 2016.
  * Jordan & Xu (1995) Michael I Jordan and Lei Xu.  Convergence results for the em approach to mixtures of experts architectures.  _Neural Networks_ , 8(9):1409–1431, 1995.
  * Kass & Raftery (1995) Robert E Kass and Adrian E Raftery.  Bayes factors.  _Journal of the American Statistical Association_ , 90(430):773–795, 1995.
  * Kim et al. (2010) Seyoung Kim, Padhraic Smyth, and Hal Stern.  A bayesian mixture approach to modeling spatial activation patterns in multisite fmri data.  _IEEE Transactions on Medical Imaging_ , 29(6):1260–1274, 2010.
  * Kording et al. (2018) Konrad Paul Kording, A Benjamin, Roozbeh Farhoodi, and Joshua I Glaser.  The roles of machine learning in biomedical science.  In _Frontiers of Engineering: Reports on Leading-Edge Engineering from the 2017 Symposium_. National Academies Press, 2018.
  * Mesgarani et al. (2014) Nima Mesgarani, Connie Cheung, Keith Johnson, and Edward F Chang.  Phonetic feature encoding in human superior temporal gyrus.  _Science_ , 343(6174):1006–1010, 2014.
  * Mitchell et al. (2008) Tom M Mitchell, Svetlana V Shinkareva, Andrew Carlson, Kai-Min Chang, Vicente L Malave, Robert A Mason, and Marcel Adam Just.  Predicting human brain activity associated with the meanings of nouns.  _science_ , 320(5880):1191–1195, 2008.
  * Naselaris et al. (2011) Thomas Naselaris, Kendrick N Kay, Shinji Nishimoto, and Jack L Gallant.  Encoding and decoding in fmri.  _Neuroimage_ , 56(2):400–410, 2011.
  * Oota et al. (2018) Subba Reddy Oota, Naresh Manwani, and Raju S. Bapi.  fmri semantic category decoding using linguistic encoding of word embeddings.  _ArXiv Preprint ArXiv:1806.05177_ , 2018.
  * Pereira et al. (2018) Francisco Pereira, Bin Lou, Brianna Pritchett, Samuel Ritter, Samuel J Gershman, Nancy Kanwisher, Matthew Botvinick, and Evelina Fedorenko.  Toward a universal decoder of linguistic meaning from brain activation.  _Nature Communications_ , 9(1):963, 2018.
  * Schwarz (1978) G Schwarz.  Estimating the dimension of a model the annals of statistics 6 (2), 461–464.  _URL: http://dx. doi. org/10.1214/aos/1176344136_ , 1978.
  * Yao et al. (2009) Bangpeng Yao, Dirk Walther, Diane Beck, and Li Fei-Fei.  Hierarchical mixture of classification experts uncovers interactions between brain regions.  In _Advances in Neural Information Processing Systems_ , pp. 2178–2186, 2009.

##  Appendix A Selecting Number of Experts using BIC

In order to find the number of experts, Bayesian Information Criteria (BIC) is
one of the successful measure to approximate the Bayes factor Kass & Raftery
(1995) i.e to find a model that has maximum posterior probability or maximum
marginal likelihood. The BIC can be formulated as follows.

$\displaystyle BIC=d\log(n)-2\log(p(X|\Theta))$ (5)

in Schwarz (1978), where d is the dimensionality of the parameter space, n is
the number of data points, $p(X|\Theta)$ denotes the parametric model density.

##  Appendix B Training Mixture of Experts using EM Algorithm

Let $N$ be the number of samples, $K$ be the number of experts and
$S=\\{(\mathbf{x}_{1},\mathbf{y}_{1}),\ldots,(\mathbf{x}_{N},\mathbf{y}_{N})\\}$
be the training set where
$(\mathbf{x}^{n},\mathbf{y}^{n})\in\mathbb{R}^{k}\times\mathbb{R}^{m},\;\forall
n\in[N]$.

The EM algorithm is an iterative method for finding the maximum likelihood
(ML) of a probability model in which some random variables are observed and
others are hidden. In training the Mixture of Experts, the indicator variables
$Z=z_{j}^{n},\;j\in[K],\;n\in[N]$ are introduced to solve the model with the
EM algorithm.

$\displaystyle z_{j}^{n}=\begin{cases}1,\text{if $\mathbf{y}^{n}$ is generated
from the $j^{th}$ expert}\\\ \text{0, otherwise}\end{cases}$

With the indicator variables, the likelihood can be written as

$\displaystyle L(\Theta;S;Z)$
$\displaystyle=P(S,Z|\Theta)=\prod_{n=1}^{N}P(\mathbf{y}^{n}|\mathbf{x}^{n})$
$\displaystyle=\prod_{n=1}^{N}\sum_{j=1}^{K}[g_{j}(\mathbf{x}^{n},\theta_{0})P(\mathbf{y}^{n}|\mathbf{x}^{n},W_{j},\Sigma_{j})]^{z_{j}^{n}}$
(6)

The parameter $\Theta$ can be estimated by the iterations through the E and M
steps.

###  B.1 E-Step : Calculation of Expectation of log-likelihood

The EM algorithm is employed to average out $z_{j}$ and maximize the expected
complete data log-likelihood $E_{Z}(logP(D,Z|\Theta))$ and
$\sum_{j=1}^{K}z_{j}^{n}$ =1, for each $n$. The expectation of the log-
likelihood in (5) results in

$\displaystyle Q(\Theta|\Theta^{(p)})$ $\displaystyle=E_{Z}(\log
P(D,Z|\Theta))$
$\displaystyle=\sum_{n=1}^{N}\sum_{j=1}^{K}h_{j}^{(p)}(n)\big{\\{}log([g_{j}(\mathbf{x}^{n},\theta_{0})P(\mathbf{y}^{n}|\mathbf{x}^{n},W_{j},\Sigma_{j})])\big{\\}}$
$\displaystyle=\sum_{n=1}^{N}\sum_{j=1}^{K}h_{j}^{(p)}(n)\log(g_{j}(\mathbf{x}^{n},\theta_{0}))+\sum_{n=1}^{N}\sum_{j=1}^{K}h_{j}^{(p)}(n)\log(P(\mathbf{y}^{n}|\mathbf{x}^{n},W_{j},\Sigma_{j})))$
(7)

where $p$ is the iteration index and $h_{j}^{(p)}(n)$ is expectation of
indicator variables which is given by

$\displaystyle h_{j}^{(p)}(n)$ $\displaystyle=E[z_{j}^{n}|D]$
$\displaystyle=P(z_{j}^{n}=1|\mathbf{x}^{(n)},\mathbf{y}^{(n)})$
$\displaystyle=\frac{g_{j}(\mathbf{x}^{(n)},\theta_{0})P(\mathbf{y}^{(n)}|\mathbf{x}^{(n)},W_{j},\Sigma_{j})}{\sum_{i=1}^{K}g_{i}(\mathbf{x}^{(n)},\theta_{0})P(\mathbf{y}^{(n)}|\mathbf{x}^{(n)},W_{i},\Sigma_{i})}$
(8)

###  B.2 M-Step : Estimation of new parameters $\Theta^{(p+1)}$

  * •

From E-step, we calculated the expectation of log-likelihood $Q$, we now
investigate the implementation of M-step by estimating the new parameters.

  * •

The M step chooses a parameter $\Theta$ that increases the $Q$ function, the
expected value of the complete-data log-likelihood i.e.

$\displaystyle\Theta^{(p+1)}=\underset{\Theta}{\text{argmax}}\hskip
2.0ptQ(\Theta|\Theta^{(p)})$ (9)

  * •

An iteration of EM also increases the original log likelihood $L$. That is
$L(\Theta^{(p+1)};Y)>L(\Theta^{(p)};Y)$ i.e. the likelihood $L$ increases
monotonically along the sequence of parameter estimates generated by an EM
algorithm.

#### Update Gate Parameters

  * •

From the Equations (4) and (6)

$E_{gate}=-\sum_{t=1}^{N}\sum_{j=1}^{K}h_{j}^{(p)}(n)log(g_{j}(\mathbf{x}^{(n)},\theta_{0})$
(10)

  * •

An interpretation of the cost function in Equation (6) is the cross-entropy of
distributing a pattern $\mathbf{x}$ amongst the expert networks. This cost is
minimal if experts are mutually exclusive and increases when experts share a
pattern.

$\displaystyle\partialderivative{Q}{\theta_{0}}$
$\displaystyle=\sum_{n=1}^{N}\sum_{j=1}^{K}h_{j}^{(p)}(n)\partialderivative{g_{j}(\mathbf{x}^{(n)},\theta_{0})}{\theta_{0}}/g_{j}(\mathbf{x}^{(n)},\theta_{0})$
$\displaystyle=\sum_{n=1}^{N}\sum_{j=1}^{K}h_{j}^{(p)}(n)[\partialderivative{({\bf
v}_{j}^{T}\mathbf{x}^{n})}{\theta_{0}}-\sum_{i=1}^{K}g_{i}(\mathbf{x}^{(n)},\theta_{0})\partialderivative{({\bf
v}_{i}^{T}\mathbf{x}^{n})}{\theta_{0}}]$
$\displaystyle=\sum_{n=1}^{N}\sum_{j=1}^{K}[h_{j}^{(p)}(n)-g_{j}(\mathbf{x}^{(n)},\theta_{0})]\partialderivative{({\bf
v}_{j}^{T}\mathbf{x}^{n})}{\theta_{0}}$ (11)

As the gating network is a specific form of a generalized linear model, the
general update for $\theta_{0}$ is written as follows

$\theta_{0}^{(p+1)}=\theta_{0}^{(p)}+\eta_{g}\partialderivative{Q}{\theta_{0}}$
(12)

  * •

Calculate the log likelihood using both $\theta_{0}^{(p)},\theta_{0}^{(p+1)}$
in Equation (5), and the difference in log likelihood values is calculated as

$\displaystyle\triangledown L=L(\Theta^{p};S;Z)-L(\Theta^{p+1};S;Z)$ (13)

  * •

If $\triangledown L$ is below a threshold, then new parameters are discarded
and MOE model will use old parameters.

  * •

There is no known closed-form way to find the minimum of cost function, and
thus we’ll resort to an iterative optimization algorithm such as gradient
descent or L-BFGS.

#### Update Expert Parameters

  * •

From the Equations (3) and (6)

$E_{expert}=-\sum_{n=1}^{N}\sum_{j=1}^{K}h_{j}^{(p)}(n)log(P(\mathbf{y}^{n}|\mathbf{x}^{n},W_{j},\Sigma_{j}))$
(14)

  * •

The update equation for expert parameters $W_{j}$ as follows

$\displaystyle\partialderivative{Q}{W_{j}}$
$\displaystyle=\sum_{n=1}^{N}\sum_{j=1}^{K}h_{j}^{(p)}(n)\partialderivative{P(\mathbf{y}^{(n)}|\mathbf{x}^{(n)},W_{j})}{W_{j}}/P(\mathbf{y}^{(n)}|\mathbf{x}^{(n)},W_{j})$
$\displaystyle=\sum_{n=1}^{N}\sum_{j=1}^{K}h_{j}^{(p)}(n)\partialderivative{(W_{j}^{T}\mathbf{x}^{(n)})}{W_{j}}\frac{[\mathbf{y}^{(n)}-W_{j}^{T}\mathbf{x}^{(n)}]}{\Sigma_{j}}$
$\displaystyle=\sum_{n=1}^{N}\sum_{j=1}^{K}h_{j}^{(p)}(n)\mathbf{x}^{(n)}\frac{[\mathbf{y}^{(n)}-W_{j}^{T}\mathbf{x}^{(n)}]}{\Sigma_{j}}$
(15)

$W_{j}^{(p+1)}=W_{j}^{(p)}+\eta_{j}\partialderivative{Q}{W_{j}}$ (16)

  * •

Calculate the log likelihood using both $W_{j}^{(p)},W_{j}^{(p+1)}$ in
Equation (5), and the difference in log likelihood values is calculated as

$\displaystyle\triangledown L=L(\Theta^{p};S;Z)-L(\Theta^{p+1};S;Z)$ (17)

  * •

If $\triangledown L$ is below a threshold, then new parameters are discarded
and MOE model will use old parameters.

#### Update Expert Covariance Matrix

  * •

The update equation for covariance matrix as follows.

$\Sigma_{j}^{(p+1)}=\frac{1}{\sum_{n=1}^{N}h_{j}^{(p)}(n)}\sum_{n=1}^{N}h_{j}^{(p)}(n)[\mathbf{y}^{(n)}-W_{j}^{T}\mathbf{x}^{(n)}]^{2}$
(18)

  * •

Calculate the log likelihood using both $\Sigma_{j}^{(p)},\Sigma_{j}^{(p+1)}$
in Equation (5), and the difference in log likelihood values is calculated as

$\displaystyle\triangledown L=L(\Theta^{p};S;Z)-L(\Theta^{p+1};S;Z)$ (19)

  * •

If $\triangledown L$ is below a threshold, then new parameters are discarded
and MOE model will use old parameters.


# Automated Configuration Synthesis for Machine Learning Models: A git-Based
Requirement and Architecture Management System

Abdullatif AlShriaf  Chalmers and University of Gothenburg
Gothenburg, Sweden
Email<EMAIL_ADDRESS>Hans-Martin Heyn  Chalmers and University of
Gothenburg
Gothenburg, Sweden
Email<EMAIL_ADDRESS>Eric Knauss  Chalmers and University of Gothenburg
Gothenburg, Sweden
Email<EMAIL_ADDRESS>
(March 2024)

##  I Background

inline]Things that I feel would be really interesting to write about in this
context:
\- The idea of extracting data from semi-formal requirements. This will make
it more attractive to keep requirements updated with code.
\- The idea to keep requirements and corresponding architectural decisions
aligned with code in git.
But I am not so really sure what is the key message that we have? (Eric)

inline]We suggest the following format for the extended abstract (2 pages
max).
Context (including, e.g., users or stakeholders of the proposed
tool/method/technique);
RE problem (challenge) addressed and Motivation;
Methodology used or the tool you developed;
any Validation/Evaluation conducted or planned;
Results and Contributions,
and a discussion of Related Work

Optional and recommended:
A set of up to 3 questions that the authors would like the attendees of the
conference to answer to help them validate their work. The questions shall be
complemented with a short description on why these questions are important and
how the answers would help in the further work. During the Tool and Demo
session, each stand will have a box and printed-out questions to answer them
on a paper, as well as links to an electronic form to answer the questions
electronically. inline]An annex to the abstract (max 2 pages) with:
For posters and student research posters :
A draft of a poster to which a short description explaining clearly how the
poster is to be presented at the conference can be added.
For tool demos:
A short description or a link to video explaining clearly how the work
described in the extended abstract is to be presented at the conference. The
annex should emphasize interaction potential. It may contain images of
screenshots or URLs leading to more information about the work, e.g., a video
clip.

The design of complex distributed systems typically follows a hierarchical
process, supported by highly specialized views for decomposing the design
task. Requirements and architecture often evolve simultaneously, requiring an
architectural framework that supports integrated and collaborative design,
including non-functional requirements and quality views. The framework must
ensure the traceability of design decisions in order to build safety cases.
Integrating requirements into software development is vital for aligning
intended functionality with implemented code. However, extracting data from
semi-formal requirements and maintaining alignment poses challenges due to its
ambiguity and variability making extracting consistent information
challenging. Aligning these requirements with other project artifacts can also
be difficult due to interpretation differences, often requiring manual effort
and leading to complexity and potential inconsistencies in development [1].

##  II Introduction

We propose GRAMS111https://gitlab.com/latiif/configgen as a git-based
Requirements and Architecture Management System which we demonstrate can
automate configuration extraction, reduce manual effort, enhance
configurability, and simplify runtime configuration management. We will show
that it improves traceability by establishing clear links between
requirements, configurations, and code, and facilitates maintenance through a
structured approach. GRAMS is intended for use by i) machine learning
engineers and architects tasked with designing the structure of ML-based
software systems, and ii) system administrators/MLOps engineers responsible
for deploying and managing the runtime environment of software systems. GRAMS
builds upon the open-source tool _T-Reqs_ for managing textual requirements in
git [2]. It extends T-Reqs by implementing a structured architectural view
system based on the concept of a compositional architectural framework for the
development of AI-enabled software systems (Figure 1, [3]). Using a schema-
type template, GRAMS compiles a machine-readable YAML _intermediary_ file to
communicate AI model specifications. This file can then be forwarded and
converted as part of a deployment-toolchain, particularly for example with the
purpose of AI model optimization. inline]If we present both GRAMS and T-Reqs,
we need to be very clear about the differences. Is GRAMS simply the
implementation of the compositional architectural framework in T-Reqs? Since
we are then not arguing with agile, what other argument do we have that git is
a good idea? (Eric)
(latiif): The tool GRAMS operates on T-Reqs-implementations of Arch.
Framework: Identifies explicit runtime-scenarios, Heurestically traverses the
AF, Locates configurations/schemes & maps them to the runtime-scenarios.
For me, GRAMS is the extended TTIM (2 extra types) + Traversal Algorithm +
Aggregation(as YAML or PlantUML)

##  III Context: Compositional Architectural Framework

The system design approach supported by GRAMS follows a scalable and
compositional architectural framework that provides flexibility to manage
requirements in software projects of varying complexity. The architectural
framework organizes quality concerns, including ethical considerations such as
explainability and security. Supporting middle-out design, GRAMS integrates
existing design decisions, enhancing alignment between requirements and
architecture [3]. Guiding rules ensure traceability, vital for managing
requirements throughout development. The framework’s adaptability to diverse
abstraction levels enables GRAMS to cater to varied project needs.

Figure 1: The VEDLIoT Architectural Framework [3].

##  IV T-Reqs: Text-based Requirements Management System

GRAMS extends T-Reqs, an open-source tool for the management of textual
requirements in git. T-Reqs is developed at Chalmers and University of
Gothenburg and used in industrial settings, including at Ericsson [4]. T-Reqs
utilizes a TTIM, or Type and Trace Information Model, that extends the
functionality of the basic Traceability Information Model (TIM) by
incorporating type information. While a TIM primarily focuses on representing
trace-links and relationships between abstract trace artifacts [5], the TTIM
defines types of traced elements and their connections, which can be optional
or required [2]. T-Reqs operates as a static analysis tool on the model
described by the TTIM.

##  V Applying GRAMS

First, GRAMS identifies explicit runtime scenarios, then heuristically
traverses the architectural framework, and finally locating
configurations/schemes and mapping them to the identified runtime scenarios.
In essence, GRAMS executes the TTIM, as illustrated in Figure 2, along with a
traversal algorithm which is outlined in Figure 3. Finally, the aggregation
capabilities of the system are presented in generic formats such as YAML or
PlantUML. Leveraging the architectural structure laid by the framework, we
introduce two new elements, as shown in Figure 2:

  * •

runtime-scenario: matches particular models to specific contexts and links to
a level of abstraction which GRAMS will scan and traverse all the way to an
AI-model optimiser input referred to as OptimizerInput.

  * •

schema-type: Annotates the OptimizerInput and contains the actual schema for
the JSON object present in the instantiating OptimizerInput.

Figure 2: Extended Architectural Framework for GRAMS: New types added for
GRAMS are highlighted in blue

Multiple runtime-scenarios can exist in the framework. From a single runtime-
scenario, multiple OptimizerInputs can be traced. The traversal follows the
algorithm outlined in Figure 3.

Figure 3: Framework traversal algorithm for GRAMS in pseudo-code.

##  VI Quality Assurance Checks in GRAMS

To ensure output quality and usability, GRAMS conducts the following checks:

  * •

Consistency is crucial to prevent conflicts and maintain coherence [6]. The
initial ”dumb” check ensures alignment with the framework’s meta-model (see
Section III).

  * •

GRAMS validates the internal correctness of the configuration-relevant
requirements specification by checking JSON schema compliance with schema-type
and OptimizerInput.

  * •

GRAMS verifies semantic equivalence between relevant parts of the
configuration schema and OptimizerInput schemas to confirm configuration
relevance to runtime-scenarios.

##  VII Discussion & Conclusion

Having demonstrated the feasibility of leveraging the architectural framework
for specific runtime configurations, the necessity for explicit information
about dynamic property schema becomes apparent. Integrating architecture with
requirement context facilitates traceable configuration generation, anchored
in system design and requirements. GRAMS’ scalability and reliance on the
architectural framework provide a smoother topology for traversing all
possible variable inputs. By specifying the requirements for the system and
each model in every runtime scenario, we enable traceability and support
safety argumentation. Additionally, generating configurations for the model
optimizer from these requirements facilitates frequent changes and iterative
development, necessitating the continuous upkeep of requirements to ensure
alignment between the optimized models and their specifications. Yet,
questions arise: Is the intermediary format sufficient for tracing relevant
properties and generating final configurations? Are GRAMS checks overly
strict, and what adjustments are needed? Additionally, should GRAMS be
extended to generic systems not reliant on the compositional architectural
framework? These queries prompt further examination.

## References

  * [1] A. Zaki-Ismail, M. Osama, M. Abdelrazek, J. Grundy, and A. Ibrahim, “Requirements formality levels analysis and transformation of formal notations into semi-formal and informal notations,” ser. Proc. of the Int. Conf. on Software Engineering and Knowledge Engineering, SEKE. KSI Research Inc., 2021, pp. 303–308.
  * [2] E. Knauss, G. Liebel, J. Horkoff, R. Wohlrab, R. Kasauli, F. Lange, and P. Gildert, “T-reqs: Tool support for managing requirements in large-scale agile system development,” in _2018 IEEE 26th International Requirements Engineering Conference (RE)_ , 2018, pp. 502–503.
  * [3] H.-M. Heyn, E. Knauss, and P. Pelliccione, “A compositional approach to creating architecture frameworks with an application to distributed ai systems,” _Systems and Software (JSS)_ , vol. 198, 2023.
  * [4] N. Theodórsdóttir, Audur Katarína Ojensa, _Evaluation of Text–Based Requirements Engineering Tools_. Master Theses at Chalmers University of Technology, 2022, available online: https://odr.chalmers.se/items/678b7cf4-9d5f-4ba0-a34e-4eda0965ecd8.
  * [5] O. Gotel, J. Cleland-Huang, J. H. Hayes, A. Zisman, A. Egyed, P. Grünbacher, A. Dekhtyar, G. Antoniol, J. I. Maletic, and P. Mäder, “Traceability fundamentals,” in _Software and Systems Traceability_ , J. Cleland-Huang, O. Gotel, and A. Zisman, Eds. Springer, 2012, pp. 3–22.
  * [6] S. Lauesen, _Software Requirements-Styles and Techniques_. Addison Wesley, 01 2002.

###  -A Instantiating the Extended Architectural Framework

An important component of the extended architectural framework outlined in
Figure 2, is the schema-type. An example of which is shown in Figure 4.

Figure 4: A schema-type element that includes the JSON schema that describes
an OptimizerInput property related to _ethernet latency_.

###  -B Target System Configuration Schema

The JSON schema for the configuration of the target system is given to GRAMS,
not only for validation checking, but also to be included in the output of
GRAMS. Figure 5 illustrates such a sample schema not only for the parts of the
configuration that are to be populated from the requirement model, but for all
properties. GRAMS traverses the schema and identifies and performs checks on
the relevant properties.

Figure 5: Input to GRAMS: A sample JSON schema describing the configuration of
a target system that has two properties: _ethernet latency_ and _model
latency_.

###  -C Intermediary Output Format

The intermediarey output format is YAML and it has two main sections:

  * •

config_schema defines a JSON schema for configuring the target system,
specifying schemas for extracting properties from the architectural framework.
These properties may be nested, and the schema serves as input to GRAMS.
Additionally, it’s included in the output, enabling direct use of the
intermediary output by the target system or by a tool for specific
adjustments.

  * •

optimizer_inputs contains a list of elements specifying OptimizerInput
properties, for each element its file_name, label, placement, treqs_type, and
uid, along with traceability path pointing to requirements going all the way
up to a runtime_scenario. Additionally, it includes a schema defining the JSON
schema for the property and value indicating the specific constraint in
milliseconds. This structure repeats for each requirement.

###  -D GRAMS Output

GRAMS output can be displayed in two formats: PlantUML or YAML, the latter is
intended for consumption of the target system (or other intermediary tooling)
whereas the PlantUML version is for providing an overview. A sample PlantUML
and YAML outputs are hosted as snippets, whereas Figure 6 is the visualization
on the PlantUML code. A recording of GRAMS’ YAML output is available.

Figure 6: Rendering of GRAMS PlantUML output which illustrates two paths
stemming from the same runtime-scenario down to an OptimizerInput property.
The legend shows the types of the two identified properties.

###  -E GRAMS Checks

As outlined in VI, the following recordings demonstrate how GRAMS detects and
reacts to the various checks:

  * •

Meta-model inconsistencies

  * •

Internal schema correctness

  * •

Semantic equivalence between configuration schema and framework’s schema.