# Non-linear bi-algebraic curves and surfaces in moduli spaces of Abelian
differentials

Bertrand Deroin & Carlos Matheus  Bertrand Deroin
CNRS & Université de Cergy-Pontoise (UMR CNRS 8088), 95302, Cergy-Pontoise,
France<EMAIL_ADDRESS>Carlos Matheus
CNRS & École Polytechnique (UMR CNRS 7640), 91128, Palaiseau, France.
<EMAIL_ADDRESS>
> Abstract. The strata of the moduli spaces of Abelian differentials are non-
> homogenous spaces carrying natural bi-algebraic structures. Partly inspired
> by the case of homogenous spaces carrying bi-algebraic structures (such as
> torii, Abelian varieties and Shimura varieties), Klingler and Lerer recently
> showed that any bi-algebraic curve in a stratum of the moduli space of
> Abelian differentials is linear provided that the so-called condition
> $(\star)$ is fulfilled.
>
> In this note, we construct a non-linear bi-algebraic curve, resp. surface,
> of Abelian differentials of genus $7$, resp. $10$.

##  1\. Introduction

The study of transcendence properties and unlikely intersections in
Diophantine geometry is a fascinating topic possessing a vast literature
developing in many directions including those related to the interplay between
Hodge theory and the geometry of homogenous spaces such as Abelian varieties
and their moduli spaces (that is, Shimura varieties). In this context, an
impressive number of heuristic principles and rigorous results111Such as
Ax–Schanuel conjectures, Ax–Lindemann type theorems, André–Oort and
Zilber–Pink conjectures. was discovered by many authors, and, as a way to
unify these statements and also suggest new ones, the point of view of _bi-
algebraic structures_ became increasingly popular: see, for instance, the
survey [11] and the articles [3] and [4].

In their recent work [10], Klingler and Lerer proposed222This is a natural
goal because the moduli spaces of Abelian differentials seem to “behave” like
homogenous spaces: for example, the celebrated breakthroughs by Eskin,
Mirzakhani and Mohammadi [6], [7] show that this is the case from the point of
view of Dynamical Systems. to extend the bi-algebraic point of view to the
non-homogenous setting of moduli spaces of Abelian differentials. More
concretely, the moduli space of Abelian differentials of genus $g$ is
stratified into complex quasi-projective algebraic orbifolds $H(\kappa)$
parametrising non-trivial Abelian differentials whose zeroes have
multiplicities prescribed by a list $\kappa=(k_{1},\dots,k_{\sigma})$ such
that $k_{1}+\dots+k_{\sigma}=2g-2$. As it was proved by Veech and Masur, the
relative periods of the elements of $H(\kappa)$ can be used to define the so-
called _period charts_ inducing a linear integral structure on the
analytification of $H(\kappa)$. In particular, after projectivising the
stratum $H(\kappa)$, $\kappa=(k_{1},\dots,k_{\sigma})$, we obtain a quasi-
projective orbifold $\mathcal{H}(\kappa)$ of dimension $2g-2+\sigma$ whose
analytification possesses a linear projective structure. In this setting, a
closed, irreducible, algebraic subvariety $W$ of $\mathcal{H}(\kappa)$ is _bi-
algebraic_ if its analytification $W^{\textrm{an}}$ is algebraic in period
charts333That is, the relative periods of Abelian differentials projectively
lying in $W^{\textrm{an}}$ satisfy exactly
$\textrm{codim}_{\mathcal{H}(\kappa)}(W)$ independent algebraic relations
(over $\mathbb{C}$)., cf. [10, Def. 1.1], and Klingler and Lerer proved that
all bi-algebraic curves in the strata $\mathcal{H}(2)$ and $\mathcal{H}(1,1)$
of Abelian differentials of genus two are (projectively) linear (in period
charts), and, in general, any bi-algebraic curve in $\mathcal{H}(\kappa)$ is
linear _provided_ their condition $(\star)$ is satisfied (cf. [10, Thm. 2.8 &
2.10]). Furthermore, they asked whether the bi-algebraicity of a subvariety of
$\mathcal{H}(\kappa)$ is already enough to automatically ensure its linearity
without extra conditions (cf. [10, Conjecture 2.7]). The main results of this
paper say that some bi-algebraic subvarieties can be non-linear:

######  Theorem 1.1.

The projectivisation of the family
$\\{(C_{u},\omega_{u})\\}_{u\in{\mathbb{C}}\setminus\\{0,\pm 1\\}}$ of Abelian
differentials defined by

(1.1) $C_{u}:=\overline{\\{y^{6}=x(x-1)(x+1)(x-u)\\}}\text{ and
}\omega_{u}:=x^{2}dx/y^{5},$

is a bi-algebraic curve in $\mathcal{H}(12)$ which is not linear.

######  Remark 1.2.

Note that $C_{u}$ is a branched cover of
$\overline{\\{z^{2}=x(x-1)(x+1)(x-u)\\}}$. In particular, the family of
Abelian differentials
$\\{(C_{u},[dx/y^{3}])\\}_{u\in\mathbb{C}\setminus\\{0,\pm 1\\}}$ corresponds
to an arithmetic Teichmüller curve (in the sense of [9, §5]).

######  Theorem 1.3.

Given a generic algebraic curve $C\subset\mathbb{C}^{3}$, the projectivisation
of the family $\\{(C_{a,b,c},\omega_{a,b,c})\\}_{\begin{subarray}{c}(a,b,c)\in
C,\\\ a,b,c\in{\mathbb{C}}\setminus\\{0,1\\}\\\ distinct\end{subarray}}$ of
Abelian differentials defined by

(1.2) $C_{a,b,c}:=\overline{\\{y^{6}=x(x-1)(x-a)(x-b)(x-c)\\}}\text{ and
}\omega_{a,b,c}:=dx/y^{5},$

is a bi-algebraic curve in $\mathcal{H}(18)$ which is not linear.

In particular, observe that if $S\subset\mathbb{C}^{3}$ is a generic algebraic
surface, then the projectivisation of the family
$\\{(C_{a,b,c},\omega_{a,b,c})\\}_{\begin{subarray}{c}(a,b,c)\in S,\\\
a,b,c\in{\mathbb{C}}\setminus\\{0,1\\}\\\ distinct\end{subarray}}$ is a bi-
algebraic surface in $\mathcal{H}(18)$ which is not linear. We do not have any
example of a non linear bi-algebraic subvariety of dimension at least $3$ in
some stratum of abelian differentials.

The proofs of these statements occupy the rest of the paper. More precisely,
after a brief discussion of the variations of Hodge structures associated to
the families of curves $C_{u}$ and $C_{a,b,c}$ in Section 3, we establish
Theorem 1.1, resp. 1.3, in Section 5, resp. 6.

## Acknowledgements

This work began at the workshop _Finiteness results for special subvarieties:
Hodge theory, o-minimality, dynamics_ organised by S. Filip, D. Fisher and B.
Klingler held in Como (Italy) from October 16 to 22, 2022. We are indebted to
several participants of this workshop (including S. Cantat, S. Filip, B.
Klingler, L. Lerer and M. Möller) for many interesting discussions.

##  2\. Eigencohomology of cyclic covers

Let $t=(t_{1},\ldots,t_{n})\in\mathbb{C}^{n}$ be a collection of distinct
complex numbers. The plane algebraic curve defined by the equation

$y^{d}=(x-t_{1})\ldots(x-t_{n})$

compactifies as a smooth compact Riemann surface $C_{t}$ by adding
$a=\textrm{gcd}(d,n)$ points at infinity. The function
$X_{t}:C_{t}\rightarrow\mathbb{P}^{1}$ defined by $X_{t}(x,y)=x$ is a ramified
covering over the sphere having $n$ critical points of order $d$ and $a$
critical point(s) of order $d/a$. So by Riemann-Hurwitz, the genus $g$ of
$C_{t}$ is given by the formula:

$g=\frac{(n-1)(d-1)-(a-1)}{2}.$

