# Bi-selective pulses for large-area atom interferometry

Jack Saywell<EMAIL_ADDRESS>Max Carey  School of Physics &
Astronomy, University of Southampton, Highfield, Southampton, SO17 1BJ, UK
Ilya Kuprov  School of Chemistry, University of Southampton, Highfield,
Southampton, SO17 1BJ, UK  Tim Freegarde  School of Physics & Astronomy,
University of Southampton, Highfield, Southampton, SO17 1BJ, UK

###### Abstract

We present designs for the augmentation ‘mirror’ pulses of large-momentum-
transfer atom interferometers that maintain their fidelity as the wavepacket
momentum difference is increased. These bi-selective pulses, tailored using
optimal control methods to the evolving bi-modal momentum distribution, should
allow greater interferometer areas and hence increased inertial measurement
sensitivity, without requiring elevated Rabi frequencies or extended frequency
chirps. Using an experimentally validated model, we have simulated the
application of our pulse designs to large-momentum-transfer atom
interferometry using stimulated Raman transitions in a laser-cooled atomic
sample of 85Rb at 1 $\mu$K. After the wavepackets have separated by 42 photon
recoil momenta, our pulses maintain a fringe contrast of 90% whereas, for
adiabatic rapid passage and conventional $\pi$ pulses, the contrast is less
than 10%. Furthermore, we show how these pulses may be adapted to suppress the
detrimental off-resonant excitation that limits other broadband pulse schemes.

##  I Introduction

Atom interferometers Berman (1997) reverse the optical interferometry roles of
light and matter by using pulses of laser light to split, redirect, recombine
and interfere atomic matter waves, allowing the precise measurement of
gravitational fields Peters  _et al._ (2001); Mc Guirk  _et al._ (2002); Rosi
_et al._ (2014), inertial motions Gustavson  _et al._ (1997); Barrett  _et
al._ (2014); Hoth  _et al._ (2016) and more esoteric fields Hamilton  _et al._
(2015); Jaffe  _et al._ (2017) to which optical interferometers would have
little or no sensitivity.

Most atom interferometers for inertial sensing use Bragg Müller  _et al._
(2008) or Raman Kasevich and Chu (1991); Kasevich  _et al._ (1991) transitions
driven by counter-propagating laser pulses as the beamsplitters and mirrors
that split and direct the atomic wavepackets. The two-photon recoil
accompanying these stimulated scattering processes imparts a momentum
difference of 2$\hbar$k between the two interferometer paths, where k is the
single-photon wave number. As with an optical interferometer, the measurement
sensitivity depends upon the spatial area enclosed. This is proportional to
the momentum separation, and can hence be increased by using additional mirror
pulses - known as augmentation pulses - to impart further impulses to the
atomic wavepackets forming a large momentum transfer (LMT) interferometer
McGuirk  _et al._ (2000).

In practice, inhomogeneities such as variations in beam intensity and atomic
velocity reduce the fidelity of the augmentation pulses. The accrued effect of
such imperfections limits the fringe visibility and can reverse the LMT
sensitivity gains McGuirk  _et al._ (2000); Butts  _et al._ (2013); Kotru  _et
al._ (2015). Cooling the atoms or filtering their velocity distribution can
improve the initial fidelity at the expense of a weaker signal, but a velocity
spread is inevitable in LMT interferometry because of the increasing momentum
separation introduced between the primary interfering arms. Techniques such as
increasing the Rabi frequency Rudolph  _et al._ (2020) and employing beam
shaping optics Mielec  _et al._ (2018) increase the experimental complexity
and power requirements of the interferometer.

An alternative approach is to adopt robust atom-optics that retain high fringe
visibility for large momentum separations. Composite Levitt and Freeman
(1981); Dunning  _et al._ (2014) and adiabatic rapid passage (ARP) Bateman and
Freegarde (2007) Raman pulses have been investigated, achieving momentum
splittings of 18$\hbar$k Butts  _et al._ (2013) and 30$\hbar$k Kotru  _et al._
(2015) respectively. Both these techniques can increase the velocity
acceptance of the augmentation pulses at the expense of an increase in pulse
duration, but in practice fringe visibility was lost after 4-7 augmentation
pulses. Alternative atom-optics to Raman transitions such as multiphoton Bragg
pulses Müller  _et al._ (2008); Chiow  _et al._ (2011) and Bloch oscillations
within optical lattices Cladé  _et al._ (2009); Müller  _et al._ (2009);
McDonald  _et al._ (2013) have achieved impressive momentum transfer, but
these put even more stringent demands on the initial momentum distribution
Szigeti  _et al._ (2012) and, when used for interferometry, similarly limit
the signal-to-noise ratio (SNR) Butts  _et al._ (2013); Kotru  _et al._
(2015).

We have previously used optimal control techniques Khaneja  _et al._ (2005);
De Fouquieres  _et al._ (2011) to design robust high fidelity pulses for small
momentum transfer interferometry Saywell  _et al._ (2018, shed). In this
paper, we address LMT, by designing “bi-selective” pulses that offer high
fidelity for two ranges of velocity that track the atoms in the two arms of
the LMT interferometer, allowing high fringe contrast to be maintained for
larger interferometer areas. We compare our pulses with conventional composite
pulses and ARP through simulations using an previously validated approach
Dunning  _et al._ (2014); Saywell  _et al._ (shed); Kotru  _et al._ (2015),
and consider an extension to reduce unwanted double diffraction Lévèque  _et
al._ (2009); Malossi  _et al._ (2010). Our method is a departure from previous
approaches in which the same robust pulse has been used for every augmentation
pulse in the interferometer sequence, ultimately limiting the achievable
momentum splitting to the velocity acceptance of the pulse. The technique has
general applicability because many interferometer arrangements are already set
up to implement similar modulation sequences and, as the algorithm optimises
tolerance of variations, the designs do not depend critically upon
experimental parameters.

##  II LMT atom interferometers

Figure 1: Atomic state trajectories (left) as a function of position within an
LMT interferometer sequence showing the case where the beamsplitter and mirror
operations are extended by a sequence of 4 augmentation pulses (18$\hbar$k
trajectory). The initial momentum distribution (right), represented by the red
shaded region, is separated into two arms. The momentum distribution seen by
each pulse throughout the LMT sequence for 85Rb atoms with a Maxwell-Boltzmann
temperature of $\sim 1~{}\mu$K is shown by the blue shaded regions on the
right. As more pulses are added to increase the interferometer area, the
separation in the resonance conditions for the arms begins to exceed the
velocity acceptance of a $\pi$-pulse, shown bottom right for a Rabi frequency
of 200 kHz.

Figure 1 shows the space-time diagram of a typical Raman LMT interferometer.
While the standard $\pi/2-\pi-\pi/2$ interferometer sequence produces a
momentum splitting between the interferometer arms of
$\hbar\mathrm{k}_{\mathrm{eff}}$, where ${k}_{\mathrm{eff}}=2$k is the
effective wavevector of the atom-optics, a greater momentum splitting can be
achieved by augmenting the mirror and beamsplitter pulses with additional
pulses with alternating wavevectors McGuirk  _et al._ (2000). These
augmentation pulses are designed to transfer atomic population between two
internal states and increase the momentum splitting between the wavepackets in
the interferometer. By extending the beamsplitter operation by $\mathrm{N}$
augmentation pulses, the momentum splitting may be increased to
$(2\mathrm{N}+1)\hbar\mathrm{k}_{\mathrm{eff}}$, increasing the intrinsic
phase sensitivity of the interferometer.

However, this increased splitting also introduces differential Doppler shifts
between the interfering arms, quantised in multiples of the two-photon recoil
shift $\delta_{\mathrm{recoil}}=\hbar\mathrm{k}_{\mathrm{eff}}^{2}/2m$, that
depend on the atomic mass $m$. Indeed, this is automatic, since the Doppler
shift and inertial sensitivity have the same origin. For the $n$th
augmentation pulse ($n=1,2,3,...,\mathrm{N}$), the arms of the interferometer
have Raman resonance conditions separated by 4$n\delta_{\mathrm{recoil}}$
Butts  _et al._ (2013); Kotru (2016); Jaffe (2019) (for Raman transitions on
the 85Rb D2 line, $\delta_{\mathrm{recoil}}\approx 2\pi\times 15.4$ kHz). The
distribution of detunings for the $n$th augmentation pulse may thus be
visualised as two Gaussian distributions, corresponding to the Maxwell-
Boltzmann temperature of the atomic source, separated in frequency space by
4$n\delta_{\mathrm{recoil}}$. As illustrated in Figure 1, the augmentation
pulses are only efficient as long as this split distribution fits within the
velocity acceptance of the pulses. For conventional $\pi$-pulses this can span
just a few recoil momenta, fundamentally limiting the momentum transfer
achievable before fringe visibility is lost McGuirk  _et al._ (2000); Butts
_et al._ (2013).

Many composite and shaped pulses have been developed in the field of nuclear
magnetic resonance (NMR) spectroscopy to improve the control of nuclear spins
in the presence of experimental inhomogeneities. Composite pulses Levitt and
Freeman (1981) replace single rectangular pulses with sequences of pulses of
varying durations and phases. The overall effect of the pulse sequence is to
replicate the operation of a single pulse, but with an increased tolerance of
unwanted variations in coupling strength, detuning from resonance, or both.
Dunning et al. Dunning  _et al._ (2014) investigated the potential fidelity
improvements various established composite pulses might bring in atom
interferometry compared with rectangular $\pi$ pulses. Of the pulses tested,
the WALTZ pulse Shaka  _et al._ (1983), a relatively short three step robust
state-transfer pulse, had the highest fidelity and velocity acceptance.
Similarly, Butts et al. Butts  _et al._ (2013) employed WALTZ to improve the
contrast in an LMT interferometer, doubling the sensitivity.

An alternative class of robust pulses is known as frequency swept adiabatic
rapid passage (ARP). ARP can transfer atomic population between the states of
a two-level atom with high efficiency Bateman and Freegarde (2007); Kovachy
_et al._ (2012); Kotru  _et al._ (2015); Jaffe  _et al._ (2018). During an ARP
pulse, the driving field frequency is swept slowly through resonance such that
the atomic state follows the evolution of the field and may be moved with high
precision anywhere on the surface of the Bloch sphere Feynman  _et al._
(1957). In this picture, the quantum state vector precesses about an
instantaneous rotation axis, the field vector, defined by the amplitude,
detuning, and phase of the driving field. For ARP to be efficient, the pulse
must satisfy a condition of adiabaticity: the motion of the field vector on
the Bloch sphere must be slower than the rate at which the atomic state
precesses about it Bateman and Freegarde (2007). The simplest example of ARP
is the linear frequency chirp, where the pulse amplitude, or Rabi frequency,
is fixed, and the laser detuning is swept linearly through resonance. This
pulse can be efficient and robust, but it is necessarily long in comparison to
rectangular and composite pulses, thereby increasing the risk that coherence
is lost through spontaneous emission.

By allowing the pulse amplitude to vary, various schemes have been developed
which outperform the linear frequency chirp Hardy  _et al._ (1986); Baum  _et
al._ (1985); Tannús and Garwood (1996); Hwang  _et al._ (1998); Garwood and
DelaBarre (2001). These pulses have high efficiency and robustness,
maintaining adiabaticity while reducing pulse duration. An example is the
Tanh/Tan ARP pulse Hwang  _et al._ (1998), where the frequency sweep follows a
tangential function of time and the pulse amplitude follows a hyperbolic
tangential function of time. The Tanh/Tan pulse was implemented in a large-
area atom interferometer by Kotru et al. Kotru  _et al._ (2015) increasing the
interferometer contrast over that obtained using simple $\pi$ pulses and
achieving a momentum splitting of 30 $\hbar$k in a 9 $\mu$K 133Cs atomic
sample.

