&=\frac{1}{\hbar^{2}}\sum_{i}\frac{I_{i}(I_{i}+1)}{3}\phi_{i}A_{i}^{2}\frac{v_{0}}{V_{D}}.\end{split}$
(111)

The parameter $\kappa$ has the units of time and describes the restoring force
that keeps the random walk bounded. Specifically, the expectation value of the
distribution in Eq. (109a) is

$\overline{\omega}\equiv\int\mathrm{d}\omega K(\omega_{0},\omega,\delta
t)=\omega_{0}e^{-\delta t/\kappa}.$ (112)

The Taylor expansion at short time $\delta t\ll\kappa$ gives

$\overline{\omega-\omega_{0}}=-\frac{\omega_{0}}{\kappa}\delta t.$ (113)

Since the expected average change is proportional to the time interval, the
proportionality factor corresponds to a polarization rate. Further, taking the
limit $\delta t\to 0$,

$\overline{\partial_{t}\omega}=-\frac{\omega}{\kappa},$ (114)

the equation expresses a restoring force, since the detuning frequency is
pulled back to the ‘equilibrium’ $\omega=0$ in proportion to its instantaneous
deviation from the equilibrium.

Similarly, the variance growth at short times,

$\overline{(\omega-\omega_{0})^{2}}=\sigma_{\delta
t}^{2}+\omega_{0}^{2}(e^{-2\delta
t/\kappa}-1)^{2}\approx\sigma_{\Omega}^{2}\frac{2\delta t}{\kappa},$ (115)

shows that the process is a diffusion with the diffusion constant

$D_{\Omega}=\frac{2\sigma_{\Omega}^{2}}{\kappa}.$ (116)

With these results, we can convert the ‘intrinsic’ thermal fluctuations of the
Overhauser field, which are bounded and well described by a Gaussian
distribution at long times, into the corresponding parameters of the above
model. As we already noted, there are several experimental measurements of
$D_{\Omega}$ in gated GaAs quantum dots. Based on the values given in Footnote
27, we take $D_{\Omega}=(2\pi\times 7\,\mathrm{kHz})^{2}/1\,\mu$s as a
representative value. Equation (116) then gives $\kappa=0.2$ s (we evaluated
$\sigma_{\Omega}^{2}$ from Eq. (111) using our parameters), predicting the
Overhauser field equilibration scale in seconds.

In silicon, we are not aware of a direct experimental measurement of the
diffusion constant of the quantum dot Overhauser field, $D_{\Omega}$. To
arrive at an estimate, we use the result of Ref. [98], which, using the
methods of Refs. [99, 100], derives the time-correlation of the Overhauser
field, converted to angular frequency units as

$\overline{\omega(t)\omega(t+\delta
t)}=\frac{\sigma_{\Omega}^{2}}{\sqrt{2\pi}}\prod_{\alpha\in\\{x,y,z\\}}\left(1+2Dl_{\alpha}^{-2}|\delta
t|\right)^{-1/2},$ (117)

where $D$ is the material bulk nuclear spin diffusion constant. Taylor-
expanding for short times $\delta t\to 0$, we get

$D_{\Omega}=\frac{\sigma_{\Omega}^{2}}{\sqrt{2\pi}}D(l_{x}^{-2}+l_{y}^{-2}+l_{z}^{-2}),$
(118)

and finally

$\kappa=\frac{\sqrt{8\pi}}{D}\left(l_{x}^{-2}+l_{y}^{-2}+l_{z}^{-2}\right)^{-1}.$
(119)

Cross-checking the value for GaAs, using the bulk diffusion $D=7$ nm2/s
estimated theoretically [101, 102] gives $\kappa=0.001$ s, implying
equilibration time of the order of a minute. The two values delimitate the
range for the expected value of the intrinsic rate $\Gamma_{0}$, which we use
in the caption of Fig. 6 as $0.01$-$0.1$ s-1. Assuming that in silicon the
spin diffusion is slower, with $D=2$ nm2/s measured in Ref. [103], we use an
order of magnitude smaller rates, $\Gamma_{0}\sim 0.001$-$0.01$ s-1, as an
orientation value in Figs. 11 and 13.

##  Appendix I DNSP in hole qubits

We now apply the results of Tab. 2 to quantum dots with holes. Aiming at rough
estimates, we consider the heavy-hole (HH) limit with the hyperfine
interaction of the Ising form[104]464646Going beyond this simplest limit might
require numerics to evaluate the hole wave function. The hyperfine interaction
tensor is non-generic, given by the details of the confinement potential [39].

$H_{\mathrm{hf}}=\sum_{n}A_{||,n}v_{0}|\Psi(\mathbf{r}_{n},z_{n})|^{2}I_{z}s_{z}.$
(120)

This limit leads to a simple result for the matrix element $X$. Namely, with
$J_{zz}$ being the only non-zero element of the hyperfine tensor,474747 The
light-hole–heavy-hole mixing results in further elements in the hyperfine
tensor. Treating the mixing perturbatively, some of these elements arising in
the first and second order were given in Refs. [81, 105]. Tab. 2 gives
$M^{(0)}_{+-}=0=M^{(-2)}_{+-}$ and
$M^{(-1)}_{+-}=J_{zz}\cos\delta\sin\gamma/4$. It means, first of all, that in
the Ising limit the ‘shaking’ mechanism is not effective, only the
‘deflection’ one contributes,

$\displaystyle X_{\mathrm{df}}^{\mathrm{HH}}$
$\displaystyle=\frac{A_{||}}{4N_{\mathrm{tot}}}\sin\delta\cos\gamma,$ (121)
$\displaystyle X_{\mathrm{sh}}^{\mathrm{HH}}$ $\displaystyle=0.$ (122)

The next difference to electrons is that for holes, apart from the hyperfine
tensor, the $g$-tensor is also strongly anisotropic and the confinement has
strong effects on the hole spin. Important here, the deflection of the
quantization axes of the hole spin and nuclear spins will be most often
dominated by the quantum dot confinement rather than the small gradients of
the magnetic field. The factor $\sin\delta$ is then not necessarily small for
holes. Specifically, consider a quasi-two-dimensional quantum dot with the
strong confinement along the $z$ axis, what fixes the heavy-hole spin along z.
With the magnetic field in the plane, the deflection angle $\delta$ is
$\pi/2$. The factor $\sin\delta$ in Eq. (121) is then 1, rather than
$l\nabla_{\perp}B/B\approx 0.03$ in Eq. (47b), boosting the rate by orders of
magnitude. On the other hand, the nuclei are still polarized in the plane, so
that the arising polarization is not visible as a change in the hole Larmor
frequency. Additional NMR pulses would be needed to detect this polarization
through the hole.

Concerning the material, recent progress with hole qubits [106] motivates us
to consider silicon and germanium atoms for a possible DNSP. The hyperfine
constants in the valence band for the two are similar,484848Our Eq. (120)
corresponds to Eq. (17) of Ref. [107] with $A_{n}=A_{||}$ and the
perpendicular components $A_{\perp}$ neglected. The reference gives
$A_{||}=-2.5$ neV for 29Si and $A_{||}=-1.1$ neV for 73Ge, with $A_{\perp}$
two orders of magnitude smaller. while the 9/2 nuclear spin of 73Ge is much
larger than the 29Si spin 1/2. These numbers would suggest germanium as more
perspective to search for the DNSP signal. However, its low g-factor makes the
nuclear Larmor frequency low, in turn the detection of the Hartmann-Hahn
resonance challenging.

For a SiGe hole qubit, we summarize as follows. Since the hyperfine constant
in the silicon valence band is similar to the one in the conduction band (see
Tab. 1 in Ref. 107), taking a heavy hole with spin along z and the magnetic
field also along z, the resulting DNSP rate is similar to the one for an
electron qubit in silicon. It was plotted in Fig. 10, where only the
deflection mechanism applies for a hole. The resulting rate is low. A somewhat
larger rate arises in germanium atoms, because of the larger nuclear spin.
However, the resonance frequency is low (below 1 MHz for $B=1$ T). Concerning
a possible observation of the DNSP with holes, the most favorable scenario
then looks to be searching for it in silicon atoms with a heavy-hole quantum
dot in an in-plane magnetic field.

Figure 13: Stabilization by feedback in a hole qubit in SiGe. The plot is
analogous to Fig. 6 (see its caption for the description) and adopts the same
parameters except for the atomic parameters of Si and Ge.

We illustrate this case with Fig. 13, plotting the induced rate $\Gamma^{*}$.
For 29Si atoms, the rate can be compared to the analogous plot for an electron
quantum dot, Fig. 11. The lack of suppression due to the factor $\sin\delta$
(with $\sin\delta=1$ for the hole) boosts the rate by three orders of
magnitude. On the other hand, one order of magnitude is offset by a smaller
size of the electron-qubit quantum dot, due to a larger effective mass. As a
result, the difference between the curves for Si in Figs. 11 and 13 is
approximately two orders of magnitude.

As seen in Fig. 13, the rate for 73Ge atoms can become larger than for 29Si.
It is due to a larger nuclear spin of germanium. However, the resonance
happens at a low frequency, so that the resonant peak is not discernible for
Ge in Fig. 13, overwhelmed by the rate behavior at zero frequency.494949We
note that Eq. (58) diverges in the limit $f_{RR}\to 0$. This divergence is
spurious, and stems from the assumption $f_{RR}\gg f_{\Delta}$, which we
adopted in deriving Eqs. (58) and (59). At detunings larger than $f_{RR}$, the
assumption is violated, Eq. (59) does not hold, and the quantity $\Gamma^{*}$,
though still well defined, is of little use. For this reason, we limit the
lowest frequency on the horizontal axis in Figs. 6, 11, and 13 to an ad-hoc
value of 1 MHz. A discernible peak appears for $B=2$ T or higher (not shown),
but such fields might be too high for holes in SiGe to be useful as spin
qubits.

##  Appendix J Notation: list of defined quantities

We collect the definitions of the main symbols used throughout the text for
easier reference and lookup. We group them in the three parts of Tab. J.

