# Tensor clustering fossils in modified gravity and high-redshift
gravitational-wave sound speed

Daisuke Yamauchi  d-yamauchi”at”ous.ac.jp Department of Physics, Faculty of
Science, Okayama University of Science, 1-1 Ridaicho, Okayama, 700-0005, Japan

###### Abstract

We investigate the tensor clustering fossils as a possible probe to constrain
the theory of gravity, in particular the deviation of the sound speed of
gravitational waves from the speed of light at high redshifts. We develop the
formalism of the effective Poisson equation to include the novel
phenomenological model of the scalar-tensor tidal interactions that are
expected to be induced by the modification of the theory of gravity. We show
that the tensor clustering fossils can arise from the propagation of
gravitational waves, the growth of the large-scale structures, and the second-
order contributions from the effective Poisson equation. We construct the
small-scale effective Lagrangian from the Horndeski scalar-tensor theory and
derive the formula applicable to the tensor clustering fossils in the language
of the effective field theory of dark energy. As a demonstration, we apply the
formalism to the constraint on the sound speed of gravitational waves in the
futuristic survey.

##  I Introduction

The clustering of large-scale structure has been intensively studied using the
two-point correlation function in real space, or the power spectrum in Fourier
space, which are determined under the assumption of statistical anisotropy.
When large-scale tensor perturbations (gravitational waves: GWs) couple to the
scalar sector, they tidally imprint local anisotropy in the small-scale
distribution of matter and the effect appears as local deviations from
statistical anisotropy Masui:2010cz ; Jeong:2012df ; Dai:2013kra ;
Schmidt:2013gwa ; Masui:2017fzw . Interestingly, this imprint constitutes a
fossilized map of the primordial tensor field since it persists even after the
tensor mode has decayed by redshifting. This fossil effect can arise from the
coupling between the scalar and tensor perturbations not only during the
inflationary epoch Dimastrogiovanni:2014ina ; Dimastrogiovanni:2015pla ;
Dimastrogiovanni:2019bfl  but also during the post-inflationary epoch
Dai:2013kra . It can be shown that the detection of primordial GWs from the
tensor clustering fossils in the standard cosmological model requires
measurements with a huge dynamic range, which is probably beyond the reach of
current galaxy surveys, but such measurements would be achievable using the
21-cm line probes of neutral hydrogen during the dark ages Masui:2010cz ;
Jeong:2012df . There are many projects to measure the 21-cm line in the dark
ages from the far side of the Moon or from a satellite orbiting around the
Moon, such as FARSIDE Burns:2021pkx , DAPPER Burns:2021ndk , NCLE
2020AdSpR..65..856B , LCRT LCRT , DSL Chen:2019xvd , TREED TREED , TSUKUYOMI
TSUKUYOMI  and so on.

In order to accurately evaluate the tensor clustering fossils, we must
correctly take into account the nonlinear scalar-tensor tidal interactions
induced during the late-time Universe. Even if the initial density
fluctuations are well described by linear theory, the nonlinearity of the
gravitational dynamics eventually dominates the large-scale structure and the
nonlinear mode coupling is naturally induced even during the post-inflationary
evolution, since the continuity and Euler equations have nonlinear terms.
Furthermore, the modification of the underlying theory of gravity alters the
nonlinear growth history of structure through the Poisson equation. The quasi-
nonlinear growth of large-scale structures can provide new insights into the
theory of gravity that would not be imprinted in the linear perturbation
theory Yamauchi:2017ibz ; Yamauchi:2022fss ; Yamashita:2024mdn ;
Sugiyama:2023tes ; Sugiyama:2023zvd ; Yamauchi:2021nxw ; Arai:2022ilw .
Although the detailed predictions depend on models of modified gravity, we can
introduce phenomenological functions in the Poisson equation to characterize
the deviations from the standard cosmological model. The phenomenological
model allows us to discuss general constraints on the modifications of gravity
independently of the details of the gravity model.

One of the most stringent constraints on the theory of gravity is obtained
from the nearly simultaneous detection of the GW event GW170817 and its
optical counterpart GRB170817A, which gave the constraint that the sound speed
of GWs propagating from a neutron star binary should not deviate from that of
light TheLIGOScientific:2017qsa ; Monitor:2017mdv ; GBM:2017lvd . This
observation contradicts any extension of general relativity that predicts a
large deviation of the GW sound speed and implies that the sound speed of GWs
is strictly coincident with the speed of light Ezquiaga:2017ekz ;
Baker:2017hug ; Sakstein:2017xjx ; Creminelli:2017sry ; Langlois:2017dyl .
However, the propagation of GWs from GW170817 as well as the local test of
gravity can place any constraints on the theory of gravity at the relatively
low redshift, $z\lesssim 0.01$. There is still a large viability of theory of
gravity whose deviation from general relativity becomes apparent at high
redshifts (see e.g., TerenteDiaz:2023iqk ; Lee:2022cyh ).

In this paper, we consider the tensor clustering fossils as a probe of the
theory of gravity, in particular the deviation of the GW sound speed from the
speed of light at high redshifts. To do this, we extend the formulation of the
tensor clustering fossils to include the phenomenological model of the scalar-
tensor tidal interactions in the effective Poisson equation. We then consider
the specific theory of gravity, which naturally induces the deviation of the
GW sound speed from the speed of light, and explore the nature of GWs at high
redshifts by using the resulting formula of the tensor clustering fossils.

This paper is organized as follows. In Sec. II, we derive the formula of the
tensor clustering fossils, which can be applied even in the situation that the
nonlinear effective Poisson equation is modified. In Sec. III, we compute the
tensor clustering fossils in the modified theory of gravity, in particular the
Horndeski scalar-tensor theory. We construct the small-scale effective
Lagrangian with the nonlinear mode-coupling between the scalar and tensor
modes and show the explicit expression of the tensor clustering fossils
induced during the post-inflationary evolution. We then qualitatively estimate
the impact of the deviation of the GW speed from the speed of light by using
the resulting formula of the tensor clustering fossils. Finally, Sec. IV is
devoted to the summary and discussion.

##  II Effective Poisson equation and tensor clustering fossils

In this section, we will derive the general expression for the tensor
clustering fossils to include the effect of the nonlinear mode-coupling
between the scalar and tensor modes in the late-time Universe. In particular,
we extend the analysis of Ref. Dai:2013kra  by incorporating the effect of the
phenomenological model in the effective Poisson equation. Throughout this
paper, we consider a metric with a Friedmann-Lemaître-Robertson-Walker (FLRW)
background in the Newton gauge, which is given by

$\displaystyle\mathrm{d}s^{2}=-\Bigl{[}1+2\Phi(t,{\bm{x}})\Bigr{]}\mathrm{d}t^{2}+a^{2}(t)\biggl{\\{}\Bigl{[}1-2\Psi(t,{\bm{x}})\Bigr{]}\delta_{ij}+h_{ij}(t,{\bm{x}})\biggr{\\}}\mathrm{d}x^{i}\mathrm{d}x^{j}\,.$
(1)

Here, $h_{ij}$ denotes the transverse-traceless tensor mode. We treat scalar
and matter perturbations as independent small parameters from tensor
perturbations, and we only keep their cross terms at second order. Matter
perturbations including the matter density contrast $\delta$ and peculiar
velocity $v_{i}$ grow from the primordial scalar perturbations. Assuming the
pressure and anisotropic stress are neglected, and the matter is minimally
coupled to gravity, the energy-momentum for the pressureless nonrelativistic
matter is given by

$\displaystyle T^{0}{}_{0}=-\rho_{\rm
m}(t)\Bigl{[}1+\delta(t,{\bm{x}})\Bigr{]}\,,$ (2) $\displaystyle
T^{0}{}_{i}=a(t)\rho_{\rm m}(t)v_{i}(t,{\bm{x}})\,,$ (3) $\displaystyle
T^{i}{}_{0}=-\frac{1}{a(t)}\rho_{\rm
m}(t)\Bigl{[}\delta^{ij}-h^{ij}(t,{\bm{x}})\Bigr{]}v_{j}(t,{\bm{x}})\,,$ (4)
$\displaystyle T^{i}{}_{j}=0\,.$ (5)

To simplify the analysis, throughout this paper, we apply the quasi-static
approximation. In particular, the wavenumbers of the scalar and tensor
perturbations, which represent $k$ and $K$ respectively, are much larger than
the comoving Hubble length scale. The continuity and Euler equations for the
pressureless nonrelativistic matter are given by

$\displaystyle\dot{\delta}+\frac{1}{a}\left(\delta^{ij}-h^{ij}\right)\partial_{i}v_{j}=0\,,$
(6) $\displaystyle\dot{v}_{i}+Hv_{i}+\frac{1}{a}\partial_{i}\Phi=0\,,$ (7)

where the dot denotes the derivative with respect to the cosmic time $t$.
Although these fluid equations are the same as the standard ones in general
relativity, there appears the effect of modification of the theory of gravity
through the gravitational potential $\Phi$. To obtain the closed-form equation
for the matter perturbations, we need to add the information for the relation
between the gravitational potential and the matter density contrast, i.e., the
Poisson equation. In this paper, we formally consider the following form of
the effective Poisson equation 111 The effective Poisson equation may include
terms such as $\dot{\delta}$ and $\ddot{\delta}$. When the Degenerate Higher-
Order Scalar-Tensor (DHOST) theory of gravity Langlois:2015cwa ;
Crisostomi:2016czh ; Achour:2016rkg ; BenAchour:2016fzp ; Langlois:2017dyl
beyond the Horndeski theory Horndeski:1974wa ; Deffayet:2011gz ;
Kobayashi:2011nu  is considered as gravity theory governing the structure
formation, there can be shown to appear such time-derivatives of density
contrast Hirano:2019nkz ; Hirano:2020dom . In this paper, we simply neglect
such terms for simplicity. :

$\displaystyle\frac{1}{a^{2}H^{2}}\partial^{2}\Phi(t,{\bm{x}})=\frac{3}{2}\Omega_{\rm
m}(t)\mu(t)\delta(t,{\bm{x}})+\left(-\Gamma_{1}(t)\frac{1}{a^{2}H^{2}}\partial^{2}h_{ij}(t,{\bm{x}})+\Gamma_{2}(t)\frac{1}{H}\dot{h}_{ij}(t,{\bm{x}})\right)\frac{\partial_{i}\partial_{j}}{\partial^{2}}\delta(t,{\bm{x}})+\cdots\,,$
(8)

where $\Omega_{\rm m}$ denotes the energy fraction of the matter field, and we
have introduced the time-dependent phenomenological functions $\Gamma_{1}$ and
$\Gamma_{2}$, which are related to the underlying gravity model. When general
relativity is considered, the second-order contribution to the effective
Poisson equation (8) is given by
$h^{ij}(\partial_{i}\partial_{j}/\partial^{2})\delta$. Since this term is
suppressed by the factor $a^{2}H^{2}/K^{2}(\ll 1)$ compared with the terms
appearing in Eq. (8) in the subhorizon limit, we can drop such term. In the
subsequent analysis, we will show such contribution can be safely neglected
when one computes the tensor clustering fossils. Therefore, in the case of the
standard cosmology with general relativity, we take $\Gamma_{1}=\Gamma_{2}=0$.
The nonvanishing phenomenological functions $\Gamma_{1}$ and $\Gamma_{2}$ can
be treated as the signals of the deviation from general relativity.

In this paper, we will solve these equations of motion by using the perturbed
expansion defined as

$\displaystyle\delta=\sum_{n=1}\delta^{(n)}\,,\ \ \
\Phi=\sum_{n=1}\Phi^{(n)}\,,\ \ \ \cdots\,,$ (9)

where $\delta^{(n)}$ , $\Phi^{(n)}$ , $\cdots={\cal O}(\epsilon^{n})$ are the
$n$-th order quantities with $\delta^{(1)}$ being ${\cal O}(\epsilon)$
quantity. At linear order, the closed-form equation for the density contrast
is given by

$\displaystyle\ddot{\delta}^{(1)}+2H\dot{\delta}^{(1)}-\frac{3}{2}H^{2}\Omega_{\rm
m}\mu\delta^{(1)}=0\,,$ (10)

When the time dependence of the effective gravitational coupling $\mu$ is
given, one can numerically solve Eq. (10). In general, Eq. (10) has two
independent solutions, which we call $D_{\pm}(t)$. We take the boundary
condition of the growing mode as $\delta^{(1)}(t,{\bm{x}})=D_{+}(t)\delta_{\rm
prim}({\bm{x}})$, where $\delta_{\rm prim}({\bm{x}})$ denotes the primordial
density contrast defined at the very early epoch. The linear peculiar velocity
field $v_{i}^{(1)}$ can be described in terms of the linear density contrast
through Eq. (6) as
$v_{i}^{(1)}=-(\partial_{i}/\partial^{2})\dot{\delta}^{(1)}=-Hf(\partial_{i}/\partial^{2})\delta^{(1)}$,
where $f=\mathrm{d}\ln D_{+}/\mathrm{d}\ln a$ denotes the growth rate of the
density contrast. On the other hand, as for the tensor mode, we assume the
form of the evolution equation as

$\displaystyle\ddot{h}_{ij}+\left(3+\nu\right)H\dot{h}_{ij}-c_{\rm
T}^{2}\frac{\partial^{2}}{a^{2}}h_{ij}=0\,,$ (11)

where $\nu$ and $c_{\rm T}^{2}$ represent the time-dependent phenomenological
functions characterizing the properties of the large-scale tensor mode.

Let us consider the second-order density contrast induced by the scalar-tensor
tidal interactions. Substituting the first-order solutions of the scalar and
tensor modes into Eqs. (6)–(7) as the source terms, we obtain the second-order
equation as

