text:
\section{Introduction~\label{section:Introduction}}

The fate of quarkonia -- for example the $J/\psi$ and the $\Upsilon$ meson families
-- in a thermal medium, such as the quark-gluon plasma (QGP) created in heavy
ion collisions (HIC), can help us characterize its properties. In particular, quarkonia are
sensitive to the space-time temperature profile and  transport coefficients of the QGP,
see~\cite{Mocsy:2013syh,Datta:2014wga,Andronic:2015wma} for
recent reviews. Experimentally, a key observable that carries such information is the nuclear
modification factor of the yields of quarkonia in nucleus-nucleus ($AA$) collisions, when
compared to their yields
in nucleon-nucleon ($NN$) collisions scaled with the number of binary interactions
\begin{equation}
  R_{AA}= \frac{1}{\langle N_{\rm coll.} \rangle }
  \frac{d \sigma_{AA}^{\rm Quarkonia}/dy dp_T}{d \sigma_{pp}^{\rm Quarkonia}/dy dp_T}
\label{raa}
\end{equation}

In HICs one expects that the short distance formation dynamics of a $Q\bar{Q}$
pair is not affected
since $m_Q\gg T$ where $T$ is the typical temperature of the QGP. To simplify
the calculations, it is often also assumed that the matrix
elements for transition of $Q\bar{Q}$ to quarkonia is not modified, and
for every binary collision the formation of a specific quarkonium state
happens with the same probabilities as in $NN$ collisions. This takes a time scale close
to its inverse of its binding energy.  However, due to the screening of the
color interaction between $Q$ and $\bar{Q}$ in a deconfined
QGP~\cite{Matsui:1986dk}, as well as processes leading to the
dissociation~\cite{Xu:1995eb} of quarkonium states, we expect the yields of
quarkonia to be suppressed in heavy ion collisions ($R_{AA}<1$).

Several methods have been used to estimate screening and dissociation effects
encountered by a quarkonium in a thermal medium. A widely used approach is
based on the intuitive idea that the real part of the ``finite temperature
potential'' between two (nearly) static heavy quarks captures the screened
$Q\bar{Q}$ interaction while the imaginary part of the potential captures
dissociation. For $T=0$, the real part can be quantitatively obtained  by calculating the
Polyakov loop correlation functions ~\cite{Kaczmarek:2005ui,Bazavov:2012bq,Burnier:2015tda}.
For $T>0$, the connection between various correlators
 calculated on the lattice
and the potential between $Q$ and $\bar{Q}$ is
subtle~\cite{Kaczmarek:2005ui,Mocsy:2007jz,Bazavov:2012bq,Burnier:2015tda}.
While the singlet free energy $F_1(r)$ of the $Q\bar{Q}$ state as a function of
the separation $r$, and internal energy $U(r)$ can be extracted from the
lattice data, the connection of either of the two with the $Q\bar{Q}$ potential
is indirect. An important step in clarifying this connection was taken in
Ref.~\cite{Laine:2006ns}, which showed that
the ``potential'' between two heavy quarks is complex, with the imaginary part
connected to thermal processes that can lead to dissociation of quarkonia.
Significant progress has been made in the perturbative calculation of the real and
imaginary part of the $Q\bar{Q}$ potential for $Q\bar{Q}$ at rest in the
medium~\cite{Laine:2006ns,Margotta:2011ta,Brambilla:2013dpa}
or moving slowly in the medium~\cite{Escobedo:2013tca}. Analytic calculations can be performed in
certain regimes by considering various hierarchies of energy scales. For example, one has to assume
that $T$ and $m_Q$ are large enough so that perturbation theory is valid all the way down to
energy scales $\pi T$ and the binding energy $E_b$.

While the real part of the potential for $T\gg \Lambda_{QCD}$ can be obtained
using perturbation theory,
non-perturbative effects are substantial near the crossover temperature
and it is better to estimate this quantity
from lattice calculations by matching the Euclidean correlators measured
on the
lattice~\cite{Mocsy:2007jz,Rothkopf:2009pk,Petreczky:2010tk,Burnier:2015tda}
with those evaluated using the potential. The extraction of the
imaginary part of the potential using this technique is challenging and often
perturbative estimates are used. Extensive
phenomenological study of quarkonium suppression by using this approach has
been performed~\cite{Strickland:2011mw,Song:2014qoa,Krouppa:2015yoa,Hoelck:2016tqf}.
Recently, approaches treating the $Q\bar{Q}$ as an open quantum system were
developed, where a stochastic equation is written for the evolution of the $Q\bar{Q}$
wavefunction~\cite{Akamatsu:2011se,Blaizot:2015hya,Brambilla:2016wgg,Kajimoto:2017rel}.
We also note that the connection between the
heavy quark correlators measured
on the lattice and in a non-perturbative framework using the
in medium T-matrix has been made in Refs.~\cite{Riek:2010fk,Riek:2010py}.
Quarkonium suppression~\cite{Rapp:2009my,Emerick:2011xu,Du:2017qkv}, as well as
low $p_T$ observables like heavy quark flow have been studied using this approach.
Thermal properties of quarkonia have also been investigated in the strong
coupling regime using AdS/CFT techniques, see Ref.~\cite{CasalderreySolana:2011us}.

In this paper we calculate the differential $R_{AA}$ as a function of $p_T$ by
solving rate equations~\cite{Adil:2006ra,Sharma:2009hn,Sharma:2012dy}
describing the change in the yields as a function of time in HICs.
Conceptually, our approach resembles treating the $Q\bar{Q}$ as an open quantum
system. It was, however,  introduced earlier to describe the attenuation of open heavy
flavor~\cite{Adil:2006ra}. Gluon exchanges with the medium lead to the modification of the
$Q\bar{Q}$ state and, hence, a reduction in the overlap with the quarkonium
wavefunction. However we do not connect the imaginary part of the potential
calculated in approaches cited above to the decay rate. This is because we
focus on $R_{AA}$ at high $p_T$. For high $p_T$ partons traversing the medium,
a very successful picture is that the interactions with the medium lead to
transverse momentum broadening. The decay rate in our calculation is related to
the accumulation of relative momenta between $Q$ and $\bar{Q}$.

The form of the rate equations is the same as used in our previous
work~\cite{Sharma:2012dy} and they involve the dissociation time and the formation
time as inputs. The formation time in our formalism is a measure of the time
scale on which the proto-quarkonium $Q\bar{Q}$ state develops interactions
with the medium  and its overlap  with the
quarkonia becomes substantial. We vary it in the neighborhood of ${{\cal{O}}}(1\,
{\rm{fm}})$. The dissociation time scale is computed as in~\cite{Sharma:2012dy}
by calculating the survival probability of the quarkonia.

One conceptual change in our framework, when compared to~\cite{Sharma:2012dy},
is that we use the real part of the lattice motivated thermal
potentials~\cite{Mocsy:2007jz,Petreczky:2010tk,Bazavov:2012bq} to solve for the
quarkonium wavefunction and square the overlap with the thermal wavefunction to
get the survival probability. This is justified if the time scales on which the
medium screens the $Q\bar{Q}$ interaction is smaller than the formation time as
well as the dissociation time. A rough estimate of the Debye screening time is
the inverse of the Debye screening mass ($\mu_D$), which is $\sim gT$. For
$g=1.85$, ($1/\mu_D$) is numerically smaller than $1/2$~fm for most of the
evolution of the medium at the LHC.  On the other hand, the formation is taken
to be $\sim 1$~fm. Finally, this puts a constraint that the minimum
dissociation time consistent within our formalism is roughly $1$~fm. We see
that while there is no substantial hierarchy between $1/\mu_D$ and the
formation and dissociation time scales in our calculation, numerically
$1/\mu_D$ is smaller and, hence, we work in this approximation.

In addition, we make two technical improvements. First, we consider a $2+1$ dimensional
viscous hydrodynamic medium, rather than a simplified $1+1$ dimensional Bjorken
expansion, as a model for the QGP~\cite{Shen:2014vra}.  An important element of
our approach is the use of non-relativistic QCD (NRQCD)~\cite{Bodwin:1994jh} to
obtain the  baseline nucleon-nucleon cross sections for quarkonia  and understand the
$p_T$-dependent feed down. The second technical improvement is a refit of the
long distance matrix elements for $\chi_c$ and $\psi(2S)$ to obtain a
better description of the high
$p_T$ part of their differential spectra. We are further
motivated by the  fact that new experimental results that extend the
measurements of quarkonia to high $p_T$ have appeared. Improved measurements
of the  higher excited states, such as the $\psi(2S)$ and the $\Upsilon(2S)$
and $\Upsilon(3S)$, that are differential in transverse momentum are
particularly useful in constraining the mechanisms of quarkonium suppression.

In this paper we focus on the nuclear modification of the prompt $J/\psi$ and $\Upsilon$
states around mid-rapidity at the LHC and at high transverse
momentum~\cite{Adam:2015rba,Abelev:2013ila,Aad:2010aa,Khachatryan:2016ypw,Khachatryan:2016xxp,Sirunyan:2016znt}. In section~\ref{section:formalism} we describe the theoretical model that is
employed in the calculation of quarkonium production in  heavy ion
collisions. Phenomenological results for the nuclear modification of the
$J/\psi$ and $\Upsilon$ states are presented in section~\ref{section:Results}.
Our conclusions and outlook are given in section~\ref{section:Conclusions}.

\section{Theoretical formalism~\label{section:formalism}}

Quarkonia are bound states of heavy quarks ($Q$) and antiquarks ($\bar{Q}$),
and a concrete picture of the dynamics of the heavy quark pair ($Q{\bar{Q}}$) in
$NN$ collisions is given by non-relativistic quantum chromodynamics
(NRQCD)~\cite{Bodwin:1994jh}. In this effective theory, the initial hard
collision produces a short distance ($\sim1/m_Q$) $Q{\bar{Q}}$ pair in a
color-singlet or an octet state with a specific spin and orbital structure. The
production cross-section for this short-distance state can be calculated using
perturbative QCD. This $Q\bar{Q}$ state evolves into a quarkonium state with
probabilities that are given by long distance matrix elements (LDMEs). For
color-octet states, this evolution process also involves the emission of soft
partons to form a net color-singlet object which we assume occurs on a
time scale which is shorter than ${\cal O}(1\, {\rm fm})$.

This framework has been successfully used to calculate the
unpolarized differential yields of quarkonia versus the transverse momentum
($p_T$) in proton-proton ($pp$) or proton-antiproton ($p\bar{p}$)]
collisions~\cite{Cho:1995vh,Braaten:2000cm,Butenschoen:2011yh,Wang:2012is}.
An accurate description of both the cross sections and  polarization of
quarkonia in hadronic reactions still remains a challenge~\cite{Bodwin:2014gia}.
Recently, it has been suggested~\cite{Baumgart:2014upa,Kang:2017yde,Bain:2017wvk}
that new experimental measurements of quarkonium production inside
jets~\cite{Aaij:2017fak} may help  better constrain the relevant LDMEs.  The
focus of  this paper is, however, different.  Our study concentrates on the
production, propagation, and dissociation of quarkonium states in strongly
interacting matter.  We follow the NRQCD calculation outlined
in~\cite{Sharma:2012dy} and use the LDMEs extracted there to give good
description of the cross sections for bottomonia for $pp$ and
$p\bar{p}$ collisions for $p_T$ in the range of $5$ to $30$~GeV. For charmonia
we improve the $\chi_c$ fitting procedure by allowing the singlet matrix
element as a free parameter. We also refit LDMEs for the $\psi(2S)$ by fitting
the LHC 7~TeV and CDF 1.8~TeV data. Both these changes make the spectra for
$\chi_c$ and $\psi(2S)$ in the $p_T\sim 10-20$~GeV region softer and improve the
description of data without spoiling the agreement at lower
$p_T$.

In order to address quarkonium attenuation in heavy ion reactions, we need to
understand the $J/\psi$ and $\Upsilon$ states' behavior and melting at finite
temperature in the QGP, and the dissociation processes
due to collisional interactions with the quasi-particles of the QCD medium.
 To accomplish this,  a detailed knowledge of the wavefunctions
 at zero and finite temperatures is necessary. We start by solving  the Schr\"{o}dinger
 equation  by  separating  the radial and angular parts of the wavefunction,
 $\psi({\bf{r}})=Y_l^m({\hat{r}})  R_{nl}(r)$.
The reduced equation for the  radial part  can be written as
\begin{eqnarray}
&& \left[-\frac{1}{2\mu_{\rm{red}}}\frac{\partial^2}{\partial r^2}
 +\frac{l(l+1)}{2\mu_{\rm{red}} r^2}
 +V(r)\right]rR_{nl}(r)  \nonumber \\[1ex]
 && = (E_{nl}) rR_{nl}(r)  \; ,