Zeeman energies, Larmor frequencies, and related quantities
---
electron | nucleus | description | sign | definition
$q_{e}$ | $q_{n}$ | sign of electric charge | signed |  $q_{e}=-1$, $q_{n}=+1$
$g_{e}$ | $g_{n}$ |  $g$ factor | signed | material parameter
$\mu_{B}$ | $\mu_{N}$ | magneton | positive | nature constant
$s$ | $I$ | spin magnitude | positive |  $s=1/2$, $I=\mathrm{(half)integer}$
$\mathbf{s}$ | $\mathbf{I}$ | spin operator | vector |  $\mathbf{I}^{2}=I(I+1)$, $\mathbf{s}^{2}=s(s+1)$
$\mathbf{B}_{e}$ | $\mathbf{B}_{n}$ | magnetic field | vector | tunable parameter
$\mathbf{z}_{e}$ | $\mathbf{z}_{n}$ | spin ground-state direction | unit vector | $\mathbf{z}=\mathrm{sgn}(qg)\mathbf{B}/B$
$\hbar\omega_{e}$ | $\hbar\omega_{n}$ | Zeeman energy | positive |  $\hbar\omega_{n}=|g_{n}\mu_{n}B|$; for $\hbar\omega_{e}$, see Eq. (10)
$f_{e}$ | $f_{n}$ | Larmor frequency | positive | $f=\omega/2\pi$
$\delta,\delta^{\prime}$ | angles relating $\mathbf{z}_{e}$ and $\mathbf{z}_{n}$ | signed | see Fig. 2
Table III-1:  Quantities related to the Larmor precession speed and the spin
orientation. Atomic and nuclear quantities
---
quantity | description | sign | definition
$a_{0}$ | lattice constant | positive | material parameter
$v_{0}$ | volume per atom | positive | $v_{0}=a_{0}^{3}/8$
$V_{D}$ | quantum dot volume | positive | $V_{D}=1/\int\mathrm{d}V|\Psi|^{4}$
$N_{\mathrm{tot}}$ | number of atoms in the dot | positive | $N_{\mathrm{tot}}=V_{D}/v_{0}$
$\phi_{i}$ | isotopic fraction | positive | material parameter
$N_{i}$ | number of atoms for $i$-th isotope | positive | $N_{i}=N_{\mathrm{tot}}\phi_{i}$
$A_{i}$ | hyperfine constant for $i$-th isotope | signed | $\frac{4}{3v_{0}}\mu_{N}\mu_{B}g_{i}|\eta_{i}|^{2}$
$J_{n}$ | hyperfine coupling strength for nucleus $n$ | signed | $J_{n}=A_{n}v_{0}|\Psi_{n}|^{2}$
$J$ | average hyperfine coupling strength | signed | $J\equiv\langle J_{n}\rangle=A_{i}/N_{\mathrm{tot}}$
Table III-2:  Quantities related to nuclear spins. The relation
$v_{0}=a_{0}^{3}/8$ applies for zinc-blende and diamond crystals. For our
wave-function choice in Eq. (1), the quantum dot volume is $V_{D}=2\pi
l^{2}l_{z}$, see App. A. Finally, $\eta_{i}$ is the amplitude of the electron
Bloch wave function at the atomic nucleus. EDSR related quantities
---
quantity | description | sign | definition
$\mathbf{d}$ | dot shift in space | in-plane vector | $\mathbf{d}=e\mathbf{E}_{0}l^{2}/(\hbar^{2}/ml^{2})$
$\mathbf{b}$ | direction of the EDSR field | unit vector | see Eq. (13)
$\gamma$ | detuning angle | signed | $\sin\gamma=-f_{\Delta}/f_{R}$
$\phi_{\mathrm{rf}}$ | phase shift of the EDSR signal | signed | $\mathbf{E}(t)=\mathbf{E}_{0}\cos(\omega_{\mathrm{rf}}t-\phi_{\mathrm{rf}})$
freq. |  | ang. freq. | $(\omega=2\pi f)$
$f_{\mathrm{rf}}$ |  | $\omega_{\mathrm{rf}}$ | frequency of the EDSR drive | positive | tunable parameter
$f_{\Delta}$ |  | $\omega_{\Delta}$ | detuning frequency | signed | $f_{\Delta}=f_{\mathrm{rf}}-f_{e}$
$f_{RR}$ |  | $\omega_{RR}$ | Rabi frequency at resonance | positive | see Eq. (13)
$f_{R}$ |  | $\omega_{R}$ | Rabi frequency | positive | $f_{R}=\sqrt{(f_{RR})^{2}+f_{\Delta}^{2}}$
Table III-3:  Quantities related to the EDSR drive.