$\displaystyle\ddot{\delta}^{(2)}+2H\dot{\delta}^{(2)}-\frac{\partial^{2}}{a^{2}}\Phi^{(2)}=-Hf\dot{h}^{ij}\frac{\partial_{i}\partial_{j}}{\partial^{2}}\delta^{(1)}+\frac{3}{2}H^{2}\Omega_{\rm
m}\mu\,h^{ij}\frac{\partial_{i}\partial_{j}}{\partial^{2}}\delta^{(1)}\,.$
(12)

Since the first and second terms in the right hand side of the above equation
are estimated as ${\cal O}(H\dot{h}\delta)\sim{\cal O}((K^{2}/a^{2})h\delta)$
and ${\cal O}(H^{2}h\delta)$ respectively, the first term is ${\cal
O}(K^{2}/a^{2}H^{2})$ larger than the second term in the subhorizon limit. We
therefore neglect the second term in the right hand side of Eq. (12). We then
use the second-order effective Poisson equation (8) to obtain the closed-form
equation for the second-order density contrast as

$\displaystyle\ddot{\delta}^{(2)}+2H\dot{\delta}^{(2)}-\frac{3}{2}H^{2}\Omega_{\rm
m}\mu\delta^{(2)}=\Bigl{\\{}-\Gamma_{1}\frac{1}{a^{2}}\partial^{2}h^{ij}-\left(f-\Gamma_{2}\right)H\dot{h}^{ij}\,\Bigr{\\}}\frac{\partial_{i}\partial_{j}}{\partial^{2}}\delta^{(1)}\,.$
(13)

Introducing the following quantity:

$\displaystyle\widetilde{\delta}^{(2)}=\delta^{(2)}-h^{ij}\frac{\partial_{i}\partial_{j}}{\partial^{2}}\delta^{(1)}\,,$
(14)

and using Eq. (11), we can rewrite Eq. (15) as

$\displaystyle\ddot{\widetilde{\delta}}{}^{(2)}+2H\dot{\widetilde{\delta}}{}^{(2)}-\frac{3}{2}H^{2}\Omega_{\rm
m}\mu\,\widetilde{\delta}^{(2)}=\Bigl{\\{}\left(c_{\rm
T}^{2}-\Gamma_{1}\right)\frac{1}{a^{2}}\partial^{2}h^{ij}+\left(f-1-\nu+\Gamma_{2}\right)H\dot{h}^{ij}\,\Bigr{\\}}\frac{\partial_{i}\partial_{j}}{\partial^{2}}\delta^{(1)}\,,$
(15)

where we have dropped the sub-leading terms. The solution of the above
equation can be divided into two pieces:
$\widetilde{\delta}^{(2)}(t,{\bm{x}})=\delta^{(2)}_{\rm
homo}(t,{\bm{x}})+\delta^{(2)}_{\rm spec}(t,{\bm{x}})$, where
$\delta^{(2)}_{\rm homo}$ denotes the solution of the homogeneous part of Eq.
(15), and $\delta^{(2)}_{\rm spec}$ is a special solution that solves the full
equation with the vanishing boundary condition at the very early epoch. The
special solution can be constructed by using the Green function method. Using
the conformal time defined as $\eta=\int^{t}\mathrm{d}\bar{t}/a(\bar{t})$ and
transforming to Fourier space, we obtain the special solution of the form

$\displaystyle\delta_{\rm spec}^{(2)}(\eta,{\bm{k}})=-h_{ij}^{\rm
prim}({\bm{K}})\widehat{k}^{i}\widehat{k}^{j}\delta^{(1)}(\eta,{\bm{k}}){\cal
S}(\eta,K)\,,$ (16)

where the prime denotes the derivative with respect to the conformal time,
$\widehat{k}^{i}\equiv k^{i}/k$, and ${\cal S}$ is given by

$\displaystyle{\cal
S}(\eta,K)=\int_{0}^{\eta}\mathrm{d}\bar{\eta}\biggl{\\{}\Bigl{[}c_{\rm
T}^{2}(\bar{\eta})-\Gamma_{1}(\bar{\eta})\Bigr{]}K^{2}{\cal
T}_{h}(\bar{\eta},K)-\Bigl{[}f(\bar{\eta})-1-\nu(\bar{\eta})-\Gamma_{2}(\bar{\eta})\Bigr{]}{\cal
H}(\bar{\eta}){\cal
T}_{h}^{\prime}(\bar{\eta},K)\biggr{\\}}\frac{D_{+}(\bar{\eta})}{D_{+}(\eta)}G_{\rm
ret}(\eta,\bar{\eta})\,,$ (17)

with ${\cal H}=aH$. We have introduced the transfer function of the tensor
mode as $h_{ij}(\eta,{\bm{K}})={\cal T}_{h}(\eta,K)h_{ij}^{\rm
prim}({\bm{K}})$ with $h_{ij}^{\rm prim}({\bm{K}})$ being the primordial
tensor mode. Here, $G_{\rm ret}$ denotes the retarded Green function which can
be written in terms of the two independent solutions of the homogeneous part
as

$\displaystyle G_{\rm
ret}(\eta,\bar{\eta})=\frac{D_{+}(\eta)D_{-}(\bar{\eta})-D_{+}(\bar{\eta})D_{-}(\eta)}{D_{+}^{\prime}(\bar{\eta})D_{-}(\bar{\eta})-D_{+}(\bar{\eta})D_{-}^{\prime}(\bar{\eta})}\Theta(\eta-\bar{\eta})\,,$
(18)

where $\Theta(x)$ denotes the step function. Combining the resultant
expressions derived in this section, the solution of the matter density
contrast valid up to the second order can be expressed as

$\displaystyle\delta(\eta,{\bm{k}})=\delta^{(1)}(\eta,{\bm{k}})-h_{ij}^{\rm
prim}({\bm{K}})\widehat{k}^{i}\widehat{k}^{j}\biggl{[}\frac{1}{2}\frac{\mathrm{d}\ln{\cal
T}_{\delta}}{\mathrm{d}\ln k}{\cal T}_{h}(\eta,K)+{\cal
S}(\eta,K)\biggr{]}\delta^{(1)}(\eta,{\bm{k}})\,.$ (19)

where we have introduced the linear transfer function for the density
contrast, defined as $\delta^{(1)}(\eta,{\bm{k}})={\cal
T}_{\delta}(\eta,k)\Phi_{\rm prim}({\bm{k}})$. In our prescription,
$\mathrm{d}\ln{\cal T}_{\delta}/\mathrm{d}\ln k=2$ in the subhorizon limit. As
we will show later, even in the case of the modified gravity, ${\cal S}$ is
expected to persist after the tensor mode has decayed. Therefore, we still
refer to this effect as tensor clustering fossils. We find from Eq. (17) that
the tensor clustering fossils can generally be induced by three parts: the
properties of the large-scale GWs ($c_{\rm T}^{2}$ and $\nu$), the deviation
of the structure growth rate from unity ($f-1$), and the second-order scalar-
tensor tidal interactions from the effective Poisson equations ($\Gamma_{1}$
and $\Gamma_{2}$). These effects are expected to increase the amplitude of the
tensor clustering fossils. Even when the contributions from $(f-1)$ and
$\Gamma_{1,2}$ can be neglected, the modification of the tensor clustering
fossils still exists. In other words, we can use the tensor clustering fossils
(17) to constrain the properties of large-scale GWs. One can easily find that
if we consider general relativity and matter-dominated era, namely
$D_{+}\propto\eta^{2}$ , $D_{-}\propto 1/\eta^{3}$ , $c_{\rm T}^{2}=1$ , and
$\nu=\Gamma_{1}=\Gamma_{2}=0$, Eq. (17) reduces to the standard form of the
tensor clustering fossils, which was previously derived in Ref. Dai:2013kra .

We further assume that the primordial fields are generated through the
standard inflationary scenario, in which the primordial scalar-scalar-tensor
bispectrum satisfies the consistency relation in the squeezed limit. In the
limit of $K\ll k$, the primordial correlations between large-scale tensor and
small-scale scalar modes lead to an anisotropic primordial scalar power
spectrum as Giddings:2010nc ; Giddings:2011zd

$\displaystyle\widetilde{P}_{\Phi}({\bm{k}})\Bigl{|}_{h_{ij}({\bm{K}})}=P_{\Phi}(k)\biggl{[}1-\frac{1}{2}h_{ij}^{\rm
prim}\widehat{k}^{i}\widehat{k}^{j}\frac{\mathrm{d}\ln P_{\Phi}}{\mathrm{d}\ln
k}\biggr{]}\,.$ (20)

Here, $|_{h_{ij}({\bm{K}})}$ represents an ensemble average over realizations
of the scalar perturbations while holding the tensor perturbations constant.
We then obtain the anisotropic power spectrum for matter perturbations
Masui:2010cz ; Jeong:2012df ; Dai:2013kra ; Schmidt:2013gwa ; Masui:2017fzw :

$\displaystyle\widetilde{P}(\eta,{\bm{k}})\Bigl{|}_{h_{ij}({\bm{K}})}=P(\eta,k)\biggl{\\{}1-h_{ij}^{\rm
prim}({\bm{K}})\widehat{k}^{i}\widehat{k}^{j}\biggl{[}\frac{1}{2}\frac{\mathrm{d}\ln
P}{\mathrm{d}\ln k}+2{\cal S}_{N}(\eta,K)\biggr{]}\biggr{\\}}\,,$ (21)

where

$\displaystyle{\cal S}_{N}(\eta,K)={\cal
S}(\eta,K)-\frac{1}{2}\frac{\mathrm{d}\ln{\cal T}_{\delta}}{\mathrm{d}\ln
k}\Bigl{[}1-{\cal T}_{h}(\eta,K)\Bigr{]}\,.$ (22)

Once the underlying gravity model is specified, we can calculate the tensor
clustering fossils by using the formula derived above. In the subsequent
analysis, we focus on the Horndeski class of modified gravity which naturally
induces the deviation of the sound speed of GWs from the speed of light.

##  III Tensor clustering fossils in modified gravity

###  III.1 Horndeski scalar-tensor theory

In this section, let us consider the Horndeski scalar-tensor theory
Horndeski:1974wa ; Deffayet:2011gz ; Kobayashi:2011nu  (see Langlois:2018dxi ;
Kobayashi:2019hrl ; Arai:2022ilw  for a review) as the theory of gravity
governing the clustering of the matter density contrast. The Horndeski theory
is a most general scalar-tensor gravity theory with the second-order
equations-of-motion, to which attention is paid as one of the attractive
modified gravity theories. The Lagrangian of the Horndeski theory is given by

$\displaystyle{\cal L}_{\rm
grav}=\sqrt{-g}\biggl{\\{}G_{2}(\phi,X)-G_{3}(\phi,X)\Box\phi+G_{4}(\phi,X)R+G_{4X}\Bigl{[}(\Box\phi)^{2}-\nabla^{\mu}\nabla^{\nu}\phi\nabla_{\mu}\nabla_{\nu}\phi\Bigr{]}\biggr{\\}}\,,$
(23)

where $G_{2}$ , $G_{3}$ , and $G_{4}$ are arbitrary functions of the scalar
field $\phi$ and its kinetic term
$X:=-g^{\mu\nu}\nabla_{\mu}\phi\nabla_{\nu}\phi/2$, and we have introduced the
notation: $g_{\phi}=\partial g/\partial\phi$, $g_{X}:=\partial g/\partial X$,
and so on. A number of scalar-tensor theories such as the Jordan-Brans-Dicke
theory Brans:1961sx , the $f(R)$ gravity Starobinsky:1980te , the $k$-essence
Chiba:1999ka , the kinetic gravity braiding Deffayet:2010qz , and the non-
minimal coupling to the Gauss-Bonnet term Kobayashi:2011nu  belong to the
Horndeski family. In particular, the Horndeski theory naturally includes the
models with the deviation of the GW sound speed from the speed of light. We
note that we have simply dropped the “$G_{5}$” terms in Eq. (23), but the
generic features presented in the subsequent analysis are expected to be the
same, and a wider class of modified gravity such as DHOST theory
Langlois:2015cwa ; Crisostomi:2016czh ; Achour:2016rkg ; BenAchour:2016fzp ;
Langlois:2017dyl  would be straightforward.

To describe cosmological perturbations around the FLRW background in theories
with a preferred slicing induced by a time-dependent scalar field, it is
convenient to use the language of the effective field theory (EFT) of dark
energy. In this context, instead of considering the explicit form of the
Horndeski functions $G_{a}$, we will introduce the EFT parameters to determine
completely the dynamics of cosmological perturbations. Moreover, we focus only
on the operators that contribute in the quasi-static limit. In this limit, the
minimum set to specify the total amount of cosmological information up to the
linear order is a set of three independent functions of time that can be
labelled by $\\{\alpha_{\rm T},\alpha_{\rm M},\alpha_{\rm B}\\}$ in addition
to the Hubble parameter $H$ and the effective Planck mass $M$. The EFT
parameters are related explicitly to the Horndeski functions through
Bellini:2014fua

$\displaystyle M^{2}=2G_{4}-4XG_{4X}\,,$ (24) $\displaystyle MH^{2}\alpha_{\rm
M}=\frac{\mathrm{d}}{\mathrm{d}t}M^{2}\,,$ (25) $\displaystyle
HM^{2}\alpha_{\rm B}=2\dot{\phi}\left(XG_{3X}-G_{4\phi}-2XG_{4\phi
X}\right)+8HX\left(G_{4X}+2XG_{4XX}\right)\,,$ (26) $\displaystyle
M^{2}\alpha_{\rm T}=4XG_{4X}\,.$ (27)

To generalize this formalism to the second-order perturbations in the
subhorizon scales, one other function of time is required Yamauchi:2017ibz ;
Bellini:2015wfa . The explicit form is given by

$\displaystyle M^{2}\alpha_{\rm V}=2X\left(G_{4X}+2XG_{4XX}\right)\,.$ (28)

Here, the right hand sides of these EFT parameters are evaluated at the
cosmological background.

