capbtabboxtable[][]

Present address: ]Institute for Research in Electronics and Applied Physics,
Univ. of Maryland, College Park, 20899-6202, MD, USA Present address: ]GFMC,
Departamento de Física de Materiales, Universidad Complutense de Madrid, 28040
Madrid, Spain

# Flicker and random telegraph noise between gyrotropic and dynamic C-state of
a vortex based spin torque nano oscillator

Steffen Wittrock<EMAIL_ADDRESS>Unité Mixte de Physique
CNRS, Thales, Univ. Paris-Saclay, 1 Avenue Augustin Fresnel, 91767 Palaiseau,
France  Philippe Talatchian  [  Unité Mixte de Physique CNRS, Thales, Univ.
Paris-Saclay, 1 Avenue Augustin Fresnel, 91767 Palaiseau, France  Miguel
Romera  [  Unité Mixte de Physique CNRS, Thales, Univ. Paris-Saclay, 1 Avenue
Augustin Fresnel, 91767 Palaiseau, France  Mafalda Jotta Garcia  Unité Mixte
de Physique CNRS, Thales, Univ. Paris-Saclay, 1 Avenue Augustin Fresnel, 91767
Palaiseau, France  Marie-Claire Cyrille  Univ. Grenoble Alpes, CEA-LETI,
MINATEC-Campus, 38000 Grenoble, France  Ricardo Ferreira  International
Iberian Nanotechnology Laboratory (INL), 471531 Braga, Portugal  Romain Lebrun
Unité Mixte de Physique CNRS, Thales, Univ. Paris-Saclay, 1 Avenue Augustin
Fresnel, 91767 Palaiseau, France  Paolo Bortolotti  Unité Mixte de Physique
CNRS, Thales, Univ. Paris-Saclay, 1 Avenue Augustin Fresnel, 91767 Palaiseau,
France  Ursula Ebels  Univ. Grenoble Alpes, CEA, INAC-SPINTEC, CNRS, SPINTEC,
38000 Grenoble, France  Julie Grollier  Unité Mixte de Physique CNRS, Thales,
Univ. Paris-Saclay, 1 Avenue Augustin Fresnel, 91767 Palaiseau, France
Vincent Cros<EMAIL_ADDRESS>Unité Mixte de Physique CNRS,
Thales, Univ. Paris-Saclay, 1 Avenue Augustin Fresnel, 91767 Palaiseau, France

(August 27, 2024)

###### Abstract

Vortex based spin torque nano oscillators (STVOs) can present more complex
dynamics than the spin torque induced gyrotropic (G) motion of the vortex
core. The respective dynamic modes and the transition between them can be
controlled by experimental parameters such as the applied dc current. An
interesting behavior is the stochastic transition from the G- to a dynamic
C-state occurring for large current densities. Moreover, the C-state
oscillations exhibit a constant active magnetic volume. We present noise
measurements in the different dynamic states that allow accessing specific
properties of the stochastic transition, such as the characteristic state
transition frequency. Furthermore, we confirm, as theoretically predicted, an
increase of flicker noise with $I_{dc}^{2}$ when the oscillation volume
remains constant with the current. These results bring insight into the
potential optimization of noise properties sought for many potential rf
applications with spin torque oscillators. Furthermore, the investigated
stochastic characteristics open up new potentialities, for instance in the
emerging field of neuromorphic computing schemes.

††preprint: AIP/123-QED

##  I Introduction

Spin torque nano oscillators (STOs) exploit magnetoresistive effects in order
to convert magnetization dynamics into rf electrical signals. Exhibiting
interesting properties, such as a nanometric size, complete CMOS
compatibilityMakarov  _et al._ (2016), radiation hardnessHughes  _et al._
(2012) and high frequency tunabilityArun  _et al._ (2020), they are considered
the next-generation devices for rf applications and communicationsLocatelli,
Cros, and Grollier (2013); Ebels  _et al._ (2017). However, their
potentialities go even beyond this and include exploitation for hard disk
readingSato  _et al._ (2012), or, based on the STOs’ high responsiveness to
external stimuli, broadband frequency detectionLitvinenko  _et al._ (2020) and
microwave energy harvestingFang  _et al._ (2019). Furthermore, along with the
possibility to efficiently synchronize two or more STOs togetherKaka  _et al._
(2005); Lebrun  _et al._ (2017); Tsunegi  _et al._ (2018) in order to leverage
novel neuro-inspired computing architectures for e.g. pattern
recognitionRomera  _et al._ (2018), also more complex phenomena have been
demonstrated in STOs, such as stochastic Jenkins  _et al._ (2019); Wittrock
_et al._ (2021) or even chaoticPetit-Watelot  _et al._ (2012); Devolder  _et
al._ (2019) behavior.

It has recently been shownWittrock  _et al._ (2021) that the transition
between two precession modes in a vortex based STO (STVO) exhibits stochastic
characteristics that can be controlled by the external parameter, i.e. the dc
current or applied magnetic field. The two involved modes are the vortex core
gyrotropic state and beyond, the dynamic C-state motion of the magnetization
distribution with particularly interesting characteristics, such as a nearly
constant magnetic oscillation volumeWittrock  _et al._ (2021).

So far, the main limitation of all the above mentioned possible STO
applications is their large noise characteristics, which is necessary to be
reduced and better understood. Flicker noise, that is dominant at longer
timescales of sought stability, has been studied and classified mainly in the
gyrotropic regimeWittrock  _et al._ (2019, 2020) of a STVO. We recently found
in particular a dependence of the noise level on the active magnetic volume
$V$ of the oscillation, along with the nonlinear dynamical parametersWittrock
_et al._ (2019), and predict a Hooge-like increaseHooge and Hoppenbrouwers
(1969) with $I_{dc}^{2}$ when the volume $V$ is constant. In this general
sense, the described C-state oscillations with constant oscillation volume $V$
represent an interesting case to study the system’s flicker noise level.
Furthermore, along with the characterization of the flicker noise in the
C-state, we here perform noise measurements in the transition between
gyrotropic and C-state. Typical for a two-level stochastic system, we
demonstrate the occurrence of random telegraph noise (RTN) within the
transition, which turns out to be useful in order to better characterize the
stochastic transition.

##  II Experiment

(a)  (b)  (c)

Figure 1:  (a) Measured frequency vs. current evolution of the described G-
and C-state regime with sketched motion of the vortex core inside/outside the
nanodisk characterizing G- and C-state. (b) Phase noise curve in the C-state
at $17\,$mA with fitted $\alpha/f^{3}$ flicker noise contribution. (c)
Corresponding flicker noise level $\alpha$ vs. dc current in the C-state
regime. The applied perpendicular-to-plane magnetic field is
$\mu_{0}H_{\perp}=520\,$mT.

The studied samples are the same as those used in Ref. Wittrock  _et al._
(2021), i.e. circularly shaped nanopillars of $300\,$nm diameter consisting of
a synthetic antiferromagnetic stack (SAF) as a pinning layer. In order to
exploit the tunnel magnetoresistive (TMR) effect, a MgO tunnel barrier
separates the SAF and a $7\,$nm NiFe free layer in the magnetic vortex
configuration. The TMR ratio lies around $50\,$% at room temperature and the
complete layer stack is
SAF/""MgO($1$)/""Co40Fe40B20($1.5$)/""Ta($0.2$)/""Ni80Fe20(7)/""Ta(5)/Ru(7)
with the SAF composed of
PtMn($20$)/""Co70Fe30($2$)/""Ru($0.85$)/""Co40Fe40B20($2.2$)/""Co70Fe30($0.5$)
and the layer thicknesses in brackets. The TMR effect converts the spin torque
driven vortex dynamics into a rf electrical signal which is simultaneously
measured on a spectrum analyzer and an oscilloscope. A perpendicular-to-plane
magnetic field $\mu_{0}H_{\perp}$ tilts the magnetization of the SAF’s top
layer in order to create a perpendicular-to-plane spin polarization providing
the necessary spin torque component to reach self-sustained dynamicsDussaux
_et al._ (2012).

Noise data are gathered from single-shot oscilloscope voltage time traces and
processing via the Hilbert transform methodBianchini  _et al._ (2010); Quinsat
_et al._ (2010); Wittrock  _et al._ (2019). To obtain noise data at low offset
frequencies close to the carrier and furthermore resolve small frequency
variations, a heterodyne detection techniqueKeller  _et al._ (2010); Eklund
_et al._ (2014); Wittrock  _et al._ (2019) (signal down-conversion via high-
side injection and low-pass filtering (bandwidth DC to $22\,$MHz)) is used.

##  III Flicker noise in the C-state

The dynamic C-state describes the motion of the magnetic vortex core in a
circular disk beyond the gyrotropic (G) state. At one point in the periodic
trajectory, the vortex core reaches the disk boundary, disappears and an in-
plane magnetic C-state distribution forms and precesses for a part of the
oscillation period. We have recently shownWittrock  _et al._ (2021) that this
behavior is mainly triggered by the presence of a field-like torque and can be
identified by a rather abrupt decrease of the oscillation frequency at larger
applied dc currents $I_{dc}$. In fig. 1, a characteristic measurement with a
perpendicular-to-plane magnetic field $\mu_{0}H_{\perp}=520\,$mT is shown:
Recognized from the frequency vs. current curve (fig. 1a), the STVO at smaller
dc currents oscillates in the G-state up to a transition from G- to C-state
which stabilizes above $\sim 14\,$mA. In fig. 1a’s inset, a sketch of the
vortex core trajectories is presented.

In fig. 1b, we show a phase noise measurementBianchini  _et al._ (2010);
Quinsat  _et al._ (2010); Wittrock  _et al._ (2019) in the C-state, recorded
at room temperature at $I_{dc}=17\,$mA. At higher frequency offsets, a
$1/f^{2}$ noise characteristics due to thermal fluctuations is observed
Tiberkevich, Slavin, and Kim (2008) while at lower frequency offsets a
$1/f^{3}$ flicker noise process is detectedWittrock  _et al._ (2019).
Performing $\alpha/f^{3}$ flicker noise fits, we present in fig. 1c the
evolution of the flicker noise level $\alpha$ as a function of the applied dc
current $I_{dc}$. A flicker noise minimum is reached at $I_{dc}\approx
15.5\,$mA, where we assume the C-state to be stabilized after the transition.
For larger applied currents, the noise level indeed increases. Performing a
linear fit on the curve as depicted in fig. 1c, we find a curve slope of
$\sim(1.6\pm 0.1)\,$dBc/(Hz$\cdot$mA). The Hooge formulaHooge and
Hoppenbrouwers (1969); Fermon and Pannetier-Lecoeur (2013)
$\alpha\sim\alpha_{H}I_{dc}^{2}/V$ predicts a value of 2 (note the logarithmic
scale) at constant active volumeWittrock  _et al._ (2019). The fact that we do
not find an increase with slope of $2$ but slightly lower might be explained
by the complex current dependence of the Hooge-parameter $\alpha_{H}$, a sort
of quality factor of the systemHooge and Hoppenbrouwers (1969), which usually
decreases for higher applied biasGokce  _et al._ (2006); Almeida, Wisniowski,
and Freitas (2008); Aliev  _et al._ (2007); Wittrock  _et al._ (2019) linked
to the simultaneously decreasing TMR value. Furthermore, the dynamical
parameters of the oscillation, such as the damping rate back to the limit
cycle, might play a roleWittrock  _et al._ (2019); Slavin and Tiberkevich
(2009). In general, in the gyrotropic state the flicker noise level is
decreasing with increasing current $I_{dc}$ Wittrock  _et al._ (2019). Our
measurement confirms that once the oscillation volume saturates, what is the
case also in the G-state at high currents, the flicker noise level evolution
however turns and increases with increasing $I_{dc}$, hence the Hooge formula
description dominates the evolution. Thus, it means that there is a minimum in
the noise level which can be exploited in order to minimize the noise level in
the STO operation.

##  IV Random Telegraph noise in the state transition

(a)  (b)  (c)

Figure 2:  (a) Via spectrum analyzer measured frequency vs. $I_{dc}$ with
stochastic G$\rightarrow$C-state transition. The inset shows a corresponding
time trace of the heterodyne signal at $I_{dc}=11.5\,$mA highlighting the
temporal change of state frequency, here in the interval
$\sim[11.8;12.8]\,\upmu$s. The frequency difference of this interval to the
rest of the shown curve is about $8\,$MHz. (b) Noise PSD with RTN fit on the
phase noise data at $I_{dc}=11.8\,$mA. (c) Parameters of the RTN fits as a
function of the current. The applied perpendicular-to-plane magnetic field is
$\mu_{0}H_{\perp}=425\,$mT.

In fig. 2a, we show the frequency vs. current evolution of a (different)
measurement at $\mu_{0}H_{\perp}=425\,$mT showing a transition from G- to
C-state from $\sim 9\,$mA on. With similar characteristics of G- and C-state
as before, the frequency decrease in comparison seems rather abrupt. However,
traces from both of the two states are detected in the entire range from
$8.9\,$mA up to $12.5\,$mA. At $\sim 9.3\,$mA, inside the transition regime,
the G-state oscillation is even stable at least for the duration of the
measurement ($\sim 10\,$ms). However, within the transition the system is
constantly changing between the two states with frequency differences of
typically a few MHz. This is represented in the heterodyne voltage signal in
fig. 2a’s inset with the local oscillator frequency chosen such that it is
$\sim 2\,$MHz larger than the determined principal frequency shown in fig. 2a.
At $\sim 11.8\,\upmu$s, the oscillator changes from $343\,$MHz into another
state of $8\,$MHz higher frequency for a couple of periods. By performing
phase noise measurements inside the transition (fig. 2b), we observe a noise
PSD which is no more solely described by only thermal and flicker
contributions. Instead, an additional random telegraph noise (RTN) process
with a typical characteristic offset frequency $f_{0}$ of a few kHz adds to
the noise PSD in fig. 2b. The appearance of such RTN contribution to the noise
is found in a large range of $I_{dc}$ from $\sim 8.9\,$mA to $\sim 12.5\,$mA
and corresponds to the bistability of the two states, i.e. a change of
frequency as also presented in fig. 2a’s inset. This type of noise arises from
fluctuations between two statesFermon and Pannetier-Lecoeur (2013). In order
to obtain characteristic parameters of the two level time signal, we perform
fits to the noise data according toAlmeida, Wisniowski, and Freitas (2008)

$\displaystyle
S_{\delta\phi}=10\cdot\log\left(\frac{\alpha}{f^{\beta}}+\frac{S_{0}}{\left(1+(f/f_{0})^{2}\right)f^{2}}\right)\leavevmode\nobreak\
\leavevmode\nobreak\ \leavevmode\nobreak\ ,$

as represented in fig. 2b. Here, $\alpha/f^{\beta}$ with $\beta\approx 3$
reflects the flicker noise at low offset frequencies. The second term
describes the RTN contribution with $S_{0}$ the RTN amplitude and $f_{0}$ the
characteristic RTN frequency, both chosen as fitting parameters.

(a)  (b)

Figure 3:  (a) Wigner-Ville time-frequency spectrogram at $I_{dc}=10.25\,$mA.
(b) Determined C-state probability vs. current through the G- to C-state
transition regime from $\sim 8.9\,$mA to $\sim 12.5\,$mA. Data correspond to
those shown in fig. 2.

The obtained RTN parameters as a function of the applied current are shown in
fig. 2c. The characteristic frequency $f_{0}$ reveals that the stochastic
state switching events occur on a kHz scale. It firstly decreases with
increasing $I_{dc}$ for $I_{dc}\lesssim 10.5\,$mA and subsequently increases
in the range $[10.6;12.5]\,$mA until the C-state oscillation becomes stable
for $I\gtrsim 12.5\,$mA. Note that in the range $[10;10.6]\,$mA, the
characteristic frequency is partly too low for a proper parameter
determination and furthermore, the frequency of the local oscillator might be
between the states’ frequency levels making them difficult to distinguish.
Along with the probability to be in the G- or C-state, also the switching
speed represented by $f_{0}$ can therefore be manipulated by the applied
current, additionally adding a further system handling. In fig. 3, we analyze
the state probability inside the transition more in detail. By performing
Wigner-ville time-frequency spectraCohen (1995) of the heterodyne signal, such
as shown in fig. 3a for $I_{dc}=10.25\,$mA, the frequency and hence the state
probability can be analyzed. In fig. 3a, it is clearly seen that at $\sim
80\,\upmu$s, the oscillation frequency jumps from $\sim 1.7\,$MHz to $\sim
3.5\,$MHz corresponding to a state change from G- to C-state. With the local
oscillator at $339.3\,$MHz, the related real frequencies at $I_{dc}=10.25\,$mA
can consequently be identified as $341\,$MHz for the predominant G-state and
$335.8\,$MHz for the C-state. Note that these frequency differences are not
properly resolved in the frequency determination by a spectrum analyzer as
shown in fig. 2a (typical measurement duration $\sim 10\,$ms). The probability
to be in the C-state can now be analyzed and is shown in fig. 3b through the
transition range. Its evolution describes a classical stochastic phase
transition with the dc current as control parameter. We recognize that a
$50\,$% probability is reached at $\sim 10.6\,$mA, the current value where
also the slope of the characteristic RTN frequency evolution is roughly
estimated to invert (see fig. 2c, disregarding points with large errors
between $10$ and $10.6\,$mA). In general, random telegraph noise is a major
limitation especially for magnetic sensor applications Almeida, Wisniowski,
and Freitas (2008); Weitensfelder  _et al._ (2018) and can be removed here by
stabilization of the G- or C-mode. However, in the sense of the G- to C-state
transition, it can serve as a sort of easily accesible indicator of the
probabilistic properties.

##  V Discussion & Conclusion

We present here measurements of the dynamic C-state in vortex based spin
torque nano oscillators and analyze the system’s phase noise characteristics.
Regarding the evolution of the system’s flicker noise, we find its increase
with increasing applied dc current in the measured C-state regime that is
dominated by Hooge’s law. The measurement confirms that a minimum in the
flicker noise level exists also in the G-state when the active oscillation
volume saturates with the current.

At the transition between G- and C-state, we demonstrate the system
oscillating for several periods in the G- and several periods in the C-state,
stochastically switching between both states at a characteristic frequency in
the kHz range. The occurrence of random telegraph noise in the stochastic
regime is demonstrated. Its determined characteristic parameters are compared
with the state probability evaluated from a time-frequency analysis and can
serve as a sort of indicator of the stochastics. The parameters show a
controllable dependence of the probabilistic properties with the applied
current.

We believe, firstly, that the presented results bring more insight into the
noise characteristics of spin torque nano oscillators, specifically those
based on the texture of a magnetic vortex, but also partly generalizable to
other types of STOs. Secondly, the characterization of the stochastic
transition between the two dynamic states in vortex based spin torque
oscillators is an important step in order to add another feature to the
exploitation of STVOs, namely stochasticity, demonstrating the STVO’s
capability of real multifunctionality sought in multiple technologies, such as
neuromorphic systems.

###### Acknowledgements.

Gilles Cibiel, Serge Galliou and Enrico Rubiola are acknowledged for fruitful
discussions. S.W. acknowledges funding from Labex FIRST-TF under contract
number ANR-10-LABX-48-01. The work is supported by the French ANR project
”SPINNET” ANR-18-CE24-0012. P.T. acknowledges support under the Cooperative
Research Agreement Award No. 70NANB14H209, through the University of Maryland.
M.R. acknowledges support from Spanish MICINN (PGC2018-099422-A-I00) and
Comunidad de Madrid (Atracción de Talento Grant No. 2018-T1/IND-11935).

## Data Availability Statement

The data that support the findings of this study are available from the
corresponding author upon reasonable request.

## References

  * Makarov  _et al._ (2016) A. Makarov, T. Windbacher, V. Sverdlov, and S. Selberherr, “CMOS-compatible spintronic devices: a review,” Semiconductor Science and Technology 31, 113006 (2016).
  * Hughes  _et al._ (2012) H. Hughes, K. Bussmann, P. J. McMarr, S.-F. Cheng, R. Shull, A. P. Chen, S. Schafer, T. Mewes, A. Ong, E. Chen, M. H. Mendenhall, and R. A. Reed, “Radiation studies of spin-transfer torque materials and devices,” IEEE Transactions on Nuclear Science 59, 3027–3033 (2012).
  * Arun  _et al._ (2020) R. Arun, R. Gopal, V. K. Chandrasekar, and M. Lakshmanan, “Frequency enhancement and power tunability in tilted polarizer spin-torque nano-oscillator,” Journal of Applied Physics 127, 153903 (2020).
  * Locatelli, Cros, and Grollier (2013) N. Locatelli, V. Cros, and J. Grollier, “Spin-torque building blocks,” Nature Materials 13, 11–20 (2013).
  * Ebels  _et al._ (2017) U. Ebels, J. Hem, A. Purbawati, A. R. Calafora, C. Murapaka, L. Vila, K. J. Merazzo, E. Jimenez, M.-C. Cyrille, R. Ferreira, M. Kreissig, R. Ma, F. Ellinger, R. Lebrun, S. Wittrock, V. Cros, and P. Bortolotti, “Spintronic based RF components,” in  _2017 Joint Conference of the European Frequency and Time Forum and IEEE International Frequency Control Symposium (EFTF/IFC)_ (IEEE, 2017).
  * Sato  _et al._ (2012) R. Sato, K. Kudo, T. Nagasawa, H. Suto, and K. Mizushima, “Simulations and experiments toward high-data-transfer-rate readers composed of a spin-torque oscillator,” IEEE Transactions on Magnetics 48, 1758–1764 (2012).
  * Litvinenko  _et al._ (2020) A. Litvinenko, V. Iurchuk, P. Sethi, S. Louis, V. Tyberkevych, J. Li, A. Jenkins, R. Ferreira, B. Dieny, A. Slavin, and U. Ebels, “Ultrafast sweep-tuned spectrum analyzer with temporal resolution based on a spin-torque nano-oscillator,” Nano Letters 20, 6104–6111 (2020).
  * Fang  _et al._ (2019) B. Fang, M. Carpentieri, S. Louis, V. Tiberkevich, A. Slavin, I. N. Krivorotov, R. Tomasello, A. Giordano, H. Jiang, J. Cai, Y. Fan, Z. Zhang, B. Zhang, J. A. Katine, K. L. Wang, P. K. Amiri, G. Finocchio, and Z. Zeng, “Experimental demonstration of spintronic broadband microwave detectors and their capability for powering nanodevices,” Physical Review Applied 11 (2019), 10.1103/physrevapplied.11.014022.
  * Kaka  _et al._ (2005) S. Kaka, M. R. Pufall, W. H. Rippard, T. J. Silva, S. E. Russek, and J. A. Katine, “Mutual phase-locking of microwave spin torque nano-oscillators,” Nature 437, 389–392 (2005).
  * Lebrun  _et al._ (2017) R. Lebrun, S. Tsunegi, P. Bortolotti, H. Kubota, A. S. Jenkins, M. Romera, K. Yakushiji, A. Fukushima, J. Grollier, S. Yuasa, and V. Cros, “Mutual synchronization of spin torque nano-oscillators through a long-range and tunable electrical coupling scheme,” Nature Communications 8 (2017), 10.1038/ncomms15825.
  * Tsunegi  _et al._ (2018) S. Tsunegi, T. Taniguchi, R. Lebrun, K. Yakushiji, V. Cros, J. Grollier, A. Fukushima, S. Yuasa, and H. Kubota, “Scaling up electrically synchronized spin torque oscillator networks,” Scientific Reports 8 (2018), 10.1038/s41598-018-31769-9.
  * Romera  _et al._ (2018) M. Romera, P. Talatchian, S. Tsunegi, F. Abreu Araujo, V. Cros, P. Bortolotti, J. Trastoy, K. Yakushiji, A. Fukushima, H. Kubota, S. Yuasa, M. Ernoult, D. Vodenicarevic, T. Hirtzlin, N. Locatelli, D. Querlioz, and J. Grollier, “Vowel recognition with four coupled spin-torque nano-oscillators,” Nature (2018), 10.1038/s41586-018-0632-y.
  * Jenkins  _et al._ (2019) A. S. Jenkins, L. S. E. Alvarez, P. P. Freitas, and R. Ferreira, “Nanoscale true random bit generator based on magnetic state transitions in magnetic tunnel junctions,” Scientific Reports 9 (2019), 10.1038/s41598-019-52236-z.
  * Wittrock  _et al._ (2021) S. Wittrock, P. Talatchian, M. Romera, S. Menshawy, M. J. Garcia, M.-C. Cyrille, R. Ferreira, R. Lebrun, P. Bortolotti, U. Ebels, J. Grollier, and V. Cros, “Beyond the gyrotropic motion: Dynamic c-state in vortex spin torque oscillators,” Applied Physics Letters 118, 012404 (2021), https://aip.scitation.org/doi/pdf/10.1063/5.0029083 .
  * Petit-Watelot  _et al._ (2012) S. Petit-Watelot, J.-V. Kim, A. Ruotolo, R. M. Otxoa, K. Bouzehouane, J. Grollier, A. Vansteenkiste, B. V. de Wiele, V. Cros, and T. Devolder, “Commensurability and chaos in magnetic vortex oscillations,” Nature Physics 8, 682–687 (2012).
  * Devolder  _et al._ (2019) T. Devolder, D. Rontani, S. Petit-Watelot, K. Bouzehouane, S. Andrieu, J. Létang, M.-W. Yoo, J.-P. Adam, C. Chappert, S. Girod, V. Cros, M. Sciamanna, and J.-V. Kim, “Chaos in magnetic nanocontact vortex oscillators,” Physical Review Letters 123 (2019), 10.1103/physrevlett.123.147701.
  * Wittrock  _et al._ (2019) S. Wittrock, S. Tsunegi, K. Yakushiji, A. Fukushima, H. Kubota, P. Bortolotti, U. Ebels, S. Yuasa, G. Cibiel, S. Galliou, E. Rubiola, and V. Cros, “Low offset frequency 1/f flicker noise in spin-torque vortex oscillators,” Physical Review B 99 (2019), 10.1103/physrevb.99.235135.
  * Wittrock  _et al._ (2020) S. Wittrock, P. Talatchian, S. Tsunegi, D. Crété, K. Yakushiji, P. Bortolotti, U. Ebels, A. Fukushima, H. Kubota, S. Yuasa, J. Grollier, G. Cibiel, S. Galliou, E. Rubiola, and V. Cros, “Influence of flicker noise and nonlinearity on the frequency spectrum of spin torque nano-oscillators,” Scientific Reports 10 (2020), 10.1038/s41598-020-70076-0.
  * Hooge and Hoppenbrouwers (1969) F. Hooge and A. Hoppenbrouwers, “$1/f$ noise in continuous thin gold films,” Physica 45, 386–392 (1969).
  * Dussaux  _et al._ (2012) A. Dussaux, A. V. Khvalkovskiy, P. Bortolotti, J. Grollier, V. Cros, and A. Fert, “Field dependence of spin-transfer-induced vortex dynamics in the nonlinear regime,” Physical Review B 86 (2012), 10.1103/physrevb.86.014402.
  * Bianchini  _et al._ (2010) L. Bianchini, S. Cornelissen, J.-V. Kim, T. Devolder, W. van Roy, L. Lagae, and C. Chappert, “Direct experimental measurement of phase-amplitude coupling in spin torque oscillators,” Applied Physics Letters 97, 032502 (2010), http://dx.doi.org/10.1063/1.3467043 .
  * Quinsat  _et al._ (2010) M. Quinsat, D. Gusakova, J. F. Sierra, J. P. Michel, D. Houssameddine, B. Delaet, M.-C. Cyrille, U. Ebels, B. Dieny, L. D. Buda-Prejbeanu, J. A. Katine, D. Mauri, A. Zeltser, M. Prigent, J.-C. Nallatamby, and R. Sommet, “Amplitude and phase noise of magnetic tunnel junction oscillators,” Applied Physics Letters 97, 182507 (2010), http://dx.doi.org/10.1063/1.3506901 .
  * Keller  _et al._ (2010) M. W. Keller, M. R. Pufall, W. H. Rippard, and T. J. Silva, “Nonwhite frequency noise in spin torque oscillators and its effect on spectral linewidth,” Physical Review B 82 (2010), 10.1103/physrevb.82.054416.
  * Eklund  _et al._ (2014) A. Eklund, S. Bonetti, S. R. Sani, S. M. Mohseni, J. Persson, S. Chung, S. A. H. Banuazizi, E. Iacocca, M. Östling, J. Åkerman, and B. G. Malm, “Dependence of the colored frequency noise in spin torque oscillators on current and magnetic field,” Applied Physics Letters 104, 092405 (2014).
  * Tiberkevich, Slavin, and Kim (2008) V. S. Tiberkevich, A. N. Slavin, and J.-V. Kim, “Temperature dependence of nonlinear auto-oscillator linewidths: Application to spin-torque nano-oscillators,” Physical Review B 78 (2008), 10.1103/physrevb.78.092401.
  * Fermon and Pannetier-Lecoeur (2013) C. Fermon and M. Pannetier-Lecoeur, “Noise in gmr and tmr sensors,” in  _Giant Magnetoresistance (GMR) Sensors: From Basis to State-of-the-Art Applications_, Smart Sensors, Measurement and Instrumentation, edited by C. Reig, S. Cardoso, and S. Mukhopadhyay (Springer Berlin Heidelberg, 2013).
  * Gokce  _et al._ (2006) A. Gokce, E. R. Nowak, S. H. Yang, and S. S. P. Parkin, “1/f noise in magnetic tunnel junctions with MgO tunnel barriers,” Journal of Applied Physics 99, 08A906 (2006).
  * Almeida, Wisniowski, and Freitas (2008) J. M. Almeida, P. Wisniowski, and P. P. Freitas, “Low-frequency noise in mgo magnetic tunnel junctions: Hooge’s parameter dependence on bias voltage,” IEEE Transactions on Magnetics 44, 2569–2572 (2008).
  * Aliev  _et al._ (2007) F. G. Aliev, R. Guerrero, D. Herranz, R. Villar, F. Greullet, C. Tiusan, and M. Hehn, “Very low 1/f noise at room temperature in fully epitaxial fe/MgO/fe magnetic tunnel junctions,” Applied Physics Letters 91, 232504 (2007).
  * Slavin and Tiberkevich (2009) A. Slavin and V. Tiberkevich, “Nonlinear auto-oscillator theory of microwave generation by spin-polarized current,” IEEE Transactions on Magnetics 45, 1875–1918 (2009).
  * Cohen (1995) L. Cohen,  _Time-frequency Analysis_, Electrical engineering signal processing (Prentice Hall PTR, 1995).
  * Weitensfelder  _et al._ (2018) H. Weitensfelder, H. Brueckl, A. Satz, K. Pruegl, J. Zimmer, S. Luber, W. Raberg, C. Abert, F. Bruckner, A. Bachleitner-Hofmann, R. Windl, and D. Suess, “Comparison of sensitivity and low-frequency noise contributions in giant-magnetoresistive and tunneling-magnetoresistive spin-valve sensors with a vortex-state free layer,” Physical Review Applied 10 (2018), 10.1103/physrevapplied.10.054056.


# Differentiable Stochastic Halo Occupation Distribution

Benjamin Horowitz,1,2 ChangHoon Hahn,1 Francois Lanusse,3 Chirag Modi,4,5
Simone Ferraro2
1Department of Astrophysical Sciences, Princeton University, Princeton, NJ
08544, USA
2Lawrence Berkeley National Lab, 1 Cyclotron Road, Berkeley, CA 94720, USA
3AIM, CEA, CNRS, Université Paris-Saclay, Université Paris Diderot, Sorbonne
Paris Cité, F-91191 Gif-sur-Yvette, France
4Center for Computational Astrophysics, Flatiron Institute, 162 Fifth Ave.,
New York, NY 10010, USA
5Center for Computational Mathematics, Flatiron Institute, 162 Fifth Ave., New
York, NY 10010, USA E-mail<EMAIL_ADDRESS>
(Accepted XXX. Received YYY; in original form ZZZ)

###### Abstract

In this work, we demonstrate how differentiable stochastic sampling techniques
developed in the context of deep Reinforcement Learning can be used to perform
efficient parameter inference over stochastic, simulation-based, forward
models. As a particular example, we focus on the problem of estimating
parameters of Halo Occupancy Distribution (HOD) models which are used to
connect galaxies with their dark matter halos. Using a combination of
continuous relaxation and gradient re-parameterisation techniques, we can
obtain well-defined gradients with respect to HOD parameters through discrete
galaxy catalogs realisations. Having access to these gradients allows us to
leverage efficient sampling schemes, such as Hamiltonian Monte-Carlo, and
greatly speed up parameter inference. We demonstrate our technique on a mock
galaxy catalog generated from the Bolshoi simulation using the Zheng et al.
(2007) HOD model and find near identical posteriors as standard Markov Chain
Monte Carlo techniques with an increase of $\sim 8$x in convergence
efficiency. Our differentiable HOD model also has broad applications in full
forward model approaches to cosmic structure and cosmological analysis.

###### keywords:

methods: numerical – cosmology: theory – galaxies: haloes – galaxies:
fundamental parameters

††pubyear: 2021††pagerange: Differentiable Stochastic Halo Occupation
Distribution–A

##  1 Introduction

Over the past twenty years, there has been significant observational and
theoretical progress in connecting galaxies to their cosmic environments
(Peacock & Smith, 2000; van den Bosch et al., 2003; Kravtsov et al., 2004;
Wechsler et al., 2006; Neistein et al., 2011). Understanding this connection
is critical for understanding galaxy formation/evolution (Crain et al., 2009;
Zehavi et al., 2011) as well as using galaxies as bias tracers of the
underlying mass density for cosmological analyses (Benson et al., 2000a;
Desjacques et al., 2018). Studying this connection is a key component of many
upcoming galaxy surveys including the Prime Focus Spectrograph (Takada et al.,
2014; Tamura et al., 2016), the Dark Energy Spectroscopic Instrument (DESI
Collaboration et al., 2016; Abareshi et al., 2022), and the Nancy Grace Roman
Space Telescope (Spergel et al., 2015; Wang et al., 2022).

A key theoretical tool for these studies has been the Halo Occupation
Distribution (HOD; Lemson & Kauffmann, 1999; Benson et al., 2000b; Seljak,
2000; Scoccimarro et al., 2001; Berlind & Weinberg, 2002; Wechsler & Tinker,
2018), a framework that specifies how collapsed dark matter halos (Press &
Schechter, 1974; Bond et al., 1991; Cooray & Sheth, 2002) are populated with
galaxies. This is in contrast to “environmental" biasing schemes, such as
Eulerian or Lagrangian biasing schemes (Mann et al., 1998; Desjacques et al.,
2018), common in cosmological analyses of galaxy survey data (Cuesta et al.,
2016; Ivanov et al., 2020; To et al., 2021; Beutler et al., 2017). Unlike
environmental biasing schemes that only model summary statistics like power
spectrum of the galaxy field, a well-formulated HOD model provides direct
physical insight into galaxy formation physics through its parameters, which
are related to critical mass scales in galaxy-halo relation. For example, this
allows direct measurement of HOD parameters by comparing observed galaxy
populations with dynamical mass measurements, such as x-ray clusters (Zheng et
al., 2009; Mehrtens et al., 2016).

In standard HOD implementations (_e.g_ Zheng et al., 2007), the HOD model
specifies the probability distribution of the number of galaxies, $N$, hosted
by a dark matter halo given its properties, such as halo mass: $P(N|M_{\rm
halo})$. A semi-analytical halo-model approach can include HOD parameters to
predict two and higher point function (Cooray & Sheth, 2002), but those
predictions are often not accurate enough for analysing modern datasets.
Alternatively, a more precise Monte Carlo approach is often used to
stochastically assign galaxies to halos in a large simulation box following
the HOD prescription, and then the galaxy power spectrum or other quantities
of interest are directly measured from the simulation.

In practice, Markov-Chain Monte-Carlo methods have primarily been used to fit
HOD parameters from mock or actual data (_e.g._ White et al., 2011; Rodríguez-
Torres et al., 2016; Sinha et al., 2018). However these methods scale poorly
with the number of parameters that need to be fit. As novel decorated HOD
models increasingly add more assembly bias parameters to accurately capture
small scale observations, this exercise can become challenging, especially if
there are unforeseen degeneracies in the parameter space. These challenges can
be overcome more easily with parameter inference methods that rely on the
gradient information _i.e._ where we can estimate the response of the
observations with respect to the underlying parameters of the model, such as
Hamiltonian Monte Carlo (Duane et al., 1987a), Variational Inference
(Peterson, 1987; Beal, 2003; Blei et al., 2016) or combinations thereof
(Gabrié et al., 2022; Modi et al., 2022). However these methods are not
applicable in current HOD models as the implementation of their stochastic
galaxy assignment schemes make them classically non-differentiable.

A differentiable HOD framework will enhance dynamical forward modelling
frameworks that seek to reconstruct latent cosmological fields (_e.g._ Seljak
et al., 2017), which are constrained to use gradient based methods for
optimization due to the high dimensionality of the inference problem. These
frameworks generally rely on perturbative bias models that are accurate only
on large scales (Schmidt et al., 2019; Modi et al., 2019) or heuristic neural
network models with a large number of latent parameters (Modi et al., 2018). A
differentiable HOD approach will allow one to push to smaller scales with a
well-understood, physically developed model that has only a handful of
parameters.

An alternative to HOD models that maintains the requisite differentiability is
to use differentiable emulators (Kwan et al., 2015; Wibking et al., 2020) or
fitting functions of the observables like the one proposed in Hearin et al.
(2021) for galaxy assembly bias. However these are efficient only for the
particular summary statistics and cosmological parameters on which they are
trained. Hence, they require a new training set once these are varied.
Depending on the parameter space of interest this could be of prohibitive
computationally cost. In addition, separate emulators must be trained for each
summary statistic of interest as such methods do not match the galaxy
observations at the field level.

Motivated by this, here we adopt a different approach and aim to make the HOD
sampling itself differentiable. Our aim is to be able to compute gradients of
any observable with respect to HOD parameters through a particular realisation
of a galaxy catalog. Common wisdom states that differentiating through
stochastically sampled discrete random variables, such as the number of
satellites in a given halo, is not possible. However, modern Reinforcement
Learning have spurred the development of techniques to deal with these types
of categorical variables in the context of deep neural network training via
back propagation. In particular, we use the Gumbel-Softmax or CONCRETE method
(Jang et al., 2016; Maddison et al., 2016) that utilises continuous
distributions to approximate the sampling process of discrete stochastic
variables, such as galaxies, in a differentiable fashion. It relies on two
insights: 1) a re-parameterization for a discrete (or categorical)
distribution in terms of the Gumbel distribution (referred to as the “Gumbel
trick”; Maddison et al., 2014) and 2) making the corresponding function
continuous by using a continuous approximation that depends on a temperature
parameter, which in the zero-temperature case degenerates to the
discontinuous, original expression.