## References

  * [1] Loss, D. & DiVincenzo, D. P.  Quantum computation with quantum dots.  _Physical Review A_ 57, 120–126 (1998).
  * [2] Burkard, G., Ladd, T. D., Nichol, J. M., Pan, A. & Petta, J. R.  Semiconductor Spin Qubits (2021).  eprint 2112.08863.
  * [3] Stano, P. & Loss, D.  Review of performance metrics of spin qubits in gated semiconducting nanostructures.  _Nature Reviews Physics_ 4, 672–688 (2022).
  * [4] Golovach, V. N., Borhani, M. & Loss, D.  Electric-dipole-induced spin resonance in quantum dots.  _Physical Review B_ 74 (2006).
  * [5] Nowack, K. C., Koppens, F. H. L., Nazarov, Y. V. & Vandersypen, L. M. K.  Coherent Control of a Single Electron Spin with Electric Fields.  _Science_ 318, 1430–1433 (2007).
  * [6] Tokura, Y., van der Wiel, W. G., Obata, T. & Tarucha, S.  Coherent Single Electron Spin Control in a Slanting Zeeman Field.  _Physical Review Letters_ 96 (2006).
  * [7] Pioro-Ladrière, M. _et al._ Electrically driven single-electron spin resonance in a slanting Zeeman field.  _Nature Physics_ 4, 776–779 (2008).
  * [8] Ono, K. & Tarucha, S.  Nuclear-Spin-Induced Oscillatory Current in Spin-Blockaded Quantum Dots.  _Physical Review Letters_ 92 (2004).
  * [9] Koppens, F. H. L. _et al._ Control and Detection of Singlet-Triplet Mixing in a Random Nuclear Field.  _Science_ 309, 1346–1350 (2005).
  * [10] Koppens, F. H. L. _et al._ Driven coherent oscillations of a single electron spin in a quantum dot.  _Nature_ 442, 766–771 (2006).
  * [11] Khaetskii, A. V., Loss, D. & Glazman, L.  Electron Spin Decoherence in Quantum Dots due to Interaction with Nuclei.  _Physical Review Letters_ 88 (2002).
  * [12] Merkulov, I. A., Efros, A. L. & Rosen, M.  Electron spin relaxation by nuclei in semiconductor quantum dots.  _Physical Review B_ 65 (2002).
  * [13] Coish, W. A. & Loss, D.  Hyperfine interaction in a quantum dot: Non-Markovian electron spin dynamics.  _Physical Review B_ 70 (2004).
  * [14] Petta, J. R. _et al._ Coherent manipulation of coupled electron spins in semiconductor quantum dots.  _Science_ 309, 2180–2184 (2005).
  * [15] Johnson, A. C. _et al._ Triplet–singlet spin relaxation via nuclei in a double quantum dot.  _Nature_ 435, 925–928 (2005).
  * [16] Chesi, S., Yang, L.-P. & Loss, D.  Dephasing due to Nuclear Spins in Large-Amplitude Electric Dipole Spin Resonance.  _Physical Review Letters_ 116, 066806 (2016).
  * [17] Schliemann, J., Khaetskii, A. & Loss, D.  Electron spin dynamics in quantum dots and related nanostructures due to hyperfine interaction with nuclei.  _Journal of Physics: Condensed Matter_ 15, R1809–R1833 (2003).
  * [18] Hanson, R., Kouwenhoven, L. P., Petta, J. R., Tarucha, S. & Vandersypen, L. M. K.  Spins in few-electron quantum dots.  _Reviews of Modern Physics_ 79, 1217–1265 (2007).
  * [19] Coish, W. A. & Baugh, J.  Nuclear spins in nanostructures.  _physica status solidi (b)_ 246, 2203–2215 (2009).
  * [20] Chekhovich, E. A. _et al._ Nuclear spin effects in semiconductor quantum dots.  _Nature Materials_ 12, 494–504 (2013).
  * [21] Klauser, D., Coish, W. A. & Loss, D.  Nuclear spin state narrowing via gate-controlled Rabi oscillations in a double quantum dot.  _Physical Review B_ 73 (2006).
  * [22] Danon, J. & Nazarov, Y. V.  Nuclear Tuning and Detuning of the Electron Spin Resonance in a Quantum Dot: Theoretical Consideration.  _Physical Review Letters_ 100, 056603 (2008).
  * [23] Klauser, D., Coish, W. A. & Loss, D.  Nuclear spin dynamics and Zeno effect in quantum dots and defect centers.  _Physical Review B_ 78 (2008).
  * [24] Rudner, M. S., Koppens, F. H. L., Folk, J. A., Vandersypen, L. M. K. & Levitov, L. S.  Nuclear spin dynamics in double quantum dots: Fixed points, transients, and intermittency.  _Physical Review B_ 84 (2011).
  * [25] Gullans, M., Krich, J. J., Taylor, J. M., Halperin, B. I. & Lukin, M. D.  Preparation of nonequilibrium nuclear spin states in double quantum dots.  _Physical Review B_ 88 (2013).
  * [26] Foletti, S., Bluhm, H., Mahalu, D., Umansky, V. & Yacoby, A.  Universal quantum control of two-electron spin quantum bits using dynamic nuclear polarization.  _Nature Physics_ 5, 903–908 (2009).
  * [27] Vink, I. T. _et al._ Locking electron spins into magnetic resonance by electron–nuclear feedback.  _Nature Physics_ 5, 764–768 (2009).
  * [28] Bluhm, H., Foletti, S., Mahalu, D., Umansky, V. & Yacoby, A.  Enhancing the Coherence of a Spin Qubit by Operating it as a Feedback Loop That Controls its Nuclear Spin Bath.  _Physical Review Letters_ 105, 216803 (2010).
  * [29] Petta, J. R. _et al._ Dynamic Nuclear Polarization with Single Electron Spins.  _Physical Review Letters_ 100, 067601 (2008).
  * [30] Rudner, M. S. & Levitov, L. S.  Phase transitions in dissipative quantum transport and mesoscopic nuclear spin pumping.  _Physical Review B_ 82 (2010).
  * [31] Barthel, C. _et al._ Relaxation and readout visibility of a singlet-triplet qubit in an Overhauser field gradient.  _Physical Review B_ 85, 035306 (2012).
  * [32] Nichol, J. M. _et al._ Quenching of dynamic nuclear polarization by spin–orbit coupling in GaAs quantum dots.  _Nature Communications_ 6, 7682 (2015).
  * [33] Nadj-Perge, S. _et al._ Disentangling the effects of spin-orbit and hyperfine interactions on spin blockade.  _Physical Review B_ 81 (2010).
  * [34] Shulman, M. D. _et al._ Suppressing qubit dephasing using real-time Hamiltonian estimation.  _Nature Communications_ 5, 5156 (2014).
  * [35] Malinowski, F. K. _et al._ Notch filtering the nuclear environment of a spin qubit.  _Nature Nanotechnology_ 12, 16–20 (2016).
  * [36] Nakajima, T. _et al._ Coherence of a Driven Electron Spin Qubit Actively Decoupled from Quasistatic Noise.  _Physical Review X_ 10, 011060 (2020).
  * [37] Camenzind, L. C. _et al._ A spin qubit in a fin field-effect transistor.  _Nature Electronics_ 5, 178–183 (2022).  eprint 2103.07369.
  * [38] Yoneda, J. _et al._ Noise-correlation spectrum for a pair of spin qubits in silicon (2022).
  * [39] Bosco, S. & Loss, D.  Fully Tunable Hyperfine Interactions of Hole Spin Qubits in Si and Ge Quantum Dots.  _Physical Review Letters_ 127, 190501 (2021).
  * [40] Gangloff, D. A. _et al._ Quantum interface of an electron and a nuclear ensemble.  _Science_ 364, 62–66 (2019).
  * [41] Bodey, J. H. _et al._ Optical spin locking of a solid-state qubit.  _npj Quantum Information_ 5, 95 (2019).
  * [42] Gangloff, D. A. _et al._ Witnessing quantum correlations in a nuclear ensemble via an electron spin qubit.  _Nature Physics_ 17, 1247–1253 (2021).
  * [43] Jackson, D. M. _et al._ Quantum sensing of a coherent single spin excitation in a nuclear ensemble.  _Nature Physics_ 17, 585–590 (2021).
  * [44] Jackson, D. M. _et al._ Optimal Purification of a Spin Ensemble by Quantum-Algorithmic Feedback.  _Physical Review X_ 12, 031014 (2022).
  * [45] Rudner, M. S. & Levitov, L. S.  Electrically Driven Reverse Overhauser Pumping of Nuclear Spins in Quantum Dots.  _Physical Review Letters_ 99, 246602 (2007).
  * [46] Danon, J. _et al._ Multiple Nuclear Polarization States in a Double Quantum Dot.  _Physical Review Letters_ 103 (2009).
  * [47] Hartmann, S. R. & Hahn, E. L.  Nuclear Double Resonance in the Rotating Frame.  _Physical Review_ 128, 2042–2053 (1962).
  * [48] Pines, A., Gibby, M. G. & Waugh, J. S.  Proton-enhanced NMR of dilute spins in solids.  _The Journal of Chemical Physics_ 59, 569–590 (1973).
  * [49] Henstra, A., Dirksen, P., Schmidt, J. & Wenckebach, W.  Nuclear spin orientation via electron spin locking (NOVEL).  _Journal of Magnetic Resonance (1969)_ 77, 389–393 (1988).
  * [50] Tenberg, S., McNeil, R. P. G., Rubbert, S. & Bluhm, H.  Narrowing of the Overhauser field distribution by feedback-enhanced dynamic nuclear polarization.  _Physical Review B_ 92, 195428 (2015).
  * [51] Noiri, A. & et al.  (unpublished) (unpublished).
  * [52] Markmann, S., Reichl, C., Wegscheider, W. & Salis, G.  Universal nuclear focusing of confined electron spins.  _Nature Communications_ 10, 1097 (2019).
  * [53] Merkulov, I. A., Alvarez, G., Yakovlev, D. R. & Schulthess, T. C.  Long-term dynamics of the electron-nuclear spin system of a semiconductor quantum dot.  _Physical Review B_ 81, 115107 (2010).
  * [54] Henstra, A. & Wenckebach, W.  The theory of nuclear orientation via electron spin locking (NOVEL).  _Molecular Physics_ 106, 859–871 (2008).
  * [55] Weis, V., Bennati, M., Rosay, M. & Griffin, R. G.  Solid effect in the electron spin dressed state: A new approach for dynamic nuclear polarization.  _The Journal of Chemical Physics_ 113, 6795–6802 (2000).
  * [56] Weis, V. & Griffin, R.  Electron-nuclear cross polarization.  _Solid State Nuclear Magnetic Resonance_ 29, 66–78 (2006).
  * [57] Jain, S. K., Mathies, G. & Griffin, R. G.  Off-resonance NOVEL.  _The Journal of Chemical Physics_ 147, 164201 (2017).
  * [58] Shchepetilnikov, A. V. _et al._ Nuclear magnetic resonance and nuclear spin relaxation in AlAs quantum well probed by ESR.  _Physical Review B_ 94, 241302 (2016).
  * [59] Yang, W. & Sham, L. J.  Collective Nuclear Stabilization by Optically Excited Hole in Quantum Dot.  _Physical Review B_ 85, 235319 (2012).  eprint 1012.0060.
  * [60] Müller, L. & Ernst, R.  Coherence transfer in the rotating frame: Application to heteronuclear cross-correlation spectroscopy.  _Molecular Physics_ 38, 963–992 (1979).
  * [61] Greilich, A. _et al._ Nuclei-Induced Frequency Focusing of Electron Spin Coherence.  _Science_ 317, 1896–1899 (2007).
  * [62] Redfield, A. G.  Nuclear Magnetic Resonance Saturation and Rotary Saturation in Solids.  _Physical Review_ 98, 1787–1809 (1955).
  * [63] Laird, E. A. _et al._ Hyperfine-Mediated Gate-Driven Electron Spin Resonance.  _Physical Review Letters_ 99, 246601 (2007).
  * [64] Henstra, A. & Wenckebach, W.  Dynamic nuclear polarisation via the integrated solid effect I: Theory.  _Molecular Physics_ 112, 1761–1772 (2014).
  * [65] Demco, D. E., Tegenfeldt, J. & Waugh, J. S.  Dynamics of cross relaxation in nuclear magnetic double resonance.  _Physical Review B_ 11, 4133–4151 (1975).
  * [66] Vega, S., Shattuck, T. W. & Pines, A.  Double-quantum cross-polarization NMR in solids.  _Physical Review A_ 22, 638–661 (1980).
  * [67] McARTHUR, D. A., Hahn, E. L. & Walstedt, R. E.  Rotating-Frame Nuclear-Double-Resonance Dynamics: Dipolar Fluctuation Spectrum in Ca F 2.  _Physical Review_ 188, 609–638 (1969).
  * [68] Lurie, F. M. & Slichter, C. P.  Spin Temperature in Nuclear Double Resonance.  _Physical Review_ 133, A1108–A1122 (1964).
  * [69] Ota, T. _et al._ Decoherence of nuclear spins due to dipole-dipole interactions probed by resistively detected nuclear magnetic resonance.  _Applied Physics Letters_ 91, 193101 (2007).
  * [70] Chekhovich, E., Hopkinson, M., Skolnick, M. & Tartakovskii, A.  Suppression of nuclear spin bath fluctuations in self-assembled quantum dots induced by inhomogeneous strain.  _Nature Communications_ 6, 6348 (2015).
  * [71] Sundfors, R. K.  Exchange and Quadrupole Broadening of Nuclear Acoustic Resonance Line Shapes in the III-V Semiconductors.  _Physical Review_ 185, 458–472 (1969).
  * [72] Hester, R. K., Sher, A., Soest, J. F. & Weisz, G.  Nuclear-magnetic-resonance detection of charge defects in gallium arsenide.  _Physical Review B_ 10, 4262–4273 (1974).
  * [73] Yusa, G., Muraki, K., Takashina, K., Hashimoto, K. & Hirayama, Y.  Controlled multiple quantum coherences of nuclear spins in a nanometre-scale device.  _Nature_ 434, 1001–1005 (2005).
  * [74] Noorhidayati, A. _et al._ Resistively detected NMR in a triple-gate quantum point contact: Magnetic field dependence.  _Physical Review B_ 101, 035425 (2020).
  * [75] Kawamura, M., Ono, K., Stano, P., Kono, K. & Aono, T.  Electronic Magnetization of a Quantum Point Contact Measured by Nuclear Magnetic Resonance.  _Physical Review Letters_ 115, 036601 (2015).
  * [76] Urbaszek, B. _et al._ Nuclear spin physics in quantum dots: An optical investigation.  _Reviews of Modern Physics_ 85, 79–133 (2013).
  * [77] Delbecq, M. R. _et al._ Quantum Dephasing in a Gated GaAs Triple Quantum Dot due to Nonergodic Noise.  _Physical Review Letters_ 116, 046802 (2016).
  * [78] Laucht, A. _et al._ A dressed spin qubit in silicon.  _Nature Nanotechnology_ 12, 61–66 (2016).
  * [79] Takeda, K. _et al._ A fault-tolerant addressable spin qubit in a natural silicon quantum dot.  _Science Advances_ 2, e1600694–e1600694 (2016).
  * [80] Yoneda, J. _et al._ A quantum-dot spin qubit with coherence limited by charge noise and fidelity higher than 99.9%.  _Nature Nanotechnology_ 13, 102–106 (2017).
  * [81] Xu, X. _et al._ Optically controlled locking of the nuclear field via coherent dark-state spectroscopy.  _Nature_ 459, 1105–1109 (2009).
  * [82] Malinowski, F. K. _et al._ Spectrum of the Nuclear Environment for GaAs Spin Qubits.  _Physical Review Letters_ 118, 177702 (2017).
  * [83] Jang, W. _et al._ Wigner-molecularization-enabled dynamic nuclear field programming (2022).
  * [84] Högele, A. _et al._ Dynamic Nuclear Spin Polarization in the Resonant Laser Excitation of an InGaAs Quantum Dot.  _Physical Review Letters_ 108, 197403 (2012).
  * [85] Shafiei, M., Nowack, K. C., Reichl, C., Wegscheider, W. & Vandersypen, L. M. K.  Resolving Spin-Orbit- and Hyperfine-Mediated Electric Dipole Spin Resonance in a Quantum Dot.  _Physical Review Letters_ 110, 107601 (2013).
  * [86] Bulaev, D. V. & Loss, D.  Electric Dipole Spin Resonance for Heavy Holes in Quantum Dots.  _Physical Review Letters_ 98, 097202 (2007).
  * [87] Froning, F. N. M. _et al._ Ultrafast hole spin qubit with gate-tunable spin–orbit switch functionality.  _Nature Nanotechnology_ 16, 308–312 (2021).
  * [88] Bosco, S. _et al._ Phase driving hole spin qubits (2023).  eprint 2303.03350.
  * [89] Press, W. H., Teukolsky, S. A., Vetterling, W. T. & Flannery, B. P.  _Numerical Recipes: The Art of Scientific Computing_ (Cambridge University Press, Cambridge, UK ; New York, 2007), 3rd ed edn.
  * [90] Huang, C.-W. & Hu, X.  Theoretical study of nuclear spin polarization and depolarization in self-assembled quantum dots.  _Physical Review B_ 81, 205304 (2010).
  * [91] Petrov, M. Y. _et al._ Coupled electron-nuclear spin dynamics in quantum dots: A graded box model approach.  _Physical Review B_ 80, 125318 (2009).
  * [92] Glazov, M. M., Yugova, I. A. & Efros, A. L.  Electron spin synchronization induced by optical nuclear magnetic resonance feedback.  _Physical Review B_ 85, 041303 (2012).
  * [93] Neder, I. _et al._ Semiclassical model for the dephasing of a two-electron spin qubit coupled to a coherently evolving nuclear spin bath.  _Physical Review B_ 84 (2011).
  * [94] Shirley, J.  Solution of the Schrödinger Equation with a Hamiltonian Periodic in Time.  _Physical Review_ 138, B979–B987 (1965).
  * [95] Uhlenbeck, G. E. & Ornstein, L. S.  On the Theory of the Brownian Motion.  _Physical Review_ 36, 823–841 (1930).
  * [96] Rabenstein, K., Sverdlov, V. A. & Averin, D. V.  Qubit decoherence by Gaussian low-frequency noise.  _Journal of Experimental and Theoretical Physics Letters_ 79, 646–649 (2004).
  * [97] Gutierrez-Rubio, A., Stano, P. & Loss, D.  Optimal frequency estimation and its application to quantum dots (2020).
  * [98] Rojas-Arias, J. S. _et al._ Spatial noise correlations beyond nearest-neighbor in 28Si/SiGe spin qubits 11 (unpublished).
  * [99] Taylor, J. M.  _Hyperfine Interactions and Quantum Information Processing in Quantum Dots_.  Ph.D. thesis, Harvard University, Cambridge, Massachusetts (2006).
  * [100] Reilly, D. J. _et al._ Measurement of Temporal Correlations of the Overhauser Field in a Double Quantum Dot.  _Physical Review Letters_ 101, 236803 (2008).
  * [101] Deng, C. & Hu, X.  Nuclear spin diffusion in quantum dots: Effects of inhomogeneous hyperfine interaction.  _Physical Review B_ 72 (2005).
  * [102] Gong, Z.-X., Yin, Z.-q. & Duan, L.-M.  Dynamics of the Overhauser field under nuclear spin diffusion in a quantum dot.  _New Journal of Physics_ 13, 033036 (2011).
  * [103] Hayashi, H., Itoh, K. M. & Vlasenko, L. S.  Nuclear magnetic resonance linewidth and spin diffusion in S 29 i isotopically controlled silicon.  _Physical Review B_ 78 (2008).
  * [104] Fischer, J., Coish, W. A., Bulaev, D. V. & Loss, D.  Spin decoherence of a heavy hole coupled to nuclear spins in a quantum dot.  _Physical Review B_ 78, 155329 (2008).
  * [105] Eble, B. _et al._ Hole–Nuclear Spin Interaction in Quantum Dots.  _Physical Review Letters_ 102, 146601 (2009).
  * [106] Hendrickx, N. W. _et al._ A four-qubit germanium quantum processor.  _Nature_ 591, 580–585 (2021).
  * [107] Fang, Y., Philippopoulos, P., Culcer, D., Coish, W. A. & Chesi, S.  Recent advances in hole-spin qubits.  _Materials for Quantum Technology_ 3, 012003 (2023).