In particular, if $(d,n)=(6,4)$, $g=7$, and if $(d,n)=(6,5)$, $g=10$.

The covering $X_{t}$ is an abelian covering whose Galois group is generated by
the transformation $\pi_{t}(x,y)=(x,\zeta y)$ where $\zeta=\exp(2i\pi/d)$. We
denote by $H^{*}_{\zeta}(C_{t},\mathbb{C})$ the eigenspace
$\text{Ker}(\pi_{t}^{*}-\zeta)\subset H^{*}(C_{t},\mathbb{C})$. Unless
specified, a holomorphic/meromorphic eigenform on $C_{t}$ is a form $\eta$
such that $\pi_{t}^{*}\eta=\zeta\eta$. The following result is well-known,
see, e.g., [13, §3] for more details. We provide the proof for completeness.

######  Lemma 2.1.

In the regime $n<d$, we have

$H^{1,0}_{\zeta}(C_{t},\mathbb{C})=\\{U(x)\frac{dx}{y^{d-1}}\,|\,U\text{
polynomial of degree }\leq n-2\\}.$

Moreover, the eigenspace $H^{1}_{\zeta}(C_{t},\mathbb{C})$ is made of
cohomology classes of holomorphic eigenforms, i.e.

$H^{1}_{\zeta}(C_{t},\mathbb{C})\simeq H^{1,0}_{\zeta}(C_{t},\mathbb{C})\text{
and
}H^{0,1}_{\zeta}(C_{t},\mathbb{C})=\overline{H^{1,0}_{\zeta^{d-1}}(C_{t},\mathbb{C})}=0.$

###### Proof.

A form in $H^{1,0}_{\zeta}(C_{t})$ can be expressed as
$\eta=U(x)\frac{dx}{y^{d-1}}$ where $U$ is a meromorphic function on
$\mathbb{P}^{1}$, which is holomorphic except possibly at the points $t_{i}$’s
and the point at infinity. In the sequel assume that $U$ does not vanish
identically.

Around the point $t_{i}$, we can write $U(x)\sim\text{cst}\,(x-t_{i})^{k_{i}}$
for some integer $k_{i}$ and a non zero constant. We can also take $y$ as a
holomorphic coordinates and we have $x-t_{i}\sim\text{cst}\,y^{d}$, so

$\eta\sim_{(t_{i},0)}\text{cst}\,y^{dk_{i}}dy$

We deduce for $\eta$ to be holomorphic, $k_{i}$ must be non negative for every
$i=1,\ldots,n$, hence $U$ is a polynomial of $x$.

It will be convenient to introduce the integers $b,c$ such that $n=ab$ and
$d=ac$. Around a point $\infty$ of $C_{t}$ at infinity, we have a coordinate
$z$ so that $x\sim\text{cst}\,z^{-c}$ and $y\sim\text{cst}\,z^{-b}$. We then
have (denoting $\text{deg}(U)$ the degree of $U$)

$\eta\sim_{\infty}\text{cst}\,z^{-c\text{deg}(U)-c-1+b(d-1)}$

which shows that $\eta$ is holomorphic iff
$\text{deg}(U)\leq\frac{b(d-1)-c-1}{c}=n-1-\frac{(b+1)}{c}$, or equivalently,
since $b<c$, iff $\text{deg}(U)\leq n-2$. In particular
$\text{dim}\,H^{1,0}_{\zeta}(C_{t},\mathbb{C})=n-1=\text{dim}\,H^{1}_{\zeta}(C_{t},\mathbb{C})$
so we are done. ∎

##  3\. Finite monodromies

The condition $(\star)$ of Klingler and Lerer (cf. [10, Def. 6.5]) indicates
that the families of curves whose Jacobians possess non-trivial fixed parts
are potential sources of non-linear bi-algebraic subvarieties: see §6.4 of
[10] for more explanations. In this direction, we note that such families of
curves were previously studied by several authors in the context of
Teichmüller dynamics (e.g., [8], [12], [14], [13], [2] among many others),
and, in fact, the families of curves $C_{u}$ and $C_{a,b,c}$ are particular
instances of the families of curves described in [13, Thm. 8.3]. For later
reference, we recall below some of the remarkable properties of the families
of Jacobians of $C_{u}$ and $C_{a,b,c}$.

######  Proposition 3.1.

If $(d,n)=(6,4)$ or $(d,n)=(6,5)$, the monodromy of the Gauss-Manin connection
on the bundle $H^{1}_{\zeta}$ is finite. In particular, the families of forms
$\\{(C_{u},[\omega_{u}])\\}_{u\neq 0,\pm 1}$, lying in $\mathcal{H}(12)$, and
$\\{(C_{a,b,c},[\omega_{a,b,c}])\\}_{\begin{subarray}{c}a,b,c\neq 0,1\\\
distinct\end{subarray}}$, lying in $\mathcal{H}(18)$, are bi-algebraic.

Here $[\omega]$ is the projective class of $\omega$ in the projectivization of
the Hodge bundle.

###### Proof.

This statement is included in Theorem 8.3 of [13]. Let us provide the proof
for completeness. Lemma 2.1 ensures that the Gauss-Manin invariant hermitian
form $\frac{i}{2}\int\omega\wedge\overline{\omega}$ on $H^{1}_{\zeta}$ (resp.
$H^{1}_{\zeta^{5}}=\overline{H^{1}_{\zeta}}$) is positive definite, so its
monodromy is unitary. Moreover, the bundle
$H^{1}(C_{t},\mathbb{C})_{\zeta}\oplus H^{1}(C_{t},\mathbb{C})_{\zeta^{5}}$ is
defined over $\mathbb{Q}$ (because $\zeta$, $\zeta^{5}$ are the sole primitive
sixth roots of unity). So the monodromy is integral (it is the restriction of
the symplectic integral representation to a rational subspace) and preserves a
unitary form; hence it is finite. In particular, as it is explained in [13,
pp. 929] for instance, since the periods of the corresponding families of
forms are holomorphic (multivalued) functions of $t$ which are algebraic when
the monodromy is finite, one has that the corresponding families of forms are
bi-algebraic. ∎

##  4\. Picard-Fuchs theory: computation of the Gauss-Manin connection

The vector bundle $H^{1}_{\zeta}$ over the configuration space $B$ of $n$
distinct complex numbers is isomorphic, thanks to Lemma 2.1, to the product
$B\times\mathbb{C}_{n-2}[x]$, where $\mathbb{C}_{n-2}[x]$ stands for the space
of polynomials of degree $\leq n-2$. More precisely, for
$U\in\mathbb{C}_{n-2}[x]$, we denote

(4.3) $\omega_{t,U}=U(x)\frac{dx}{y^{d-1}}$

the corresponding holomorphic eigenform on $C_{t}$. The map
$(t,U)\mapsto\omega_{t,U}$ gives the desired trivialization. In the following,
we compute the Gauss-Manin connection on $H^{1}_{\zeta}$ in this latter: it is
defined as follows

$\nabla_{\cdot}(\omega_{t,U})=\omega_{t,\nabla_{\cdot}U}.$

We denote

(4.4) $P(x)=(x-t_{1})\ldots(x-t_{n}).$

######  Proposition 4.1.

In the trivialization $H^{1}_{\zeta}\simeq B\times\mathbb{C}_{n-2}[x]$ given
by Lemma 2.1, the Gauss-Manin connection is given by

$\nabla_{\partial_{t_{k}}}U=-\frac{U(t_{k})}{P^{\prime}(t_{k})}\frac{P(x)-P(t_{k})-P^{\prime}(t_{k})(x-t_{k})}{(x-t_{k})^{2}}+\frac{1}{d}\frac{U(t_{k})}{P^{\prime}(t_{k})}\frac{P^{\prime}(x)-P^{\prime}(t_{k})}{x-t_{k}}+$
$+\frac{d-1}{d}\frac{U(x)-U(t_{k})}{x-t_{k}}$

for any $U\in\mathbb{C}_{n-2}[x]$.