For the purpose of studying the clustering fossils in the modified gravity, we
will derive the small-scale effective Lagrangian for the scalar and tensor
modes, following Refs. Kobayashi:2014ida ; Koyama:2013paa ; Hirano:2020dom ;
Yamauchi:2021nxw . Without loss of generality, we take
$\phi(t,{\bm{x}})=t+\pi(t,{\bm{x}})$. We now expand the full Horndeski
Lagrangian Eq. (23) around the FLRW background to obtain the effective
Lagrangian using the following assumptions: As for the scalar perturbations,
we employ the quasi-static approximation so that
$(\partial\epsilon)^{2}\gg(\dot{\epsilon})^{2}$ in the Lagrangian, where
$\epsilon$ stands for any of $\Phi$, $\Psi$ , and $H\pi$. On the other hand,
as for the tensor perturbations, its time derivative provides the same order
of magnitude as its spatial derivative. Hence, we will keep the scalar
perturbations with the highest spatial derivatives and the tensor
perturbations with all the time and spatial derivatives. Moreover, in this
paper, we only take into account the scalar-scalar-tensor three-point
interaction terms with the highest derivatives. By doing this expansion, we
schematically find the following terms:

$\displaystyle(\partial\epsilon)^{2}\,,\ \ (\partial h)^{2}\,,\ \
\dot{h}^{2}\,,\ \ \partial^{2}h(\partial\epsilon)^{2}\,\ \
\dot{h}(\partial\epsilon)^{2}\,,\ \ \ddot{h}(\partial\epsilon)^{2}\,.$ (29)

By explicitly expanding the full Horndeski Lagrangian Eq. (23) under these
prescriptions, we find the small-scale effective Lagrangian valid up to the
third order as

$\displaystyle{\cal L}_{\rm eff}=$
$\displaystyle\frac{M^{2}a}{2}\biggl{[}4\Psi\partial^{2}\Phi-2\left(1+\alpha_{\rm
T}\right)\Psi\partial^{2}\Psi-a^{2}(\dot{h}_{ij})^{2}+\left(1+\alpha_{\rm
T}\right)(\partial_{k}h_{ij})^{2}$
$\displaystyle\quad\quad+4H\left(\alpha_{\rm M}-\alpha_{\rm
T}\right)\Psi\partial^{2}\pi-4H\alpha_{\rm
B}\Phi\partial^{2}\pi+H^{2}c_{\pi\pi}\pi\partial^{2}\pi$
$\displaystyle\quad\quad+2\alpha_{\rm
T}\dot{h}^{ij}\Psi\partial_{i}\partial_{j}\pi-4\alpha_{\rm
V}\dot{h}^{ij}\Phi\partial_{i}\partial_{j}\pi$
$\displaystyle\quad\quad-\frac{5}{a^{2}}\alpha_{\rm
T}\partial^{2}h^{ij}\pi\partial_{i}\partial_{j}\pi+Hc_{\pi\pi\dot{h}}\dot{h}^{ij}\pi\partial_{i}\partial_{j}\pi-2\alpha_{\rm
V}\ddot{h}^{ij}\pi\partial_{i}\partial_{j}\pi\biggr{]}\,,$ (30)

where the coefficients $c_{\pi\pi}$ and $c_{\pi\pi\dot{h}}$ are defined as

$\displaystyle
c_{\pi\pi}=-2\biggl{\\{}\,\frac{\dot{H}}{H^{2}}+\frac{3}{2}\Omega_{\rm
m}+\alpha_{\rm T}-\alpha_{\rm M}+\frac{(aM^{2}H\alpha_{\rm
B})^{\cdot}}{aM^{2}H^{2}}\biggr{\\}}\,,$ (31) $\displaystyle
c_{\pi\pi\dot{h}}=2\biggl{\\{}\left(\alpha_{\rm T}-\alpha_{\rm M}-\alpha_{\rm
B}\right)-\frac{(aM^{2}\alpha_{\rm V})^{\cdot}}{aM^{2}H}\biggr{\\}}\,.$ (32)

Based on this effective Lagrangian, we then study the quasi-static behavior of
those perturbations deep inside the cosmological horizon. Varying the
Lagrangian Eq. (30) with respect to $\Phi$, $\Psi$, and $\pi$, we derive the
first-order equations of motion for the scalar perturbations as

$\displaystyle\partial^{2}\bigl{(}\Psi^{(1)}-\alpha_{\rm
B}H\pi^{(1)}\bigr{)}-\frac{3}{2}a^{2}H^{2}\Omega_{\rm m}\delta^{(1)}=0\,,$
(33) $\displaystyle\partial^{2}\Bigl{[}\Phi^{(1)}-(1+\alpha_{\rm
T})\Psi^{(1)}+(\alpha_{\rm M}-\alpha_{\rm T})H\pi^{(1)}\Bigr{]}=0\,,$ (34)
$\displaystyle\partial^{2}\biggl{[}(\alpha_{\rm M}-\alpha_{\rm
T})\Psi^{(1)}-\alpha_{\rm
B}\Phi^{(1)}+\frac{1}{2}c_{\pi\pi}H\pi^{(1)}\biggr{]}=0\,,$ (35)

which can be easily solved analytically to obtain the following form:

$\displaystyle\partial^{2}\epsilon^{(1)}=\frac{3}{2}a^{2}H^{2}\Omega_{\rm
m}\mu_{\epsilon}\delta^{(1)}\,,$ (36)

where $\epsilon=\Phi$ , $\Psi$ , $H\pi$ , and their coefficients are given by
(see also Pogosian:2016pwr ; Gleyzes:2015rua ; Hirano:2019nkz )

$\displaystyle\mu_{\Phi}=1+\alpha_{\rm T}+\frac{2\xi^{2}}{c_{\rm
s}^{2}\alpha}\,,$ (37) $\displaystyle\mu_{\Psi}=1+\frac{2\alpha_{\rm
B}\xi}{c_{\rm s}^{2}\alpha}\,,$ (38)
$\displaystyle\mu_{H\pi}=\frac{2\xi}{c_{\rm s}^{2}\alpha}\,,$ (39)

with $\xi=\alpha_{\rm B}(1+\alpha_{\rm T})+\alpha_{\rm T}-\alpha_{\rm M}$.
Here, we have introduced the following quantity:

$\displaystyle c_{\rm s}^{2}\alpha=$ $\displaystyle-2\biggl{\\{}(1+\alpha_{\rm
B})\biggl{[}\frac{\dot{H}}{H^{2}}+\alpha_{\rm B}(1+\alpha_{\rm T})+\alpha_{\rm
T}-\alpha_{\rm M}\biggr{]}+\frac{\dot{\alpha}_{\rm
B}}{H}+\frac{3}{2}\Omega_{\rm m}\biggr{\\}}\,,$ (40)

with $c_{\rm s}^{2}$ representing the propagation speed of the scalar mode and
$\alpha>0$. We note that $\alpha$ is also written in terms of other EFT
parameters, but we don’t show it explicitly. For the tensor perturbations, the
equation of motion can be derived by varying the effective Lagrangian with
respect to $h_{ij}$:

$\displaystyle\ddot{h}_{ij}+\left(3+\alpha_{\rm
M}\right)H\dot{h}_{ij}-\left(1+\alpha_{\rm
T}\right)\frac{\partial^{2}}{a^{2}}h_{ij}=0\,,$ (41)

which corresponds to $\nu=\alpha_{\rm M}$ and $c_{\rm T}^{2}=1+\alpha_{\rm T}$
in the context of the phenomenological model Eq. (11).

We next consider the second-order perturbations induced by the scalar-tensor
tidal interactions. To obtain the second-order solutions of the gravitational
potentials, we need to substitute the first-order solutions Eq. (36) into the
interaction terms between the scalar and tensor modes derived from Eq. (30).
The second-order equations of motion are given by

$\displaystyle\frac{2}{3a^{2}H^{2}\Omega_{\rm
m}}\partial^{2}\bigl{(}\Psi^{(2)}-\alpha_{\rm
B}H\pi^{(2)}\bigr{)}-\delta^{(2)}=\Pi_{\Phi}\frac{1}{H}\dot{h}^{ij}\frac{\partial_{i}\partial_{j}}{\partial^{2}}\delta^{(1)}\,,$
(42) $\displaystyle\frac{2}{3a^{2}H^{2}\Omega_{\rm
m}}\partial^{2}\Bigl{[}\Phi^{(2)}-(1+\alpha_{\rm T})\Psi^{(2)}+(\alpha_{\rm
M}-\alpha_{\rm
T})H\pi^{(2)}\Bigr{]}=\Pi_{\Psi}\frac{1}{H}\dot{h}^{ij}\frac{\partial_{i}\partial_{j}}{\partial^{2}}\delta^{(1)}\,,$
(43)

for the gravitational potentials, and

$\displaystyle\frac{2}{3a^{2}H^{2}\Omega_{\rm
m}}\partial^{2}\biggl{[}(\alpha_{\rm M}-\alpha_{\rm T})\Psi^{(2)}-\alpha_{\rm
B}\Phi^{(2)}+\frac{1}{2}c_{\pi\pi}H\pi^{(2)}\biggr{]}=\biggl{[}-\Pi_{\pi,1}\frac{1}{a^{2}H^{2}}\partial^{2}h^{ij}+\Pi_{\pi,2}\frac{1}{H}\dot{h}^{ij}\biggr{]}\frac{\partial_{i}\partial_{j}}{\partial^{2}}\delta^{(1)}\,,$
(44)

for the scalar field perturbations. The coefficients for the interaction terms
are

$\displaystyle\Pi_{\Phi}=\alpha_{\rm V}\mu_{H\pi}\,,$ (45)
$\displaystyle\Pi_{\Psi}=-\frac{1}{2}\alpha_{\rm T}\mu_{H\pi}\,,$ (46)
$\displaystyle\Pi_{\pi,1}=-\frac{1}{2}\Bigl{\\{}5\alpha_{\rm T}+\alpha_{\rm
V}\left(1+\alpha_{\rm T}\right)\Bigr{\\}}\mu_{H\pi}\,,$ (47)
$\displaystyle\Pi_{\pi,2}=\alpha_{\rm V}\mu_{\Phi}-\frac{1}{2}\alpha_{\rm
T}\mu_{\Psi}-\biggl{\\{}\frac{1}{2}c_{\pi\pi\dot{h}}+\alpha_{\rm
V}\left(3+\alpha_{\rm M}\right)\biggr{\\}}\mu_{H\pi}\,.$ (48)

Combining these equations, we finally obtain the solution of the second-order
gravitational potential induced by the scalar-tensor tidal interactions as

$\displaystyle\frac{1}{a^{2}H^{2}}\partial^{2}\Phi^{(2)}=\frac{3}{2}\Omega_{\rm
m}\mu_{\Phi}\delta^{(2)}+\frac{3}{2}\Omega_{\rm
m}\biggl{[}-\mu_{H\pi}\Pi_{\pi,1}\frac{1}{a^{2}H^{2}}\partial^{2}h^{ij}+\Bigl{\\{}\mu_{\Phi}\Pi_{\Phi}+\kappa_{\Psi}\Pi_{\Psi}+\mu_{H\pi}\Pi_{\pi,2}\Bigr{\\}}\frac{1}{H}\dot{h}^{ij}\biggr{]}\frac{\partial_{i}\partial_{j}}{\partial^{2}}\delta^{(1)}\,.$
(49)

In summary, the phenomenological functions $\mu$ , $\Gamma_{1}$ , and
$\Gamma_{2}$, which were defined in the effective Poisson equation (8), can be
written in terms of the EFT parameters as

$\displaystyle\mu=\mu_{\Phi}=1+\alpha_{\rm T}+\frac{2[\alpha_{\rm
B}(1+\alpha_{\rm T})+\alpha_{\rm T}-\alpha_{\rm M}]^{2}}{c_{\rm
s}^{2}\alpha}\,,$ (50) $\displaystyle\Gamma_{1}=\frac{3}{2}\Omega_{\rm
m}\,\mu_{H\pi}\Pi_{\pi,1}\,,$ (51)
$\displaystyle\Gamma_{2}=\frac{3}{2}\Omega_{\rm
m}\Bigl{\\{}\mu_{\Phi}\Pi_{\Phi}+\mu_{\Psi}\Pi_{\Psi}+\mu_{H\pi}\Pi_{\pi,2}\Bigr{\\}}\,,$
(52)

where $\mu_{i}$ and $\Pi_{i}$ were defined in Eqs. (37)–(39) and (45)–(48).
These expressions are one of the main results of this paper. Once the time
dependence of the EFT parameters is explicitly given, we can compute the
coefficients of the effective Poisson equation with the use of the above
expressions and the tensor clustering fossils by using the formula Eq. (17)
derived in the previous section.

###  III.2 $\alpha_{\rm T}$ model

In this subsection, we apply the resulting formula derived in the previous
subsection to the constraint on a specific parameter. In particular, we focus
on the deviation of the GW sound speed from the speed of light. To do this,
let us consider the specific setup such that all the EFT parameters except for
$\alpha_{\rm T}$ are taken to be zero and the Universe is dominated by the
matter, namely $\alpha_{\rm M}=\alpha_{\rm B}=\alpha_{\rm V}=0$ and
$\Omega_{\rm m}=1$. We call this parametrization “$\alpha_{\rm T}$ model”.
Under these assumptions, one finds that the phenomenological functions
characterizing the first-order effective Poisson equation [Eqs. (37)–(39)] in
the $\alpha_{\rm T}$ model reduce to

$\displaystyle\mu_{\Phi}=\mu_{\Psi}=-\mu_{H\pi}=1\,.$ (53)

This fact implies that the first-order equations for the scalar perturbations
are the same as the standard ones. Hence, the gravitational potentials
coincide with each other and these are conserved through matter domination,
i.e., $\Phi^{(1)}=\Psi^{(1)}=\Phi_{\rm prim}$, and the matter density contrast
grows in proportion to the scale factor, $D_{+}(t)\propto a(t)$.