# Linearizing Large Language Models

Jean Mercat&Igor Vasiljevic11footnotemark: 1&Sedrick Keh11footnotemark:
1Kushal Arora&Achal Dave&Adrien Gaidon&Thomas Kollar
Toyota Research Institute
<EMAIL_ADDRESS>Equal contribution.

###### Abstract

Linear transformers have emerged as a subquadratic-time alternative to softmax
attention and have garnered significant interest due to their fixed-size
recurrent state that lowers inference cost. However, their original
formulation suffers from poor scaling and underperforms compute-matched
transformers. Recent linear models such as RWKV and Mamba have attempted to
address these shortcomings by proposing novel time-mixing and gating
architectures, but pre-training large language models requires significant
data and compute investments. Thus, the search for subquadratic architectures
is limited by the availability of compute and quality pre-training datasets.
As a cost-effective alternative to pre-training linear transformers, we
propose Scalable UPtraining for Recurrent Attention (SUPRA).111We borrow the
term “uptraining” from Ainslie et al. (2023) to refer to continued training
with a modified architecture, as opposed to fine-tuning, which usually refers
to continued training on a different dataset. We present a method to uptrain
existing large pre-trained transformers into Recurrent Neural Networks (RNNs)
with a modest compute budget. This allows us to leverage the strong pre-
training data and performance of existing transformer LLMs, while requiring 5%
of the training cost. We find that our linearization technique leads to
competitive performance on standard benchmarks, but we identify persistent in-
context learning and long-context modeling shortfalls for even the largest
linear models. Our code and models can be found at https://github.com/TRI-
ML/linear_open_lm.

Figure 1: We reuse pre-trained LLMs (gray) and convert them to RNNs with
minimal uptraining (SUPRA), outperforming linear attention models (RWKV) on
natural language tasks like HellaSwag, attaining the same memory advantages
(center), but also inheriting the limitations of RNNs on tasks like MMLU
(right).

##  1 Introduction

Over the last few years, Transformers (Vaswani et al., 2017) have displaced
Recurrent Neural Networks (RNNs) in sequence modeling tasks, owing to their
highly parallel training efficiency and unmatched scaling performance (Kaplan
et al., 2020). However, this training efficiency comes at the cost of
inference cost that scales linearly with the number of tokens, compared to the
fixed-cost inference of RNNs. The memory-intensive nature of transformers has
led to renewed interest in recurrence—the fixed-size hidden state remains an
attractive modeling proposition to reduce the cost of inference for language
and multimodal models.

Several recent works, starting with Linear Transformers (Katharopoulos et al.,
2020), have observed a relationship between a linearized form of attention and
recurrence, leading to a duality between transformers and RNNs: models can be
trained with sequence parallelism (i.e. as transformers, avoiding
backpropagation through time), but can operate as RNNs at inference time.
Although this architecture allows efficient training of RNNs, softmax
transformers continue to outperform linear transformers across natural
language understanding benchmarks. A number of novel RNN architectures have
attempted to bridge this performance gap. These include RWKV (Peng et al.,
2023a), Retentive Networks (Sun et al., 2023), TransNormer (Qin et al.,
2022a), and more recently, Griffin (De et al., 2024) and RecurrentGemma
(Griffin Team et al., 2024). These models are pre-trained on the same pre-
training datasets as transformers and show promising results.

State-space models (Gu et al., 2021) (SSMs) are another recurrent alternative
to softmax transformers, combining RNNs and convolutional networks to
efficiently model long sequences. The Mamba (Gu & Dao, 2023) architecture is a
SSM that shows impressive performance at smaller scales, matching or exceeding
the performance of softmax transformers on a number of natural language
understanding (NLU) benchmarks. However, a gap remains for long-context NLU
tasks, showing a persistent advantage of softmax attention.

Architecture search at the scale of large language models is expensive. Rather
than pre-training linear models, another approach is to convert an existing
transformer into an RNN; Kasai et al. (2021) proposed to uptrain encoder-
decoder transformers into RNNs by introducing an approximating MLP attention
module. Zhang et al. (2024) improved on this method by adding a loss to match
softmax attention to approximate more closely the base transformer.

While approximating attention is an intriguing approach to re-using pre-
trained transformers, it leads to instability and poor performance when
uptraining large-scale models. We instead take a different approach: rather
than approximate softmax attention, we replace it with a linear kernel and a
normalization strategy to uptrain the most performant LLMs into RNNs (see
Figure 2). We take advantage of models trained on high-quality, proprietary
datasets for trillions of tokens (e.g. Mistral (Jiang et al., 2023) and Llama2
(Touvron et al., 2023)). Fine-tuning these models on publicly available data
for a small fraction of pre-training tokens (see Figure 1), we obtain linear
models that are competitive with the best linear transformers for a fraction
of the compute. We call our approach Scalable UPtraining for Recurrent
Attention (SUPRA).

Figure 2: Our linearization strategy: we replace the softmax normalization
with GroupNorm (GN) and introduce a small MLP to project the queries and keys,
converting a pre-trained attention block (left) to a linear attention (right).
The model can be be trained in parallel as a transformer and used recurrently
at inference time with a mathematically equivalent reformulation.

Our contributions are as follows:

  * •

We propose Scalable UPtraining for Recurrent Attention (SUPRA), a
linearization strategy to uptrain state-of-the-art LLMs into performant RNNs.

  * •