\end{eqnarray}
where $\mu_{\rm{red}}=\frac{m_Q}{2}$ is the reduced mass, and $n$ and $l$
are the principal and orbital quantum numbers, respectively.
$V(r)$ is the potential between the two heavy quarks, which can be estimated
from the lattice~\cite{Bazavov:2012bq}.  The binding energy of the meson is
then $E^b_{nl}=V(\infty) - E_{nl}$.  For $T=0$ we take the form of the potential to
be of the standard Cornell form
\begin{equation}
\begin{split}
V(r) = - \frac{\xi}{r} + \sigma r , \;\; \xi=0.385,  \;  \sigma=0.224\;{\rm{GeV}}^2.
\end{split}
\end{equation}
The form of the potential has been validated in multiple lattice calculations
by evaluating the Polyakov loop correlator as a function of the quark-antiquark
separation. The Cornell long-distance  part is cut off at
$r_{\rm{max}}=1.1$~GeV~\cite{Mocsy:2007jz} to
model string breaking on this length scale. This defines the value of
$V(\infty)\approx 1.2$~GeV. For the mass of the heavy quark we take
$m_c=1.34$ for the charm quark and $m_b=4.5$ for the bottom quark.
For $T>0$ the form of the real part of the  potential we use is the internal
energy found in~\cite{Kaczmarek:2005ui}.
We have checked that the potentials we use are quite close to the internal energies
computed in the more recent work~\cite{Bazavov:2012bq} and the potentials used
in~\cite{Petreczky:2010tk}. Internal energies provide a stronger binding potential
compared to the single
free energies~\cite{Bazavov:2012bq}. The use of the internal energies can be
justified if the dynamics of the quarkonia
are too rapid for the $Q\bar{Q}$ potential to fully thermalize, which might be
especially likely for high $p_T$ quarkonia in the medium.  With this setup,
solutions are obtained for a variety of temperatures and for the  S-wave and
P-wave states.  Even though we are primarily interested in the $J/\psi$,
$\psi(2S)$, and $\Upsilon(nS) $, the   P-wave $\chi_c$ and $\chi_b$ contribute
via feed down.

\begin{table}[h]
\begin{tabular}{cccccc}
$l$ & $n$ & $E^b_{nl}$~(GeV) & $\sqrt{\langle r^2\rangle}$~(GeV$^{-1}$)  & $k^2$~(GeV$^2$) & Meson \\
\hline \hline
0 & 1 & 0.700 & 2.24 & 0.30 & $J/\psi$ \\
0 & 2 & 0.086 & 5.39 & 0.05 & $\psi(2S)$ \\
1 & 1 & 0.268 & 3.50 & 0.20  & $\chi_{c}$ \\[1ex]
\hline
0 & 1 & 1.122  & 1.23 & 0.99 & $\Upsilon(1S)$ \\
0 & 2 & 0.578  & 2.60 & 0.22 & $\Upsilon(2S)$ \\
0 & 3 & 0.214  & 3.89 & 0.10 & $\Upsilon(3S)$ \\
1 & 1 & 0.710  & 2.07 & 0.58 & $\chi_{b}(1P)$ \\
1 & 2 & 0.325  & 3.31 & 0.23 & $\chi_{b}(2P)$ \\
1 & 3 & 0.051  & 5.57 & 0.08 & $\chi_{b}(3P)$\\
\end{tabular}
\caption{Charmonia and bottomonia wavefunctions at zero temperature. $l$ refers to the angular momentum of the
$Q{\bar{Q}}$ state, while $n$ is the radial quantum number. $E^b_{nl}$ is the binding energy, $\sqrt{\langle r^2\rangle}$ is the root mean square  (RMS)
radius of the quarkonium state, and  $k^2$ is the mean square momentum. }\label{table:Tzero}
\end{table}

\begin{table}[h]
\begin{tabular}{cccccc}
$l$ & $n$ & $E^b_{nl}$~(GeV) & $\sqrt{\langle r^2\rangle}$~(GeV$^{-1}$)  & $k^2$~(GeV$^2$) & Meson \\
\hline \hline
0 & 1 & 0.366 & 2.34  & 0.27  & $J/\psi$  \\
0 & 2 & - & -  & -  & $\psi(2S)$ \\
1 & 1 & 0.003 & 8.15  & 0.04  & $\chi_{c}$  \\[1ex]
\hline
0 & 1 & 0.782  & 1.23  &  0.98  & $\Upsilon(1S)$ \\
0 & 2 & 0.244  & 2.72  &  0.20       & $\Upsilon(2S)$ \\
0 & 3 & -      & -     &  -       & $\Upsilon(3S)$ \\
1 & 1 & 0.371  & 2.09  &  0.57        & $\chi_{b}(1P)$ \\
1 & 2 & 0.040  & 4.56  &  0.12        & $\chi_{b}(2P)$ \\
1 & 3 & -      & -     &  -       & $\chi_{b}(3P)$ \\
\end{tabular}
\caption{Charmonia and bottomonia  wavefunctions at finite temperature. We chose a temperature  of 192~MeV  to
illustrate the  disappearance of the weakly bound states and the changes in the quarkonium wavefunctions.}
\label{table:Tfinite}
\end{table}

Pertinent results  for the  quarkonium and bottomonium wavefunctions are
presented in Table~\ref{table:Tzero} and Table~\ref{table:Tfinite} for zero
temperature ($T$) and $T=192$~MeV, respectively. We show the binding energy
$E^b_{nl}$ and  the root mean square  (RMS) size $\sqrt{\langle r^2\rangle}$ of
the quarkonium state. The RMS size changes only slowly with $T$, except near
the  dissociation temperature. We have chosen the finite temperature value
$\sim 190$~MeV to illustrate that several states, i.e.  $\psi(2S)$,
$\Upsilon(3S)$, and  $\chi_{b}(3P)$, cease to exist even with our selected  strong
binding potential. The $\chi_{c}$  and $\chi_{b}(2P)$ states are very near
dissociation.  If we Fourier transform to momentum space, the mean squared
momentum  $k^2$ can also be evaluated and is given in  Tables~\ref{table:Tzero}
and~\ref{table:Tfinite}. We note that while there is a correlation between the
width of the wavefunctions  of quarkonia and their binding energies, this
correlation is highly non-linear. The widths change rapidly only near
dissociation when $E^b_{nl} \rightarrow 0$.

In this work we are interested in large transverse momentum quarkonia and it is convenient to work in light
cone momenta and with light-cone wavefunctions. The relation between the instant form and the light-cone form
of the momentum space wavefunctions for mesons was discussed in detail  in~\cite{Sharma:2009hn,Sharma:2012dy}.
For the case of quarkonia,  the color-singlet contribution  can be understood  as one  matching to the lowest order ($n=2$)
Fock component of the state.  A color-octet initial state  must emit  at least one gluon for a color neutral hadron to be
produced. In either case the heavy meson or proto-quarkonium  state of momentum $\vec{P^+} = (P^+,{\bf P})$ can be
approximated as:
\begin{eqnarray}
 |\vec{P}^+ \rangle
&=& \int \frac{d^2{\bf k}}{(2\pi)^{3}} \frac{dx}{ 2\sqrt{ x(1-x)}}
  \frac{\delta_{c_1c_2}}{\sqrt{3}} \,
\psi(x,{\bf k})  \nonumber \\
&& \times a_Q^{\dagger\;  c_1 }(x\vec{P}^++{\bf k})  b_{\bar{Q}}^{\dagger \;  c_2  }
((1-x)\vec{P}^+-{\bf k})  |0 \rangle  \; ,
\label{Mp1}
\end{eqnarray}
where  $a^\dagger$ ($b^\dagger$) represent an
``effective'' heavy quark (anti-quark) in the $3$ ($\bar{3}$) state, $c_1,c_2$ being the color
indices~\cite{Sharma:2009hn,Sharma:2012dy}.
The   light cone wavefunction $\psi(x, {\bf{k}})$ in Eq.~(\ref{Mp1}), which describe the longitudinal momentum
fraction $x$ ($1-x$) distribution  and the transverse momentum ${\bf k}$ ($-{\bf k}$) distribution of
heavy quarks (antiquarks)
is given by
\begin{eqnarray}
  && \psi(x,{\bf k}) = {\rm Norm} \times \exp\left( - \frac{{\bf k}^2 + m_Q^2 }{2 \Lambda^2(T) x(1-x) }
  \right) \;   ,  \nonumber \\
&& \frac{1}{2 (2\pi)^{3} } \int dx d^2{\bf k}  \;
| \psi(x,{\bf k}) |^2 = 1 \; .
\label{lonorm}
\end{eqnarray}
In Eq.~(\ref{lonorm}) $\Lambda(T)$ is the transverse momentum width of the light-cone wavefunction which
needs to be constrained at 0 and finite temperature to the mean transverse momentum squared
from the solution to the Schr\"{o}dinger equation, which remains invariant under boost.
If we introduce the notation $\Delta {\bf k} = {\bf k}_1-{\bf k}_2 = 2 {\bf
k}$, the equation for  $\Lambda(T)$ reads
\begin{equation}
\frac{1}{2 (2\pi)^{3} } \int dx d^2{\bf k}  \;
{\Delta \bf k}^2  | \psi(x,{\bf k}) |^2 =  4\langle {\bf k^2} \rangle =  \frac{2}{3} \kappa^2  \; .
\label{lonorm2}
\end{equation}
The factor $2/3$ comes from the 2D projection of the mean squared  transverse
momentum $\kappa^2$ from the instant-form wavefunction.

\begin{figure}[!t]
\vspace*{-.1in}
\hspace*{-.15in}\includegraphics[width=3.5in,angle=0]{tempGridX.eps}
\caption{(Color online) Example of the temperature profile of the QGP at a typical time $t=2$~fm
on 0-10\% central Pb+Pb collisions at $\sqrt{s} = 2.76$~TeV the LHC. Glauber initial conditions and $\eta/s=0.08$ are used
in a 2+1D hydrodynamic simulation based on~\cite{Shen:2014vra}. }
\label{hydro}
\end{figure}

The temperature and/or density profiles of the medium, which play an important role in the
dissociation of quarkonia, can be obtained from hydrodynamic simulations of the
QGP~\cite{Schenke:2010nt,Habich:2014jna,Shen:2014vra} in $2.76$~TeV and $5.02$~TeV
Pb+Pb collisions at the LHC. Specifically, we use the iEBE-VISHNU (2+1)-dimensional
event-by-event viscous hydrodynamic package~\cite{Shen:2014vra}. A sample
temperature distribution at time $t=2$~fm when the interplay between the formation and dissociation of quarkonia
is important in setting the final observed $J/\psi$ and $\Upsilon$ multiplicities  is shown in
Fig.~\ref{hydro}.  By comparing the temperature in the different points in the $(x,y)$ plane
perpendicular to the collision axis to the results in Table~\ref{table:Tfinite} one can get a sense of
how the different quarkonium states  will be attenuated in heavy ion collisions relative to proton collisions.

The propagation of a  $Q\bar{Q}$ state in matter is accompanied by collisional
interactions mediated at the partonic level, as long as the momentum exchanges between the
medium quasi-particles and the heavy quarks can resolve the partonic structure of the  meson.  The related
modification of the quarkonium wavefunction in Eq.~(\ref{lonorm})  can lead to the dissociation of
$J/\psi$s and $\Upsilon$s in addition to the thermal
effects. The cumulative one dimensional momentum transfer for a quarkonium state that
starts at transverse position ${\bf x}_0$ and propagates with velocity ${\bm \beta}$,
such that  ${\bf x}(\tau) = {\bf x}_0 + {\bm \beta} (\tau-t_0)$,  reads
\begin{equation}
\chi \mu_D^2 \xi  = \int_{t_0}^t  d\tau
\frac{\mu_D^2(\bf{x}(\tau),\tau) }{\lambda_q({\bf x}(\tau),\tau) } \xi\, .
~\label{eq:xibar}
\end{equation}
Here $\mu_D^2 = g^2T^2(1+N_f/6)$   is the  Debye screening scale and we use 2 active light quark flavors
$N_f=2$.  The scattering inverse length of the quark is ${1}/{\lambda_q}  = \sigma_{qq}\rho_q + \sigma_{qg} \rho_g$,
where $\rho_q$ and $\rho_g$ are the partial densities of light quarks and gluons in the QGP.   We label the
cumulative one dimensional momentum transfer $\chi \mu_D^2 \xi$ in analogy with a uniform static medium
 where the Debye scale is fixed and $\chi = L/\lambda_q$ is the opacity. The elastic scattering cross sections are given by