In this paper, we will implement the Gumbel-Softmax method in the context of
HOD models and apply it to mock datasets. In Sec. 2 we will describe our HOD
model and the methods used to allow differentiability of its categorical
outputs. In Sec. 3, we apply this technique to a Monte Carlo analysis of a
mock galaxy catalog constructed from the Planck-Bolshoi simulation. In Sec. 4
we compare the differentiable HOD model to that from a standard approach and
discuss its applications.

##  2 Method

In this section, we provide some background on the various components of our
HOD model, and detail our strategy to make this model differentiable. We
implement our model using Tensorflow Probability (Tran et al., 2016; Morgan,
2018), particularly the Tensorflow Distribution package (Dillon et al., 2017).

###  2.1 HOD Model

To describe the population of galaxies in our halos we use the standard Zheng
et al. (2007) HOD model. In the Zheng et al. (2007) model, the probability of
a given halo hosting $N$ galaxies is dictated solely by its mass — $P(N|M)$.
The model separately populates central and satellite galaxies, motivated by
theoretical studies (Kravtsov et al., 2004; Zheng et al., 2005), and has five
free parameters with some physical significance that can be related back to
well studied mass-luminosity relationships.

####  2.1.1 Central occupation

For central galaxies, the mean occupation function is step-like with a soft
cutoff to account for natural scatter between galaxy luminosity the halo host
mass. There are two free parameters controlling this function, the
characteristic minimum mass of halos hosting central galaxies above some
luminosity threshold, $M_{\textrm{min}}$, and the width of the cutoff profile,
$\sigma_{\log M}$:

$\langle
N_{\textrm{cen}}(M)\rangle=\frac{1}{2}\left[1+\textrm{erf}\left(\frac{\log
M-\log M_{\textrm{min}}}{\sigma_{\log M}}\right)\right],$ (1)

where erf is the standard error function and $M$ is the halo mass. Given the
mean occupation for a halo of a given mass, central galaxies are assigned to
halos by sampling a Bernoulli distribution:

$N_{\textrm{cen}}\sim\textrm{Bernoulli}\big{(}p=\langle
N_{\textrm{cen}}(M)\rangle\big{)}.$ (2)

####  2.1.2 Satellite occupation

Simulations suggest that satellites follow an approximately power-law
distribution with a slope close to unity at the high mass end. At lower
masses, the shape of the distribution changes and the overall distribution can
be parameterized as

$\langle N_{\textrm{sat}}(M)\rangle=\langle N_{\rm
cen}(M)\rangle\left(\frac{M-M_{0}}{M_{1}^{\prime}}\right)^{\alpha},$ (3)

where $\alpha$ is the power law slope at high masses, $M_{0}$ is the
characteristic mass of the change-over and $M_{1}^{\prime}$ is the
characteristic amplitude. This mean number of satellites for a given mass is
then used to define the intensity $\lambda$ of a Poisson distribution, from
which a particular number of satellites are drawn for each halo.