###### Proof.

Let $N_{i}$ be simply connected neighborhoods of the $t_{i}$’s, chosen
sufficiently small to be disjoint. Denoting by $N=\cup N_{i}$, and
$Y=\mathbb{C}\setminus N$, the Riemann surfaces $X_{t^{\prime}}^{-1}(Y)$ for
$t^{\prime}=(t_{1}^{\prime},\ldots,t_{n}^{\prime})\in\prod_{i}N_{i}$, are
naturally identified to $X_{t}^{-1}(Y)$ via a holomorphic family of
biholomorphisms $\Phi_{t^{\prime},t}:X_{t}^{-1}(Y)\rightarrow
X_{t^{\prime}}^{-1}(Y)$ that satisfy
$X_{t^{\prime}}\circ\Phi_{t^{\prime},t}=X_{t}$ and $\Phi_{t,t}=\text{id}$.
Using this family it is possible to derivate a holomorphic family of
holomorphic $1$-forms $\omega_{t}$ on $X_{t}^{-1}(Y)$ with respect to the
$t_{k}$-variables, by the formula

$\partial_{t_{k}}\omega_{t}:=\left(\partial_{t^{\prime}_{k}}\Phi_{t^{\prime},t}^{*}\omega_{t^{\prime}}\right)_{|t^{\prime}=t};$

by shrinking $N$, this defines a meromorphic $1$-form on $C_{t}$ which is
holomorphic apart from the points $(t_{i},0)\in C_{t}$ but that might develop
a pole at each $(t_{i},0)\in C_{t}$. These poles have no residue at the points
$(t_{i},0)$, so the period of $\partial_{t_{k}}\omega_{t}$ is well-defined as
an element of $H^{1}(C_{t},{\mathbb{C}})$: for more discussion about this
construction (leading to the Kodaira-Spencer map and the Gauss-Manin
connection), see Clemens’ book [5, §2.10] and Voisin’s book [15, Ch. 9].

Applied to the family of eigenforms
$\\{(C_{t},\omega_{t,U})\\}_{t\in\mathbb{C}^{n}\setminus\Delta}$, where
$U\in\mathbb{C}_{n-2}[x]$, the derivative $\partial_{t_{k}}\omega_{t}$ is a
meromorphic eigenform, which is cohomologous to a holomorphic eigenform
$\omega_{t,V_{t}}$ with $V_{t}\in\mathbb{C}_{n-2}[x]$ by Lemma 2.1. By
definition

(4.5) $\nabla_{t_{k}}U:=V_{t}.$

This equation is characterized by the existence of a (unique) meromorphic
eigenfunction $f:C_{t}\rightarrow\mathbb{P}^{1}$ which is such that

$\partial_{t_{k}}\omega_{t,U}-\omega_{t,V_{t}}=df.$

Here eigenfunction means $f\circ\pi_{t}=\zeta f$. Let us now make the
computations and find the function $f$.

Derivating the multi-valued function $y=\prod_{j}(x-t_{j})^{1/d}$ with respect
to the variables $t_{k}$ yields $\partial_{t_{k}}y=\frac{-y}{d(x-t_{k})}$ from
which we get

(4.6)
$\partial_{t_{k}}\left(U(x)\frac{dx}{y^{d-1}}\right)=\frac{d-1}{d}\frac{U(x)}{x-t_{k}}\frac{dx}{y^{d-1}}.$

This meromorphic eigenform has a pole at the point $(t_{k},0)$, which in the
coordinates $y$ of $C_{t}$ at $(t_{k},0)$ has the following local behaviour
(assuming that $U(t_{k})\neq 0$)

(4.7)
$\partial_{t_{k}}\left(U(x)\frac{dx}{y^{d-1}}\right)\sim_{(t_{k},0)}(d-1)\frac{U(t_{k})}{y^{d}}dy.$

The meromorphic function $f:C_{t}\rightarrow\mathbb{P}^{1}$ defined by

(4.8) $f(x,y):=\frac{y}{x-t_{k}}$

has a unique pole at the point $(t_{k},0)$ and it has the following local
behaviour

(4.9) $f\sim_{(t_{k},0)}\frac{P^{\prime}(t_{k})}{y^{d-1}}$

(it has no pole at infinity since $d\geq n$). In particular, the eigenform

$\partial_{t_{k}}\left(U(x)\frac{dx}{y^{d-1}}\right)+\frac{U(t_{k})}{P^{\prime}(t_{k})}df$

is holomorphic on $C_{t}$, we thus have

(4.10)
$\partial_{t_{k}}\left(U(x)\frac{dx}{y^{d-1}}\right)+\frac{U(t_{k})}{P^{\prime}(t_{k})}df=\omega_{t,V_{t}},$

with $\nabla_{t_{k}}U=V_{t}$ by (4.5), so it remains to compute (4.10). We
have

$df=f\frac{df}{f}=\frac{y}{x-t_{k}}\left(\frac{dy}{y}-\frac{dx}{x-t_{k}}\right)=\frac{y}{x-t_{k}}\left(\frac{1}{d}\sum_{j\neq
k}\frac{dx}{x-t_{j}}+\left(\frac{1}{d}-1\right)\frac{dx}{x-t_{k}}\right)=$
$=\frac{P(x)}{x-t_{k}}\left(\frac{1}{d}\sum_{j\neq
k}\frac{1}{x-t_{j}}-\frac{d-1}{d}\frac{1}{x-t_{k}}\right)\frac{dx}{y^{d-1}}.$

So using the identity

$\frac{P(x)}{x-t_{k}}\sum_{j\neq
k}\frac{1}{x-t_{j}}=\left(\frac{P(x)}{x-t_{k}}\right)^{\prime}=\frac{P^{\prime}(x)}{x-t_{k}}-\frac{P(x)}{(x-t_{k})^{2}}$

we deduce

$V_{t}(x)=\frac{d-1}{d}\left(\frac{U(x)}{x-t_{k}}-\frac{U(t_{k})}{P^{\prime}(t_{k})}\frac{P(x)}{(x-t_{k})^{2}}\right)+\frac{1}{d}\frac{U(t_{k})}{P^{\prime}(t_{k})}\left(\frac{P^{\prime}(x)}{x-t_{k}}-\frac{P(x)}{(x-t_{k})^{2}}\right).$

Rearranging terms and using $P(t_{k})=0$ we find the expression of the Gauss-
Manin connection given by Proposition 4.1. ∎

##  5\. Proof of Theorem 1.1

In view of Proposition 3.1, our task is reduced to check that the lines
spanned by $\omega_{u}$, $u\neq 0,1$, sit “generically” on
$H^{1}_{\zeta}(C_{u},\mathbb{C})$. A direct application of Proposition 4.1
shows that the Gauss-Manin on $H^{1}_{\zeta}$ in the trivialization given by
the sections $p_{k}=x^{k}\frac{dx}{y^{5}}$ for $k=0,1,2$ is

(5.11) $\nabla=d+\frac{dt}{6t(1-u^{2})}\begin{pmatrix}5u^{2}-4&u&u^{2}\\\
5u&5u^{2}&5u\\\ 2&2u&2u^{2}\end{pmatrix}$

By Proposition 3.1, it suffices to prove that the family of periods
$\\{p_{2,u}\\}_{u\neq 0,\pm 1}$ does not lie in a flat strict subbundle of
$H^{1}_{\zeta}$, which will be achieved by proving that at a generic parameter
$u$ the periods $p_{2,u},\nabla_{u}p_{2,u},\nabla_{u}^{2}p_{2,u}$ are linearly
independent in $H^{1}_{\zeta}(C_{u},\mathbb{C})$. We use the basis $p_{k,u}$
of $H^{1}_{\zeta}(C_{u},\mathbb{C})$ to perform the computations, together
with the explicit expression of the Gauss-Manin connection (5.11). We have in
this basis