We show that this simple uptraining technique is competitive with the
strongest pre-trained recurrent LLMs.

  * •

We investigate the limitations of recurrent LLMs, comparing pre-trained and
uptrained RNNs to transformers, revealing a persistent gap for in-context
learning and long-context tasks.

##  2 Methodology

In this section we review linear transformers, and describe the linearization
technique of Kasai et al. (2021) as it lays the groundwork for our approach.
Finally, we present SUPRA, our method for uptraining large transformers into
RNNs.

###  2.1 Background: Linear Attention

Linear Transformers (Katharopoulos et al., 2020) establish a connection
between transformers and RNNs, generalizing the definition of attention by
replacing the softmax dot-product attention $\mathbf{v}^{\prime}$ with a more
generic similarity function $\text{sim}(\mathbf{q},\mathbf{k})$ between the
queries $\mathbf{q}$ and keys $\mathbf{k}$:

$\displaystyle\mathbf{v}^{\prime}_{i}=\frac{\sum_{j=1}^{i}\text{sim}(\mathbf{q}_{i},\mathbf{k}_{j})\mathbf{v}_{j}}{\sum_{j=1}^{i}\text{sim}(\mathbf{q}_{i},\mathbf{k}_{j})}.$
(1)

Standard softmax attention is a special case, using
$\text{sim}(\mathbf{q},\mathbf{k})=\exp\left(\frac{\mathbf{q}^{T}\mathbf{k}}{\sqrt{d}}\right)$.

The authors explore several alternative functions for
$\text{sim}(\mathbf{q},\mathbf{k})$, including a linear kernel. Their main
architecture uses the similarity function
$\text{sim}(\mathbf{q},\mathbf{k})=\phi(\mathbf{q})\cdot\phi(\mathbf{k})$ with
a fixed exponential linear unit kernel $\phi(x)=\text{elu}(x)+1$. They show
the computational benefits of linear attention and, more importantly for this
work, they demonstrate how such a model can be expressed as an RNN in the case
of attention with causal masking.

#### Recurrent Inference

Linear attention can be expressed as an RNN that updates a state
$\mathbf{s}_{i}$ and a normalization factor $\mathbf{z}_{i}$ at each time
step. Katharopoulos et al. (2020) call these terms the attention memory and
normalized memory. This RNN formulation is mathematically equivalent to linear
attention, allowing the user to choose the most efficient one for a given task
and hardware. Consider a stream of tokens we want to generate
$X=[\mathbf{x}_{1},\mathbf{x}_{2},\mathbf{x}_{3},\dotsc]$. At inference time,
we use the following update rule, where subscripts denote timestep in the
recurrence (calling $\mathbf{k}_{i}=W_{K}\mathbf{x}_{i}$, etc):

$\displaystyle\mathbf{s}_{0}=0~{}~{}~{}\mathbf{z}_{0}=0$ (2)
$\displaystyle\mathbf{s}_{i}=\mathbf{s}_{i-1}+\phi(\mathbf{k}_{i})\mathbf{v}_{i}^{T}$
(3) $\displaystyle\mathbf{z}_{i}=\mathbf{z}_{i-1}+\phi(\mathbf{k}_{i})$ (4)
$\displaystyle\mathbf{v}’_{i}=\frac{\phi(\mathbf{q}_{i})^{T}\mathbf{s}_{i}}{\phi(\mathbf{q}_{i})^{T}\mathbf{z}_{i}}$
(5)

The state $\mathbf{s}_{i}$ acts as a constant-size KV cache. Instead of
appending new values to the cache, the state is updated. This allows for
inference cost that is constant in the number of generated tokens.

###  2.2 Finetuning a Transformer into an RNN

Kasai et al. (2021) introduced a linear transformer uptraining procedure that
converts a pre-trained softmax transformer into an RNN by approximating the
attention computation with multi-layer perceptrons (MLPs). The method (T2R)
starts with a softmax attention model, and linearizes the softmax operation.
Recall the kernel linear attention similarity function:

$\displaystyle\text{sim}(\mathbf{x},\mathbf{y})=\phi(\mathbf{x})\cdot\phi(\mathbf{y}).$
(6)

Instead of choosing $\phi$ as a simple non-linearity, the authors use a
trainable layer:

$\displaystyle\phi(\mathbf{x})=\text{relu}(W\mathbf{x}+\mathbf{b}).$ (7)

The weights are shared between keys and queries for a given attention head. By
using $\phi$ and rearranging the operations, attention can be written as:

$\displaystyle\mathbf{v}’_{i}=\frac{\phi(\mathbf{q}_{i})^{T}\sum_{j=1}^{i}\phi(\mathbf{k}_{j})\mathbf{v}_{j}^{T}}{\phi(\mathbf{q}_{i})^{T}\sum_{j=1}^{i}\phi(\mathbf{k}_{j})}.$
(8)

This allows the recurrent inference described in Section 2.1. However, this
formulation has a number of drawbacks. First, it requires a significant re-
training of the model, using approximately $20\%$ of pre-training tokens for
conversion, while suffering a $5-10\%$ drop in performance on language
benchmarks. Furthermore, this approach was tested on relatively small models
($\approx 100M$ scale). Because it mimics the attention formulation closely,
it suffers from stability issues at larger scales. To address these issues, we
modify the approach to adapt it to large-scale model uptraining.

###  2.3 SUPRA: Scalable UPtraining for Recurrent Attention

Rather than pre-training linear models from scratch, we choose to instead
uptrain state-of-the-art transformers. Leveraging models that take advantage
of high-quality (but proprietary) pre-training datasets, we linearize them
using a modest fraction of pre-training data (see Figure 1). We build on T2R,
identifying two major issues and proposing SUPRA, an approach to fine-tuning
very large transformers into RNNs.

We first follow the literature in linear transformers and identify the
normalization factor in linear attention as unstable (e.g. TransNormer (Qin et
al., 2022a)). In Section 3.3 we show that uptraining a 1B model following the
procedure in T2R causes a large drop in performance. We instead follow
Retentive Networks (Sun et al., 2023) and replace the normalization with a
GroupNorm operation.

Next we note that linear attention suffers more with absolute positional
encoding than softmax attention, and a modern relative positional encoding
scheme like RoPE (Su et al., 2021) is crucial for competitive performance.
Rather than training a linear transformer from scratch incorporating these
findings (RetNet, TransNormer) we use MLP kernels to convert large language
models into RNNs.

Starting with the pre-trained model, we add weights shared between keys and
queries $W$, $\phi(\mathbf{x})=\text{relu}(W\mathbf{x}+\mathbf{b})$ and use
the rotary positional embedding (RoPE (Su et al., 2021)) such that the
similarity function becomes

$\displaystyle\text{sim}(\mathbf{q}_{i},\mathbf{k}_{j})=\text{RoPE}(\phi(\mathbf{q}_{i}))\cdot\text{RoPE}(\phi(\mathbf{k}_{j})).$
(9)

We normalize the output with a GroupNorm (Wu & He, 2018) instead of dividing
by the sum of $\text{sim}(\mathbf{q}_{i},\mathbf{k}_{j})$ (as in T2R). We use
a fixed decay vector $\gamma\in(0,1)^{h}$, with $h$ heads, as in Sun et al.
(2023). This leads to the following attention formulation (see Figure 2 for a
graphical representation):

$\displaystyle\mathbf{v}’_{i}=\text{GroupNorm}\left(\sum_{j=1}^{i}\gamma^{i-j}\text{sim}(\mathbf{q}_{i},\mathbf{k}_{j})\mathbf{v}_{j}\right).$
(10)

These new parameters are trained jointly with the rest of the network; at test
time, we use the recurrent formulation for inference.

##  3 Experiments

We uptrain a variety of models from the 1B to 7B range into RNNs (Llama2
(Touvron et al., 2023) and Mistral (Jiang et al., 2023)), and evaluate our
models in two settings: standard language understanding benchmarks and long-
context evaluations. We compare the results of different architectural choices
and training strategies, and then show the limitations of linear models on
various benchmarks, describing the persistent gap between vanilla attention
and recurrence. We choose Llama2-7B and Mistral-7B as our base models for
uptraining, but our recipe is general to any transformer model.

We compare our procedure to a variety of pre-trained recurrent models. Given
that the largest available state-space models are at the 2.8B scale, we also
train a Mamba model on the RefinedWeb (Penedo et al., 2023) dataset from
scratch for 1.2T tokens, to serve as a strong baseline for a pre-trained
recurrent model 222The Mistral-SUPRA and Mamba-7B models are released along
with the code..

We use a fork of OpenLM (Gururangan et al., 2023) for all training and fine-
tuning. Please see Section 7 for hyperparameters and further details on
reproducibility.

Model | Size | Tokens | HellaSwag | PIQA | WG | ARC-E | ARC-C | MMLU | Average
---|---|---|---|---|---|---|---|---|---
StableLM2 | 1.6B | 2000 | 69.0 | 76.7 | 63.6 | 68.6 | 38.9 | 38.4 | 59.2
StableLM | 3B | 1000 | 73.8 | 79.3 | 65.8 | 72.1 | 40.0 | 44.2 | 62.5
Gemma | 2B | 2000 | 71.4 | 78.6 | 64.4 | 74.0 | 41.5 | 41.2 | 61.9
Mamba | 1.4B | 600 | 59.0 | 73.9 | 61.4 | 65.5 | 32.9 | 25.2 | 53.0
RWKV-5 | 1.5B | 1100 | 53.1 | 71.6 | 59.0 | 62.2 | 32.7 | 26.2 | 50.8
Mamba | 2.8B | 600 | 66.2 | 75.8 | 63.4 | 69.7 | 36.3 | 26.3 | 56.3
Llama2 | 7B | 2000 | 76.0 | 79.1 | 69.1 | 76.3 | 46.3 | 45.9 | 65.4
Gemma | 7B | 6000 | 80.7 | 81.9 | 73.7 | 81.1 | 53.2 | 62.9 | 72.2
Mistral | 7B | 8000(?) | 81.0 | 82.1 | 74.0 | 80.9 | 53.8 | 62.4 | 72.4
RetNet | 6.7B | 200 | 60.7 | 75.4 | 58.1 | – | – | – | –
RWKV-5 | 7B | 1100 | 70.9 | 77.2 | 67.4 | 71.8 | 43.6 | 31.0 | 60.3
RWKV-5-1.7T | 7B | 1700 | 73.0 | 78.6 | 72.9 | 75.8 | 45.6 | 34.9 | 63.5
Mamba (ours) | 7B | 1200 | 77.9 | 81.0 | 71.8 | 77.5 | 46.7 | 33.3 | 64.7
Llama2-SUPRA | 7B | +20 | 71.8 | 78.6 | 65.8 | 71.1 | 39.5 | 24.9 | 58.6
Mistral-SUPRA | 7B | +20 | 74.8 | 80.1 | 67.4 | 74.6 | 42.3 | 28.0 | 61.2
Mistral-SUPRA | 7B | +100 | 77.1 | 80.4 | 70.3 | 75.9 | 45.8 | 34.2 | 64.0
Table 1: Linear models (RNNs and SSMs) highlighted in gray. 5-shot results are
used for MMLU. Norm results are used for PIQA, HellaSwag, ARC-C. RetNet
results taken from RetNet paper.