$N_{\textrm{sat}}\sim\textrm{Poisson}\big{(}\lambda=\langle
N_{\textrm{sat}}(M)\rangle\big{)}$ (4)

####  2.1.3 Spatial satellite distribution

In the Zheng et al. (2007) HOD, central galaxies are located at the center of
its host halo and satellite galaxies are distributed according to a Navarro et
al. (1997) profile (hereafter NFW). To sample the satellite galaxy positions,
we utilize the Robotham & Howlett (2018) implementation, which constructs an
efficient mapping from a random sample and the full NFW profile via the
calculation of the quantile function, _i.e._ the inverse Cumulative
Distribution Function (CDF). This can be written analytically as

$q(p;c,M_{\rm vir})=-\frac{1}{c}\left[1+\frac{1}{W_{0}(-e^{-pM_{\rm
vir}-1})}\right],$ (5)

where $W_{0}$ is the Lambert-W function, $M_{\rm vir}$ is the virial mass and
$c$ is the concentration parameter. Using this inverse CDF, we can now
randomly draw radial distances for satellites by sampling $p$ from $U[0,1]$
and mapping it to radii as $r/r_{\rm vir}=q(p,c)$. The angular position of the
satellite is sampled uniformly in this isotropic model.

###  2.2 Differentiable stochastic sampling