Although ARP can obtain impressive population transfer efficiency, with a
detuning bandwidth that increases with the pulse duration, its potential
utility in interferometry is limited by the dynamic phase imprinted on the
diffracting wavepackets and the variation in effective time origin Bateman and
Freegarde (2007); Kotru (2016); Jaffe  _et al._ (2018); Jaffe (2019). The
dynamic phase depends on the optical intensity, and rapid dephasing is
therefore inevitable when the atom cloud expands through variations in laser
intensity. In practice this effect is limited because ARP pulses applied in
quick succession approximately cancel the dynamic phase, but it leads to a
trade-off between longer and theoretically more efficient pulses, and
dephasing caused by imperfect dynamic phase cancellation when the beam quality
is non-ideal Kotru  _et al._ (2015). There is therefore a need for pulses
which match or improve upon ARP in terms of state transfer efficiency, but
which are robust to temporal variations in the Rabi frequency during the
interferometer sequence.

LMT interferometry requires augmentation pulses that provide efficient
population transfer across the atom cloud and throughout the sequence. To
maintain sensitivity and prevent loss of fringe visibility, the interferometer
sequence itself should impart the same phase to all atoms, and all atoms
should have the same sensitivity to the external influence being measured.
Pulses need not satisfy these conditions individually, provided that
subsequent cancellation achieves them for the sequence as a whole. ARP pulses,
for example, have detuning-dependent effective time origins and dynamic phases
which can be cancelled by a later pulse provided there is no variation in
optical intensity.

##  III Optimal control techniques

Robust pulses may be generated with optimal control techniques Kobzar  _et
al._ (2004, 2012) by dividing a pulse into discrete time slices and treating
the phase and/or amplitude of each slice as control parameters that may be
adjusted to optimise the fidelity of a desired optimisation.

In a typical Raman atom interferometer, two-photon transitions are driven by
lasers far from single-photon resonance so that the intermediate level can be
adiabatically eliminated to leave an effective two-level system between stable
states $\ket{\mathrm{g}}$ and $\ket{\mathrm{e}}$. The atom optics then effect
a rotation of the quantum state on the surface of the Bloch sphere for this
basis at a rate $\tilde{\Omega}_{R}=\sqrt{\Omega_{R}^{2}+\delta^{2}}$, about
an axis (the field vector)

$\mathbf{\Omega}=\Omega_{R}\cos(\phi_{L})\hat{x}+\Omega_{R}\sin(\phi_{L})\hat{y}+\delta\hat{z},$
(1)

determined by the relative laser phase $\phi_{L}$, the two-photon Rabi
frequency on resonance $\Omega_{R}$, and the Raman detuning

$\delta=(\omega_{1}-\omega_{2})-\omega_{\mathrm{eg}}+\delta_{\mathrm{Doppler}}+\delta_{\mathrm{recoil}}$
(2)

that includes terms for the two-photon recoil shift
$\delta_{\mathrm{recoil}}=\hbar\mathrm{k}_{\mathrm{eff}}^{2}/2m$ and the
Doppler shift
$\delta_{\mathrm{Doppler}}=\mathbf{k}_{\mathrm{eff}}\cdot\mathbf{p}/m$, which
depends on the initial momentum $\mathbf{p}$ of each atom in the
interferometer. $\omega_{1,2}$ are the laser frequencies and
$\omega_{\mathrm{eg}}$ is the frequency splitting of the levels
$\ket{\mathrm{g}}$ and $\ket{\mathrm{e}}$. This rotation can be written in
terms of a propagator

$\hat{U}=\begin{pmatrix}C^{*}&-iS^{*}\\\ -iS&C\end{pmatrix}$ (3)

acting on the ($\ket{g}$, $\ket{e}$) basis, where $C$ and $S$ are the
‘continuing’ and ‘scattering’ amplitudes defined as Stoner  _et al._ (2011)

$\displaystyle C$
$\displaystyle\equiv\cos\big{(}{\tilde{\Omega}_{R}\tau}\big{/}{2}\big{)}+i\big{(}{\delta}\big{/}{\tilde{\Omega}_{R}}\big{)}\sin\big{(}{\tilde{\Omega}_{R}\tau}\big{/}{2}\big{)}$
$\displaystyle S$ $\displaystyle\equiv
e^{i\phi_{L}}\big{(}{{\Omega}_{R}}\big{/}{\tilde{\Omega}_{R}}\big{)}\sin\big{(}{\tilde{\Omega}_{R}\tau}\big{/}{2}\big{)}.$
(4)

The conventional beamsplitter and mirror pulses of atom interferometry have a
fixed amplitude and phase, and the durations of the interactions are set so
that, on resonance, they perform $\pi/2$ and $\pi$ rotations respectively.
Composite and ARP pulses depend on varying the pulse parameters
$\Omega_{R}(t)$, $\phi_{L}(t)$, and $\delta(t)$ as a functions of time $t$ so
that, for a given atom, the rotation axis and rate change during the pulse.

The action of such a pulse on the quantum state can be evaluated efficiently
by dividing the pulse into discrete timesteps of duration
$\mathop{}\\!\mathrm{d}t$, with the propagator for an entire pulse given by
the time-ordered product of propagators for each slice, calculated according
to Equation (3). A pulse is then described by piece-wise constant waveforms
$\Omega_{R}(t)$, $\phi_{L}(t)$, and $\delta(t)$ that constitute a finite
number of control parameters. Upon defining a suitable fidelity, such as a
measure of how accurately a given initial state $\ket{\psi(t=0)}$ is driven to
a target $\ket{\psi_{\mathrm{T}}}$ by the pulse, optimal waveforms may be
found using numerical routines from optimal control theory. In this work we
focus on producing pulses with optimised phase and amplitude profiles
$\phi_{L}(t)$ and $\Omega_{R}(t)$.

One efficient and popular quantum control algorithm, originally developed for
NMR applications, is gradient ascent pulse engineering, or GRAPE Khaneja  _et
al._ (2005). GRAPE computes derivatives of the pulse fidelity with respect to
the control parameters without the need for computationally expensive finite-
differencing methods, and can be used in conjunction with the limited-memory
Broyden-Fletcher-Goldfarb-Shanno (L-BFGS) quasi-Newton method De Fouquieres
_et al._ (2011). By computing the gradient and estimating the Hessian of the
fidelity landscape, and provided local maxima with low-fidelity are avoided,
the algorithm can find an optimum solution efficiently.

We consider the following two possible fidelities for augmentation pulses,

$\displaystyle\mathcal{F}_{\mathrm{square}}$
$\displaystyle=|\braket{\mathrm{e}}{\hat{U}}{\mathrm{g}}|^{2}$ (5)
$\displaystyle\mathcal{F}_{\mathrm{real}}$
$\displaystyle=\mathrm{Re}(\braket{\mathrm{e}}{\hat{U}}{\mathrm{g}}).$ (6)

Both of these fidelities, if maximised, will yield pulses that efficiently
transfer population from one basis state to the other, the essential
requirement of an augmentation pulse. If the fidelities are averaged over an
ensemble of detunings and a range of Rabi frequencies, the resulting pulse
will be made robust to these specific errors. Maximising
$\mathcal{F}_{\mathrm{square}}$ leads to pulses where the phase of the overlap
$\braket{\mathrm{e}}{\hat{U}}{\mathrm{g}}$ is unimportant but the quantum
state is rotated by 180∘ from pole-to-pole on the Bloch sphere. Conversely,
maximising $\mathcal{F}_{\mathrm{real}}$ leads to pulses where the phase of
the overlap $\braket{\mathrm{e}}{\hat{U}}{\mathrm{g}}$ is well-defined.

If the phase of the overlap $\braket{\mathrm{e}}{\hat{U}}{\mathrm{g}}$ varies
from pulse-to-pulse across an atomic sample in an LMT sequence, the resulting
interference fringes can be washed out when the contributions from each atom
in the cloud are averaged at the end of the interferometer. Providing the Rabi
frequency does not change much between augmentation pulses, these phases
approximately cancel in LMT interferometers but if the Rabi frequency varies
temporally throughout the interferometer, rapid dephasing can occur. This is
readily observed with ARP, which requires a high degree of cancellation in the
phase factors introduced by each pulse in the sequence Kotru  _et al._ (2015);
Jaffe  _et al._ (2018); Jaffe (2019). Optimising $\mathcal{F}_{\mathrm{real}}$
for a range of Rabi frequencies will mean the phase of the overlap
$\braket{\mathrm{e}}{\hat{U}}{\mathrm{g}}$ remains fixed even if the coupling
strength varies within that range. Therefore, we expect
$\mathcal{F}_{\mathrm{real}}$ to yield augmentation pulses that are
insensitive to dephasing caused by temporal changes in the Rabi frequency due,
for example, to noisy beam intensity profiles and/or the ballistic expansion
of the atom cloud.

##  IV Bi-selective Raman pulses

For the nominal three-pulse $\pi/2-\pi-\pi/2$ interferometer, the Raman
detunings arise from the velocity components parallel to the Raman beam axis
and hence the temperature of the atomic cloud. As the sequence is extended by
augmentation pulses with alternating effective wavevector directions, the
resonant frequencies of each interferometer arm separate due to the imparted
photon momentum (shown on the right-hand side of Figure 1). This means the
detuning distribution splits in two, and continues to separate as more
momentum is imparted. Optimal control can optimise a pulse for any assumed
velocity distribution Kobzar  _et al._ (2005). We use this feature to design
“bi-selective” pulses, tailored to yield efficient population transfer only in
the frequency ranges occupied by the primary interferometer arms. By
optimising each pulse in the sequence individually, as opposed to employing a
single augmentation pulse design throughout the interferometer McGuirk  _et
al._ (2000); Butts  _et al._ (2013); Kotru  _et al._ (2015), we can therefore
take advantage of the fact that their transfer efficiency in the frequency
space in between the two arms is unimportant. This allows bi-selective pulses
to be shorter than ARP, composite or optimised shaped pulses of the equivalent
efficiency.

We express the bi-selective fidelity for the $n$th augmentation pulse as

$\mathcal{F}^{\mathrm{A}}_{n}=\sum_{\delta,\Omega_{R}\in\mathrm{L}_{n}}\mathcal{F}(\delta,\Omega_{R})+\sum_{\delta,\Omega_{R}\in\mathrm{U}_{n}}\mathcal{F}(\delta,\Omega_{R}),$
(7)

where $\mathcal{F}=\mathcal{F}_{\mathrm{real},\ \mathrm{square}}$ is the
single-atom augmentation pulse fidelity (Equations 5 and 6) and
$\mathrm{L}_{n}$ and $\mathrm{U}_{n}$ are the ensembles representing the
atomic frequency distributions of the lower and upper arms during the $n$th
augmentation pulse respectively. In order to normalise the fidelity such that
the maximum value is unity, we divide Equation 7 by the number of detunings
and amplitude errors included in the entire ensemble.

We compose the detuning ensembles for the lower and upper interferometer arms
for the $n$th augmentation pulse from two uniform discrete distributions
centred at $\pm 2n\delta_{\mathrm{recoil}}$. The ensembles used in the
optimisation should represent the velocity distribution of the atoms. However,
in order to reduce computation time we approximate the true distribution using
two uniform distributions each with a sample size of 20. The range spanned by
each distribution is given by 4 standard deviations of detuning arising from
the velocity distribution along the Raman beam axis, which we assume follows a
Maxwell Boltzmann distribution.