$p_{2,u}=\begin{pmatrix}0\\\ 0\\\ 1\end{pmatrix},\ \
\nabla_{u}p_{2,u}=\frac{1}{6u(1-u^{2})}\begin{pmatrix}u^{2}\\\ 5u\\\
2u^{2}\end{pmatrix}$

and

$\nabla_{u}\left(6(1-u^{2})\nabla_{u}p_{2,u}\right)=\frac{1}{6(1-u^{2})}\begin{pmatrix}7+u^{2}\\\
40u\\\ 24-8u^{2}\end{pmatrix}.$

Those three vectors are generically linearly independent, and by Leibniz rule,
so are the vectors $p_{2,u},\nabla_{u}p_{2,u},\nabla_{u}^{2}p_{2,u}$.

######  Remark 5.1.

The family $(\\{y^{10}=x(x-1)(x-v)\\},[dx/y^{9}])_{v\neq 0,1}$ is a bi-
algebraic curve in $\mathcal{H}(16)$ because $dx/y^{9}$ lies in a fixed part
of dimension four (cf. [13, Thm. 8.3]). Nonetheless, a direct calculation
using Proposition 4.1 reveals that this bi-algebraic curve is linear.

##  6\. Proof of Theorem 1.3

In view of Proposition 3.1, our task is reduced to check that the lines
spanned by $\omega_{a,b,c}$, $a,b,c\neq 0,1$, sit “generically” on the fibers
of the vector bundle $H^{1}_{\zeta}$. As it turns out, this fact is a
particular case of Theorem 6.1 in [13], where a result of Veech and Masur
(about the period charts) is cleverly explored to derive that the (period) map
taking the configuration of points $\\{0,1,a,b,c\\}$ to $[dx/y^{5}]$ is a
holomorphic local diffeomorphism.444Furthermore, McMullen also gets in [13,
Thm. 9.1] that the periods of $dx/y^{5}$ are algebraic functions of $(a,b,c)$.

In the sequel, we give an alternative proof of the non-linearity of the bi-
algebraic surface
$\\{(C_{a,b,c},\omega_{a,b,c})\\}_{\begin{subarray}{c}a,b,c\in{\mathbb{C}}\setminus\\{0,1\\}\\\
distinct\end{subarray}}$ via Proposition 4.1: in particular, we do not rely on
the results of Veech and Masur, but rather perform direct calculations (which
might be of independent interest).

Proposition 4.1 shows that

$\nabla_{\partial_{u}}\frac{dx}{y^{5}}=\frac{1}{P^{\prime}(u)}\left(\frac{1}{6}\frac{P^{\prime}(x)-P^{\prime}(u)}{x-u}-\frac{P(x)-P^{\prime}(u)(x-u)}{(x-u)^{2}}\right)\frac{dx}{y^{5}}$

(where $u\in\\{a,b,c\\}$ and $P(x)=x(x-1)(x-a)(x-b)(x-c)$).

At this point, it remains to check that

$\frac{dx}{y^{5}},\ \ \nabla_{\partial_{a}}\frac{dx}{y^{5}},\ \
\nabla_{\partial_{b}}\frac{dx}{y^{5}},\ \
\nabla_{\partial_{c}}\frac{dx}{y^{5}}$

are linearly independent. For this sake, let us write them in terms of the
basis $\\{x^{k}dx/y^{5}:k=0,1,2,3\\}$. If we denote $\\{u,v,w\\}=\\{a,b,c\\}$,
then

$\displaystyle\frac{P^{\prime}(x)-P^{\prime}(u)}{x-u}$ $\displaystyle=$
$\displaystyle(u^{3}-u^{2}+uv-u^{2}v+uw-u^{2}w-2vw+uvw)+(u^{2}-u+3v-uv+3w-uw+3vw)x$
$\displaystyle+$ $\displaystyle(u-4-4v-4w)x^{2}+5x^{3},$
$\displaystyle\frac{P(x)-P^{\prime}(u)(x-u)}{(x-u)^{2}}$ $\displaystyle=$
$\displaystyle(u^{3}-u^{2}+uv-u^{2}v+uw-u^{2}w-vw+uvw)$ $\displaystyle+$
$\displaystyle(u^{2}-u+v-uv+w-uw+vw)x+(u-1-v-w)x^{2}+x^{3},$

so that

$\displaystyle\frac{1}{6}\frac{P^{\prime}(x)-P^{\prime}(u)}{x-u}-\frac{P(x)-P^{\prime}(u)(x-u)}{(x-u)^{2}}$
$\displaystyle=$
$\displaystyle\frac{1}{6}(5u^{2}-5u^{3}-5uv+5u^{2}v-5uw+5u^{2}w+4vw-5uvw)$
$\displaystyle+$ $\displaystyle\frac{1}{6}(5u-5u^{2}-3v+5uv-3w+5uw-3vw)x$
$\displaystyle+$ $\displaystyle\frac{1}{6}(2-5u+2v+2w)x^{2}-\frac{x^{3}}{6}.$

In particular, the matrix $M$ whose column vectors are $\frac{dx}{y^{5}},\ \
P^{\prime}(a)\nabla_{\partial_{a}}\frac{dx}{y^{5}},\ \
P^{\prime}(b)\nabla_{\partial_{b}}\frac{dx}{y^{5}},\ \
P^{\prime}(c)\nabla_{\partial_{c}}\frac{dx}{y^{5}}$ written in the basis
$dx/y^{5},xdx/y^{5},x^{2}dx/y^{5},x^{3}dx/y^{5}$ is

$M=\tiny{\left(\begin{array}[]{cccc}1&\ast&\ast&\ast\\\
0&\frac{1}{6}(5a-5a^{2}-3b+5ab-3c+5ac-3bc)&\frac{1}{6}(5b-5b^{2}-3a+5ab-3c-3ac+5bc)&\frac{1}{6}(5c-5c^{2}-3a+5ac-3b-3ab+5bc)\\\
0&\frac{1}{6}(2-5a+2b+2c)&\frac{1}{6}(2+2a-5b+2c)&\frac{1}{6}(2+2a+2b-5c)\\\
0&-\frac{1}{6}&-\frac{1}{6}&-\frac{1}{6}\end{array}\right)}$

Since the determinant of $M$ is

$\det(M)=-\frac{91}{216}(a-b)(a-c)(b-c),$

the proof of Theorem 1.3 is now complete.

######  Remark 6.1.

The determinant of $\frac{dx}{y^{5}},\ \
\nabla_{\partial_{a}}\frac{dx}{y^{5}},\ \
\nabla_{\partial_{b}}\frac{dx}{y^{5}},\ \
\nabla_{\partial_{c}}\frac{dx}{y^{5}}$ in the basis
$dx/y^{5},xdx/y^{5},x^{2}dx/y^{5},x^{3}dx/y^{5}$, which is somehow more
natural, is given by the expression

$-\frac{91}{216}\frac{(a-b)(a-c)(b-c)}{P^{\prime}(a)P^{\prime}(b)P^{\prime}(c)}=\frac{91}{216}\frac{1}{abc(a-1)(b-1)(c-1)(a-b)(a-c)(b-c)}$