In this section, we review the key ideas behind differentiable stochastic
sampling, which will form the building blocks of DiffHOD.

####  2.2.1 Stochastic backpropagation by reparametrisation

One of the most common approaches for backpropagation through stochastic
sampling is the so-called reparametrisation trick, extensively used for
instance in Variational Auto-Encoders (Kingma & Welling, 2013; Jimenez Rezende
et al., 2014).

The key idea of this approach is to rewrite samples $z$ from a given
paramteric distribution $\mathbb{P}_{\theta}$ as a deterministic and
differentiable transformation $f$ applied to a fixed distribution
$\mathbb{P}_{\epsilon}$:

$z=f(\theta,\epsilon)\mbox{ with }\epsilon\sim\mathbb{P}_{\epsilon}$ (6)

This reparametrisation of the samples allows to side-step having to take
derivatives of the stochastic variable $\epsilon$ when computing derivatives
of some downstream function $h$ with respect to the distribution parameters
$\theta$:

$\frac{\partial}{\partial\theta}\mathbb{E}_{z\sim
p_{\theta}}\left[h(z)\right]=\mathbb{E}_{\epsilon\sim
p_{\epsilon}}\left[\frac{\partial}{\partial\theta}h(f(\theta,\epsilon))\right]$
(7)

In the right-hand side of this expression, the derivative now only involves
taking gradients of a deterministic function of $\theta$.

To provide a simple concrete example of such reparameterisation, let us
consider a Gaussian distribution of mean $\mu$ and standard deviation
$\sigma$. One can express a sample $z\sim\mathcal{N}(\mu,\sigma^{2})$ as
$z=\mu+\sigma\epsilon$ with $\epsilon\sim\mathcal{N}(0,I)$, making it trivial
to take derivatives of the samples with respect to the parameters of the
distribution ($\mu$ and $\sigma$).

####  2.2.2 Gumbel-Softmax trick for categorical variables

The reparameterisation trick as presented above requires the samples to be
expressible as a deterministic and differentiable function of a random
variable. While this can often be achieved for continuous distributions, it is
typically not directly possible for discrete categorical variables. To
overcome this limitation, the Gumbel-Softmax trick (Jang et al., 2016;
Maddison et al., 2016) introduces a relaxation of a categorical distribution
to a continuous distribution, which can then be handled with the
reparameterisation trick.

Let $z$ be a categorical variable with class probabilities
$\pi_{1},\pi_{2},...\pi_{j}$ that we wish to sample. We assume that
categorical samples are encoded as $N$-dimensional one-hot vectors, _i.e._
they are 1$\times N$ vectors with all elements 0 except the the element
corresponding to the sampled class which is 1. The simplest way to sample z is
by

$z=\textrm{onehot}(\textrm{max}{\,i|\pi_{1}+...+\pi_{i-1}\leq U}),\quad
U\sim\textrm{Uniform}(0,1)$ (8)

A first step towards making these samples differentiable is to use the Gumbel-
Max trick (Gumbel, 1954; Maddison et al., 2014), which reparametrizes
categorical sampling as

$z=\textrm{onehot}\left(\textrm{argmax}_{i}[g_{i}+\log{(\pi_{i})}]\right),$
(9)