Our LMT pulses are individually optimised, which assumes there are no
correlations in residual errors between pulses. However, It may be possible to
develop alternative measures of performance which reflect the fidelity of the
entire interferometer and allow pulses within a sequence to be optimised
cooperatively Braun and Glaser (2014), for example by compensating each
other’s imperfections.

An alternative approach to achieve bi-selectivity involves concatenating two
inversion pulses with the appropriate frequency shifts to address each
interferometer arm separately. While this makes the individual optimisations
simpler, the resulting pulses are longer, give different velocity classes
different pulse origins, and lose fidelity as the classes began to overlap. A
simple superposition of the two components would cure the problems of length
and origin, though not overlap, and requires amplitude modulation in addition
Emsley  _et al._ (1990).

The concept of bi-selective pulses for LMT interferometry bears resemblance to
the technique of band-selective pulses in NMR spectroscopy Emsley  _et al._
(1990); Geen and Freeman (1991); Freeman (1998); Janich  _et al._ (2011),
where pulses have been designed to excite or invert nuclear spins within
single or multiple frequency bands, but suppress the response of spins outside
the desired frequency range. Typically, these pulses require smooth waveforms
because the response at large resonance offsets is determined by the Fourier
transform of the pulse shape Warren (1984). As a result, composite pulses,
which are composed of concatenated sequences of constant amplitude pulses with
discrete phase shifts, lead to non-negligible excitation far off-resonance.

The suppression of the action of a pulse outside a specific frequency range
may be useful in interferometer geometries where, for example, the counter-
propagating Raman beams are obtained by retroreflection of a single beam with
both Raman frequencies and there are necessarily four frequency components
that may interact with the atom cloud. When there is no acceleration of the
atoms along the beam axis then double-diffraction interferometry schemes can
be be employed that make use of all of the frequency components simultaneously
Lévèque  _et al._ (2009); Malossi  _et al._ (2010) . However, in vertically-
orientated, ground-based atom interferometers, the Doppler shift caused by
gravitational acceleration is commonly used to isolate a single frequency pair
by shifting the other pair off resonance Butts  _et al._ (2013). However,
broadband pulses such as composite or ARP pulses often have non-zero transfer
efficiency at large detunings Kotru (2016); Warren (1984), meaning one must
wait longer to ensure negligible excitations from the off-resonant pair,
costing time which may otherwise be used to enhance the sensitivity.

By directly suppressing the transfer efficiency outside a specific range of
frequencies, such off-resonant excitation with broadband pulses may be
avoided. We achieve this suppression of unwanted excitation by modifying our
bi-selective pulse fidelity (Equation 7), adding the following penalty term
which represents the suppression band of detunings where the transfer
efficiency should be minimised,

$\mathcal{F}_{\mathrm{suppress}}=\sum_{\Omega_{R},\delta=\delta_{\mathrm{min}}}^{\delta_{\mathrm{max}}}|\braket{\mathrm{g}}{\hat{U}}{\mathrm{g}}|^{2}+\sum_{\Omega_{R},\delta=-\delta_{\mathrm{min}}}^{-\delta_{\mathrm{max}}}|\braket{\mathrm{g}}{\hat{U}}{\mathrm{g}}|^{2}.$
(8)

Here, $\delta_{\mathrm{min},\mathrm{max}}$ represent the initial and final
Raman detuning values for the two suppression bands.

##  V Results

Figure 2: Bi-selective phase-modulated augmentation pulses optimised with the
phase-sensitive fidelity $\mathcal{F}_{\mathrm{real}}$ for a temperature of 1
$\mu$K. The resulting phase profiles for the 1st, 4th, 7th, and 10th
augmentation pulses are shown in the left panel. The right panel shows the
simulated transfer efficiency of each bi-selective pulse (blue solid curve),
the ARP Tanh/Tan pulse of equivalent length (orange dashed curve), and the
WALTZ composite pulse (purple dotted curve) as a function of the Raman
detuning. The corresponding detuning distributions seen by each augmentation
pulse for the primary interferometer arms are shown by the shaded regions. The
range of detunings included in each optimisation is shown by the hatched
regions.

We have optimised phase-modulated, constant amplitude, bi-selective
augmentation pulses with GRAPE, using both the phase insensitive fidelity
$\mathcal{F}_{\mathrm{square}}$ and the phase sensitive fidelity
$\mathcal{F}_{\mathrm{real}}$. We have investigated their performance through
simulation of the resulting interferometer contrast in a laser-cooled sample
of 85Rb at 1 $\mu$K. In all optimisations, the timestep of the pulses was 25
$n$s, and the effective Rabi frequency was 200 kHz, meaning the duration of a
rectangular $\pi$ pulse ($t_{\pi}$) should be 2.5 $\mu$s in the absence of any
detuning or amplitude error. The pulses were optimised for an atomic
temperature of 1 $\mu$K, and a range of amplitude errors of $\pm$10 % the
effective Rabi frequency. The length of the augmentation pulses was fixed to
be 4$t_{\pi}$. This duration allowed for a sufficiently good pulse to be
obtained when using both fidelities, although optimising
$\mathcal{F}_{\mathrm{real}}$ typically requires a longer pulse to reach an
equivalent fidelity to the phase-insensitive case. The initial guess for the
phase profile of the pulses was in each case a sequence of random phases. We
keep the pulse amplitude constant when optimising the bi-selective fidelity
$\mathcal{F}^{\mathrm{A}}_{n}$ (Equation 7) because allowing the amplitude to
vary does not lead to a higher terminal fidelity in this case.

The NMR spin simulation software suite for MATLAB, Spinach Hogben  _et al._
(2011), and its optimal control module, were modified to optimise bi-selective
pulses using the L-BFGS GRAPE method De Fouquieres  _et al._ (2011). Each
pulse optimisation was set to terminate following 300 iterations or when
either the norm of the fidelity gradient became smaller than 10-7 or the norm
of the step size dropped below 10-3. The resulting waveforms, optimising the
fidelity $\mathcal{F}_{\mathrm{real}}$, are shown in Figure 2, for pulses
tailored to the velocity distributions expected for the 1st, 4th, 7th, and
10th augmentation pulses of the LMT sequence. Efficient population transfer is
achieved in each case. Interestingly, smooth and symmetrical waveforms are
found despite there being no constraint on symmetry or waveform smoothness
included in the optimisation.

Figure 3 shows the phase profile and simulated transfer efficiency of the 10th
augmentation pulse when the fidelity $\mathcal{F}_{\mathrm{square}}$ is
instead optimised. Results are also shown for the WALTZ composite pulse and
the Tanh/Tan ARP pulse with durations of 4$t_{\pi}$ and 8$t_{\pi}$. The
corresponding detuning distributions for the upper and lower interferometer
arms are also shown as shaded regions. Whereas the efficiencies of the WALTZ
and ARP augmentation pulses are limited by the finite velocity acceptances,
the bi-selective pulse is tailored for the split momentum distribution of the
two primary interferometer arms, allowing it to achieve a higher efficiency
than a Tanh/Tan ARP pulse of twice its duration. Many of the profiles found
optimising $\mathcal{F}_{\mathrm{square}}$ have a strong parabolic component
corresponding to a frequency chirp, although we are far from the adiabatic
regime; and some include a $\pi$ phase step near the temporal midpoint: we
find that this combination alone gives the main features but does not produce
the flat regions of high efficiency achieved using optimal control. All the
pulses found optimising $\mathcal{F}_{\mathrm{real}}$ are time-symmetric.

Figure 3: Phase profile (A) and simulated transfer efficiency (B) for the 10th
bi-selective pulse found optimising the fidelity
$\mathcal{F}_{\mathrm{square}}$ (blue solid line). The transfer efficiency for
the ARP Tanh/Tan pulse (orange dashed line and green dot-dashed line) and the
WALTZ pulse (purple solid line) are shown in panel C. The detuning
distributions of the two primary interferometer arms, corresponding to a
temperature of 1 $\mu$K during the 10th augmentation pulse are shown by shaded
regions. The hatched regions represent the bands of detuning where the
transfer efficiency is maximised using GRAPE.

###  V.1 LMT contrast

We have simulated LMT interferometers of different LMT orders, and compared
the performance using our bi-selective augmentation pulses with that using
rectangular $\pi$ and Tanh/Tan ARP augmentation pulses. The contrast in each
case is averaged over a thermal distribution of atoms and a uniform
distribution of Rabi frequencies that is either kept temporally constant – to
represent a non-uniform laser intensity distribution – or else varied from
pulse-to-pulse – to represent the motion of atoms across such a distribution.
Following the approach taken by Kotru Kotru (2016), only the primary
interfering paths are included in the calculation to reduce computation time.

The simulation results for an atomic temperature of 1 $\mu$K are shown in
Figure 4. To meaningfully compare the performance of the GRAPE bi-selective
pulses with ARP, the total duration of the sequences is kept the same. GRAPE
bi-selective pulses far outperform basic rectangular $\pi$ pulses, and
Tanh/Tan ARP pulses of equivalent duration. When no thermal expansion is
modelled, the dynamic phase cancellation of the ARP pulses is perfect, meaning
the reduction in contrast with increasing LMT order is purely due to the
limited velocity acceptance of the pulses as the arms separate in detuning.
For the ARP contrast to match that achievable with bi-selective pulses, the
ARP pulse duration must be increased, thus increasing the susceptibility to
spontaneous emission (not modelled currently) and dynamic phase dephasing
Kotru  _et al._ (2015). This highlights one key advantage of our adaptive
approach.

Figure 4: Simulated interferometer contrast for a cloud of 85Rb at 1$\mu$K as
a function of LMT order for different pulse sequences. The contrast is
numerically averaged over Raman detunings due to a velocity distribution
corresponding to a temperature of 1$\mu$K and a uniform distribution of static
Rabi rate errors of $\pm$10% (A) or a random temporal variation of Rabi rate
errors of $\pm$10% from pulse-to-pulse (B). Contrast values for bi-selective
GRAPE pulses found with fidelities $\mathcal{F}_{\mathrm{square}}$ and
$\mathcal{F}_{\mathrm{real}}$ (blue dotted line with diamonds and red solid
line with diamonds respectively) are shown. The results for $\pi$ augmentation
pulses (black dotted line with rectangles) and the Tanh/Tan ARP pulse (dashed
orange line with circles) are also shown.

If the Rabi frequency varies from pulse-to-pulse for different atoms, the
dynamic phase imprinted on the atoms by ARP is not cancelled in the
interferometer, meaning different atoms obtain different phases and the
interference is washed out. Although ARP pulses have a velocity acceptance
which increases with pulse duration, longer ARP pulses become more susceptible
to errors in dynamic phase cancellation Kotru  _et al._ (2015).

When the Rabi rate is instead varied randomly between pulses in the simulation
to emulate noise in the Raman beam intensity profile (panel B in Figure 4),
the bi-selective pulses found optimising the phase-sensitive fidelity
$\mathcal{F}_{\mathrm{real}}$ are able to maintain significantly higher
contrast than other pulse sequences. This is because the phase variation of
the wavepackets is minimised with respect to variations in the Rabi frequency.
This is not the case with pulses found using $\mathcal{F}_{\mathrm{square}}$
and the contrast is significantly reduced in such cases.