## References

  * [1]
  * [2] A. Avila, C. Matheus and J.-C. Yoccoz, The Kontsevich-Zorich cocycle over Veech-McMullen family of symmetric translation surfaces, J. Mod. Dyn. 14 (2019), 21–54.
  * [3] B. Bakker, B. Klingler and J. Tsimerman, Tame topology of arithmetic quotients and algebraicity of Hodge loci, J. Amer. Math. Soc. 33 (2020), 917–939.
  * [4] G. Baldi, B. Klingler and E. Ullmo, On the distribution of the Hodge locus, preprint (2021) available at http://arxiv.org/abs/2107.08838.
  * [5] C. Herbert Clemens, A scrapbook of complex curve theory, University Series in Mathematics. Plenum Press, New York-London, 1980. ix+186 pp.
  * [6] A. Eskin and M. Mirzakhani, Invariant and stationary measures for the $\textrm{SL}(2,\mathbb{R})$ action on moduli space, Publ. Math. Inst. Hautes Études Sci. 127 (2018), 95–324.
  * [7] A. Eskin, M. Mirzakhani and A. Mohammadi, Isolation, equidistribution, and orbit closures for the $\textrm{SL}(2,\mathbb{R})$ action on moduli space, Ann. of Math. (2) 182 (2015), 673–721.
  * [8] G. Forni, C. Matheus and A. Zorich, Square-tiled cyclic covers, J. Mod. Dyn. 5 (2011), 285–318.
  * [9] E. Gutkin and C. Judge, Affine mappings of translation surfaces: geometry and arithmetic, Duke Math. J. 103 (2000), 191–213.
  * [10] B. Klingler and L. Lerer, Abelian differentials and their periods: the bi-algebraic point of view, preprint (2022) available at arXiv:2202.06031.
  * [11] B. Klingler, E. Ullmo and A. Yafaev, Bi-algebraic geometry and the André-Oort conjecture, Algebraic geometry: Salt Lake City 2015, 319–359, Proc. Sympos. Pure Math., 97.2, Amer. Math. Soc., Providence, RI, 2018.
  * [12] C. Matheus and J.-C. Yoccoz, The action of the affine diffeomorphisms on the relative homology group of certain exceptionally symmetric origamis, J. Mod. Dyn. 4 (2010), 453–486.
  * [13] C. McMullen, Braid groups and Hodge theory, Math. Ann. 355 (2013), 893–946.
  * [14] M. Möller, Shimura and Teichmüller curves, J. Mod. Dyn. 5 (2011), 1–32.
  * [15] C. Voisin Théorie de Hodge et géométrie algébrique complexe, Cours Spécialisés, 10. Société Mathématique de France, Paris, 2002. viii+595 pp.


# Retrieval Head Mechanistically Explains Long-Context Factuality

Wenhao Wu${}^{\text{\textlambda}}$ Yizhong Wang${}^{\text{\textdelta}}$
Guangxuan Xiao${}^{\text{\textsigma}}$ Hao Peng${}^{\text{\textpi}}$ Yao
Fu${}^{\text{\textmu}}$
${}^{\text{\textlambda}}$Peking University ${}^{\text{\textdelta}}$University
of Washington ${}^{\text{\textsigma}}$MIT ${}^{\text{\textpi}}$UIUC
${}^{\text{\textmu}}$University of Edinburgh
<EMAIL_ADDRESS><EMAIL_ADDRESS><EMAIL_ADDRESS>
https://github.com/nightdessert/Retrieval_Head

###### Abstract

Despite the recent progress in long-context large language models (LLMs), it
remains elusive how these transformer-based language models acquire the
capability to retrieve relevant information from arbitrary locations within
the long context. This paper aims to address this question. Our systematic
investigation across 4 model families, 6 model scales, and 3 types of
finetuning reveals that a special type of attention heads are largely
responsible for retrieving relevant information from long context, which we
dub _retrieval heads_. We identify important and intriguing properties of
retrieval heads: (1) universal: all the explored models with long-context
capability have a set of retrieval heads; (2) sparse: only a small portion
(less than 5%) of the attention heads are retrieval. (3) intrinsic: retrieval
heads already exist in models pretrained with short context. When extending
the context length to 32-128K by continual pretraining, it is still the same
set of heads that perform information retrieval. (4) dynamically activated:
take Llama-2 7B for example, 12 retrieval heads always attend to the required
information no matter how the context is changed. The rest of the retrieval
heads are activated in different contexts. (5) causal: completely pruning
retrieval heads leads to failure in retrieving relevant information and
results in hallucination, while pruning random non-retrieval heads does not
affect the model’s retrieval ability. We further show that retrieval heads
strongly influence chain-of-thought (CoT) reasoning, where the model needs to
frequently refer back the question and previously-generated context.
Conversely, tasks where the model directly generates the answer using its
intrinsic knowledge are less impacted by masking out retrieval heads. These
observations collectively explain which internal part of the model seeks
information from the input tokens. We believe our insights on retrieval heads
foster future research on reducing hallucination, improving reasoning, and
compressing the KV cache.

##  1 Introduction

This work studies the internal mechanism of how long-context language models
can utilize information at arbitrary locations within the input. Recent
advances in long-context language modeling [1, 20, 6] show inspiring results,
particularly on the Needle-in-a-Haystack test [14], which asks the model to
precisely retrieve the information of a short sentence (the needle) within a
long context (the haystack). Such capability is the basis of more advanced
long-context tasks, which usually interleaves retrieval and reasoning in a
multi-step fashion [17]. Based on extensive experiments across 4 model
families, 6 model scales, and 3 types of finetuning, we show that across the
models’ attention layers, there exist a small number of retrieval heads that
search the information being asked, and redirect the relevant tokens from the
input to the output. Activation of retrieval heads explains whether the output
is factual or hallucinated. When such heads are activated, the model behaves
faithful to the input document. When they are not activated, or intentionally
masked out in controlled experiments (Fig. 1), the model cannot find the
relevant information and hallucinate instead.

Figure 1: Retrieval heads are the heads that redirect information from the
input to the output. Left: masking out the top retrieval heads of LLaMA 2 7B
80K, its Needle-in-a-Haystack performance drops significantly, and the model
hallucinates during decoding. Right: masking out random non-retrieval heads
does not influence the model’s needle-in-a-haystack behavior. We further note
that the retrieval head only influence factuality but not the language
capability: when they are masked, the model hallucinates by saying “go to the
beach”, which is still a fluent sentence, but not factual (as the correct
answer is “eating a sandwich at Dolores park.”).

The discovery of the retrieval head is motivated by the question of what the
attention mechanism is doing when the model can or cannot find the given
needle. We take important inspiration from two existing works: the CopyNet
[10] and the Induction Head [19]. The CopyNet is essentially a single-layer,
single-head attention mechanism in the age of RNNs that copy-paste tokens from
the input to the output. Induction Heads [19] are special heads within a
multi-layer, multi-head attention network that implements an implicit program
induction algorithm. Combining the observation from the two works, we natually
hypothesize that, just like induction heads are accountable for in-context
learning, there might exist special heads that are accountable for information
retrieval and implement a conditional copy-paste algorithm.

We design algorithms to detect retrieval heads within the transformer
architecture (Sec. 2), and conduct large-scale experiments to demonstrate
important properties of them (Sec. 3): (1) retrieval heads are universal and
sparse: for any model family (LLaMA [21], Yi [25], QWen [2] and Mistral [12]),
at any scale (6B, 14B, and 34B and 8$\times$7B), either base or chat, either
dense or MoE, as long as the model can precisely recite the input information,
they have a small number of retrieval heads (Fig. 1); (2) they are intrinsic:
the base model (e.g., LLaMA2 base) already contains retrieval heads (as a
consequence of large-scale pretraining). Subsequent derivations, such as the
long-context continue pretraining (LLaMA2 7B 80K), chat fine-tuning (Qwen
Chat), or even sparse upcycling [16, 13] uses the same retrieval heads as the
base model (Fig. 5); (3) they are dynamically activated according to the
context: the strongest retrieval heads (e.g., 13 for LLaMA 2 7B) are always
activated no matter what the required information is, while weaker retrieval
heads are activated on different parts of the required information;
consequently these heads compensate each other’s functionality: removing a
subset of the heads, the model at least retrieve part of the required
information; (4) the retrieval heads are causal: say we put a needle "the best
thing to do in San Francisco is to eat a sandwich in Dolores Park on a sunny
day", completely masking out retrieval heads, the model hallucinates (by
saying the best thing is to visit Golden Gate bridge); partially masking out
the heads, the model retrieves part of the needle (e.g., it gets the sandwich
but forget the Dolores Park); masking out random non-retrieval heads, the
model still find full needle; when we do not mask the head yet the model still
hallucinate in some cases, the retrieval heads are not activated. We further
note that chain-of-thought reasoning also heavily relies on retrieval heads
because the model needs to refer back the input information, indicating a
complex relationship between the model’s retrieval and reasoning capability.