To discuss the effect of $\alpha_{\rm T}$ in the context of the tensor
clustering fossils, we consider the second-order phenomenological functions
given in Eqs. (51) and (52). In the $\alpha_{\rm T}$ model, these functions
can be easily shown to be

$\displaystyle\Gamma_{1}=-\frac{15}{4}\alpha_{\rm T}\,,\ \ \ \Gamma_{2}=0\,.$
(54)

Substituting these into the general formula of the tensor fossils Eq. (17), we
have

$\displaystyle{\cal
S}(\eta,K)=\frac{1}{5}\int_{0}^{\eta}\mathrm{d}\bar{\eta}\left(1+\frac{19}{4}\alpha_{\rm
T}(\bar{\eta})\right)K^{2}\bar{\eta}\,{\cal T}_{h}(\bar{\eta},K;\alpha_{\rm
T})\Biggl{[}1-\left(\frac{\bar{\eta}}{\eta}\right)^{5}\,\Biggr{]}\,.$ (55)

This expression provides the generic result in the $\alpha_{\rm T}$ model. The
time-dependent parameter $\alpha_{\rm T}(\eta)$ affects not only the
coefficient but also the transfer function of the tensor mode. As we will
show, the modification of the transfer function provides non-negligible
contributions to the tensor clustering fossils. We note that the value of
$\alpha_{\rm T}$ should be chosen in the range of $-1\leq\alpha_{\rm T}\leq 0$
to avoid instability in the perturbation for the scalar and tensor modes.

Figure 1:  The amplitude of the tensor clustering fossils ${\cal S}$ as a
function of $K\eta$ while keeping $K\eta_{\ast}$ constant in the $\alpha_{\rm
T}$ model. We show the standard one ${\cal S}_{0}$ (black) and the correction
factor $|\Delta{\cal S}|$ with $K\eta_{\ast}=1$ (red), $2.7437$ (orange), $5$
(green), $10$ (blue), $30$ (purple), and $50$ (magenta). The solid and dashed
curves represent the positive and negative values, respectively.

To continue the analysis, we now consider the specific time-dependence of
$\alpha_{\rm T}$. As mentioned in Sec. I, the observation of GW170817
TheLIGOScientific:2017qsa ; Monitor:2017mdv ; GBM:2017lvd  as well as the
local test of gravity put a strong constraint on the GW sound speed at the
relatively low redshift. There is still a large viability of the theory of
gravity where the deviation of the GW sound speed from the speed of light
appears only at high redshifts. With this in mind, we focus on the case where
the deviation of the GW sound speed from the speed of light occurs only for a
very short period during the (high-redshift) matter-dominated phase.
Specifically, we consider the following form of $\alpha_{\rm T}$:

$\displaystyle\alpha_{\rm T}(\eta)=\alpha_{\rm T}^{\ast}\,\delta_{\rm
D}(\ln\eta-\ln\eta_{\ast})\,.$ (56)

Hereafter, we treat two constant model parameters $\alpha_{\rm T}^{\ast}$ and
$\eta_{\ast}$ as the characteristic amplitude and time of the modification. In
this setup, the equation of motion for the tensor mode in the Fourier space is
given by

$\displaystyle
h_{ij}^{\prime\prime}+\frac{4}{\eta}h_{ij}^{\prime}+\Bigl{[}1+\alpha_{\rm
T}^{\ast}\,\delta_{\rm D}(\ln\eta-\ln\eta_{\ast})\Bigr{]}K^{2}h_{ij}=0\,.$
(57)

Since for $\eta<\eta_{\ast}$ there are no modifications, the transfer function
of the tensor mode is described by the standard one, that is ${\cal
T}_{h}(K\eta)=3j_{1}(K\eta)/K\eta=:{\cal T}_{h,0}(K\eta)$, with $j_{n}(x)$ and
$y_{n}(x)$ being the spherical Bessel functions of the first and second kinds.
To obtain the solution at $\eta>\eta_{\ast}$, we need to solve the junction
conditions at $\eta=\eta_{\ast}$. We then find the transfer function of the
tensor mode valid even after $\eta>\eta_{\ast}$ as

$\displaystyle{\cal T}_{h}(\eta,K)={\cal T}_{h,0}(K\eta)+\alpha_{\rm
T}^{\ast}\,\Delta{\cal T}_{h}(\eta,K)\,\Theta(\eta-\eta_{\ast})\,,$ (58)

with

$\displaystyle\Delta{\cal
T}_{h}(\eta,K)=3(K\eta_{\ast})^{3}j_{1}(K\eta_{\ast})\left(y_{1}(K\eta_{\ast})\frac{j_{1}(K\eta)}{K\eta}-j_{1}(K\eta_{\ast})\frac{y_{1}(K\eta)}{K\eta}\right)\,.$
(59)

Substituting the transfer function Eq. (58) into Eq. (55), the tensor
clustering fossils with the instantaneous change of the GW sound speed in the
$\alpha_{\rm T}$ model can be explicitly expressed as

$\displaystyle{\cal S}(\eta,K)={\cal S}_{0}(K\eta)+\alpha_{\rm
T}^{\ast}\,\Delta{\cal S}(\eta,K)\,\Theta(\eta-\eta_{\ast})\,,$ (60)

where ${\cal S}_{0}$ represents the standard part, which is the same as that
in the standard cosmology with general relativity Dai:2013kra :

$\displaystyle{\cal
S}_{0}(x)=\frac{1}{5}\int_{0}^{x}\mathrm{d}\bar{x}\,\bar{x}\,{\cal
T}_{h,0}(\bar{x})\Biggl{[}1-\biggl{(}\frac{\bar{x}}{x}\biggr{)}^{5}\,\Biggr{]}\,,$
(61)

and $\Delta{\cal S}$ is the correction part defined as

$\displaystyle\Delta{\cal S}(\eta,K)=\frac{19}{20}(K\eta_{\ast})^{2}\,{\cal
T}_{h,0}(K\eta_{\ast})\Biggl{[}1-\biggl{(}\frac{\eta_{\ast}}{\eta}\biggr{)}^{5}\,\Biggr{]}+\frac{1}{5}K^{2}\int_{\eta_{\ast}}^{\eta}\mathrm{d}\bar{\eta}\,\bar{\eta}\,\Delta{\cal
T}_{h}(\bar{\eta},K)\Biggl{[}1-\biggl{(}\frac{\bar{\eta}}{\eta}\biggr{)}^{5}\,\Biggr{]}\,.$
(62)

To see the behaviour of ${\cal S}$ in the $\alpha_{\rm T}$ model, in Fig. 1 we
plot the function ${\cal S}_{0}$ (black) and $\Delta{\cal S}$ for
$K\eta_{\ast}=1$ (red), $2.7437$ (orange), $5$ (green), $10$ (blue), $30$
(purple), and $50$ (magenta) as a function of $K\eta$. We find from Fig. 1
that not only the standard part ${\cal S}_{0}$ but also the correction part
$\Delta{\cal S}$ persists after the tensor mode has decayed. To understand the
asymptotic behaviour of $\Delta{\cal S}$, taking the late-time limit while
keeping $K\eta_{\ast}$ constant, we obtain

$\displaystyle\Delta{\cal
S}(K\eta\to\infty)\to\frac{9}{4}K\eta_{\ast}\,j_{1}(K\eta_{\ast})\,.$ (63)

Moreover, the effect of the second-order contribution to the modified Poisson
equation, namely $\Gamma_{1}$, on the tensor clustering fossils is much larger
than other effects. This fact implies that to extract the information of the
modified GW sound speed in the context of the tensor clustering fossils, it is
important to consider the dynamics of the scalar and tensor modes
simultaneously.

Figure 2:  The smallest value of the tensor-to-scalar ratio $r$ detectable at
the 3$\sigma$ level as a function of the amplitude of the instantaneous change
in the GWs sound speed, $\alpha_{\rm T}^{\ast}$. We show the results for
survey parameters of $\\{z_{\rm obs},k_{\rm max},V_{\rm
sur}\\}=\\{10,10\,[h/{\rm Mpc}],200\,[({\rm Gpc}/h)^{3}]\\}$ and
$\\{50,200\,[h/{\rm Mpc}],1000\,[({\rm Gpc}/h)^{3}]\\}$ with $z_{\ast}=100$
(red solid), $200$ (green dotted), and $500$ (blue dot-dashed), respectively.

Before closing this section, we would like to discuss the detectability of the
instantaneous change in the sound speed of GWs. To do this, we quantify the
signal-to-noise ratio for the tensor power spectrum $P_{h}(K)$ estimated from
the tensor clustering fossils as Masui:2010cz

$\displaystyle\frac{S}{N}=\Biggl{[}V_{\rm
sur}\int\frac{\mathrm{d}^{3}{\bm{K}}}{(2\pi)^{3}}\left(\frac{P_{h}(K)}{P_{h}(K)+N_{h}(K)}\right)^{2}\Biggr{]}^{1/2}\,,$
(64)

where $V_{\rm sur}$ is the volume of the survey, and $N_{h}(K)$ is the tensor
noise power spectrum. In practice, the tensor mode can be reconstructed by
applying the quadratic estimators to the density field. Following Jeong:2012df
, we have the reconstructed tensor noise power spectrum $N_{h}$ as

$\displaystyle
N_{h}(K)=\Biggl{[}\frac{1}{8}\int\frac{\mathrm{d}^{3}{\bm{k}}}{(2\pi)^{3}}\frac{(1-\widehat{\bm{K}}\cdot\widehat{\bm{k}})^{2}P^{2}(k)}{2P^{\rm
tot}(k)P^{\rm
tot}(|{\bm{K}}-{\bm{k}}|)}\biggl{\\{}\frac{1}{2}\frac{\mathrm{d}\ln
P}{\mathrm{d}\ln k}+2{\cal S}_{N}(K)\biggr{\\}}^{2}\Biggr{]}^{-1}\,,$ (65)

where $P^{\rm tot}$ denotes the power spectrum of the matter field including
noise. We assume that the factor $\frac{1}{2}\frac{\mathrm{d}\ln
P}{\mathrm{d}\ln k}+2{\cal S}_{N}$ can be approximately treated as a constant
value for large $k$ and introduce the largest wavenumber $k_{\rm max}$ for
which the matter power spectrum can be measured with high signal-to-noise,
$P(k)/P^{\rm tot}(k)\approx 1$ for $k<k_{\rm max}$ and $P(k)/P^{\rm
tot}(k)\approx 0$ for $k>k_{\rm max}$. With these assumptions, the noise power
spectrum reduces to Masui:2017fzw

$\displaystyle N_{h}(K)\approx\frac{45(2\pi)^{2}}{k_{\rm
max}^{3}}\biggl{[}\frac{1}{2}\frac{\mathrm{d}\ln P}{\mathrm{d}\ln k}+2{\cal
S}_{N}(K)\biggr{]}^{-2}\,.$ (66)

We adopt the standard inflationary form for the tensor power spectrum as
$P_{h}(K)=2\pi^{2}rA_{\rm s}/K^{3}$, where $A_{\rm s}$ denotes the amplitude
of the primordial scalar perturbation. Once the survey parameters $\\{z_{\rm
obs},k_{\rm max},V_{\rm sur}\\}$ and model parameters $\\{\alpha_{\rm
T}^{\ast},\eta_{\ast}(z_{\ast})\\}$ are specified, we can evaluate the
significance of the detection by using these expressions. We set $A_{\rm
s}=2.12\times 10^{-9}$ and $\frac{\mathrm{d}\ln P}{\mathrm{d}\ln k}=-2.75$ as
the fiducial model. In Fig. 2, we show the projected detection limit of the
tensor-to-scalar ratio $r$ at the 3$\sigma$ level as a function of the
amplitude of the instantaneous change in GW sound speed, $\alpha_{\rm
T}^{\ast}$, for surveys with various parameters. One can easily find that when
$|\alpha_{\rm T}^{\ast}|$ is large, the detection limit of $r$ becomes small.
Assuming the survey with $z_{\rm obs}=10$ and $V_{\rm sur}=200\,[({\rm
Gpc}/h)^{3}]$ measuring scales down to $k_{\rm max}=10\,[h/{\rm Mpc}]$, we
find that if $r=10^{-2}$, $\alpha_{\rm T}^{\ast}<-0.17$, $-0.25$, and $-0.50$
for $z_{\ast}=100$, $200$, and $500$, respectively, could be detected. For the
futuristic survey that can be achieved by the 21-cm line survey of the dark
ages, we consider the survey parameters $z_{\rm obs}=50$, $k_{\rm
max}=200\,[h/{\rm Mpc}]$ and $V_{\rm sur}=1000\,[({\rm Gpc}/h)^{3}]$. In this
case, all the allowed parameter ranges of $\alpha_{\rm T}^{\ast}$ can be
reached even if $r=5\times 10^{-7}$.

##  IV Conclusion

We have studied the tensor clustering fossils as a possible probe of the
theory of gravity, in particular the properties of large-scale tensor
perturbations. We have considered the effect of the modification of the theory
of gravity on the tensor clustering fossils and applied the resulting
expressions to the model with the deviation of the sound speed of GWs from the
speed of light. Using the perturbative treatment of gravitational clustering
and introducing the effective description of the Poisson equation, which
includes the novel phenomenological model of the scalar-tensor tidal
interactions, we have reformulated the tensor clustering fossils without
assuming the explicit form of the gravity model [Eq. (17)]. We found that the
tensor clustering fossils can arise from the propagation of the large-scale
tensor mode, the growth of the large-scale structure, and the scalar-tensor
tidal interactions from the effective Poisson equation. This fact implies that
the tensor clustering fossils can be used to constrain the underlying theory
of gravity in the late-time Universe. To apply the derived expression, we
considered the Horndeski scalar-tensor theory as the underlying gravity model,
and we have constructed the small-scale effective Lagrangian with the tidal
interactions between the scalar and tensor modes in the language of the
effective field theory of dark energy [Eq. (30)]. On the basis of the
resulting effective Lagrangian, we have consistently solved the second-order
equations of motion for the gravitational potentials and have derived the
formula applicable to the tensor clustering fossils in the Horndeski theory
[Eqs. (50)–(52)]. Finally, we have applied the resulting formula to constrain
the sound speed of GWs as a demonstration. Considering the specific model with
the instantaneous change in the GW sound speed, we have discussed the
detectability of the tensor clustering fossils by using the quadratic
estimator of the density field. We found that the measurement with the large
survey volume, which can be achieved by the futuristic 21-cm line survey, can
detect the large-scale GWs with the modified sound speed.