We have repeated the bi-selective optimisation and contrast simulation for a
hotter atomic temperature of 5 $\mu$K and for a longer pulse duration of
$5t_{\pi}$. The results are summarised in Table 1. We also show the results
for the WALTZ composite pulse and a phase-modulated pulse obtained using GRAPE
maximising $\mathcal{F}_{\mathrm{square}}$ for a large range of detunings
centred on resonance. This non-selective pulse was optimised following the
procedure outlined in Saywell  _et al._ (2018). Longer pulses can achieve
higher terminal fidelities, and using the phase-sensitive fidelity
$\mathcal{F}_{\mathrm{real}}$ requires a longer duration than
$\mathcal{F}_{\mathrm{square}}$ to reach an equivalent terminal fidelity. This
is shown most clearly when optimising for the hotter cloud. There is no
guarantee that we have found the global maximum for each choice of duration
and temperature. Even so, pulses were found which outperform the composite and
ARP alternatives tested.

Table 1:  Simulated contrast values at temperatures of 1 and 5 $\mu$K for different LMT orders with no Rabi frequency variation between pulses. The $\pi$, WALTZ, and $\mathcal{F}_{\mathrm{square}}$ non-selective pulses are previous results with either defined ($\pi$, WALTZ) or previously chosen (non-selective) durations. The results for the ARP and bi-selective pulses are presented for two different durations. Bold values indicate the best performing pulse at each temperature.  | Length | 1 $\mathbf{\mu}$K Contrast | 5 $\mathbf{\mu}$K Contrast
---|---|---|---
Pulse |  ($T/t_{\pi}$) |  10$\hbar$k |  26$\hbar$k |  42$\hbar$k |  10$\hbar$k |  26$\hbar$k |  42$\hbar$k
$\pi$ | 1 | 0.50 | - | - | 0.35 | - | -
WALTZ | 3 | 0.73 | 0.51 | 0.00 | 0.72 | 0.41 | 0.00
$\mathcal{F}_{\mathrm{square}}$ non-selective | 4 | 0.92 | 0.72 | 0.38 | 0.85 | 0.67 | 0.34
| 4 | 0.95 | 0.62 | 0.06 | 0.88 | 0.52 | 0.04
ARP | 5 | 0.96 | 0.79 | 0.19 | 0.90 | 0.69 | 0.15
| 4 | 0.91 | 0.78 | 0.69 | 0.79 | 0.43 | 0.22
$\mathcal{F}_{\mathrm{real}}$ bi-selective | 5 | 0.95 | 0.87 | 0.82 | 0.85 | 0.60 | 0.50
| 4 | 0.97 | 0.92 | 0.90 | 0.83 | 0.64 | 0.57
$\mathcal{F}_{\mathrm{square}}$ bi-selective | 5 | 0.96 | 0.94 | 0.93 | 0.89 | 0.81 | 0.76

###  V.2 Suppression of off-resonant excitation

When the Doppler shift from gravitational acceleration is used to discriminate
between retroreflected frequency pairs in vertically-orientated
interferometers, off-resonant excitations from broadband atom-optics can lead
to unwanted double-diffraction. For example, following a 10 ms drop time Butts
_et al._ (2013), the resonance conditions of the two frequency pairs will
shift by $2g\mathrm{k_{\mathrm{eff}}}\times 10\ \mathrm{ms}\approx$ 500 kHz
for 85Rb, where $g$ is the local gravitational acceleration. Conventional
robust pulses can result in a non-zero transfer at comparable detunings
(Figure 5), potentially leading to double-diffraction if the unwanted
frequency pair has not been shifted far enough from resonance Butts  _et al._
(2013); Kotru (2016).

We have optimised bi-selective pulses which suppress the transfer efficiency
outside the detuning bands of interest using a suitable modification of our
bi-selective fidelity (Equation 8). Figure 5 shows the preliminary results. In
order to achieve a good level of suppression at large detunings we have found
it necessary to allow both the phase and amplitude of the pulse to vary in the
optimisation procedure. We have also added penalties in the optimisation to
limit the maximum intensity during the pulse and enforce smooth waveforms
Goodwin and Kuprov (2016).

Figure 5: Panels A and C show the phase and amplitude profiles for two bi-
selective pulses maximising $\mathcal{F}_{\mathrm{square}}$. The dashed blue
curve shows the case where no suppression of off-resonant excitation is
included in the optimisation and the solid red curve shows the case where off-
resonant excitation is suppressed. The transfer efficiency of these two pulses
as a function of detuning is shown in panel B, and panel D shows the transfer
efficiency of the WALTZ and Tanh/Tan pulses. The diagonally hatched regions
coinciding with the shaded detuning distributions show the optimisation bands
of detuning during the 8th augmentation pulse. The outer hatched regions show
the range of the detunings for which the state transfer is suppressed
($\delta_{\mathrm{min}},\delta_{\mathrm{max}}=500,2500$ kHz). The response is
unconstrained outside the hatched regions.

##  VI Conclusions

We have presented a new Raman pulse scheme for the augmentation pulses in
large-area atom interferometry, whereby individually tailored Raman pulses
found using optimal control techniques maintain resonance with the diverging
wavepackets as the Raman detuning increases between the interferometer arms.
We optimise our bi-selective pulses to provide maximum state-transfer while
minimising variation in the interferometer phase across the atomic ensemble.
The pulses can be made robust to large spatial and temporal variations in the
Rabi rate, potentially allowing LMT interferometry in non-ideal experimental
environments such as those with warmer atom clouds and inhomogeneous laser
beam fronts. Our simulations show that our pulses can maintain contrast at
significantly higher momentum splittings than interferometers that have
employed the best augmentation pulses demonstrated to-date, including the
WALTZ composite pulse and the Tanh/Tan adiabatic rapid passage pulse, whose
finite velocity acceptances limit the LMT momentum range. For large LMT
orders, our pulses can be considerably shorter than ARP or composite
equivalents of the same efficiency, reducing the susceptibility to spontaneous
emission.

*

##  Appendix A Numerical model

We model 85Rb atoms undergoing Raman transitions as two level systems,
described by the basis states $\ket{\mathrm{g},\mathbf{p}}$ and
$\ket{\mathrm{e},\mathbf{p}+\hbar\mathbf{k}_{\mathrm{eff}}}$ with
corresponding time-dependent amplitudes $c_{\mathrm{g,e}}(t)$. Pulses are
described by propagators of the form given by Equation 3, and are given by
time-ordered products of the propagators for individual slices in the case of
shaped and composite pulses, where the amplitude, detuning (frequency), and
phase may be varied for each step.

The basis states for each pulse in an LMT sequence vary. For example, during
the initial $\pi/2$ pulse the basis is
$\ket{\mathrm{g},\mathbf{p}}$,$\ket{\mathrm{e},\mathbf{p}+\hbar\mathbf{k}_{\mathrm{eff}}}$,
but for the first augmentation pulse (where the wavevector is reversed) the
upper arm has a basis described by the states
$\ket{\mathrm{g},\mathbf{p}+2\hbar\mathbf{k}_{\mathrm{eff}}}$,$\ket{\mathrm{e},\mathbf{p}+\hbar\mathbf{k}_{\mathrm{eff}}}$,
and the lower arm has a basis described by the states
$\ket{\mathrm{g},\mathbf{p}}$,$\ket{\mathrm{e},\mathbf{p}-\hbar\mathbf{k}_{\mathrm{eff}}}$.
This means the resonance conditions are separated for the two interferometer
arms by 4$\delta_{\mathrm{recoil}}$. Furthermore, to reduce computation time,
only the amplitudes following the two primary interferometer arms are included
in the calculation, as shown in Figure 6. This is a good assumption when the
pulses providing the large-momentum transfer are efficient, or when there is
no interference from atoms that have not followed the primary interferometer
paths Kotru (2016).

Figure 6: Diagram of the first half of an N=1 LMT pulse sequence, shown up to
the central mirror pulse, indicating how the momentum and atomic state vary in
each arm of the interferometer during the sequence. The resonant Raman
frequency depends upon the pair of states coupled at each stage. ‘A’
represents an augmentation pulse. Single and double arrows represent the
internal states $\ket{\mathrm{g}}$ and $\ket{\mathrm{e}}$ respectively.

We simulate how the interferometer contrast is affected by the atomic
temperature and variations in the Rabi rate across the atom cloud. In order to
do this, we draw a sample of atomic velocities (and hence Raman detunings)
from a Maxwell-Boltzmann distribution for 85Rb at a given temperature and a
uniform distribution of Rabi frequencies in the range $\pm$10% of the intended
Rabi rate. Alternatively, the Rabi rate may be randomly varied from pulse-to-
pulse within $\pm$10% of the intended Rabi rate in order to explore the
robustness of sequences to temporal variation throughout the interferometer.

The contrast following an interferometer pulse sequence may be calculated by
evolving the state amplitudes for an atom initially in the ground internal
state, $\ket{\mathrm{g}}$, by applying the relevant pulse propagators in the
correct order with the correct Raman detunings for each pulse and
interferometer arm. Fringes were calculated for each atom in the ensemble by
repeating the evolution of the final rectangular beamsplitter and varying the
phase, $\phi_{bs}$, between 0 and $2\pi$. The ensemble average interferograms
were fitted to the sinusoidal function $0.5(A+B\cos(\phi_{bs}+C))$, where $A$
is an offset, $B$ is the contrast, and $C$ is a possible phase shift. The
effect of spontaneous emission is not modelled at present but will ultimately
limit the achievable contrast with extended pulse sequences in LMT
interferometers.

The sweep parameters used to define the Tanh/Tan ARP pulse are the same as
those used by Kotru et al. Kotru  _et al._ (2015). The Tanh/Tan ARP pulse
waveforms were divided into 2500 time-steps in the simulations to ensure the
piece-wise constant approximation was sufficiently accurate to model the
frequency and amplitude sweep.

###### Acknowledgements.

This work was supported by Dstl (Grant Nos. DSTLX-1000091758 and
DSTLX-1000097855) and the UK Engineering and Physical Sciences Research
Council (Grant Nos. EP/M013294/1 and EP/L015382/1).