#### Language Modeling.

In Table 1 we report results on standard NLU evaluations using the Eleuther
evaluation harness (Gao et al., 2023). We primarily compare to transformers
and linear models at the 7B scale, and we train a Mamba model at 7B for
comparison with RWKV-5. As our model is initialized from strong pre-trained
transformers (Llama2 and Mistral-7B), it preserves performance on most
benchmarks (except MMLU; see Section 4 for a discussion below). Our technique
outperforms RWKV-5 with minimal uptraining and is competitive with our 7B
Mamba trained from scratch on 1.2T tokens.

Model | Size | Train | Qasper (2-shot) | NarrativeQA (0-shot)
---|---|---|---|---
|  | Context | 2048 | 4096 | 8192 | 16384 | 2048 | 4096 | 8192 | 16384
Llama1-7B | 7B | 2048 | 24.43 | 7.23 | 5.08 | 4.88 | 21.44 | 1.03 | 0.0 | 0.0
Llama2-7B | 7B | 4096 | 23.26 | 27.26 | 6.20 | 5.49 | 21.32 | 22.61 | 0.0 | 0.0
Llama2-7B* | 7B | 4096 | 23.26 | 27.26 | 31.46 | 25.52 | 21.32 | 22.61 | 23.0 | 14.27
Mistral-7B | 7B | 8196 | 21.53 | 25.50 | 33.61 | 6.88 | 24.94 | 26.90 | 25.93 | 0.63
RecurrentGemma-2B | 2.7B | 8192 | 22.44 | 13.16 | 13.42 | 12.66 | 19.80 | 11.59 | 12.93 | 12.95
RWKV-5-1.7T | 7B | 2048 | 22.28 | 23.87 | 22.30 | 20.35 | 17.77 | 18.65 | 17.81 | 16.00
Mamba (ours) | 7B | 2048 | 19.68 | 5.58 | 5.90 | 6.32 | 19.70 | 0.28 | 0.0 | 0.0
Mistral-SUPRA | 7B | 2048 | 19.44 | 17.13 | 17.11 | 17.22 | 18.99 | 17.76 | 17.75 | 17.74
Table 2: Long context evaluations. Performance at various context size cutoffs
for Qasper (2-shot) and NarrativeQA (0-shot). * denotes linear RoPE scaling
with YaRN (Peng et al., 2023b).

#### Long Context.

Recurrent models were thought to perform well on long-context tasks because of
their ability to preserve performance beyond their training sequence size.
However, their downstream performance on long-context tasks has not been well-
documented. Prior studies either do not conduct long-context evaluations
(Katharopoulos et al., 2020; Kasai et al., 2021), evaluate only on perplexity
(Sun et al., 2023; De et al., 2024; Gu & Dao, 2023), or evaluate on datasets
which require task-specific training (Peng et al., 2023a). Instead, we
consider downstream natural language tasks from the SCROLLS benchmark (Shaham
et al., 2022a). Specifically, in Table 2 we present two tasks – Qasper (Dasigi
et al., 2021) and NarrativeQA (Kočiský et al., 2018) – from the set of tasks
evaluated in the Llama2-Long report (Xiong et al., 2023). We evaluate both
tasks with an input context cut-off at different lengths. A strong long-
context model should perform better given more context. However, the training
context lengths for these models do not go beyond 8k tokens. Transformer
models show the strongest results up to the context length they were trained
for but degrade beyond that. Interestingly, applying the YaRN trick (Peng et
al., 2023b) enables transformers to scale beyond their training context quite
well. RWKV shows a strong ability to handle much longer context than its
training. Our Mamba model on the contrary is not able to generalize beyond its
training context length. Surprisingly, the RecurrentGemma model (Griffin Team
et al., 2024) shows degrading performances even within its training context
length. Finally, our Mistral-SUPRA model preserves some performance at larger
context lengths but we believe it to result from the decay along the context
length that shortens the effective context. This is discussed in more details
below. We find a significant gap in performance between transformers and
available linear models, including models uptrained from strong long-context
transformers. We speculate that more sophisticated recurrent state update
rules may be required to perform well at this task. Ideas such as gating
strategies (De et al., 2024), higher order linear attention (Mercat, 2020), or
associative binding (Munkhdalai et al., 2024) could be explored.

#### Decay factors.

The default decay factors proposed in Qin et al. (2024) gives better results
than no decay on short context benchmarks but at a long range, the decay
cancels out the influence of the context
($\text{max}(\gamma)^{2048}=3.35e^{-4}$). This can be related to a smooth
version of window attention (Beltagy et al., 2020). However, as more context
is given to the model, the long-range evaluation performances plateau. When
using the values proposed in Sun et al. (2023), that allow longer range
attention, we observe a performance drop on short-context benchmarks and no
substantial improvement on long-context evaluation.

Model | Size | Tokens | HellaSwag | ARC-E | ARC-C
---|---|---|---|---|---
Mamba | 1B | 100B | 58.3 | 62.7 | 29.1
SUPRA (from scratch) | 1B | 100B | 54.7 | 59.6 | 27.8
T2R (from scratch) | 1B | 100B | 55.2 | 61.0 | 28.4
Transformer | 1B | 100B | 55.9 | 60.4 | 29.3
Transformer* | 1B | 1.6T | 62.1 | 66.2 | 34.3
Fine-tune only new weights | 1B | (1.6T)+10B | 33.2 | 37.8 | 22.6
2-step fine-tune | 1B | (1.6T)+10B+10B | 56.1 | 62.2 | 30.4
T2R | 1B | (1.6T)+10B | 40.6 | 53.8 | 27.9
SUPRA | 1B | (1.6T)+10B | 57.0 | 62.4 | 31.6
SUPRA | 1B | (100B)+10B | 51.7 | 57.5 | 27.4
Table 3: Ablating different choices for linear uptraining: note the importance
of normalization. For the second half of the table, we uptrain a transformer
trained on 1.6T tokens on 10B further tokens. *This model was trained on a
different mix of data.

#### Ablations.

Table 3 compares transformers pre-trained on 100B tokens to Mamba (Gu & Dao,
2023), T2R (Kasai et al., 2021), and our approach. At this scale, with 100B
tokens of training, the Mamba model performs best and other models show
similar performance. The second half of Table 3 shows results for uptraining
from a pre-trained transformer. TheT2R (Kasai et al., 2021) uptraining was
unstable, yielding poor results compapred to SUPRA. This confirms that
normalization is key for maintaining performance of the base LLM when
uptraining.

To test the hypothesis that linear attention approximates the softmax
attention, we experimented with a 2-step approach. The first step trains only
the new parameters such that the model could learn to approximate the softmax.
The second step fine-tunes all the weights. The results show no benefit from
the two steps approach and indicates that the softmax is not approximated. See
Appendix A for a different approach to compare softmax attention and linear
attention.

Finally, we compare the results of SUPRA uptrainings from two pre-trained
softmax models. It appears that pre-training a linear model for a 100B token
yields better results than fine-tuning a softmax model that was trained with
the same budget. These results also shows, along with the comparison of
LLama2-SUPRA and Mistral-SUPRA in Table 1, that SUPRA benefits significantly
from a stronger pre-trained transformer. Thus, given a limited training
budget, using SUPRA from a strong pre-trained model is the best option.

##  4 Discussion

Comparison to pre-training SSMs/RNNs. With only 20B tokens of training, which
represents $2-10\%$ of RWKV and RetNet training cost, we obtain a model that
outperforms both on HellaSwag and that is competitive on other benchmarks (see
Table 1). Given the existing performance gap between the strongest transformer
models and the most performant linear models, SUPRA is a simple recipe for
conversion, allowing the study of strong RNNs with limited uptraining.

Comparison to Transformers on Short-Context Tasks. Our approach does not
explicitly approximate attention from the base transformer model (see Appendix
A), we do see a modest drop in performance across all benchmarks compared to
softmax transformers. This could be partially explained by the lower quality
of our data compared to the pre-training mix used to train models like
Mistral-7B. It is also likely that linear transformers are inherently less
expressive. However, the performance drop is relatively modest on most
benchmarks, and significantly smaller than the drop from T2R uptraining, which
shows the relevance of our approach.

Long Context Comparisons. Prior work on linear attention showcased similar or
better validation set perplexity to transformer models over long context (e.g.
Sun et al. (2023)) but did not evaluate linear models on natural language
long-context evaluations like SCROLLS (Shaham et al., 2022b). The results in
Table 2 show that recurrent models generally maintain performance beyond their
training context (except for Mamba-7b) while transformers (without
modification) do not. However, Table 2 also demonstrates that simple linear
scaling of the rotary positional embedding (Peng et al., 2023b; emozilla,
2023) can allow for context scaling beyond the window used for training a
given transformer model, effectively nullifying the performance edge of these
linear models. Furthermore, transformers generally outperform linear models at
their maximum training context length. Further research is needed into
extending linear models to long-context inference to take full advantage of
the lower inference cost relative to vanilla transformers.