The discovery of retrieval head has profound implications on long-context
modeling: (1) it marks a significant step forward in the field of mechanistic
interpretability [3, 19] because for the first time we pin point a particular
subnet implementing the conditional retrieval algorithm; (2) it explains why
certain context-compression methods fail to keep factuality (because they
removes the retrieval head, e.g., in Xiao et al. 24), and suggests future
research on KV cache compression [7, 15], a key problem for deploying long-
context models, should consider the influence of retrieval heads.

##  2 Detecting Retrieval Head

Figure 2: An attention head that performs a copy-paste operation if the token that the head attend to is the same the token being generated. The retrieval score of a head is defined as the frequency of this head’s copy-paste behavior when answering quesions that asks for raw information from the input.  Table 1:  We consider a wide range of language model families and show that the basic properties of retrieval heads are universal and consistent across all language models we study.  Base Model | Variant | Variation Type
---|---|---
Llama-2-7B | Llama-2-7B-80K | Length Extension via Continue Pretrain
| Llama-2-13B-64K | Model Scaling and Length Extension
Mistral-7B-v0.2 | Mistral-7B-Instruct-v0.2 | SFT and RLHF
| Mixtral-8x7B-v0.1 | Sparse Upcycling to Mixture of Experts
Yi-6B | Yi-6B-200K | Length Extension via Continue Pretrain
| Yi-34B-200K | Model Scaling and Length Extension
Qwen1.5-14B | Qwen1.5-14B-Chat | SFT and RLHF

To detect which head is implementing the retrieval algorithm, we introduct a
retrieval score to measures the frequency of a head’s copy-paste behavior
during autoregressive decoding. An attention head with high retrieval score
suggests that statistically across various contexts, this head is frequently
copying the input tokens from the input to the output.

Needle-in-a-Haystack Our retrieval head detection algorithm roots from the
needle-in-a-Haystack test, which asks the model to copy-paste the input tokens
to the output. Given a question ${\bm{q}}$ and its corresponding answer
${\bm{k}}$ (the needle), we insert ${\bm{k}}$ in a given context ${\bm{x}}$
(the haystack) at a random position index range ${\bm{i}}_{q}$. The language
model is then tasked with answering ${\bm{q}}$ based on the haystack with the
inserted needle. We set ${\bm{q}}$ and ${\bm{k}}$ unique and irrelevant with
the given long context, ensuring that if an answer is correctly generated, it
is indeed copied from the context, not from the model’s internal knowledge.

Retrieval Score for Attention Heads We define the retrieval score as the
frequency of a head’s copy-paste operations. Specifically, during auto-
regressive decoding (we use greedy decoding by default), denote the current
token being generated as $w$ and the attention scores of a head as
$\bm{a}\in\mathcal{R}^{|{\bm{x}}|}$. As demonstrated in Fig. 2, we say an
attention head $h$ copies and pastes a token from the needle to the output
sentence if it follows two criteria: (1) $w\in{\bm{k}}$, i.e., $w$ is a token
within the needle sentence. (2)
${\bm{x}}_{j}=w,j=\arg\max(\bm{a}),j\in{\bm{i}}_{q}$, i.e., the input token
that receives the most attention probability mass by this head is a token
within the needle and is the same token as the currently generated token. Let
${\bm{g}}_{h}$ be the set containing all tokens copy and pasted by a given
head $h$, we define:

$\text{Retrieval score for
head}\;\;h=\frac{|{\bm{g}}_{h}\cap{\bm{k}}|}{|{\bm{k}}|},$ (1)

Intuitively, retrieval score represents a token-level recall rate of the most
attended tokens by an attention head. For example, when retrieving a needle of
10 tokens, a retrieval score of 0.9 indicates that the attention head has
copies and pasted 9 tokens in the 10-token target answer.

Retrieval Head Detection Algorithm We calculate the retrieval score for all
attention heads under a diverse set of input contexts. For each language model
we consider, we compile three sets of Needle-in-a-Haystack samples, each
consisting of a unique tuple $({\bm{q}},{\bm{k}},{\bm{x}})$. For each sample,
we make sure $({\bm{q}},{\bm{k}})$ is semantically irrelevant with ${\bm{x}}$
and that ${\bm{q}}$ cannot be answered using the model’s existing knowledge by
manually inspecting the model output. Then for each
$({\bm{q}},{\bm{k}},{\bm{x}})$ sample, we perform Needle-in-a-Haystack on 20
different length values uniformly sampled from 1K-50K, where in each length,
${\bm{q}}$ is inserted in 10 different depth uniformly ranging from the start
to the end of ${\bm{x}}$. We note that this scale of tests gives stable
outputs as the average retrieval score converges after just a few samples. In
total, each language model is subjected to approximately 600 instances of
retrieval testing. We calculate the retrieval score for each attention head in
each test and use the average of these scores as the head’s final retrieval
score. The attention heads with relatively larger retrieve score can be
considered as retrieval head. In our case (Fig. 3), we set the threshold as
0.1, meaning that as long as the head performs copy-paste in 10% of the times,
we consider it a retrieval head.

##  3 Basic Properties of Retrieval Heads

Figure 3: Retrieval heads are universal and sparse across model family and
scale. For all models we consider, less than 5% of the attention heads are
activated more than 50% of the time (with a retrieval score higher than 0.5)
when retrieval is required.

This section discusses important properties of retrieval heads: (1) universal
and sparse: any model that exhibits long-context capability has a small set of
retrieval heads; (2) dynamic: most of retrieval heads are activated under
different contexts; (3) intrinsic: retrieval heads are already within the base
model as a consequence of large-scale pretraining. Subsequent models reuse the
same set of heads. Our results are supported by extensive experiments on a
large spectrum of models (Table 1). To investigate the influence of continued
pretraining for context length extension, we compare Llama-2-7B 4K to
Llama-2-7B-80K and Llama-2-13B-60K [6]. To examine the effect of alignment, we
have study Mistral-7B-Instruct-v0.2 and Qwen-1.5-14B-Chat [2] and compare them
to their base versions. We further choose Mixtral-8x7B-v0.1 [13], a mixture of
expert versions derived from Mistral-7B-v0.2, presumably via sparse upcycling
[16], to study retrieval heads in different architectures.

###  3.1 Universal and Sparse

Figure 3 demonstrate that a sparse set of retrieval heads exist in all models
we consider, regardless of the various pertraining, fine-tuning recipes and
the underlying architecture. Between 25% and 52% of the heads exhibit copy-
paste behaviors at a very low frequency, with a score of between 0 and 0.1.
Approximately 45% to 73% of attention heads have 0 retrieval score, meaning
that they have other functionality than retrieval. Only about 3% to 6% of the
attention heads have a retrieval score larger than 0.1 (recall that a
retrieval score 0.1 means a head retrieves at least 10% of the tokens being
asked). It is also intriguing that although all model parameters, particularly
the total number to attention heads, are largely different from each other,
their ratio of retrieval heads lays in the same interval (about 5%).

Figure 4: Retrieval score (blue): average number of activated tokens;
Activation Frequency (green): frequency of at least activated on one token.
The gap between the two curves shows context-sensitivity: a head of high
activation frequency but low retrieval score means it is only activated on
certain tokens and contexts; a head of high retrieval score means it is
activated on almost any context. For both LLaMA and Yi, there exist strong
heads that are not sensitive to context and always activated.

###  3.2 Dynamically Activated Based on Tokens and Contexts

Now we study how sensitive a retrieval head is to its input context, i.e.,
whether a head is consistently activated no matter what the context is, or if
a head is activated only on specific contexts. For the needle sentences "the
best thing to do in San Francisco is eating a sandwich in Dolores park in a
sunny day", some heads are activated on the full sentence, whereas other heads
only activated on certain tokens like “eating a sandwich” or “in Dolores
park’. We define activation frequency, the frequency of a head being activated
on at least one token (v.s., the retrieval score measures the average number
of activated tokens). A head of high activation frequency but low retrieval
score means it is only activated on certain tokens and contexts. As is shown
in Fig. 4, Llama-2-7B-80K and Yi-6B-200K have 12 and 36 strongest retrieval
heads, respectively, that are always activated (activation frequency equal to
1) under all the contexts we consider. Weaker heads only activate on certain
tokens and contexts.