## References

  * Berman (1997) P. Berman, ed.,  _Atom Interferometry_ (Academic Press, 1997) p. 363.
  * Peters  _et al._ (2001) A. Peters, K. Y. Chung, and S. Chu, Metrologia 38, 25 (2001).
  * Mc Guirk  _et al._ (2002) J. M. Mc Guirk, G. T. Foster, J. B. Fixler, M. J. Snadden, and M. A. Kasevich, Phys. Rev. A 65, 033608 (2002).
  * Rosi  _et al._ (2014) G. Rosi, F. Sorrentino, L. Cacciapuoti, M. Prevedelli, and G. M. Tino, Nature 510, 518 (2014).
  * Gustavson  _et al._ (1997) T. L. Gustavson, P. Bouyer, and M. A. Kasevich, Phys. Rev. Lett. 78, 2046 (1997).
  * Barrett  _et al._ (2014) B. Barrett, R. Geiger, I. Dutta, M. Meunier, B. Canuel, A. Gauguet, P. Bouyer, and A. Landragin, C. R. Phys. 15, 875 (2014).
  * Hoth  _et al._ (2016) G. W. Hoth, B. Pelle, S. Riedl, J. Kitching, and E. A. Donley, Appl. Phys. Lett. 109, 071113 (2016).
  * Hamilton  _et al._ (2015) P. Hamilton, M. Jaffe, P. Haslinger, Q. Simmons, H. Muller, and J. Khoury, Science 349, 849 (2015).
  * Jaffe  _et al._ (2017) M. Jaffe, P. Haslinger, V. Xu, P. Hamilton, A. Upadhye, B. Elder, J. Khoury, and H. Müller, Nat. Phys. 13, 938 (2017).
  * Müller  _et al._ (2008) H. Müller, S.-w. Chiow, Q. Long, S. Herrmann, and S. Chu, Phys. Rev. Lett. 100, 180405 (2008).
  * Kasevich and Chu (1991) M. Kasevich and S. Chu, Phys. Rev. Lett. 67, 181 (1991).
  * Kasevich  _et al._ (1991) M. Kasevich, D. S. Weiss, E. Riis, K. Moler, S. Kasapi, and S. Chu, Phys. Rev. Lett. 66, 2297 (1991).
  * McGuirk  _et al._ (2000) J. M. McGuirk, M. J. Snadden, and M. A. Kasevich, Phys. Rev. Lett. 85, 4498 (2000).
  * Butts  _et al._ (2013) D. L. Butts, K. Kotru, J. M. Kinast, A. M. Radojevic, B. P. Timmons, and R. E. Stoner, J. Opt. Soc. Am. B 30, 922 (2013).
  * Kotru  _et al._ (2015) K. Kotru, D. L. Butts, J. M. Kinast, and R. E. Stoner, Phys. Rev. Lett. 115, 103001 (2015).
  * Rudolph  _et al._ (2020) J. Rudolph, T. Wilkason, M. Nantel, H. Swan, C. M. Holland, Y. Jiang, B. E. Garber, S. P. Carman, and J. M. Hogan, Phys. Rev. Lett. 124, 083604 (2020).
  * Mielec  _et al._ (2018) N. Mielec, M. Altorio, R. Sapam, D. Horville, D. Holleville, L. A. Sidorenkov, A. Landragin, and R. Geiger, Appl. Phys. Lett. 113, 161108 (2018).
  * Levitt and Freeman (1981) M. H. Levitt and R. Freeman, J. Magn. Reson. (1969) 43, 65 (1981).
  * Dunning  _et al._ (2014) A. Dunning, R. Gregory, J. Bateman, N. Cooper, M. Himsworth, J. A. Jones, and T. Freegarde, Phys. Rev. A 90, 033608 (2014).
  * Bateman and Freegarde (2007) J. Bateman and T. Freegarde, Phys. Rev. A 76, 013416 (2007).
  * Chiow  _et al._ (2011) S. W. Chiow, T. Kovachy, H. C. Chien, and M. A. Kasevich, Phys. Rev. Lett. 107, 130403 (2011).
  * Cladé  _et al._ (2009) P. Cladé, S. Guellati-Khélifa, F. Nez, and F. Biraben, Phys. Rev. Lett. 102, 240402 (2009).
  * Müller  _et al._ (2009) H. Müller, S.-w. Chiow, S. Herrmann, and S. Chu, Phys. Rev. Lett. 102, 240403 (2009).
  * McDonald  _et al._ (2013) G. D. McDonald, C. C. N. Kuhn, S. Bennetts, J. E. Debs, K. S. Hardman, M. Johnsson, J. D. Close, and N. P. Robins, Phys. Rev. A 88, 053620 (2013).
  * Szigeti  _et al._ (2012) S. S. Szigeti, J. E. Debs, J. J. Hope, N. P. Robins, and J. D. Close, New J. Phys. 14, 023009 (2012).
  * Khaneja  _et al._ (2005) N. Khaneja, T. Reiss, C. Kehlet, T. Schulte-Herbrüggen, and S. J. Glaser, J. Magn. Reson. 172, 296 (2005).
  * De Fouquieres  _et al._ (2011) P. De Fouquieres, S. G. Schirmer, S. J. Glaser, and I. Kuprov, J. Magn. Reson. 212, 412 (2011).
  * Saywell  _et al._ (2018) J. C. Saywell, I. Kuprov, D. Goodwin, M. Carey, and T. Freegarde, Phys. Rev. A 98, 023625 (2018).
  * Saywell  _et al._ (shed) J. C. Saywell, M. Carey, M. Belal, I. Kuprov, and T. Freegarde, J. Phys. B (to be published).
  * Lévèque  _et al._ (2009) T. Lévèque, A. Gauguet, F. Michaud, F. Pereira Dos Santos, and A. Landragin, Phys. Rev. Lett. 103, 080405 (2009).
  * Malossi  _et al._ (2010) N. Malossi, Q. Bodart, S. Merlet, T. Lévèque, A. Landragin, and F. Pereira Dos Santos, Phys. Rev. A 81, 013617 (2010).
  * Kotru (2016) K. Kotru,  _Timekeeping and Accelerometry with Robust Light Pulse Atom Interferometers_ , Ph.D. thesis, Massachusetts Institute of Technology (2016).
  * Jaffe (2019) M. Jaffe,  _Atom Interferometry in an Optical Cavity_ , Ph.D. thesis, University of California, Berkeley (2019).
  * Shaka  _et al._ (1983) A. J. Shaka, J. Keeler, T. Frenkiel, and R. Freeman, J. Magn. Reson. (1969) 52, 335 (1983).
  * Kovachy  _et al._ (2012) T. Kovachy, S.-w. Chiow, and M. A. Kasevich, Phys. Rev. A 86, 011606(R) (2012).
  * Jaffe  _et al._ (2018) M. Jaffe, V. Xu, P. Haslinger, H. Müller, and P. Hamilton, Phys. Rev. Lett. 121, 040402 (2018).
  * Feynman  _et al._ (1957) R. P. Feynman, F. L. Vernon, and R. W. Hellwarth, J. Appl. Phys. 28, 49 (1957).
  * Hardy  _et al._ (1986) C. Hardy, W. Edelstein, and D. Vatis, J. Magn. Reson. (1969) 66, 470 (1986).
  * Baum  _et al._ (1985) J. Baum, R. Tycko, and A. Pines, Phys. Rev. A 32, 3435 (1985).
  * Tannús and Garwood (1996) A. Tannús and M. Garwood, J. Magn. Reson., Series A 120, 133 (1996).
  * Hwang  _et al._ (1998) T.-L. Hwang, P. C. van Zijl, and M. Garwood, J. Magn. Reson. 133, 200 (1998).
  * Garwood and DelaBarre (2001) M. Garwood and L. DelaBarre, J. Magn. Reson. 153, 155 (2001).
  * Kobzar  _et al._ (2004) K. Kobzar, T. E. Skinner, N. Khaneja, S. J. Glaser, and B. Luy, J. Magn. Reson. 170, 236 (2004).
  * Kobzar  _et al._ (2012) K. Kobzar, S. Ehni, T. E. Skinner, S. J. Glaser, and B. Luy, J. Magn. Reson. 225, 142 (2012).
  * Stoner  _et al._ (2011) R. Stoner, D. Butts, J. Kinast, and B. Timmons, J. Opt. Soc. Am. B 28, 2418 (2011).
  * Kobzar  _et al._ (2005) K. Kobzar, B. Luy, N. Khaneja, and S. J. Glaser, J. Magn. Reson. 173, 229 (2005).
  * Braun and Glaser (2014) M. Braun and S. J. Glaser, New J. Phys. 16, 115002 (2014).
  * Emsley  _et al._ (1990) L. Emsley, I. Burghardt, and G. Bodenhausen, J. Magn. Reson. (1969) 90, 214 (1990).
  * Geen and Freeman (1991) H. Geen and R. Freeman, J. Magn. Reson. (1969) 93, 93 (1991).
  * Freeman (1998) R. Freeman, Prog. Nucl. Mag. Res. Spectrosc. 32, 59 (1998).
  * Janich  _et al._ (2011) M. A. Janich, R. F. Schulte, M. Schwaiger, and S. J. Glaser, J. Magn. Reson. 213, 126 (2011).
  * Warren (1984) W. S. Warren, J. Chem. Phys. 81, 5437 (1984).
  * Hogben  _et al._ (2011) H. J. Hogben, M. Krzystyniak, G. T. P. Charnock, P. J. Hore, and I. Kuprov, J. Magn. Reson. 208, 179 (2011).
  * Goodwin and Kuprov (2016) D. L. Goodwin and I. Kuprov, J. Chem. Phys. 144, 204107 (2016).


#  gzip Predicts Data-dependent Scaling Laws

Rohan Pandey
Reworkd
San Francisco, CA
<EMAIL_ADDRESS>

###### Abstract

Past work has established scaling laws that predict the performance of a
neural language model (LM) as a function of its parameter count and the number
of tokens it’s trained on, enabling optimal allocation of a fixed compute
budget. Are these scaling laws agnostic to training data as some prior work
suggests? We generate training datasets of varying complexities by modulating
the syntactic properties of a PCFG, finding that 1) scaling laws are sensitive
to differences in data complexity and that 2) gzip, a compression algorithm,
is an effective predictor of how data complexity impacts scaling properties.
We propose a new data-dependent scaling law for LM’s that accounts for the
training data’s gzip-compressibility; its compute-optimal frontier increases
in dataset size preference (over parameter count preference) as training data
becomes harder to compress.

##  1 Introduction

A neural network’s performance generally increases as more compute is
allocated for training. When scaling compute, one must decide whether to
increase a model’s parameter count or increase the dataset’s size—these must
trade-off within a fixed compute budget. Scaling laws can tell us what
specific allocation (i.e. parameters v.s. data) will maximize performance
given a compute budget. Much work has explored scaling laws for neural LM’s
[Kaplan et al., 2020], generally concluding that parameter & training token
count should be scaled 1-to-1 [Hoffmann et al., 2022].

However, most prior work on scaling laws for LM’s have been estimated from
transformers trained on scraped web text. Of course, this is quite a specific
data distribution, so we may naturally ask whether the scaling laws
extrapolated from such web text datasets generalize to other distributions.
Furthermore, it is generally understood that the art of training data mixture
is the ‘secret sauce’ that enables frontier industry labs to continually
deliver state-of-the-art LLMs [Penedo et al., 2024, 2023, Xie et al., 2024].
Considering that improving data quality can significantly raise LM performance
[Gunasekar et al., 2023] and that scaling laws in reinforcement learning have
been shown to scale with game difficulty [Jones, 2021], we may hypothesize
that current LM scaling laws (e.g. Chinchilla [Hoffmann et al., 2022]) are
individual web-text-specific cases of a more general scaling law conditioned
on properties of the training data.

Then, what properties of a training dataset of token sequences are neural
scaling laws sensitive to? In other words, what can we measure about our data
to more accurately predict the optimal compute allocation for training?
Furthermore, is data-dependence of scaling laws only of theoretical interest,
or can laws be considerably different for real-world datasets?

In order to study these questions, we seek a textual data setting where we can
intuitively control its complexity as well as open avenues for information-
theoretic understandings of why scaling laws are data-dependent. We therefore
settle on probabilistic context-free grammars (PCFG) [Chomsky, 1956] which are
relatively naturalistic (can model natural language, code, etc.), controllable
in syntactic complexity, and follow some well-understood information-theoretic
principles [Chi, 1999].

We generate 6 datasets of varying complexity by modulating syntactic
properties of a PCFG. For each of these datasets, we train LMs of 6 different
sizes (4.4M to 1.4B parameters) and record results at 6 different train step
counts (100K to 100M tokens). We then fit a scaling law for each dataset
[Muennighoff et al., 2023], finding meaningful shifts in the law’s parameters
as syntactic complexity increases. Following prior work on entropy of formal
grammars [Arora et al., 2022], we use median compressibility of each token
sequence in a dataset as a measure of complexity that is straightforward to
compute with gzip [Gailly and Adler, 1992].