Limitations. Since our method relies on initializing with strong pre-trained
transformers, our models inherit any of the biases and weaknesses of their
base models. Additionally, models that are already instruct-tuned do not
linearize as well as base models. Our models suffer from poor performance on
MMLU which requires in-context learning (5-shot), a weakness of linear models
(Akyürek et al., 2024). We leave the investigation of these weaknesses of
linear models to future work and hope that our proposed uptraining approach
can help facilitate and accelerate the research in this area.

##  5 Related Work

#### Linear Transformers.

The linear transformers introduced in Katharopoulos et al. (2020) lagged
behind vanilla transformers in downstream performance, and subsequent
architectures such as TransNormer (Qin et al., 2022a) and RetNet (Sun et al.,
2023) narrow the gap, but do not demonstrate competitive results with modern
transformers at scale. RWKV (Peng et al., 2023a), a linear transformer that
takes inspiration from LSTM (Hochreiter & Schmidhuber, 1997), is competitive
with compute-matched transformer-based models, but lags behind on a number of
NLU benchmarks. Griffin (De et al., 2024) is a concurrent model that takes a
hybrid approach, combining a sliding window with linear attention shows
impressive performance relative to vanilla transformers, but is trained on a
high-quality proprietary dataset.

Another thread in the literature focuses on efficient attention alternatives
(Performers (Choromanski et al., 2020), Cosformer (Qin et al., 2022b), LUNA
(Ma et al., 2021), RFA (Peng et al., 2021), Attention-free Transformer (Zhai
et al., 2021)). All of these approaches sacrifice performances for efficiency.
Efficiency improvements for vanilla transformers have narrowed the
capabilities gap between vanilla and linear transformers. The KV cache Pope et
al. (2023)greatly narrows the inference efficiency gap between linear and
vanilla transformers. RingAttention Liu et al. (2023) allows for very long
context scaling of vanilla attention without approximation.

#### State Space Models.

State-space models (SSMs) such as H3 (Dao et al., 2022), Hyena (Poli et al.,
2023), and Mamba (Gu & Dao, 2023) are recent alternatives to vanilla
transformers, combining the strengths of convolutional and recurrent models
with efficient hardware implementations. Instead of parallelizing training
over the sequence, they produce an efficient way to train the sequential RNN.
While these models are competitive with vanilla transformers on some tasks, we
show that SSMs share the limitations of linear transformers on several in-
context learning and long-context tasks.

#### Uptraining Linear Transformers.

Hedgehog (Zhang et al., 2024) builds on the work of Kasai et al. (2021),
identifying three different ways of training linear transformers – from
scratch, uptraining quadratic transformers for a specific task, and uptraining
generally. The authors focus on the first two, and we focus on the third.
Moreover, they aim at approximating the softmax attention matrices with linear
alternatives. In this work, we do not aim to approximate softmax attention, we
replace it with a linear alternative (see ablation above and appendix A).
Their method is only tested for smaller scale models and with parameter-
efficient fine-tuning for larger models, but presents challenges for scaling
for two reasons: (1) their training strategy involves comparing full attention
matrices which is computationally expensive, and not feasible for full fine-
tuning of large models with long sequences and (2) their method also inherits
the gradient instabilities of linear transformers studied in Sun et al.
(2023), while our normalization setup leads to stable uptraining of large
models.

##  6 Conclusion

We introduced SUPRA, a technique for converting large-scale pre-trained
softmax transformers into recurrent neural networks, enabling the study of the
strengths and limitations of recurrent models at scale with minimal compute
cost. Compared to pre-training linear models from scratch, the SUPRA strategy
produces competitive models comparable to the best available recurrent LLMs
(RWKV and Mamba) at the 7B scale.

We identify the strengths of linear models on standard NLU benchmarks but also
the enduring limitations on in-context (i.e. MMLU) and long-context
(NarrativeQA, Qasper) tasks, showing that linearized models do not inherit
these capabilities from the base softmax transformers.

One possible path to rectifying these limitations is explicitly training for
in-context learning (Akyürek et al., 2024). We leave explorations of
specialized and instruct data in the context of linear transformers to future
work. More sophisticated gating mechanisms as in in Peng et al. (2023a) and De
et al. (2024) are promising alternatives to our simple linearization. Using
our uptraining method would greatly reduce the necessary time and cost of such
experimentation.

##  7 Reproducibility

#### Codebase

We train our linear models using our fork of OpenLM (Gururangan et al., 2023)
that we modify to include a linear attention function (printed below). We use
Lightning Attention 2 (Qin et al., 2024) that offers a fast Triton (Tillet et
al., 2019) kernel for linear attention computation. Evaluations are done with
the Eleuther evaluation harness (Gao et al., 2023).

#### Data

We train and uptrain models on RefinedWeb (Penedo et al., 2023)(with 2 epochs
for our Mamba training), which we tokenize with the pre-trained model’s
tokenizers. When training from scratch, we used the GPT-NeoX-20B (Black et
al., 2022) tokenizer. We tokenize with sequence packing and use a default
sequence length of 2048.

#### Hyperparameters

We use square matrices with biases for the linear layers in the kernel $\phi$
functions to keep the same feature dimension in the queries and keys. We use
the same kernel, with the same weights for both queries and keys and apply a
ReLU activation. We use 1000 steps of linear learning rate warmup and a cosine
learning rate decay from $3e^{-5}$ to $1e^{-5}$ for our 7B uptrainings and
from $3e^{-4}$ to $1e^{-5}$ for our 1B uptrainings and for trainings from
scratch. We use the Adam optimizer with $\beta_{1}=0.9$ and $\beta_{2}=0.95$.
We trained our models with mini-batches totaling 2M tokens each. Our default
RoPE frequency uses the Llama value of $10^{4}$. For longer sequence lengths,
we use a RoPE frequency of $10^{6}$.

#### Training

Depending on the model size and the availability, we use from 4 to 32 nodes of
8 GPUs Nvidia H100 with Pytorch FSDP. We use a mixed precision strategy from
OpenLM that automatically selects between bfloat 16 and float 32 for different
operations. A linear 7B parameter model uptraining throughput is around $4300$
tokens per second per GPU.

#### Models

Our results can be reproduced by following the same training recipe or using
the model weights that we release: Mistral-SUPRA and Mamba-7b.

⬇

1def linear_attn_func(q, k, v, qk_scale: float, use_decay: bool = True,
normalize: bool = False) -> torch.Tensor:

2 """

3 Args:

4 q: queries, shape (batch_size, num_heads, seq_len, dim_qk)

5 k: keys, shape (batch_size, num_heads, seq_len, dim_qk)

6 v: values, shape (batch_size, num_heads, seq_len, dim_v)

7 qk_scale: scale factor for queries and keys

8 """

9 h = q.shape[1]

10 if use_decay:

11 s = slope_tensor(h, q.device, q.dtype)

12 else:

13 s = no_slope_tensor(h, q.device, q.dtype)

14

15 output = lightning_attn_ops(q, k * qk_scale, v, s)

16 if normalize:

17 norm = torch.clamp_min(

18 torch.einsum(

19 "nhld, nhld->nhl", q, k.cumsum(2) * qk_scale

20 ), 1e-6

21 )

22 return output / norm.unsqueeze(-1)

23 else:

24 return output