where $g_{i}$ are i.i.d. random variables drawn from the Gumbel distribution
between 0 and 1, $\textrm{Gumbel}(0,1)$111The standard Gumbel distribution is
defined by cumulative distribution function $\textrm{CDF}(x)=\exp(-\exp(x))$
and probability density function, $\textrm{PDF}(x)=\exp(-(x+\exp(-x))$.. This
reparameterization trick refactors the sampling of $z$ into a deterministic
function of the parameters ($\pi$) and some independent noise with a fixed
distribution.

However the reparametrized function is still non-differentiable due to the
argmax function. A continuous, differentiable approximation to this is given
by a softmax function,

$\mathrm{softmax}(\mathbf{z},\tau)_{i}=\frac{\mathrm{e}^{z_{i}/\tau}}{\sum_{j=1}^{k}\mathrm{e}^{z_{j}/\tau}},\quad\mathbf{z}=(z_{1},z_{2},...,z_{k})$
(10)

where $\tau$ is a free parameter sometimes referred to as the “temperature".

Using this approximation relaxes the discreteness of the Gumbel-Max trick and
generates a $k$-dimensional vector $\mathbf{z}$

$\hat{z}_{i}=\frac{\exp((\log(\pi_{i})+g_{i})/\tau)}{\sum_{j}\exp((\log(\pi_{j})+g_{j})/\tau)},$
(11)

We recover the true discrete function in the limit of $\tau\rightarrow 0$. As
this function is analytical in class probabilities $\pi$ for $\tau>0$, we can
estimate the gradients of the observed samples $\mathbf{z}$ with respect to
the parameters parameterizing $\pi$.

In the remaining of this work, we will make use of the special case when the
number of classes is 2, i.e. when the categorical distribution reduces to a
Bernoulli distribution. In this special binary case Eq. 11 can be simplified.
Maddison et al. (2016) refer to the resulting distribution as BinConcrete; we
will refer to it in this work as a Relaxed Bernoulli distribution. Using the
fact that the difference of two Gumbel variables follows a Logistic
distribution222The standard Logistic distribution follows the following
probability density function: $p(x)=\frac{\exp(-x)}{(1-\exp(-x))^{2}}$,
(Maddison et al., 2016) shows the Relaxed Bernoulli can be reparameterized as:

$z=\frac{1}{1+\exp(-(\log\pi+\epsilon)/\tau)}\mbox{ with
}\epsilon\sim\mathrm{Logistic}(0,1)\;.$ (12)

where in this expression $\pi$ is the odds-ratio $\pi=p/(1-p)$ if $p$ is the
probability of corresponding Bernoulli distribution.

###  2.3 DiffHOD Implementation

We now have all the elements needed to build a differentiable HOD (hereafter
DiffHOD) model. We describe in this section our strategy for sampling central
and satellites occupation, and satellites positions.

####  2.3.1 Differentiable central occupation sampling

As described in Sec. 2.1, the central occupation is defined by a Bernoulli
distribution, with a parameter $p=\left\langle N_{\rm
cen}(M|M_{\mathrm{min}},\sigma_{\log M})\right\rangle$ defined as a
deterministic function.

Here, we can directly apply the Gumbel-Softmax trick introduced above in the
specical case of a binary variable. We therefore sample the central occupation
of each halo using the Relaxed Bernoulli distribution:

$N_{\rm cen}=\frac{1}{1+\exp(-(\log(\frac{p}{1-p})+\epsilon)/\tau)}\mbox{ with
}\epsilon\sim\mathrm{Logistic}(0,1)\;.$ (13)

at given temperature $\tau$. Fig. 1 illustrates the high agreement between the
halo occupancy obtained by sampling centrals with this Relaxed Distribution
compared to the analytical expectation at our fiducial choice of $\tau=0.1$.

####  2.3.2 Differentiable satellite occupation sampling

For satellites, we aim to define a differentiable approach to sampling from a
Poisson distribution with intensity $\lambda=\left\langle N_{\rm
sat}(M|M_{0},M_{1}^{\prime},\alpha)\right\rangle$, also a deterministic and
differentiable function. To build on the Gumbel-Softmax trick, we propose to
replace conventional Poisson sampling of the total number of satellites by
sampling each satellite individually from a Bernoulli distribution.

Let us consider a halo with a Poisson rate $\lambda$ for its satellite
occupancy. We assume the halo can have a maximum of $N$ satellites, then for
each potential satellite we sample from a Bernoulli distribution with
probability $p=\lambda/N$ whether this satellite will be included in the halo.
The resulting statistics of the number of satellites with this procedure will
be Binomial (as N draws from i.i.d. Bernoulli).

More formally, we propose to approximate the Poisson distribution with
intensity $\lambda$ of a standard HOD by a Binomial distribution with $N$
trials and probability $p=\lambda/N$:

$N_{\rm sat}\sim\mathrm{Binomial}\left(N,p=\frac{\left\langle N_{\rm
sat}\right\rangle}{N}\right)$ (14)

By construction, this Binomial distribution will yield the same mean number of
satellites as the Poisson distribution, however the variance of both
distributions is different:

$\displaystyle\mathrm{Var}(N_{\rm sat}^{\rm Pois.})$
$\displaystyle=\left\langle N_{\rm sat}\right\rangle$ (15)
$\displaystyle\mathrm{Var}(N_{\rm sat}^{\rm Bin.})$
$\displaystyle=\left\langle N_{\rm
sat}\right\rangle*\left(1-\frac{\left\langle N_{sat}\right\rangle}{N}\right)$
(16)

From this expression, one can foresee that the Binomial distribution will be a
close approximation to a Poisson distribution when the ratio
$\frac{\left\langle N_{\rm sat}\right\rangle}{N}$ is small, i.e. when the
number of trials is large compared to the expected number of satellites. This
is actually known as the law of rare events, and at fixed $N*p$ the Binomial
distribution ${\rm Binomial}(N,p)$ converges to a Poisson distribution when
the number of trials $N\rightarrow\infty$.

Figure 1:  Halo occupancy distribution for central (top) and satellite
galaxies (bottom) as a function of halo virial mass of our differentiable HOD
model (DiffHOD solid). Different colors indicate different temperature values
used in the Gumbel-Softmax approximation (see Eq. 11). We include the
occupancy distribution from the standard Zheng et al. (2007) HOD model for
reference (star). In this work, we use DiffHODw̃ith $\tau=0.1$, which is in
good agreement with the standard HOD model throughout the full halo mass
range.

In practice, we will need to limit the number of trials $N$ to some finite
value and Fig. 2 compares the shapes of satellites distributions with a
Poisson model versus a Binomial model for two different choices of $N$, and
for different halo masses. A higher value of $N$ can improve accuracy but will
also increases memory costs, which scales linearly with the maximum number of
satellite galaxies encoded in our one-hot embedding. Meanwhile, too low of a
$N$ can bias results by artificially reducing the variance of the satellite
occupation distribution, or worse, truncating the satellite galaxies of the
most massive halos.

Our fiducial choice in this work is $N=48$. In Fig. 1 (bottom), we can see
that the satellite population reaches a maximum of $\sim 40$ galaxies for the
most massive $\sim 10^{15}M_{\odot}$ halos. For the mass range considered
here, we find $N>40$ does not significantly improve the statistical match in
summary statistics (correlation function, power spectra, etc.) of the
resulting galaxy fields and larger $N$ further will increase memory
requirements and computational time. If the end statistic of interest is
particularly sensitive to galaxy populations in the most massive clusters, a
higher $N$ might be needed. If one is limited by memory or implementing HOD
for halos with a broader mass range, then it may be more efficient to have
multiple halo mass bins with different maximum number of allowed satellites
$N$.

This Binomial assumption for the sampling of satellites brings two concrete
advantages:

  1. 1.

Having restated satellite sampling as draws from Bernoulli distributions, we
can make the procedure differentiable by using the Relaxed Bernoulli,
similarly to centrals.

  2. 2.

Using a fixed number $N$ of potential satellites gives us a practical way to
handle varying number of satellites per halos.

Concretely, for each candidate satellite $i\in\llbracket 1,N\rrbracket$ of a
halo, we sample whether the satellite will be included in the halo using:

$z_{i}=\frac{1}{1+\exp(-(\log(\frac{p}{1-p})+\epsilon_{i})/\tau)}\mbox{ with
}\epsilon\sim\mathrm{Logistic}(0,1)\;.$ (17)

where $p=\frac{\left\langle N_{sat}\right\rangle}{N}$ and $\tau$ is the
temperature. As a result, for each halo we obtain a vector $\mathbf{z}$ of
size $N$ which encodes active satellites for the halo.

In all downstream computations, this vector $\mathbf{z}$ can be interpreted as
a weight between 0 and 1 to apply to each of the $N$ satellites of each halo,
for instance in the computation of two-point correlation functions.

Figure 2:  Comparison of satellite occupation distributions for different halo
masses, under various assumed distributions: Poisson (solid line), Binomial
with 100 trials (dashed line), Binomial with 30 trials (dotted line). By
construction, the Binomial approximation recovers the mean number of
satellites, but for massive halos limiting the number of samples N will lead
to departure in the spread of the distribution compared to a Poisson
distribution. Note that for the $m_{\rm halo}=10^{14.8}$ halo the mean number
of satellites is above $30$, so the Binomial distribution is not well defined.

####  2.3.3 Differentiable sampling from NFW distribution

The last step to complete our HOD implementation is to sample the position of
satellites based on an NFW profile centered at the halo position. There are
two difficulties here: 1. sampling positions for a varying number of
satellites, 2. making the sampled positions differentiable with respect to the
NFW parameters.

The first question of dealing with varying number of satellites is solved by
our Binomial model with fixed number of trials $N$. For each halo, we will
sample the same number of $N$ sets of coordinates, one for each potential
satellite. Whether these coordinates will actually contribute in downstream
computations will depend on the per halo satellite occupation vector
$\mathbf{z}$ introduced above.

The second question, of the differentiability of stochastic coordinates, is
again solved by applying the re-parameterisation trick to the NFW profile.
From Eq. (5) we know the CDF of the NFW profile, meaning that for each
satellite $i\in\llbracket 1,N\rrbracket$ of a given halo we can sample the
halo-centric satellite radial distance as:

$r_{i}=r_{\rm vir}\ q(\epsilon,c)\mbox{ with
}\epsilon\sim\mathrm{Uniform}(0,1)$ (18)

where $q$ is a differentiable function of parameters $c$.

In our adaptation of the Zheng et al. (2007) model, we assume an isotropic NFW
distribution for satellite, so to retrieve halo-centric cartesian coordinates
$x_{i},y_{i},z_{i}$ of a given satellite, we first sample $x_{i},y_{i},z_{i}$
on the unit sphere and then multiply these coordinates by the $r_{i}$ value
sampled above. This is nothing more than an another reparameterisation step
and the resulting cartesian coordinates remain fully differentiable with
respect to the NFW parameters.

We note that contrary to the sampling of central of satellites which are
differentiable approximation to a standard HOD (due to the discrete variables
involved), this differentiable implementation of NFW sampling is exact.

####  2.3.4 Impact of temperature parameter $\tau$

In the DiffHOD model, we introduce temperature, $\tau$, as a free parameter
(Eq. 13). Depending on the context/implementation, it may be beneficial to
anneal (_i.e._ reduce) $\tau$ over the course of the optimization such that
$\tau\rightarrow 0$. For instance, in the original papers (Jang et al., 2016;
Maddison et al., 2016), the Gumbel-Softmax trick was used in the context of
training a neural network where only the optimal network weights were of
interest and hence $\tau$ was reduced to nearly zero over the course of the
training. However in our case, we are interested in stochastic sampling, and
not optimization, so we will pick a single temperature that well approximates
the target distribution (_i.e._ unbiased) while maintaining reasonable
derivative properties (i.e. less noisy).

In Fig. 1, we present the central (top) and satellite (bottom) occupancy
distributions of our DiffHOD model for different temperatures: $\tau=0.02$
(red) 0.1 (green), 0.5 (orange), and 1 (blue). We include the occupancy
distributions for the standard HOD model for comparison (star). For this work,
we take an experimental approach for determining $\tau$, as advocated in
Maddison et al. (2016): the temperature should set as high as possible while
maintaining the desired accuracy of the target distribution. We find that a
fixed $\tau=0.1$ provides high accuracy while maintaining stable gradients for
both sampling the number of galaxies in each halo as well as the positions of
the satellites from the NFW profile (see later Fig. 3).

##  3 Experimentation

Figure 3:  _Top_ : Derivatives of the central and satellite occupancy
functions with respect to the HOD parameters for various temperatures:
$\tau=0.02$ (left) 0.1 (center), and 1 (right). The derivatives are calculated
at the fiducial HOD parameter values. We include derivatives analytically
derived from Eqs. 1 and 3 for comparison (star). _Bottom_ : Derivatives of the
galaxy power spectrum with respect to the HOD parameters, evaluated at the
fiducial HOD parameter values. We include derivatives calculated using the
standard HOD with finite difference for comparison (star). Our DiffHOD model
with $\tau=0.1$ provides sufficiently smooth derivatives that are in good
agreement with analytical derivatives for the occupancy function and with
standard HOD derivatives for the power spectrum.

To test our implementation, we construct a fiducial mock galaxy catalog from
the Planck Bolshoi simulation halo catalog (Klypin et al., 2011) at $z=0$. We
treat this catalog as our mock observation. This simulation has a side length
of 250 $h^{-1}$ Mpc, and contains 1,367,493 unique halos ranging in mass from
$1.1\times 10^{15}M_{\odot}$ down to $2.7\times 10^{8}M_{\odot}$. We use
halotools (Hearin et al., 2016) to generate our fiducial catalog with HOD
parameter values:

$\displaystyle\log{M_{\textrm{min}}}=12.02,\sigma_{\log{M}}=0.26,\log
M_{0}=11.38,$ $\displaystyle\log{M_{1}}=13.31,\alpha=1.06.$

These parameters correspond to the best fit values from Zheng et al. (2007)
using an SDSS-based galaxy catalog (Zehavi et al., 2005) with $r$-band
absolute magnitude threshold of $M_{r}<-20$.

###  3.1 Summary Statistics and Derivatives with DiffHOD

Next, we compare the mock observations to parallel galaxy catalogs constructed
using DiffHOD (Section 2). The galaxy catalogs are then painted onto a grid
using a differentiable Cloud-In-Cell (CIC) painting method (Modi et al., 2020)
and its real-space power-spectra calculated via a differentiable Tensorflow
power spectrum implementation (Horowitz et al., 2021). For DiffHOD, all steps
are differentiable so the overall mapping from original halo catalog to end
power spectrum is also differentiable via the chain rule. In Fig. 3, we
present the DiffHOD derivatives of the central and satellite occupancy
distribution functions (top) and resulting power spectra (bottom) with respect
to the HOD parameters for different temperatures: $\tau=0.02$ (left), $0.1$
(middle), and $1$ (right). The derivatives are evaluated at the fiducial HOD
parameter values. In the center panels, for comparison, we include derivatives
of the central and satellite occupancy distribution functions derived
analytically and derivatives of the power spectrum derived using the standard
HOD with finite differences (star).

With $\tau=0.02$, the power spectrum derivatives have significant numerical
noise. The $\tau=1.0$ derivatives are smoother but we find inaccurate
occupancy distributions (Fig. 1) and a significantly biased power spectrum.
Meanwhile, with $\tau=0.1$ there is still some noticeable numerical noise in
the derivatives, however, this is sufficiently smooth for our application and
for our optimization to be well behaved. We also find that the $\tau=0.1$
DiffHOD derivatives are in good agreement with analytical derivatives for the
occupancy function and with standard HOD derivatives estimated using finite
differences for the power spectrum.

We compare the differentiable power spectrum from DiffHOD to the power
spectrum from the standard HOD model in Fig. 4. We use the fiducial HOD
parameter values and estimate the error bars from 500 realizations of the
DiffHOD model. We find good agreement (better than $1\%$) across all scales,
with particularly good agreement for central galaxies, which are not sampled
from the NFW profile. Low $k$ modes are most sensitive to the most massive
halos whose satellite galaxy populations are truncated by our one-hot
distribution ($N=48$; Section 2.3); those who are interested in this regime
can further increase $N$. However, even at $k<0.15h/{\rm Mpc}$, we find good
agreement between the power spectra.

Figure 4:  Comparison of power spectrum derived using a standard HOD model
($P_{\rm HT}$) versus our DiffHOD model. We plot the ratios for central and
satellite galaxies, as well as for all galaxies. Error bars represent the
standard deviation of the DiffHOD model estimated from 500 realizations of the
galaxy sampling. Grey band represents the standard deviation of $P_{\rm HT}$
for all galaxies. We find good agreement between the differentiable power
spectrum and $P_{\rm HT}$ across all scales ($k<3h/{\rm Mpc}$), as well as
similar distributional properties.

###  3.2 Monte Carlo Analysis with DiffHOD

Lastly, we demonstrate that we can derive unbiased inference using DiffHOD, by
comparing the posteriors on HOD parameters derived using DiffHOD to posteriors
derived using standard methods with the same mock observations and likelihood.
We use the power spectrum measured from the fiducial galaxy catalog as our
mock observations, and construct a covariance matrix using 100 galaxy catalog
realizations at the fiducial HOD catalogs. To this covariance, we added a
small constant diagonal term ($8.0\times 10^{-5}$) to improve numerical
stability. We limit our comparison to $k<1.0$. In both DiffHOD and standard
cases, we use the same halo catalog used to construct the mock observations,
so that the only source of error is variation caused by the HOD model. For
this analysis we impose wide Gaussian priors, $\mathcal{N}(\mu,\sigma^{2})$,
where $\mu$ is the mean, and $\sigma^{2}$ is the variance, on our parameter
values as follows:

$\displaystyle\log{M_{\textrm{min}}}\sim\mathcal{N}(12.0,0.5),\sigma_{\log{M}}\sim\mathcal{N}(0.25,0.2),$
$\displaystyle\log
M_{0}\sim\mathcal{N}(11.25,0.5),\log{M_{1}}\sim\mathcal{N}(13.20,0.5),$
$\displaystyle\alpha\sim\mathcal{N}(1.0,0.2).$

To derive the posteriors using DiffHOD, we sample over the HOD parameter using
Hamiltonian Monte Carlo (HMC) (Duane et al., 1987b; Neal et al., 2011). We use
the NoUTurn HMC implementation (Hoffman et al., 2014) in Tensorflow
Probability. We use three chains initialized around our fiducial HOD
parameters with over 1,000 steps (300 steps of burn-in).

For the standard approach, we run the same analysis using the standard HOD.
However, since this implementation does not allow easy differentiation, we
cannot use HMC instead use a Markov-Chain Monte Carlo analysis. We use the
emcee (Foreman-Mackey et al., 2013) implementation with 10 walkers and 6,000
steps. We present the posteriors on the HOD parameters for the DiffHOD (black)
and standard (blue) analyses in Fig. 5 and list the median posterior values in
Table 1. We mark the fiducial (“true”) HOD values in green. As we are using a
single galaxy realization for our mock observations, we expect some variation
between the best fit parameters and the true values. The posteriors derived
using DiffHOD and HMC is in excellent agreement with the posteriors derived
using the standard HOD and MCMC.

Our HMC analysis takes approximately 10 hours on a single Tesla V100-PCIE-32Gb
GPU. Meanwhile, the standard approach takes substantially more time to get
comparable results: $\sim$200 hours on 1 CPU — ${\sim}20\times$ slower than
our DiffHOD analysis. Some of this improvement is due to the fact that our
DiffHOD implementation is faster per iteration than the standard HOD
implementation (${\sim}1$ and ${\sim}4$ seconds per iteration, respectively).
Most of the improvement, however, comes from the fact that DiffHOD allows us
to exploit a more efficient gradient-based method to derive the posterior.

We compare the DiffHOD and standard approaches in more detail by comparing the
effective sample size of each chain per function evaluation (Gelman et al.,
2013). The effective sample size incorporates information about the auto-
correlations within a chain; _i.e._ it accounts for the dependent
relationships between the samples. We calculate it from the output of the
Markov chain:

$N_{\textrm{eff}}=\frac{N}{1+\sum\limits_{t=1}^{\infty}\rho_{t}}$ (19)

where $N$ is the number of samples in the chain and $\rho_{t}$ is the
autocorrelation of length $t$. Averaging over all HOD parameters, we find an
mean effective sample size of $524.2$ for our DiffHOD HMC evaluation and
$403.9$ for the standard MCMC evaluation. This corresponds to an effective
sample of 0.05 per step for the HMC and 0.006 for the MCMC evaluation.

Figure 5: Posterior distributions of the Zheng et al. (2007) HOD parameters
derived using DiffHOD with Hamiltonian Monte Carlo (black) and the standard
HOD with Markov Chain Monte Carlo (blue). We use a single mock galaxy catalog
realization at the fiducial HOD parameter values (green) as our observations.

##  4 Conclusions

In this work we have constructed a differentiable stochastic HOD model going
from a halo catalog to an observed galaxy power spectrum. This allows us to
use derivative-based optimization methods to quickly optimize for the
underlying model parameters. This is the first time differentiable stochastic
models have been used in the astrophysics literature. We find the DiffHOD
model provides a $4\times$ increase of speed versus the same analysis
performed via MCMC with the standard HOD implementation. DiffHOD is an
alternative to a number of recent works focusing on emulating galaxy
clustering statistics (Kwan et al., 2015; Wibking et al., 2019; Kobayashi et
al., 2020; Wibking et al., 2020; Hearin et al., 2022). Unlike emulator based
methods, DiffHOD works at the level of the halo catalog and allows fast,
differentiable, generation of any summary statistic with respect to the HOD
parameters.

In this work, we have focused on the Zheng et al. (2007) HOD model, but our
methods can be easily extended to a broad class of models. While standard HODs
are based only on halo mass, in general various properties of the halos
environment and formation history could effect the galaxy properties (Zhu et
al., 2006; Croton et al., 2007). Galaxy assembly bias has been argued
(Feldmann & Mayer, 2015; Hadzhiyska et al., 2020) to cause significant
deviations between predictions of standard HOD models and those from
hydrodynamical simulations. Decorated HOD models have been introduced to
account for assembly bias (Hearin et al., 2016) and has been extended to
include other possible effects (Yuan et al., 2018). These models still rely on
stochastic discrete sampling for assigning centrals and satellites so they can
be modelled in a differentiable way using the techniques described in this
work. As the dimensionality of our problem increases, either with extended HOD
models or joint analysis with cosmological parameters, we expect the relative
performance of derivative based methods, like HMC, over pure sampling based
methods to further improve (Neal et al., 2011).

| DiffHOD | Standard HOD
---|---|---
$\log(M_{\textrm{min}})$ | $12.03^{+0.15}_{-0.03}$ | $12.01^{+0.14}_{-0.02}$
$\sigma_{\log{M}}$ | $0.28^{+0.19}_{-0.16}$ | $0.27^{+0.18}_{-0.16}$
$\log M_{0}$ | $11.25^{+0.38}_{-0.43}$ | $11.27^{+0.49}_{-0.50}$
$\log{M_{1}}$ | $13.32^{+0.23}_{-0.23}$ | $13.34^{+0.29}_{-0.22}$
$\alpha$ | $0.96^{+0.19}_{-0.19}$ | $0.96^{+0.18}_{-0.19}$
Table 1: Posterior values from the HOD analyses using the DiffHOD and standard
HOD model. Uncertainties are estimate from the $16\%$ and $84\%$ quantiles.

Differentiable HOD models have even more apparent applications in the case of
dynamical forward model large scale reconstructions (Seljak et al., 2017) when
paired with efficient differentiable halo finding methods (Modi et al., 2018,
2020; Kodi Ramanah et al., 2019). While it is possible to perform these
reconstructions by interpreting the galaxy field as a biased version of the
dark matter field (i.e. in Horowitz et al. (2021)), inaccuracies in this
prescription will result in biases that would be difficult to account for in
cosmological constraints. Through joint inference of the HOD parameters with
the initial density field, these astrophysical uncertainties can be rigorously
marginalized out. Differentiable models are critical for this application as
the optimization is highly multidimensional (approximately number of particles
in the simulation) and would be computationally infeasible without gradient-
based methods.

While in this work we have highlighted using our DiffHOD model inside an HMC
framework, one can exploit its automatic differentiation for a variety of
first order optimization and parameter inference methods. For example standard
Variational Inference relies on having well defined derivatives for the
optimization of latent space parameters describing the likelihood surface
(Peterson, 1987; Beal, 2003; Blei et al., 2016). Variational inference could
further accelerate parameter inference when compared to Hamiltonian Monte
Carlo or nested sampling methods (Gunapati et al., 2018).

## Acknowledgements

We thank Andrew Hearin and Matt Becker for very useful discussions and
suggestions. BH and CH are supported by the AI Accelerator program of the
Schmidt Futures Foundation. SF is supported by the Physics Division of
Lawrence Berkeley National Laboratory. This research used resources of the
National Energy Research Scientific Computing Center, a DOE Office of Science
User Facility supported by the Office of Science of the U.S. Department of
Energy under Contract No. DEC02-05CH11231.

## References

  * Abareshi et al. (2022) Abareshi B., et al., 2022, Overview of the Instrumentation for the Dark Energy Spectroscopic Instrument (arXiv:2205.10939), doi:10.48550/arXiv.2205.10939
  * Beal (2003) Beal M. J., 2003, PhD thesis, University of London, University College London (United Kingdom)
  * Benson et al. (2000a) Benson A. J., Cole S., Frenk C. S., Baugh C. M., Lacey C. G., 2000a, MNRAS, 311, 793
  * Benson et al. (2000b) Benson A. J., Cole S., Frenk C. S., Baugh C. M., Lacey C. G., 2000b, MNRAS, 311, 793
  * Berlind & Weinberg (2002) Berlind A. A., Weinberg D. H., 2002, ApJ, 575, 587
  * Beutler et al. (2017) Beutler F., et al., 2017, Monthly Notices of the Royal Astronomical Society, 466, 2242
  * Blei et al. (2016) Blei D. M., Kucukelbir A., McAuliffe J. D., 2016, arXiv e-prints, p. arXiv:1601.00670
  * Bond et al. (1991) Bond J. R., Cole S., Efstathiou G., Kaiser N., 1991, ApJ, 379, 440
  * Cooray & Sheth (2002) Cooray A., Sheth R., 2002, Phys. Rep., 372, 1
  * Crain et al. (2009) Crain R. A., et al., 2009, MNRAS, 399, 1773
  * Croton et al. (2007) Croton D. J., Gao L., White S. D. M., 2007, MNRAS, 374, 1303
  * Cuesta et al. (2016) Cuesta A. J., et al., 2016, MNRAS, 457, 1770
  * DESI Collaboration et al. (2016) DESI Collaboration et al., 2016, arXiv e-prints, p. arXiv:1611.00036
  * Desjacques et al. (2018) Desjacques V., Jeong D., Schmidt F., 2018, Phys. Rep., 733, 1
  * Dillon et al. (2017) Dillon J. V., et al., 2017, arXiv e-prints, p. arXiv:1711.10604
  * Duane et al. (1987a) Duane S., Kennedy A. D., Pendleton B. J., Roweth D., 1987a, Physics Letters B, 195, 216
  * Duane et al. (1987b) Duane S., Kennedy A. D., Pendleton B. J., Roweth D., 1987b, Physics letters B, 195, 216
  * Feldmann & Mayer (2015) Feldmann R., Mayer L., 2015, MNRAS, 446, 1939
  * Foreman-Mackey et al. (2013) Foreman-Mackey D., Hogg D. W., Lang D., Goodman J., 2013, PASP, 125, 306
  * Gabrié et al. (2022) Gabrié M., Rotskoff G. M., Vanden-Eijnden E., 2022, Proceedings of the National Academy of Sciences, 119, e2109420119
  * Gelman et al. (2013) Gelman A., Carlin J. B., Stern H. S., Dunson D. B., Vehtari A., Rubin D. B., 2013, Bayesian data analysis. CRC press
  * Gumbel (1954) Gumbel E. J., 1954, Statistical theory of extreme values and some practical applications: a series of lectures. Vol. 33, US Government Printing Office
  * Gunapati et al. (2018) Gunapati G., Jain A., Srijith P. K., Desai S., 2018, arXiv e-prints, p. arXiv:1803.06473
  * Hadzhiyska et al. (2020) Hadzhiyska B., Bose S., Eisenstein D., Hernquist L., Spergel D. N., 2020, MNRAS, 493, 5506
  * Hearin et al. (2016) Hearin A. P., Zentner A. R., van den Bosch F. C., Campbell D., Tollerud E., 2016, MNRAS, 460, 2552
  * Hearin et al. (2021) Hearin A. P., Chaves-Montero J., Becker M. R., Alarcon A., 2021, arXiv e-prints, p. arXiv:2105.05859
  * Hearin et al. (2022) Hearin A. P., Ramachandra N., Becker M. R., DeRose J., 2022, The Open Journal of Astrophysics, 5, 3
  * Hoffman et al. (2014) Hoffman M. D., Gelman A., et al., 2014, J. Mach. Learn. Res., 15, 1593
  * Horowitz et al. (2021) Horowitz B., Zhang B., Lee K.-G., Kooistra R., 2021, ApJ, 906, 110
  * Ivanov et al. (2020) Ivanov M. M., Simonović M., Zaldarriaga M., 2020, J. Cosmology Astropart. Phys., 2020, 042
  * Jang et al. (2016) Jang E., Gu S., Poole B., 2016, arXiv e-prints, p. arXiv:1611.01144
  * Jimenez Rezende et al. (2014) Jimenez Rezende D., Mohamed S., Wierstra D., 2014, arXiv e-prints, p. arXiv:1401.4082
  * Kingma & Welling (2013) Kingma D. P., Welling M., 2013, arXiv e-prints, p. arXiv:1312.6114
  * Klypin et al. (2011) Klypin A. A., Trujillo-Gomez S., Primack J., 2011, ApJ, 740, 102
  * Kobayashi et al. (2020) Kobayashi Y., Nishimichi T., Takada M., Takahashi R., Osato K., 2020, Phys. Rev. D, 102, 063504
  * Kodi Ramanah et al. (2019) Kodi Ramanah D., Charnock T., Lavaux G., 2019, Phys. Rev. D, 100, 043515
  * Kravtsov et al. (2004) Kravtsov A. V., Berlind A. A., Wechsler R. H., Klypin A. A., Gottlöber S., Allgood B., Primack J. R., 2004, ApJ, 609, 35
  * Kwan et al. (2015) Kwan J., Heitmann K., Habib S., Padmanabhan N., Lawrence E., Finkel H., Frontiere N., Pope A., 2015, ApJ, 810, 35
  * Lemson & Kauffmann (1999) Lemson G., Kauffmann G., 1999, MNRAS, 302, 111
  * Maddison et al. (2014) Maddison C. J., Tarlow D., Minka T., 2014, arXiv e-prints, p. arXiv:1411.0030
  * Maddison et al. (2016) Maddison C. J., Mnih A., Whye Teh Y., 2016, arXiv e-prints, p. arXiv:1611.00712
  * Mann et al. (1998) Mann R. G., Peacock J. A., Heavens A. F., 1998, MNRAS, 293, 209
  * Mehrtens et al. (2016) Mehrtens N., et al., 2016, MNRAS, 463, 1929
  * Modi et al. (2018) Modi C., Feng Y., Seljak U., 2018, J. Cosmology Astropart. Phys., 2018, 028
  * Modi et al. (2019) Modi C., White M., Slosar A., Castorina E., 2019, arXiv e-prints, p. arXiv:1907.02330
  * Modi et al. (2020) Modi C., Lanusse F., Seljak U., 2020, arXiv e-prints, p. arXiv:2010.11847
  * Modi et al. (2022) Modi C., Li Y., Blei D., 2022, arXiv e-prints, p. arXiv:2206.15433
  * Morgan (2018) Morgan P., 2018, Probabilistic Programming in TensorFlow
  * Navarro et al. (1997) Navarro J. F., Frenk C. S., White S. D. M., 1997, ApJ, 490, 493
  * Neal et al. (2011) Neal R. M., et al., 2011, Handbook of markov chain monte carlo, 2, 2
  * Neistein et al. (2011) Neistein E., Weinmann S. M., Li C., Boylan-Kolchin M., 2011, MNRAS, 414, 1405
  * Peacock & Smith (2000) Peacock J. A., Smith R. E., 2000, MNRAS, 318, 1144
  * Peterson (1987) Peterson C., 1987, Complex systems, 1, 995
  * Press & Schechter (1974) Press W. H., Schechter P., 1974, ApJ, 187, 425
  * Robotham & Howlett (2018) Robotham A. S. G., Howlett C., 2018, Research Notes of the American Astronomical Society, 2, 55
  * Rodríguez-Torres et al. (2016) Rodríguez-Torres S. A., et al., 2016, MNRAS, 460, 1173
  * Schmidt et al. (2019) Schmidt F., Elsner F., Jasche J., Nguyen N. M., Lavaux G., 2019, J. Cosmology Astropart. Phys., 2019, 042
  * Scoccimarro et al. (2001) Scoccimarro R., Sheth R. K., Hui L., Jain B., 2001, ApJ, 546, 20
  * Seljak (2000) Seljak U., 2000, MNRAS, 318, 203
  * Seljak et al. (2017) Seljak U., Aslanyan G., Feng Y., Modi C., 2017, J. Cosmology Astropart. Phys., 2017, 009
  * Sinha et al. (2018) Sinha M., Berlind A. A., McBride C. K., Scoccimarro R., Piscionere J. A., Wibking B. D., 2018, MNRAS, 478, 1042
  * Spergel et al. (2015) Spergel D., et al., 2015, arXiv e-prints, p. arXiv:1503.03757
  * Takada et al. (2014) Takada M., et al., 2014, Publications of the Astronomical Society of Japan, 66, R1
  * Tamura et al. (2016) Tamura N., et al., 2016, in Evans C. J., Simard L., Takami H., eds, Society of Photo-Optical Instrumentation Engineers (SPIE) Conference Series Vol. 9908, Ground-based and Airborne Instrumentation for Astronomy VI. p. 99081M (arXiv:1608.01075), doi:10.1117/12.2232103
  * To et al. (2021) To C., et al., 2021, Phys. Rev. Lett., 126, 141301
  * Tran et al. (2016) Tran D., Kucukelbir A., Dieng A. B., Rudolph M., Liang D., Blei D. M., 2016, arXiv e-prints, p. arXiv:1610.09787
  * Wang et al. (2022) Wang Y., et al., 2022, The Astrophysical Journal, 928, 1
  * Wechsler & Tinker (2018) Wechsler R. H., Tinker J. L., 2018, ARA&A, 56, 435
  * Wechsler et al. (2006) Wechsler R. H., Zentner A. R., Bullock J. S., Kravtsov A. V., Allgood B., 2006, ApJ, 652, 71
  * White et al. (2011) White M., et al., 2011, ApJ, 728, 126
  * Wibking et al. (2019) Wibking B. D., et al., 2019, MNRAS, 484, 989
  * Wibking et al. (2020) Wibking B. D., Weinberg D. H., Salcedo A. N., Wu H.-Y., Singh S., Rodríguez-Torres S., Garrison L. H., Eisenstein D. J., 2020, MNRAS, 492, 2872
  * Yuan et al. (2018) Yuan S., Eisenstein D. J., Garrison L. H., 2018, MNRAS, 478, 2019
  * Zehavi et al. (2005) Zehavi I., et al., 2005, ApJ, 630, 1
  * Zehavi et al. (2011) Zehavi I., et al., 2011, ApJ, 736, 59
  * Zheng et al. (2005) Zheng Z., et al., 2005, ApJ, 633, 791
  * Zheng et al. (2007) Zheng Z., Coil A. L., Zehavi I., 2007, ApJ, 667, 760
  * Zheng et al. (2009) Zheng Z., Zehavi I., Eisenstein D. J., Weinberg D. H., Jing Y. P., 2009, ApJ, 707, 554
  * Zhu et al. (2006) Zhu G., Zheng Z., Lin W. P., Jing Y. P., Kang X., Gao L., 2006, ApJ, 639, L5
  * van den Bosch et al. (2003) van den Bosch F. C., Yang X., Mo H. J., 2003, MNRAS, 340, 771

##  Appendix A Distributional Properties of DiffHOD Model

Figure 6: Histogram showing the halo occupancy distribution from 100
independent samples from the complete halo catalog. We compare halos populated
by our DiffHODmodel as well as a standard HOD implementation. We show the
median satellite occupation number, and quote error bars representing the 32%
and 68% quartiles. We find excellent quantitative and qualitative agreement
between the two distributions.

In the main body of the work we sampled from a Bernoulli Distribution rather
than a pure Poisson Distribution due to the existing analytical tools to relax
the Bernoulli Distribution via the Gumbel-Softmax trick. This was demonstrated
to be a valid approximation at the level of various summary statistics, such
as halo occupancy distribution functions and resulting galaxy power spectrum.
In this section we show the halo occupancy distribution as a function of halo
mass.

We sample our satellite DiffHOD, at $\tau=0.1$, and the standard satellite HOD
using the parameters in the main text 100 times in order to attain reasonable
number statistics at the high mass bins. We show our results in 6, finding
quantitative good agreement between the models as calculated by their modal
and variance properties. We calculate for each mass bin the 32%, 50% and 68%
percentiles. Since DiffHOD uses a relaxed distribution instead of sampling, it
is possible to get non-integer number of satellites while for the standard HOD
model all sampling is discretized. Qualitatively we see a slight broadening of
the distributions at the extreme high mass end, however this does not
noticeably impact any of our resulting analysis due to the very small
population of these extreme high mass halos. Additional optimizations in terms
of maximum satellite population and choice of temperature could be performed
if this populations is of high interest.