We find that as the training data becomes less compressible (more complex),
the scaling law’s compute-optimal frontier gradually increases its preference
for dataset size over parameter count. We then measure the compressibility of
real-world code & natural language datasets, showing that the former is
considerably more compressible and thus subject to a predictably different
scaling law. With some napkin math, we estimate that we could reach the
performance of StarCoder [Li et al., 2023] with 24% fewer FLOPs ($278,000 in
H100 hours) since our compute-optimal scaling adjusts for data complexity.

##  2 Related Work

###  2.1 Scaling Laws

Early work established that a neural network’s test error is a power law of
training dataset size [Cortes et al., 1993], model parameter count [Rosenfeld
et al., 2019], and that the relationship holds over many orders of magnitude
[Hestness et al., 2017]. Kaplan et al. [2020] applied scaling laws to
transformer-based language models and identified a compute-optimal frontier
along which parameter & dataset size should be scaled. Hoffmann et al. [2022]
propose the Chinchilla scaling laws, finding that Kaplan et al. [2020] and Rae
et al. [2021] overparameterized their models and that the compute-optimal
frontier requires parameter & dataset size to be scaled equally (rather than
parameters scaling at 3x the rate of data).

Sorscher et al. [2022] find that we can reach exponential (rather than power
law) scaling on dataset size by pruning out redundant examples that do not
provide much information to learn from. Liu and Tegmark [2023] identify a
mechanistic explanation of why scaling follows a power law on model width.

Aghajanyan et al. [2023] extend the Chinchilla scaling laws to a variety of
modalities such as speech, image-text, and code. Caballero et al. [2023]
propose a novel functional form for scaling laws that better models complex
non-monotonic behavior and also apply it to several modalities including code.
Both these works and Hoffmann et al. [2022] (in Appendix C) find that scaling
behavior for code is different than for natural language—code is easier to
learn and its compute-optimal frontier prefers parameters slightly over data.
Bi et al. [2024] cursorily investigate scaling laws across datasets of
different qualities, finding that cleaner & higher quality data results in the
“model scaling exponent increasing” because “high-quality data usually implies
logical clarity and less predictive difficulty after sufficient training”.
However, none of these works identify an underlying general principle that
explains why scaling behavior varies across data modalities & complexities (or
even just between code and natural language).

Meanwhile, Jones [2021] explores scaling laws in the context of board games,
finding that “the compute required for a desired level of performance can be
calculated directly from the board size”. If scaling laws smoothly scale with
data complexity in the board game setting, do they also smoothly scale with
textual data complexity? And how might we measure a text dataset’s complexity?

###  2.2 Syntax & Information Theory

A long history of work in linguistics has sought to apply information-
theoretic measures to natural language [Shannon, 1951, Cherry et al., 1953,
Harris, 1991, Piantadosi et al., 2011]. Specifically, entropy (a fairly
abstract information-theoretic measure) can be operationalized in a number of
concrete ways to measure the informational complexity of a linguistic
distribution [Arora et al., 2022]. One way to determine the ‘goodness’ of an
entropy measure over sequences is whether it grows as the syntactic complexity
of a language increases (e.g. more production rules, more non-terminals,
longer right-hand sides). While entropy is straightforward to compute in
closed-form on simple formal languages (i.e. Type 3 Chomskyan grammars,
generated by Finite State Automata) [Grenander, 1967, Sánchez et al., 2018],
for Context-free (Type 2) grammars, we must estimate entropy from a set of
generated samples [Chi, 1999, Corazza and Satta, 2007].

Leveraging the widely-recognized relationship between entropy and compression
[Shannon, 1948, Huffman, 1952, Ziv and Lempel, 1977], we can use gzip [Gailly
and Adler, 1992], a lossless compression utility that implements DEFLATE
[Deutsch, 1996] (a combination of Huffman coding & the Lempel-Ziv algorithm)
to estimate the entropy of sampled linguistic sequences.

Delétang et al. [2023] adapt language models as lossless compressors of token
sequences and provide new insights on scaling laws from this compression
perspective, finding that beyond a critical model size, the compression rate
(accounting for parameter count) reverses its improvement. In a somewhat
similar vein, Jiang et al. [2023] use gzip with k-Nearest-Neighbors to beat
neural text embeddings. Importantly, they both brought to attention the
effectiveness of compression algorithms for heuristically measuring the
structure learned by language models [Huang et al., 2024]. Another line of
recent work investigates the learnability of different languages by LMs,
finding that natural languages vary in their learnability [Cotterell et al.,
2018] and that LMs struggle to model syntactically ‘impossible’ languages
[Kallini et al., 2024].

##  3 Modulating Data Complexity via Syntactic Properties of a PCFG

Probabilistic Context-Free Grammars (PCFGs) are a fundamental tool in
computational linguistics for modeling the syntax of natural languages. A PCFG
extends the concept of a standard Context-Free Grammar (CFG) by associating
probabilities with its production rules, enabling the representation of
language ambiguity and variability in a quantifiable manner. These grammars
generate trees where each node represents a syntactic category, and the edges
represent production rules applied to generate sentences. When generating
sentences from a PCFG, we probabilistically sample sequences of production
rules to apply until all leaves of the tree are terminals (actual vocab
tokens).

We can control the syntactic properties of a PCFG to naturalistically modulate
the complexity of a textual dataset that we sample from its generated
sentences. Specifically, our PCFG creation function accepts arguments for the
number of terminals, number of non-terminals, the maximum length of the right-
hand side of a production rule, and the maximum number of production rules
allowed for any non-terminal (i.e. if this is 1, a given non-terminal will
always lead to the same right-hand side). Intuitively, as each of these values
increase, the language’s syntactic complexity also increases.

To create a PCFG from the above arguments, for each non-terminal, we randomly
choose its number of productions (RHS options), each of those production’s
length, instantiate a production rule by randomly sampling from the terminals
& non-terminals, and assign it a probability normalized by the total RHS
options for the non-terminal. We then collect all the generated production
rules for all the non-terminals and instantiate a grammar using the PCFG
package [Breydo, 2021] built on NLTK [Bird and Loper, 2004].

We can then use the grammar that we have randomly (within given constraints)
created to probabilistically sample sentences to construct a dataset of token
sequences. To make it easier later to compare training across grammars that
produce sentences of different average lengths, we decide to sample sentences
into documents of the same token count (our LM’s context length). We sample
sentences from our grammar until we’ve filled up the context length, and if we
overflow, we simply truncate the sequence.

Sentences are composed of terminals that are just integers and can thus be
treated as token IDs for our LM; we concatenate sentences with the unused
integer 0, effectively corresponding to a period in natural language. To
clarify, we do not generate strings that ‘look’ like natural language and must
be tokenized—the PCFG generates sequences of token IDs themselves. Now, we can
generate 6 datasets of token sequences with varying complexities from just 6
sets of initial grammatical constraints.

###  3.1 gzip-compressibility Measures Syntactic Complexity

Now we must choose a metric to estimate the complexity of our datasets. While
for our PCFG-generated data, we could use a trivial function of its syntactic
properties (e.g. number of non-terminals plus median production rule length),
we would only be able to apply this metric to datasets where we precisely know
the grammar that generated the data. Instead, we need a measure of complexity
that we can apply to real-world token sequence datasets so that we can see
whether our data-sensitive scaling results on PCFGs generalize to real-world
datasets. While one might try to use grammar induction to identify the
underlying syntax of a real-world dataset and then compute complexity from its
syntactic properties, this is computationally intensive and difficult with
noisy web-text data.

Therefore, as explained in Sec. 2.2, we choose to use a compression algorithm,
gzip, to estimate the complexity of a dataset by virtue of considerable
theoretical work establishing $\text{compressibility}^{-1}\propto$ entropy
[Shannon, 1948] and entropy $\propto$ syntactic complexity [Chi, 1999].
Specifically, for each token sequence in a sample of 1000 from the dataset, we
apply gzip and compute the ratio of the size (in bytes) of the compressed data
to the original data. We then compute the median and standard deviation in
compressibility, confirming that grammars with higher syntactic complexity
result in datasets that are more difficult to compress.

Figure 1: Plot of different linguistic distributions and their compressibility
by gzip on their raw string & their token sequence. The large grey and orange
clusters are real-world code and natural language, respectively. The smaller
clusters are sampled from PCFG’s of increasingly complex syntax (with specific
values for properties provided in legend). Observe that 1) natural language is
more difficult to compress than code, and 2) as our PCFGs’ syntactic
complexity increases, the token sequences become less compressible. PCFG data
is more tightly clustered due to length normalization.

In Tab. 1, we list the syntactic parameters for each grammar and the
compression ratio we measured of token sequences sampled from the grammar.
Observe that as non-terminals (grammatical categories), terminals (tokens),
right-hand side options, and right-hand side length increase, the gzip-
compressibility also increases (i.e. it becomes harder to compress). We plot
these datasets alongside natural language & code in Fig. 1, showing how some
PCFG datasets are more similar in complexity to code (the ones that are easier
to compress) while others are more similar to natural language.

Table 1: Compressibility of each dataset and the syntactic properties of its PCFG gzip Compressibility | Non-terminals | Terminals | RHS Options | RHS Length
---|---|---|---|---
0.11 | 3 | 20 | 2 | 2
0.22 | 10 | 150 | 5 | 3
0.35 | 20 | 300 | 10 | 5
0.42 | 30 | 400 | 10 | 8
0.51 | 50 | 500 | 20 | 15
0.61 | 100 | 2000 | 100 | 30

##  4 Are Scaling Laws Sensitive to Data Complexity?

To identify the scaling law for a dataset, we train a set of models of varying
sizes (4.2M, 8.8M, 20.3M, 59.0M, 275.3M, 1.4B parameters; architectural
specifics in Tab. 6) on varying size subsets of the data (100K, 1M, 5M, 20M,
50M, 100M tokens) and fit a power law on the resulting final losses of all
training runs. Most experiments were run on a cluster of 4 Nvidia A100’s with
80 GB VRAM each using PyTorch FSDP [Paszke et al., 2019, Zhao et al., 2023].
We use a batch size of 32, single epoch training, the $AdamW$ optimizer
[Loshchilov and Hutter, 2017], and a learning rate starting from
$5\mathrm{e}{-5}$ and cosine decayed over the number of train steps
[Loshchilov and Hutter, 2016].

As intuitively expected, the more compressible a dataset is (lower
compressibility ratio), the faster that models regardless of size will
converge (Fig. 2). While this shows that we need more compute to model more
complex datasets (which is in and of itself notable), we need more evidence to
determine if the compute-optimal frontier directionally shifts based on data
complexity. To establish such a non-trivial sensitivity of scaling laws to
data complexity, we need to compute the law for each dataset and examine its
fitted parameters.

(a) 4.2M

(b) 8.8M

(c) 20.3M

(d) 59.0M

(e) 275.3M

(f) 1.4B

Figure 2: Training curves for models of various sizes trained on 100M token
datasets with gzip-compressibilities of 0.11, 0.22, 0.35, 0.42, 0.51, and
0.61. The harder the data is to compress, the longer it takes to train to
convergence.

###  4.1 Computing Data-sensitive Scaling Laws from gzip-compressibility

The scaling law functional form proposed by Hoffmann et al. [2022] predicts
training loss as a function of model & dataset size:

$L(N,D)=E+\dfrac{A}{N^{\alpha}}+\dfrac{B}{D^{\beta}}$ (1)