\begin{equation}
\sigma_{qq} = \frac{1}{18 \pi} \frac{g^4}{\mu_D^2} \, ,  \quad \sigma_{qg} = \frac{1}{8 \pi} \frac{g^4}{\mu_D^2} \, .
\end{equation}
Last but not least, $\xi$ is a parameter related to the heavy quark broadening from multiple scattering in the
QGP~\cite{Gyulassy:2002yv,Adil:2006ra}.  In the limit of strictly soft interactions $\xi=1$ and an enhancement
of $\xi  \sim$few  may arise from the power law tails of the Moliere
multiple scattering.

We initialize the wavefunction $\psi_{i}(\Delta {\bf k}, x)$ of the proto-quarkonium $Q\bar{Q}$  state with a
width $\Lambda_0 \equiv \Lambda(T=0) $.
This is a natural choice since in the absence of a medium it will evolve on the time-scale of
${\cal O}(1{\rm fm})$ or greater into the
observed heavy meson. By propagating in the medium this initial wavefunction  accumulates transverse momentum
broadening  $\chi \mu_D^2 \xi$.  The probability that this  $Q\bar{Q}$ configuration will transition into a
final-state heavy meson
with thermal wavefunction $\psi_{f}(\Delta {\bf k}, x)$ with $\Lambda(T)$ is given by
\begin{eqnarray}
  P_{f\leftarrow i} (\chi\mu_D^2 \xi,T) & = & \left|  \frac{1}{2 (2\pi)^{3} }  \int d^{2}{\bf k} dx \,
\psi_{f}^* (\Delta {\bf k},x)\psi_{i}(\Delta {\bf k}, x) \right|^{2}
\nonumber \\
&& \hspace*{-1.in}= \left| \frac{1}{2 (2\pi)^{3} }  \int dx \; {\rm Norm}_f {\rm Norm}_i \, \pi
  \,  e^{-\frac{ m_{Q}^{2} }{ x(1-x)\Lambda(T)^{2} } }  e^{-\frac{ m_{Q}^{2} }{ x(1-x)\Lambda_0^{2} } }  \right.   \nonumber \\
&&  \hspace*{-.9in}  \times \left .  \,   \frac{ 2 [ x(1-x)\Lambda(T)^{2}]
[\chi\mu_D^{2}\xi+x(1-x)\Lambda_0^{2}] }
{   [ x(1-x)\Lambda(T)^{2}] + [\chi\mu_D^{2}\xi+x(1-x)\Lambda_0^{2}]  }  \; \right|^2 \, . \;\; \quad
\label{sprob}
\end{eqnarray}
In Eq.~(\ref{sprob}) ${\rm Norm}_i$  is the normalization of the initial state,
including the transverse momentum broadening
from collisional interactions,  and  ${\rm Norm}_f$
is the normalization of the final state.
The dissociation rate for the specific quarkonium state  can then be expressed as
\begin{equation}
\frac{1}{t_{\rm diss.}}  =  -  \frac{1}{P_{f\leftarrow i} (\chi\mu_D^2 \xi,T)}
 \frac{d P_{f\leftarrow i} (\chi\mu_D^2 \xi,T) }{dt}  \, .
\end{equation}
It will enter the time evolution of the $J/\psi$s and $\Upsilon$s in the medium.
To visualize the dissociation rate,  we  present  $1/\tau_{\rm diss.}$ as a function of time in Fig.~\ref{rate}
for 0-10\% central Pb+Pb collisions at the LHC at $\sqrt{s}$  per nucleon pair of 2.76 TeV.
We have weighed the dissociation rates with the binary collision density in the plane perpendicular to the
collision axis and averaged over multiple hydrodynamic events. The top panel shows the dissociation rates of
charmonium states and the bottom panel shows the corresponding rates for bottomonium state.
We find that the modification of the ground states, such as $J/\psi$, $\Upsilon(1S)$ is dominated
by dissociation at the early stages of the evolution of the QGP,  whereas excited weakly bound states, such as  are
$\psi(2S)$, $\chi_b(3P)$, can be strongly modified until they escape the medium.

\begin{figure}[!t]
\vspace*{-.0in}
\hspace*{-.15in}\includegraphics[width=3.4in,angle=0]{Rate.eps}
\caption{(Color online)  The dissociation rate $1/\tau_{\rm diss.}$  for charmonium and bottomonium states
in 0-10\% central Pb+Pb collisions at $\sqrt{s}=2.76$~TeV is shown in the top and bottom panels, respectively.
We select the coupling between the quarks and the medium $g=1.85$ and the broadening parameter $\xi=1$.
}
\label{rate}
\end{figure}

Finally, we note that effects that are suppressed at high transverse momentum,
such as recombination~\cite{Thews:2006ia} of unbound $Q$ and $\bar{Q}$ because
of the Boltzmann suppression factor~\cite{Gupta:2014ova}, the Cronin
effect~\cite{Qiu:2003pm}, and power corrections~\cite{Qiu:2004qk} do not play a
role. Cold nuclear matter (CNM) energy loss might affect production cross
sections~\cite{Vitev:2006bi},  but its effects become significant at very high
$p_T$ near the kinematic threshold~\cite{Kang:2015mta}.  Last but not least, it
was also found that leading-twist shadowing effects near mid-rapidity at high
transverse momentum at the LHC are small~\cite{Vogt:2015uba}.  While
experimental results in $p+$Pb collisions at the LHC cannot exclude CNM effects
at the $5-10\%$ level at $p_T>5$~GeV around mid-rapidity, they are also consistent
with no nuclear
modification~\cite{Adam:2015iga,Sirunyan:2017mzd,Aaboud:2017cif}.  For
these reasons, for our study at $|y|<2.4$, $p_T > 5$~GeV we neglect these
effects.

\section{Phenomenological results~\label{section:Results}}

\begin{figure}[!t]
\vspace*{.0in}
\includegraphics[width=3.38in,height=3.5in]{RaaRcpNpartfinX.eps}
\caption{(Color online) Comparison of theoretical results for prompt
  $J/\psi$ suppression for $p_T> 6.5$~GeV to LHC Pb+Pb results at
  $\sqrt{s}=2.76$~TeV. Top panel: $N_{\rm part.}$ dependence versus
  ALICE~\cite{Adam:2015rba} and CMS measurements~\cite{Khachatryan:2016ypw}.
  Bottom panel: $R_{CP}$ as a function of $N_{part}$ versus ATLAS
  measurements~\cite{Aad:2010aa}.}
\label{276Npart}
\end{figure}

In this section we present the phenomenological results of our theoretical
model for quarkonium  dissociation due to thermal wavefunction effects and
collisional breakup.  For every centrality class and hydrodynamic event we
distribute the production of the proto-quarkonium states according to the
binary collision density in the 2D plane perpendicular to the collision axis.
The azimuthal distribution  in the directions of quarkonium propagation is
uniform and we also average over multiple  fluctuating hydrodynamic events that
describe the QGP background.

For each quarkonium state the dynamics of production and propagation through the QCD
medium described above is given by a set of differential equations:
\begin{eqnarray}
\label{rateq11}
\frac{d}{dt} \left( \frac{d\sigma^{Q\bar{Q}}(t;p_T)}{dp_T} \right)
&=& - \frac{1}{t_{\rm form.} } \frac{ d\sigma^{Q\bar{Q}}(p_T)}{dp_T}  \, , \\
\frac{d}{dt} \left(  \frac{d\sigma^{\rm meson}(t;p_T)}{dp_T} \right)
\label{rateq12}
&=&  \frac{1}{t_{\rm form.}} \frac{ d\sigma^{Q\bar{Q}}(t;p_T)}{dp_T}  \nonumber \\
&&- \frac{1}{t_{\rm diss.}}  \frac{d\sigma^{\rm meson}(t;p_T)}{dp_T}   \, , \\
\frac{d}{dt} \left(  \frac{d \sigma^{\rm diss.} (t;p_T)}{dp_T} \right)
&=&  \frac{1}{t_{\rm diss.}}   \frac{d\sigma^{\rm meson}(t;p_T)}{dp_T}  \, .
\label{rateq13}
\end{eqnarray}
Here we denote by ${d\sigma^{Q\bar{Q}}(t;p_T)}/{dp_T}$ the cross section to produce the
proto-quarkonium states that evolve into an interacting with the medium meson on the time scale of $t_{\rm form.}$.
This time scale for the heavy quarks to interact with the QGP is taken in the relativistic limit.
The initial condition at $t\approx 0$ includes the short distance perturbative  $Q\bar{Q}$ production cross sections and the
long-distance matrix elements for the particular quarkonium state.   ${d\sigma^{\rm meson}(t;p_T)}{dp_T}$ is the
cross section for that quarkonium state as a function of time $t$. Finally,  ${d \sigma^{\rm diss.} (t;p_T)}/{dp_T}$ is the cross section of the dissociated
$Q\bar{Q}$ pairs that will not produce quarkonia. In the absence of a medium $t_{\rm diss.} \rightarrow \infty$ and the proto-quarkonium
states evolve into the corresponding meson with unit probability.  The system of equations Eqs.~(\ref{rateq12})-(\ref{rateq13})
has to be solved for each of the quarkonium states  listed in Table~\ref{table:Tzero}.
Feed-down is then performed using the NRQCD cross sections and branching ratios.

\begin{figure}[!t]
\vspace*{.0in}
\includegraphics[width=3.38in,height=2.2in]{LHC2.76TeVUpsfinX.eps}
\caption{(Color online) Comparison of theoretical results for the
  $\Upsilon(nS)$ $R_{AA}$ in 2.76~TeV minimum bias Pb+Pb collisions versus $p_T$
to CMS experimental measurements~\cite{Khachatryan:2016xxp}.
}
\label{U276}
\end{figure}

\begin{figure}[!t]
\vspace*{.0in}
\includegraphics[width=3.38in,height=2.2in]{RaaUNpartfinX.eps}
\caption{(Color online) Suppression of $\Upsilon(1S)$, $\Upsilon(2S)$ in Pb+Pb collisions
at  2.76~TeV is shown versus the number of participants for  $p_T=5$~GeV.
Superimposed are  CMS experimental  data~\cite{Khachatryan:2016xxp} on bottomonia suppression
versus centrality.
}
\label{U276cent}
\end{figure}

We start by first discussing results at the lower energy of $\sqrt{s}=2.76$~TeV.
In Fig.~\ref{276Npart} we present the centrality dependence of prompt $J/\psi$s
in Pb+Pb collisions. The bands reflect the combined uncertainty
of the interaction onset time $t_{\rm form.}$ and the collisional dissociation of
the quarkonium states. In the evaluation of the latter we keep the coupling between
the heavy quarks and the medium fixed at $g=1.85$~\cite{Sharma:2012dy} but vary the
broadening parameter $\xi$. The upper edge of the uncertainty band corresponds to
$t_{\rm form.} = 1.5$~fm, $\xi =1$. The lower edge of the uncertainty band corresponds
to $t_{\rm form.} = 1$~fm, $\xi =2$. The upper panel of  Fig.~\ref{276Npart} shows
comparison to the ALICE~\cite{Adam:2015rba} and CMS~\cite{Khachatryan:2016ypw} prompt
$J/\psi$ measurements. We find improved description in the most central $N_{\rm part.}$
bins relative to the case when thermal screening effects were not
included~\cite{Sharma:2012dy}. Around $N_{\rm part.} =100$ the data lies on the
upper edge of the theoretical error band. The bottom panel of Fig.~\ref{276Npart} shows the $J/\psi$
$R_{CP}$, where the 40\%-80\% peripheral collisions are used as a baseline.
The  ATLAS collaboration measured inclusive $J/\psi$~\cite{Aad:2010aa}. However, in the
$p_T < 10$~GeV interval which dominates the cross section, the non-prompt
$B\rightarrow J/\psi$ contribution is limited to 20-30\%~\cite{Sharma:2012dy} and will
not noticeably  affect the theoretical results.

\begin{figure}[!t]
\vspace*{.0in}
\includegraphics[width=3.38in,height=3.5in]{LHC5.1TeVcentfinX.eps}
\caption{(Color online) Theoretical model predictions for the $R_{AA}$ of the ground
  and excited $J/\psi$ (top panel) and $\Upsilon$ (bottom panel) states
  in 0-10\% central  Pb+Pb collisions  at  $\sqrt{S}=5.02$~TeV at the LHC. The
  coupling between the heavy quarks and the medium $g=1.85$ and the bands correspond to
$t_{\rm form.}=1.5$~fm, $\xi=1$ - $t_{\rm form.}=1$~fm, $\xi=2$. }
\label{502cen}
\end{figure}