In this paper, we have neglected several realistic effects such as the
redshift-space distortion and the nonlinear growth of the 21-cm fluctuations.
In particular, the nonlinear growth of structure affects the small-scale 21-cm
spectrum and the cutoff scale due to the nonlinear mode coupling should be
taken into account Yamauchi:2022fri . The large-scale tensor mode considered
in this paper can also contribute to the intrinsic alignment of galaxy shape
(see e.g. Schmidt:2012nw ; Akitsu:2022lkl ; Philcox:2023uor ). It would be
interesting to study the observational constraints on the properties of GWs
through such observations. These are left to be investigated in future work.

###### Acknowledgements.

We thank Tomoki Katayama for useful discussions. This work was supported by
JSPS KAKENHI Grant Number 22K03627, 23K25868.

## References

  * (1) K. W. Masui and U. L. Pen, Phys. Rev. Lett. 105, 161302 (2010) doi:10.1103/PhysRevLett.105.161302 [arXiv:1006.4181 [astro-ph.CO]].
  * (2) D. Jeong and M. Kamionkowski, Phys. Rev. Lett. 108, 251301 (2012) doi:10.1103/PhysRevLett.108.251301 [arXiv:1203.0302 [astro-ph.CO]].
  * (3) L. Dai, D. Jeong and M. Kamionkowski, Phys. Rev. D 88, no.4, 043507 (2013) doi:10.1103/PhysRevD.88.043507 [arXiv:1306.3985 [astro-ph.CO]].
  * (4) F. Schmidt, E. Pajer and M. Zaldarriaga, Phys. Rev. D 89, no.8, 083507 (2014) doi:10.1103/PhysRevD.89.083507 [arXiv:1312.5616 [astro-ph.CO]].
  * (5) K. W. Masui, U. L. Pen and N. Turok, Phys. Rev. Lett. 118, no.22, 221301 (2017) doi:10.1103/PhysRevLett.118.221301 [arXiv:1702.06552 [astro-ph.CO]].
  * (6) E. Dimastrogiovanni, M. Fasiello, D. Jeong and M. Kamionkowski, JCAP 12, 050 (2014) doi:10.1088/1475-7516/2014/12/050 [arXiv:1407.8204 [astro-ph.CO]].
  * (7) E. Dimastrogiovanni, M. Fasiello and M. Kamionkowski, JCAP 02, 017 (2016) doi:10.1088/1475-7516/2016/02/017 [arXiv:1504.05993 [astro-ph.CO]].
  * (8) E. Dimastrogiovanni, M. Fasiello and G. Tasinato, Phys. Rev. Lett. 124, no.6, 061302 (2020) doi:10.1103/PhysRevLett.124.061302 [arXiv:1906.07204 [astro-ph.CO]].
  * (9) J. Burns, G. Hallinan, T. C. Chang, M. Anderson, J. Bowman, R. Bradley, S. Furlanetto, A. Hegedus, J. Kasper and J. Kocz, et al. [arXiv:2103.08623 [astro-ph.IM]].
  * (10) J. Burns, S. Bale, R. Bradley, Z. Ahmed, S. W. Allen, J. Bowman, S. Furlanetto, R. MacDowall, J. Mirocha and B. Nhan, et al. [arXiv:2103.05085 [astro-ph.CO]].
  * (11) Bentum, M. J., Verma, M. K., Rajan, R. T., et al. 2020, Advances in Space Research, 65, 856. doi:10.1016/j.asr.2019.09.007
  * (12) S. Bandyopadhyay et al., ”Conceptual Design of the Lunar Crater Radio Telescope (LCRT) on the Far Side of the Moon,” 2021 IEEE Aerospace Conference (50100), 2021, pp. 1-25, doi: 10.1109/AERO50100.2021.9438165.
  * (13) X. Chen, J. Burns, L. Koopmans, H. Rothkaehi, J. Silk, J. Wu, A. J. Boonstra, B. Cecconi, C. H. Chiang and L. Chen, et al. [arXiv:1907.10853 [astro-ph.IM]].
  * (14) http://ska-jp.org/treed/
  * (15) S. Iguchi et al., SPIE AS24-AS101, 92 (2024).
  * (16) D. Yamauchi, S. Yokoyama and H. Tashiro, Phys. Rev. D 96, no.12, 123516 (2017) doi:10.1103/PhysRevD.96.123516 [arXiv:1709.03243 [astro-ph.CO]].
  * (17) D. Yamauchi and N. S. Sugiyama, Phys. Rev. D 105, no.6, 063515 (2022) doi:10.1103/PhysRevD.105.063515 [arXiv:2108.02382 [astro-ph.CO]].
  * (18) D. Yamauchi, S. Ishimaru, T. Matsubara and T. Takahashi, Phys. Rev. D 107, no.4, 043526 (2023) doi:10.1103/PhysRevD.107.043526 [arXiv:2211.13453 [astro-ph.CO]].
  * (19) S. Yamashita, T. Matsubara, T. Takahashi and D. Yamauchi, [arXiv:2407.01221 [astro-ph.CO]].
  * (20) N. S. Sugiyama, D. Yamauchi, T. Kobayashi, T. Fujita, S. Arai, S. Hirano, S. Saito, F. Beutler and H. J. Seo, doi:10.1093/mnras/stad1505 [arXiv:2302.06808 [astro-ph.CO]].
  * (21) N. S. Sugiyama, D. Yamauchi, T. Kobayashi, T. Fujita, S. Arai, S. Hirano, S. Saito, F. Beutler and H. J. Seo, Mon. Not. Roy. Astron. Soc. 524, no.2, 1651-1667 (2023) doi:10.1093/mnras/stad1935 [arXiv:2305.01142 [astro-ph.CO]].
  * (22) S. Arai, K. Aoki, Y. Chinone, R. Kimura, T. Kobayashi, H. Miyatake, D. Yamauchi, S. Yokoyama, K. Akitsu and T. Hiramatsu, et al. PTEP 2023, no.7, 072E01 (2023) doi:10.1093/ptep/ptad052 [arXiv:2212.09094 [astro-ph.CO]].
  * (23) B. P. Abbott et al. [LIGO Scientific and Virgo Collaborations], Phys. Rev. Lett. 119 (2017) no.16, 161101 doi:10.1103/PhysRevLett.119.161101 [arXiv:1710.05832 [gr-qc]].
  * (24) B. P. Abbott et al. [LIGO Scientific and Virgo and Fermi-GBM and INTEGRAL Collaborations], Astrophys. J. 848 (2017) no.2, L13 doi:10.3847/2041-8213/aa920c [arXiv:1710.05834 [astro-ph.HE]].
  * (25) B. P. Abbott et al., Astrophys. J. 848 (2017) no.2, L12 doi:10.3847/2041-8213/aa91c9 [arXiv:1710.05833 [astro-ph.HE]].
  * (26) J. M. Ezquiaga and M. Zumalacárregui, Phys. Rev. Lett. 119, no.25, 251304 (2017) doi:10.1103/PhysRevLett.119.251304 [arXiv:1710.05901 [astro-ph.CO]].
  * (27) T. Baker, E. Bellini, P. Ferreira, M. Lagos, J. Noller and I. Sawicki, Phys. Rev. Lett. 119, no.25, 251301 (2017) doi:10.1103/PhysRevLett.119.251301 [arXiv:1710.06394 [astro-ph.CO]].
  * (28) J. Sakstein and B. Jain, Phys. Rev. Lett. 119, no.25, 251303 (2017) doi:10.1103/PhysRevLett.119.251303 [arXiv:1710.05893 [astro-ph.CO]].
  * (29) P. Creminelli and F. Vernizzi, Phys. Rev. Lett. 119, no.25, 251302 (2017) doi:10.1103/PhysRevLett.119.251302 [arXiv:1710.05877 [astro-ph.CO]].
  * (30) D. Langlois, R. Saito, D. Yamauchi and K. Noui, Phys. Rev. D 97 (2018) no.6, 061501 doi:10.1103/PhysRevD.97.061501 [arXiv:1711.07403 [gr-qc]].
  * (31) J. J. Terente Díaz, K. Dimopoulos, M. Karčiauskas and A. Racioppi, JCAP 10, 031 (2023) doi:10.1088/1475-7516/2023/10/031 [arXiv:2307.06163 [astro-ph.CO]].
  * (32) B. H. Lee, W. Lee, E. Ó. Colgáin, M. M. Sheikh-Jabbari and S. Thakur, JCAP 04, no.04, 004 (2022) doi:10.1088/1475-7516/2022/04/004 [arXiv:2202.03906 [astro-ph.CO]].
  * (33) D. Langlois and K. Noui, JCAP 02 (2016), 034 doi:10.1088/1475-7516/2016/02/034 [arXiv:1510.06930 [gr-qc]].
  * (34) M. Crisostomi, K. Koyama and G. Tasinato, JCAP 04 (2016), 044 doi:10.1088/1475-7516/2016/04/044 [arXiv:1602.03119 [hep-th]].
  * (35) J. Ben Achour, D. Langlois and K. Noui, Phys. Rev. D 93 (2016) no.12, 124005 doi:10.1103/PhysRevD.93.124005 [arXiv:1602.08398 [gr-qc]].
  * (36) J. Ben Achour, M. Crisostomi, K. Koyama, D. Langlois, K. Noui and G. Tasinato, JHEP 12 (2016), 100 doi:10.1007/JHEP12(2016)100 [arXiv:1608.08135 [hep-th]].
  * (37) G. W. Horndeski, Int. J. Theor. Phys. 10 (1974), 363-384 doi:10.1007/BF01807638
  * (38) C. Deffayet, X. Gao, D. A. Steer and G. Zahariade, Phys. Rev. D 84 (2011), 064039 doi:10.1103/PhysRevD.84.064039 [arXiv:1103.3260 [hep-th]].
  * (39) T. Kobayashi, M. Yamaguchi and J. Yokoyama, Prog. Theor. Phys. 126 (2011), 511-529 doi:10.1143/PTP.126.511 [arXiv:1105.5723 [hep-th]].
  * (40) S. Hirano, T. Kobayashi, D. Yamauchi and S. Yokoyama, Phys. Rev. D 99, no.10, 104051 (2019) doi:10.1103/PhysRevD.99.104051 [arXiv:1902.02946 [astro-ph.CO]].
  * (41) S. Hirano, T. Kobayashi, D. Yamauchi and S. Yokoyama, Phys. Rev. D 102, no.10, 103505 (2020) doi:10.1103/PhysRevD.102.103505 [arXiv:2008.02798 [gr-qc]].
  * (42) S. B. Giddings and M. S. Sloth, JCAP 01, 023 (2011) doi:10.1088/1475-7516/2011/01/023 [arXiv:1005.1056 [hep-th]].
  * (43) S. B. Giddings and M. S. Sloth, Phys. Rev. D 84, 063528 (2011) doi:10.1103/PhysRevD.84.063528 [arXiv:1104.0002 [hep-th]].
  * (44) D. Langlois, Int. J. Mod. Phys. D 28 (2019) no.05, 1942006 doi:10.1142/S0218271819420069 [arXiv:1811.06271 [gr-qc]].
  * (45) T. Kobayashi, Rept. Prog. Phys. 82 (2019) no.8, 086901 doi:10.1088/1361-6633/ab2429 [arXiv:1901.07183 [gr-qc]].
  * (46) C. Brans and R. H. Dicke, Phys. Rev. 124, 925-935 (1961) doi:10.1103/PhysRev.124.925
  * (47) A. A. Starobinsky, Phys. Lett. B 91, 99-102 (1980) doi:10.1016/0370-2693(80)90670-X
  * (48) T. Chiba, T. Okabe and M. Yamaguchi, Phys. Rev. D 62, 023511 (2000) doi:10.1103/PhysRevD.62.023511 [arXiv:astro-ph/9912463 [astro-ph]].
  * (49) C. Deffayet, O. Pujolas, I. Sawicki and A. Vikman, JCAP 10, 026 (2010) doi:10.1088/1475-7516/2010/10/026 [arXiv:1008.0048 [hep-th]].
  * (50) E. Bellini and I. Sawicki, JCAP 07, 050 (2014) doi:10.1088/1475-7516/2014/07/050 [arXiv:1404.3713 [astro-ph.CO]].
  * (51) E. Bellini, R. Jimenez and L. Verde, JCAP 05, 057 (2015) doi:10.1088/1475-7516/2015/05/057 [arXiv:1504.04341 [astro-ph.CO]].
  * (52) T. Kobayashi, Y. Watanabe and D. Yamauchi, Phys. Rev. D 91, no.6, 064013 (2015) doi:10.1103/PhysRevD.91.064013 [arXiv:1411.4130 [gr-qc]].
  * (53) K. Koyama, G. Niz and G. Tasinato, Phys. Rev. D 88, 021502 (2013) doi:10.1103/PhysRevD.88.021502 [arXiv:1305.0279 [hep-th]].
  * (54) L. Pogosian and A. Silvestri, Phys. Rev. D 94, no.10, 104014 (2016) doi:10.1103/PhysRevD.94.104014 [arXiv:1606.05339 [astro-ph.CO]].
  * (55) J. Gleyzes, D. Langlois, M. Mancarella and F. Vernizzi, JCAP 02, 056 (2016) doi:10.1088/1475-7516/2016/02/056 [arXiv:1509.02191 [astro-ph.CO]].
  * (56) D. Yamauchi, PTEP 2022, no.7, 073E02 (2022) doi:10.1093/ptep/ptac095 [arXiv:2203.15599 [astro-ph.CO]].
  * (57) F. Schmidt and D. Jeong, Phys. Rev. D 86, 083513 (2012) doi:10.1103/PhysRevD.86.083513 [arXiv:1205.1514 [astro-ph.CO]].
  * (58) K. Akitsu, Y. Li and T. Okumura, Phys. Rev. D 107, no.6, 063531 (2023) doi:10.1103/PhysRevD.107.063531 [arXiv:2209.06226 [astro-ph.CO]].
  * (59) O. H. E. Philcox, M. J. König, S. Alexander and D. N. Spergel, Phys. Rev. D 109, no.6, 063541 (2024) doi:10.1103/PhysRevD.109.063541 [arXiv:2309.08653 [astro-ph.CO]].