where $N$ is the model’s parameter count and $D$ is the training dataset’s
token count. They claim that $E$ captures the “entropy of natural text” (their
Sec. 3.3) and that scaling laws are “independent of dataset” (their App. C).
However, when we fit this function on the training results for each of our
PCFG datasets (leveraging the helpful open-source implementation of
Muennighoff et al. [2023]), we find considerably different laws for each
dataset (Tab. 2).

Table 2: Fitted scaling law parameters for each dataset with gzip compressibility gzip Compressibility | A | B | E | $\alpha$ | $\beta$
---|---|---|---|---|---
0.12 | 18.205 | 15.803 | -0.734 | 1.053 | 1.307
0.23 | 15.553 | 13.912 | -0.262 | 0.881 | 1.064
0.32 | 16.236 | 8.829 | -1.446 | 0.930 | 0.597
0.45 | 15.506 | 9.002 | -1693.171 | 0.882 | 0.596
0.60 | 9.061 | 3.525 | 1.174 | 0.559 | 0.140

The scaling law induces a compute-optimal frontier for parameter count derived
by Kaplan et al. [2020] & Hoffmann et al. [2022] and simplifiable as:

$N_{opt}(C)=\left(\dfrac{\alpha A}{\beta
B}\left(\dfrac{C}{6}\right)^{\beta}\right)^{\dfrac{1}{\alpha+\beta}}$ (2)

where $C$ is a compute budget in FLOPs. We plot this compute-optimal frontier
for Chinchilla as well as our fitted law for each PCFG dataset in Fig. 3. From
a cursory glance, we can see that the frontier of our fitted law progressively
becomes more data-preferent as data becomes harder to compress, crossing over
Chinchilla’s 1-to-1 frontier at some point
$0.23<\texttt{gzip}\text{-compressibility}<0.45$.

(a) gzip-compressibility = 0.12

(b) gzip-compressibility = 0.23

(c) gzip-compressibility = 0.45

(d) gzip-compressibility = 0.60

Figure 3: The stationary grey dashed line is the compute-optimal frontier (Eq.
2) of the allegedly data-agnostic Chinchilla scaling law, and the black line
is the same frontier computed from our fitted scaling law for each PCFG
dataset. Points are individual training runs colored by their loss. As the
data becomes more difficult to compress, the scaling law shifts from being
more parameter-preferent (a) than Chinchilla to far more data-preferent (d).
Compare Appendix Fig. 7 for parameter-data plot.

In order to predict scaling law parameters from a dataset’s compressibility,
let us fit a simple linear regression on the fitted scaling law parameters for
each dataset. Recall from Sec. 3.1 that we compute the compressibility $H$ of
a dataset $\boldsymbol{D}$ by taking the average ratio of compressed bits to
original bits for each element $d$:

$H(\boldsymbol{D})=\dfrac{1}{||\boldsymbol{D}||}\sum_{d\in\boldsymbol{D}}\dfrac{||\texttt{gzip}(d)||}{||d||}$
(3)

So once we’ve fit lines to predict each parameter (E, A, B, $\alpha$, $\beta$)
from $H$, we can re-define each parameter as a function of compressibility:

$\forall x\in\\{E,A,B,\alpha,\beta\\}:x(H)=m_{x}H+n_{x}$ (4)

where $m_{x}$ and $n_{x}$ are just our fitted linear regressions’ parameters.
These fitted values (along with the regressions’ $p$-values) are presented in
Tab. 3 and the linear regressions are visualized in Fig. 4. They’re all pretty
much monotonically decreasing at different rates with an interesting
$\alpha$-$\beta$ intercept at $H\approx 0.27$. It’s also interesting to note
that $E$, originally set as the constant ‘entropy of natural language’, is the
only parameter that loosely increases with $H$ (though not significantly)—our
proxy for complexity or entropy.

Table 3: Fitted values for $x(H)$ that make the scaling law parameters ($x\in\\{E,A,B,\alpha,\beta\\}$) a linear function of data compressibility ($H$). $x$ | $m_{x}$ | $n_{x}$ | $p_{x}$
---|---|---|---
E | 3.92 | -1.56 | .272
A | -16.20 | 20.48 | .048
B | -24.77 | 18.73 | .009
$\alpha$ | -0.87 | 1.16 | .043
$\beta$ | -2.34 | 1.55 | .008

Figure 4: As data becomes harder to compress, the scaling law’s fitted
parameters trend downwards at different rates ($p<0.05$), with $\alpha$ and
$\beta$ intersecting around 0.27—the compressibility point at which parameters
& data should be scaled roughly equiproportionally.

Now, we can reparameterize Eq. 1 as a function of compressibility $H$:

$L(N,D,H)=E(H)+\dfrac{A(H)}{N^{\alpha(H)}}+\dfrac{B(H)}{D^{\beta(H)}}$ (5)

However, since our experiments are considerably smaller scale and primarily on
PCFG data, we also present a form of our data-dependent scaling law as an
adjustment of Chinchilla, where $\varepsilon$ is the weightage of our
adjustment for the gzip-compressibility of the training data (and prime (′)
parameters are the Chinchilla constants).

$L(N,D,H)=(1-\varepsilon)E^{\prime}+\varepsilon
E(H)+\dfrac{(1-\varepsilon)A^{\prime}+\varepsilon
A(H)}{N^{(1-\varepsilon)\alpha^{\prime}+\varepsilon\alpha(H)}}+\dfrac{(1-\varepsilon)B^{\prime}+\varepsilon
B(H)}{D^{(1-\varepsilon)\beta^{\prime}+\varepsilon\beta(H)}}$ (6)

###  4.2 Eliminating Syntactic Parameters as a Confounder of Compressibility

Given just the above experiments, we have not addressed the possibility that
our compressibility measure is confounded by some underlying syntactic
property (e.g. vocab size). To address this concern, in Fig. 5 we provide
results showing that when holding vocab size steady and changing other
syntactic properties (Tab. 4), gzip-compressibility still predicts scaling law
parameter shifts (with an even stronger correlation than in the increasing
vocab size setting).

Table 4: Syntactic parameters for PCFG datasets holding terminal count the same Compressibility | Nonterminals | Terminals | RHS Options | RHS Length
---|---|---|---|---
0.11 | 3 | 300 | 2 | 2
0.25 | 10 | 300 | 5 | 3
0.36 | 20 | 300 | 10 | 5
0.47 | 50 | 300 | 20 | 10

Figure 5: When holding terminal count (vocab size) the same and varying other
syntactic properties, scaling laws still shift across gzip-compressibility,
with even stronger correlation.

We also empirically show the contrapositive in Fig. 6, demonstrating that when
we widely vary syntactic properties (Tab. 5) but such that the datasets’ final
gzip-compressibility are all the same, there is no significant shift in
scaling law parameters.

Table 5: Syntactic parameters for holding gzip-compressibility the same gzip-Compressibility | Nonterminals | Terminals | RHS Options | RHS Length
---|---|---|---|---
0.32 | 10 | 500 | 5 | 10
0.36 | 20 | 300 | 10 | 5
0.40 | 30 | 200 | 15 | 20

Figure 6: When non-monotonically varying syntactic parameters but maintaining
roughly the same gzip-compressibility (0.32 - 0.40), scaling law parameters
(especially $\alpha$ and $A$) do not meaningfully shift.

Although the intersection behavior observed in Fig. 4 is not observed in the
iso-vocab case (Fig. 5), the steeper negative slope of $\alpha$ v.s. $\beta$
(and $A$ v.s. $B$) implies the same phenomenon of increased data-preference as
gzip-compressibility grows.

Thus we have shown that scaling laws are dependent on training data and gzip-
compressibility is a strong predictor of how data complexity will impact
scaling properties. We provide fitted coefficients (Tab. 3) for predicting
scaling parameters as a function of $H$ (gzip-compressibility), and offer an
adjustment-based formulation (Eq. 6) of the data-dependent scaling law to make
our small scale synthetic experiments more adaptable to training frontier
models.

##  5 Discussion

We have shown that the Chinchilla scaling law [Hoffmann et al., 2022] and
similar laws that claim to be data-agnostic are in fact web-text-specific
instances of a broader family of scaling laws sensitive to information-
theoretic measures of training data. Not only do datasets that are difficult
to compress (regardless of syntactic specifics) require more compute, they
also require a different trade-off between model size & data.

Why does a dataset with a certain level of compressibility $H$ result in the
specific scaling parameters that it does? And why does increased gzip-
compressibility result in increased preference for data over parameter count?
These questions warrant further exploration from a theoretical perspective—we
hope to leverage connections between syntax & information theory to answer
this question in future work.

On the empirical side, an especially promising avenue is in LM’s for code
since code datasets have significantly lower gzip-compressibility than natural
language. Therefore, we can expect the compute-optimal scaling law for code to
have stronger parameter-preference than Chinchilla. We are currently running
experiments at 1.92e19 FLOPs (some gzip-compressibility analysis in Fig. 8 and
preliminary runs in Fig. 9) to compare performance based on Chinchilla’s
allocation and our data-adjusted scaling law’s (Eq. 6) allocation. As
highlighted in Sec. 1, scaling up this finding could save $278,000 in H100
hours when training a single relatively small (6B parameter) code-generation
model.

Beyond a deeper theoretical understanding of scaling laws & more optimal
compute allocation for code-generation LMs, our findings also open
possibilities for using compressibility as a metric in training data
filtering, curriculum learning ordering, and fine-tuning data requirement
estimation. Anecdotally, we have heard from researchers at frontier AI labs
that compression algorithms are already used in production for filtering LLM
training data.

The call-to-action here is not per se to use our data-dependent scaling law to
decide how to spend your next million dollars of compute, but rather to begin
investigating whether performance discrepancies between models trained on
different datasets are explainable by simple information-theoretic measures
(e.g. gzip-compressibility) of training data. If so (as experienced by Bi et
al. [2024]), it may be worth building systems to dynamically determine compute
allocation for each training dataset, with the general rule of thumb that
easy-to-compress datasets prefer parameters and hard-to-compress datasets
prefer data.

##  6 Limitations

The work presented here is restricted to the synthetic PCFG setting and does
not yet explicitly show generalization to predicting different scaling
properties of real-world datasets. Our method of modulating training data
complexity by varying syntactic properties of a PCFG, though naturalistic, is
fairly specific and may still be prone to confounders. In addition, there are
features of training data that affect scaling but appear to not be modeled
well by gzip-compressibility. Alternative compression algorithms &
information-theoretic measures of training data were not tested though they
likely all broadly correlate with gzip-compressibility.

##  7 Acknowledgements

We are grateful to Alon Albalak, Rylan Schaeffer, Charlie Snell, Ethan
Caballero, Martin Ziqiao Ma, Sankeerth Rao Karingula, Guillaume Lample, Vibhu
Sapra, Aryaman Arora, Tristan Thrush, and several anonymous Twitter users for
their thorough reviews & feedback on drafts of this paper.

This work was also sponsored by generous compute grants from Luke Piette
(RunPod), Kye Gomez (Swarms), Vincent Weisser & Johannes Hagemann (Prime
Intellect), and of course the rest of my team at Reworkd—Asim Shrestha & Adam
Watkins.