###  3.3 Intrinsic

Figure 5: Retrieval head is intrinsic and already within the base model. The
subsequent model derivations, either by continue pretraining (LLaMA 2 7B 80K)
or chat finetuning (Qwen 1.5 14B Chat) or sparse upcycling (Mixtral
8$\times$7B), use the same set of retrieval head as the base model, as
demonstrated by a high level of similarity between the heatmap patterns.

We show that the retrieval heads, thus the ability of utilizing information at
arbitrary location of the input, is an intrinsic property [6] of the base
model as a consequence of large-scale pretraining, with subsequent small-scale
training exerting only minor alterations to these head activation patterns. In
Figure 5, we present the retrieval score distributions for a range of base
models in the initial row, followed by their corresponding variants in the
subsequent row. We see that regardless of the models being continuously pre-
trained, chat fine-tuned, or sparsely upcycled, there is a notable consistency
in their retrieval scores heatmaps. Figure 7 offers a more direct and strict
examination, where we compute the statistical correlations between different
models. The data reveal a high degree of correlation in the retrieval score
distributions between base models and their respective variants, with a
Pearson correlation coefficient exceeding 0.8. Models from different families
exhibit a correlation coefficient of less than 0.1, indicative of their
distinct pretraining recipes.

Figure 6: The retrieval heads of models of the same family are strongly
correlated, i.e., the chat model and base model typically uses the same set of
retrieval heads. The retrieval heads of models of different families are
clearly different.

Figure 7: Masking out top K retrieval heads vs K random heads. For all models
are consider, the removal of retrieval heads clearly reduces the Needle-in-a-
Haystack performance, versus the removal of non-retrieval heads have much
weaker influence.

##  4 Influence on Downstream Tasks

This section examines how retrieval heads influence downstream tasks. Across
the experiments we use Mistrial-7B-Instruct-v0.2 [18] as it is a popular and
strong open language model with 32K context length. We first show that
retrieval heads explains the factuality of Needle-in-a-Haystack test. When the
model can retrieve the needle, retrieval heads are always activated. When the
model cannot retrieve the needle and hallucinate instead, retrieval heads are
either partially activated or not activated. Then we show that retrieval heads
significantly influence question answering that requires extracting the
information from the input, but does not strongly influence tasks where the
model directly produce answers based on its internal knowledge. We further
explore how retrieval heads influence more sophisticated reasoning behaviors
like chain-of-thought [23].

###  4.1 Retrieval Heads Explains Factuality in Needle-in-a-Haystack

We start with a closer look at Needle-in-a-Haystack and construct an
additional set of needle tests that are different from the three sets used for
retrieval head detection. We gradually mask out the number of retrieval/
random heads and see how the model’s behavior changes. As is shown in Fig. 7,
masking out retrieval heads severely damages the model’s Needle-in-a-Haystack
performance, while masking out random heads shows much smaller performance
impact. Notably, when increasing the number of masked heads K to 50 (about 5%
of the full number of heads), all models’ needle test performance drop to
below 50, showing that the top retrieval heads are responsible for most of the
needle retrieval behavior.

Figure 8: Masking out retrieval heads severely damages ExtractQA performance.

Figure 9: When the model fails to retrieve the full needle, there are two
typical errors: (1) incomplete retrieval, where the retrieval heads miss part
of the information “eat a sandwich”; (2) hallucination, where the retrieval
heads attend to the initial tokens.

We observe three types of error: (1) Incomplete retrieval, where models only
captured partial information of the target and omitted key details; (2)
Hallucination, where models generated hallucinated sentences; and (3) Wrong
extraction, where models incorrectly retrieved irrelevant content from the
haystack. Without any masking, instances of wrong extraction occurred when
retrieval heads were active but focused on incorrect sections. During
hallucinated generation, retrieval heads predominantly attended to the initial
token of the input, which is often known as "attention sink" [24], presumably
dummy tokens that provide less meaningful signal.

As we increase the number of masked heads, initially, a small subset of the
most robust retrieval heads are masked, and incomplete retrieval began to
appear. In the absence of the strongest retrieval heads, the remaining weaker
heads only managed to retrieve a fraction of the target information.
Metaphorically, each retrieval head holds a small piece of the "needle," yet
these pieces cannot form a complete one, resulting in an incomplete final
output. This phenomenon typically begins when the mask out heads of retrieval
score larger than 0.4. As we further increase the number of mask,
hallucinations become more prevalent, signaling a complete failure of the
retrieval capability.

###  4.2 Influence on Extractive QA

Now we study how retrieval heads influence more realistic tasks beyond Needle-
in-a-Haystack. We use extractive QA as a test bed, as common usecase of long-
context model where the user typically upload a pdf (research papers,
finantial reports, legal documents, .etc) and ask questions about specific
information within the document. To make sure the knowledge being asked does
not exist in the model’s internal knowledge, we synthesize an extractive QA
dataset by selecting a set of up-to-date news articles, extract a paragraph
from it, and asking GPT-4 to produce a question-answer pair based on the
extracted paragraph, similar to the evaluation conducted in Anthropic [1]. As
illustrated in Figure 9, randomly masking out non-retrieval heads demonstrated
no significant impact on performance. Masking out retrieval heads led to a
substantial decrease in F1 scores, with reductions of 9.2% and 23.1%. These
observations demonstrate that real-world document QA tasks heavily rely on the
functionality of retrieval heads.

###  4.3 Chain-of-Thought Reasoning also Requires Retrieval Heads

Figure 10: Retrieval heads significantly influence tasks that require chain-
of-thought reasoning. This is because typically in a reasoning chain, the next
step reasoning requires the model to refer to previous information. See Fig.
11 for examples.  Figure 11:  When we mask out retrieval heads, the model
becomes “blind” to important information in the question description resulting
in incorrect reasoning chains.

We test Mistrial-7B-Instruct-v0.2’s performance on MMLU [11], MuSiQue and
GSM8K [4], with and without chain-of-thought reasoning. As is shown in Fig.
10, if we use answer-only prompting (without CoT), masking out either
retrieval or random heads do not really influence the performance, presumably
because the model’s generation is based on its internal knowledge primarily
stored in the FFN layers [8]. For CoT styled reasoning, masking out retrieval
heads signifantly influence the model’s performance. Upon inspecting typically
error cases (Fig. 11), we find out the model becomes “blind” to important
input information and hallucinate instead. We find the relationship between
CoT and retrieval heads particularly intriguing as it may offers deeper
insights into model’s complex reasoning performance. We leave more in-depth
studies to future research.

##  5 Discussions

General Functionalities of Attention Heads For transformer language models, we
tend to view the functionality of FNNs layers to be the place for storing
knowledge [8], and the attention layers to be the place for implementing
algorithms [19]. The induction head discussed in Olsson et al. [19] typically
searches repeated patterns of the input, which is at a certain level similar
to the retrieval heads (as it also searches and repeats information).
Different than the induction heads, the retrieval heads are typically
responsible for redirecting the information according to the context, but do
not for inferring programs. We tend to believe that there exist more algorithm
and functionalities implemented by other types of attention heads to be
discovered by future research.

Relationship to Local and Linear Attention and State-Space Models Although
there exist numerous works about local [24] / linear [22] attention, state
space models [9], and hybrid architectures [5] achieving inspiring efficiency
in long-context modeling, so far there is no linear attention / SSM
architecture that passes the Needle-in-a-Haystack test to the best of our
knowledge, suggesting that the full attention might be a must for long-context
information retrieval. One example is that the Mistral v0.1 [12] uses sliding
window attention but cannot pass needle-in-a-haystack, and their authors
changes the attention to full in v0.2 [18], then it can pass the needle test.
Our results showing strong evidence why full attention is a must. For the
model to precisely utilize input information at arbitrary location, it is
crutial for the retrieval heads to work on the full KV cache.