Recently, experimental results for the differential  suppression of the $\Upsilon(nS)$
family have appeared at high $p_T$~\cite{Khachatryan:2016xxp}. Theoretical calculations
for the $\Upsilon(1S)$ (red band) and $\Upsilon(2S)$ (blue band) in minimum bias
$\sqrt{s}=2.76$~TeV Pb+Pb reactions are shown in Fig.~\ref{U276}. We have evaluated
the cross sections for quarkonia in 10 centrality classes (labeled $i$) and
\begin{eqnarray}
  R_{AA}^{\rm min.\ bias}(p_T)& =&
  \frac{\sum_i R_{AA}(\langle b_i \rangle) W_i}{\sum_i W_i} \quad {\rm where}
  \nonumber \\
  W_i &=&\int_{b_{i\, \min}}^{b_{i\, \max}}  N_{\rm coll.}(b)\,\pi \, b \,db  \, .
\label{cent}
\end{eqnarray}
The experimental data is described well, including its magnitude and $p_T$ dependence.
We note that collisional dissociation mostly
affects the ground  $\Upsilon$ state, while thermal wavefunction effects dominate
the attenuation pattern of the excited $\Upsilon$ states. The CMS collaboration also
put an upper limit on the  $\Upsilon(3S)$ cross section in Pb+Pb reactions, corresponding to an
upper limit on its $R_{AA}$~\cite{Khachatryan:2016xxp}. Our calculated $\Upsilon(3S)$
cross section is consistent with this limit.  While the theoretical approach presented in this Letter  is  applicable at large
transverse momenta, we observe in Fig.~\ref{U276} that the nuclear modification
factor is approximately constant. This allows us to compare in Fig.~\ref{U276cent}  the centrality dependence
in the lowest $p_T=5$~GeV bin,
not very different form the mean $p_T$ of bottomonia at the LHC, to the experimentally measured
$\Upsilon(1S)$ and $\Upsilon(2S)$  $R_{AA}$ dependence  on the number of participants~\cite{Khachatryan:2016xxp}.
Very good agreement between data and theory is observed.

\begin{figure}[!t]
\vspace*{.0in}
\includegraphics[width=3.38in,height=3.5in]{LHC5.1TeVminbiasfinX.eps}
\caption{(Color online) Same as in Fig.~\ref{502cen}, but for minimum bias collisions.}
\label{502mb}
\end{figure}

In Figs.~\ref{502cen} and~\ref{502mb}, we present theoretical predictions for
the $R_{AA}$ of various quarkonium species as a function of $p_T$ in Pb+Pb
collisions at $\sqrt{s} = 5.02$~TeV.  The top and bottom panels display results
for charmonium and bottomonium states, respectively. We find a clear separation
in suppression based on how tightly bound the quarkonium state is. We also find
a flat or slightly increasing $R_{AA}$ with $p_T$. By comparing
Fig.~\ref{502cen} to Fig.~\ref{502mb}, we observe that the attenuation of
quarkonia in minimum bias collisions is only slightly smaller than in the most
central collisions. The reason for that behavior is that minimum bias
collisions are strongly dominated by the first $3$ most central classes, as given
by the weights $W_i$ in Eq.~(\ref{cent}).

