# Reconstructing a non-linear interaction in the dark sector with cosmological
observations

Jiangang Kang<EMAIL_ADDRESS>National Astronomical Observatories, Chinese
Academy of Sciences, Beijing 100101, China  CAS Key Laboratory of
FAST,National Astronomical Observatories, Chinese Academy of Sciences, Beijing
100101, China  School of Astronomy and Space Science,University of Chinese
Academy of Sciences, Beijing 100049, China

###### Abstract

In this work we model two non-linear directly interacting scenarios in dark
sector of the universe with the dimensionless parameter $\alpha$ and $\beta$,
which dominate the energy exchange between dark energy and dark matter. The
central goal of this investigation is to research the interacting model and
discuss the cosmological implications based on the current observational
datasets. The class of the interaction is generally characterized by a
coupling function $Q\propto H(z)\rho_{x}$, $x$ denotes the energy density of
dark matter or dark energy. The constrained results we obtained indicate that
the direct interaction in cosmic dark sector is favored by various
observational data and the key effects on CMB power spectrum and linear matter
power spectrum appear compared to $\Lambda$CDM standard paradigm. Finally, we
discuss in depth the effect of different neutrino mass hierarchy on matter
power spectrum and the variation of the ratio of CMB temperature power
spectrum $C_{\ell}^{TT}$ and matter power spectrum $P(k)$ when the $\Delta
N_{eff}=N-3.046$ from 0.5 to 2, respectively.

Keywords: dark energy; dark matter; cosmological observations; cosmological
parameters; neutrinos.

##  I Introduction

The accelerated expansion of the late-time Universe since was discoveried by
observations of type Ia supernovae (Riess et al., 1998; Perlmutter et al.,
1999), understanding the physical nature behind the acceleration is one of the
fundamental topics in modern cosmology.

An exotic species of the energy with the equation of state
$w_{d}^{eff}=P_{d}/\rho_{d}<-\frac{1}{3}$ dubbed as dark energy (DE) has been
introduced (Peebles and Ratra, 2003; Copeland et al., 2006) to explain the
expansion of the Universe. Based on Einstein’s theory of General
Relativity(GR), the dark energy is represented by the cosmological constant
$\Lambda$, which combinates with the pressureless cold dark matter (CDM) have
regarded as a standard model for modern cosmology.

Many predictions from the standard $\Lambda$ cold dark matter($\Lambda$CDM)
theory have been successfully verified by a variety of astronomical
observations. However, there are still several challenging puzzles which need
to be addressed, (i) the so-call fine tuning problem, i.e. the result of the
vacuum energy density based on quantum theory reaches some 120 orders of
magnitude larger than that inferred from the cosmological constant $\Lambda$
(Li et al., 2011; Wetterich, 1995); (ii) coincidence puzzle: is that the
energy density of DM and DE are approximate one order of magnitude, although
the two quantities do not share similar evolution laws and have different
rates of evolution as the universe expands, so why do they happen to be of the
same order of magnitude right now? (Copeland et al., 2006; Wang et al., 2005;
Tocchini-Valentini and Amendola, 2002; Chimento et al., 2003; Rudra, 2016;
Paul and Pavicic, 1997; Zlatev et al., 1999); (iii) the Hubble constant
tension: the$\ H_{0}=67.4\pm 0.5kms^{-1}Mpc^{-1}$ derived from Planck 2018(
P18) Cosmic Microwave Background (CMB) (Collaboration et al., 2018; Di
Valentino, 2017) determination in $\Lambda$CDM is less than 4.4$\sigma$ error
that$\ H_{0}=74.03\pm 1.82kms^{-1}Mpc^{-1}$ obtained directly from the local
distance ladder measurements of SN Ia (R19) (Riess et al., 2019); (iv) the
estimated root-mean-square mass fluctuation amplitude $\sigma_{8}$ in
8$h^{-1}Mpc$ spherical volume and matter energy density $\Omega_{m}$ from the
Planck 2018 CMB measurements (Collaboration et al., 2018; Kilbinger et al.,
2013) under the $\Lambda$CDM cosmology is different from that derived from
weak lensing effects based on the KiDS survey (Köhlinger et al., 2017;
Hildebrandt et al., 2017; Joudaki et al., 2018). Therefore by measuring the
growth rate of structure in the distribution of galaxies as function of
redshift, one can place constraints on gravity, and test if dark energy could
be due to deviations from GR (Guzzo et al., 2008).

In order to address these discrepancies, a reasonable solution is recently
introduced as a novel mechanism beyond the canonical $\Lambda$CDM model. This
modification to dark energy, which is treated as a smooth non-clustering
perfect fluid, may directly interact with dark matter without gravity action.
Such an idea has been extensively investigated by a number of phenomenological
models, which show that the interaction between DE and DM can alleviate the
coincidence problem, and the model predictions have been passed various
observational tests. In particular, many parameterized models assume dynamical
couplings between DE and DM, and show that the current observational
measurements support a nonzero interplay behavior for the cosmic dark
sector(Salvatelli et al., 2014; Yang and Xu, 2014a, b; Nunes et al., 2016; van
de Bruck et al., 2017; Yang et al., 2016; van de Bruck and Thomas, 2019; Yang
et al., 2019; Martinelli et al., 2019; Pan et al., 2020; Bi et al., 2005; Guo
et al., 2018; van de Bruck and Mifsud, 2018). Generally, these dynamical
coupling scenarios are described with a functional coupling term $Q$. In this
paper, we propose two novel coupling scenarios of DE and DM, and then put
constraints on the relevant parameters within the coupling model $Q_{1}$ and
$Q_{2}$ with the help of mainstream cosmological the combination of datasets
and finally we research the impact of two interaction models on the
$C_{\ell}^{TT}$ and matter power $P(k)$ spectra on large scale of universe.
Beyond that, we comment on different the neutrino mass hierarchy problem
(Vagnozzi et al., 2018; Betts et al., 2013; Lesgourgues and Pastor, 2006),
i.e. either $m_{1}<m_{2}\ll m_{3}$, normal hierarchy(NH) or $m_{3}\ll
m_{1}<m_{2}$, inverse hierarchy(IH) and analyse the ratios with $\Lambda$CDM
when setting $\Sigma m_{\nu}=0.06eV$ and comment on some crucial hints about
the neutrino mass splittings can be demonstrated when investigating these IDE
models with the current cosmological observations. From the viewpoint of
phenomenology, the given assumptions made for $Q$ function should be tested
with a number of cosmological probes (Feng et al., 2020).

This paper is structured as follows: The scenarios of the interaction IDE1 and
IDE2 in dark sector we propose are introduced and from which we derive the
relevant evolution equations in Section II. In Section III, we present the SN
Ia, BAO, CMB and H(z) data combinations adopted in this work as well as the
analysis method in each fitting process with these datasets. The eventual
results about the parameters and the $C_{\ell}$ and $P(k)$ spectra are
obtained from the IDE1 and IDE2 models in Section IV. Finally, we draw our
conclusions on two interacting schemes in Section V. We use natural units
throughout $c=8\pi G=1$.

##  II The model of the interaction in dark sector

Considering a homogeneous and isotropic Universe with a spatially flat
Friedmann-Lematre-Robertson-Walker(FLRW) spacetime whose line element follows
the form:

$ds^{2}=-dt^{2}+a^{2}(t)\left(dr^{2}+r^{2}d\theta^{2}+r^{2}sin^{2}\theta
d\phi^{2}\right),$ (1)

where $a(t)$ is cosmic scale factor, $t$ is cosmic time and ($r,\theta,\phi$)
is comoving spherical coordinate system. For a uncoupled standard cosmological
model (SM) usually contains radiation, baryons, dark matter and dark energy,
their perfect fluid energy–momentum tensor satisfies
$T_{\mu\nu}=pg_{\mu\nu}+(\rho+p)U_{\mu}U_{\nu}$ respectively, where
$U_{\mu,\nu}$ stands for four velocity vector. When introducing the directly
interaction of dark constituents, one can understand the conservation equation
of total energy momentum tensor as
$\bigtriangledown_{\mu}(T^{\mu\nu}_{c}+T^{\mu\nu}_{d})=0$ based on the
Einstein’s field equation $G_{\mu\nu}=T_{\mu\nu}^{c}+T_{\mu\nu}^{d}$, $c,d$
denotes dark matter and dark energy while the radiation and baryons
independently evolve as own mode, respectively. Thus, the FLRW evolution
equation for this framework can be written as

$\displaystyle 3H^{2}=\rho_{d}+\rho_{c},$ (2) $\displaystyle
2\dot{H}+H^{2}=-p_{d},$ (3)

where an overdot refers to a derivative w.r.t the cosmic time $t$,
$H=\dot{a}/a$ is the Hubble parameter, $\rho_{c}$, $\rho_{d}$ are the matter
density of dark matter and dark energy and $\ p_{c}$,$\ p_{d}$ represents the
corresponding pressure of dark matter and dark energy, respectively. Dark
matter is usually treated as a pressureless dust matter,$\ p_{c}=0$ while dark
energy is taken as a negative pressure fluid $\ p_{d}=w_{d}\rho_{d}$ with the
effective equation of state(EoS) $w_{d}$. For instance, if the DE is the
smooth vacuum or $\Lambda$, i.e. $w_{d}$ = -1. The physical behavior of the
effective dark fluid when the presence of interaction in dark sector will be
different in light of the sign of the coupling function $Q$. Thus, the whole
energy conservation equation for the interaction can be read as the following
forms:

$\displaystyle\dot{\rho}_{c}+3H(1+w_{c}^{eff})\rho_{c}$ $\displaystyle=Q\,$
(4) $\displaystyle\dot{\rho}_{d}+3H(1+w_{d}^{eff})\rho_{d}$
$\displaystyle=-Q\,$ (5)

For a interacting dark fluid of Universe, the sign and strength of the
coupling function $Q\neq 0$ determines the effective rate of energy conversion
between DM and DE. In the case of $Q<0$, meaning that the energy transfer from
dust DM to DE, while for $Q>0$, oppositely which implys the DE likewise decay
into pressureless DM. As a phenomenological model, it is usually assumed that
the interaction term $Q$ is proportional to the product of energy density of a
dark fluid and the Hubble parameter, $Q\propto\rho_{x}H$. Conventionally, the
coupling strength $Q$ can be tested by a number of observational results with
a linear dependence on the energy density of dark elements as presented in the
mathematical formula (Gavela et al., 2009; Quartin et al., 2008; Böhmer et
al., 2008, 2009; Zimdahl and Pavón, 2007; Bachega et al., 2020):

$Q=3H(\alpha_{c}\rho_{c}+\alpha_{d}\rho_{d}),$ (6)

where $\alpha_{c}$ and $\alpha_{d}$ are dimensionless constants dominating the
strength of the DM-DE interaction. In absence of the interaction,namely $Q=0$,
the model restore to $\Lambda$CDM cosmology. The redshifted energy density
evolution equations for any coupling term $Q$ and the effective EoS $w^{eff}$
for dark sector can be deduced as:

$\displaystyle\rho_{c}=\rho_{c0}\left[\int_{a_{0}}^{a}(\frac{Q}{aH})da\right],$
(7) $\displaystyle\rho_{d}=$
$\displaystyle\rho_{d0}\left[a^{-3(1+w_{d})}-\int_{a_{0}}^{a}(\frac{Q}{aH})\right]da,$
(8) $\displaystyle\ w_{c}^{eff}$ $\displaystyle=\frac{Q}{3H\rho_{c}},\quad
w_{d}^{eff}=w_{d}-\frac{Q}{3H\rho_{d}},$ (9)