## References

  * Ainslie et al. (2023) Joshua Ainslie, James Lee-Thorp, Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai.  Gqa: Training generalized multi-query transformer models from multi-head checkpoints.  _arXiv preprint arXiv:2305.13245_ , 2023.
  * Akyürek et al. (2024) Ekin Akyürek, Bailin Wang, Yoon Kim, and Jacob Andreas.  In-context language learning: Arhitectures and algorithms.  _arXiv preprint arXiv:2401.12973_ , 2024.
  * Ba et al. (2016) Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton.  Layer normalization.  _arXiv preprint arXiv:1607.06450_ , 2016.
  * Beltagy et al. (2020) Iz Beltagy, Matthew E Peters, and Arman Cohan.  Longformer: The long-document transformer.  _arXiv preprint arXiv:2004.05150_ , 2020.
  * Black et al. (2022) Sid Black, Stella Biderman, Eric Hallahan, Quentin Anthony, Leo Gao, Laurence Golding, Horace He, Connor Leahy, Kyle McDonell, Jason Phang, Michael Pieler, USVSN Sai Prashanth, Shivanshu Purohit, Laria Reynolds, Jonathan Tow, Ben Wang, and Samuel Weinbach.  Gpt-neox-20b: An open-source autoregressive language model, 2022.  URL https://arxiv.org/abs/2204.06745.
  * Choromanski et al. (2020) Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al.  Rethinking attention with performers.  _arXiv preprint arXiv:2009.14794_ , 2020.
  * Dao et al. (2022) Tri Dao, Daniel Y Fu, Khaled K Saab, Armin W Thomas, Atri Rudra, and Christopher Ré.  Hungry hungry hippos: Towards language modeling with state space models.  _arXiv preprint arXiv:2212.14052_ , 2022.
  * Dasigi et al. (2021) Pradeep Dasigi, Kyle Lo, Iz Beltagy, Arman Cohan, Noah A. Smith, and Matt Gardner.  A dataset of information-seeking questions and answers anchored in research papers.  In _Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies_ , pp. 4599–4610, Online, June 2021. Association for Computational Linguistics.  doi: 10.18653/v1/2021.naacl-main.365.  URL https://aclanthology.org/2021.naacl-main.365.
  * 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.
  * emozilla (2023) emozilla.  Dynamically scaled rope further increases strength of retaining walls, 2023.  URL https://www.reddit.com/r/LocalLLaMA/comments/14mrgpr/dynamically_scaled_rope_further_increases/.  Reddit post.
  * Gao et al. (2023) Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Alain Le Noac’h, Haonan Li, Kyle McDonell, Niklas Muennighoff, Chris Ociepa, Jason Phang, Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang Sutawika, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou.  A framework for few-shot language model evaluation, 12 2023.  URL https://zenodo.org/records/10256836.
  * Griffin Team et al. (2024) Alexsandar Botev Griffin Team, Soham De, Samuel L Smith, Anushan Fernando, George-Christian Muraru, Ruba Haroun, and Leonard Berrada et al.  Recurrentgemma.  _arXiv preprint arXiv:2404.07839_ , 2024.
  * Gu & 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. (2021) Albert Gu, Karan Goel, and Christopher Ré.  Efficiently modeling long sequences with structured state spaces.  _arXiv preprint arXiv:2111.00396_ , 2021.
  * Gururangan et al. (2023) Suchin Gururangan, Mitchell Wortsman, Samir Yitzhak Gadre, Achal Dave, Maciej Kilian, Weijia Shi, Jean Mercat, Georgios Smyrnis, Gabriel Ilharco, Matt Jordan, Reinhard Heckel, Alex Dimakis, Ali Farhadi, Vaishaal Shankar, and Ludwig Schmidt.  OpenLM: a minimal but performative language modeling (lm) repository, 2023.  URL https://github.com/mlfoundations/open_lm/.  GitHub repository.
  * Hochreiter & Schmidhuber (1997) Sepp Hochreiter and Jürgen Schmidhuber.  Long short-term memory.  _Neural computation_ , 9(8):1735–1780, 1997.
  * 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.
  * Kaplan et al. (2020) Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei.  Scaling laws for neural language models.  _arXiv preprint arXiv:2001.08361_ , 2020.
  * Kasai et al. (2021) Jungo Kasai, Hao Peng, Yizhe Zhang, Dani Yogatama, Gabriel Ilharco, Nikolaos Pappas, Yi Mao, Weizhu Chen, and Noah A Smith.  Finetuning pretrained transformers into rnns.  _arXiv preprint arXiv:2103.13076_ , 2021.
  * Katharopoulos et al. (2020) Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret.  Transformers are rnns: Fast autoregressive transformers with linear attention.  In _International conference on machine learning_ , pp. 5156–5165. PMLR, 2020.
  * Kočiský et al. (2018) Tomáš Kočiský, Jonathan Schwarz, Phil Blunsom, Chris Dyer, Karl Moritz Hermann, Gábor Melis, and Edward Grefenstette.  The NarrativeQA reading comprehension challenge.  _Transactions of the Association for Computational Linguistics_ , 6:317–328, 2018.  doi: 10.1162/tacl˙a˙00023.  URL https://aclanthology.org/Q18-1023.
  * Liu et al. (2023) Hao Liu, Matei Zaharia, and Pieter Abbeel.  Ring attention with blockwise transformers for near-infinite context.  _arXiv preprint arXiv:2310.01889_ , 2023.
  * Ma et al. (2021) Xuezhe Ma, Xiang Kong, Sinong Wang, Chunting Zhou, Jonathan May, Hao Ma, and Luke Zettlemoyer.  Luna: Linear unified nested attention.  _Advances in Neural Information Processing Systems_ , 34:2441–2453, 2021.
  * Mercat (2020) Jean Mercat.  Higher order linear transformer.  _arXiv preprint arXiv:2010.14816_ , 2020.
  * Munkhdalai et al. (2024) Tsendsuren Munkhdalai, Manaal Faruqui, and Siddharth Gopal.  Leave no context behind: Efficient infinite context transformers with infini-attention.  _arXiv preprint arXiv:2404.07143_ , 2024.
  * Penedo et al. (2023) Guilherme Penedo, Quentin Malartic, Daniel Hesslow, Ruxandra Cojocaru, Alessandro Cappelli, Hamza Alobeidli, Baptiste Pannier, Ebtesam Almazrouei, and Julien Launay.  The refinedweb dataset for falcon llm: outperforming curated corpora with web data, and web data only.  _arXiv preprint arXiv:2306.01116_ , 2023.
  * Peng et al. (2023a) Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Huanqi Cao, Xin Cheng, Michael Chung, Matteo Grella, Kranthi Kiran GV, et al.  Rwkv: Reinventing rnns for the transformer era.  _arXiv preprint arXiv:2305.13048_ , 2023a.
  * Peng et al. (2023b) Bowen Peng, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole.  Yarn: Efficient context window extension of large language models.  _arXiv preprint arXiv:2309.00071_ , 2023b.
  * Peng et al. (2021) Hao Peng, Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah A Smith, and Lingpeng Kong.  Random feature attention.  _arXiv preprint arXiv:2103.02143_ , 2021.
  * Poli et al. (2023) Michael Poli, Stefano Massaroli, Eric Nguyen, Daniel Y Fu, Tri Dao, Stephen Baccus, Yoshua Bengio, Stefano Ermon, and Christopher Ré.  Hyena hierarchy: Towards larger convolutional language models.  _arXiv preprint arXiv:2302.10866_ , 2023.
  * Pope et al. (2023) Reiner Pope, Sholto Douglas, Aakanksha Chowdhery, Jacob Devlin, James Bradbury, Jonathan Heek, Kefan Xiao, Shivani Agrawal, and Jeff Dean.  Efficiently scaling transformer inference.  _Proceedings of Machine Learning and Systems_ , 5, 2023.
  * Qin et al. (2022a) Zhen Qin, Xiaodong Han, Weixuan Sun, Dongxu Li, Lingpeng Kong, Nick Barnes, and Yiran Zhong.  The devil in linear transformer.  _arXiv preprint arXiv:2210.10340_ , 2022a.
  * Qin et al. (2022b) Zhen Qin, Weixuan Sun, Hui Deng, Dongxu Li, Yunshen Wei, Baohong Lv, Junjie Yan, Lingpeng Kong, and Yiran Zhong.  cosformer: Rethinking softmax in attention.  _arXiv preprint arXiv:2202.08791_ , 2022b.
  * Qin et al. (2024) Zhen Qin, Weigao Sun, Dong Li, Xuyang Shen, Weixuan Sun, and Yiran Zhong.  Lightning attention-2: A free lunch for handling unlimited sequence lengths in large language models.  _arXiv preprint arXiv:2401.04658_ , 2024.
  * Shaham et al. (2022a) Uri Shaham, Elad Segal, Maor Ivgi, Avia Efrat, Ori Yoran, Adi Haviv, Ankit Gupta, Wenhan Xiong, Mor Geva, Jonathan Berant, and Omer Levy.  SCROLLS: Standardized CompaRison over long language sequences.  In _Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing_ , pp. 12007–12021, Abu Dhabi, United Arab Emirates, December 2022a. Association for Computational Linguistics.  URL https://aclanthology.org/2022.emnlp-main.823.
  * Shaham et al. (2022b) Uri Shaham, Elad Segal, Maor Ivgi, Avia Efrat, Ori Yoran, Adi Haviv, Ankit Gupta, Wenhan Xiong, Mor Geva, Jonathan Berant, et al.  Scrolls: Standardized comparison over long language sequences.  _arXiv preprint arXiv:2201.03533_ , 2022b.
  * Su et al. (2021) Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu.  Roformer: Enhanced transformer with rotary position embedding.  _arXiv preprint arXiv:2104.09864_ , 2021.
  * Sun et al. (2023) Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma, Yuqing Xia, Jilong Xue, Jianyong Wang, and Furu Wei.  Retentive network: A successor to transformer for large language models.  _arXiv preprint arXiv:2307.08621_ , 2023.
  * Tillet et al. (2019) Philippe Tillet, H. T. Kung, and David Cox.  Triton: an intermediate language and compiler for tiled neural network computations.  In _Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages_ , MAPL 2019, pp. 10–19, New York, NY, USA, 2019. Association for Computing Machinery.  ISBN 9781450367196.  doi: 10.1145/3315508.3329973.  URL https://doi.org/10.1145/3315508.3329973.
  * 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.
  * Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  _Advances in neural information processing systems_ , 30, 2017.
  * Wu & He (2018) Yuxin Wu and Kaiming He.  Group normalization.  In _Proceedings of the European conference on computer vision (ECCV)_ , pp. 3–19, 2018.
  * Xiong et al. (2023) Wenhan Xiong, Jingyu Liu, Igor Molybog, Hejia Zhang, Prajjwal Bhargava, Rui Hou, Louis Martin, Rashi Rungta, Karthik Abinav Sankararaman, Barlas Oguz, Madian Khabsa, Han Fang, Yashar Mehdad, Sharan Narang, Kshitiz Malik, Angela Fan, Shruti Bhosale, Sergey Edunov, Mike Lewis, Sinong Wang, and Hao Ma.  Effective long-context scaling of foundation models, 2023.
  * Zhai et al. (2021) Shuangfei Zhai, Walter Talbott, Nitish Srivastava, Chen Huang, Hanlin Goh, Ruixiang Zhang, and Josh Susskind.  An attention free transformer.  _arXiv preprint arXiv:2105.14103_ , 2021.
  * Zhang et al. (2024) Michael Zhang, Kush Bhatia, Hermann Kumbong, and Christopher Ré.  The hedgehog & the porcupine: Expressive linear attentions with softmax mimicry.  _arXiv preprint arXiv:2402.04347_ , 2024.

##  Appendix A Attention Approximation

In this section we investigate whether our up-training procedure leads to
linear attention that approximates the softmax attention from the base model,
as might be expected.

There are many possible ways to compare attention matrices. Moreover, some
architecture changes such as attention decay and lack of normalization in the
linear attention make a meaningful comparison difficult. We represent non-
normalized comparisons in Figure 3. It represents the cosine similarities and
singular value distances between the attention matrices at every layer and for
every head of the Mistral model compared with our Mistral-SUPRA. Each pixel of
these images is a scalar similarity measure between two matrices represented
by a color scale. In Figure 3, we see large differences between the matrices.

Since we removed the attention matrix normalization and replaced it with a
LayerNorm Ba et al. (2016), we want to compare normalized attention matrices
instead. We divide each line of the matrix by the absolute value of the sum of
its elements such that the softmax attention matrix is unaffected and the
linear attention matrix is normalized. In Figure 4, we see significantly
higher between most matrices with some exceptions. These observations indicate
that the linear attention matrices derived from SUPRA are not an approximation
of the softmax matrices.

Figure 3: Representation of the cosine similarity and the distance between the
singular values of the softmax attention matrices compared to the SUPRA
attention matrices.

Figure 4: Representation of the cosine similarity and the distance between the
singular values of the _normalized_ softmax attention matrices compared to the
_normalized_ SUPRA attention matrices.