\begin{figure}[!t]
\includegraphics[width=3.38in,height=2.2in]{LHC5.02TeVPStoJfinX.eps}
\caption{(Color online) Theoretical model predictions for the
  double $(\psi(2S)_{AA}/(\psi(2S)_{pp})/ (J/\psi_{AA}/(J/\psi_{pp})$ ratio in  minimum bias
  Pb+Pb collisions
at  $\sqrt{S}=2.76$~TeV at the LHC. Data is from CMS~\cite{Sirunyan:2016znt}. }
\label{J2toJ1}
\end{figure}

Very recently, at $\sqrt{s} =5.02$~TeV,  measurements of relative suppression
ratios of excited to ground quarkonium states have
appeared~\cite{Sirunyan:2016znt,Sirunyan:2017lzi}. The data for $\psi(2S) / J/\psi $
is publicly available and shown in Fig.~\ref{J2toJ1}.  Theoretically, the double suppression
ratio can be obtained from the results in Fig.~\ref{502mb} and is compatible
with the experimental data within the statistical and systematic error bars. \\

\section{Conclusions~\label{section:Conclusions}}

In summary, we presented theoretical results for the $p_T$-differential
suppression of charmonia and bottomonia in Pb+Pb collisions at the LHC. The
dynamics of $Q\bar{Q}$ pairs, which evolve into the observed quarkonium states,
 is governed in HICs by the formation and dissociation time scales. A key
element of our formalism that addresses this dynamics is that the formation time
of proto-quarkonia is $\sim 1$~fm. We assume that this time scale is long enough that
the $Q\bar{Q}$ interact via a color screened thermal
potential~\cite{Mocsy:2007jz,Kaczmarek:2005ui}. Therefore we employ the
wavefunctions obtained by solving the Schr\"{o}dinger equation for $Q\bar{Q}$
interacting via a screened potential to calculate the dissociation time scale,
using the theoretical setup described in~\cite{Sharma:2012dy}. The technical
advances that are further incorporated in the calculation are better
constraints on the NRQCD matrix elements that are relevant for the production
of high-$p_T$ $\psi(2S)$ and $\chi_c$ states, and their feed-down  to $J/\psi$,
and 2+1 dimensional event-by-event hydrodynamic modeling of the QGP
background~\cite{Shen:2014vra}.

We explored the phenomenological implications of this theoretical model for
quarkonium production and propagation in the QGP background created in heavy
ion collisions, first around mid-rapidity in Pb+Pb collisions at the LHC. We
found good separation in the magnitude of the suppression between the ground
and excited charmonium and bottomonium states, compatible with recent
experimental measurements. Our results indicate that effects of thermal
screening of the  $Q\bar{Q}$ attractive potential fully dominate the
attenuation of $\psi(2S)$, $\Upsilon(2S)$ and $\Upsilon(3S)$. On the other
hand, $J/\psi$ and $\Upsilon(1S)$ are also sensitive to the dissociation
processes due to collisional interactions. The approximately constant or
slightly decreasing $R_{AA}$ with $p_T$  predicted by this model arises from
the early ${\cal O}(1{\rm \ fm})$ formation of the interacting quarkonium
state.  The uncertainty of the phenomenological results was estimated by
varying the formation time and the strength of the collisional broadening of
the $Q\bar{Q}$ pair.  We found that the charmonium suppression measurements are
better described by the upper edge of the $R_{AA}$ uncertainty band, whereas
bottomonium suppression measurements are better described by its lower edge.
While an illustrative subset of results was presented here, detailed
predictions are available that will allow to test this model versus upcoming
experimental measurements of quarkonium suppression at $\sqrt{s}=5.02$~TeV.

In the future, we plan to address data at finite rapidity with the same parameters
and test the model further. This will likely require inclusion of CNM effects since
$p-$Pb data at finite rapidity seems to show non-trivial nuclear modification patterns.
It will also be interesting and instructive to investigate non-prompt $J/\psi$
production, which probes the complementary but different physics of in-medium
modification of heavy-quark parton showers~\cite{Kang:2016ofv,Huang:2013vaa}. \\

{\bf Acknowledgments:} This research is supported by  the US Department of Energy, Office of
Science, under Contract No.~DE-AC52-06NA25396 and in part by the DOE
Early Career Program. We are grateful to C. Shen for help with the iEBE-VISHNU code.
RS thanks J. P. Blaizot, S. Datta, R. Gavai, S. Gupta, and
A. Tiwari for rewarding discussions on the topic.


text:
\subsubsection*{\bibliography_phd2}}
\renewcommand{\bibsection}{\subsubsection*{\bibname}}

\bibliographystyle{apalike}

\usepackage[utf8]{inputenc}
\usepackage[T1]{fontenc}
\usepackage{hyperref}
\usepackage{url}
\usepackage{booktabs}
\usepackage{amsfonts}
\usepackage{nicefrac}
\usepackage{microtype}

\usepackage[page]{appendix}

\RequirePackage{epsfig}
\RequirePackage{amssymb}
\RequirePackage{natbib}
\RequirePackage{graphicx}

\usepackage{amsthm}

\newcommand{\BlackBox}{\rule{1.5ex}{1.5ex}}
\newtheorem{example}{Example}
\newtheorem{theorem}{Theorem}
\newtheorem{lemma}[theorem]{Lemma}
\newtheorem{proposition}[theorem]{Proposition}
\newtheorem{remark}[theorem]{Remark}
\newtheorem{corollary}[theorem]{Corollary}
\newtheorem{definition}[theorem]{Definition}
\newtheorem{conjecture}[theorem]{Conjecture}
\newtheorem{axiom}[theorem]{Axiom}
\newtheorem{experiment}{Experiment}

\usepackage{epstopdf}
\usepackage{graphicx}
\usepackage{subcaption}

\usepackage{amsmath}

\usepackage{bbm}
\usepackage{bm}

\usepackage{algorithm}
\usepackage{algpseudocode}
\usepackage{mathtools}

\newcommand{\begin{fig}}{\begin{fig}}
	\newcommand{\end{fig}}{\end{fig}}
\newcommand{\begin{tab}}{\begin{tab}}
	\newcommand{\end{tab}}{\end{tab}}
\newcommand{\begin{array}}{\begin{array}}
	\newcommand{\end{array}}{\end{array}}
\newcommand{\begin{equation}}{\begin{equation}}
\newcommand{\end{equation}}{\end{equation}}
\newcommand{\begin{eqnarray*}}{\begin{eqnarray*}}
	\newcommand{\end{eqnarray*}\noindent}{\end{eqnarray*}\noindent}
\newcommand{\begin{eqnarray}}{\begin{eqnarray}}
\newcommand{\end{eqnarray}\noindent}{\end{eqnarray}\noindent}
\newcommand{\begin{displaymath}}{\begin{displaymath}}
\newcommand{\end{displaymath}\noindent}{\end{displaymath}\noindent}

\newcommand{\mathbb{N}}{\mathbb{N}}
\newcommand{\mathbb{Z}}{\mathbb{Z}}
\newcommand{\mathbb{R}}{\mathbb{R}}
\newcommand{\mathbb{C}}{\mathbb{C}}
\newcommand{\mathbb{D}}{\mathbb{D}}
\newcommand{\mathbb{Q}}{\mathbb{Q}}

\newcommand{\stackrel{P}{\rightarrow}}{\stackrel{P}{\rightarrow}}
\newcommand{\stackrel{\mathscr{D}}{\rightarrow}}{\stackrel{\mathscr{D}}{\rightarrow}}
\newcommand{\stackrel{\rm a.s.}{\rightarrow}}{\stackrel{\rm a.s.}{\rightarrow}}
\newcommand{\stackrel{\mathscr{D}}{=}}{\stackrel{\mathscr{D}}{=}}
\newcommand{{x\to\infty}}{{x\to\infty}}
\newcommand{{n\to\infty}}{{n\to\infty}}
\newcommand{{t\to\infty}}{{t\to\infty}}
\newcommand{\quad\hfill\mbox{$\Box$}}{\quad\hfill\mbox{$\Box$}}
\newcommand{\displaystyle\lim}{\displaystyle\lim}
\newcommand{{\cal S}}{{\cal S}}
\newcommand{{\cal L}}{{\cal L}}
\newcommand{{\cal R}}{{\cal R}}

\newcommand{{\alpha}}{{\alpha}}
\newcommand{{\lambda}}{{\lambda}}
\newcommand{{\Lambda}}{{\Lambda}}
\newcommand{{\epsilon}}{{\epsilon}}
\newcommand{{\varepsilon}}{{\varepsilon}}
\newcommand{{\gamma}}{{\gamma}}
\newcommand{{\Gamma}}{{\Gamma}}
\newcommand{{\varphi}}{{\varphi}}
\newcommand{{\sigma}}{{\sigma}}

\newcommand{{\mathrm{card}}}{{\mathrm{card}}}
\newcommand{{\mathrm{var}}}{{\mathrm{var}}}
\newcommand{{\mathrm{Var}}}{{\mathrm{Var}}}
\newcommand{{\mathrm{cov}}}{{\mathrm{cov}}}
\newcommand{{\mathrm{Cov}}}{{\mathrm{Cov}}}
\newcommand{{\mathrm{corr}}}{{\mathrm{corr}}}

\newcommand{\overline}{\overline}
\newcommand{\widehat}{\widehat}
\newcommand{\widetilde}{\widetilde}

\DeclareMathOperator{\e}{e}
\newcommand{\mathop{\mathbb E}}{\mathop{\mathbb E}}

\begin{document}

\twocolumn[

\aistatstitle{Scalable Bayesian Learning for State Space Models using Variational Inference with SMC Samplers}

\aistatsauthor{ Marcel Hirt \And Petros Dellaportas  }
\aistatsaddress{University College of London, UK\\ \And
	University College of London, UK,\\
	Athens University of Economics and Business, Greece	\\
	and The Alan Turing Institute, UK  }

]

\begin{abstract}
	We present a scalable approach to performing approximate fully Bayesian inference in generic state space models. The proposed method is an alternative to particle MCMC that provides fully Bayesian inference of both the dynamic latent states and the static parameters of the model. We build up on recent advances in computational statistics that combine variational methods with sequential Monte Carlo sampling and we demonstrate the advantages of performing  full Bayesian inference over the static parameters rather than just performing variational EM approximations. We illustrate how our approach enables scalable inference in multivariate stochastic volatility models and self-exciting point process models that allow for flexible dynamics in the latent intensity function.
\end{abstract}

\section{Introduction}

We deal with generic state-space models (SSM) which may be nonlinear and non-Gaussian.  Inference for this important and popular family of statistical models presents tremendous challenges that has prohibited their widespread applicability.  The key difficulty is that inference on the latent process of the model depends crucially on unknown static parameters that need to be also estimated.  While MCMC samplers are unsatisfactory because they both  fail to produce high dimensional, efficiently mixing Markov chains and because they are inappropriate for on-line inference, sequential Monte Carlo (SMC) methods \citep{kantas2015particle} provide the tools to construct successful viable implementation strategies.  In particular, particle MCMC \citep{andrieu2010particle} utilises SMC to build generic efficient MCMC algorithms that provide inferences for both static parameters and latent paths.  We provide a scalable alternative to these methods via an approximation that combines SMC  and variational inference.

We introduce a new variational distribution that unlike recent strand of literature \citep{maddison2017filtering,naesseth2017variational,le2017auto} performs variational inference also on the static parameters of the SSM. This is essential for various reasons. First, when there is dependency between static and dynamic parameters posterior inference may be inaccurate if the joint posterior density is approximated by conditioning on fixed values of static parameters. Second, inferring the static parameter is often the primary problem of interest: for example, for biochemical networks and models involving Lotka Voltera equations, we are not interested in the population of the species per se, but we want to infer some chemical rate constants (such as reaction rates or predation/growth rates), which are parameters of the transition density;  in neuroscience, Bayesian decoding of neural spike trains is often made via a state-space representation of point processes in which inference for static parameters is of great importance.
Finally, for complex dynamic systems it is often advisable to improve model compression or interpretability by encouraging sparsity and such operations may require inference for the posterior densities of the static parameters.

Sampling from the new variational distribution involves running a SMC algorithm which yields an unbiased estimate of the likelihood for a fixed static parameter value. Importantly, we show that the SMC algorithm constructs a computational graph that allows for optimisation of the variational bound using stochastic gradient descent. We provide some empirical evidence that variational inference on static parameters can give better predictive performance, either out-of sample in the linear Gaussian state space  model or in-sample for predictive distributions in a multivariate stochastic volatility model.  We also illustrate our method by modelling fairly general intensity functions in a multivariate Hawkes process model.

\section{Background}

Let us begin by introducing the standard inference problem in a generic SSM, followed by a review of the SMC approach to sample from a sequence of distributions arising in such probabilistic structures. SSMs are characterized by a latent Markov state process $\{X_n\}_{n\geq 0}$ on $\mathbb{R}^{d_x}$ and an observable process $\{Y_n\}_{n\geq 0}$ on $\mathbb{R}^{d_y}$. We follow the standard convention of using capital letters for random variables and the corresponding lower case letter to denote their values. The dynamics of the latent states is determined, conditional on a static parameter vector $\theta \in \Theta$, by a transition probability density
\[X_{n}| (\theta, X_{n-1}=x_{n-1},Y_{n-1}=y_{n-1}) \sim f_{\theta}(\cdot | x_{n-1},y_{n-1}), \]
along with an initial density $X_0  \sim f_{\theta}(\cdot)$. The observations are assumed to be conditionally iid given the states with density given by
\[ Y_n | (\theta, X_{0:n}=x_{0:n},Y_{0:n-1}=y_{0:n-1})\sim g_{\theta}(\cdot | x_n), \]
for any $n\geq 0$ with the generic notation $x_{0:n}=(x_0,...,x_n)$.

We consider a Bayesian framework and assume $\theta$ has a prior density $p(\theta)$. Consequently, for observed data $y_{0:M}$, we perform inference using the posterior density
\begin{equation} \pi(\theta,x_{0:M})\coloneqq p(\theta,x_{0:M}| y_{0:M})\propto p(\theta) p_{\theta}(x_{0:M},y_{0:M}),\label{posterior}\end{equation}
where the joint density of the latent states
and observations given a fixed static parameter value $\theta$ writes as
\begin{align} &p_{\theta}(x_{0:M},y_{0:M})=\gamma_{\theta}(x_{0:M}) \nonumber\\
\coloneqq &f_{\theta}(x_0) \prod_{n=1}^M f_{\theta}(x_n| x_{n-1},y_{n-1})\prod_{n=0}^Mg_{\theta}(y_n| x_n).\label{joint_density}
\end{align}
The posterior density $p(\theta,x_{0:M}| y_{0:M})$ is in general intractable, as is
\begin{equation} p_{\theta}(x_{0:M}| y_{0:M})=\frac{\gamma_{\theta}(x_{0:M})}{p_{\theta}(y_{0:M})},\label{filtering_density_pf} \end{equation}
where $p_{\theta}(y_{0:M})=\int p_{\theta}(x_{0:M},y_{0:M})dx_{0:M}$.
However, an SMC algorithm can be used to approximate $(\ref{filtering_density_pf})$. A brief review of how this sampling algorithm proceeds is as follows and further details can be found in \citet{doucet2000sequential,doucet2009tutorial}.\\
SMC methods approximate $p_{\theta}(x_{0:n}| y_{0:n})$ using a set of $K$ weighted random samples $X_{0:n}^{1:K}=(X_{0:n}^1,...,X_{0:n}^K)$, also called particles, having positive weights $W_n=W_n^{1:K}$, so that $p_{\theta}(x_{0:n}| y_{0:n})\approx \hat{p}_{\theta}(x_{0:n}| y_{0:n})=\sum_{k=1}^K W_{n}^k \delta_{X_{0:n}^k}(x_{0:{n}})$. Here, $\delta$ denotes the Dirac delta function. To do so, one starts at $n=0$ by sampling $X_0^k$ from an importance density $M_0^{\phi}(\cdot | y_0)$, parametrized with $\phi$, where $\phi$ can depend on the static parameters $\theta$. For any $n\geq 1$, we first resample an ancestor variable $A_{n-1}^k$ that represents the 'parent' of particle $X_{0:n}^k$ according to $A_{n-1}^k\sim r(\cdot | W_{n-1})$, where $r$ is a categorical distribution on $\{1,...,K\}$ with probabilities $W_{n-1}$. We then set $W_{n-1}=\frac{1}{K}$ and proceed by extending the path of each particle by sampling from a transition kernel $X_n^k\sim M_n^{\phi}(\cdot | y_n,X_{0:n-1}^{A_{n-1}^k})$. This yields an updated latent path $X_{0:n}^k=(X_{0:n-1}^{A_{n-1}^k},X_n^k)$ for which we compute the incremental importance weight

\[\alpha_n(X_{0:n}^k)=\frac{\gamma_{\theta}(X_{0:n}^k)}{\gamma_{\theta}(X_{0:n-1}^k)M_n^{\phi}(X_n^k | y_n, X_{0:n-1}^{A_{n-1}^k})}.\]
We set $w_n(X_{0:n}^k)=W_{n-1}^k \alpha_n(X_{0:n}^k)$ as well as
$W_n^k=\frac{w_n(X_{0:n}^k)}{\sum_l w_n(X_{0:n}^l)}$ and define
\[\hat{Z}_n^{\theta,\phi}\coloneqq \prod_{m=0}^n\sum_{k=1}^Kw_m(X_{0:m}^k),\]
which is an unbiased and strongly consistent estimator of $p_{\theta}(y_{0:n})$, see \citet{del1996non}. A pseudo-code (Algorithm \ref{SMCAlgo}) for this standard SMC sampler can be found in Appendix \ref{AppendixSmcSampler}. It is possible to perform the resampling step only if some condition on $W_{n-1}$ is satisfied, see Algorithm \ref{SMCAlgo}. For simplicity, we assume that the particles are resampled at every step. The density of all variables generated by this SMC sampler for a fixed static parameter value $\theta$ is given by
\begin{align*}&q_{\phi}(x_{0:M}^{1:K},a_{0:M-1}^{1:K},l| \theta)
=w_M^l\prod_{k=1}^K M_0^{\phi}(x_0^k| y_0) \\& \cdot \prod_{n=1}^M \prod_{k=1}^K r(a_{n-1}^k | w_{n-1})M_n^{\phi}({x_n^k} | y_n, x_{0:n-1}^{a_{n-1}^k}),
\end{align*}
where $l$ is a final particle index drawn from a categorical distribution with weights $W_M$.
Since $\hat{Z}_n^{\theta, \phi}$ is unbiased, we have
\begin{equation} {\mathop{\mathbb E}}_{q_{\phi}(x_{0:M}^{1:K},a_{0:M-1}^{1:K},l| \theta)}\left[ \hat{Z}_M^{\theta,\phi}\right]=p_{\theta}(y_{0:M}).\label{SMC_unbiased}\end{equation}

\section{Variational bounds for state space models using SMC samplers}

Variational inference \citep{jordan1999introduction,wainwright2008graphical,blei2017variational} allows Bayesian inference to scale to large data sets \citep{hoffman2013stochastic} and is applicable to a wide range of models \citep{ranganath2014black,kucukelbir2017automatic}. It generally postulates a family of approximating distributions with variational parameters that minimize some divergence, most commonly the KL divergence, between the approximating distribution and the posterior. The quality of the approximation hinges on the expressiveness of the variational family.

Let $q_{\psi}(\theta)$ be a distribution on $\Theta$ with variational parameters $\psi$. We aim to approximate the posterior density $p(\theta,x_{0:M}| y_{0:M})$ in $(\ref{posterior})$  with a variational distribution that results as an appropriate marginal of auxiliary variables arising from an SMC sampler of the form
\begin{equation} q_{\psi,\phi}(\theta,x_{0:M}^{1:K},a_{0:M-1}^{1:K},l)\coloneqq q_{\psi}(\theta)q_{\phi}(x_{0:M}^{1:K},a_{0:M-1}^{1:K},l| \theta),\label{variational_distribution}\end{equation}
defined precisely below. Note that sampling from the extended variational distribution (\ref{variational_distribution}) just means sampling $\theta\sim q_{\psi}(\theta)$ and then running a particle filter using the sampled value $\theta$ as the static parameter.

We introduce the proposed variational bound first as a lower bound on $\log p(y_{0:M})-\text{KL}(q_{\psi}(\theta)||p(\theta | y_{0:M}))$. We then show that optimizing the proposed bound means minimizing the KL-divergence between the extended variational distribution $(\ref{variational_distribution})$ and an extended target density that resembles closely the density targeted in particle MCMC methods

We can write $p(\theta | y_{0:M})=p(\theta)p_{\theta}(y_{0:M})/p(y_{0:M})$. Hence, using the fact that the likelihood estimator is unbiased $(\ref{SMC_unbiased})$ and due to Jensen's inequality,
\begin{align*}
	&-\text{KL}(q_{\psi}(\theta)||p(\theta | y_{0:M})) +\log p(y_{0:M})\\
	=&{\mathop{\mathbb E}}_{q_{\psi}(\theta)}\left[\log p_{\theta}(y_{0:M})+\log p(\theta) - \log q_{\psi}(\theta)\right] \\
	=&{\mathop{\mathbb E}}_{q_{\psi}(\theta)}\left[\log {\mathop{\mathbb E}}_{q_{\phi}(x_{0:M}^{1:K},a_{0:M-1}^{1:K},l| \theta)}\left[\hat{Z}_M^{\theta,\phi}\right] +\log \frac{p(\theta)}{q_{\psi}(\theta)}\right]\\
	\geq & {\mathop{\mathbb E}}_{q_{\psi}(\theta)}\left[ {\mathop{\mathbb E}}_{q_{\phi}(x_{0:M}^{1:K},a_{0:M-1}^{1:K},l| \theta)}\left[\log \hat{Z}_M^{\theta,\phi}\right] +\log \frac{p(\theta)}{ q_{\psi}(\theta)}\right]\\
	=:& \mathcal{L}(\psi,\phi).
\end{align*}
In particular, $\mathcal{L}(\psi,\phi)$ is a lower bound on $p(y_{0:M})- \text{KL}(q_{\psi}(\theta)||p(\theta | y_{0:M}))$.

\begin{remark}[\bfseries Inference for multiple independent time series]
	Instead of considering one latent process $\{X\}$ and observable process $\{Y\}$, we can also consider $S$ independent latent processes $\{X^s\}_{s=1,...,S}$ with corresponding observable processes $\{Y^s\}_{s=1,...,S}$ described by the same static parameter $\theta$. We obtain a lower bound on $p(y_{0:M})-\text{KL}(q_{\psi}(\theta)||p(\theta | y^1_{0:M},..., y^S_{0:M}))$ given by
	\begin{align*}
	&{\mathop{\mathbb E}}_{q_{\psi}(\theta)}\Bigg[ {\mathop{\mathbb E}}_{\prod_sq_{\phi}(x_{0:M}^{s,1:K},a_{0:M-1}^{s,1:K},l^s| \theta)}\left[\sum_{s=1}^S\log \hat{Z}_{M,s}^{\theta,\phi}\right] \\&+\log p(\theta) - \log q_{\psi}(\theta)\Bigg],
	\end{align*}
	where $\hat{Z}^{\theta,\phi}_{M,s}$ is the estimator of $p_{\theta}(y_{0:M}^s)$. Note that we can obtain an unbiased estimate of this bound by sampling an element $s\in\{1,...,S\}$ and using $S\cdot \log \hat{Z}_{M,s}^{\theta,\phi}$ as an estimate of $\sum_{s'=1}^S\log \hat{Z}_{M,s'}^{\theta,\phi}$, thereby allowing our method to scale to a large number of independent time series. For ease of exposition, we formulate our results for a single time series only.
\end{remark}
Next, we show that the variational bound can be represented as the difference between the log-evidence and the KL divergence between the variational distribution and an extended target density. More concretely, following \citet{andrieu2010particle}, we consider a target density on the extended space $ \Theta \times \mathcal{X}$, $\mathcal{X} \coloneqq (\mathbb{R}^{d_x})^{(M+1)K}\times \{1,...,K\}^{MK+1}$,
\begin{align*}
&\tilde{\pi}(\theta,x_{0:M}^{1:K},a_{0:M-1}^{1:K},l)  \coloneqq \frac{\pi(\theta,x^l_{0:M})}{K^{M+1}} \\ &\cdot  \frac{q_{\phi}(x_{0:M}^{1:K},a_{0:M-1}^{1:K},l| \theta)}{M_0^{\phi}(x_0^{b_0^l}| y_0) \prod_{n=1}^M r(b_{n-1}^l| w_{n-1}) M_n^{\phi}(x_n^{b_n^l}| y_n, x_{0:n-1}^{b_{n-1}^l})}.
\end{align*}
Here, we have defined $b_M^l=l$ and $b_n^l=a_n^{b_{n+1}^l}$ for $n=M-1,...,1$, i.e. $b_n^l$ is the index that the ancestor of particle $X_{0:M}^l$ at generation $n$ had.
It follows, using $r(b_n^l| w_{n-1})=w_{n-1}^{b_{n-1}^l}$, that the ratio between the extended target density and the variational distribution is given by
\begin{align}
&\frac{\tilde{\pi}(\theta,x_{0:M}^{1:K},a_{0:M-1}^{1:K},l)}{ q_{\phi,\psi}(\theta,x_{0:M}^{1:K},a_{0:M-1}^{1:K},l)}\nonumber \\
=&\frac{
	K^{-(M+1)}p(\theta) p_{\theta}(x_{0:M}^l,y_{0:M})/p(y_{0:M})}
{q_{\psi}(\theta)W_M^lM_0^{\phi}(x_0^{b_0^l}| y_0)\prod_{n=1}^M W_{n-1}^{b_{n-1}^l}M_{n}^{\phi}(x_n^{b_n^l} | y_n,x_{0:n-1}^{b_{n-1}^l})}.\label{ratio_extended_target}
\end{align}

\begin{proposition}[\bfseries KL divergence in extended space]\label{proposition_extended_space}
	It holds that
	\[ \mathcal{L}(\psi,\phi)=-\text{KL}(q_{\psi,\phi}||\tilde{\pi})+\log p(y_{0:M}).\]
\end{proposition}
The proof can be found in Appendix \ref{AppendixExtendedSpaceRepresentation}. Recall that we have introduced $\mathcal{L}(\psi,\phi)$ so that its maximisation pushes the variational approximation of the static parameter $\theta$ closer to its true posterior as measured by the KL divergence. The above proposition shows that this objective also minimizes the KL divergence between densities on an extended space that includes multiple latent paths. To elucidate further the relation between the variational distribution of a single latent path and its posterior, we need to introduce a further distribution. Consider the density under $\tilde{\pi}$ of the variables generated by a SMC algorithm conditional on a fixed latent path $(x_{0:M}^{l},b_{0:M-1}^{l})$. This is known as a conditional SMC algorithm \citep{andrieu2010particle}, with distribution given by

\begin{align*}
&\tilde{\pi}_{\text{CSMC}}(x_{0:M}^{\neg b_{0:M}^l},a_{0:M-1}^{\neg b_{0:M-1}^l} | \theta,x_{0:M}^l,b_{0:M}^l)\\
=&\frac{ q_{\phi}(x_{0:M}^{1:K},a_{0:M-1}^{1:K},l| \theta )}{W_M^lM_0^{\phi}(X_0^{b_0^l}|y_0)\prod\limits_{n=1}^M r(b_{n-1}^l |W_{n-1})M_{n}^{\phi}(x_n^{b_n^l} | y_n, x_{0:n-1}^{b_{n-1}^l})},
\end{align*}

where $\neg b_{0:M}^l$ are the indices of all particles that are not equal to $b_{0:M}^l$. We obtain the following corollary proved in Appendix \ref{AppendixMarginalKL}.
\begin{corollary}[\bfseries Marginal KL divergence and marginal ELBO]\label{corollary_marginal_bound}
	The KL divergence in the extended space is an upper bound on the KL divergence between the marginal variational approximation and the posterior, with the gap between bounds being
	\begin{align*}&\text{KL}\left(q_{\psi,\phi}(\theta,x_{0:M}^{1:K},a_{0:M-1}^{1:K},l)||\tilde{\pi}(\theta,x_{0:M}^{1:K},a_{0:M-1}^{1:K},l)\right)\\
	&-\text{KL}\left(q_{\psi,\phi}(\theta,x_{0:M})||\pi(\theta,x_{0:M}^l)\right)\\
		=&{\mathop{\mathbb E}}_{q_{\psi,\phi}(\theta,x_{0:M}^l,b_{0:M}^l)}\bigg[ \\
		& \quad \text{KL}(q_{\phi}(x_{0:M}^{\neg b_{0:M}^l},a_{0:M-1}^{\neg b_{0:M-1}^l})| \theta,x_{0:M}^l,b_{0:M}^l) || \\
		& \quad \quad \tilde{\pi}_{\text{CSMC}}(x_{0:M}^{\neg b_{0:M}^l},a_{0:M-1}^{\neg b_{0:M-1}^l} | \theta,x_{0:M}^l,b_{0:M}^l))\bigg].
	\end{align*}
	Particularly, $\mathcal{L}$ is a lower bound compared to the standard ELBO using the marginal $q_{\psi,\phi}(\theta,x_{0:M})$ with $x^l_{0:M}=x_{0:M}$ as the variational distribution:
	\[\mathcal{L}(\psi,\phi)\leq -\text{KL}\left(q_{\psi,\phi}(\theta,x_{0:M})||\pi(\theta,x_{0:M})\right)+\log p(y_{0:M}).\]