## References

  * Aghajanyan et al. [2023] Armen Aghajanyan, Lili Yu, Alexis Conneau, Wei-Ning Hsu, Karen Hambardzumyan, Susan Zhang, Stephen Roller, Naman Goyal, Omer Levy, and Luke Zettlemoyer.  Scaling laws for generative mixed-modal language models.  In _International Conference on Machine Learning_ , pages 265–279. PMLR, 2023.
  * Arora et al. [2022] Aryaman Arora, Clara Meister, and Ryan Cotterell.  Estimating the entropy of linguistic distributions.  In _Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 2: Short Papers)_ , pages 175–195, 2022.
  * Bi et al. [2024] Xiao Bi, Deli Chen, Guanting Chen, Shanhuang Chen, Damai Dai, Chengqi Deng, Honghui Ding, Kai Dong, Qiushi Du, Zhe Fu, et al.  Deepseek llm: Scaling open-source language models with longtermism.  _arXiv preprint arXiv:2401.02954_ , 2024.
  * Bird and Loper [2004] Steven Bird and Edward Loper.  NLTK: The natural language toolkit.  In _Proceedings of the ACL Interactive Poster and Demonstration Sessions_ , pages 214–217, Barcelona, Spain, July 2004. Association for Computational Linguistics.  URL https://aclanthology.org/P04-3031.
  * Breydo [2021] Thomas Breydo.  _thomasbreydo/pcfg_.  8 2021.  URL https://github.com/thomasbreydo/pcfg.
  * Caballero et al. [2023] Ethan Caballero, Kshitij Gupta, Irina Rish, and David Krueger.  Broken neural scaling laws.  In _ICLR 2023 Workshop on Mathematical and Empirical Understanding of Foundation Models_ , 2023.
  * Cherry et al. [1953] E Colin Cherry, Morris Halle, and Roman Jakobson.  Toward the logical description of languages in their phonemic aspect.  _Language_ , pages 34–46, 1953.
  * Chi [1999] Zhiyi Chi.  Statistical properties of probabilistic context-free grammars.  _Computational Linguistics_ , 25(1):131–160, 1999.  URL https://aclanthology.org/J99-1004.
  * Chomsky [1956] Noam Chomsky.  Three models for the description of language.  _IRE Transactions on information theory_ , 2(3):113–124, 1956.
  * Corazza and Satta [2007] Anna Corazza and Giorgio Satta.  Probabilistic context-free grammars estimated from infinite distributions.  _IEEE transactions on pattern analysis and machine intelligence_ , 29(8):1379–1393, 2007.
  * Cortes et al. [1993] Corinna Cortes, Lawrence D Jackel, Sara Solla, Vladimir Vapnik, and John Denker.  Learning curves: Asymptotic values and rate of convergence.  _Advances in neural information processing systems_ , 6, 1993.
  * Cotterell et al. [2018] Ryan Cotterell, Sabrina J. Mielke, Jason Eisner, and Brian Roark.  Are all languages equally hard to language-model?  In Marilyn Walker, Heng Ji, and Amanda Stent, editors, _Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 2 (Short Papers)_ , pages 536–541, New Orleans, Louisiana, June 2018. Association for Computational Linguistics.  doi: 10.18653/v1/N18-2085.  URL https://aclanthology.org/N18-2085.
  * Delétang et al. [2023] Grégoire Delétang, Anian Ruoss, Paul-Ambroise Duquenne, Elliot Catt, Tim Genewein, Christopher Mattern, Jordi Grau-Moya, Li Kevin Wenliang, Matthew Aitchison, Laurent Orseau, et al.  Language modeling is compression.  _arXiv preprint arXiv:2309.10668_ , 2023.
  * Deutsch [1996] Peter Deutsch.  Deflate compressed data format specification version 1.3.  Technical report, 1996.
  * Gailly and Adler [1992] Jean-loup Gailly and Mark Adler.  Gnu gzip.  _GNU Operating System_ , 1992.
  * Grenander [1967] Ulf Grenander.  _Syntax-controlled probabilities_.  Division of Applied Mathematics, Brown University, 1967.
  * Gunasekar et al. [2023] Suriya Gunasekar, Yi Zhang, Jyoti Aneja, Caio César Teodoro Mendes, Allie Del Giorno, Sivakanth Gopi, Mojan Javaheripi, Piero Kauffmann, Gustavo de Rosa, Olli Saarikivi, et al.  Textbooks are all you need.  _arXiv preprint arXiv:2306.11644_ , 2023.
  * Harris [1991] Zellig Harris.  _A theory of language and information: a mathematical approach_.  Oxford University Press, 1991.
  * Hestness et al. [2017] Joel Hestness, Sharan Narang, Newsha Ardalani, Gregory Diamos, Heewoo Jun, Hassan Kianinejad, Md Mostofa Ali Patwary, Yang Yang, and Yanqi Zhou.  Deep learning scaling is predictable, empirically.  _arXiv preprint arXiv:1712.00409_ , 2017.
  * Hoffmann et al. [2022] Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al.  Training compute-optimal large language models.  _arXiv preprint arXiv:2203.15556_ , 2022.
  * Huang et al. [2024] Yuzhen Huang, Jinghan Zhang, Zifei Shan, and Junxian He.  Compression represents intelligence linearly.  _arXiv preprint arXiv:2404.09937_ , 2024.
  * Huffman [1952] David A Huffman.  A method for the construction of minimum-redundancy codes.  _Proceedings of the IRE_ , 40(9):1098–1101, 1952.
  * Jiang et al. [2023] Zhiying Jiang, Matthew Yang, Mikhail Tsirlin, Raphael Tang, Yiqin Dai, and Jimmy Lin.  “low-resource” text classification: A parameter-free classification method with compressors.  In Anna Rogers, Jordan Boyd-Graber, and Naoaki Okazaki, editors, _Findings of the Association for Computational Linguistics: ACL 2023_ , pages 6810–6828, Toronto, Canada, July 2023. Association for Computational Linguistics.  doi: 10.18653/v1/2023.findings-acl.426.  URL https://aclanthology.org/2023.findings-acl.426.
  * Jones [2021] Andy L Jones.  Scaling scaling laws with board games.  _arXiv preprint arXiv:2104.03113_ , 2021.
  * Kallini et al. [2024] Julie Kallini, Isabel Papadimitriou, Richard Futrell, Kyle Mahowald, and Christopher Potts.  Mission: Impossible language models.  _arXiv preprint arXiv:2401.06416_ , 2024.
  * 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.
  * Li et al. [2023] Raymond Li, Loubna Ben Allal, Yangtian Zi, Niklas Muennighoff, Denis Kocetkov, Chenghao Mou, Marc Marone, Christopher Akiki, Jia Li, Jenny Chim, et al.  Starcoder: may the source be with you!  _arXiv preprint arXiv:2305.06161_ , 2023.
  * Liu and Tegmark [2023] Ziming Liu and Max Tegmark.  A neural scaling law from lottery ticket ensembling.  _arXiv preprint arXiv:2310.02258_ , 2023.
  * Loshchilov and Hutter [2016] Ilya Loshchilov and Frank Hutter.  Sgdr: Stochastic gradient descent with warm restarts.  _arXiv preprint arXiv:1608.03983_ , 2016.
  * Loshchilov and Hutter [2017] Ilya Loshchilov and Frank Hutter.  Decoupled weight decay regularization.  _arXiv preprint arXiv:1711.05101_ , 2017.
  * Muennighoff et al. [2023] Niklas Muennighoff, Alexander Rush, Boaz Barak, Teven Le Scao, Nouamane Tazi, Aleksandra Piktus, Sampo Pyysalo, Thomas Wolf, and Colin A Raffel.  Scaling data-constrained language models.  _Advances in Neural Information Processing Systems_ , 36, 2023.
  * Paszke et al. [2019] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al.  Pytorch: An imperative style, high-performance deep learning library.  _Advances in neural information processing systems_ , 32, 2019.
  * 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.
  * Penedo et al. [2024] Guilherme Penedo, Hynek Kydlíček, Leandro von Werra, and Thomas Wolf.  Fineweb, 2024.  URL https://huggingface.co/datasets/HuggingFaceFW/fineweb.
  * Piantadosi et al. [2011] Steven T Piantadosi, Harry Tily, and Edward Gibson.  Word lengths are optimized for efficient communication.  _Proceedings of the National Academy of Sciences_ , 108(9):3526–3529, 2011.
  * Rae et al. [2021] Jack W Rae, Sebastian Borgeaud, Trevor Cai, Katie Millican, Jordan Hoffmann, Francis Song, John Aslanides, Sarah Henderson, Roman Ring, Susannah Young, et al.  Scaling language models: Methods, analysis & insights from training gopher.  _arXiv preprint arXiv:2112.11446_ , 2021.
  * Rosenfeld et al. [2019] Jonathan S Rosenfeld, Amir Rosenfeld, Yonatan Belinkov, and Nir Shavit.  A constructive prediction of the generalization error across scales.  _arXiv preprint arXiv:1909.12673_ , 2019.
  * Sánchez et al. [2018] Joan Andreu Sánchez, Martha Alicia Rocha, Verónica Romero, and Mauricio Villegas.  On the derivational entropy of left-to-right probabilistic finite-state automata and hidden Markov models.  _Computational Linguistics_ , 44(1):17–37, April 2018.  doi: 10.1162/COLI_a_00306.  URL https://aclanthology.org/J18-1002.
  * Shannon [1951] Claude E Shannon.  Prediction and entropy of printed english.  _Bell system technical journal_ , 30(1):50–64, 1951.
  * Shannon [1948] Claude Elwood Shannon.  A mathematical theory of communication.  _The Bell system technical journal_ , 27(3):379–423, 1948.
  * Sorscher et al. [2022] Ben Sorscher, Robert Geirhos, Shashank Shekhar, Surya Ganguli, and Ari Morcos.  Beyond neural scaling laws: beating power law scaling via data pruning.  _Advances in Neural Information Processing Systems_ , 35:19523–19536, 2022.
  * Xie et al. [2024] Sang Michael Xie, Hieu Pham, Xuanyi Dong, Nan Du, Hanxiao Liu, Yifeng Lu, Percy S Liang, Quoc V Le, Tengyu Ma, and Adams Wei Yu.  Doremi: Optimizing data mixtures speeds up language model pretraining.  _Advances in Neural Information Processing Systems_ , 36, 2024.
  * Zhao et al. [2023] Yanli Zhao, Andrew Gu, Rohan Varma, Liang Luo, Chien-Chin Huang, Min Xu, Less Wright, Hamid Shojanazeri, Myle Ott, Sam Shleifer, et al.  Pytorch fsdp: experiences on scaling fully sharded data parallel.  _arXiv preprint arXiv:2304.11277_ , 2023.
  * Ziv and Lempel [1977] Jacob Ziv and Abraham Lempel.  A universal algorithm for sequential data compression.  _IEEE Transactions on information theory_ , 23(3):337–343, 1977.

##  Appendix A Appendix

Table 6: Architectural properties for each model size Parameter Count | Hidden Size | Intermediate Size | Hidden Layers | Attention Heads
---|---|---|---|---
4.2M | 64 | 128 | 2 | 1
8.8M | 128 | 256 | 4 | 2
20.3M | 256 | 512 | 6 | 4
59.0M | 512 | 1024 | 10 | 8
275.3M | 1024 | 2048 | 20 | 16
1.4B | 2048 | 4096 | 30 | 32

Figure 7: Scaling contours in parameter-data space, showing how for dataset
with higher gzip-compressibility (0.60), not only is more compute required,
but the optimal frontier increasingly prefers data. Figure 8: Real-world code
datasets of different languages have slightly different compressibilities, but
all considerably less than natural language.

Figure 9: In real-world code datasets with even slightly different
compressibilities (HTML, C, & Python), slightly more compute is required
(regardless of model size) as data becomes harder to compress. Further work is
required to compare their scaling laws in relation to natural language.