where $w_{c}^{eff}$ and $w_{d}^{eff}$ are respectively termed as the effective
equation of state for CDM and DE, $\rho_{c0}$ and $\rho_{d0}$ denote cold dark
matter and dark energy current matter density, and $a_{0}$ is today cosmic
factor, $Q$ stands for any coupling function of dark sector,respectively,

The origin of instabilities depend on different signs of $\alpha_{c}$ and
$\alpha_{d}$ and effective EoS $w^{eff}$ under the interacting dark energy
(IDE) scenarios(Li et al., 2014a, b; Guo et al., 2017; Zhang, 2017; Dai and
Xia, 2019; Yang et al., 2018). In this work, we propose originally two
interacting scenarios as following physical mechanism: in which the coupling
function $Q(t)$ can be defined as:

$\displaystyle IDE1:Q_{1}=3H\alpha\rho_{d}\left[1+\beta
ln(1-\frac{\rho_{c}}{\rho_{d}})\right],$ (10) $\displaystyle
IDE2:Q_{2}=3H\alpha\rho_{d}\left[1+\beta
sin(\frac{\rho_{c}}{\rho_{d}})\right],$ (11)

where $\alpha$ and $\beta$ are dimensionless coupling parameter governing the
interaction strength of the dark sector and then one can expands $Q$ function
around the $\rho_{c}/\rho_{d}=1$ in the manner of Taylor series and the first
order approximation we take for the interaction model IDE1 and IDE2 and thus
$Q_{1}$ and $Q_{2}$ can be rewritten as a non-linear expression:

$Q_{1}\approx
3H\alpha\rho_{d}(1-\frac{\beta\rho_{c}}{\rho_{d}})=3H\alpha(\rho_{d}-\beta\rho_{c}),\\\
$ (12)

and

$Q_{2}\approx
3H\alpha\rho_{d}(1+\frac{\beta\rho_{c}}{\rho_{d}})=3H\alpha(\rho_{d}+\beta\rho_{c}),$
(13)

Then we further derive the effective EoS Eq.9 in term of above interaction
function via $r=\rho_{c}/\rho_{d}$ for IDE1 model:

$\ w_{c}^{eff}=\alpha(\frac{1}{r}-\beta),\quad
w_{d}^{eff}=w_{d}-\alpha(1-\beta r),$ (14)

meanwhile for IDE2:

$\ w_{c}^{eff}=\alpha(\frac{1}{r}+\beta),\quad
w_{d}^{eff}=w_{d}-\alpha(1+\beta r),$ (15)

Figure 1: The figure shows the redshift evolution of the effective equations
of state for DE(top) and CDM (below) under the interaction models $Q_{1}$ and
$Q_{2}$ when taking different two groups value for $\alpha$ and
$\beta$,respectively.

Fig.1 shows the evolution behavior of $w_{c}^{eff}$ and $w_{d}^{eff}$ with
respect to redshift $z$ for the two interaction schemes $Q_{1}$ and
$Q_{2}$,respectively. The left panel of Fig.1 displays the evolution of
$w_{c}^{eff}$ as well as the evolution of $w_{d}^{eff}$ in the right panel
taking same values of the $\alpha$ and $\beta$ for IDE1 as well as IDE2 when
fixing $w_{d}=-0.99$, one can find the effective equations of state for dust
DM and DE are sensitive to the interaction parameter $\alpha$ and $\beta$ and
the evolution of $w_{c}^{eff}$ and $w_{d}^{eff}$ for IDE1 and IDE2 maintain
stable for $z>1.5$ around the $w_{d}=-1$ from the picture. As for IDE1 we see,
the energy transition occurs from DM to DE $Q<0$ ($\alpha<0$ and $\beta>0$)
while $Q>0$ ($\alpha>0$ and $\beta<0$) energy flows from DM to DE after
$z>0.5$, while for IDE2 case, the energy flows between DE and DM opposite
direction compared to IDE1 case in term of the various sign of $\alpha$ and
$\beta$.

Now, the main perturbative equations of energy density $\delta$ and velocity
$v$ under the any interaction $Q$ of DE and DM model in synchronous gauge for
dark energy become (Ma and Bertschinger, 1995; Jackson et al., 2009; Gavela et
al., 2010, 2009; Salvatelli et al., 2013; Feng et al., 2019; Ballesteros and
Lesgourgues, 2010; Xia et al., 2008; Marcondes et al., 2016):

$\displaystyle\delta^{\prime}_{d}=\dfrac{a}{\bar{\rho}_{d}}{Q}_{d}-\dfrac{a{Q}}{\bar{\rho}_{d}}\left[\delta_{d}+3\mathcal{H}\left(c^{2}_{(s)d}-w_{d}\right)\delta
v\right],$ (16) $\displaystyle
v^{\prime}_{d}=\dfrac{a{Q}}{\bar{\rho}_{d}}\left[v_{d}-\left(1+c^{2}_{(s)d}\right)v_{d}\right]-\dfrac{a}{{\rho}_{d}}\left[\dfrac{c^{2}_{(s)d}\rho_{d}}{a}\delta_{d}-{Q}{v}\right].$
(17)

Similarly, as for the cold dark matter perturbations, $w_{c}=0$ leading to
$c^{2}_{(a)c}={c}^{2}_{(s)c}=0$, its perturbative formula of energy
$\delta_{c}^{\prime}$ and speed $v^{\prime}_{c}$ can be read respectively as:

$\displaystyle\delta^{\prime}_{c}=k^{2}v_{c}-\dfrac{h^{\prime}}{2}-\dfrac{a}{\bar{\rho}_{c}}{Q}+\dfrac{a{Q}}{\bar{\rho}_{c}}\delta_{c}.$
(18) $\displaystyle
v^{\prime}_{c}=-\mathcal{H}v_{c}-\dfrac{a{Q}}{\bar{\rho}_{c}}\left(v-v_{c}\right).$
(19)

here $\delta_{i}=\delta\rho/\rho$ ($i=c,d$) and $\mathcal{H}$ = $aH$ is
comoving Hubble parameter.

Besides, we discuss the contribution of matter overdensity $\delta$ that
resulted from the interaction between dark energy and dark matter. In term of
the general definition of the growth rate of matter perturbations:
$f_{c}=\frac{dln\delta_{c}}{dlna}$=$\delta^{\prime}$/$\delta_{c}$ and for an
arbitrary interaction function $Q$ (Jiménez et al., 2020; Barros et al., 2020;
Pan et al., 2020), the growth rate can be understood as:

$\displaystyle\delta^{\prime\prime}_{c}+\left(1-\frac{Q}{H\rho_{c}}\right)\mathcal{H}\delta^{\prime}_{c}=\frac{3}{2}\mathcal{H}^{2}\Omega_{b}\delta_{b}+\frac{3}{2}\mathcal{H}^{2}\Omega_{c}\delta_{c}\Bigg{\\{}1+\frac{2}{3}\frac{\rho_{d}}{\rho_{c}}\frac{Q}{H\rho_{c}}$
$\displaystyle\Bigg{[}\frac{\mathcal{H}^{\prime}}{\mathcal{H}^{2}},+1-3w_{d}+\frac{w_{d}}{\mathcal{H}(1+w_{d})}+\frac{Q}{H\rho_{c}}\left(1+\frac{\rho_{d}}{\rho_{c}}\right)\Bigg{]}\Bigg{\\}},$
(20)

It’s difficult to solve the analysis expression of $f_{c}$ in Eq.II but we can
fit it using the format $f(z)=\Omega_{\rm m}^{\gamma}$ ($\gamma\approx 0.55$
in $\Lambda$CDM cosmology(Linder and Cahn, 2007; Hudson and Turnbull, 2012))
with observations combinations as shown in Fig.2 and from the plot, we find
that growth formation of the CDM is much sensitive to various interaction
model $Q$ and coupling factors especially after redshfit $z=2$ when DE begins
to dominate cosmic energy ingredient and drives the accelerated expansion of
the universe .

Figure 2: The linear evolution of $f\sigma_{8}$ as the function of redshift
$z$ for interaction models $Q_{1}$ and $Q_{2}$ (dash blue and lime line,black
solid line stands for $\Lambda$CDM model, respectively). Errorbar datapoints
are employed in Table I of (Sagredo et al., 2018)

In this work, we investigate cosmological dynamics with the flat spacetime due
to the above Eqs.1–13, neglecting the radiation density $\Omega_{r0}\approx
10^{-5}$, therefore assuming the matter density relation $\Omega_{m}$
+$\Omega_{d}$ = 1 for IDE1 case can be rewritten:

$\displaystyle
H(z)^{2}/H_{0}^{2}=\Omega_{d}\left(\frac{w}{\alpha+\beta}(1+z)^{3}+\frac{1}{\alpha+\beta}(1+z)^{3(1+w+\alpha)}\right)+\Omega_{m}(1+z)^{3},$
(21)

For the IDE2 scheme:

$\displaystyle
H(z)^{2}/H_{0}^{2}=\Omega_{m}\left(\frac{w}{\alpha+\beta}(1+z)^{3(1+w)}+\frac{1}{\alpha+\beta}(1+z)^{3(1-w-\alpha)}\right)\
+\Omega_{d}(1+z)^{3(1+w)}$ (22)

where $\Omega_{m}$ and $\Omega_{d}$ are current energy density of dark matter
and dark energy.

##  III Data and Methodology

In this section we introduce the mainstreams cosmological probes and the
likelihood function $\mathcal{L}\propto e^{-\chi^{2}/2}$ used in this work to
extract the posterior distribution functions (PDFs) of each parameter within
IDE1 and IDE2 framework.

###  III.1 SN Ia data

The Type Ia Supernovae (SN Ia) play a key role in unveiling the character of
the dark energy as the ”standard candles”. We utilize the Pantheon compilation
consisting of 1048 SN Ia samples covering redshift range $0.01<z<2.3$ (Scolnic
et al., 2018). The $\chi^{2}$ distribution for the SN Ia can be written as

$\chi^{2}_{\rm SN}=\Delta{\bf u}^{T}\cdot{\bf C^{-1}}\cdot\Delta{\bf u}.$ (23)

Here $\Delta{\bf u}={\bf u}_{\rm obs}-{\bf u}_{\rm th}$, ${\bf u}_{\rm obs}$
and ${\bf u}_{\rm th}$ are the vectors of observational and theoretical
distance modulus, respectively,and ${\bf C}$ is the covariance matrix.$\
u_{\rm th}=m_{\rm th}-M$ , where $u_{\rm th}$ is the theoretical apparent
magnitude, it can be obtained as

$\displaystyle u_{th}=5\rm{log}_{10}D_{\rm L}(z)+25,$ (24)

where $M$ is the absolute magnitude of SN Ia as a nuisance parameter, $D_{\rm
L}(z)=(1+z)D_{\rm C}(z)$ is the luminosity distance in unit of $\ Mpc$. The
covariance matrix $\bf C$ takes the form as

${\bf C}={\bf C}_{\rm stat}+{\bf C}_{\rm sys},$ (25)