\end{corollary}

The proposed surrogate objective resembles variational bounds with auxiliary variables \citep{salimans2015markov,maaloe2016auxiliary,ranganath2016hierarchical} where the gap between the two bounds is expressed by the KL-divergence between the variational approximation of the auxiliary variable given the latent variable of interest and a so-called reverse model. Here, this reverse model is specified by the conditional SMC algorithm.
The above corollary implies that the variational bound is looser than the standard ELBO with the auxiliary variables integrated out. This marginal variational distribution cannot in general be evaluated analytically. However, we can obtain unbiased estimates of it by computing the log-likelihood estimate under a conditional SMC algorithm, resembling a particle Gibbs update. This constitutes an extension of Proposition 1 in \citet{naesseth2017variational}. We present a proof in Appendix \ref{AppendixMarginalDist}.

\begin{proposition}[\bfseries Marginal variational distribution]\label{proposition_marginal_distribution}
	We have
	\begin{align*}&q_{\psi,\phi}(\theta, x^l_{0:M},b_{0:M}^l)=q_{\psi}(\theta)\gamma_{\theta}(x^l_{0:M}) \\
	& \cdot {\mathop{\mathbb E}}_{\tilde{\pi}_{\text{CSMC}}(x_{0:M}^{\neg b_{0:M}^l},a_{0:M-1}^{\neg b_{0:M-1}^l} | \theta, x_{0:M}^l)}\left[\left({\hat{Z}^{\theta,\phi}_M}\right)^{-1}\right]
	\end{align*}
	and there exists $c(\theta,\phi)<\infty$ so that
	\begin{align*}&\text{KL}(q_{\psi,\phi}(\theta,x_{0:M})||p(\theta,x_{0:M}|y_{0:M})\\\leq &{\mathop{\mathbb E}}_{q_{\psi}(\theta)}\left[\frac{c(\theta,\phi)}{K}\right]+\text{KL}(q_{\psi}(\theta)||p(\theta|y_{0:M})).
	\end{align*}
\end{proposition}
The last inequality in Proposition \ref{proposition_marginal_distribution} is a straightforward extension of an analogous result in the EM setting \citep{naesseth2017variational}. It implies that, for fixed variational parameters $\psi$ and $\phi$, the approximation becomes more accurate for increasing $K$. Sampling from this distribution can be seen as an extension of visualizing the expected importance weighted approximation in Importance Weighted Auto-Encoders \citep{cremer2017reinterpreting}. Since this distribution can be high-dimensional, the preceding proposition gives an alternative to kernel-density estimation.

Lastly, from a different angle, the variational objective can be seen as a sequential variational-autoencoding (VAE) bound. Indeed, as a consequence of Proposition \ref{proposition_extended_space} and equation (\ref{ratio_extended_target}), we obtain immediately the following result. We elaborate on it further in the next section.

\begin{corollary}[\bfseries Sequential VAE representation]\label{proposition_VAE_representation} The variational bound can be written as
	\begin{align*}\mathcal{L}(\psi,\phi)= {\mathop{\mathbb E}}_{q_{\psi}(\theta)}\Bigg[ &
		{\mathop{\mathbb E}}_{q_{\phi}(x_{0:M}^{1:K},a_{0:M-1}^{1:K},l| \theta)}\bigg[ \\
		 & \quad \sum_{n=0}^M \log g_{\theta}(y_n| x_n^{b_n^l})  -\log  W_n^{b_n^l} \\
		& \quad +\log \frac{ f_{\theta}(x_n^{b_n^l} | x_{n-1}^{b_{n-1}^l},y_{n-1})}{M^{\phi}(x_n^{b_n^l}| y_n, x_{0:n}^{b_{n-1}^l})}\bigg] \Bigg]\\
		&
		-(M+1)\log K-\text{KL}(q_{\psi}(\theta)||p(\theta)).
	\end{align*}

\end{corollary}

\section{Related Work}
The representation in Corollary \ref{proposition_VAE_representation} allows us to contrast the variational bound to previously considered sequential VAE frameworks  \citep{chung2015recurrent,archer2015black, fraccaro2016sequential,krishnan2017structured,goyal2017z}. The introduced bound contains the cross-entropy between the proposal distribution and the likelihood common to sequential VAE bounds. However, this reconstruction error is only evaluated for surviving particles. Similarly, while a sequential VAE framework includes a KL-divergence between the proposal distribution and the prior transition probability, the log-ratio of these two densities is only evaluated for a surviving path. Most work using sequential VAEs have considered observation and state transition models parametrised by neural networks, and given the high-dimensionality of the static parameters, have confined their analysis to variational EM inferences. This is also the case for the approaches in \citet{maddison2017filtering,naesseth2017variational,le2017auto}, to which this work is most closely related. They have demonstrated that resampling increases the variational bound compared to a sequential IWAE \citep{burda2015importance} approach. \citet{rainforth2018tighter} demonstrated that increasing the number of particles leads to a worse signal to noise ratio of the gradient estimate of the proposal parameters in an IWAE setting. \citet{le2017auto} suggested to use fewer particles without resampling for calculating the proposal gradient. A possible approach left for future work would be to consider a different resampling threshold for the proposal gradients. Finally, the objective in this work differs from adaptive SMC approaches optimizing the reverse KL-divergence (or $\chi^2$-divergence) between the posterior and the proposal, cf. \citet{cornebise2008adaptive,gu2015neural}.

\section{Optimization of the variational bound}

The gradient of the variational bound is given by
\begin{align} &\nabla_{\psi,\phi} \mathcal{L}(\psi,\phi) \label{grad_bound} \\
=&\nabla_{\psi,\phi} \left({\mathop{\mathbb E}}_{q_{\psi}(\theta)}\left[ {\mathop{\mathbb E}}_{q_{\phi}(x_{0:M}^{1:K},a_{0:M-1}^{1:K},l| \theta)}\left[\log \hat{Z}_M^{\theta,\phi}\right] \right] \right) \nonumber \\
& \phantom{=}+ \nabla_{\psi} \left( {\mathop{\mathbb E}}_{q_{\psi}(\theta)}\left[\log \frac{p(\theta)}{q_{\psi}(\theta)}\right] \right).\nonumber
\end{align}

We focus on the gradient of the first expectation and note that the gradient of the second expectation can be estimated by standard (black-box) approaches in variational inference, depending of course on the chosen variational approximation. If for instance the variational distribution over the static parameters is continuously reparametrisable, one can use standard low-variance reparametrised gradients \citep{kingma2014auto,rezende2014stochastic,titsias2014doubly}. This is the gradient estimator that we use in our experiments in combination with mean-field variational families. We assume that the proposals $X_n^k\sim M_n^{\phi}(\cdot | y_n, x_{0:n-1}^{a_{n-1}^k})$ are reparametrisable, i.e. there exists a differentiable deterministic function $h_{\phi}$ such that $X_n^k=h_{\phi}(X_{0:n-1}^{A_{n-1}^k},{\epsilon}_n^k)$, with ${\epsilon}_n^k\sim p(\cdot)$ continuous and independent of $\phi$. Similarly, we assume that the variational distribution of the static parameters is reparametrisable, i.e. there exists a differentiable deterministic function $h_{\psi}$ such that $\theta=h_{\psi}(\eta)$, with $\eta \sim p(\cdot)$ continuous and independent of $\psi$.
We abbreviate $\bm{{\epsilon}}={\epsilon}_{0:M}^{1:K}$, $\bm{x}=x_{0:M}^{1:K}$ and $\bm{a}=a_{0:M-1}^{1:K}$. Using the product rule, observe that the first gradient in $(\ref{grad_bound})$ is

\begin{align*}
	& \nabla_{\psi,\phi} \int p(\eta) p(\bm{{\epsilon}})q_{\phi}( \bm{a} | \theta,\bm{x})\\
	 &\phantom{=}\cdot \log \hat{Z}_M^{\theta,\phi}  d(\eta, \bm{a},\bm{{\epsilon}}) \bigg|_{\theta=h_{\psi}(\eta),\bm{x}=h_{\phi}(\bm{{\epsilon}})}\\
	&= \int p(\eta) p(\bm{{\epsilon}}) \nabla_{\psi,\phi}q_{\phi}( \bm{a} | \theta,\bm{x})\\
	&\phantom{=}\cdot \log \hat{Z}_M^{\theta,\phi}  d(\eta, \bm{a},\bm{{\epsilon}}) \bigg|_{\theta=h_{\psi}(\eta),\bm{x}=h_{\phi}(\bm{{\epsilon}})}\\
	&={\mathop{\mathbb E}}_{p(\eta)p(\bm{{\epsilon}})q_{\phi}(\bm{a}| h_{\psi}(\eta),h_{\phi}(\bm{{\epsilon}}))} \Bigg[ \nabla_{\psi,\phi}\log\hat{Z}_M^{h_{\psi}(\eta),\phi} \\
	&\quad +\nabla_{\psi,\phi} \log q_{\phi}(\bm{a}| h_{\psi}(\eta),h_{\phi}(\bm{{\epsilon}})) \log \hat{Z}_M^{h_{\psi}(\eta),\phi} \Bigg].
\end{align*}
Analogously to \citet{maddison2017filtering,le2017auto,naesseth2017variational} in a variational EM framework, we have also ignored  the second summand in the gradient due to its high variance in our experiments. We take Monte Carlo samples of the expectation above and optimize the bound using Adam \citep{kingma2014adam}. It is also possible to use natural gradients \citep{amari1998natural}, see Appendix \ref{AppendixNaturalGradients}.

\section{Experiments}

\subsection{Linear Gaussian state space models}

\paragraph{Regularisation in a high-dimensional model.}
We illustrate potential benefits of a fully Bayesian approach in a standard linear Gaussian state space model
\begin{align}
	f_{\theta}(x_{n}| x_{n-1})&=\mathcal{N}(A x_{n-1},\Sigma_x),\label{lgss_state_transition}\\
	g_{\theta}(y_n| x_n)&=\mathcal{N}(Bx_n,\Sigma_y)\label{lgss_observation},
\end{align}
with initial state distribution $X_0\sim \mathcal{N}(A^0,\Sigma_x^0)$
and parameters $A,\Sigma_x, \Sigma_x^0\in  \mathbb{R}^{d_x\times d_x}$, $A^0\in \mathbb{R}^{d_x}$, $B \in \mathbb{R}^{d_x \times d_y}$, and $C,\Sigma_y \in \mathbb{R}^{d_x \times d_y}$.
\citet{naesseth2017variational} have shown in a linear Gaussian model that learning the proposal yields a higher variational lower bound compared to proposing from the prior and the variational bound is close to the true log-marginal likelihood for both sparse and dense emission matrices $B$. However, an EM approach might easily over-fit, unless one employs some regularisation, such as stopping early if the variational bound decreases on some test set. We demonstrate this effect by re-examining one of the experiments in \citet{naesseth2017variational}, setting $(d_x,d_y)=(10,3)$, $M=10$ and assume that $\Sigma_x$, $\Sigma_x^0$ and $\Sigma_y$ are all identity matrices. Furthermore, $A^0=0$ and $(A_{ij})=\alpha^{|i-j|+1}$ with $\alpha=0.42$, and $B$ has randomly generated elements with $B_{ij}\sim \mathcal{N}(0,1)$.
We assume that the proposal density is
\[M_{n+1}^{\phi}(x_{n+1}| x_{n},y_{n+1})=\mathcal{N}(x_{n+1}|  A_{\phi}x_n+B_{\phi}y_{n+1},\Sigma_{\phi}),\]
and $M_{0}^{\phi}(x_{0}| y_{0})=\mathcal{N}(x_{0}|  A^0_{\phi}+B_{\phi}y_{0},\Sigma^0_{\phi})$, with $\Sigma_{\phi}$ and $\Sigma^0_{\phi}$ diagonal matrices.
We perform both a variational EM approach and a fully Bayesian approach over the static parameters using $K=4$ particles. In the latter case, we place Normal priors $B_{ij}\sim \mathcal{N}(0,10)$ and $A_{ij}\sim \mathcal{N}(0,1)$. Furthermore, we suppose that a priori $ \Sigma_y$ is diagonal with variances drawn independently from an Inverse Gamma distribution with shape and scale parameters of $0.01$ each. A mean-field approximation for the static parameters is assumed. We suppose that the variational distribution over each element of  $A$ and $B$ is a normal distribution and the approximation over the diagonal elements of $\Sigma_y$ is log-normal. For identifiability reasons, we assume that $\Sigma_x$, $\Sigma_x^0$ and $A^0$ are known. We compare the EM and VB approach in terms of log-likelihoods on out-of-sample data assuming training and testing on $10$ iid sequences. Figure \ref{LGSS_LLH} shows that in contrast to the VB approach, the EM approach attains a higher log-likelihood on the training data with a lower log-likelihood on the test set as the training progresses.\\

\begin{figure}[htb]
	\begin{subfigure}{.23\textwidth}
		\centering
		\includegraphics[width=1.0\linewidth]{"Figures/aistats_llh_training__10"}
		\caption{Log-likelihood on training data.}
		\label{LGSS_LLH_10_train}
	\end{subfigure}
	\begin{subfigure}{.23\textwidth}
		\centering
		\includegraphics[width=1.0\linewidth]{"Figures/aistats_llh_testing__10"}
		\caption{Log-likelihood on testing data.}
		\label{LGSS_LLH_10_test}
	\end{subfigure}
	\caption{Log-likelihood for linear Gaussian state space models. Log-likelihood values are computed using Kalman filtering. The static parameters used in the VB case are the mean of the variational distribution (VB mean) or the samples from the variational distribution (VB samples) as they are drawn during training.}
	\label{LGSS_LLH}
\end{figure}

\paragraph{Approximation bias in a low-dimensional model.}
Variational approximations for the latent path can yield biased estimates of the static parameters, see \citet{turner2011two}. We illustrate that this bias decreases for increasing $K$ in a two-dimensional linear Gaussian model, both in an EM and VB setting.
We therefore consider inference in a linear Gaussian state space model (\ref{lgss_state_transition}-\ref{lgss_observation}) with two-dimensional latent states and one-dimensional observations. The state transition matrix is assumed to be determined by the autoregressive parameter $\lambda$ with $A=\begin{pmatrix}
\lambda & 0 \\ 0 & \lambda
\end{pmatrix}.$ We consider inference over $\lambda$ as the static parameter and fix $B=(1,1)$ with $\Sigma_x$ and $\Sigma_y$ being identity matrices. We simulate $30$ realisations of length $M=100$ each using $\lambda=0.9$. Inference is performed with different initialisations and learning rates over the simulated datasets. It has been documented in such a linear Gaussian model, see \citet{turner2011two}, that Gaussian variational approximations of the latent path that factorise over the state components underestimate $\lambda$. We observe the same effect in Figure \ref{LGSS_Boxplot} when using just $K=1$ particle. However, increasing the number of particles used during inference reduces this bias.
Furthermore, we find that point estimates of the static parameters show some variation over different simulations, while an approximate Bayesian approach can be argued to better account for this uncertainty.
The variational distributions for $\theta$ for each of the simulations using $K=100$ particles is shown in Figure \ref{LGSS_AR_Variaional}, confirming that they all put significant mass on the ground truth. Let us remark that these experiments also complement those in \citet{le2017auto}, where it is illustrated that increasing $K$ improves learning point estimates of the static parameters in a Gaussian model with a one-dimensional latent state. Indeed, as shown next, the marginal variational distribution allows not just for dependencies in the latent states across time, but also across different state dimensions, even if they are independent under the proposal.

\begin{figure}[htb]
	\begin{subfigure}{.23\textwidth}
		\centering
		\includegraphics[width=1.0\linewidth]{"Figures/aistats_lgss_boxplot"}
		\caption{Point estimate of the autoregressive parameter $\lambda$ in the EM case or the variational mean in the VB case over $30$ simulations for $K\in\{1,10,100\}$ particles.}
		\label{LGSS_Boxplot}
	\end{subfigure}
	\begin{subfigure}{.23\textwidth}
		\centering
		\includegraphics[width=1.0\linewidth]{"Figures/aistats_lgss_vb_density_100"}
		\caption{Variational distribution of the autoregressive parameter $\lambda$ using $K=100$ particles for each of the $30$ simulations.}
		\label{LGSS_AR_Variaional}
	\end{subfigure}

	\caption{Inference on the autoregressive parameter $\lambda$ over $30$ simulations of length $M=100$. Ground truth values are $\lambda=0.9$.}
	\label{LGSS_AR_param}
\end{figure}

\paragraph{Marginal variational distribution in a low-dimensional model.}
In an additional experiment, we evaluate if the variational approximation from Proposition $\ref{proposition_marginal_distribution}$ of the latent path matches the distribution of its true posterior. We consider the above state space model over $2$ time steps as in \citet{turner2011two}. Note that for given static parameters, the posterior is Gaussian. Indeed, for $\bm{x}=(x_{0}^{(0)},x_{0}^{(1)},x_{1}^{(0)},x_{1}^{(1)})$, where $x_{n}^{(i)}$ denotes dimension $i$ of $x_n$, we have $p(\bm{x}|y_{0:1},\lambda)=\mathcal{N}(\mu_{x|y},\Sigma_{x|y})$ with
\[\Sigma_{x|y}^{-1}=
\begin{pmatrix}
2&1&-\lambda& 0 \\
1&2&0&-\lambda\\
-\lambda& 0& 2&1\\
0&-\lambda &1 &2 &\\
\end{pmatrix},
\mu_{x|y}=\Sigma_{x|y}\begin{pmatrix}
y_0\\y_0\\y_1\\y_1\end{pmatrix}, \]
assuming $X_0\sim \mathcal{N}(0,\frac{1}{1-\lambda^2} I)$ is drawn from its stationary distribution. We visualise the posterior distribution along with the marginal variational distribution
\begin{align*}
&q_{\phi}(x^l_{0:M}| \theta)\\=&\gamma_{\theta}(x^l_{0:M}) {\mathop{\mathbb E}}_{\tilde{\pi}_{\text{CSMC}}(x_{0:M}^{\neg b_{0:M}^l},a_{0:M-1}^{\neg b_{0:M-1}^l} | \theta, x_{0:M}^l)}\left[\left({\hat{Z}^{\theta,\phi}_M}\right)^{-1}\right]
\end{align*}
in Figure \ref{csmc_vis_text} using $K=100$ particles and $50$ samples for the expectation. We find that the approximation mirrors the true posterior. In particular, it accounts for explaining-away between different dimensions of the latent state, although we have used isotropic proposals.

\begin{figure}[htb]
	\begin{subfigure}{.23\textwidth}
		\centering
		\includegraphics[width=1.0\linewidth]{"Figures/csmc__lgss1_second_timestep09_100"}
		\caption{Joint distribution of the latent states at the second time step. Top: variational approximation, bottom: true posterior.}
		\label{csmct1_text}
	\end{subfigure}
	\begin{subfigure}{.23\textwidth}
		\centering
		\includegraphics[width=1.0\linewidth]{"Figures/csmc__lgss1_first_dim09_100"}
		\caption{Joint distribution of the first state component at the first and second time step. Top: variational approximation, bottom: true posterior.}
		\label{csmcc1_text}
	\end{subfigure}

	\caption{Two-dimensional contour plots of the distribution of the latent path over two time steps and two state components. Function arguments are set to the ground truth state values as simulated if they are not shown.}
	\label{csmc_vis_text}
\end{figure}

\subsection{Stochastic volatility models}
To show that our method allows inference of latent states and static parameters of higher dimensions, we consider a multivariate stochastic volatility model,
\begin{align*} f_{\theta}(x_n| x_{n-1})&=\mathcal{N}(\mu+\text{diag}(a)(x_{n-1}-\mu),\Sigma_x),\\
	g_{\theta}(y_n| x_n)&= \mathcal{N}(0,\exp(\text{diag}(x_n)),
\end{align*}
where $X_0\sim \mathcal{N}(\mu,\Sigma_x^0)$ with $x_n,y_n,\mu,a\in \mathbb{R}^D$, and covariance matrix $\Sigma_x\in \mathbb{R}^{D\times D}$, $\theta=(\mu,a,\Sigma_x,\Sigma_x^0)$. This model has been considered in \citet{guarniero2017iterated} using particle MCMC methods under the restriction that $\Sigma_x$ is band-diagonal to reduce the number of parameters. It is also more general than that entertained in \citet{naesseth2017variational} with $\Sigma_x$ assumed diagonal, see also \citet{chib2009multivariate} for a review on stochastic volatility models. We consider a fully Bayesian treatment as in \citet{guarniero2017iterated}, applied to the same data set of $90$ monthly returns (9/2008 to 2/2016) of $20$ exchange rates with respect to the US dollar as reported by the Federal Reserve System. The specification of the prior and variational forms of the static parameters are explained in Appendix \ref{appendix_stoch_vol_approximations}. We consider proposals of the form
\[M_{\phi}(x_{n+1}| y_{n+1},x_n)= \mathcal{N}(\mu+\text{diag}(a)(x_{n}-\mu),\Sigma^{\phi}),\]
where $\Sigma^{\phi}$ is diagonal and using $K=50$ particles. Densities of the variational approximation that correspond to the GBP exchange rate can be found in Appendix \ref{appendix_stoch_vol_approximations}, Figure \ref{GBP_variational}, which are largely similar to those obtained in \citep{guarniero2017iterated}. Furthermore, we approximate the one- and two-step predictive distributions
\[p(y_{m+p}| y_{0:m}) \approx \frac{1}{S}\sum_{s=1}^S \sum_{k=1}^K W_m^{k,s} \delta_{X^{k,s}_{m+p}}p_{\theta_s}(y_{m+p}| X^{k,s}_{m+p})  \]
for $p\in \{1,2\}$,where $\theta_1,...,\theta_S \sim q_{\psi}(\theta)$, $\sum_{k=1}^K W_{m}^{k,s}\delta_{X_{m}^k}$ is the approximation of $p_{\theta_s}(x_m| y_{0:m})$ by the particle filter and $X_n^s\sim p_{\theta_s}(x_n^{k,s}| X_{n-1}^{s},Y^{k,s}_{n-1})$ with $Y^s_{n} \sim p_{\theta_s}(y^{k,s}_{n}| X^{k,s}_{n})$ for $n=m+1,...,m+p$ simulated from the generative model. The predictive distributions are evaluated using a log scoring rule \citep{gneiting2007strictly,geweke2010comparing} to arrive at the predictive log-likelihoods in Table \ref{pred_llh_table}. The full variational approach attains higher predictive log-likelihoods.

\begin{table}
	\caption{Average $p$-step predictive log-likelihoods per observation for the stochastic volatility model with different number of particles $K$ and number of samples $S$ from the variational distribution. In the EM case, we run $S$ particle filters with the same optimal static values. Mean estimates with standard deviation in parentheses based on 100 replicates.}
	\label{pred_llh_table}
	\centering
	\begin{tabular}{lll}
		\toprule
		&\multicolumn{2}{c}{$(S,K)=(4,50)$}\\
		Method & $p=1$     & $p=2$  \\
		\midrule
		EM & 9.697 (0.008)  & 9.716 (0.008)       \\
		VB     & 9.707 (0.011) & 9.728 (0.015)       \\
		\midrule[0.25ex]
		&\multicolumn{2}{c}{$(S,K)=(20,100)$}\\
		Method & $p=1$     & $p=2$  \\
		\midrule
		EM  & 9.690 (0.003)  & 9.713 (0.003)      \\
		VB  & 9.701 (0.004)  & 9.727 (0.005)    \\
		\bottomrule
	\end{tabular}
\end{table}

\subsection{Non-linear stochastic Hawkes processes}

There has been an increasing interest in modelling asynchronous sequential data using point processes in various domains, including social networks \citep{linderman2014discovering,wang2017linking}, finance \citep{bacry2015hawkes}, and electronic health \citep{lian2015multitask}. Recent work \citep{du2016recurrent,mei2017neural,xiao2017joint,xiao2017wasserstein} have advocated the use of neural networks in a black-box treatment of point process dynamics.

We illustrate that our approach allows scalable probabilistic inference for continuous-time event data $\{T_n,C_n\}_{n>0}$, $T_n<T_{n+1}$, where $T_n$ is the time when the $n$-th event occurs and $C_n\in \{1,...,D\}$ is an additional discrete mark associated with the event. We consider describing such a realisation as a $D$-variate point process with intensities $\lambda_{t}=h_{\theta}(\mu+\sum_{b=1}^B \Xi^b_t)$, driven by $B$ continuous time processes \[\Xi_t^{b}=\sum_{n\geq 1}  \beta_b A_n^b\e^{-\beta_b (t-T_n)}1_{\left[0,t\right)}(T_n),\quad t>0,\]
and a non-negative monotone function $h_{\theta}$. Moreover, $\mu,A_n \in\mathbb{R}^D$ and $\beta^b>0$. Importantly, we allow $A_n^b$ to depend on $C_n$, and the $i$-th component of $A_n^b$ describes by how much the $n$-th event excites, if $(A_n^b)^i>0$, or inhibits, if $(A_n^b)^i<0$, subsequent events of type $i$. It is possible to view the dynamics as a discrete-time SSM; the essential idea being that $\Xi^b$ is piecewise-deterministic between events, see Appendix \ref{AppendixHawkesLiterature} for details along with related work on Hawkes point processes \citep{hawkes1971point}.
Let us define the discrete-time latent process $X_{n+1}=(Z_{n},A_{n})$ with $Z_n=\Xi_{T_n}$, $A_n=\text{vec}(A_n^1,...,A_n^B)$. Standard theory about point processes, see \citet{daleyintroduction1}, implies that the observation density is given by $g_{\theta}(t_n,c_n| z_{n-1})=\lambda_{t_n}^{c_n}\exp\left(-\sum_{i=1}^D\int_{t_{n-1}}^{t_n} \lambda_s^{i} ds\right)$,
where our model specification yields $\lambda_s$ as a deterministic function between $T_{n-1}$ and $T_n$ given $Z_{n-1}$. Similar to \citet{mei2017neural}, we set $h_{\theta}(y)=\nu \text{ softplus}(y/\nu)=\nu \log(1+\exp(y/\nu))$ as a scaled softplus function with $\nu$ a static parameter. Next, we specify the dynamics of $A_n$. We take the arguable most simple model, assuming $f_{\theta}(a_n| a_{n-1},z_{n-1},c_n)=\mathcal{N}(\sum_d \alpha_d \delta_{c_nd},\sum_d \sigma^2_d  \delta_{c_nd})$ with $\alpha_1,...,\alpha_D\in \mathbb{R}^{BD}$ and $\sigma^2_1,...,\sigma^2_D$ positive diagonal matrices, while remarking in passing that our approach allows readily for extensions that could include temporal dynamics between successive intensity jumps or intensity jumps instantaneously correlated across different marks and time scales. Due to the piecewise deterministic decay of $\Xi$, note that $Z_n^b| Z_{n-1}^b, A_n^b=\e^{-\beta^b(T_{n}-T_{n-1})}Z_{n-1}^b + \beta^b A_{n}^b$, so the state transition of the process $X$ is fully specified.\\
We apply our model to 20 days of high-frequency financial data for the BUND futures contract. The data is available as part of the tick library \citep{bacry2017tick} with $4$ event types: (i) mid-price up moves, (ii) mid-price down moves, (iii) buyer-initiated trades leaving the mid-price unchanged and (iv) seller-initiated trades not changing the mid. We train our model on 15 days and evaluate how well it predicts the type of the next event on out of sample data from the remaining 5 days.\\
Table \ref{hawkes_prediction_table} reports better predictive performance of the proposed model in comparison with two benchmark models. First, a linear Hawkes process model estimated using maximum likelihood. Second, to illustrate that improved predictions might not be just explained due to inhibitory effects, we also compare against a non-linear Hawkes model. The latter can be seen, and has been implemented, as a limiting case of our generative model letting $\sigma_d^2\to 0$, with inference thus performed using stochastic gradient descent of the negative log-likelihood.
Predictions are Monte Carlo samples of the next event realisation from the generative model. Further details including assumptions on the variational distributions and the predictive performance using a smaller training set are given in Appendix \ref{AppendixHawkesInference}.

\begin{table}
	\caption{Prediction metric for different Hawkes process models on the test set of around 206k events. The stochastic Hawkes model is trained with $20$ particles and uses $K\in\{20,80\}$ particles during testing.}
	\label{hawkes_prediction_table}
	\centering
	\begin{tabular}{ll}
	\toprule
	Method & Error rate      \\
		&		next mark \\
	\midrule
	Linear Hawkes & 43.3 \%         \\
	Non-linear Hawkes & 40.9 \%         \\
	Non-linear stochastic Hawkes ($K=20$) &40.0\%        \\
	Non-linear stochastic Hawkes ($K=80$)& 39.3\%         \\
	\bottomrule
	\end{tabular}
\end{table}

\section{Conclusion}

This paper has explored an inference approach that merges the scalability of variational methods with SMC sampling.
We would like to emphasize that our approach is completely complementary to many recent advances in variational inference that can be used to parametrize $q_{\psi}(\theta)$. For instance, one can consider more expressive variational families \citep{rezende2015variational,kingma2016improved,salimans2015markov,maaloe2016auxiliary,ranganath2016hierarchical}. Similarly, our Bayesian approach naturally allows us to incorporate prior knowledge. For instance, one could place sparsity-inducing priors and impose corresponding variational approximations \citep{ingraham2017variational,ghosh2017model,louizos2017bayesian}. Applying such variational approximations to more expressive autoregressive models would be an interesting avenue to explore in future work.

\subsection*{Acknowledgements}
This research has been partly financed by the Alan Turing Institute under the EPSRC grant EP/N510129/1. The authors acknowledge the use of the UCL Legion High Performance Computing Facility (Legion@UCL), and associated support services, in the completion of this work.

redpajama_set_name: RedPajamaArXiv