Applications to KV Cache Compression The problem that the KV cache is too
large and occupies a large chunk of the GPU memory severely hinders the
deployment of long-context models. For example, for LLaMA 2 7B, the KV cache
of 100K tokens requires more than 50GB memory, while 2K context requires less
than 1GB memory. If we serve this model on one 80G A100, then the concurrency
of 100K context can be 50 times less than 2K context queries, which is
prohibitively expensive. The results from this work indicates that we might be
possible to radically prune out the KV cache corresponding to the non-
retrieval heads (recall in Fig. 3 shows only 5% of the heads are retrieval)
and significantly reducing the deployment cost of long-context models. We
leave this study to future research.

##  6 Conclusions

This paper discovers retrieval heads, a special set of attention heads that
are responsible for implementing the conditional copy algorithm and redirect
information from the input to the output. Retrieval heads are the primarily
reason why a successful long-context model can pass the Needle-in-a-Haystack
test, and their activation explains why a language model is faithful to the
input or hallucinate. Compared to non-retrieval heads, retrieval heads have a
stronger influence on downstream tasks that require the model to precisely
recall the input information, either in extractive question answering or
chain-of-thought reasoning. We believe this work will foster future research
on reducing hallucination, improving reasoning, and compressing the KV cache.

## References

  * Anthropic [2023] Anthropic.  Model card and evaluations for claude models, July 2023.  URL https://www.anthropic.com/product.
  * Bai et al. [2023] Jinze Bai, Shuai Bai, Yunfei Chu, Zeyu Cui, Kai Dang, Xiaodong Deng, Yang Fan, Wenbin Ge, Yu Han, Fei Huang, et al.  Qwen technical report.  _arXiv preprint arXiv:2309.16609_ , 2023.
  * Bricken et al. [2023] Trenton Bricken, Adly Templeton, Joshua Batson, Brian Chen, Adam Jermyn, Tom Conerly, Nick Turner, Cem Anil, Carson Denison, Amanda Askell, Robert Lasenby, Yifan Wu, Shauna Kravec, Nicholas Schiefer, Tim Maxwell, Nicholas Joseph, Zac Hatfield-Dodds, Alex Tamkin, Karina Nguyen, Brayden McLean, Josiah E Burke, Tristan Hume, Shan Carter, Tom Henighan, and Christopher Olah.  Towards monosemanticity: Decomposing language models with dictionary learning.  _Transformer Circuits Thread_ , 2023.  https://transformer-circuits.pub/2023/monosemantic-features/index.html.
  * Cobbe et al. [2021] Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman.  Training verifiers to solve math word problems, 2021.
  * De et al. [2024] Soham De, Samuel L Smith, Anushan Fernando, Aleksandar Botev, George Cristian-Muraru, Albert Gu, Ruba Haroun, Leonard Berrada, Yutian Chen, Srivatsan Srinivasan, et al.  Griffin: Mixing gated linear recurrences with local attention for efficient language models.  _arXiv preprint arXiv:2402.19427_ , 2024.
  * Fu et al. [2024] Yao Fu, Rameswar Panda, Xinyao Niu, Xiang Yue, Hannaneh Hajishirzi, Yoon Kim, and Hao Peng.  Data engineering for scaling language models to 128k context.  _arXiv preprint arXiv:2402.10171_ , 2024.
  * Ge et al. [2023] 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.  _arXiv preprint arXiv:2310.01801_ , 2023.
  * Geva et al. [2020] Mor Geva, Roei Schuster, Jonathan Berant, and Omer Levy.  Transformer feed-forward layers are key-value memories.  _arXiv preprint arXiv:2012.14913_ , 2020.
  * Gu and Dao [2023] Albert Gu and Tri Dao.  Mamba: Linear-time sequence modeling with selective state spaces.  _arXiv preprint arXiv:2312.00752_ , 2023.
  * Gu et al. [2016] Jiatao Gu, Zhengdong Lu, Hang Li, and Victor O.K. Li.  Incorporating copying mechanism in sequence-to-sequence learning.  In Katrin Erk and Noah A. Smith, editors, _Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)_ , pages 1631–1640, Berlin, Germany, August 2016. Association for Computational Linguistics.  doi: 10.18653/v1/P16-1154.  URL https://aclanthology.org/P16-1154.
  * Hendrycks et al. [2020] Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt.  Measuring massive multitask language understanding.  In _International Conference on Learning Representations_ , 2020.
  * Jiang et al. [2023] Albert Q Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, et al.  Mistral 7b.  _arXiv preprint arXiv:2310.06825_ , 2023.
  * Jiang et al. [2024] Albert Q Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, et al.  Mixtral of experts.  _arXiv preprint arXiv:2401.04088_ , 2024.
  * Kamradt [2023] Greg Kamradt.  Needle in a haystack - pressure testing llms.  https://github.com/gkamradt/LLMTest_NeedleInAHaystack, 2023.
  * Kang et al. [2024] Hao Kang, Qingru Zhang, Souvik Kundu, Geonhwa Jeong, Zaoxing Liu, Tushar Krishna, and Tuo Zhao.  Gear: An efficient kv cache compression recipefor near-lossless generative inference of llm.  _arXiv preprint arXiv:2403.05527_ , 2024.
  * Komatsuzaki et al. [2022] Aran Komatsuzaki, Joan Puigcerver, James Lee-Thorp, Carlos Riquelme Ruiz, Basil Mustafa, Joshua Ainslie, Yi Tay, Mostafa Dehghani, and Neil Houlsby.  Sparse upcycling: Training mixture-of-experts from dense checkpoints.  _arXiv preprint arXiv:2212.05055_ , 2022.
  * Kuratov et al. [2024] Yuri Kuratov, Aydar Bulatov, Petr Anokhin, Dmitry Sorokin, Artyom Sorokin, and Mikhail Burtsev.  In search of needles in a 10m haystack: Recurrent memory finds what llms miss.  _arXiv preprint arXiv:2402.10790_ , 2024.
  * Mistral [2024] Mistral.  Model card for mistral-7b-instruct-v0.2, April 2024.  URL https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.2.
  * Olsson et al. [2022] Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, et al.  In-context learning and induction heads.  _arXiv preprint arXiv:2209.11895_ , 2022.
  * Reid et al. [2024] Machel Reid, Nikolay Savinov, Denis Teplyashin, Dmitry Lepikhin, Timothy Lillicrap, Jean-baptiste Alayrac, Radu Soricut, Angeliki Lazaridou, Orhan Firat, Julian Schrittwieser, et al.  Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context.  _arXiv preprint arXiv:2403.05530_ , 2024.
  * Touvron et al. [2023] Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al.  Llama 2: Open foundation and fine-tuned chat models.  _arXiv preprint arXiv:2307.09288_ , 2023.
  * Wang et al. [2020] Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma.  Linformer: Self-attention with linear complexity.  _arXiv preprint arXiv:2006.04768_ , 2020.
  * Wei et al. [2022] Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Brian Ichter, Fei Xia, Ed H. Chi, Quoc V. Le, and Denny Zhou.  Chain-of-thought prompting elicits reasoning in large language models.  In Sanmi Koyejo, S. Mohamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh, editors, _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/9d5609613524ecf4f15af0f7b31abca4-Abstract-Conference.html.
  * Xiao et al. [2023] Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis.  Efficient streaming language models with attention sinks.  _arXiv preprint arXiv:2309.17453_ , 2023.
  * Young et al. [2024] Alex Young, Bei Chen, Chao Li, Chengen Huang, Ge Zhang, Guanwei Zhang, Heng Li, Jiangcheng Zhu, Jianqun Chen, Jing Chang, et al.  Yi: Open foundation models by 01. ai.  _arXiv preprint arXiv:2403.04652_ , 2024.