where ${\bf C}_{\rm stat}$ is the vector of statistic error, ${\bf C}_{\rm
sys}$ is the systematic covariance matrix for the data (Scolnic et al., 2018).

###  III.2 BAO data

Table 1: The 12 BAO data used in this work redshift from $z=0.1$ to 2.4. Redshift | Measurement | Value | $r_{\rm s,fid}$ | Survey | Reference
---|---|---|---|---|---
0.106 | $r_{\rm s}/D_{\rm V}$ | 0.3360$\pm$0.015 | – | 6dFGS | (Beutler et al., 2011)
0.15 | $r_{\rm s}/D_{\rm V}$ | 0.2239$\pm$0.0084 | – | SDSS DR7 | (Ross et al., 2015)
0.32 | $r_{\rm s}/D_{\rm V}$ | 0.1181$\pm$0.0024 | – | BOSS LOW-Z | (Anderson et al., 2014)
0.57 | $r_{\rm s}/D_{\rm V}$ | 0.0726$\pm$0.0007 | – | BOSS CMASS | (Padmanabhan et al., 2012)
0.44 | $r_{\rm s}/D_{\rm V}$ | 0.0870$\pm$0.0042 | – | WiggleZ | (Blake et al., 2012)
0.60 | $r_{\rm s}/D_{\rm V}$ | 0.0672$\pm$0.0031 | – | WiggleZ | (Blake et al., 2012)
0.73 | $r_{\rm s}/D_{\rm V}$ | 0.0593$\pm$0.0020 | – | WiggleZ | (Blake et al., 2012)
2.34 | $r_{\rm s}/D_{\rm V}$ | 0.0320$\pm$0.0013 | – | SDSS-III DR11 | (Delubac et al., 2015)
2.36 | $r_{\rm s}/D_{\rm V}$ | 0.0329$\pm$0.0009 | – | SDSS-III DR11 | (Hell et al., 2015)
0.38 | $D_{\rm M}(r_{\rm s,fid}/r_{\rm s})$ | 1518$\pm$22 | 147.78 | SDSS DR12 | (Alam et al., 2017)
0.51 | $D_{\rm M}(r_{\rm s,fid}/r_{\rm s})$ | 1977$\pm$27 | 147.78 | SDSS DR12 | (Alam et al., 2017)
2.40 | $D_{\rm H}/r_{\rm s}$ | 8.94$\pm$0.22 | – | SDSS DR12 | (du Mas des Bourboux et al., 2017)

Table.1 lists the adopted quantities that derived from the different BAO
measurements as the ”standard ruler”, $r_{\rm s}$ is the radius of the
comoving sound horizon at the drag epoch as following:

$r_{\rm s}=\int_{0}^{t_{\rm s}}c_{\rm s}\frac{{\rm d}t}{a},$ (26)

where $c_{\rm s}$ is the sound speed, $t_{\rm s}$ is the epoch of last
scattering, $a$ is the scale factor. Since $r_{\rm s}$ are not sensitive to
physics at low redshifts, we fix $r_{\rm s}=147.09\pm 0.26$ measured from the
Planck 2018 release (Collaboration et al., 2018). The spherically averaged
distance $D_{\rm V}$ is given by

$D_{\rm V}(z)=\left[D_{\rm M}^{2}(z)\frac{cz}{H(z)}\right]^{1/3},$ (27)

where $D_{\rm M}(z)=(1+z)D_{\rm A}$ is the comoving angular diameter distance,
and $D_{\rm A}=D_{\rm C}/(1+z)$ is the physical angular diameter distance.
$D_{\rm H}=c/H(z)$ is the Hubble distance.

The $\chi^{2}$ for the BAO data can be calculated by

$\chi^{2}_{\rm BAO}=\Delta{\bf D}^{T}\cdot{\bf C_{\rm D}}^{-1}\cdot\Delta{\bf
D},$ (28)

where ${\Delta D}={\bf D}_{\rm obs}-{\bf D}_{\rm th}$, ${\bf D}_{\rm obs}$ and
${\bf D}_{\rm th}$ are the observational and theoretical quantities shown in
Table . 1 and $\bf C_{\rm D}$ is corresponding covariance matrix.

###  III.3 CMB data

For the CMB dataset, we exploit the parameter of acoustic scale $l_{A}$, the
shift parameter $R$, the decoupling redshift $z_{\ast}$ inferred from the
Planck CMB measurement and the distance priors from the Planck 2018
results(Collaboration et al., 2018; Planck Collaboration et al., 2020), they
are defined respectively as

$R=\sqrt{\Omega_{\rm m}H^{2}_{0}}r(z_{\ast}),$ (29) $\ell_{\rm A}=\pi
r(z_{\ast})/r_{\rm s}(z_{\ast}),$ (30)

where $\Omega_{\rm{m}}$ is current energy density of dark matter.
$r(z_{\ast})$ is the comoving size of the sound horizon at the redshift of the
decoupling epoch of photons $z_{\ast}$, $z_{\ast}$ can be expressed as (Hu and
Sugiyama, 1996; Cao et al., 2011; Lazkoz and Majerotto, 2007):

$z_{\ast}=1048[1+0.00124(\Omega_{\rm
b}h^{2})^{-0.738}][1+g_{1}(\Omega_{\rm{m}}h^{2})^{g_{2}}],$ (31)

where

$g_{1}=\frac{0.0783(\Omega_{\rm{b}}h^{2})^{-0.238}}{1+39.5(\Omega_{\rm{b}}h^{2})^{0.763}},g_{2}=\frac{0.560}{1+21.1(\Omega_{\rm{b}}h^{2})^{1.81}}.$
(32)

The $\chi^{2}$ distribution for the CMB data can be estimated as

$\chi^{2}_{\rm CMB}=\Delta{\bf X}^{T}C^{-1}_{\rm CMB}\Delta{\bf X},$ (33)

where $\omega_{b}=\Omega_{b}h^{2}$ is the baryon energy density, $h$ is the
dimensionless Hubble constant. $\Delta X=\bf X_{obs}-\bf X_{th}$, $\bf
X=(R,\ell_{\rm A},\omega_{\rm b})$ and $C_{\rm CMB}$ is the covariance matrix.

###  III.4 H(z) Data

The $H(z)$ data contain 30 data points in the redshift range from 0 to 2 in
Table 4 of (Moresco et al., 2016), which are obtained using the differential-
age method, that is to compare the ages of passively-evolving galaxies with
similar metallicity, separated in a small redshift interval (Jimenez and Loeb,
2002) serving as cosmic chronometers and yielding a model-independent
measurement of the expansion rate of the Universe at various redshifts. The
$\chi^{2}$ distribution for the $H(z)$ data can be expressed as (Cao et al.,
2011; Lazkoz and Majerotto, 2007)

$\chi^{2}_{H}=\sum_{i=1}^{N=30}\frac{\left[H_{\rm obs}(z_{i})-H_{\rm
th}(z_{i})\right]^{2}}{\sigma_{H}^{2}}.$ (34)

Here $H_{\rm obs}$ and $H_{\rm th}$ are the observational and theoretical
Hubble parameters, respectively, and $\sigma_{H}$ is the error.

Finally, the joint $\chi^{2}$ of above four datasets employed in this work can
be calculated by

$\chi^{2}=\chi^{2}_{SN}+\chi^{2}_{BAO}+\chi^{2}_{CMB}+\chi^{2}_{H(z)}.$ (35)

Table 2: The table lists the priors on the parameters space. Parameter | Prior range
---|---
$\Omega_{m}$ | [0, 0.7]
$h$ | [0.5, 1]
$\alpha$ | [-1, 1]
$\beta$ | [-1, 1]
$\ w_{d}^{eff}$ | [-3,-0.3]
$\sigma_{8}$ | [0.2,1.4]

##  IV Results

In order to analyse the IDE1 and IDE2 model introduced in Sec.II and place
constraints on the free parameters within the scenarios, we use CMB data and
the combination of CMB+Pantheon, CMB+BAO+H(z), and CMB+Pantheon++BAO+H(z)
datasets and modify the Boltzmann code CLASS (Blas et al.,
2011)111https://github.com/lesgourg/class_public to solve the background and
perturbation equations of IDE1 and IDE2 model and then we utilize MontePython,
a Markov chain Monte Carlo package (MCMC)(Audren et al., 2013; Brinckmann and
Lesgourgues, 2018; Katzgraber,
2011)222https://github.com/brinckmann/montepython_public with the Metropolis-
Hastings algorithm in accordance with the Gelman-Rubin convergence criterion,
requiring $\left|R-1\right|<$ 0.01 to extract the chains of each parameter by
computing the reduced chi-square $\chi^{2}_{red}=\chi^{2}_{min}/(N-n)$ in the
fitting process for each dataset combination, where $\chi^{2}_{min}$ is the
minimum $\chi^{2}$, $N$ and $n$ are the number of data point and free
parameter,respectively. The priors on the parameter space are shown in
Table.2.