# Efficient LLM Training and Serving with Heterogeneous Context Sharding among
Attention Heads

Xihui Lin1, Yunan Zhang1∗, Suyu Ge2, Barun Patra1, Vishrav Chaudhary1, Xia
Song1
1Microsoft, 2UIUC
<EMAIL_ADDRESS>
Leading authors. Xihui Lin is the major contributor of the kernel. Code is
available is at https://github.com/linxihui/dkernel

###### Abstract

Existing LLM training and inference frameworks struggle in boosting efficiency
with sparsity while maintaining the integrity of context and model
architecture. Inspired by the sharding concept in database and the fact that
attention parallelizes over heads on accelerators, we propose Sparsely-Sharded
(S2) Attention, an attention algorithm that allocates heterogeneous context
partitions for different attention heads to divide and conquer. S2-Attention
enforces each attention head to only attend to a partition of contexts
following a strided sparsity pattern, while the full context is preserved as
the union of all the shards. As attention heads are processed in separate
thread blocks, the context reduction for each head can thus produce end-to-end
speed-up and memory reduction. At inference, LLMs trained with S2-Attention
can then take the KV cache reduction as free meals with guaranteed model
quality preserve. In experiments, we show S2-Attentioncan provide as much as
(1) 25.3X wall-clock attention speed-up over FlashAttention-2, resulting in 6X
reduction in end-to-end training time and 10X inference latency, (2) on-par
model training quality compared to default attention, (3)perfect needle
retrieval accuracy over 32K context window. On top of the algorithm, we build
DKernel, an LLM training and inference kernel library that allows users to
customize sparsity patterns for their own models. We open-sourced DKerneland
make it compatible with Megatron, Pytorch, and vLLM.

##  1 Introduction

The emergence of transformer-based large language models (OpenAI, 2023;
Touvron et al., 2023) has brought about revolutionary opportunities to
research and user experience. However, the remarkable success of these models
across diverse applications also underscores the pressing need for training
and serving these models economically.

The core of the transformer architecture is the self-attention mechanism.
However, self-attention costs quadratic time and memory complexity with
respect to the context length, making it frustratingly expensive to scale to
longer context. For example, with 4096 context length, training Llama 2 70b on
2 trillion tokens Touvron et al. (2023) needs 23 days to finish with 2048 A100
GPUs Rucinski (2024), which is estimated to cost 2 million dollars.

(a) Speedup over FlashAttention-2.

(b) Needle in a haystack evaluation.

Figure 1: Training Efficiency and long-context analysis of S2-Attention.

During inference, the forward pass of LLMs typically involves storing the
previously computed key and value vectors (KV Cache) to avoid recomputing. For
instance, serving Llama 2 70b requires a staggering 43 GB of KV cache when
setting the batch size to 16 and the context length to 2048.

Our study aims to improve the training and inference efficiency of LLM. We
starts from the observation of the sparsity nature of transformers (Wu et al.,
2024; Olsson et al., 2022; Ge et al., 2024). As illustrated in Figure 3,
despite densely trained, the majority of attention mass still concentrates on
a few tokens. For most cases, over $95\%$ attention mass can be recalled with
less than $1\%$ tokens. Also, different attention heads could have dynamically
different attention mass distribution patterns. As inference works with a
sparse subset of contexts for each head, can we accelerate training by
enforcing this sparsity on each head? More specifically, we want to examine
how to maintain the model quality while maximally reducing the context load of
every attention head in training. From a hardware perspective, as each
attention head is processed independently by a distinct thread block,
assigning a subset of the context to each thread block is similar to the
sharding concept in database. Thus, we design Sparsely-Sharded (S2) Attention
built on three key premises:

Figure 2: Illustration of S2-Attentionwith four attention heads on a
hypothetical GPU with 4 thread blocks. Each attention head is allocated with a
shard of the context.

P1: In training, we can get equivalent model quality while allocating only a
partition of the context to each attention head, as long as the union of all
the shards equals the complete context.

P2: Following P1, attention heads should attend to diverse shards. Under the
same compute budget, models trained with heterogeneous sparsity pattern on
attention heads would have a better quality over those trained with
homogeneous sparsity pattern.

P3: Dense layers at certain point are necessary for preserving model quality.

In this way, S2-Attentioncan be trained with significant fewer FLOPs while
maintaining the model quality, including long context capabilities, compared
to those trained with default attention. Compared to previous sparsity method
(Child et al., 2019; Ho et al., 2019; Zaheer et al., 2020), sharding the
context among numerous attention head can achieve higher sparsity per head,
while explicitly maintaining the full context. Also, the heterogeneous
assignment could better utilize the multi-head expressiveness compared to
previous homogeneous designs. During serving, S2-Attentioncan then take KV
Cache saving as a free meal while having a better quality guarantee, as each
attention head is originally trained on reduced contexts.

In this paper, we provide an instantiation of S2-Attentionthat satisfies all
the arguments. We first chunk the sequence of tokens into blocks. For an
attention head $h$, it only attends to the blocks that are adjacent to the
current token, and blocks whose ids are divisible by a stride after offset by
$h$ . A larger stride means each attention head receives a smaller, and more
sparse shard. The stride size is between 1 and number of attention head, which
secure the union can preserve full context. Lastly, we keep the first two
dense layers dense while the rest of layers sparse, as we find it’s crucial to
maintain the initial layers intact.

We validate the effectiveness of S2-Attentionby evaluating the model quality
on a wide range of tasks covering context retrieval, common sense reasoning,
and question answering. Across benchmarks, S2-Attentionachieves on par, or
even better performance compared to densely trained models. For training
efficiency benefit, compared to Flash-Attention 2, S2-Attentioncan achieve as
much as 8.8X wall-clock speed-up at 128K sequence length for 1.3B model. The
boost increases as the model and context length scales up. On a 70B model with
64 heads, S2-Attentioncan bring 25.3X speed-up. For inference speed-up,
S2-Attentioncan achieves up to 12X speed up over FlashAttention-2 on a 1M
context window. Lastly, we back S2-Attentionwith an optimized kernel library,
DKernel, designed to be a plug-in-and-play replacement of the popular
FlashAttnetion-2 and PyTorch SPDA used in most modern training frameworks. It
allows users to freely customized the interested pattern they want to study in
both training and inference. Lastly, we’ve integrated DKernelinto vLLM for
easy adoption in serving.

##  2 Related Work

Efficient Transformers Numerous of attempts have been made to make the
training and serving of LLMs more efficient. FlashAttention family (Dao et
al., 2022; Dao, 2023) are the most widely-adopted attention acceleration
framework. Dao et al. (2022) breaks down the softmax operation into smaller
blockwise computation to reduce the IO between SRAM and HBM. Dao (2023)
further improve the performance by reducing the non-matmul operations, extra
parallelization over sequence length, and better warp organization. Our
implementation is based on the FlashAttention kernel, and leverage its
parallelization over the number of heads. Liu et al. (2023a) further scales
the blockwise computation into multi-accelerator setting by using each GPU as
cache. However, these algorithm does not tackle the redundancy in the
attention structure, which leaves great headroom to further improve
efficiency. The FlashAttention family calculates the exact self-attention. On
the other hand, computation-efficient attention approximation mainly leverages
low-rank and sparsity approximation to reduce the computational complexity of
self-attention (Child et al., 2019; Katharopoulos et al., 2020; Kitaev et al.,
2020; Zaheer et al., 2020; Beltagy et al., 2020). Zaheer et al. (2020)
replaces the one-to-all attention in transformers with sliding window
attention and random attention. Qiu et al. (2020) chunks the sequence into
blocks, which reduces FLOPs by performing attention on a larger granularity.
However, many of these methods can’t bring wall-clock speed-up due to
ignorance of the IO overhead Dao et al. (2022). Also, restricted by receptive
field and the homogeneous sparsity patterns among attention heads, these
methods under-utilize the multi-head structure and limit the sparsity level.
Random attention on the other hand, makes KV cache undecidable during decoding
for generative models.

Attention Pattern Analysis Several studies have sought to understand the roles
individual attention heads play in the Transformer. For encoder-only models,
Voita et al. (2019); Michel et al. (2019); Clark et al. (2019) analyze the
self-attention heads in BERT Devlin et al. (2019) and find a significant
portion of heads attends separator tokens, next or previous tokens and a
hybrid of them. Even heads in the same layer could have different patterns and
impact the downstream tasks in a diverse way. Ge et al. (2024); Fu (2024); Wu
et al. (2024); Olsson et al. (2022) study the attention head behavior for
decoder-only models. They reveal most of the attention heads concentrate on a
few tokens and categorize these patterns to facilitate KV Cache compression.

Infer-time KV Cache management Along with these studies, a line of inference
optimization work emerges which evict part of the KV cache at inference based
on attention score. (Zhang et al., 2023; Liu et al., 2023b) proposes an
eviction policy to retain a combination of recent tokens and important tokens
decided by frequency statistics. Xiao et al. (2023) proposes to only keep the
first quarter and local contexts as they accounts for the predominantly amount
of attention scores. Ge et al. (2024) studies the head level attention pattern
at inference, and proposed to use an adaptive hybrid eviction policies to
better present the attention pattern for each attention head. However, these
methods cannot guarantee performance preserving, especially long context
capabilities Li et al. (2024).

##  3 Observation

Figure 3: Left: Attention mass sparsity in different heads during decoding the
same sequence. Right: Sparsity in different layers. Except for the bottom
layer, all other layers display significant attention sparsity. Experiments
are done on Llama3-8B with different task prompts, including needle retrieval,
summarization and math.

In this section, we present an empirical study to see what each attention head
actually attends to at inference time. We conduct experiments on Llama-3 8B
(Meta, 2024) using a 2048 context length. We first start with the overall
layer by layer pattern, then showcase specific sparsity patterns we find.

Scattered Distribution in Initial Layers. In the initial layers, attention
score is generally uniform for all heads. We call this pattern Uniform
Scatter, which is illustrated at the right of Figure 3. Such pattern suggests
the model does not heavily prioritize any specific part of the input. Thus,
it’s risky to reduce the context when the head show such patterns.
Intuitively, these attention layers should be kept dense.

Sparse Patterns in Middle Layers. As we move deeper into the network, the
attention mechanisms exhibit a marked departure from uniformity. We start with
identifying simpler patterns. Attention Sink Xiao et al. (2023) and Locality
can be seen in a small subset of heads. For these heads, nearly all attention
probability concentrate on the first few tokens or tokens adjacent to the
current position, as shown in the upper left and upper right of Figure 3.

Middle Spikes, as shown in Figure 3, can also be identified in a few heads.
This pattern puts significant probability mass on a few tokens that are in the
middle of the context.

However, the vast majority of attention heads actually show a more complex
hybrid pattern that could be decomposed into patterns mentioned above. For
example, hybrids like Attention Sink + Local, or Attention Sink + Middle
Spikes + Local are the most prevalent patterns, which is also demonstrated in
Ge et al. (2024); Fu (2024).

From the observation, we can conclude that (1) individual attention heads do
not necessarily take the full context to function, (2) the actual context
needed by different heads can be quite different, and (3) the initial
attention layers are usually denser than any other layer. This inspires us to
design S2-Attention, where the context is diversely sharded and distributed to
each attention head. We will discuss the details of our method in the next
section.

##  4 Methodology

###  4.1 Preliminary

In this section, we first present formulation of the context sharding process
for Sparsely Sharded Attention (S2-Attention). Then we discuss context-
sharding policies we designed to instantiate S2-Attention. Lastly, we analyze
why S2-Attentioncan achieve significant wall-clock speed-up for both training
and inference.

###  4.2 Context Sharding

Let $\mathcal{X}=\\{x_{1},x_{2},\dots,x_{N}\\}$ be the input sequence where
$N$ is the length of the sequence,
$\mathcal{H}=\\{h_{1},h_{2},\dots,h_{K}\\}$, where $K$ is the number of
attention heads. Each attention head $h_{k}$ is associated with a shard
$S_{k}$ of the sequence $\mathcal{X}$, determined by a context sharding policy
$f$. The policy can be viewed as a mapping function $f:\mathcal{H}\to
2^{\mathcal{X}}$, where $2^{\mathcal{X}}$ is the power set of $\mathcal{X}$.
Each head receives a subset of the sequence elements: $f(h_{k})=S_{k}$.

We design the sharding process to meet the following properties:

  1. 1.

The union of all the shards equals the full context. Namely, each element
$x_{i}$ in $\mathcal{X}$ must be assigned to at least one shard. Formally,
this can be expressed as:

$\bigcup_{k=1}^{K}S_{k}=\mathcal{X}$

  2. 2.

Different attention heads receive heterogeneous context shards. At least two
heads should receive different subsets of $\mathcal{X}$. Specifically:

$\exists i,j\in\\{1,\dots,K\\},\,i\neq j:S_{i}\neq S_{j}$

  3. 3.

Initial layers should receive full context for optimal performances. For
simplicity, we set the initial two layers to be dense and sparsely shared
contexts among the rest layers.

###  4.3 Context Sharding Policy Choices

From our experiments in Figure 3 and previous studies (Ge et al., 2024; Wu et
al., 2024), we summarize 6 basic attention patterns in trained transformers
detailed in Section 3. As we found the Local pattern is most common across
attention heads, we first ensure each shard includes local context. For
contexts out of this local window, termed as remote context, we would like the
sharding policy to cover the Attention Sink, Middle Spikes and their hybrid
combinations Ge et al. (2024).

To capture the huge combinations while maintaining sparsity, we take advantage
of the inherent expressiveness in the multi-head architecture. Each attention
head is assigned with strided partitions of the remote contexts, with the
starting points of these partitions uniquely offset by the index of the
corresponding head. The heterogeneous assignment allows for a nuanced and
scalable coverage of remote contexts. By combining the two designs, we
introduce a heterogeneous strided sharding policy as formulated below.

###  4.4 Heterogeneous Strided Sharding

KV-cache efficient sharding has a vertical pattern. For $l\geq 1$ and $j\geq
i$, we say a sharding $S$ is KV cache efficient if $k_{i},v_{i}$ is attended
by $q_{j+l}$, then it must be also attended by $q_{j}$. Otherwise,
$k_{i},v_{i}$ has to be stored at token $j$ for generating token $j+l$, but
not used to generate token $j$. In this case, we can just let $q_{j}$ also
attend to $k_{i},v_{i}$ without additional memory for KV cache. However, KV
cache can be dropped after generating certain tokens. This also means that the
sharding $S$ should be based on absolute token positions.

Figure 4: Left: KV-cache efficient sharding. Right: sharding that is not KV-
cache efficient since token 2 is not attended by tokens 2 - 11, but needs to
be stored at token 12 to generate token 13 and onward. The upper right corner
is masked due in case of causal attention.

Vertical pattern with one fixed stride and a local window. For a transformer
with $H$ heads, given a sequence of $N$ tokens, we first chunk context into a
list of blocks of size $S$, which gives us $B=ceil(N/S)$ blocks. We use
$q_{i}$ and $k_{i}$ to denote the query block and key block index. Note that
$q_{i}$ corresponds to the $i$ th indexed block, instead of the $i$ th token,
and one block includes $S$ tokens.

Among the total $B$ blocks, We take the most recent $N_{l}$ blocks as local
blocks and set the rest as remote blocks. Since local blocks are more
frequently attended, we use different stride sizes for local and remote blocks
for fine-grained context modeling. Denote the local stride as $v_{l}$ and
remote stride as $v_{r}$. For simplicity, we use full attention for local
blocks by setting $v_{l}=1$, since local blocks only accounts for a small
portion of the total context and brings negligible change to the overall
efficiency. In practice, $v_{r}$ is set to be a factor number of attention
heads to satisfy the “union as full context” constraint. For each head
$h\in\\{1,...,H\\}$, we set a unique position offset $o_{h}$. For simplicity,
we set $o_{h}=h$ here, which means the offset of each head equals its head
index.

Then for each head $h$, its block attention mask $M^{h}$ in $B\times B$
dimension is as follows:

$M^{h}_{i,j}=\left\\{\begin{array}[]{ll}1&q_{i}-k_{j}<N_{l},\\\
1&(k_{j}-o_{h})\in v_{r}\mathbb{Z}_{\geq 0}\,\&\,q_{i}-k_{j}\in[N_{l},B)\\\
0&\text{otherwise},\end{array}\right.$ (1)

where $x\in m\mathbb{Z}_{\geq 0}$ mean $x$ is 0 or a positive multiple of $m$.

The first constraint determines the shards within the local window. The second
constraint controls what shard each head attends to outside the local window.
Note that if the distance is beyond $B$, the attention is dropped, thus
extends similar to a sliding window. As an example, setting
$B=8,N_{l}=2,v_{l}=1,v_{r}=3,o_{h}=h$ will give us the sharding as presented
in the left part of Figure 2. $B=8,N_{l}=3,v_{l}=2,v_{r}=3,o_{h}=h$ will give
us the right part of Figure 2.

Figure 5: Example Sparsity Patterns.

We pick this neat instantiation but the formulation is quite general. The
constraints can also be met by choosing different offset $o_{h}$, or different
vertical stride $v_{i}$, or local attention window size $N_{l}$ for each head,
which are all supported by our library.

Vertical pattern with multiple fixed strides and a local window. This local-
stride pattern can be easily extended to have multiple strides in different
remote blocks. With blocks $N_{l}<N_{r_{1}}<N_{r_{2}}<B$, and
$v_{r_{1}}<v_{r_{2}}$, we define block attention mask as

$M^{h}_{i,j}=\left\\{\begin{array}[]{llll}1&q_{i}-k_{j}\in[0,N_{l}),&&\\\
1&q_{i}-k_{j}\in[N_{l},N_{r_{1}})&\&&(k_{j}-o^{1}_{h})\in
v_{r_{1}}\mathbb{Z}_{\geq 0},\\\
1&q_{i}-k_{j}\in[N_{r_{2}},B)&\&&(k_{j}-o^{2}_{h})\in
v_{r_{2}}\mathbb{Z}_{\geq 0},\\\ 0&\text{otherwise}.\end{array}\right.$ (2)

To ensure that $M_{i,j}^{h}$ is KV-cache efficient, we requires that
$v_{r_{2}}\in v_{r_{1}}\mathbb{Z}_{\geq 0}$ and $(o^{2}_{h}-0^{1}_{h})\in
v_{r_{1}}\mathbb{Z}_{\geq 0}$.

###  4.5 Training and Inference Advantages

With the formulation of S2-Attention, we can achieve significant FLOPs
reduction in training. To train a default self-attention layer, the total
FLOPs needed is $O(TN)$, where $T$ is the number of tokens to train, $N$ is
the max sequence length. When using S2-Attention, in case where$v_{l}=1$, we
have a dense local window of $L$ tokens, and the vertical stride $v_{r}=v$ for
the remote context. For each attention head, $L+(N-L)/v$ is number of tokens
attention to, i.e. the dense equivalent sequence length. Then the total FLOPS
reduced from $O(TN)$ to $O(T(L+(N-L)/v)$ at an attention layer. For a 1.3B
Llama-like model with 16 attention heads, trained on 8192 sequence length, set
$L=64$ and $v=16$ will generate a 14.32X FLOPs reduction. As the stride $v$ is
bounded by the number of heads for the union constraint, we can have larger v
choices as the model size increases, which will further boost the reduction.
It can be derived that the upper bound of the speed-up grows linearly with the
number of attention heads. For example, for a 70B model with 64 attention
heads, S2-Attentioncan generate 42.89X reduction.

During inference, the prompt prefilling stage has the same FLOPs saving as the
forward pass in training. In the decoding phase, despite the FLOPs saving,
S2-Attentionwill have extra memory footprint reduction. As each attention head
is processed in one Streaming Multiprocessor(SM), skipping the blocks outside
the shard of a given head can thus reduce the memory usage, which further
boosts throughput.

###  4.6 System Implementation

Our S2-Attentiontraining kernel is developed upon the Triton flash attention
code. We’ve compared it to the CUDA version of FlashAttention-2. It accepts a
sparse sparse pattern tensors specified the per head sparsity pattern on a
shard level. The underline algorithm turn the sparse matrix to compressed
sparse row (CSR) format as input, which contains k shard indices for each q
shard, then scan over the k shard indices to do the FlashAttention in a thread
block. We see that when vertical stride is set to 1, e.g., normal dense
attention, the latency is comparable to the original Triton flash attention.
Compared to the CUDA version of Flash Attention 2, the forward pass has
similar latency while the backward pass lags slightly. When the vertical
stride increases, the performance increase almost linearly as shown in Figure
6. The memory usage in S2-attention is the same as FlashAttention-2 in
training regardless of the vertical stride. However at inference, during the
decoding phase, the KV cache is about $1/v$ of a dense attention if local
window is considerably small compared to the current sequence length.

Both the CUDA version and the Triton version of flash attention load the whole
head dimension(D) at once to SRAM, which is intuitive as the whole vector is
needed to compute the attention scores. However, in our experiment, we found
that split at the D dimension in many cases is beneficial. We suspect that
this is because D-split reduce SRAM usage and enable a larger degree of
software pipelining. Interestingly, we also find that with shard size of 64,
D-split only helps when head dimension is 128, the most commonly used case,
while has no additional benefit when head dimension is 64 or 256. We suspect
that this is due to under-tuned hyper-parameters like kernel block sizes.

Algorithm 1 Scratch of D-split at the forward pass in a thread block. See Dao
(2023) for more detail

Init SRAM $b,p,q_{0},q_{1},o_{1},o_{2}$

Load $q_{0}\leftarrow q[...,:D/2]$, $q_{1}\leftarrow q[...,D/2:]$

for j in attn_shard_indices do

Load $b\leftarrow k_{0}^{j}$

$p\leftarrow q_{0}@b$

Load $b\leftarrow k_{1}^{j}$

$p\leftarrow p+q_{1}@b$

$p\leftarrow\mathbf{softmax}(p)$

Load $b\leftarrow v_{0}^{j}$

$o_{1}=update(o_{1},p@b)$

Load $b\leftarrow v_{1}^{j}$

$o_{2}=update(o_{2},p@b)$

end for

Figure 6: Left: Comparison between dense and S2-attention with different
sparsity. All experiment use (batch, heads, head_dim) = (4, 16, 128) on A100
80GB-SXM. Right: Impact of D-split. All experiments use (batch, heads,
head_dim) = (4, 16, 128) on A100 80GB-SXM. We see D-Split=2 consistently
outperforms D-Split=1 in most implementation of flash attention on the forward
pass. However. the impact on the backward pass is negligible.

##  5 Experiment

We conduct comprehensive experiments to validate the significant efficiency
advantage and quality preserving S2-Attentioncan achieve. We first show
S2-Attentioncan achieve on par, if not better, downstream task performance and
context retrieval capability compared to default attention in Section 5.1 and
Section 5.2. We then demonstrate the wall-clock speed-up of S2-Attentionover
FlashAttention-2 across different context length settings in Section 5.3. We
then present the end-to-end throughput and latency improvement of
S2-Attentionover the default attention implementation and infer-time KV cache
management in Section 5.4. Lastly, we test our hypotheses for S2-Attentionwith
comprehensive ablation studies.

Table 1: Training quality evaluation. SWA refers to sliding window attention, while S2 refers to S2-Attention. L refers to number of local blocks. V refers to the vertical stride size. Dense refers to the idex of dense attention layers. Model | KV reduction | Passkey | WinoGrande | piqa | race | wikitext103(ppl)
---|---|---|---|---|---|---
Dense | 0% | 0.865 | 0.592 | 0.733 | 0.403 | 15.88
SWA | 92.9% | 0.334 | 0.566 | 0.721 | 0.380 | 21.03
SWA + Dense 1,2 | 85.4% | 0.771 | 0.577 | 0.728 | 0.381 | 17.45
S2-L1V15 | 92.70% | 0.782 | 0.571 | 0.724 | 0.361 | 19.55
S2-L1V15 + Dense 1,2 | 85.0% | 0.941 | 0.587 | 0.725 | 0.397 | 17.18
S2-L8V15 | 87.5% | 0.740 | 0.586 | 0.721 | 0.379 | 17.87
S2-L8V15 + Dense 1,2 | 80.4% | 0.884 | 0.586 | 0.721 | 0.379 | 17.49

###  5.1 Benchmarking Model Training Quality

Experiment Settings. To study the S2-Attention, we train a range of 1.3B
models with 24 layers, 2048 hidden size with 16 heads, with max sequence
length as 8192. We use the FineWeb-Edu-350B Penedo et al. (2024) as the pre-
training corpus. An OpenAI Tiktoken tokenizer with 100K vocabulary size is
used to process the raw text. All model variations use batch size of 4M tokens
for all sequence lengths and train for a total of 300 billion tokens. For
hyperparameters, we use $\mu$P Yang et al. (2022) with a base shape of 256. A
$\mu$P learning rate of 0.02 is used with linear decay and 0.5% of total
training tokens for warmup. All models are evaluated after training on the
total 300B tokens for one epoch.

We use a model whose attention layers are dense, denoted as Dense, as our
baseline. We use the standard dense attention as our baseline, which is
implemented as flash attention in To study the effect of union constraint and
heterogeneous sharding, we include sliding window attention (SWA) for
comparison. For SWA, we control the FLOPs to be equivalent to S2-Attention.
For S2-Attentionwith 1 local block and a vertical stride of 15, the total
FLOPs can be translated to SWA whose local window size is 9*64=576. To study
the effect of initial dense layers, we also train a SWA with the first two
layers dense.

From Table 1, we can observe S2-Attentionshows promising results. Compared to
fully dense models, S2-Attentioncan achieve on par performance for all the
downstream tasks. With a delta below 0.008, S2-Attentioncan achieves $95.0\%$
KV reduction compared to Dense. The minor gap can be explained by the scaling
law, as S2-Attentiononly has around $10\%$ attention FLOPs compared to Dense.
Notably, in the Passkey Retrieval task, S2-Attentioncan achieve much better
performance compared to Dense, despite each head only has a reduced context.
This observation demonstrate the context understanding ability of the
S2-Attentiondesign.

Compared to SWA, we can see S2-Attentionwith the same FLOPs, can consistently
achieve better performance across all the downstream tasks. Among all the
tasks, the Passkey Retrieval task is where S2-Attentionand SWA have the
largest performance margin. The superior performance indicate the benefit of
securing the union constraint and heterogeneous constraint. Moreover, we can
see that for both S2-Attentionand SWA can benefit from the initial dense
layers. If changing initial layers to sparse, all downstream tasks will have a
significant downgradation.

###  5.2 Context Retrieval

In this section, we evaluate the long context understanding capability on the
Needle in a Haystack task. In this task, a factual sentence, “needle”, is
inserted into texts of 512, 1024, 2048, 4096, 8192, 16384, 32768 tokens. The
needle is positioned at 0-$100\%$ depth of the original context, where 0 is
the start of the context and $100\%$ being the end of the context.

Following the settings in previous section, we pretrain an 1.3B model on a
350B subset of FineWeb-Edu. The first two attention layers are dense while the
rest of layers deploy S2-Attention, with vertical stride as 15 and local
window as 1. We modified the RoPE base to 1,00,000 to adapt to the longer
context window, and continue train the model on BookS3.

As shown in Figure1(b), S2-Attentioncan achieves perfect recall on the 32k
context window. The results validate the long context understanding capability
of S2-Attentiondesign.

###  5.3 Training Speed-up

Benchmark Settings We measure the attention runtime of S2-Attentionand
FlashAttention-2 on an A100 80GB GPU for different context length, number of
head, and head dimension settings.

Figure 7: End-to-end latency speed-up for 1.3B Llama architecture models.

We benchmark 3 model sizes to showcase the scalability of our system. Figure
1(a) shows the 1.3B model results. Figure 7 shows the 70B speedup. For all the
model sizes, S2-Attentioncan achieve multiple times of speed-up over
FlashAttention-2. For 1.3B models with 16 heads, S2-Attentioncan achieve as
much as 8.79X speed-up as the max sequence length grows longer. For 70B models
with 64 heads, S2-Attentioncan give 25.3X end-to-end speed-up. Taken forward
pass alone, S2-Attentioncan generate 37.58X reduction. The overall boost is
hedged due to our less optimized backward kernel, which leaves room for
further improvement.

###  5.4 Inference Speed-up

Figure 8: End-to-end latency speed-up for 7B size Llama architecture models.

In order to demonstrate the inference improvments of S2-Attention, we measure
the end-to-end prefilling and decoding latency over different context length
settings in Figure 8. We choose the FlashAttention-2 backend in vLLM as
baseline for fair comparison. Both methods are deployed on a single node with
8 A100 80GPU, with tensor parallel size setting as 8. For prefilling speed-up,
we set output length as 1, and vary input length between 16k to 512k. As shown
in Figure 8, S2-Attentioncan achieves 1.2X, 2.7X, 3X, 5.7X speed-up on 16k,
128K, 256K, and 512K context.

##  6 Conclusion

In this paper, we introduce S2-Attention, a framework that sparsely shards the
context for different attention heads to divide-and-conquer. Experimental
results show model trained with S2-Attentioncan achieve promising performance
on long context tasks with reduced context for each head. We back
S2-Attentionwith a highly optimized kernel library, we can get equivalent
model quality while achieving speed-up over FlashAttention-2 linearly
increasing over the number of attention heads. We open-sourced our kernel
library and make it a plug-in-and-play alternative for FlashAttention-2 module
in popular training frameworks like Megatron and Pytorch. We also integrated
S2-Attentioninto vLLM backend for instant serving. Both the training and
inference kernels allow users to freely customize their sparsity pattern,
facilitating the whole community to study the topic in the future.

##  7 Acknowledgement

We would like to thank Liyuan Liu and Yao Fu from Microsoft Research and
University of Edinburgh for insightful discussions. We also would like to
thank Zhuohan Li, Simon Mo, and Kaichao You from UC Berkeley and Tsinghua
University for sharing expertise on the vLLM codebase.

## References

  * Beltagy et al. (2020) Iz Beltagy, Matthew E. Peters, and Arman Cohan.  Longformer: The long-document transformer.  _CoRR_ , abs/2004.05150, 2020.  URL https://arxiv.org/abs/2004.05150.
  * Child et al. (2019) Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever.  Generating long sequences with sparse transformers.  _CoRR_ , abs/1904.10509, 2019.  URL http://arxiv.org/abs/1904.10509.
  * Clark et al. (2019) Kevin Clark, Urvashi Khandelwal, Omer Levy, and Christopher D. Manning.  What does BERT look at? an analysis of BERT’s attention.  In _Proceedings of the 2019 ACL Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP_ , pp. 276–286, Florence, Italy, August 2019. Association for Computational Linguistics.  doi: 10.18653/v1/W19-4828.  URL https://aclanthology.org/W19-4828.
  * Dao (2023) Tri Dao.  Flashattention-2: Faster attention with better parallelism and work partitioning.  _CoRR_ , abs/2307.08691, 2023.  doi: 10.48550/ARXIV.2307.08691.  URL https://doi.org/10.48550/arXiv.2307.08691.
  * Dao et al. (2022) Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré.  Flashattention: Fast and memory-efficient exact attention with io-awareness.  In Sanmi Koyejo, S. Mohamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh (eds.), _Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Processing Systems 2022, NeurIPS 2022, New Orleans, LA, USA, November 28 - December 9, 2022_ , 2022.  URL http://papers.nips.cc/paper_files/paper/2022/hash/67d57c32e20fd0a7a302cb81d36e40d5-Abstract-Conference.html.
  * Devlin et al. (2019) Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova.  BERT: pre-training of deep bidirectional transformers for language understanding.  In Jill Burstein, Christy Doran, and Thamar Solorio (eds.), _Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, NAACL-HLT 2019, Minneapolis, MN, USA, June 2-7, 2019, Volume 1 (Long and Short Papers)_ , pp. 4171–4186. Association for Computational Linguistics, 2019.  doi: 10.18653/v1/n19-1423.  URL https://doi.org/10.18653/v1/n19-1423.
  * Fu (2024) Yao Fu.  How do language models put attention weights over long context?  _Yao Fu’s Notion_ , Mar 2024.  URL https://yaofu.notion.site/How-Do-Language-Models-put-Attention-Weights-over-Long-Context.
  * Ge et al. (2024) Suyu Ge, Yunan Zhang, Liyuan Liu, Minjia Zhang, Jiawei Han, and Jianfeng Gao.  Model tells you what to discard: Adaptive KV cache compression for llms.  In _The Twelfth International Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024_. OpenReview.net, 2024.  URL https://openreview.net/pdf?id=88nT0j5jAn.
  * Ho et al. (2019) Jonathan Ho, Nal Kalchbrenner, Dirk Weissenborn, and Tim Salimans.  Axial attention in multidimensional transformers.  _CoRR_ , abs/1912.12180, 2019.  URL http://arxiv.org/abs/1912.12180.
  * Katharopoulos et al. (2020) Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret.  Transformers are rnns: Fast autoregressive transformers with linear attention.  In _Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event_ , volume 119 of _Proceedings of Machine Learning Research_ , pp. 5156–5165. PMLR, 2020.  URL http://proceedings.mlr.press/v119/katharopoulos20a.html.
  * Kitaev et al. (2020) Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya.  Reformer: The efficient transformer.  In _8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020_. OpenReview.net, 2020.  URL https://openreview.net/forum?id=rkgNKkHtvB.
  * Li et al. (2024) Yuhong Li, Yingbing Huang, Bowen Yang, Bharat Venkitesh, Acyr Locatelli, Hanchen Ye, Tianle Cai, Patrick Lewis, and Deming Chen.  Snapkv: LLM knows what you are looking for before generation.  _CoRR_ , abs/2404.14469, 2024.  doi: 10.48550/ARXIV.2404.14469.  URL https://doi.org/10.48550/arXiv.2404.14469.
  * Liu et al. (2023a) Hao Liu, Matei Zaharia, and Pieter Abbeel.  Ring attention with blockwise transformers for near-infinite context.  _CoRR_ , abs/2310.01889, 2023a.  doi: 10.48550/ARXIV.2310.01889.  URL https://doi.org/10.48550/arXiv.2310.01889.
  * Liu et al. (2023b) Zichang Liu, Aditya Desai, Fangshuo Liao, Weitao Wang, Victor Xie, Zhaozhuo Xu, Anastasios Kyrillidis, and Anshumali Shrivastava.  Scissorhands: Exploiting the persistence of importance hypothesis for LLM KV cache compression at test time.  _CoRR_ , abs/2305.17118, 2023b.  doi: 10.48550/arXiv.2305.17118.  URL https://doi.org/10.48550/arXiv.2305.17118.
  * Meta (2024) AI Meta.  Introducing meta llama 3: The most capable openly available llm to date.  _Meta AI_ , 2024.
  * Michel et al. (2019) Paul Michel, Omer Levy, and Graham Neubig.  Are sixteen heads really better than one?  In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett (eds.), _Advances in Neural Information Processing Systems_ , volume 32. Curran Associates, Inc., 2019.  URL https://proceedings.neurips.cc/paper_files/paper/2019/file/2c601ad9d2ff9bc8b282670cdd54f69f-Paper.pdf.
  * Olsson et al. (2022) Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds, Danny Hernandez, Scott Johnston, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah.  In-context learning and induction heads.  _CoRR_ , abs/2209.11895, 2022.  doi: 10.48550/ARXIV.2209.11895.  URL https://doi.org/10.48550/arXiv.2209.11895.
  * OpenAI (2023) OpenAI.  Gpt-4 technical report, 2023.
  * Penedo et al. (2024) Guilherme Penedo, Hynek Kydlícek, Loubna Ben Allal, Anton Lozhkov, Margaret Mitchell, Colin Raffel, Leandro von Werra, and Thomas Wolf.  The fineweb datasets: Decanting the web for the finest text data at scale.  _CoRR_ , abs/2406.17557, 2024.  doi: 10.48550/ARXIV.2406.17557.  URL https://doi.org/10.48550/arXiv.2406.17557.
  * Qiu et al. (2020) Jiezhong Qiu, Hao Ma, Omer Levy, Wen-tau Yih, Sinong Wang, and Jie Tang.  Blockwise self-attention for long document understanding.  In Trevor Cohn, Yulan He, and Yang Liu (eds.), _Findings of the Association for Computational Linguistics: EMNLP 2020, Online Event, 16-20 November 2020_ , volume EMNLP 2020 of _Findings of ACL_ , pp. 2555–2565. Association for Computational Linguistics, 2020.  doi: 10.18653/V1/2020.FINDINGS-EMNLP.232.  URL https://doi.org/10.18653/v1/2020.findings-emnlp.232.
  * Rucinski (2024) Szymon Rucinski.  Efficient language adaptive pre-training: Extending state-of-the-art large language models for polish.  _CoRR_ , abs/2402.09759, 2024.  doi: 10.48550/ARXIV.2402.09759.  URL https://doi.org/10.48550/arXiv.2402.09759.
  * Touvron et al. (2023) Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cristian Canton Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu, Wenyin Fu, Brian Fuller, Cynthia Gao, Vedanuj Goswami, Naman Goyal, Anthony Hartshorn, Saghar Hosseini, Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez, Madian Khabsa, Isabel Kloumann, Artem Korenev, Punit Singh Koura, Marie-Anne Lachaux, Thibaut Lavril, Jenya Lee, Diana Liskovich, Yinghai Lu, Yuning Mao, Xavier Martinet, Todor Mihaylov, Pushkar Mishra, Igor Molybog, Yixin Nie, Andrew Poulton, Jeremy Reizenstein, Rashi Rungta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric Michael Smith, Ranjan Subramanian, Xiaoqing Ellen Tan, Binh Tang, Ross Taylor, Adina Williams, Jian Xiang Kuan, Puxin Xu, Zheng Yan, Iliyan Zarov, Yuchen Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurelien Rodriguez, Robert Stojnic, Sergey Edunov, and Thomas Scialom.  Llama 2: Open foundation and fine-tuned chat models, 2023.
  * Voita et al. (2019) Elena Voita, David Talbot, Fedor Moiseev, Rico Sennrich, and Ivan Titov.  Analyzing multi-head self-attention: Specialized heads do the heavy lifting, the rest can be pruned, July 2019.  URL https://aclanthology.org/P19-1580.
  * Wu et al. (2024) Wenhao Wu, Yizhong Wang, Guangxuan Xiao, Hao Peng, and Yao Fu.  Retrieval head mechanistically explains long-context factuality.  _CoRR_ , abs/2404.15574, 2024.  doi: 10.48550/ARXIV.2404.15574.  URL https://doi.org/10.48550/arXiv.2404.15574.
  * Xiao et al. (2023) Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis.  Efficient streaming language models with attention sinks.  _CoRR_ , abs/2309.17453, 2023.  doi: 10.48550/ARXIV.2309.17453.  URL https://doi.org/10.48550/arXiv.2309.17453.
  * Yang et al. (2022) Greg Yang, Edward J. Hu, Igor Babuschkin, Szymon Sidor, David Farhi, Jakub Pachocki, Xiaodong Liu, Weizhu Chen, and Jianfeng Gao.  Tensor programs v: Tuning large neural networks via zero-shot hyperparameter transfer.  In _NeurIPS 2021_ , March 2022.  URL https://www.microsoft.com/en-us/research/publication/tuning-large-neural-networks-via-zero-shot-hyperparameter-transfer/.
  * Zaheer et al. (2020) Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontañón, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, and Amr Ahmed.  Big bird: Transformers for longer sequences.  In Hugo Larochelle, Marc’Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin (eds.), _Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual_ , 2020.  URL https://proceedings.neurips.cc/paper/2020/hash/c8512d142a2d849725f31a9a7a361ab9-Abstract.html.
  * Zhang et al. (2023) Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Ré, Clark W. Barrett, Zhangyang Wang, and Beidi Chen.  H${}_{\mbox{2}}$o: Heavy-hitter oracle for efficient generative inference of large language models.  _CoRR_ , abs/2306.14048, 2023.  doi: 10.48550/arXiv.2306.14048.  URL https://doi.org/10.48550/arXiv.2306.14048.

##  Appendix A Appendix

You may include other additional sections here.