Figure 3:  The one-dimensional posterior distributions and two-dimensional contours in 1-$\sigma$ (68.3%) and 2-$\sigma$ (95.5%) confidence level(C.L.) of the free parameters for the $Q_{1}=3H\alpha\rho_{d}(1+\beta ln(1-\frac{\rho_{c}}{\rho_{d}})$ interacting dark energy scenario. The red, blue, green, and orange plots stand for the constraint results from CMB, CMB+SN Ia, CMB+BAO+H(z) and joint datasets, respectively. Table 3: The table summarizes the mean values and 1-$\sigma$(68.3%) uncertainties of the parameters under the scenario of the IDE1:$Q_{1}=3H\alpha\rho_{d}(1+\beta ln(1-\frac{\rho_{c}}{\rho_{d}})$ case in Fig.3. Here the parameter $\Omega_{m}$ equal the energy fractions of baryons plus dark matter, i.e. $\Omega_{m}=\Omega_{c}+\Omega_{b}$. Parameter | CMB | CMB+SN Ia | CMB+BAO+H(z) | CMB+SN Ia+BAO+H(z)
---|---|---|---|---
$\Omega_{m}$ | $0.322_{-0.009}^{+0.01}$ | $0.303_{-0.007}^{+0.007}$ | $0.297_{-0.008}^{+0.008}$ | $0.298_{-0.006}^{+0.006}$
$h$ | $0.677_{-0.008}^{+0.007}$ | $0.726_{-0.0067}^{+0.0061}$ | $0.691_{-0.007}^{+0.007}$ | $0.684_{-0.006}^{+0.0051}$
$\alpha$ | $-0.457_{-0.005}^{+0.005}$ | $-0.465_{-0.004}^{+0.004}$ | $-0.47_{-0.004}^{+0.004}$ | $-0.468_{-0.004}^{+0.003}$
$\beta$ | $0.861_{-0.014}^{+0.012}$ | $0.886_{-0.012}^{+0.013}$ | $0.876_{-0.013}^{+0.013}$ | $0.874_{-0.013}^{+0.013}$
$\ w_{d}^{eff}$ | $-0.986_{-0.189}^{+0.026}$ | $-1.005_{-0.019}^{+0.019}$ | $-1.006_{-0.024}^{+0.027}$ | $-1.015_{-0.022}^{+0.020}$
$\sigma_{8}$ | $0.811_{-0.009}^{+0.01}$ | $0.814_{-0.009}^{+0.008}$ | $0.827_{-0.01}^{+0.01}$ | $0.822_{-0.01}^{+0.01}$
$\chi^{2}_{red}$ | 0.989 | 1.067 | 1.173 | 1.114
Figure 4: The one-dimensional posterior distributions and two-dimensional contours in 1- $\sigma$ (68.3%) and 2-$\sigma$ (95.5%) confidence level(C.L.) of the free parameters for the $Q_{2}=3H\alpha\rho_{d}(1+\beta sin(\frac{\rho_{c}}{\rho_{d}}))\rho_{c}$ interacting dark energy scenario. The red, blue, green, and orange plots show the constraint results from CMB, CMB+SN Ia, CMB+BAO+H(z) and joint datasets, respectively. Table 4: The table shows mean values and 1-$\sigma$(68.3/%) confidence level(C.L.) of the cosmological parameters under the IDE2:$Q_{2}=3H\alpha\rho_{d}(1+\beta sin(\frac{\rho_{c}}{\rho_{d}}))$ case in Fig.4 fitted by above datasets,here $\Omega_{m}=\Omega_{c}+\Omega_{b}$. Parameter | CMB | CMB+SN Ia | CMB+BAO+H(z) | CMB+SN Ia+BAO+H(z) |
---|---|---|---|---|---
$\Omega_{m}$ | $0.317_{-0.022}^{+0.025}$ | $0.298_{-0.018}^{+0.015}$ | $0.300_{-0.007}^{+0.006}$ | $0.294_{-0.006}^{+0.006}$ |
$h$ | $0.680_{-0.028}^{+0.022}$ | $0.701_{-0.020}^{+0.020}$ | $0.694_{-0.006}^{+0.006}$ | $0.686_{-0.006}^{+0.006}$ |
$\alpha$ | $-0.866_{-0.0049}^{+0.0049}$ | $-0.866_{-0.005}^{+0.005}$ | $-0.867_{-0.0045}^{+0.0047}$ | $-0.864_{-0.005}^{+0.005}$ |
$\beta$ | $0.765_{-0.019}^{+0.011}$ | $0.775_{-0.013}^{+0.014}$ | $0.775_{-0.012}^{+0.013}$ | $0.7727_{-0.013}^{+0.012}$ |
$w_{d}^{eff}$ | $-1.005_{-0.017}^{+0.019}$ | $-0.995_{-0.027}^{+0.024}$ | $-0.998_{-0.013}^{+0.018}$ | $-1.007_{-0.013}^{+0.016}$ |
$\sigma_{8}$ | $0.815_{-0.017}^{+0.013}$ | $0.8236_{-0.013}^{+0.014}$ | $0.820_{-0.009}^{+0.009}$ | $0.8278_{-0.009}^{+0.009}$ |
$\chi^{2}_{min}$ | 0.973 | 1.082 | 1.154 | 1.203 |

Firstly, Fig.3 and Fig.4 show the 68.3% and 95.5% confidence level(C.L.)
contour map and 1-D PDFs of the free parameter of IDE1:
$Q_{1}=3H\alpha\rho_{d}(1+\beta ln(1-\frac{\rho_{c}}{\rho_{d}})$ and IDE2:
$Q_{2}=3H\alpha\rho_{d}(1+\beta sin(\frac{\rho_{c}}{\rho_{d}}))$ models
explored by using CMB and CMB+SN Ia and CMB+BAO+H(z) and joint datasets
respectively because of these datasets can break the degeneracy of the among
quantities. We find that the model can be constrained by these datasets well
due to $\chi^{2}_{red}$ is close to 1 for the above four data combinations. As
listed on table 3, the best-fit value of the coupling factor $\alpha$ and
$\beta$ in Fig 3 yield $\alpha\simeq-0.468_{-0.004}^{+0.003}$ and $\beta\simeq
0.874_{-0.013}^{+0.013}$ for the joint datasets and the
$Q_{1}\approx-0.00002H(z)<0$ manifests a difference from zero, i.e. energy
transfers from dark matter into dark energy. In addition, the equation of
state $w_{d}^{eff}<-1$ for DE is preferred from the fitting consequences of
CMB+SN Ia+BAO+H(z) combination. The results of other cosmological parameters
from the fitting are also perfectly consistent with observational values.
Otherwise, concerning Fig.4 and Table.4 present the results of IDE2:
$Q_{2}=3H\alpha\rho_{d}(1+\beta sin(\frac{\rho_{c}}{\rho_{d}}))$ scenario, the
mean values of the interaction factors $\alpha$ $\simeq-0.866$ and
$\beta\simeq 0.775$ and so $Q_{2}\approx 0.00004H(z)>0$ meaning the dark
energy decays into dark matter in light of the results from observational
dataset combinations.

Furthermore, we analyse either issues of $H_{0}$ tension or the $\sigma_{8}$
discrepancy (Di Valentino and Bridle, 2018; Kilbinger et al., 2013) mentioned
in Section.I and both issues can be slightly alleviate due to the constraints
values under the two interacting IDE1 and IDE2 framework. From Table.3 we see,
the IDE1 model gives the dimensionless Hubble constant $h$ from 0.666 to 0.686
using the four data combinations, which are almost consistent in 1$\sigma$
level and similarly, the resulting values $\sigma_{8}$ are tightly obtained
from 0.811 to 0.827. While in the IDE2 model, the tensions are also decreased
to some extent, the values of $h$ and $\sigma_{8}$ show no serious deviation
from the CMB data only to the joint dataset fitting. In short, these results
of parameters from solely Planck CMB data to four observations combinations
imply that the non-linear interactions of DE and DM considered in this work
can be supported well by the available observational data which determined
from $\Lambda$CDM cosmology.

Figure 5: The evolution for the ratio of matter and radiation $\Omega_{\rm m}$
and $\Omega_{\rm r}$ with the variation of the interaction parameter $\alpha$
and $\beta$ for IDE1 and IDE2, here $\Omega_{\rm m}=\Omega_{c}+\Omega_{\rm
b}$. The horizontal gray thick line corresponds to the case of $\Omega_{\rm
m}=\Omega_{\rm r}$, namely the era of the matter-radiation equality and we fix
these parameters for the plot
$\Omega_{m}=0.278$,$\Omega_{b}=0.04$,$\Omega_{d}=0.682$,$\Omega_{\rm
r}=0.00005$.

.

In order to further understand the relation between the coupling factor
$\alpha$ and $\beta$ and the time of matter-radiation equality for the
existing of an interaction in cosmic dark sector, we draw the evolution curves
of $\Omega_{\rm m}/\Omega_{\rm r}$ in redshift space $z$, the time of matter-
radiation equality comes later in both IDE1 and IDE2 case with increasing
value of $\alpha$ and $\beta$ around redshift $z$=3400 to 3500 compared to
non-interacting model as shown in Fig.5, indicating the increasing of
interaction parameter $\alpha$ and $\beta$ will lead to a addition of energy
density of the dust matter $\Omega_{\rm m}$, the epoch of matter-radiation
balance will become late and the sound horizon will scales up.

Figure 6: The effects on the CMB temperature power spectra(left panel) and the
matter power spectra (right panel)for two case of the interaction term $Q_{1}$
and $Q_{2}$. The red and blue curves lines for the IDE1: $Q_{1}$ and IDE2:
$Q_{2}$ reference to $\Lambda$CDM in green solid line, respectively.

Fig.6 presents some key signatures on the CMB temperature power spectrum in
left and the linear matter power spectrum in right panel when taking the
different values for $\alpha$ and $\beta$ describing interacting models IDE1
and IDE2. Concerning the CMB temperature power spectra in the left panel, we
deem that actually either dark energy or dark matter gains additional energy
density and influences the microwave background temperature anisotropy
especially at low multiple $l<30$ under the IDE1 and IDE2 respectively. The
increase of dark energy density within IDE1 compared to IDE2, especially
alters integrated Sachs-Wolfe(ISW) level on large scale due to the decay of
gravitational potential over expansion history as well as changes era of
matter-radiation balance and the height of the first peak of CMB temperature
power spectra is lifted for the tight coupling of dark sector. As for the
right panel, the dominant effects on the linear matter power spectrum $P(k)$
show up at large scales($k<0.042h/Mpc$), hinting the enhancement of the
content of dark matter that forms larger cosmic structures and earlier matter
clusters of galaxies under its gravity. Consequently, the shape of matter
power spectra $P(k)$ somewhat is more intensified due to larger horizon
boundary in IDE2 than IDE1 model. Moreover, the first peak evidently is rised
for IDE1 in small multipole $\ell$ ¡10 since the dark energy fluid flows to
dark matter to change the energy of CMB photons when across it. The presence
of interaction of dark sector changes the regular evolution law and which is
imprinted on the integrated Sachs-Wolfe (ISW) effect in large scale effect of
the CMB TT spectrum.

Figure 7:  The ratio of the CMB temperature power spectra to $\Lambda$CDM in
presence of neutrino for the interaction term IDE1(left panel) and IDE2(right
panel) with the variation of $\Delta N_{eff}$ from 0.5 to 2.

Figure 8: The ratio of the matter power spectra to $\Lambda$CDM of the
interaction term IDE1(left panel) and IDE2(right panel) when $\Delta N_{eff}$
from 0.5 to 2.

In this paragraph, we comment on neutrino hierarchies problem under the
framework IDE1 and IDE2 model. The experiment of neutrino oscillation has
suggested that neutrinos have mass size and mass splittings among three
flavors neutrino, though either absolute mass of neutrinos or mass hierarchies
remian puzzling(see (King, 2007; Vagnozzi et al., 2018; Betts et al., 2013; Bi
et al., 2005; Feng et al., 2020) for a review). Actually, both the properties
of massive neutrinos and dark energy could imprint important signatures on
expansion history and large scale structure (LSS) in the evolution history of
the universe(Yoshikawa et al., 2020; Coloma et al., 2020; Lesgourgues and
Pastor, 2006; Coloma et al., 2020; Betts et al., 2013). Therefore, It’s
expected that some significant information can be found on $C_{\ell}$ and
$P(k)$ spectrum in the direct interaction of dark energy and dark matter
scenario when considering mass hierarchies. From the ratio of $C_{\ell}^{TT}$
to $C_{\ell}^{TT}$ in $\Lambda$CDM model in Fig.7, one can find there is mass
size and mass hierarchies of neutrino species and showing bigger deviation in
IDE1 than IDE2 related to $\Lambda$CDM regardless of NH or IH mode and
demonstrating the gradual suppression for the larger free degree of $\Delta
N_{eff}=N-3.046$ from 0.5 to 2 occurring in IDE1 and IDE2 due to the smoothing
effect of more neutrinos free screaming. Similarly, concerning the ratio of
matter power $P(k)/P(k)_{\Lambda}$CDM is dramatically intensified in
$k>0.02h/Mpc$ in the wake of $\Delta N_{eff}$ turns bigger from 0.5 to 2 in
Fig.8, the attribute of neutrinos may strongly degeneracy with dark matter
halos to cluster of baryon particles and host a vast of galaxies or other
objects in both IDE1 and IDE2 models, but in $k<0.02h/Mpc$ phase, the
magnitude of ratio amplitude of IDE2 than IDE1 is less than and tend to 1 due
to the density of neutrinos are diluted since in larger scales and difficult
to cluster.

##  V Conclusions

In this analysis, we have investigated two scenarios of direct interaction
between the dark energy and dark matter, in which the interaction term $Q$ is
reconstructed as Eq.10 and Eq.11 with coupling parameter $\alpha$ and $\beta$
determining the conversion between DE and DM.

Firstly, we come up with two novel interacting model IDE1 and IDE2, and then
we use (i) the cosmic microwave background from Planck 2018 result, (ii)
Pantheon sample of Supernovae Type Ia, (iii) baryon acoustic oscillations
distance measurements and (iv) direct H(z) observation to put constraints on
the two interacting scenarios $Q_{1}=3H\alpha\rho_{d}(1+\beta
ln(1-\frac{\rho_{c}}{\rho_{d}})$ and $Q_{2}=3H\alpha\rho_{d}(1+\beta
sin(\frac{\rho_{c}}{\rho_{d}}))$ by calling the MontePython sampler and Class
code. All results summarized in Table.3 and 4 and Fig.3 and 4, we find that

(1) our analysis with all combined datasets support non-zero value for the
interacting factors $\alpha$ and $\beta$ while the dark energy equation of
state(EoS) $w<-1$ at 2$\sigma$ confidence level.

(2) With IDE1 or IDE2 model, we point out that the results from the
interactions are compatible with a general model such as the uncoupled $w$CDM
cosmology and fit the latest observations. Specially, it is evident that the
interacting models relaxed the existing tensions of $H_{0}$ and $\sigma_{8}$
arised from different cosmological measurements relying on the $\Lambda$CDM
picture at about 2$\sigma$ confidence level, though it does not seem to be
able to completely alleviate them.

(3) When the addition of neutrinos mass normal hierarchy(NH) and inverse
hierarchy(IH) mode to the interaction theory of IDE1 and IDE2 scenario, we
discuss that the influence on the ratio of CMB temperature spectrum,
indicating the neutrinos obviously smooth temperature fluctuation especially
in $k>100$ with the free degree error $\Delta N_{eff}$ from 0.5 to 2, while
the effect of dark matter to drag baryons particles to cluster and spark the
physical processes are reinforced due to the more neutrinos behaving like dark
matter to form larger structures of halos.

As a whole, we confirm that the interaction of dark energy and dark matter in
this work can be explored by the cosmological observation combinations and
eliminate some discrepancies to some extent that from theory and determination
and even can produce some key signals on the $C_{\ell}$ and $P(k)$ spectrum
can be detected in a number of cosmology surveys in future.

## Acknowlegment

We thank Tongjie Zhang for constructive comments on the manuscript.

## References

  * Riess et al. (1998) A. G. Riess, A. V. Filippenko, P. Challis, A. Clocchiatti, A. Diercks, P. M. Garnavich, R. L. Gilliland, C. J. Hogan, S. Jha, R. P. Kirshner, et al., Astron. J. 116, 1009 (1998), eprint astro-ph/9805201.
  * Perlmutter et al. (1999) S. Perlmutter, G. Aldering, G. Goldhaber, R. A. Knop, P. Nugent, P. G. Castro, S. Deustua, S. Fabbro, A. Goobar, D. E. Groom, et al., Astrophys. J. 517, 565 (1999), eprint astro-ph/9812133.
  * Peebles and Ratra (2003) P. J. Peebles and B. Ratra, Reviews of Modern Physics 75, 559 (2003), eprint astro-ph/0207347.
  * Copeland et al. (2006) E. J. Copeland, M. Sami, and S. Tsujikawa, International Journal of Modern Physics D 15, 1753 (2006), eprint hep-th/0603057.
  * Li et al. (2011) M. Li, X.-D. Li, S. Wang, and Y. Wang, Communications in Theoretical Physics 56, 525 (2011), eprint 1103.5870.
  * Wetterich (1995) C. Wetterich, Astron. Astrophys. 301, 321 (1995), eprint hep-th/9408025.
  * Wang et al. (2005) B. Wang, Y. Gong, and E. Abdalla, Physics Letters B 624, 141 (2005), eprint hep-th/0506069.
  * Tocchini-Valentini and Amendola (2002) D. Tocchini-Valentini and L. Amendola, Phys. Rev. D 65, 063508 (2002), eprint astro-ph/0108143.
  * Chimento et al. (2003) L. P. Chimento, A. S. Jakubi, D. Pavón, and W. Zimdahl, Phys. Rev. D 67, 083513 (2003), eprint astro-ph/0303145.
  * Rudra (2016) P. Rudra, in _41st COSPAR Scientific Assembly_ (2016), vol. 41 of _COSPAR Meeting_ , pp. H0.2–11–16.
  * Paul and Pavicic (1997) H. Paul and M. Pavicic, Journal of the Optical Society of America B Optical Physics 14, 1275 (1997), eprint quant-ph/9908023.
  * Zlatev et al. (1999) I. Zlatev, L. Wang, and P. J. Steinhardt, Phys. Rev. Lett. 82, 896 (1999), eprint astro-ph/9807002.
  * Collaboration et al. (2018) P. Collaboration, N. Aghanim, Y. Akrami, M. Ashdown, J. Aumont, C. Baccigalupi, M. Ballardini, A. J. Banday, R. B. Barreiro, N. Bartolo, et al., _Planck 2018 results. vi. cosmological parameters_ (2018), eprint 1807.06209.
  * Di Valentino (2017) E. Di Valentino, Nature Astronomy 1, 569 (2017), eprint 1709.04046.
  * Riess et al. (2019) A. G. Riess, S. Casertano, W. Yuan, L. M. Macri, and D. Scolnic, Astrophys. J. 876, 85 (2019), eprint 1903.07603.
  * Kilbinger et al. (2013) M. Kilbinger, L. Fu, C. Heymans, F. Simpson, J. Benjamin, T. Erben, J. Harnois-Déraps, H. Hoekstra, H. Hildebrand t, T. D. Kitching, et al., Mon. Not. R. Astron. Soc. 430, 2200 (2013), eprint 1212.3338.
  * Köhlinger et al. (2017) F. Köhlinger, M. Viola, B. Joachimi, H. Hoekstra, E. van Uitert, H. Hildebrandt, A. Choi, T. Erben, C. Heymans, S. Joudaki, et al., Mon. Not. R. Astron. Soc. 471, 4412 (2017), eprint 1706.02892.
  * Hildebrandt et al. (2017) H. Hildebrandt, M. Viola, C. Heymans, S. Joudaki, K. Kuijken, C. Blake, T. Erben, B. Joachimi, D. Klaes, L. Miller, et al., Mon. Not. R. Astron. Soc. 465, 1454 (2017), eprint 1606.05338.
  * Joudaki et al. (2018) S. Joudaki, C. Blake, A. Johnson, A. Amon, M. Asgari, A. Choi, T. Erben, K. Glazebrook, J. Harnois-Déraps, C. Heymans, et al., Mon. Not. R. Astron. Soc. 474, 4894 (2018), eprint 1707.06627.
  * Guzzo et al. (2008) L. Guzzo, M. Pierleoni, B. Meneux, E. Branchini, O. Le Fèvre, C. Marinoni, B. Garilli, J. Blaizot, G. De Lucia, A. Pollo, et al., Nature (London) 451, 541 (2008), eprint 0802.1944.
  * Salvatelli et al. (2014) V. Salvatelli, N. Said, M. Bruni, A. Melchiorri, and D. Wands, Physical Review Letters 113, 181301 (2014), eprint 1406.7297.
  * Yang and Xu (2014a) W. Yang and L. Xu, Phys. Rev. D 90, 083532 (2014a), eprint 1409.5533.
  * Yang and Xu (2014b) W. Yang and L. Xu, Phys. Rev. D 89, 083517 (2014b), eprint 1401.1286.
  * Nunes et al. (2016) R. C. Nunes, S. Pan, and E. N. Saridakis, Phys. Rev. D 94, 023508 (2016), eprint 1605.01712.
  * van de Bruck et al. (2017) C. van de Bruck, J. Mifsud, and J. Morrice, Phys. Rev. D 95, 043513 (2017), eprint 1609.09855.
  * Yang et al. (2016) W. Yang, H. Li, Y. Wu, and J. Lu, JCAP. 10, 007 (2016), eprint 1608.07039.
  * van de Bruck and Thomas (2019) C. van de Bruck and C. C. Thomas, Phys. Rev. D 100, 023515 (2019), eprint 1904.07082.
  * Yang et al. (2019) W. Yang, S. Pan, and A. Paliathanasis, Mon. Not. R. Astron. Soc. 482, 1007 (2019), eprint 1804.08558.
  * Martinelli et al. (2019) M. Martinelli, N. B. Hogg, S. Peirone, M. Bruni, and D. Wands, Mon. Not. R. Astron. Soc. 488, 3423 (2019), eprint 1902.10694.
  * Pan et al. (2020) S. Pan, W. Yang, and A. Paliathanasis, Monthly Notices of the Royal Astronomical Society 493, 3114 (2020), ISSN 0035-8711, eprint https://academic.oup.com/mnras/article-pdf/493/3/3114/32890756/staa213.pdf, URL https://doi.org/10.1093/mnras/staa213.
  * Bi et al. (2005) X.-J. Bi, B. Feng, H. Li, and X. Zhang, Phys. Rev. D 72, 123523 (2005), eprint hep-ph/0412002.
  * Guo et al. (2018) R.-Y. Guo, J.-F. Zhang, and X. Zhang, Chinese Physics C 42, 095103 (2018), eprint 1803.06910.
  * van de Bruck and Mifsud (2018) C. van de Bruck and J. Mifsud, Physical Review D 97 (2018), ISSN 2470-0029, URL http://dx.doi.org/10.1103/PhysRevD.97.023506.
  * Vagnozzi et al. (2018) S. Vagnozzi, S. Dhawan, M. Gerbino, K. Freese, A. Goobar, and O. Mena, Phys. Rev. D 98, 083501 (2018), eprint 1801.08553.
  * Betts et al. (2013) S. Betts, W. R. Blanchard, R. H. Carnevale, C. Chang, C. Chen, S. Chidzik, L. Ciebiera, P. Cloessner, A. Cocco, A. Cohen, et al., arXiv e-prints arXiv:1307.4738 (2013), eprint 1307.4738.
  * Lesgourgues and Pastor (2006) J. Lesgourgues and S. Pastor, 307 (2006), eprint astro-ph/0603494.
  * Feng et al. (2020) L. Feng, H.-L. Li, J.-F. Zhang, and X. Zhang, Science China Physics, Mechanics, and Astronomy 63, 220401 (2020), eprint 1903.08848.
  * Gavela et al. (2009) M. B. Gavela, D. Hernández, L. Lopez Honorez, O. Mena, and S. Rigolin, JCAP. 2009, 034 (2009), eprint 0901.1611.
  * Quartin et al. (2008) M. Quartin, M. O. Calvão, S. E. Jorás, R. R. R. Reis, and I. Waga, JCAP. 2008, 007 (2008), eprint 0802.0546.
  * Böhmer et al. (2008) C. G. Böhmer, G. Caldera-Cabral, R. Lazkoz, and R. Maartens, Phys. Rev. D 78, 023505 (2008), eprint 0801.1565.
  * Böhmer et al. (2009) C. Böhmer, G. Caldera-Cabral, R. Lazkoz, and R. Maartens, in _American Institute of Physics Conference Series_ , edited by K. E. Kunze, M. Mars, and M. A. Vázquez-Mozo (2009), vol. 1122 of _American Institute of Physics Conference Series_ , pp. 197–200.
  * Zimdahl and Pavón (2007) W. Zimdahl and D. Pavón, Classical and Quantum Gravity 24, 5461 (2007), eprint astro-ph/0606555.
  * Bachega et al. (2020) R. R. Bachega, A. A. Costa, E. Abdalla, and K. Fornazier, Journal of Cosmology and Astroparticle Physics 2020, 021–021 (2020), ISSN 1475-7516, URL http://dx.doi.org/10.1088/1475-7516/2020/05/021.
  * Li et al. (2014a) Y.-H. Li, J.-F. Zhang, and X. Zhang, Phys. Rev. D 90, 063005 (2014a), eprint 1404.5220.
  * Li et al. (2014b) Y.-H. Li, J.-F. Zhang, and X. Zhang, Phys. Rev. D 90, 123007 (2014b), eprint 1409.7205.
  * Guo et al. (2017) R.-Y. Guo, Y.-H. Li, J.-F. Zhang, and X. Zhang, JCAP. 2017, 040 (2017), eprint 1702.04189.
  * Zhang (2017) X. Zhang, Science China Physics, Mechanics, and Astronomy 60, 50431 (2017), eprint 1702.04564.
  * Dai and Xia (2019) J.-P. Dai and J.-Q. Xia, Astrophys. J. 876, 125 (2019), eprint 1904.04149.
  * Yang et al. (2018) W. Yang, S. Pan, E. Di Valentino, R. C. Nunes, S. Vagnozzi, and D. F. Mota, JCAP. 2018, 019 (2018), eprint 1805.08252.
  * Ma and Bertschinger (1995) C.-P. Ma and E. Bertschinger, Astrophys. J. 455, 7 (1995), eprint astro-ph/9506072.
  * Jackson et al. (2009) B. M. Jackson, A. Taylor, and A. Berera, Phys. Rev. D 79, 043526 (2009), eprint 0901.3272.
  * Gavela et al. (2010) M. B. Gavela, L. Lopez Honorez, O. Mena, and S. Rigolin, JCAP. 2010, 044 (2010), eprint 1005.0295.
  * Salvatelli et al. (2013) V. Salvatelli, A. Marchini, L. Lopez-Honorez, and O. Mena, Phys. Rev. D 88, 023531 (2013), eprint 1304.7119.
  * Feng et al. (2019) L. Feng, J.-F. Zhang, and X. Zhang, Physics of the Dark Universe 23, 100261 (2019), eprint 1712.03148.
  * Ballesteros and Lesgourgues (2010) G. Ballesteros and J. Lesgourgues, JCAP. 2010, 014 (2010), eprint 1004.5509.
  * Xia et al. (2008) J.-Q. Xia, Y.-F. Cai, T.-T. Qiu, G.-B. Zhao, and X. Zhang, International Journal of Modern Physics D 17, 1229 (2008), eprint astro-ph/0703202.
  * Marcondes et al. (2016) R. J. F. Marcondes, R. C. G. Landim, A. A. Costa, B. Wang, and E. Abdalla, JCAP. 2016, 009 (2016), eprint 1605.05264.
  * Jiménez et al. (2020) J. B. Jiménez, L. Heisenberg, T. Koivisto, and S. Pekar, Phys. Rev. D 101, 103507 (2020), eprint 1906.10027.
  * Barros et al. (2020) B. J. Barros, T. Barreiro, T. Koivisto, and N. J. Nunes, Physics of the Dark Universe 30, 100616 (2020), eprint 2004.07867.
  * Linder and Cahn (2007) E. V. Linder and R. N. Cahn, Astroparticle Physics 28, 481 (2007), eprint astro-ph/0701317.
  * Hudson and Turnbull (2012) M. J. Hudson and S. J. Turnbull, The Astrophysical Journal 751, L30 (2012), ISSN 2041-8213, URL http://dx.doi.org/10.1088/2041-8205/751/2/L30.
  * Sagredo et al. (2018) B. Sagredo, S. Nesseris, and D. Sapone, Phys. Rev. D 98, 083543 (2018), eprint 1806.10822.
  * Scolnic et al. (2018) D. M. Scolnic, D. O. Jones, A. Rest, Y. C. Pan, R. Chornock, R. J. Foley, M. E. Huber, R. Kessler, G. Narayan, A. G. Riess, et al., Astrophys. J. 859, 101 (2018), eprint 1710.00845.
  * Beutler et al. (2011) F. Beutler, C. Blake, M. Colless, D. H. Jones, L. Staveley-Smith, L. Campbell, Q. Parker, W. Saunders, and F. Watson, Mon. Not. R. Astron. Soc. 416, 3017 (2011), eprint 1106.3366.
  * Ross et al. (2015) A. J. Ross, L. Samushia, C. Howlett, W. J. Percival, A. Burden, and M. Manera, Mon. Not. R. Astron. Soc. 449, 835 (2015), eprint 1409.3242.
  * Anderson et al. (2014) L. Anderson, É. Aubourg, S. Bailey, F. Beutler, V. Bhardwaj, M. Blanton, A. S. Bolton, J. Brinkmann, J. R. Brownstein, A. Burden, et al., Mon. Not. R. Astron. Soc. 441, 24 (2014), eprint 1312.4877.
  * Padmanabhan et al. (2012) N. Padmanabhan, X. Xu, D. J. Eisenstein, R. Scalzo, A. J. Cuesta, K. T. Mehta, and E. Kazin, Mon. Not. R. Astron. Soc. 427, 2132 (2012), eprint 1202.0090.
  * Blake et al. (2012) C. Blake, S. Brough, M. Colless, C. Contreras, W. Couch, S. Croom, D. Croton, T. M. Davis, M. J. Drinkwater, K. Forster, et al., Mon. Not. R. Astron. Soc. 425, 405 (2012), eprint 1204.3674.
  * Delubac et al. (2015) T. Delubac, J. E. Bautista, N. G. Busca, J. Rich, D. Kirkby, S. Bailey, A. Font-Ribera, A. Slosar, K.-G. Lee, M. M. Pieri, et al., Astron. Astrophys. 574, A59 (2015), eprint 1404.1801.
  * Hell et al. (2015) S. W. Hell, S. J. Sahl, M. Bates, X. Zhuang, R. Heintzmann, M. J. Booth, J. Bewersdorf, G. Shtengel, H. Hess, P. Tinnefeld, et al., Journal of Physics D Applied Physics 48, 443001 (2015), eprint 1711.04999.
  * Alam et al. (2017) S. Alam, M. Ata, S. Bailey, F. Beutler, D. Bizyaev, J. A. Blazek, A. S. Bolton, J. R. Brownstein, A. Burden, C.-H. Chuang, et al., Mon. Not. R. Astron. Soc. 470, 2617 (2017), eprint 1607.03155.
  * du Mas des Bourboux et al. (2017) H. du Mas des Bourboux, J.-M. Le Goff, M. Blomqvist, N. G. Busca, J. Guy, J. Rich, C. Yèche, J. E. Bautista, É. Burtin, and K. S. Dawson, Astron. Astrophys. 608, A130 (2017), eprint 1708.02225.
  * Planck Collaboration et al. (2020) Planck Collaboration, N. Aghanim, Y. Akrami, M. Ashdown, J. Aumont, C. Baccigalupi, M. Ballardini, A. J. Banday, R. B. Barreiro, N. Bartolo, et al., Astron. Astrophys. 641, A5 (2020), eprint 1907.12875.
  * Hu and Sugiyama (1996) W. Hu and N. Sugiyama, Astrophys. J. 471, 542 (1996), eprint astro-ph/9510117.
  * Cao et al. (2011) S. Cao, N. Liang, and Z.-H. Zhu, Monthly Notices of the Royal Astronomical Society 416, 1099–1104 (2011), ISSN 0035-8711, URL http://dx.doi.org/10.1111/j.1365-2966.2011.19105.x.
  * Lazkoz and Majerotto (2007) R. Lazkoz and E. Majerotto, Journal of Cosmology and Astroparticle Physics 2007, 015–015 (2007), ISSN 1475-7516, URL http://dx.doi.org/10.1088/1475-7516/2007/07/015.
  * Moresco et al. (2016) M. Moresco, L. Pozzetti, A. Cimatti, R. Jimenez, C. Maraston, L. Verde, D. Thomas, A. Citro, R. Tojeiro, and D. Wilkinson, Journal of Cosmology and Astroparticle Physics 2016, 014–014 (2016), ISSN 1475-7516, URL http://dx.doi.org/10.1088/1475-7516/2016/05/014.
  * Jimenez and Loeb (2002) R. Jimenez and A. Loeb, Astrophys. J. 573, 37 (2002), eprint astro-ph/0106145.
  * Blas et al. (2011) D. Blas, J. Lesgourgues, and T. Tram, JCAP. 7, 034 (2011), eprint 1104.2933.
  * Audren et al. (2013) B. Audren, J. Lesgourgues, K. Benabed, and S. Prunet, JCAP 1302, 001 (2013), eprint 1210.7183.
  * Brinckmann and Lesgourgues (2018) T. Brinckmann and J. Lesgourgues (2018), eprint 1804.07261.
  * Katzgraber (2011) H. G. Katzgraber, _Introduction to monte carlo methods_ (2011), eprint 0905.1629.
  * Di Valentino and Bridle (2018) E. Di Valentino and S. Bridle, Symmetry 10 (2018), ISSN 2073-8994, URL https://www.mdpi.com/2073-8994/10/11/585.
  * King (2007) S. F. King, Contemporary Physics 48, 195 (2007), eprint https://doi.org/10.1080/00107510701770539, URL https://doi.org/10.1080/00107510701770539.
  * Yoshikawa et al. (2020) K. Yoshikawa, S. Tanaka, N. Yoshida, and S. Saito, arXiv e-prints arXiv:2010.00248 (2020), eprint 2010.00248.
  * Coloma et al. (2020) P. Coloma, M. C. Gonzalez-Garcia, and M. Maltoni, arXiv e-prints arXiv:2009.14220 (2020), eprint 2009.14220.


# LAST: Scalable Lattice-Based Speech Modelling in JAX

###### Abstract

We introduce LAST, a LAttice-based Speech Transducer library in JAX. With an
emphasis on flexibility, ease-of-use, and scalability, LAST implements
differentiable weighted finite state automaton (WFSA) algorithms needed for
training & inference that scale to a large WFSA such as a recognition lattice
over the entire utterance. Despite these WFSA algorithms being well-known in
the literature, new challenges arise from performance characteristics of
modern architectures, and from nuances in automatic differentiation. We
describe a suite of generally applicable techniques employed in LAST to
address these challenges, and demonstrate their effectiveness with benchmarks
on TPUv3 and V100 GPU.

Index Terms—  ASR, weighted finite state machines

##  1 Introduction

Weighted finite state automata (WFSA) have been a vital tool in building
automatic speech recognition (ASR) systems for both model training &
inference. Common neural ASR models (CTC [1], RNN-T [2]) exhibit a close
connection to the WFSA formalism, and there has been a renewed interest in
casting these models into WFSAs, e.g. frameworks such as k2 [3], GTN [4], and
PyTorch-Struct [5]. More recently, [6] demonstrates globally normalized
sequence models built within the WFSA formalism can greatly reduce the word
error rate of streaming ASR models.

Motivated by the need for efficient yet flexible and easy-to-use sequence
modelling tools for implementing the aforementioned methods, especially the
globally normalized models such as [6], we developed LAST, a lattice-based
speech transducer library in JAX [7]. LAST provides a simple interface for
operations over a recognition lattice, and is fully compatible with core
features in JAX such as automatic differentiation and function
transformations. Being written in pure JAX allows LAST to run efficiently on
CPUs, GPUs, and TPUs without any code modification, whereas the existing WFSA
based sequence modelling toolkits [4, 3] employ handwritten kernels targeting
a limited set of hardware.

In this paper, we introduce the design of LAST, discuss the technical
challenges that drive crucial design choices, and present benchmark results.
LAST is open source software under the Apache License 2.0. The complete code
can be found at https://github.com/google-research/last.

##  2 Background

###  2.1 Finite State Automata

We refer readers to [8, 9] for a comprehensive introduction to finite state
automata and related algorithms. A weighted finite state automaton (WFSA)
$A=(\Sigma,Q,i,f,E)$ over a semiring
$(\mathbb{K},\oplus,\otimes,\bar{0},\bar{1})$ is specified by  a finite
alphabet $\Sigma$,  a finite set of states $Q$,  an initial state $i\in Q$,  a
final state $f$,  and a finite set of arcs $E\subseteq
Q\times(\Sigma\cup\\{\epsilon\\})\times\mathbb{K}\times Q$ ($\epsilon$ denotes
the empty label sequence).  We denote  the set of accepting paths $\Pi$;  an
arc $e$’s weight $w[e]$;  $\mathrm{next[e]}$and $\mathrm{prev[e]}$ the source
and destination states of arc $e$.  WFSA are also commonly refered to as
_lattices_ in the ASR literature. The most commonly used WFSA algorithms in
building ASR systems are

Intersection

A WFSA is a compact representation of a set of weighted strings from
$\Sigma^{*}\times\mathbb{K}$. The intersection of two weighted sets
representable as WFSAs is equivalent to the intersection of the corresponding
WFSAs.

Shortest distance

The shortest distance of a WFSA is the $\oplus$-sum of the weights of all
paths from state $i$ to state $f$ (_accepting paths_). Most ASR lattices are
acyclic WFSAs by construction, and their shortest distance can be computed by
visiting the states in topological order.

Shortest path

The shortest path of a WFSA is the accepting path with lowest cost or highest
score. It can be computed for acyclic WFSAs in a similar fashion as shortest
distance.

For example, the computation of the RNN-T loss [2] can be broken down into two
WFSA operations:  (1) Intersecting the (infinite) complete recognition lattice
with the reference output string (Figure 1 of [2]);  (2) Shortest distance on
the intersection.  Another example is the computation of the globally
normalized loss in [6]:  (1) Shortest distance on the complete recognition
lattice to obtain the denominator;  (2) Similar to RNN-T, intersection with
the reference and then shortest distance to obtain the numerator.

During training, we differentiate through shortest distance & intersection to
obtain the gradients of the training loss with respect to the arc weights.

###  2.2 JAX and Automatic Differentiation

We chose to implement LAST in JAX because of the performance, ease of use, and
flexibility JAX offers. JAX [7] is a high-performance machine learning
research toolkit that has seen a growing popularity. JAX programs written in
the NumPy API execute efficiently on CPUs, GPUs and TPUs. Similar to [10], the
vector-Jacobian-product function needed for (reverse mode) automatic
differentiation [11, 12] can be either automatically derived for any function
expressed in differentiable JAX primitives, or manually defined.

###  2.3 Gradients of Shortest Distance

With the WFSA topology fixed, the shortest distance
$D=\bigoplus_{\pi\in\Pi}\left(\bigotimes_{e\in\pi}w[e]\right)$ is a function
of arc weights $w[e],\forall e\in E$. Under the log semiring, $\frac{\partial
D}{\partial w[e]}$ is the marginal probability of the arc $e$, viewing path
weights as unnormalized log-probabilities,

$\frac{\partial D}{\partial w[e]}=\sum_{\pi\in\Pi\land
e\in\pi}\frac{\exp\left(\sum_{e^{\prime}\in\pi}w[e^{\prime}]\right)}{\exp(D)}.$
(1)

The forward-backward algorithm [13] is commonly used to compute this quantity
for acyclic WFSAs, while automatic differentiation can also be applied to
achieve the same goal in the same time complexity [14] (more on this in
Section 4.2).

Moreover, under the tropical semiring, $\frac{\partial D}{\partial w[e]}$ is a
0/1 mask denoting if $e$ is on the shortest path. A differentiable shortest
distance algorithm thus greatly simplifies the implementation of inference: we
get the shortest path algorithm for free. This technique of deriving shortest
path by differentiating shortest distance is also adopted by k2 [3] and
PyTorch-Struct [5].

##  3 Library Design

The GNAT model outlined in [6] is a modular theoretical framework encompassing
all the common neural speech recognition models. [6] reports that the globally
normalized GNAT model greatly reduces the word error rate of streaming ASR.
The primary goal of LAST is to replicate GNAT. Therefore, the top level API of
LAST closely follows components of the GNAT model. GNAT breaks down modelling
choices into 3 components:  The context dependency is an unweighted
deterministic finite automaton that accepts $\Sigma^{*}$. The states in a
context dependency encode the output history.  The alignment lattice for a
given input length $T$ is an acyclic unweighted finite automaton. The states
in an alignment lattice encode the correspondence between input frames and
output alignment labels.  The weight function is a function $Q_{T}\times
Q_{C}\times(\Sigma\cup\\{\epsilon\\})\rightarrow\mathbb{K}$. The intersection
of the context dependency and the alignment lattice defines the topology of
the recognition lattice, whereas the weight function defines the weights of
each arc.  Figures 1(a), 1(b), and 1(c) are a concrete example of these
components from [6].

$\epsilon$$a$$b$$a$$b$$a$$b$$a$$b$

(a) Context size 1 $n$-gram context dependency with $\Sigma=\\{a,b\\}$.

0123$\Sigma$$\epsilon$$\Sigma$$\epsilon$$\Sigma$$\epsilon$

(b) Frame dependent alignment lattice with $T=3$.

$0,\epsilon$$1,a$$1,\epsilon$$1,b$$2,\epsilon$$2,a$$2,b$$3,\epsilon$$3,a$$3,b$$\boldsymbol{\epsilon}$$\boldsymbol{a}$$b$$\epsilon$$\boldsymbol{a}$$b$$\boldsymbol{\epsilon}$$a$$\boldsymbol{b}$$\epsilon$$a$$b$$\epsilon$$a$$\boldsymbol{b}$$\epsilon$$a$$\boldsymbol{b}$$\boldsymbol{\epsilon}$$a$$b$

(c) Recognition lattice: intersection of Figure 1(a) and Figure 1(b) with all
paths for output $ab$ highlighted. The weights (omitted here) are defined by
passing the states and labels to the weight function.

The core of the LAST library, class RecognitionLattice is a flax.linen module
[15] that combines these 3 components to produce the actual states and weights
for any recognition lattice induced from a sequence of input frames. The
shortest distance of either the entire RecognitionLattice, or its intersection
with an output string, can then be obtained through API methods for training
locally or globally normalized models. A low level state/arc-based API is also
provided for applications such as beam search decoding.

While class RecognitionLattice provides efficient high level operations
through a simple interface by hiding implementation details, the modelling
components are instead designed as easy to extend interfaces,

  * •

In addition to the FullNGram context dependency described in [6], there is
also a NextStateTable context dependency, enabling the use of any context
dependency topology with a look-up table.

  * •

Both FrameDependent and FrameLabelDependent alignment lattices are available.

  * •

Weight functions can be implemented as the combination of a WeightFnCacher
that performs precomputation of static data (such as context state
embeddings), and a WeightFn that actually produces arc weights.

WFSAs in LAST are not just defined with data (i.e. arrays of state/arc tables)
but also with code that describes higher level operations. This is a key
difference of LAST from prior art such as k2 [3]. For example,
ContextDependency provides ForwardReduce and BackwardBroadcast, simple yet
crucial primitives that RecognitionLattice uses to produce a highly vectorized
forward-backward implementation (Section LABEL:sec:random-access). Such higher
level APIs are designed to take advantage of special structures in certain
WFSAs to achieve higher efficiency with minimum interface complexity.

##  4 Challenges

Lattice | $|Q|$ | $|E|$
---|---|---
RNN-T, intersected | $O(TU)$ | $O(TU)$
GNAT, intersected | $O(TU)$ | $O(TU)$
GNAT, complete | $O(TC)$ | $O(TCV)$
Table 1: Lattice sizes. $T$: input length, $U$: output length, $C$: number of
context dependency states, $V$: vocabulary size.

The main challenge in scaling up differentiable WFSA algorithms for modern
accelerator architectures is memory: the memory is more limited in size
compared to CPU DRAM, and memory access can be much slower than compute.
Efficient memory use is thus behind many design details in LAST.

###  4.1 The Need for On-the-Fly Arc Weight Computation

Existing frameworks such as k2 [3] or GTN [4] store all the arc weights of a
WFSA in memory at once. Table 1 summarizes the number of states $|Q|$ and the
number arcs $|E|$ of various types of lattices. When the lattice is obtained
from an intersection with a reference output, $|Q|$ and $|E|$ are comparable.
Storing all the arc weights at once in memory is thus practical for training
locally normalized models which only involves the lattice after intersection.
However, $|E|$ can be much larger than $|Q|$ for a complete recognition
lattice needed in training globally normalized models, and storing all
$O(TCV)$ arc weights quickly become impractical.

Remedy The highly structured nature of ASR lattices is very suitable for on-
the-fly arc weight computation. The states are $(t,q)$ pairs of time $t$ and
additional state $q$. The arc weights can be computed in the increasing order
of $t$ during the topological visit in shortest distance.

###  4.2 Standard Automatic Differentiation Is Memory Inefficient under the
Log Semiring

[14] points out that automatic differentiation over shortest distance gives
rise to the classical forward-backward algorithm [13]. However [14] only
analyzed the algorithms under the real semiring. We shall show that under the
log semiring, naive automatic differentiation yields an algorithm that
requires $O(|E|+|Q|)$ space, whereas the forward-backward algorithm only
requires $O(|Q|)$. Notably, this space complexity holds when the arc weights
are either a part of the input (not included in space complexity), or computed
on-the-fly.

For an acyclic WFSA, the following forward weight recurrence computed for each
$q\in Q$ in a topological order gives the shortest distance $D=\alpha[f]$,

$\alpha[q]=\bigoplus_{\mathrm{next}[e]=q}\left(\alpha[\mathrm{prev}[e]]\otimes
w[e]\right).$ (2)

Automatic differentiation, a.k.a. back propagation [16], computes the partial
derivative $\frac{\partial D}{\partial w[e]}$by repeatedly computing
$\frac{\partial\alpha[\mathrm{next}[e]]}{\partial w[e]}$ and
$\frac{\partial\alpha[\mathrm{next}[e]]}{\partial\alpha[\mathrm{prev}[e]]}$ in
a reverse topological order and then applying the chain rule $\frac{\partial
D}{\partial w[e]}=\sum_{q\in Q}\frac{\partial
D}{\partial\alpha[q]}\frac{\partial\alpha[q]}{\partial w[e]}$.

Computing $\frac{\partial\alpha[\mathrm{next}[e]]}{\partial w[e]}$ requires
knowing the weights of arcs leading to $\mathrm{next}[e]$ and the forward
weights of the source states of these arcs. Because $w[e]$ is either a part of
the input, or computed on-the-fly, only $\alpha[q]$ is stored, taking space
$O(|Q|)$. The situation is different for
$\frac{\partial\alpha[\mathrm{next}[e]]}{\partial\alpha[\mathrm{prev}[e]]}$,
which we shall analyze separately for the real, log, and tropical semirings.
For simplicity, let $\Gamma[q]$ be the set of states with arcs leading to
state $q$, we denote the partial derivative term
$\frac{\partial\alpha[\mathrm{next}[e]]}{\partial\alpha[\mathrm{prev}[e]]}$
alternatively as $\frac{\partial\alpha[q]}{\partial\alpha[p]}$ for
$p\in\Gamma[q]$ (for states not in $\Gamma[q]$, the partial derivative is not
used in automatic differentiation).

####  4.2.1 Real semiring

In this case,
$\alpha[q]=\sum_{\mathrm{next}[e]=q}\alpha[\mathrm{prev}[e]]w[e]$, thus
$\frac{\partial\alpha[q]}{\partial\alpha[p]}=\sum_{\mathrm{prev[e]}=p\land\mathrm{next[e]}=q}w[e].$
Since we just need the arc weights to compute the relevant gradients, the
overall space complexity remains $O(|Q|)$.

####  4.2.2 Log semiring

In this case,
$\alpha[q]=\log\left(\sum_{\mathrm{next}[e]=q}\exp(\alpha[\mathrm{prev}[e]]+w[e])\right)$,
thus
$\frac{\partial\alpha[q]}{\partial\alpha[p]}=\sum_{\mathrm{prev[e]}=p\land\mathrm{next[e]}=q}\frac{\exp(\alpha[p]+w[e])}{\exp\alpha[q]}.$
Forward-backward utilizes the following fact to eliminate the need of storing
$\alpha[p]+w[e]$,

$\displaystyle\frac{\partial\alpha[q]}{\partial\alpha[p]}\frac{\partial\alpha[p]}{\partial\alpha[r]}$
$\displaystyle=$
$\displaystyle\left(\sum_{\mathrm{prev[e_{2}]}=p\land\mathrm{next[e_{2}]}=q}\frac{\exp(\alpha[p]+w[e_{2}])}{\exp\alpha[q]}\right)\times$
$\displaystyle\left(\sum_{\mathrm{prev[e_{1}]}=r\land\mathrm{next[e_{1}]}=p}\frac{\exp(\alpha[r]+w[e_{1}])}{\exp\alpha[p]}\right)$
$\displaystyle=$
$\displaystyle\sum_{\begin{subarray}{c}\mathrm{prev[e_{2}]}=p\land\mathrm{next[e_{2}]}=q\\\
\mathrm{prev[e_{1}]}=r\land\mathrm{next[e_{1}]}=p\end{subarray}}\frac{\exp(w[e_{2}]+w[e_{1}]+\alpha[r])}{\exp\alpha[q]}.$

Repeatedly chaining the partial derivatives from the final state $f$ to
$\mathrm{prev}[e]$ results in a summation over all the paths from
$\mathrm{prev}[e]$ to $f$, giving rise to Equation 1. To compute Equation 1,
the forward-backward algorithm stores the backward weights $\beta[q]$, the sum
of path weights from all states to the final state, in addition to
$\alpha[q]$, with the space complexity remaining $O(|Q|)$. In contrast, an
automatic differentiation engine can not perform the symbolic optimization
that cancels out $\exp\alpha[p]$ from the equation due to its mechanical
nature. It thus need to store all the intermediate $\alpha[p]+w[e]$ values
from the forward pass to use in the backward pass, leading to an overall
$O(|E|+|Q|)$ space complexity, regardless of whether the arc weights are part
of the input, or computed on-the-fly.

####  4.2.3 Tropical semiring

In this case,
$\alpha[q]=\max_{\mathrm{next}[e]=q}(\alpha[\mathrm{prev}[e]]+w[e])$, thus
(the sub-gradient)
$\frac{\partial\alpha[q]}{\partial\alpha[p]}=\mathbf{1}_{\alpha[p]+w[e]=\alpha[q]}.$
The vector-Jacobian-product of $\max$ is implemented by storing the index
$\arg\max_{p}\alpha[p]+w[e]$ for all $q$, taking up just $O(|Q|)$ space.

####  4.2.4 Remedy: Gradient Rematerialization

JAX offers a flexible gradient rematerialization (also called checkpointing
[17, 12]) mechanism to control what intermediate values can be saved during
the forward pass. Whatever not saved is recomputed during the backward pass,
thereby trading memory usage with compute. We can thus disallow any saving of
intermediate values within individual invocations of Equation 2 simply through
a jax.remat call. In Section 5, we compare gradient rematerialization with the
hand-written forward-backward algorithm.

##  5 Benchmarks

Vocabulary size |  32 (i.e. grapheme)
---|---
Context dependency |  FullNGram, context size 2, 1057 states
Alignment lattice |  FrameDependent
Weight function |  SharedEmb, 512 hidden units
Batch size |  16
Input length |  1024
Output length |  256
Table 2: Model configurations and input size Acclerator |  TPUv3 |  V100
---|---|---
Training |  Forward-backward |  0.92s / 217MB |  2.42s / 278MB
|  Auto-diff with remat |  1.12s / 219MB |  2.40s / 278MB
|  Auto-diff without remat |  OOM / $\geq$ 102G |  OOM / $\geq$ 68.5G
Inference |  0.41s / 171MB |  0.30s / 203MB
k2 Training |  |  1.69s / 11.9G
k2 Inference |  |  0.70s / 9.6G
Table 3: Benchmark results (speed and memory usage)

We benchmark the training and inference speed and memory usage on a single
Google TPUv3 core and a single nVidia V100 GPU. Table 2 outlines the
configurations of the globally normalized GNAT model and benchmark data. The
setup amounts to approximately 1 million parameters in the “decoder” component
of an encoder-decoder architecure. Each input in the batch induces a lattice
of about 1 million states and 33 million arcs. As a comparison, we implement
the same model in k2 [3]. Even with optimizations such as DenseFsaVec and
gradient rematerialization, our k2 implementation is not able to process batch
size 16 input without running out of memory. We thus split the input into
smaller batches (5+5+6) and accumulate gradients during training.

Table 3 lists the benchmark results. LAST only implements the “decoder”
component in an encoder-decoder architecture, we thus benchmark the gradient
computation from the loss value to the decoder parameters and the encoder
frames, using various strategies for computing the shortest distance gradients
as discussed in Section 4.2.4. Consistent with our analysis in Section 4.2,
the forward-backward algorithm uses far less memory than automatic
differentiation without gradient rematerialization. Because we ran out of
memory on both TPUv3 and V100 when trying to run automatic differentiation
without gradient rematerialization, we report the amount of projected memory
usage reported by the XLA compiler when it aborted the compilations. Forward-
backward and automatic differentiation with rematerialization use a similar
amount of memory. On TPUv3, forward-backward is 18% faster than automatic
differentiation with gradient rematerialization, whereas on GPU, there is no
significant speed difference. For inference speed, we benchmark shortest path
derived from differentiating shortest distance. Despite a drastic 2.4x
difference in training speed, the V100 GPU is faster at inference. These
differences in the speed are likely due to different hardware design trade-
offs.

The memory usage of our implementation in k2 shows that on-the-fly arc weight
computation can bring about significant memory savings. k2’s training speed is
faster than LAST because computing all arc weights at once affords more
opportunity for parallelization: we compute arc weights 1 frame at a time in
LAST, but 16 frames at a time in k2. Had we also computed arc weights for 1 or
2 frames at a time in k2, the training time would increase to 4.76s or 2.59s
respectively.

##  6 Conclusion

We presented the LAST library and discussed challenges in scaling up WFSA
algorithms in an automatic differentiation framework. These challenges led us
to design LAST differently from prior art. Our generally applicable techniques
to counter these challenges proved effective as shown in the benchmark
results.

## References

  * [1] Alex Graves, Santiago Fernández, Faustino Gomez, and Jürgen Schmidhuber,  “Connectionist temporal classification: labelling unsegmented sequence data with recurrent neural networks,”  in Proceedings of the 23rd international conference on Machine learning, 2006, pp. 369–376.
  * [2] Alex Graves,  “Sequence transduction with recurrent neural networks,”  arXiv preprint arXiv:1211.3711, 2012.
  * [3] Daniel Povey, Piotr Zelasko, and Sanjeev Khudanpur,  “Speech recognition with next-generation kaldi (k2, lhotse, icefall),”  Interspeech: tutorials, 2021.
  * [4] Awni Hannun, Vineel Pratap, Jacob Kahn, and Wei-Ning Hsu,  “Differentiable weighted finite-state transducers,”  arXiv preprint arXiv:2010.01003, 2020.
  * [5] Alexander Rush,  “Torch-struct: Deep structured prediction library,”  in Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics: System Demonstrations, Online, July 2020, pp. 335–342, Association for Computational Linguistics.
  * [6] Ehsan Variani, Ke Wu, Michael Riley, David Rybach, Matt Shannon, and Cyril Allauzen,  “Global normalization for streaming speech recognition in a modular framework,”  Advances in neural information processing systems, to appear, 2022\.
  * [7] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang,  “JAX: composable transformations of Python+NumPy programs,” 2018\.
  * [8] Mehryar Mohri, Fernando Pereira, and Michael Riley,  “Speech recognition with weighted finite-state transducers,”  in Handbook of Speech Processing, Jacob Benesty, M. Sondhi, and Yiteng Huang, Eds., chapter 28, pp. 559–582. Springer, 2008.
  * [9] Mehryar Mohri,  “Weighted automata algorithms,”  in Handbook of Weighted Automata, Manfred Droste, Werner Kuich, and Heiko Vogler, Eds., pp. 213–254. Springer, 2009.
  * [10] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al.,  “Pytorch: An imperative style, high-performance deep learning library,”  Advances in neural information processing systems, vol. 32, 2019\.
  * [11] Seppo Linnainmaa,  “Taylor expansion of the accumulated rounding error,”  BIT Numerical Mathematics, vol. 16, no. 2, pp. 146–160, 1976.
  * [12] Andreas Griewank and Andrea Walther,  Evaluating derivatives: principles and techniques of algorithmic differentiation,  SIAM, 2008.
  * [13] Leonard E Baum et al.,  “An inequality and associated maximization technique in statistical estimation for probabilistic functions of markov processes,”  Inequalities, vol. 3, no. 1, pp. 1–8, 1972.
  * [14] Jason Eisner,  “Inside-outside and forward-backward algorithms are just backprop (tutorial paper),”  in Proceedings of the Workshop on Structured Prediction for NLP, 2016, pp. 1–17.
  * [15] The Flax authors,  “Flax: A neural network library and ecosystem for jax designed for flexibility,” .
  * [16] David E Rumelhart, Geoffrey E Hinton, and Ronald J Williams,  “Learning representations by back-propagating errors,”  nature, vol. 323, no. 6088, pp. 533–536, 1986.
  * [17] Charles H Bennett,  “Logical reversibility of computation,”  IBM journal of Research and Development, vol. 17, no. 6, pp. 525–532, 1973.