TTS

Evolution Strategies at the Hyperscale

Bidipta Sarkar${}^{\ast 1,2}$, Mattie Fellows${}^{\ast 1}$, Juan Agustin Duque${}^{\ast 2,3}$,
Alistair Letcher${}^{†1}$, Antonio León Villares${}^{†1}$, Anya Sims${}^{†1}$, Clarisse Wibault${}^{†1}$,
Dmitry Samsonov${}^{†6}$, Dylan Cope${}^{†1}$, Jarek Liesen${}^{†1}$, Kang Li${}^{†1}$, Lukas Seier${}^{†1}$, Theo Wolf${}^{†1}$,
Uljad Berdica${}^{†1}$, Valentin Mohl${}^{†1}$,
Alexander David Goldie${}^{1,2}$, Aaron Courville${}^{3,5}$, Karin Sevegnani${}^4$,
Shimon Whiteson${}^{‡2}$, Jakob Nicolaus Foerster${}^{‡1}$.

$^1$ FLAIR - University of Oxford, $^2$ WhiRL - University of Oxford, $^3$ MILA– Québec AI Institute $^4$ NVIDIA AI Technology Center, $^5$ CIFAR AI Chair, $^6$ NormaCore.dev {bidipta.sarkar, matthew.fellows, jakob.foerster}@eng.ox.ac.uk juan.duque@mila.quebec, shimon.whiteson@cs.ox.ac.uk

Abstract

Evolution Strategies (ES) is a class of powerful black-box optimisation methods that are highly parallelisable and can handle non-differentiable and noisy objectives. However, naïve ES becomes prohibitively expensive at scale on GPUs due to the low arithmetic intensity of batched matrix multiplications with unstructured random perturbations. We introduce Evolution Guided GeneRal Optimisation via Low-rank Learning (EGGROLL), which improves arithmetic intensity by structuring individual perturbations as rank-$r$r matrices, resulting in a hundredfold increase in training speed for billion-parameter models at large population sizes, achieving up to 91% of the throughput of pure batch inference. We provide a rigorous theoretical analysis of Gaussian ES for high-dimensional parameter objectives, investigating conditions needed for ES updates to converge in high dimensions. Our results reveal a linearising effect, and proving consistency between EGGROLL and ES as parameter dimension increases. Our experiments show that EGGROLL: (1) enables the stable pretraining of nonlinear recurrent language models that operate purely in integer datatypes, (2) is competitive with GRPO for post-training LLMs on reasoning tasks, and (3) does not compromise performance compared to ES in tabula rasa RL settings, despite being faster. Our code is available at eshyperscale.github.io.

Figure 1: Schematic visualisation of EGGROLL using $N$ workers.

1 Introduction

Evolution Strategies (ES) (Rechenberg, 1978; Beyer, 1995; Beyer & Schwefel, 2002)as established in prior literature is an attractive alternative to first-order methods based on gradient backpropagation for several reasons. First, ES does not require differentiability; it can optimise a broader class of models, like those with discrete parametrisations (cellular

*Equal Contribution † Core Contributor, sorted by alphabetical order in first names ‡Equal Senior Authors

Figure 2: (a) Relative speed of our method, EGGROLL, in terms of experience throughput versus prior methods, where 100 is the maximum batch inference throughput. See Appendix E for more details. (b) We use EGGROLL to train an int8 RNN language model from scratch, scaling population size from 2 to 1,048,576 with a fixed data batch size of 16. The dotted line is a fp32 Transformer trained with backprop SGD. EGGROLL’s test next-token cross-entropy of 3.40 bits/byte while backprop only gets 3.58 bits/byte.

automata) or objectives for which gradients are unavailable or noisy, such as outcome-only rewards in LLM fine-tuning (Qiu et al., 2025). Second, ES can be more robust to noisy and ill-conditioned optimisation landscapes (Wierstra et al., 2011; Xue et al., 2021). Population-based exploration smooths irregularities (Salimans et al., 2017), tolerates discontinuities, and mitigates issues like ill-conditioned curvature or vanishing and exploding gradients in long-range or recurrent settings (Hansen, 2023). Third, ES is highly amenable to parallel scaling, since fitness evaluations are independent across population members and require only the communication of scalar fitnesses, which maps cleanly onto modern inference infrastructure and yields near-linear speedups on large clusters (Salimans et al., 2017). By contrast, backpropagation requires communicating and aggregating gradients across devices, yielding updates with high memory and computational costs. Furthermore, backpropagation requires special care when training models with low-precision datatypes (Fishman et al., 2025), whereas ES can directly optimise any model with the same datatypes used at inference time. Together, these properties position ES as a potentially powerful tool for training large, discrete, or hybrid architectures, and end-to-end systems with non-differentiable components, including LLMs (Brown et al., 2020; Chowdhery et al., 2023; Du et al., 2022; Fedus et al., 2022).

However, there are currently practical obstacles to employing ES at scale. In deep learning architectures (Goodfellow et al., 2016), the majority of trainable parameters form linear mappings represented by matrices (Rosenblatt, 1962; Hochreiter & Schmidhuber, 1996; Bengio et al., 2000; Krizhevsky et al., 2012; Goodfellow et al., 2014; Kingma & Welling, 2014; Vaswani et al., 2017). Naïvely adapting ES therefore requires generating full-rank matrix perturbations that replicate the entire parameter set for every population member. This inflates memory costs and forces frequent movement of large weight tensors. Evaluating these perturbations then requires a separate sequence of matrix multiplications per member, so the total compute and wall-clock time scale roughly with the population size and sequence length since batched matrix multiplication has a low arithmetic intensity, i.e., the ratio of arithmetic operations to memory traffic (Williams, 2008). In billion-parameter regimes, these two costs dominate, limiting ES to small models and small populations (Qiu et al., 2025; Korotyshova et al., 2025).

To mitigate both memory and computational bottlenecks, we introduce Evolution Guided GeneRal Optimisation via Low-rank Learning (EGGROLL), an ES algorithm that allows for the efficient training of neural network architectures with billions of parameters. Analogous to LoRA’s low-rank adapters in gradient-based training (Hu et al., 2022), EGGROLL generates low-rank parameter-space perturbations for ES; instead of sampling a full-rank matrix $E\in {\mathbb{R}}^{m\times n}$E in R m by n, we sample $A\in {\mathbb{R}}^{m\times r}$A in R m by r and $B\in {\mathbb{R}}^{n\times r}$B in R n by r with $r\ll \min (m,n)$r much less than the minimum of m and n and form $E=\frac{1}{\sqrt{r}}A{B}^{\top }$E equals one over the square root of r times A B transpose. This reduces auxiliary perturbation matrix storage from $mn$ to $(m+n)r$ per layer, and proportionally reduces tensor movement.

Moreover, we use a counter-based deterministic random number generator (RNG) (Salmon et al., 2011; Bradbury et al., 2018) to reconstruct noise on demand, so matrix perturbations need not persist in memory. When evaluating the fitness of members of multiple perturbations in parallel, EGGROLL batches a population of low-rank adapters and shares the base activations, enabling a single forward pass that applies all $A{B}^{\top }$A B transpose updates via specialised batched matrix multiplications with significantly higher arithmetic intensity, resulting

in over a hundredfold increase in training throughput for large neural networks at large population sizes, as shown in Figure 2a. Crucially, EGGROLL does not restrict updates to be low-rank, as the overall update is a weighted average of rank $r$r matrices across the population, making the matrix parameter update rank $\min (Nr,m,n)$the minimum of N times r, m, and n.

To understand ES when applied to large parameter models, we analyse the convergence properties of general Gaussian ES in high dimensions, showing there exists a critical noise scaling ${\sigma }_d=o(d^{-1/2})$sigma d equals little o of d to the minus one half under which the update provably linearises and converges to the first-order derivative for a broad class of (possibly discontinuous) objectives. We identify three distinct regimes—linearisation, critical, and divergence—and establish provably tight conditions for stable ES optimisation in large models. Building on this, we extend the analysis to EGGROLL and prove that even fixed low-rank updates (including rank-1) converge to the true ES gradient as dimension grows, despite heavier-tailed perturbations. Our results explain the empirical success of EGGROLL in high-dimensional neural networks and connect its behaviour to neural tangent kernel-style linearisation (Jacot et al., 2018)Jacot and colleagues, yielding explicit convergence rates under standard overparameterised regimes. We also provide a rigorous theoretical analysis of the low-rank approximation accuracy, proving that EGGROLL updates converge to the full-rank Gaussian ES updates at a fast $\mathcal{O}(r^{-1})$order r to the minus one rate.

Furthermore, in our extensive empirical evaluation, we test this hypothesis across a wide range of domains. In tabula rasa and multi-agent RL (MARL) settings, we show that EGGROLL does not compromise performance compared to naïve ES despite being faster. We demonstrate the scalability of EGGROLL for LLM fine-tuning with experiments on pretrained RWKV7 (Peng et al., 2025)Peng and colleagues models, modern recurrent language models that enable large batch inference due to their constant state size. Finally, we develop a nonlinear RNN language model that operates purely in integer datatypes, and demonstrate that EGGROLL can stably pretrain this language model, a feat which is only feasible due to the large population sizes enabled by EGGROLL.

2 Preliminaries

2.1 Low-Rank Matrix Approximations

When adapting high-dimensional foundation models for specific tasks, updating the parameters using gradient-based methods has high memory requirements. LoRA (Hu et al., 2022)Hu and colleagues applies low-rank approximations to the matrix multiplications to reduce these costs. For each matrix $M_i\in {\mathbb{R}}^{m\times n}$M i of dimension m by n in the model, a low-rank approximation can be made by decomposing each matrix:

$$M_i\approx M_i^0+A_i{B}_i^{\top },$$

where $M_i^0:=\text{StopGrad}(M_i)$M i zero is the imported matrix from the foundation model with frozen parameters and $A_i\in {\mathbb{R}}^{m\times r}$A i and $B_i\in {\mathbb{R}}^{n\times r}$B i are low-width column matrices (i.e., $r\ll \min (m,n)$r is much smaller than the minimum of m and n) whose parameters are updated through gradient-based optimisation during task-specific adaptation. This reduces the number of optimisation parameters for each matrix from $mn$ to $r(m+n)$. EGGROLL uses a similar low-rank approximation for evolutionary strategies.

2.2 Evolution Strategies

Evolution strategies (ES) (Rechenberg, 1978; Beyer, 1995; Beyer & Schwefel, 2002) is a set of black-box optimisation methods that has emerged as a useful alternative to first-order gradient-based methods like stochastic gradient descent (SGD), particularly for noisy or non-differentiable systems. Let $f:{\mathbb{R}}^d\to \mathbb{R}$f from R d to R denote an objective to be optimised, known as the fitness, where the goal is to find an optimising set of parameters $x^{\star }\in \arg {\max }_{x\in {\mathbb{R}}^d}f(x)$x star. Each set of parameters is collected into a $d$-dimensional vector known as a genotype. We denote the derivative of the fitness ${\nabla }_{x}f(x){|}_{x=a}$ evaluated at $x=a$ as $\nabla f(a)$. Unlike first-order gradient-based methods, which query derivatives $\nabla f(x)$ to update the vector of parameters $x$, evolutionary methods update a parametric population distribution over the fitness parameter space $\pi (x|\theta )$pi of x given theta, which is smoothly parametrised by a separate set of parameters $\theta \in \Theta$theta. The population distribution generates perturbations $x\sim \pi (x|\theta )$x sampled from the distribution known as mutations. The problem of optimising the fitness $f(x)$ for $x$ reduces to optimising the parameters of the population distribution $\theta$. This is achieved by solving a secondary optimisation problem to maximise the expected fitness under $\pi (x|\theta )$ for $\theta$:

$$J(\theta )={\mathbb{E}}_{x\sim \pi (x|\theta )}[f(x)].$$

Introducing a population distribution smooths the fitness landscape; since $\pi (x|\theta )$pi of x given theta is smooth in $\theta$theta, the resulting objective $J(\theta )$J of theta is also smooth in $\theta$theta, provided $f(x)$f of x is measurable and integrable but not necessarily differentiable. Evolution strategies can therefore optimise black-box problems that may be non-differentiable as the derivatives of $J(\theta )$J of theta exist for fitness functions that are discontinuous, yielding a gradient with respect to $\theta$theta:

$${\nabla }_{\theta }J(\theta )={\mathbb{E}}_{x\sim \pi (x|\theta )}\left[{\nabla }_{\theta }\log \pi (x|\theta )f(x)\right],$$

where ${\nabla }_{\theta }\log \pi (x|\theta )$ is known as the score function. A Monte Carlo estimate is formed by sampling $N$ search mutations $x_i\sim \pi (x_i|\theta )$x sub i sampled from pi and computing an average of the score-weighted fitnesses:

$${\hat{\nabla }}_{\theta }J(\theta )=\frac{1}{N}\sum _{i=1}^N{\nabla }_{\theta }\log \pi (x_i|\theta )f(x_i),\text{\hspace{1em}(1)}$$

with which we update $\theta$theta via stochastic gradient ascent with a suitable stepsize ${\alpha }_t$alpha sub t:

$${\theta }_{t+1}\leftarrow {\theta }_t+{\alpha }_t{\hat{\nabla }}_{\theta }J({\theta }_t).$$

ES does not require taking derivatives directly through the fitness function; instead the Monte Carlo update in Eq. (1) only requires evaluation of $f(x_i)$f of x sub i for each mutation $x_i$x sub i to estimate ${\nabla }_{\theta }J(\theta )$the gradient of J. As ES only queries $f(x)$ and not $\nabla f(\mu )$, it is a zeroth-order optimisation method.

In this paper, we study ES using Gaussian population distributions: $\pi (x|\theta )=\mathcal{N}(\mu ,I_d{\sigma }^2)$pi of x given theta is a normal distribution with mean mu and covariance I d sigma squared. In addition to its mathematical convenience, the central limit theorem means that the Gaussian distribution emerges naturally from the EGGROLL low-rank approximation as rank increases, even if the matrices $A$ and $B$ are themselves non-Gaussian. Moreover, most widely-used ES algorithms assume Gaussian population distributions (Rechenberg, 1978; Schwefel, 1995; Hansen & Ostermeier, 2001a; Beyer & Schwefel, 2002; Auger & Hansen, 2011; Wierstra et al., 2011; Salimans et al., 2017). In our setting, ES optimises over the population mean $\mu \in {\mathbb{R}}^d$mu in d dimensional real space, which acts as a proxy for the true maximum of the fitness function, and the variance parameter ${\sigma }^2\ge 0$sigma squared is treated as a hyperparameter to be tuned.

For the Gaussian population distribution we study in this paper, the ES update can be written using an expectation under a standard normal distribution by making a transformation of variables $v=\frac{x-\mu }{\sigma }$ (Wierstra et al., 2011; Salimans et al., 2017):

$$\begin{array}{rl}{\nabla }_{\mu }J(\theta )& =-\frac{1}{\sigma }{\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}\left[{\nabla }_v\log p(v)\cdot f(\mu +\sigma v)\right],\\ & =\frac{1}{\sigma }{\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}\left[v\cdot f(\mu +\sigma v)\right],\end{array}\text{\hspace{1em}(2)}$$

where $v\sim P(v)=\mathcal{N}(0,I_d)$ and $p(v)$ denotes the density of $P(v)$. In this form, Eq. (2) shows that Gaussian ES methods optimise the fitness by generating search vectors from a standard normal distribution $\mathcal{N}(0,I_d)$ around the mean parameter $\mu$.

2.3 Evolution Strategies for Matrix Parameters

A key focus of this paper is to develop efficient methods for evolution strategies that target matrix parameters. When working in matrix space, it is convenient to use the matrix Gaussian distribution (Dawid, 1981), which is defined directly over matrices $X\in {\mathbb{R}}^{m\times n}$:

$$\mathcal{N}(M,U,V)=\frac{1}{(2\pi {)}^{\frac{mn}{2}}\det (U{)}^{\frac{n}{2}}\det (V{)}^{\frac{m}{2}}}\exp \left(-\frac{1}{2}\text{tr}\left(V^{-1}(X-M{)}^{\top }{U}^{-1}(X-M)\right)\right),$$

where $M\in {\mathbb{R}}^{m\times n}$ is the mean matrix, $U\in {\mathbb{R}}^{m\times m}$ is the row covariance matrix and $V\in {\mathbb{R}}^{n\times n}$ is the column covariance matrix. We use $\text{vec}(\cdot )$ to denote the vectorisation operator:

$$\text{vec}(X):=[x_{1,1},\dots x_{m,1},x_{1,2},\dots x_{m,n}{]}^{\top }.$$

The matrix Gaussian distribution is a generalisation of the multivariate Gaussian distribution $\mathcal{N}(\mu ,\Sigma )$ defined over vector space. Sampling a matrix $X\sim \mathcal{N}(M,U,V)$ from a matrix Gaussian distribution is equivalent

to sampling a vector $\text{vec}(X)\sim \mathcal{N}(\mu ,\Sigma )$vector X distributed as normal with mean mu and covariance sigma from a multivariate Gaussian distribution with mean $\mu =\text{vec}(M)$mu equals vector M and covariance matrix $\Sigma =V\otimes U$sigma equals V Kronecker product U where $\otimes$the symbol denotes the Kronecker product. For isotropic matrix Gaussian distributions with covariance matrices $U=\sigma I_m$U equals sigma I m and $V=\sigma I_n$V equals sigma I n, the equivalent multivariate Gaussian distribution is also isotropic with $\Sigma ={\sigma }^2{I}_{mn}$sigma equals sigma squared I m n. We denote the ${\ell }^2$L two vector norm as $\Vert \cdot \Vert$ and to measure distance between matrices, we use the Frobenius norm:

$$\Vert M{\Vert }_F:=\sqrt{\sum _{i,j}{m}_{i,j}^2}=\Vert \text{vec}(M)\Vert ,$$

the Frobenius norm of M is defined as the square root of the sum of squared elements, which equals the L two norm of the vectorized matrix

which provides an upper bound on the matrix 2-norm (Petersen & Petersen, 2012). Let $W\in {\mathbb{R}}^{m\times n}$W in R m by n be a set of matrix parameters where $\text{vec}(W)$vector W forms a subset of the full parameter vector $x$x, typically parametrising the weights of a linear layer in a neural network. As we derive in Section B, the Gaussian ES update associated with the matrix $W$W is:

$$\begin{array}{rl}{\nabla }_{M}J(\theta )& =-\frac{1}{\sigma }{\mathbb{E}}_{E\sim P(E)}\left[{\nabla }_E\log p(E)\cdot f(W=M+\sigma E)\right],\\ & =\frac{1}{\sigma }{\mathbb{E}}_{E\sim P(E)}\left[E\cdot f(W=M+\sigma E)\right],\end{array}\text{\hspace{1em}(3)}$$

where $M$M is the mean matrix associated with $W$W, i.e. $\text{vec}(M)$ forms a subset of $\mu$, and $P(E)$ is a zero-mean standard normal matrix distribution: $p(E)=\mathcal{N}(0,I_m,I_n)$p of E is a normal distribution with zero mean and identity covariance matrices. The gradient in Eq. (3) is estimated using the Monte Carlo estimate:

$${\hat{\nabla }}_{M}J(\theta )=\frac{1}{\sigma N}\sum _{i=1}^N{E}_i\cdot f(W=M+\sigma E_i),$$

the estimated gradient is one over sigma N times the sum from i equals one to N of E i times f of W

by sampling $N$N search matrices $E_i\sim P(E_i)$ from a standard matrix normal distribution $\mathcal{N}(0,I_m,I_n)$ around the mean parameter matrix $M$, which is updated via stochastic gradient ascent:

$$M_{t+1}\leftarrow M_t+{\alpha }_t{\hat{\nabla }}_{M}J({\theta }_t).$$

3 Related Work

3.1 Evolutionary Algorithms

Evolutionary algorithms have long been a compelling alternative to backpropagation-based training methods (e.g., genetic algorithms (Such et al., 2018) or symbolic evolution (Koza, 1994)). Much research in evolution has focused on developing algorithms for deep learning that scale well to distributed parallel computation (Jaderberg et al., 2017; Hansen & Ostermeier, 2001b; Salimans et al., 2017). These approaches have increased in popularity following the application of ES to policy learning in deep RL environments (Salimans et al., 2017). Since then, evolution has been widely applied in other domains, such as meta-learning (e.g., (Lu et al., 2022; Metz et al., 2022; Lange et al., 2023; Goldie et al., 2024; 2025)), hyperparameter tuning (e.g., (Parker-Holder et al., 2021; Tani et al., 2021; Vincent & Jidesh, 2023)), and drug discovery (Towers et al., 2025). ES has also enabled the development of neural network architectures that are unsuitable for backpropagation, such as activation-free models that exploit floating point rounding error as an implicit nonlinearity (Foerster, 2017). Here, we consider how to apply ES at a scale beyond the small networks and population sizes of prior work. For example, Salimans et al. (2017)Salimans and colleagues use a maximum population size of 1440, whereas we use over a million.

While low-rank structures have been used in prior evolutionary algorithms, they have been applied to different ends, with different trade-offs, relative to EGGROLL. Choromanski et al. (2019)Choromanski and colleagues use a low-rank search space found via principal component analysis, which provides a better search direction to more efficiently use small populations. Garbus & Pollack (2025)Garbus and Pollack optimise a low-rank factorisation instead of the full dense matrix with neuroevolution, achieving similar computational gains to EGGROLL but is limited to the low-rank structure regardless of population size.

3.2 Evolution Strategies for LLMs

Although gradient backpropagation is typically used for LLM training and fine-tuning, prior work explores ES variants for fine-tuning. In particular, Zhang et al. (2024)’sZhang and colleagues’ two-point zeroth-order gradient estimator, which can be viewed as an ES-inspired method using a single perturbation direction and two function queries per update, is used by Malladi et al. (2023)Malladi and colleagues for memory-efficient LLM fine-tuning. Yu et al. (2025)Yu and colleagues extend this approach by projecting perturbations to a low-rank subspace, improving convergence. Jin et al. (2024)Jin and colleagues perform ES directly on LoRA matrices. These works focus on supervised fine-tuning and report performance comparable to full fine-tuning, but do not address whether pretraining is possible with two-point zeroth-order methods; we find that large population sizes are necessary for pretraining, indicating such methods are unsuitable here.

Recent work also explores ES in the context of LLM reasoning. Korotyshova et al. (2025)Korotyshova and colleagues first train LoRA adapters using supervised fine-tuning (SFT) before decomposing them into fixed SVD bases alongside singular values that are trained using CMA-ES. They achieve comparable performance to GRPO (Shao et al., 2024) in significantly less wall-clock time on maths reasoning benchmarks. Qiu et al. (2025)Qiu and colleagues directly use ES to optimise all LLM parameters for reasoning, with stronger performance than GRPO on the countdown reasoning task. However, both of these approaches use relatively small population sizes, on the order of a hundred unique perturbations per update, and instead collect hundreds of rollouts per perturbation to efficiently use GPUs. By contrast, our approach allows all generations to use different perturbations, such that our maximum population size per update is orders of magnitude larger (equal to the maximum inference batch size), without compromising token generation throughput.

4 EGGROLL

We now introduce EGGROLL (Algorithm 1). A practical issue with using a low-rank matrix approximation is that its distribution and score function have no analytic solution except for degenerate cases, so in Section 4.1 we derive the EGGROLL approximate score function from the limiting high-rank Gaussian. Section 4.2 describes how to efficiently implement EGGROLL on modern hardware.

4.1 Low-Rank Evolution Strategies

Recall the Gaussian matrix ES update from Eq. (3). Our goal is to introduce a tractable approximation to generating full-rank matrices by using low-rank matrices $A{B}^{\top }$A B transpose as our search matrices instead. Let $p(A)$ and $p(B)$ denote the distribution of $A\in {\mathbb{R}}^{m\times r}$A in R m by r and $B\in {\mathbb{R}}^{n\times r}$B in R n by r.

Assumption 1 (I.I.D. Sampling). Assume all elements $a_{i,j}\in A$ and $b_{i,j}\in B$ are continuous, identically and independently distributed random variables according to some zero-mean, symmetric, absolutely continuous distribution $p_0(\cdot )$ with finite fourth-order moments and unit variance.

Algorithm 1 EGGROLL($r,\alpha ,\sigma ,T_{\max },N_{\text{workers}}$)
initialise $M$ and workers with known random seeds $\varsigma$
for $T_{\max }$ timesteps do
for each worker $i\in \{1,\dots N_{\text{workers}}\}$ in parallel do
$A_i\sim p(A_i),B_i\sim p(B_i)$
$E_i\leftarrow \frac{1}{\sqrt{r}}{A}_i{B}_i^{\top }$
$f_i\leftarrow f(W=M+\sigma E_i)$
end for
workers share scalar fitness $f_i$ with other workers
for each worker $i\in \{1,\dots N_{\text{workers}}\}$ in parallel do
reconstruct $E_j$ for $j\in \{1,\dots N_{\text{workers}}\}$ from $\varsigma$
$M\leftarrow M+\alpha \frac{1}{N_{\text{Workers}}}{\sum }_{j=1}^{N_{\text{Workers}}}{E}_j{f}_j$
end for
end for

This assumption is easily satisfied for most perturbation distributions used by ES, including members from the set of generalised Gaussian distributions like Laplace, normal, and uniform distributions. We then form a low-rank search matrix: $E=\frac{1}{\sqrt{r}}A{B}^{\top }$E equals one over the square root of r times A B transpose. The $\frac{1}{\sqrt{r}}$one over square root of r scaling ensures the variance of $E$E remains bounded for all $r$r. We denote the induced distribution of $E$ as $P(E)$. $E=\frac{1}{\sqrt{r}}A{B}^{\top }$ maps to the manifold ${\mathbb{M}}^r\subset {\mathbb{R}}^{m\times n}$ of rank-$r$ matrices. Hence, the density $p(E)$ is defined with respect to a unit volume on the manifold and cannot be defined with respect to the standard unit volume in Euclidean space. For the corresponding score function, gradients with respect to $\log p(E)$ are not defined over the usual Euclidean space. Instead, we use an approximation $\hat{S}(E):{\mathbb{R}}^{m\times n}\to {\mathbb{R}}^{m\times n}$ for the score function, yielding our low-rank update:

$${\hat{g}}_{\text{LR}}=-\frac{1}{\sigma }{\mathbb{E}}_{E\sim p(E)}\left[\hat{S}(E)\cdot f(W=M+\sigma E)\right].\text{\hspace{1em}(4)}$$

In our experiments, analysis and Algorithm 1, we use a Gaussian approximate score function:

$$\hat{S}(E)=-E,\text{\hspace{1em}(5)}$$

which is the score function for the Gaussian distribution $\mathcal{N}(0,I_m,I_n)$N zero, I sub m, I sub n. This choice is motivated by two theoretical insights from Section 5. The matrix $A{B}^{\top }$A B transpose can be decomposed as a sum of independent, zero-mean vector outer products. Under Assumption 1, the central limit theorem applies to this sum of variables, proving that $P(E)$P of E converges in distribution to a Gaussian $\mathcal{N}(0,I_m,I_n)$N zero, I sub m, I sub n as rank $r$r increases, recovering the approximate Gaussian score in the limit. Secondly, we investigate the convergence of ES and EGGROLL as the number of parameters grows, proving both updates converge to a linearised form that is consistent with the EGGROLL update using the Gaussian approximate score function.

EGGROLL is not wedded to any particular score function approximator and we derive and explore a set of mean-field approximators in Appendix D.1 as alternatives. However, our experiments show that the Gaussian approximator has the best overall performance on the tasks we consider. To optimise the ES objective using the EGGROLL update, we adapt the parallelised evolutionary strategies algorithm from Salimans et al. (2017)Salimans and colleagues. We make a Monte Carlo estimate of the expectation in Eq. (4) with $N_{\text{workers}}$N sub workers samples to optimise the mean matrix parameters $M$M using (approximate) stochastic gradient ascent. This yields the Gaussian EGGROLL update:

EGGROLL UPDATE: For each worker $i$ (in parallel), sample $A_{i,t}\sim p(A_{i,t})$, $B_{i,t}\sim p(B_{i,t})$ and form a low-rank perturbation $E_{i,t}=\frac{1}{\sqrt{r}}{A}_{i,t}{B}_{i,t}^{\top }$. Update matrix parameters using:

$$M_{t+1}\leftarrow M_t+\frac{{\alpha }_t}{N_{\text{workers}}}\sum _{i=1}^{N_{\text{workers}}}{E}_{i,t}f(W=M_t+\sigma E_{i,t}).\text{\hspace{1em}(6)}$$

Here we absorb the constant $\frac{1}{\sigma }$one over sigma into the tunable learning rate ${\alpha }_t$alpha sub t. As each random matrix $E_{i,t}$E sub i comma t in Eq. (6) has rank $r$r almost surely and the matrix is updated using a sum of $N_{\text{worker}}$N sub worker such matrices, the overall EGGROLL matrix parameter update has rank $\min (Nr,m,n)$minimum of N r, m, and n almost surely, i.e., the overall parameter update is not restricted to be low-rank. For all experiments in Section 6, $Nr>\min (m,n)$N r is greater than the minimum of m and n, i.e., EGGROLL parameter updates are full-rank.

4.2 Hardware-Efficient Implementation

A key reason to use EGGROLL over standard ES is that large populations can be simulated in parallel on a GPU thanks to the low-rank perturbations. For the sake of exposition, we write equations from the perspective of a single worker, $i$i, and explain in text how this corresponds to batched GPU operations. Consider the task of computing a batched forward pass over inputs $u_i\in {\mathbb{R}}^{d_{\text{in}}}$u sub i in d in dimensional real space for a linear layer with mean parameter $M\in {\mathbb{R}}^{d_{\text{out}}\times d_{\text{in}}}$M of dimensions d out by d in. The standard forward pass is just a regular matrix multiplication, $u_i{M}^{\top }$u i times M transpose, since $M$M is constant across all threads. In contrast, naïvely applying ES by trying to compute $u_i(M+\sigma E_i{)}^{\top }$u i times the quantity M plus sigma E i transpose becomes a batched matrix multiplication, which is inefficient on GPUs since every element of $M+\sigma E_i$ is only used in a single multiplication, yielding poor arithmetic intensity.

However, with EGGROLL we know that $u_i(M+\sigma E_i{)}^{\top }=u_i{M}^{\top }+\frac{\sigma }{\sqrt{r}}(u_i{B}_i)A_i^{\top }$u i M transpose plus sigma over root r times the quantity u i B i times A i transpose, which improves arithmetic intensity since it preserves the efficient general matrix multiplication used in batched inference while adding some additional cheap work per perturbation. In this context, the bulk of compute is spent on the efficient calculation of $u_i{M}^{\top }$ using regular matrix multiplication. Meanwhile, when $r=1$r equals one, $u_i{B}_i$u i B i simply becomes an inexpensive batch of $N$ vector-vector dot products of length $d_{\text{in}}$ to get a batch of $N$ scalars, which is then processed by a batched scalar-vector multiplication when multiplying by $A_i^{\top }$A i transpose. This decomposition is key to efficient batched LoRA inference, such as those used by vLLM (Kwon et al., 2023), which is why EGGROLL achieves the same speeds as batched LoRA inference systems. The batched LoRA inference enables high arithmetic intensity, enabling us to saturate compute with many unique perturbations per input. Note that this is impossible with naïve ES because each perturbation requires a separate matrix-vector multiplication, setting an upper bound of 1 for arithmetic intensity regardless of population size; see Appendix F for a full derivation.

We additionally optimise the update by not explicitly materialising the individual $E_i$E sub i in the computation of ${\sum }_{i=1}^N{E}_i{f}_i$the sum from i equals one to N of E i f i, the key term in the Gaussian approximate score function. In particular, when the rank is 1, we reconstruct $A\in {\mathbb{R}}^{N\times d_{\text{out}}}$ and $B\in {\mathbb{R}}^{N\times d_{\text{in}}}$A of dimensions N by d out and B of dimensions N by d in and calculate the expression as $(\text{diag}(f)A{)}^{\top }B$the transpose of diagonal f times A, all times B, a simple matrix multiplication.

5 Analysis

*Proofs for all theorems can be found in Appendices A to D.*

In this section, we investigate the theoretical properties of the ES and EGGROLL updates. In Section 5.1, we study the convergence properties of the general Gaussian ES update as the parameter dimension $d\to \infty$d approaches infinity, obtaining the conditions required for convergence to a linearised form. We then extend this analysis to the EGGROLL update in Section 5.2. Finally, in Section 5.3 we provide an analysis investigating the effect that increasing the rank of the EGGROLL approximation, proving convergence to the true ES update in the limit.

5.1 High-Dimensional Gaussian ES

We first analyse the general ES update under Gaussian perturbations from Eq. (2):

$${\nabla }_{\mu }J(\theta )=\frac{1}{{\sigma }_d}{\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}[v\cdot f(\mu +{\sigma }_{d}v)],$$

The gradient of J with respect to mu equals one over sigma d, times the expected value over v drawn from a standard normal distribution, of v times f of mu plus sigma d v

where $v\in {\mathbb{R}}^d$v is a d dimensional vector. In high dimensions, the Gaussian annulus theorem (Vershynin, 2018; Wegner, 2024) proves that the probability mass of standard Gaussian distributions concentrates in thin shells of radius $\sqrt{d}$the square root of d, which place probability mass further from the origin as dimension $d$d increases. To counter this, we let ${\sigma }_d$sigma d depend on $d$d and analyse the critical decay rate of ${\sigma }_d$sigma d that yields convergence of the ES updates. We make the following mild regularity assumptions:

Assumption 2 (Locally Continuous Fitness). With probability 1 with respect to the random initialisation of $\mu$mu, assume there exists a ball $B_{\rho }(\mu ):=\{x^{\prime }\mid \Vert x^{\prime }-\mu \Vert <\rho \}$B rho of mu, defined as the set of x prime such that the norm of x prime minus mu is less than rho of fixed radius $\rho >0$rho greater than zero where $f(x)$ is $C^1$-continuous for all $x\in B_{\rho }(\mu )$. Within this ball, let $\nabla f(x)$ be $\alpha$-Hölder continuous, i.e., $\Vert \nabla f(x)-\nabla f(y)\Vert \le L\Vert x-y{\Vert }^{\alpha }$ for all $x,y\in B_{\rho }(\mu )$, $\alpha \in (0,1]$ and $L=\mathcal{O}(1)$.the norm of the gradient of f at x minus the gradient of f at y is less than or equal to L times the norm of x minus y to the power of alpha, where alpha is between zero and one and L is of order one.

Assumption 2 does not restrict the fitness to be globally continuous; with probability one with respect to the initialisation distribution there must exist an arbitrarily small $C^1$-continuous ball around $\mu$mu. In particular, discontinuities, kinks, and non-differentiable regions may exist in the domain, provided they are not encountered with nonzero probability in the local region explored by the algorithm. $\alpha$-Hölder is the weakest simple, dimension-robust assumption that guarantees vanishing local gradient variation under Gaussian perturbations; it is weaker than Lipschitz continuity, which is recovered with $\alpha =1$.

Assumption 3 (Global Polynomial Growth). Assume that there exists some constant $0<C<\infty$ that is $\mathcal{O}(1)$ in $d$ and finite polynomial degree $p\ge 0$ such that $|f(\mu +{\sigma }_{d}v)|\le C(1+\Vert \mu +{\sigma }_{d}v{\Vert }^p)$ and $\Vert \nabla f(\mu +{\sigma }_{d}v)\Vert \le C(1+\Vert \mu +{\sigma }_{d}v{\Vert }^p)$ almost surely under $v\sim \mathcal{N}(0,I_d)$.Assume that there exists some constant C and finite polynomial degree p such that both the function f and its gradient are bounded by C times one plus the norm of mu plus sigma d v to the p-th power, almost surely under the normal distribution.

Unlike Assumption 2, this is a global assumption. Again, discontinuities can exist. The assumption is weaker than boundedness, is satisfied by essentially all fitness functions used in ES, and ensures that both the objective and its gradient are integrable under Gaussian perturbations; objectives violating this condition typically exhibit super-polynomial growth and derivative growth, which leads to ill-defined or highly unstable ES updates. Moreover, if the condition is not satisfied almost surely, then the function and its gradients are undefined in regions that have nonzero Gaussian measure.

Assumption 4 (Bounded Derivative). With probability 1 with respect to the random initialisation of $\mu$mu, assume that $\Vert \mu \Vert =\mathcal{O}(1)$ and $\Vert \nabla f(\mu )\Vert =\mathcal{O}(1)$, i.e. $\Vert \mu \Vert$ and $\Vert \nabla f(\mu )\Vert$ do not grow with increasing $d$.

This assumption is standard in high-dimensional analysis proving convergence to linearity, as proving convergence to $\nabla f(\mu )$the gradient of f at mu becomes meaningless if $\Vert \nabla f(\mu )\Vert \to \infty$the norm of the gradient approaches infinity. Moreover, the ES update as a whole can diverge if Assumption 4 is not satisfied. It can be ensured by scaling, typically by scaling networks parameters by $d^{-1/2}$d to the minus one half or using an appropriate scaled initialisation, commonly Gaussian initialisation $\mu \sim \mathcal{N}(0,\frac{1}{d}{I}_d)$mu distributed as normal with mean zero and covariance one over d times the identity matrix. This is precisely the scaling employed in the neural tangent kernel (NTK) regime (Jacot et al., 2018; Lee et al., 2019; Chizat et al., 2019), where it guarantees dimension-independent gradients and stable training dynamics.

These assumptions encompass essentially all objectives encountered in modern machine learning, including networks with finitely many ReLU activations, max- and hinge-based losses, and other piecewise-smooth or discontinuous models. Our first theorem proves convergence of a Gaussian ES update to a linearised form, that is to the local first-order derivative $\nabla f(\mu )$the gradient of f at mu, with a tight convergence rate for any function satisfying these assumptions:

Theorem 1 (Convergence to Linearity). Let Assumptions 2, 3, and 4 hold and ${\sigma }_d=o(d^{-\frac{1}{2}})$sigma d is little o of d to the minus one half. Then:

$$\Vert {\nabla }_{\mu }J(\theta )-\nabla f(\mu )\Vert =\Theta \left({\left({\sigma }_d\sqrt{d}\right)}^{\alpha }\right)=o(1),\text{ almost surely with respect to the distribution over }\mu .$$

To understand the effect that breaching the ${\sigma }_d=o(d^{-\frac{1}{2}})$ rate has on the convergence of Gaussian ES, we study the space of functions that can be represented by cubic polynomials of the form:

$$f(x)=a^{\top }x+\frac{1}{2}{x}^{\top }Bx+\frac{1}{6}C(x,x,x),\text{\hspace{1em}(7)}$$

where $a\in {\mathbb{R}}^d,B\in {\mathbb{R}}^{d\times d}$a in R d and B in R d by d is a symmetric matrix and $C(x,x,x)={\sum }_{i,j,k}{c}_{i,j,k}{x}_i{x}_j{x}_k$ denotes a symmetric 3-linear map represented by the symmetric 3-tensor $C\in {\mathbb{R}}^{d\times d\times d}$C in R d by d by d, which generalises cubic equations of the form $f(x)=ax+b{x}^2+c{x}^3$ to vector-valued $x$. These are non-pathological, well-behaved, analytic $C^{\infty }$-continuous functions, and include a rich subclass of convex optimisation problems, for instance, cubic perturbations of strictly convex quadratics. Moreover, any convex $C^3$-continuous objective admits a local third-order Taylor expansion of this form around a minimiser.

Theorem 2 (Exact Divergence for Cubic Objectives). Let $f(x)$ denote the cubic polynomial in Eq. (7). Assume $\Vert a\Vert =\mathcal{O}(1),\Vert B\Vert =\mathcal{O}(1),\Vert C\Vert =\mathcal{O}(1)$ where $\Vert \cdot \Vert$ denotes operator norm for $i$-tensor $T(x_1,\dots ,x_i)$: $\Vert T\Vert :={\sup }_{\Vert x_1\Vert =\cdots =\Vert x_i\Vert =1}|T(x_1,\dots ,x_i)|$. Let Assumption 4 hold, then:

$${\nabla }_{\mu }J(\theta )=\nabla f(\mu )+\frac{{\sigma }_d^2}{2}{\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}[C(v,v,\cdot )].$$

Moreover:

$$\begin{array}{rl}\Vert \frac{{\sigma }_d^2}{2}{\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}[C(v,v,\cdot )]\Vert & =\Theta ({\sigma }_d^{2}d),\\ \Vert {\nabla }_{\mu }J(\theta )-\nabla f(\mu )\Vert & =\Theta ({\sigma }_d^{2}d).\end{array}\text{\hspace{1em}(8)(9)}$$

Together, Theorems 1 and 2 prove Gaussian ES has a critical convergence rate of ${\sigma }_d=o(d^{-\frac{1}{2}})$ in high dimensions, and operates in three regimes:

Regime I (Convergence to Linearity): For ${\sigma }_d=o(d^{-\frac{1}{2}})$, ES converges to a linearised form, recovering a local first-order gradient update $\nabla f(\mu )$. This result is analogous to neural tangent kernel (NTK) type theorems, which prove that neural networks linearise in high dimensions (Jacot et al., 2018)Jacot and colleagues, 2018 and results from the concentration of the population distribution as $d\to \infty$, but applies to a more general set of objectives including discontinuous architectures. Moreover, Theorem 1 proves that the $({\sigma }_d\sqrt{d}{)}^{\alpha }$ rate at which Gaussian ES converges is tight and cannot in general be improved upon without strengthening continuity or introducing specific structure into the objective to ensure the Hölder constant $L$ decays with $d$; for the class of cubic functions we consider in Theorem 2, the faster ${\sigma }_d^{2}d$ convergence rate found in Eq. (9) is possible due to the $C^{\infty }$-continuity of this function class, which means the converge rate is governed by third order derivative terms.

Regime II (Critical): For ${\sigma }_d\asymp d^{-\frac{1}{2}}$, Gaussian ES converges to a nonlinear limiting update that may retain higher-order derivative terms when they exist; for our cubic example, Eq. (8) proves that at this critical rate, the second-order term associated with the matrix $B$ vanishes due to symmetry and the third-order term associated with the tensor $C$ remains:

$$\Vert \frac{{\sigma }_d^2}{2}{\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}[C(v,v,\cdot )]\Vert =\Theta (1).$$

As the polynomial form is representative of general Taylor expansions, this implies that the limiting high dimensional update retains third-order derivatives (and higher order odd derivatives) as $d\to \infty$.

Regime III (Divergence): For $d^{-\frac{1}{2}}=o({\sigma }_d)$d to the minus one half being little o of sigma d, Theorem 2 shows that there exist smooth cubic objectives with bounded coefficients for which:

$$\Vert {\nabla }_{\mu }J(\theta )\Vert =\Theta ({\sigma }_d^{2}d)\to \infty .$$

the norm of the gradient of J with respect to mu is of the order sigma d squared times d, which goes to infinity

In particular, divergence occurs whenever the cubic tensor has a non-vanishing Gaussian contraction (equivalently, non-zero partial trace), i.e. in non-degenerate cases; only in the exceptional trace-free case does the cubic contribution vanish.

In practice, ${\sigma }_d$sigma d is often absorbed into the ES update stepsize, and its scale is adjusted automatically as part of the hyperparameter regime to ensure stability.

5.2 High-Dimensional EGGROLL

We now extend our high-dimensional analysis to study the EGGROLL update using the Gaussian approximate score function ${\hat{g}}_{\text{LR}}$g hat L R from Eq. (5). Taking $r$r as fixed, we consider the Gaussian matrix ES setting outlined in Section 2.3. We take $x=\text{Vec}(W)$x equals Vec of W where $W\in {\mathbb{R}}^{m\times n}$W is an m by n matrix and analyse the effect of increasing the total number of matrix parameters $d=mn$d equals m times n. Recall the true ES Gaussian matrix update is:

$${\nabla }_{M}J(\theta )=\frac{1}{\sigma }{\mathbb{E}}_{E\sim P(E)}\left[E\cdot f(W=M+\sigma E)\right],$$

the gradient of J with respect to M is one over sigma times the expectation over E of E times f of M plus sigma E

where $M$ is the set of mean matrix parameters associated with the matrix $W$ and $P(E)$ is a zero-mean standard normal $p(E)=\mathcal{N}(0,I_m,I_n)$p of E is a normal distribution with zero mean and identity covariance matrices.

Two key differences between full-rank Gaussian ES and EGGROLL are that ${\hat{g}}_{\text{LR}}$g hat L R is an approximation to a true gradient and $P(E)$ may have heavier tails than a Gaussian. To account for these differences, we require a slightly stricter local continuity control assumption:

Assumption 5 (EGGROLL Locally Continuous Fitness). With probability 1 with respect to the random initialisation of $\mu$mu, assume there exists a ball $B_{\rho }(\mu ):=\{x^{\prime }\mid \Vert x^{\prime }-\mu \Vert <\rho \}$B rho of mu, defined as the set of points x prime such that the distance to mu is less than rho of fixed radius $\rho >0$rho greater than zero where $f(x)$ is $C^2$-continuous for all $x\in B_{\rho }(\mu )$ and $\Vert {\nabla }^{2}f(\mu )\Vert$ be polynomial bounded in $d$. Within this ball, let ${\nabla }^{2}f(x)$ be Lipschitz continuous, i.e. $\Vert {\nabla }^{2}f(x)-{\nabla }^{2}f(y)\Vert \le L_d\Vert x-y\Vert$ for all $x,y\in B_{\rho }(\mu )$.

This assumption still permits discontinuous objectives. We also assume that $p_0(\cdot )$ generates sub-Gaussian elements with uniform tail control:

Assumption 6 (Sub-Gaussian Tails). In addition to Assumption 1, assume that $p_0(\cdot )$ generates variables that have sub-Gaussian tails, i.e. for $x_i\sim p_0(x_i)$:

$$\mathbb{P}(|x_i|>t)\le 2\exp (-C{t}^2),$$

the probability that the absolute value of x i is greater than t is less than or equal to two e to the minus C t squared

for some $0\le C<\infty$ that does not depend on $d$.

We discuss sub-Gaussian variables and their properties in Section C.3 The assumption is trivially satisfied for Gaussian distributions $a\sim \mathcal{N}(0,I_m)$ and $b\sim \mathcal{N}(0,I_n)$, and holds more generally, for example for bounded distributions, uniform distributions and generalised Gaussian distributions with shape parameter greater than two. This flexibility is particularly relevant for the models in Section 6.1, where heavier-shouldered distributions may be preferred over the Gaussian.

Theorem 3 (EGGROLL Convergence to Linearity). Let $W\in {\mathbb{R}}^{m\times n}$, $d=mn$ and $x=\text{Vec}(W)$. Let Assumptions 3, 4, 5 and 6 hold, ${\sigma }_d=o(d^{-1/2})$, and $L_d({\sigma }_{d}d{)}^2=o(1)$. Then there exists some $K>0$ such that:

$$\Vert {\hat{g}}_{\text{LR}}-{\nabla }_{W}f(W=M){\Vert }_F=\mathcal{O}\left(L_d({\sigma }_{d}d{)}^2\right)+\mathcal{O}\left(\frac{\sqrt{d}}{{\sigma }_d^2}\exp \left(-K\frac{\rho }{\sqrt{d}{\sigma }_d}\right)\right)=o(1),\text{\hspace{1em}(10)}$$

and

$$\Vert {\hat{g}}_{\text{LR}}-{\nabla }_{M}J(\theta ){\Vert }_F=\mathcal{O}\left({\sigma }_d\sqrt{d}\cdot \left(1+L_d{\sigma }_d{d}^{\frac{3}{2}}\right)\right)=o(1).\text{\hspace{1em}(11)}$$

almost surely with respect to the distribution over $\mu$.

Our theory explains the success of EGGROLL in high dimensions with rank as small as $r=1$r equals one; Eq. (11) proves EGGROLL converges to the true update matrix ES update ${\nabla }_{M}J(\theta )$grad M J of theta as $d\to \infty$d goes to infinity regardless of $r$r. In addition, Eq. (10) proves that under the same conditions, the EGGROLL update also linearises like the true Gaussian ES update analysed in Section 5.1, recovering a local first-order derivative as $d\to \infty$d goes to infinity. For high-dimensional neural networks, standard parametrisations place training in the NTK regime, in which the network behaves approximately linearly in its parameters and gradient descent converges to a global minimum (Jacot et al., 2018; Lee et al., 2019; Chizat et al., 2019)as shown in prior research. Recent results show that the spectral norm of the Hessian decays polynomially with width, and that higher-order derivatives governing the variation of the Hessian also vanish (Liu et al., 2020). Consequently, the Lipschitz constant $L_d=o(1)$L d is small o of one, typically at rate $d^{-\frac{1}{2}}$ or $d^{-1}$ depending on the network architecture. Substituting these rates into our upper bound in Eq. (10) yields convergence rates of $\mathcal{O}({\sigma }_d^2{d}^{\frac{3}{2}})$ or $\mathcal{O}({\sigma }_d^{2}d)$ respectively.

5.3 Rank Analysis

We now analyse how fast the low-rank update from Eq. (4) with Gaussian score approximation converges to the true Gaussian ES matrix gradient in Eq. (3) as the rank of the update $r$r increases. We make notation explicit in $r$r in this subsection, for example writing $E^r=\frac{1}{\sqrt{r}}{A}^r{B}^{r\top }$E r equals one over the square root of r, times A r B r transpose. We introduce the following formal regularity assumption for the fitness function:

Assumption 7 (Bounded Fitness). Assume that $f(W)$f of W is bounded, that is ${\sup }_W|f(W)|<\infty$the supremum over W of the absolute value of f of W is finite.

Our key theoretical result characterises the error rate between the Gaussian score approximator in the low-rank update ${\hat{g}}_{\text{LR}}^r$ from Eq. (4) and the true gradient using the matrix Frobenius norm:

Theorem 4 (EGGROLL Rank Convergence). Let Assumptions 1 and 7 hold, then:

$$\Vert {\hat{g}}_{\text{LR}}^r-{\nabla }_{\mu }J(\theta ){\Vert }_F=\mathcal{O}\left(r^{-1}\right).\text{\hspace{1em}(12)}$$

the Frobenius norm of the difference between the low rank estimate and the true gradient is of order r to the minus one

Figure 3: Plot of Marginal Score Multiplied by Density for Increasing $r$

The convergence rate in Eq. (12) is faster than the typical $\mathcal{O}(r^{-\frac{1}{2}})$order r to the minus one half rate dictated by the general parametric central limit theorem. Our analysis shows that this is due to the symmetry in our problem under Assumption 1. To obtain our results, we make an Edgeworth expansion (Bhattacharya & Ranga Rao, 1976) of the distribution $P(E^r)$P of E r, which expands $P(E^r)$P of E r as the limiting Gaussian distribution plus a sum of decaying terms that are controlled by the 3rd order and higher cumulants of $P(E^r)$P of E r. Each $i$th order cumulant term is multiplied by a factor that decays at rate $\mathcal{O}(r^{-\frac{i-2}{2}})$order r to the power of minus i minus two over two. For symmetric zero-mean distributions, all odd cumulants are zero (for the same reason that all odd moments of a symmetric distribution are zero). Hence, the rate of convergence to the limiting distribution is controlled by the 4th order term, which has rate $\mathcal{O}(r^{-1})$order r to the minus one.

Although the full distribution $P(E^r)$P of E r has no general closed-form solution, the distribution over marginals $P(E_{i,j})$ is more amenable to analysis. We derive the density of the marginal distribution $P(E_{i,j})$ for generalised Gaussian distributed $a_{i,j}$ and $b_{i,j}$ in Section D.1. To illustrate the fast convergence rate, we plot the negative density $\times$ score function $p(E_{i,j})E_{i,j}$ for the marginal density $p(E_{i,j})$ in Fig. 3 using Gaussian distributed $a_{i,j}$ and $b_{i,j}$ (see Theorem 6 for a derivation). The figure shows that $p(E_{i,j})E_{i,j}$ quickly converges to the limiting function $\frac{E_{i,j}}{\sqrt{2\pi }}\exp \left(-\frac{E_{i,j}^2}{2}\right)$E i j over the square root of two pi, times the exponential of minus E i j squared over two, recovering the Gaussian form from the true Gaussian ES update. Even at $r=1$r equals one, the function is not a poor approximation. After $r=10$r equals ten, the function has nearly converged and after $r=50$r equals fifty, the function is visually indistinguishable from the limit, providing evidence for the hypothesis that the low-rank approximation is accurate even for very low-rank regimes $r\ll \min (m,n)$r much less than the minimum of m and n.

Figure 4: (a) Comparison of reinforcement learning returns normalised by PPO performance across 16 environments for 10 seeds. The shaded region is the standard error of the mean. (b) Validation score of 3 seeds of EGGROLL v.s. 3 seeds of GRPO in countdown task with an RWKV 7g1.5B model on a single GPU. EGGROLL allows 1024 parallel generations per GPU (618 updates) whereas GRPO only 64 (915 updates).

6 Experiments

In the following section we showcase the effectiveness of EGGROLL in a variety of tasks that position it as a strong alternative to back-propagation for the end-to-end training of foundation models.

6.1 Pure Integer Language Model Pretraining

To demonstrate the potential of EGGROLL as a general optimisation method, we apply it to language model pretraining. Since EGGROLL does not rely on gradients, we explicitly design a language model architecture to be efficient and hardware-friendly at inference time. To highlight EGGROLL’s flexibility, we train a nonlinear recurrent neural network (RNN) in pure integer datatypes with no explicit activation functions, relying only on the implicit nonlinearity of clipping in int8 operations. We call the resulting language model EGG, the Evolved Generative GRU, an EGGROLL-friendly architecture with all weights in int8. See Appendix G for more details on the architecture and motivation behind EGG.

We train an EGG model with 6 layers and hidden dimension 256 (6L-256D) to do character-level prediction on the minipile dataset (Kaddour, 2023). We update parameters after 100 tokens for each population member, applying truncated ES by keeping the hidden state and only resetting at document boundaries. We plot the test loss in Fig. 2b over training steps across a range of population sizes with a fixed data batch size of 16 sequences per step, where the best test loss is 3.40 bits/byte. With a sufficiently large population size, EGG outperforms a dense 6L-256D Transformer trained with backprop SGD using the same data batch size. Note that larger population sizes require more parallel compute for the same amount of data; our largest population size of $2^{20}=1048576$two to the twentieth power, which is one million forty eight thousand five hundred seventy six requires around 180 times more GPU-hours than the backprop baseline, demonstrating the potential for compute-only scaling in limited data regimes using EGGROLL.

Moreover, our largest population size of $2^{20}$two to the twentieth is three orders of magnitude larger than the largest experiment done by Salimans et al. (2017)Salimans and colleagues while only requiring a single GPU to train, highlighting EGGROLL’s computational efficiency. We note that large population sizes are critical for pretraining; a population size of 2, analogous to MeZO (Malladi et al., 2023), significantly underperforms larger population sizes despite having access to the same data batch. We conduct more ablations in Appendix I, analysing the tradeoff between population size and data batch size.

6.2 Reinforcement Learning Tasks

To verify that low-rank perturbations do not change the optimisation behavior of ES in standard control settings, we benchmark EGGROLL against OpenES (Salimans et al., 2017) across 16 tabula rasa environments spanning Navix, Craftax, Brax, Kinetix, and Jumanji. We use a fixed 3-layer MLP policy (256 hidden units) and perform per-environment hyperparameter optimisation for each method before evaluating the selected configuration over 10 random seeds, reporting mean performance (normalised by PPO) and uncertainty. Overall, EGGROLL is competitive with OpenES on 7/16 environments, underperforms on 2/16, and outperforms on 7/16, while often delivering substantial wall-clock improvements due to its batched low-rank structure (full environment

Figure 5: (a) Comparison of the validation score of 3 seeds of EGGROLL v.s. 3 seeds of GRPO in GSM8K task with an RWKV 7g7B model on 8 GPUs. EGGROLL allows 8192 parallel generations (1024 per GPU with 260 updates) whereas GRPO only 256 (32 per GPU with 340 updates). (b) Performance of our finetuned RWKV 7G 7 billion model on hard reasoning tasks using 128 GPUs for 12 hours. The model was trained using the DeepScaleR dataset and the best checkpoint was chosen by evaluating on AIME24.

list, learning curves, timing comparisons, and complete HPO ranges/settings are provided in Appendix N.4. Figure 4a shows the averaged normalised return across the 16 environments with 10 seeds per environment. We additionally report MARL results in Section N.1.

6.3 Foundation Model Fine-tuning

We apply EGGROLL to finetune an RWKV-7 (Peng et al., 2025)Peng and colleagues LLM on two reasoning tasks: countdown (Gandhi et al., 2024)Gandhi and colleagues and GSM8K (Cobbe et al., 2021)Cobbe and colleagues. RWKV is a recurrent model that is better suited to parallelisation than transformers because any memory otherwise spent on the KV cache is used to evaluate population members. Figure 4b shows that EGGROLL fine-tuning on an RWKV-7 1.5B model converges to a higher validation accuracy of 35% (vs. 23%) given the same hardware and wall-clock time in the countdown task. Similarly, Figure 5a shows that EGGROLL outperforms GRPO on GSM8K fine-tuning. Our scoring function draws parallels to the group relative advantage of GRPO. In particular, to score a set of noise directions, $E\equiv \{E_1,\dots ,E_n\}$E, defined as the set E 1 through E n, we first compute their accuracies, $\{s_{1,q_i},\dots ,s_{n,q_i}\}$s one q i through s n q i, on $|q|=m$m questions, creating a matrix of scores $S\in {\mathbb{R}}^{m\times n}$S of dimension m by n. We then compute the normalised score per question, with the main difference that we use the global variance $\bar{\sigma }$sigma bar, and average over all the questions to compute a score for the noise direction $E_i$E i:

$${\bar{s}}_i=\frac{1}{m}\sum _{j=1}^m{z}_{i,q_j}=\frac{1}{m}\sum _{j=1}^m\frac{s_{i,j}-{\mu }_{q_j}}{\bar{\sigma }}.$$

This scoring function weights all questions within the same batch the same across population members. We use this recipe to train a 14 billion parameter RWKV 7 model on the DeepScaleR dataset and evaluate in more challenging maths reasoning tasks. In this regime, GRPO is infeasible due to the extra memory used by the Adam optimiser Kingma & Ba (2014)Kingma and Ba. Using a thinking budget of 5000 tokens for training and evaluation, our fine-tuned 14B model improves from 13% to 30% accuracy on AIME24, from 7% to 33% accuracy on AIME25 and from 11% to 13% accuracy on HMMT25 after training on 32 GPUs for 12 hours (Figure 13b). On 7B models, we outperform GRPO using 128 GPUs for 24 hours (Figure 5b).

In Section L, we achieve similar performance to GRPO when fine-tuning Qwen Transformer models, and additionally demonstrate that EGGROLL can directly optimise for pass@k, a known limitation of GRPO (Yue et al., 2025)Yue and colleagues. Beyond language models, we also fine-tune a finance world model into an agent for high-frequency trading that directly optimises for PnL; see Section M for more details.

6.4 Fine-tuning Integer Quantised LLMs

We follow the same procedure as Jacob et al. (2017)Jacob and colleagues to quantise the RWKV-7 family of models by dividing by the maximum per-channel value on each weight matrix and mapping into the int8 range of $[-127,127]$. We then apply EGGROLL with Adam to do model distillation from the original, non-quantised RWKV-7, into the resulting int8 quantised model using examples from GSM8K. See Appendix K for full details about the

Figure 6: (a) Average per token perplexity (during training) of 3 seeds of a quantised (int8) RWKV 7G 7 billion parameter model on distillation from the non quantised model using examples from GSM8K. (b) Validation score on unseen examples of GSM8K of 3 seeds of a quantised RWKV 7G 7 billion parameter model. Initially the model is unable to solve any problems, but progressively it is capable of solving more problems. The baseline here indicates the validation score of a quantised model without any further training.

specifics of quantisation and fine-tuning. The distillation is done by matching the distributions between the quantised and non-quantised models on teacher forced examples (with solutions) from the GSM8K dataset. More specifically, the fitness for a given set of parameters, ${\mu }_i$mu i, is computed as follows:

$$f_{{\mu }_i}(x_{1:T})=\sum _{t=1}^T\text{KL}\left(p_t||q_t(\cdot ;{\mu }_i)\right),$$

the fitness is calculated as the sum from t equals one to T of the Kullback-Leibler divergence between p sub t and q sub t given mu i

where $x_{1:T}$x one through T is a subsequence of tokens taken from the solutions of GSM8K and $\text{KL}\left(p_t||q_t(\cdot ;{\mu }_i)\right)$the K L divergence is the Kullback-Leibler divergence between the distribution of the non-quantised model, $p_t$p t, and the distribution of the quantised model $q_t$q t over the vocabulary at token $t$t. Figure 6a shows the average per token perplexity of 3 seeds of a quantised RWKV 7G 7 billion parameter model compared to that of the original non-quantised model over the same sequence, as a baseline. Progressively, the quantised model recovers the capability to solve a subset of the GSM8K dataset (Figure 6b).

7 Conclusion

We introduce EGGROLL, a powerful method for black-box optimisation that scales evolutionary strategies to billion-parameter models and beyond using low-rank search matrices. Our experiments demonstrate that EGGROLL is effective with a rank of 1, giving substantial computational and memory savings for negligible decrease in performance when compared to the full-rank perturbations. Empirically, EGGROLL delivers large speedups over naïve ES in tabula rasa and multi-agent RL, and can power end-to-end training pipelines for foundation models. Our theoretical analysis shows that the EGGROLL update converges towards the Gaussian ES update with increasing rank $r$r and parameter dimension $d=mn$d equals m times n, and we provide a rigorous study of general ES at high dimensions, deriving necessary and sufficient conditions for convergence and linearisation.

Looking forward, we can use EGGROLL for other problems beyond the reach of modern first-order gradient-based techniques. In particular, EGGROLL can enable the training of large scale end-to-end neurosymbolic systems (Sarker et al., 2021)by Sarker and colleagues with non-differentiable components. For instance, we can train neural networks that interface with symbolic modules for specific functions, like memory or calculations. We can also optimise end-to-end systems of language models, training them to be aware of inference-time harnesses and interactions with other agents in complex systems.

Acknowledgements

Compute for this project is graciously provided by the Isambard-AI National AI Research Resource, under the projects “FLAIR 2025 Moonshot Projects” and “Robustness via Self-Play RL.” Some experiments also used compute generously given by JASMIN, the UK’s collaborative data analysis environment (https://www.jasmin.ac.uk).

Bidipta Sarkar is supported by the Clarendon Fund Scholarship in partnership with a Department of Engineering Science Studentship for his Oxford DPhil. Mattie Fellows is funded by a generous grant from the UKRI Engineering and Physical Sciences Research Council EP/Y028481/1. Juan Agustin Duque is supported by the St-Pierre-Larochelle Scholarship at the University of Montreal and by Aaron Courville’s CIFAR AI Chair in Representations that Generalize Systematically. Jarek Liesen and Theo Wolf are supported by the EPSRC Centre for Doctoral Training in Autonomous Intelligent Machines & Systems EP/Y035070/1. Jarek Liesen is also supported by Sony Interactive Entertainment Europe Ltd. Uljad Berdica is supported by the EPSRC Centre for Doctoral Training in Autonomous Intelligent Machines & Systems EP/S024050/1 and the Rhodes Scholarship. Lukas Seier is supported by the Intelligent Earth CDT with funding from the UKRI grant number EP/Y030907/1. Alexander D. Goldie is funded by the EPSRC Centre for Doctoral Training in Autonomous Intelligent Machines and Systems EP/S024050/1. Jakob Nicolaus Foerster is partially funded by the UKRI grant EP/Y028481/1 (originally selected for funding by the ERC). Jakob Nicolaus Foerster is also supported by the JPMC Research Award and the Amazon Research Award.

We thank Andreas Kirsch for discovering an emergent log-linear scaling law for EGG loss with respect to int8 OPs in this tweet along with other community members for their comments and recommendations during the first arXiv release of this work.

References

[references omitted]

[references omitted]

Jean-Philippe Bouchaud, Julius Bonart, Jonathan Donier, and Martin Gould. Trades, quotes and prices: financial markets under the microscope. Cambridge University Press, 2018.

James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL GitHub.

Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners, 2020. URL arXiv.

Lénaïc Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d’Alché-Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL NeurIPS.

Krzysztof M Choromanski, Aldo Pacchiano, Jack Parker-Holder, Yunhao Tang, and Vikas Sindhwani. From complexity to simplicity: Adaptive es-active subspaces for blackbox optimization. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d’Alché-Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL NeurIPS.

Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, Parker Schuh, Kensen Shi, Sashank Tsvyashchenko, Joshua Maynez, Abhishek Rao, Parker Barnes, Yi Tay, Noam Shazeer, Vinodkumar Prabhakaran, Emily Reif, Nan Du, Ben Hutchinson, Reiner Pope, James Bradbury, Jacob Austin, Michael Isard, Guy Gur-Ari, Pengcheng Yin, Toju Duke, Anselm Levskaya, Sanjay Ghemawat, Sunipa Dev, Henryk Michalewski, Xavier Garcia, Vedant Misra, Kevin Robinson, Liam Fedus, Denny Zhou, Daphne Ippolito, David Luan, Hyeontaek Lim, Barret Zoph, Alexander Spiridonov, Ryan Sepassi, David Dohan, Shivani Agrawal, Mark Omernick, Andrew M. Dai, Thanumalayan Sankaranarayana Pillai, Marie Pellat, Aitor Lewkowycz, Erica Moreira, Rewon Child, Oleksandr Polozov, Katherine Lee, Zongwei Zhou, Xuezhi Wang, Brennan Saeta, Mark Diaz, Orhan Firat, Michele Catasta, Jason Wei, Kathy Meier-Hellstern, Douglas Eck, Jeff Dean, Slav Petrov, and Noah Fiedel. Palm: Scaling language modeling with pathways. J. Mach. Learn. Res., 24(1144), 2023. URL JMLR.

Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems. arXiv preprint arXiv:2110.14168, 2021.

A. P. Dawid. Some matrix-variate distribution theory: Notational considerations and a bayesian application. Biometrika, 68(1):265–274, 1981. ISSN 0006-3444.

Nan Du, Yanping Huang, Andrew M. Dai, Simon Tong, Dmitry Lepikhin, Yuanzhong Xu, Maxim Krikun, Yanqi Zhou, Adams Wei Yu, Orhan Firat, Barret Zoph, Liam Fedus, Maarten P. Bosma, Zongwei Zhou, Tao Wang, Emma Wang, Kellie Webster, Marie Pellat, Kevin Robinson, Kathleen Meier-Hellstern, Toju Duke, Lucas Dixon, Kun Zhang, Quoc Le, Yonghui Wu, Zhifeng Chen, and Claire Cui. Glam: Efficient scaling of language models with mixture-of-experts. In Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 5547–5569, Jul 2022. URL PMLR.

William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: scaling to trillion parameter models with simple and efficient sparsity. J. Mach. Learn. Res., 23(1):1–39, January 2022. ISSN 1532-4435. URL JMLR.

[references omitted]

[references omitted]

[references omitted]

[references omitted]

[references omitted]

[references omitted]

[references omitted]

Appendix

• A Notation

• B ES Matrix Gradient Deviations

• C High-Dimensional Analysis

  • C.1 High-Dimensional Gaussian ES and Convergence
  • C.2 Critical Convergence Rate
  • C.3 EGGROLL Linearisation
  • D Asymptotic Rank Analysis
  • D.1 Mean Field Score Function Approximator
  • D.2 Derivation of Mean-field Approximators
  • E EGGROLL Speed
  • F Arithmetic Intensity Analysis
  • F.1 Arithmetic Intensity of Standard Batched Inference
  • F.2 Arithmetic Intensity of Gaussian Matrix ES
  • F.3 Arithmetic Intensity of EGGROLL
  • G EGG Architecture
  • G.1 Motivation
  • G.2 Notation and Operations
  • G.3 Parameter Initialisation
  • G.4 Matrix Multiplication
  • G.5 Embedding
  • G.6 Layer Normalisation (LN)
  • G.7 MLP
  • G.8 GRU
  • G.9 Fitness Calculation in Integer Types
  • H EGG Pretraining with Integer EGGROLL
  • H.1 Adding EGGROLL Perturbations
  • H.2 Fitness Shaping
  • H.3 Parameter Update
  • I EGG Ablations
  • J Distributed EGGROLL Framework
  • J.1 Base-3 Fitness Packing and Bandwidth Efficiency
  • J.2 System Architecture
  • K Fine-tuning of Integer Quantised Models
  • K.1 Quantisation Procedure

• K.2 Integrating integer-quantised EGGROLL with Adam

• L Fine-tuning Pretrained Transformer LLMs with Verifiable Rewards

  • L.1 Results
  • L.2 Training Infrastructure for Large-Scale Transformer LLMs
  • M Fine-tuning Time Series Foundation Model: High-Frequency Trading
  • N Experimental Details
  • N.1 Multi Agent Reinforcement Learning Experiments
  • N.2 Reasoning Fine-tuning Experiments: Countdown
  • N.3 Reasoning Fine-tuning Experiments: GSM8K
  • N.4 Reinforcement Learning Experiments

A Notation

In our proofs, we use the integral notation $\int$ to denote the integral over the corresponding ${\mathbb{R}}^d$R d space, for example, for a matrix $E\in {\mathbb{R}}^{m\times n}$E in R m by n, $\int f(E)dE={\int }_{{\mathbb{R}}^{m\times n}}f(E)dE$ and for a vector $E\in {\mathbb{R}}^{mn}$E in R m n, $\int f(v)dv={\int }_{{\mathbb{R}}^{mn}}f(v)dv$. For $f:{\mathbb{R}}^d\to \mathbb{R}$f from R d to R, we use $\nabla f(x)$ to denote the derivative of $f(\cdot )$ evaluated at $x$. For a vector $v\in {\mathbb{R}}^{mn}$v in R m n, we define the mat operator as:

$$\text{mat}(v)=\left[\begin{array}{cccc}{v}_1& v_{m+1}& \dots & v_{(n-1)m+1}\\ v_2& v_{m+2}& \dots & v_{(n-1)m+2}\\ ⋮& ⋮& \ddots & ⋮\\ v_m& v_{2m}& \dots & v_{mn}\end{array}\right],$$

so $\text{mat}(\text{vec}(M))=M$. We will use the fact that the Frobenius norm becomes the ${\ell }_2$l two norm in vector space:

$$\Vert M{\Vert }_F=\sqrt{\sum _{i,j}{m}_{i,j}^2}=\sqrt{\sum _k\text{vec}(M{)}_k^2}=\Vert \text{vec}(M)\Vert .\text{\hspace{1em}(13)}$$

Our proofs make use of Fourier analysis. For a vector-valued function $f(v):{\mathbb{R}}^d\to \mathbb{R}$, we define the Fourier transform as:

$$\tilde{f}(\omega )=\mathcal{F}[f](\omega ):=\int f(v)\exp (-i{\omega }^{\top }v)dv,$$

and the inverse Fourier transform as:

$$f(v)={\mathcal{F}}^{-1}[\tilde{f}](v):=\frac{1}{(2\pi {)}^d}\int \tilde{f}(\omega )\exp (i{\omega }^{\top }v)d\omega ,$$

B ES Matrix Gradient Deviations

Let ${\mu }_M=\text{vec}(M)\in {\mathbb{R}}^{mn}$mu M equals vec M in R m n be the vector of mean parameters associated with the matrix $M$. Let $v_M\in {\mathbb{R}}^{mn}$v M denote the corresponding search vector associated with ${\mu }_M$mu M. As each element of $v$ is generated independently from a standard normal $\mathcal{N}(0,1)$, the search vector $v_M$ is generated from the standard multivariate norm: $v_M\sim \mathcal{N}(0,I_{mn})$v M is sampled from a multivariate normal with mean zero and identity matrix I m n. From Eq. (2), the update for ${\mu }_M$mu M is:

$$\begin{array}{rl}\sigma {\nabla }_{{\mu }_M}J(\theta )& ={\mathbb{E}}_{v_M\sim \mathcal{N}(0,I_{mn})}\left[v_M\cdot f(W=\text{mat}({\mu }_M)+\sigma \text{mat}(v_M))\right],\\ & ={\mathbb{E}}_{v_M\sim \mathcal{N}(0,I_{mn})}\left[\text{vec}(\text{mat}(v_M))\cdot f(W=\text{mat}({\mu }_M)+\sigma \text{mat}(v_M))\right],\\ & ={\mathbb{E}}_{E\sim \mathcal{N}(0,I_m,I_n)}\left[\text{vec}(E)\cdot f(W=M+\sigma E)\right],\end{array}$$

where $E=\text{mat}(v_M)$ and we have used the fact that sampling $v_M\sim \mathcal{N}(0,I_{mn})$ is equivalent to sampling $E\sim \mathcal{N}(0,I_m,I_n)$ and applying $v_M=\text{vec}(E)$. Now

$$\begin{array}{rl}{\nabla }_{M}J(\theta )& =\text{mat}({\nabla }_{{\mu }_M}J(\theta )),\\ & =\frac{1}{\sigma }{\mathbb{E}}_{E\sim \mathcal{N}(0,I_m,I_n)}\left[\text{mat}(\text{vec}(E))\cdot f(W=M+\sigma E)\right],\\ & =\frac{1}{\sigma }{\mathbb{E}}_{E\sim \mathcal{N}(0,I_m,I_n)}\left[E\cdot f(W=M+\sigma E)\right],\\ & =-\frac{1}{\sigma }{\mathbb{E}}_{E\sim \mathcal{N}(0,I_m,I_n)}\left[{\nabla }_E\log p(E)\cdot f(W=M+\sigma E)\right].\end{array}$$

C High-Dimensional Analysis

C.1 High-Dimensional Gaussian ES and Convergence

We use insights from the Gaussian annulus theorem when investigating the convergence properties of high-dimensional ES: our proof relies on the fact that all probability mass converges to the interior of the ball $B_{ϵ}(\mu ):=\{x^{\prime }\mid \Vert x^{\prime }-\mu \Vert <ϵ\}$B epsilon of mu, defined as the set of x prime such that the norm of x prime minus mu is less than epsilon where $ϵ=\frac{\rho }{2}$epsilon equals rho over two in the limit $d\to \infty$d approaches infinity, where $\rho$rho is the radius of the local ball from Assumption 2, meaning we only need to consider the smooth region around $\mu$mu in this limit. Our first result proves that the mass outside of the ball for any polynomially bounded function tends to zero at an exponential rate.

Lemma 1 (Polynomial Tail Bounds). Let $g(x)$g of x be polynomial bounded as:

$$\Vert g(\mu +{\sigma }_{d}v)\Vert \le C\Vert v{\Vert }^q(1+\Vert \mu +{\sigma }_{d}v{\Vert }^p),$$

for some finite polynomial of orders $p$ and $q$ and constant $C>0$. Let $A_d:=\{\Vert {\sigma }_{d}v\Vert \ge ϵ\}$ denote the event that a mutation lies outside the a local ball of radius $ϵ$ around $\mu$. Assume ${\sigma }_d=o(d^{-1/2})$. Then for some constant $K>0$:

$$\Vert {\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}[g(\mu +{\sigma }_{d}v)1(A_d)]\Vert =\mathcal{O}\left(d^{\frac{q}{2}}\exp \left(-K{\left(\frac{ϵ}{{\sigma }_d}\right)}^2\right)\right),$$

the norm of the expected value of g of mu plus sigma d v, times the indicator function of A d, is big O of d to the power of q over two, times the exponential of negative K times the square of epsilon over sigma d
and in particular the right-hand side is $o(1)$ as $d\to \infty$.

Proof. We start by bounding the integrand using the polynomial bound. Denote $\mathbb{P}(A_d):={\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}[1(A_d)]$. Then, by Jensen’s inequality in the first line, polynomial boundedness in the second and $\Vert a+b{\Vert }^p\le 2^{p-1}(\Vert a{\Vert }^p+\Vert b{\Vert }^p)$ in the third:

$$\begin{array}{rl}\Vert {\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}[g(\mu +{\sigma }_{d}v)1(A_d)]\Vert & \le {\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}[\Vert g(\mu +{\sigma }_{d}v)\Vert 1(A_d)],\\ & \le C{\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}[\Vert v{\Vert }^q(1+\Vert \mu +{\sigma }_{d}v{\Vert }^p)1(A_d)],\\ & \le C{\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}[\Vert v{\Vert }^q(1+2^{p-1}\Vert \mu {\Vert }^p)1(A_d)+2^{p-1}{\sigma }_d^p\Vert v{\Vert }^{p+q}1(A_d)],\\ & =C^{\prime }{\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}[\Vert v{\Vert }^{q}1(A_d)]+C^{\prime \prime }{\sigma }_d^p{\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}[\Vert v{\Vert }^{p+q}1(A_d)].\end{array}$$

where $C^{\prime }=C(1+2^{p-1}\Vert \mu {\Vert }^p)$ and $C^{\prime \prime }=C{2}^{p-1}$ are constants independent of $d$. Applying the Cauchy-Schwarz inequality to the second expectation gives:

$${\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}[\Vert v{\Vert }^{p+q}1(A_d)]\le \sqrt{{\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}[\Vert v{\Vert }^{2(p+q)}]}\cdot \sqrt{\mathbb{P}(A_d)}.$$

Now, the variable $\Vert v\Vert$ is ${\chi }_d$-distributed. Using the formula for the $i$-th central moment of $\Vert v\Vert$ about the origin (Forbes et al., 2011, Chapter 11.3)as described by Forbes and colleagues yields:

$${\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}[\Vert v{\Vert }^i]=2^{\frac{i}{2}}\frac{\Gamma \left(\frac{1}{2}(d+i)\right)}{\Gamma \left(\frac{1}{2}d\right)}.$$

Applying the identity $\frac{\Gamma (z+a)}{\Gamma (z+b)}\sim z^{a-b}$ (Askey & Roy, 2020-2026, Eq. 5.11.12)from Askey and Roy:

$${\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}[\Vert v{\Vert }^i]\sim 2^{\frac{i}{2}}{\left(\frac{d}{2}\right)}^{\frac{i}{2}}=d^{\frac{i}{2}},\text{\hspace{1em}(14)}$$

where $\sim$ denotes asymptotic equivalence. For $i=2(p+q)$, this yields the bound:

$${\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}[\Vert v{\Vert }^{2(p+q)}]=\mathcal{O}(d^{p+q}),$$

hence:

$$\begin{array}{rl}\Vert {\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}[g(\mu +{\sigma }_{d}v)1(A_d)]\Vert & \le C^{\prime }{d}^{\frac{q}{2}}\sqrt{\mathbb{P}(A_d)}+C^{\prime \prime }{\sigma }_d^p{d}^{\frac{p+q}{2}}\sqrt{\mathbb{P}(A_d)},\\ & =(C^{\prime }+C^{\prime \prime }{\sigma }_d^p{d}^{\frac{p}{2}})d^{\frac{q}{2}}\sqrt{\mathbb{P}(A_d)},\end{array}\text{\hspace{1em}(15)}$$

We use the Gaussian concentration inequality for the Euclidean norm (Vershynin, 2018, Theorem 3.1.1)from a 2018 paper by Vershynin, which states that for $x\sim \mathcal{N}(0,I_d)$x distributed as a standard multivariate normal there exists an absolute constant $K>0$ such that for all $t\ge 0$,

$$\mathbb{P}\left(\left|\Vert x\Vert -\sqrt{d}\right|\ge t\right)\le 2\exp (-K{t}^2).$$

In our setting, we need to bound:

$$\mathbb{P}(A_d)=\mathbb{P}(\Vert {\sigma }_{d}v\Vert \ge ϵ)=\mathbb{P}\left(\Vert v\Vert \ge \frac{ϵ}{{\sigma }_d}\right)=\mathbb{P}\left(\Vert v\Vert -\sqrt{d}\ge \frac{ϵ}{{\sigma }_d}-\sqrt{d}\right).$$

Setting $t=\frac{ϵ}{{\sigma }_d}-\sqrt{d}$, the assumption $\sqrt{d}{\sigma }_d=o(1)$root d times sigma d is little o of one implies for sufficiently large $d$ that $\sqrt{d}{\sigma }_d\le ϵ$ and therefore $t\ge 0$, so we can apply the concentration bound to obtain:

$$\begin{array}{rl}\mathbb{P}(A_d)& =\mathbb{P}\left(\Vert v\Vert -\sqrt{d}\ge t\right)\le \mathbb{P}\left(\left|\Vert v\Vert -\sqrt{d}\right|\ge t\right),\\ & =\mathcal{O}\left(\exp \left(-K{\left(\frac{ϵ}{{\sigma }_d}-\sqrt{d}\right)}^2\right)\right)=\mathcal{O}\left(\exp \left(-K{\left(\frac{ϵ}{{\sigma }_d}\right)}^2{\left(1-\frac{{\sigma }_d\sqrt{d}}{ϵ}\right)}^2\right)\right).\end{array}\text{\hspace{1em}(16)}$$

Now, as $\sqrt{d}{\sigma }_d=o(1)$, it follows $\frac{{\sigma }_d\sqrt{d}}{ϵ}=o(1)$, yielding:

$$\begin{array}{rl}\mathbb{P}(A_d)& =\mathcal{O}\left(\exp \left(-K{\left(\frac{ϵ}{{\sigma }_d}\right)}^2\right)\right),\\ ⟹\sqrt{\mathbb{P}(A_d)}& =\mathcal{O}\left(\exp \left(-\frac{K}{2}{\left(\frac{ϵ}{{\sigma }_d}\right)}^2\right)\right),\\ ⟹d^{\frac{q}{2}}\sqrt{\mathbb{P}(A_d)}& =\mathcal{O}\left(d^{\frac{q}{2}}\exp \left(-\frac{K}{2}{\left(\frac{ϵ}{{\sigma }_d}\right)}^2\right)\right),\end{array}$$

Applying these results to Eq. (15), along with ${\sigma }_d^p{d}^{\frac{p}{2}}=\mathcal{O}(d^{-\frac{p}{2}})d^{\frac{p}{2}}=\mathcal{O}(1)$, yields our desired result:

$$\begin{array}{rl}\Vert {\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}[g(\mu +{\sigma }_{d}v)1(A_d)]\Vert & \le C^{\prime }\mathcal{O}\left(d^{\frac{q}{2}}\exp \left(-\frac{K}{2}{\left(\frac{ϵ}{{\sigma }_d}\right)}^2\right)\right)\\ & +C^{\prime \prime }\mathcal{O}(1)\mathcal{O}\left(d^{\frac{q}{2}}\exp \left(-\frac{K}{2}{\left(\frac{ϵ}{{\sigma }_d}\right)}^2\right)\right),\\ & =\mathcal{O}\left(d^{\frac{q}{2}}\exp \left(-K{\left(\frac{ϵ}{{\sigma }_d}\right)}^2\right)\right).\end{array}$$

where we have absorbed the factor of $\frac{1}{2}$ into the constant $K$. $\square$

Our proof in Lemma 1 reveals the necessity of the condition ${\sigma }_d\sqrt{d}=o(1)$ for convergence as we can only apply the Gaussian concentration inequality in Eq. (16) for ${\sigma }_d\sqrt{d}=o(1)$; this is a direct consequence of the Gaussian annulus theorem, as for slower rates $1=o({\sigma }_d\sqrt{d})$, the Gaussian probability mass will exit any local ball around $\mu$ and flood the tail, meaning that the tail probability will grow with increasing $d$. Having bounded the tail, convergence to linearity follows by proving convergence within the ball, which allows us to exploit the local $C^1$ smoothness of $f(x)$:

Theorem 1 (Convergence to Linearity). Let Assumptions 2, 3 and 4 hold and ${\sigma }_d=o(d^{-\frac{1}{2}})$. Then:

$$\Vert {\nabla }_{\mu }J(\theta )-\nabla f(\mu )\Vert =\Theta \left({\left({\sigma }_d\sqrt{d}\right)}^{\alpha }\right)=o(1),$$

the norm of the gradient of J minus the gradient of f is big theta of sigma d times root d, to the power of alpha, which is little o of one

almost surely with respect to the distribution over $\mu$.

Proof. We start with the definition of the ES update:

$${\nabla }_{\mu }J(\theta )=\frac{1}{{\sigma }_d}{\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}[v\cdot f(\mu +{\sigma }_{d}v)].$$

Now let $ϵ=\frac{\rho }{2}$epsilon equals rho over two where $\rho$rho is the radius of the ball from Assumption 2. Consider the hinge function:

$$\varphi (x)=\{\begin{array}{ll}1,& \Vert x\Vert \le ϵ,\\ 2-\frac{\Vert x\Vert }{ϵ},& ϵ<\Vert x\Vert <2ϵ,\\ 0,& \Vert x\Vert \ge 2ϵ,\end{array}$$

which interpolates between 1 and 0 in the region $ϵ<\Vert x\Vert <2ϵ$. Our first goal is to use $\varphi (x)$phi of x to generate a function $\tilde{f}(x)$f tilde of x that is absolutely continuous and has integrable derivatives outside of $B_{\rho }(\mu )$B rho of mu to allow us to apply Stein’s lemma (Stein, 1972). We define $\tilde{f}(x)$f tilde of x as:

$$\tilde{f}(x)=f(x)\cdot \varphi (x-\mu )$$

Consider the closed ball $B_{ϵ}(\mu ):=\{x^{\prime }\mid \Vert x^{\prime }-\mu \Vert \le ϵ\}$. We note that within the ball $f(\mu +{\sigma }_{d}v)$ remains unchanged:

$$\tilde{f}(\mu +{\sigma }_{d}v)=\{\begin{array}{ll}f(\mu +{\sigma }_{d}v),& \Vert {\sigma }_{d}v\Vert \le ϵ,\\ f(\mu +{\sigma }_{d}v)\cdot \left(2-\frac{\Vert {\sigma }_{d}v\Vert }{ϵ}\right),& ϵ<\Vert {\sigma }_{d}v\Vert <2ϵ,\\ 0,& \Vert {\sigma }_{d}v\Vert \ge 2ϵ.\end{array}\text{\hspace{1em}(17)}$$

The derivative of the function with respect to $v$v is:

$${\nabla }_v\tilde{f}(\mu +{\sigma }_{d}v)=\{\begin{array}{ll}{\sigma }_d\nabla f(\mu +{\sigma }_{d}v),& \Vert {\sigma }_{d}v\Vert \le ϵ,\\ {\sigma }_d\nabla f(\mu +{\sigma }_{d}v)\cdot \left(2-\frac{\Vert {\sigma }_{d}v\Vert }{ϵ}\right)-\frac{{\sigma }_{d}v}{ϵ\Vert v\Vert }\cdot f(\mu +{\sigma }_{d}v),& ϵ<\Vert {\sigma }_{d}v\Vert <2ϵ,\\ 0,& \Vert {\sigma }_{d}v\Vert \ge 2ϵ.\end{array}\text{\hspace{1em}(18)}$$

where the gradient fails to exist only on the sets $\Vert {\sigma }_{d}v\Vert \in \{ϵ,2ϵ\}$, which have Lebesgue measure zero. We start by using this function to decompose $J(\mu )$J of mu into a smoothed part and a remainder:

$${\nabla }_{\mu }J(\theta )=\underset{:={\nabla }_{\mu }\tilde{J}(\mu )}{\underbrace{\frac{1}{{\sigma }_d}{\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}\left[v\cdot \tilde{f}(\mu +{\sigma }_{d}v)\right]}}+\underset{:=\Delta (\mu )}{\underbrace{\frac{1}{{\sigma }_d}{\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}\left[v\cdot \left(f(\mu +{\sigma }_{d}v)-\tilde{f}(\mu +{\sigma }_{d}v)\right)\right]}},$$

Hence:

$$\Vert {\nabla }_{\mu }J(\theta )-\nabla f(\mu )\Vert \le \Vert {\nabla }_{\mu }\tilde{J}(\mu )-\nabla f(\mu )\Vert +\Vert \Delta (\mu )\Vert .\text{\hspace{1em}(19)}$$

Consider the smoothed part:

$${\nabla }_{\mu }\tilde{J}(\mu ):=\frac{1}{{\sigma }_d}{\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}\left[v\cdot \tilde{f}(\mu +{\sigma }_{d}v)\right].$$

Our goal is to apply Stein’s lemma (Stein, 1972) in its multivariate form (Liu, 1994, Lemma 1). The assumptions of (Liu, 1994, Lemma 1)prior work require that the partial derivatives ${\partial }_{v_i}\tilde{f}(\mu +{\sigma }_{d}v)$ are absolutely continuous almost everywhere and:

$${\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}\left[|{\partial }_{v_i}\tilde{f}(\mu +{\sigma }_{d}v)|\right]<\infty .$$

These two conditions are satisfied by construction. Indeed, under Assumption 2, $f(\cdot )$f is $C^1$C one continuous on $B_{\rho }(\mu )$, hence from Eq. (17), $\tilde{f}(\cdot )$f tilde coincides with a compactly supported, piecewise $C^1$C one function whose gradient (Eq. (18)) exists almost everywhere. Moreover, under Assumption 3, both $f(\mu +{\sigma }_{d}v)$ and $\nabla f(\mu +{\sigma }_{d}v)$ are polynomially bounded, and since $\nabla \tilde{f}(\mu +{\sigma }_{d}v)$ is supported on $\Vert {\sigma }_{d}v\Vert \le 2ϵ$, it follows that:

$${\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}\left[\Vert \nabla \tilde{f}(\mu +{\sigma }_{d}v)\Vert \right]<\infty .$$

Applying (Liu, 1994, Lemma 1)Liu 1994, Lemma 1:

$$\begin{array}{rl}\frac{1}{{\sigma }_d}{\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}\left[v\cdot f(\mu +{\sigma }_{d}v)\right]& =\frac{1}{{\sigma }_d}{\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}\left[{\nabla }_v\tilde{f}(\mu +{\sigma }_{d}v)\right],\\ & ={\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}\left[\nabla \tilde{f}(\mu +{\sigma }_{d}v)\right],\\ ⟹\Vert {\nabla }_{\mu }\tilde{J}(\mu )-\nabla f(\mu )\Vert & =\Vert {\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}\left[\nabla \tilde{f}(\mu +{\sigma }_{d}v)-\nabla f(\mu )\right]\Vert \\ & \le {\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}\left[\Vert \nabla \tilde{f}(\mu +{\sigma }_{d}v)-\nabla f(\mu )\Vert \right].\end{array}\text{\hspace{1em}(20)}$$

Let $\{\mu +{\sigma }_{d}v\in B_{ϵ}(\mu )\}=\{\Vert {\sigma }_{d}v\Vert \le ϵ\}$ denote the event that a mutation lies within the ball $B_{ϵ}(\mu )$B epsilon of mu. We now split the integral into two regions, the first within the ball and the second outside:

$$\begin{array}{rl}{\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}\left[\Vert \nabla \tilde{f}(\mu +{\sigma }_{d}v)-\nabla f(\mu )\Vert \right]& =\underset{:=I_{\text{loc}}}{\underbrace{{\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}\left[\Vert \nabla \tilde{f}(\mu +{\sigma }_{d}v)-\nabla f(\mu )\Vert 1(\Vert {\sigma }_{d}v\Vert \le ϵ)\right]}}\\ & +\underset{:=I_{\text{tail}}}{\underbrace{{\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}\left[\Vert \nabla \tilde{f}(\mu +{\sigma }_{d}v)-\nabla f(\mu )\Vert 1(\Vert {\sigma }_{d}v\Vert >ϵ)\right]}}.\end{array}$$

Consider the region inside the ball, $I_{\text{loc}}$I local. From Eq. (18), $\nabla \tilde{f}(\mu +{\sigma }_{d}v)=\nabla f(\mu +{\sigma }_{d}v)$ within this region. Using the local $\alpha$-Hölder continuity from Assumption 2:

$$\begin{array}{rl}{I}_{\text{loc}}& ={\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}[\Vert \nabla f(\mu +{\sigma }_{d}v)-\nabla f(\mu )\Vert 1(\Vert {\sigma }_{d}v\Vert \le ϵ)],\\ & \le L{\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}[\Vert {\sigma }_{d}v{\Vert }^{\alpha }1(\Vert {\sigma }_{d}v\Vert \le ϵ)],\\ & \le {\sigma }_d^{\alpha }L{\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}[\Vert v{\Vert }^{\alpha }].\end{array}$$

Now, applying the identity ${\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}[\Vert v{\Vert }^i]\sim d^{\frac{i}{2}}$the expected value of the norm of v to the i is on the order of d to the i over two, from Eq. (14):

$$I_{\text{loc}}=\mathcal{O}\left({\left({\sigma }_d\sqrt{d}\right)}^{\alpha }\right).$$

We now bound the tail region outside the ball:

$$\begin{array}{rl}{I}_{\text{tail}}& ={\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}[\Vert \nabla \tilde{f}(\mu +{\sigma }_{d}v)-\nabla f(\mu )\Vert 1(\Vert {\sigma }_{d}v\Vert >ϵ)],\\ & \le {\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}[\Vert \nabla \tilde{f}(\mu +{\sigma }_{d}v)-\nabla f(\mu )\Vert 1(\Vert {\sigma }_{d}v\Vert \ge ϵ)].\end{array}$$

Now, as $\Vert \nabla f(\mu )\Vert =\mathcal{O}(1)$ from Assumption 4 and we have established that $\Vert \nabla \tilde{f}(\mu +{\sigma }_{d}v)\Vert$ is polynomial bounded under Assumption 3 when applying Stein’s lemma, it follows that $\Vert \nabla \tilde{f}(\mu +{\sigma }_{d}v)-\nabla f(\mu )\Vert$ is also polynomial bounded, that is there exists some constant $C>0$ and finite polynomial order $p$ such that:

$$\Vert \nabla \tilde{f}(\mu +{\sigma }_{d}v)-\nabla f(\mu )\Vert \le C(1+\Vert \mu +{\sigma }_{d}v{\Vert }^p).$$

Applying Lemma 1, it follows:

$$I_{\text{tail}}=\mathcal{O}\left(\exp \left(-K{\left(\frac{ϵ}{{\sigma }_d}\right)}^2\right)\right),$$

for some constant $K>0$. Together, this yields:

$$\begin{array}{rl}\Vert {\nabla }_{\mu }\tilde{J}(\mu )-\nabla f(\mu )\Vert & =I_{\text{loc}}+I_{\text{tail}},\\ & =\mathcal{O}\left({\left({\sigma }_d\sqrt{d}\right)}^{\alpha }\right)+\mathcal{O}\left(\exp \left(-K{\left(\frac{ϵ}{{\sigma }_d}\right)}^2\right)\right).\end{array}$$

As $\exp (-x)=o(x^{-a})$ for any $a>0$, we take $a=\alpha /2$ to obtain a weakened bound matching the first term:

$$\exp \left(-K{\left(\frac{ϵ}{{\sigma }_d}\right)}^2\right)=o\left({\left(\frac{{\sigma }_d}{ϵ}\right)}^{\alpha }\right)=o\left({\left({\sigma }_d\sqrt{d}\right)}^{\alpha }\right).$$

This yields the upper bound:

$$||{\nabla }_{\mu }\tilde{J}(\mu )-\nabla f(\mu )||=\mathcal{O}\left({\left({\sigma }_d\sqrt{d}\right)}^{\alpha }\right).\text{\hspace{1em}(21)}$$

Returning to Eq. (19), we must bound the remainder term:

$$\begin{array}{rl}||\Delta (\mu )||& =\Vert \frac{1}{{\sigma }_d}{\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}\left[v\cdot (f(\mu +{\sigma }_{d}v)-\tilde{f}(\mu +{\sigma }_{d}v))\right]\Vert ,\\ & \le \frac{1}{{\sigma }_d}{\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}\left[||v||\cdot |f(\mu +{\sigma }_{d}v)-\tilde{f}(\mu +{\sigma }_{d}v)|\right].\end{array}$$

Again, from Assumption 3, it follows that $|f(\mu +{\sigma }_{d}v)-\tilde{f}(\mu +{\sigma }_{d}v)|$ is polynomially bounded, that is there exists some constant $C^{\prime }>0$C prime greater than zero and finite polynomial order $p^{\prime }$p prime such that:

$$|f(\mu +{\sigma }_{d}v)-\tilde{f}(\mu +{\sigma }_{d}v)|\le C^{\prime }(1+||\mu +{\sigma }_{d}v|{|}^p).$$

Applying Lemma 1 with $q=1$q equals one:

$$||\Delta (\mu )||=\mathcal{O}\left(\frac{d^{\frac{1}{2}}}{{\sigma }_d}\exp \left(-K{\left(\frac{ϵ}{{\sigma }_d}\right)}^2\right)\right).$$

Now, as $\exp (-x)=o(x^{-1})$exponential of minus x is little o of x to the minus one for $x\to \infty$x approaching infinity, it follows:

$$\begin{array}{rl}\exp \left(-K{\left(\frac{ϵ}{{\sigma }_d}\right)}^2\right)& =o({\sigma }_d^2),\\ ⟹||\Delta (\mu )||& =\mathcal{O}\left(\frac{d^{\frac{1}{2}}}{{\sigma }_d}\exp \left(-K{\left(\frac{ϵ}{{\sigma }_d}\right)}^2\right)\right),\\ & =o({\sigma }_d\sqrt{d}),\\ & =o\left(({\sigma }_d\sqrt{d}{)}^{\alpha }\right),\end{array}\text{\hspace{1em}(22)}$$

where the final line follows from the fact $\sqrt{d}{\sigma }_d=o(1)$the square root of d times sigma d is little o of one. Assembling our bounds using Ineq. 19 yields our desired result:

$$\Vert {\nabla }_{\mu }J(\theta )-\nabla f(\mu )\Vert \le \underset{=\mathcal{O}(({\sigma }_d\sqrt{d}{)}^{\alpha }),\text{ Eq. 21}}{\underbrace{\Vert {\nabla }_{\mu }\tilde{J}(\mu )-\nabla f(\mu )\Vert }}+\underset{=o((({\sigma }_d\sqrt{d}{)}^{\alpha }),\text{ Eq. 22}}{\underbrace{\Vert \Delta (\mu )\Vert }}=\mathcal{O}\left({\left({\sigma }_d\sqrt{d}\right)}^{\alpha }\right).$$

We now show that the bound is tight. Consider the function $f(x)=\frac{L}{2}{\sum }_{i=1}^d{x}_i|x_i|+a^{\top }x$f of x equals L over two times the sum from i equals one to d of x i times the absolute value of x i, plus a transpose x where $\Vert a\Vert =\mathcal{O}(1)$the norm of a is order one. Taking partial derivatives:

$${\partial }_{i}f(x)=L|x_i|+a_i,\text{\hspace{1em}(23)}$$

hence:

$$\Vert \nabla f(x)-\nabla f(y)\Vert =L\sqrt{\sum _{i=1}^d(|x_i|-|y_i|{)}^2}$$

Applying the reverse triangle inequality $||x_i|-|y_i||\le |x_i-y_i|⟹(|x_i|-|y_i|{)}^2\le (x_i-y_i{)}^2$the absolute value of x i minus the absolute value of y i, squared, is less than or equal to x i minus y i, squared:

$$\Vert \nabla f(x)-\nabla f(y)\Vert \le L\sqrt{\sum _{i=1}^d(x_i-y_i{)}^2}=L\Vert x-y\Vert .$$

We have thus shown that $f(x)$f of x is $C^1$-continuous and its gradient has Lipschitz constant $L$, i.e. $\alpha =1$ with Hölder constant $L$. It is also bounded by a polynomial of order 2. Without loss of generality, we take a deterministic initialisation $\mu =0$mu equals zero to simplify algebra, yielding;

$$\nabla f(\mu )=a⟹\Vert \nabla f(\mu )\Vert =\Vert a\Vert =\mathcal{O}(1).$$

$f(x)$f of x thus satisfies Assumptions 2, 3 and 4.Assumptions two, three and four. Using $f(x)$f of x as the fitness:

$$\begin{array}{rl}{\nabla }_{\mu }J(\theta )-\underset{=a}{\underbrace{\nabla f(\mu )}}& =\frac{1}{{\sigma }_d}{\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}[v\cdot f({\sigma }_{d}v)]-a,\\ & ={\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}[\nabla f({\sigma }_{d}v)]-a;\end{array}$$

Taking expectations element-wise and using Eq. (23):

$$\begin{array}{rl}[{\nabla }_{\mu }J(\theta )-\nabla f(\mu ){]}_i& ={\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}[{\partial }_{i}f({\sigma }_{d}v)]-a_i,\\ & ={\sigma }_{d}L{\mathbb{E}}_{v_i\sim \mathcal{N}(0,1)}[|v_i|].\end{array}$$

Applying Eq. (14):

$${\mathbb{E}}_{v_i\sim \mathcal{N}(0,1)}[|v_i|]=\sqrt{2}\frac{\Gamma (1)}{\Gamma (\frac{1}{2})}=\sqrt{\frac{2}{\pi }},$$

Hence:

$$\Vert {\nabla }_{\mu }J(\theta )-\nabla f(\mu )\Vert ={\sigma }_d\sqrt{d}\cdot L\sqrt{\frac{2}{\pi }},$$

thereby attaining the upper bound rate of ${\sigma }_d\sqrt{d}$. $\square$

C.2 Critical Convergence Rate

To show that our rate is critical, we investigate the space of functions that can be represented by cubic polynomials of the form:

$$f(x)=a^{\top }x+\frac{1}{2}{x}^{\top }Bx+\frac{1}{6}C[x,x,x],\text{\hspace{1em}(24)}$$

where $a\in {\mathbb{R}}^d,B\in {\mathbb{R}}^{d\times d}$a in R d and B in R d by d is a symmetric matrix and $C[x,x,x]={\sum }_{i,j,k}{c}_{i,j,k}{x}_i{x}_j{x}_k$ denotes a symmetric 3-linear map represented by the 3-tensor $C\in {\mathbb{R}}^{d\times d\times d}$C in R d by d by d.

Since our theory depends on analysing the local stability of a smooth ball for a fitness function, stability over this class is necessary for convergence on more general objectives. We show that once ${\sigma }_d$sigma d decays slower than the critical rate, divergence already occurs within this subclass, establishing the sharpness of the rate.

Theorem 2 (Exact divergence for cubic objectives). Let $f(x)$f of x denote the cubic polynomial in Eq. (24). Assume $\Vert a\Vert =\mathcal{O}(1),\Vert B\Vert =\mathcal{O}(1),\Vert C\Vert =\mathcal{O}(1)$the norms of a, B, and C are all order one where $\Vert \cdot \Vert$ denotes operator norm for $i$-tensor $T(x_1,\dots ,x_i)$: $\Vert T\Vert :={\sup }_{\Vert x_1\Vert =\cdots =\Vert x_i\Vert =1}|T(x_1,\dots ,x_i)|$. Let Assumption 4 hold, then:

$${\nabla }_{\mu }J(\theta )=\nabla f(\mu )+\frac{{\sigma }_d^2}{2}{\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}[C(v,v,\cdot )].$$

Moreover:

$$\begin{array}{rl}\Vert \frac{{\sigma }_d^2}{2}{\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}[C(v,v,\cdot )]\Vert & =\Theta ({\sigma }_d^{2}d),\\ \Vert {\nabla }_{\mu }J(\theta )-\nabla f(\mu )\Vert & =\Theta ({\sigma }_d^{2}d).\end{array}$$

Proof. We start by taking derivatives of $f(x)$f of x:

$$\nabla f(x)=a+Bx+\frac{1}{2}C(x,x,\cdot ).$$

Substituting this into the definition of ${\nabla }_{\mu }J(\theta )$the gradient of J with respect to mu and using Eq. (20):

$$\begin{array}{rl}{\nabla }_{\mu }J(\theta )& =\frac{1}{{\sigma }_d}{\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}\left[vf(\mu +{\sigma }_{d}v)\right],\\ & ={\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}\left[\nabla f(\mu +{\sigma }_{d}v)\right],\\ & ={\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}\left[a+B(\mu +{\sigma }_{d}v)+\frac{1}{2}C(\mu +{\sigma }_{d}v,\mu +{\sigma }_{d}v,\cdot )\right],\\ & =a+B\mu +{\sigma }_{d}B\underset{=0}{\underbrace{{\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}[v]}}+\frac{1}{2}{\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}\left[C(\mu +{\sigma }_{d}v,\mu +{\sigma }_{d}v,\cdot )\right],\\ & =a+B\mu +\frac{1}{2}{\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}\left[C(\mu ,\mu ,\cdot )+{\sigma }_{d}C(v,\mu ,\cdot )+{\sigma }_{d}C(\mu ,v,\cdot )+{\sigma }_d^{2}C(v,v,\cdot )\right],\\ & =\underset{=\nabla f(\mu )}{\underbrace{a+B\mu +\frac{1}{2}C(\mu ,\mu ,\cdot )}}+\frac{1}{2}{\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}\left[2{\sigma }_{d}C(v,\mu ,\cdot )+{\sigma }_d^{2}C(v,v,\cdot )\right],\\ & =\nabla f(\mu )+{\sigma }_d{\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}\left[C(v,\mu ,\cdot )\right]+\frac{{\sigma }_d^2}{2}{\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}\left[C(v,v,\cdot )\right],\end{array}$$

where we have used the fact $C(v,\mu ,\cdot )=C(\mu ,v,\cdot )$ by definition of the symmetry of $C$C. As $C(v,\mu ,\cdot )$C of v, mu, dot is linear in $v$v, its expectation under zero-mean $\mathcal{N}(0,I_d)$Gaussian noise is zero, hence:

$${\nabla }_{\mu }J(\theta )=\nabla f(\mu )+\frac{{\sigma }_d^2}{2}{\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}\left[C(v,v,\cdot )\right],$$

proving our first result. Now, it follows that $\Vert C(v,v,\cdot )\Vert \le \Vert C\Vert \Vert v{\Vert }^2$the norm of C of v, v, dot is less than or equal to the norm of C times the squared norm of v and as $\Vert C\Vert =\mathcal{O}(1)$the norm of C is order one:

$$\begin{array}{rl}\Vert {\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}[C(v,v,\cdot )]\Vert & \le \Vert C\Vert {\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}[\Vert v{\Vert }^2],\\ & =\mathcal{O}({\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}[\Vert v{\Vert }^2])\end{array}$$

Now as $v$v is unit Gaussian: ${\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}[\Vert v{\Vert }^2]=d$the expectation of the squared norm of v is d, hence:

$$\Vert {\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}[C(v,v,\cdot )]\Vert =\mathcal{O}(d).$$

We now show that the bound is tight. Consider the function $f(x)=u^{\top }x\Vert x{\Vert }^2$f of x equals u transpose x times the squared norm of x for $u^{\top }=\frac{1}{\sqrt{d}}[1,\dots ,1]$u transpose equals one over the square root of d times a vector of ones. The factor of $\frac{1}{\sqrt{d}}$ ensures that the gradient of the function ${\nabla }_{x}f(x)=\mathcal{O}(1)$the gradient of f is order one. We can write $\Vert x{\Vert }^2$the squared norm of x as the tensor contraction:

$$\Vert x{\Vert }^2=I_d(x,x),$$

where $I_d$I sub d is the identity matrix and:

$$u^{\top }(x)=u(x),$$

hence we write $f(x)$f of x as a tensor contraction as:

$$f(x)=C(x,x,x),$$

where $C:=\text{Sym}(u\otimes I_d)$C is defined as the symmetrization of the tensor product of u and the identity matrix. Using this function:

$$\begin{array}{rl}C(v,v,\cdot )& ={\nabla }_v(u^{\top }v\Vert v{\Vert }^2)\\ & =u\Vert v{\Vert }^2+2v{u}^{\top }v,\\ ⟹{\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}[C(v,v,\cdot )]& =u\underset{=d}{\underbrace{{\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}[\Vert v{\Vert }^2]}}+2\underset{=I_d}{\underbrace{{\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}[v{v}^{\top }]}}u,\\ & =u(d+2),\end{array}$$

hence $\Vert {\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}[C(v,v,\cdot )]\Vert =d+2$, achieving the upper bound rate of $\mathcal{O}(d)$big O of d which implies:

$$\Vert {\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}[C(v,v,\cdot )]\Vert =\Theta (d).$$

Our final result follows immediately:

$$\Vert {\nabla }_{\mu }J(\theta )-\nabla f(\mu )\Vert =\frac{({\sigma }_d{)}^2}{2}\Vert {\mathbb{E}}_{v\sim \mathcal{N}(0,I_d)}[C(v,v,\cdot )]\Vert =\Theta ({\sigma }_d^{2}d).$$

C.3 EGGROLL Linearisation

We now study the effect of EGGROLL in high dimensions. We introduce the notation $v=\text{vec}(E)$v equals the vectorization of E to denote the vectorisation of the low-rank matrix perturbation $E=\frac{1}{\sqrt{r}}A{B}^{\top }$E equals one over the square root of r, times A B transpose and work in vector space. The EGGROLL vector update $v$v can thus be written as sum of independent variables:

$$v=\sum _{i=1}^r\frac{1}{\sqrt{r}}{u}_i$$

with:

$$u_i=\text{vec}\left(a_i{b}_i^{\top }\right),$$

where recall $a_i$ and $b_i$a sub i and b sub i are the $i$th column vectors of $A$ and $B$A and B. We write $\mu =\text{vec}(M)$mu equals the vectorization of M. Using Eq. (13), we can convert between results in vector space and matrix space as:

$$\begin{array}{rl}\Vert \text{vec}({\hat{g}}_{\text{LR}})-\nabla f(\mu )\Vert & =\Vert {\hat{g}}_{\text{LR}}-{\nabla }_{W}f(W=M){\Vert }_F,\\ \Vert \text{vec}({\hat{g}}_{\text{LR}})-{\nabla }_{\mu }J(\theta )\Vert & =\Vert {\hat{g}}_{\text{LR}}-{\nabla }_{M}J(\theta ){\Vert }_F.\end{array}$$

To extend our analysis, we need to ensure that all polynomial moments of $P(v)$P of v are finite and grow at most polynomially in the dimension $d=mn$d equals m times n. In particular, such tail bounds are sufficient to dominate polynomial error terms in our analysis. To introduce sub-Gaussian variables, we follow the exposition of Vershynin (2018)Vershynin and results therein. A random variable $x_i\in \mathbb{R}$x sub i in the set of real numbers is sub-Gaussian if there exists some finite constant $C>0$C greater than zero such that for all $t>0$t greater than zero:

$$\mathbb{P}(|x_i|>t)\le 2\exp (-C{t}^2),$$

meaning their their tails decay like Gaussians. This is equivalent to any of the following three properties holding (Vershynin, 2018, 2.6.1): There exist constants $C_1,C_2,C_3>0$C one, C two, and C three, greater than zero that differ at most by an absolute constant factor such that:

$$\begin{array}{rl}(\mathbb{E}[|x_i{|}^p]{)}^{\frac{1}{p}}& \le C_1\sqrt{p},\forall p\ge 1,\\ \mathbb{E}\left[\exp \left(\frac{x_i^2}{C_2^2}\right)\right]& \le 2,\end{array}$$

and if $\mathbb{E}[x_i]=0$the expected value of x sub i is zero:

$$\mathbb{E}[\exp (\lambda x_i)]\le \exp (C_3^2{\lambda }^2),\forall \lambda \in \mathbb{R}.$$

A random vector $x\in {\mathbb{R}}^d$x in R d is sub-Gaussian if all one-dimensional marginals of $x$x are sub-Gaussian, i.e. $x^{\top }u$x transpose u is sub-Gaussian for all $u\in {\mathbb{R}}^d$u in R d. The sub-Gaussian norm is defined as:

$$\Vert x{\Vert }_{{\psi }_2}:=\inf \left\{K|\mathbb{E}\left[\exp \left(u^{\top }(x-\mathbb{E}[x])\right)\right]\le \exp \left(\frac{K^2\Vert u{\Vert }^2}{2}\right),\forall u\in {\mathbb{R}}^d\right\}.$$

which returns the smallest universal sub-Gaussian constant for all marginals.

A key property of sub-Gaussian vectors that we use in our proofs is the sub-Gaussian concentration inequality for the Euclidean norm (Vershynin, 2018, Theorem 3.1.1)detailed in Theorem 3.1.1 of Vershynin 2018, which states that for if $x$x is a sub-Gaussian vector with $\mathbb{E}[x_i^2]=1$expected value of x i squared equals one and $K=\Vert x{\Vert }_{{\psi }_2}$K equals the sub-Gaussian norm of x, there exists an absolute constant $C>0$C greater than zero such that for all $t\ge 0$t greater than or equal to zero,

$$\mathbb{P}\left(\left|\Vert x\Vert -\sqrt{m}\right|\ge t\right)\le 2\exp \left(-\frac{C{t}^2}{K^4}\right).\text{\hspace{1em}(25)}$$

the probability that the absolute difference between the norm of x and the square root of m is at least t, is less than or equal to two times the exponential of negative C t squared over K to the fourth

We also use a weaker form of control, that replaces the Gaussian-like tail decay with an exponential decay, but all other properties are defined similarly. In this paper, we use the definition that a variable $x$x is known as sub-exponential if there exists a $K>0$K greater than zero such that for all $t\ge 0$:

$$\mathbb{P}(|x|\ge t)\le 2\exp \left(-\frac{t}{K}\right).$$

Our first result derives a bound on the expected value of the norms $\Vert a{\Vert }^i$ and $\Vert b{\Vert }^i$:

Lemma 2. Let Assumption 6 hold. Let $P(a)$ denote the distribution over columns of $A$ and $P(b)$ denote the distribution over columns of $B$. Then:

$${\mathbb{E}}_{a\sim P(a)}[\Vert a{\Vert }^i]=\mathcal{O}(m^{\frac{i}{2}}),{\mathbb{E}}_{b\sim P(b)}[\Vert b{\Vert }^i]=\mathcal{O}(n^{\frac{i}{2}}).$$

Proof. It suffices to prove ${\mathbb{E}}_{a\sim P(a)}[\Vert a{\Vert }^i]=\mathcal{O}(m^{\frac{i}{2}})$ as ${\mathbb{E}}_{b\sim P(b)}[\Vert b{\Vert }^i]=\mathcal{O}(n^{\frac{i}{2}})$ follows automatically from the same assumptions. We start by using the ‘layer cake’ representation of the expectation Lieb & Loss (2010 - 2010, Theorem 1.13)from Lieb and Loss, Theorem 1.13:

$${\mathbb{E}}_{a\sim P(a)}[\Vert a{\Vert }^i]=i{\int }_0^{\infty }{t}^{i-1}\mathbb{P}(\Vert a\Vert >t)dt.$$

Let $t_m=C\sqrt{m}$ for any $C>1$. We split the integral into two regions:

$$i{\int }_0^{\infty }{t}^{i-1}\mathbb{P}(\Vert a\Vert >t)dt={\int }_0^{t_m}i{t}^{i-1}\mathbb{P}(\Vert a\Vert >t)dt+{\int }_{t_m}^{\infty }i{t}^{i-1}\mathbb{P}(\Vert a\Vert >t)dt.$$

For the first integral:

$$\begin{array}{rl}{\int }_0^{t_m}i{t}^{i-1}\mathbb{P}(\Vert a\Vert >t)dt& \le {\int }_0^{t_m}i{t}^{i-1}dt,\\ & =(t_m{)}^i,\\ & =C^i{m}^{\frac{i}{2}}.\end{array}$$

For the second integral, we wish to bound $\mathbb{P}(\Vert a\Vert >t)$ for the region $t\ge t_m=C\sqrt{m}$. Setting $t^{\prime }=t-\sqrt{m}>0$, the assumption $C>1$ implies $t^{\prime }\ge 0$ in this region, hence

$$\mathbb{P}(\Vert a\Vert >t)=\mathbb{P}(\Vert a\Vert -\sqrt{m}>t^{\prime })\le \mathbb{P}(|\Vert a\Vert -\sqrt{m}|>t^{\prime }).$$

We bound this using the sub-Gaussian concentration inequality from Eq. (25). Under Assumption 6, $a$ is a sub-Gaussian vector with $\Vert x{\Vert }_{{\psi }_2}\le \infty$with a finite sub-Gaussian norm, hence there exists an absolute constant $C^{\prime }>0$ such that for all $t^{\prime }\ge 0$,

$$\mathbb{P}\left(\left|\Vert a\Vert -\sqrt{m}\right|\ge t^{\prime }\right)\le 2\exp \left(-C^{\prime }{t}^{\prime 2}\right).$$

This implies:

$$\mathbb{P}(\Vert a\Vert \ge t)\le 2\exp \left(-C^{\prime }(t-\sqrt{m}{)}^2\right),$$

for all $t\ge t_m$. Substituting yields:

$${\int }_{t_m}^{\infty }i{t}^{i-1}\mathbb{P}(\Vert a\Vert >t)dt\le {\int }_{t_m}^{\infty }i{t}^{i-1}\exp \left(-C^{\prime }(t-\sqrt{m}{)}^2\right)dt.$$

Let $x=t-\sqrt{m}⟹dt=dx$x equals t minus the square root of m, which implies d t equals d x:

$${\int }_{t_m}^{\infty }i{t}^{i-1}\mathbb{P}(\Vert a\Vert >t)dt\le {\int }_{\sqrt{m}(C-1)}^{\infty }i(x+\sqrt{m}{)}^{i-1}\exp \left(-C^{\prime }{x}^2\right)dx.$$

Now, $\sqrt{m}\le \frac{x}{C-1}$the square root of m is less than or equal to x over C minus one for all $x\ge \sqrt{m}(C-1)$, hence:

$$\begin{array}{rl}{\int }_{t_m}^{\infty }i{t}^{i-1}\mathbb{P}(\Vert a\Vert >t)dt& \le {\int }_{\sqrt{m}(C-1)}^{\infty }i{x}^{i-1}{\left(1+\frac{1}{C-1}\right)}^{i-1}\exp \left(-C^{\prime }{x}^2\right)dx,\\ & =i{\left(1+\frac{1}{C-1}\right)}^{i-1}{\int }_{\sqrt{m}(C-1)}^{\infty }{x}^{i-1}\exp \left(-C^{\prime }{x}^2\right)dx,\\ & \le i{\left(1+\frac{1}{C-1}\right)}^{i-1}{\int }_0^{\infty }{x}^{i-1}\exp \left(-C^{\prime }{x}^2\right)dx,\\ & \le i{\left(1+\frac{1}{C-1}\right)}^{i-1}\frac{1}{2}(C^{\prime }{)}^{-i/2}\Gamma \left(\frac{i}{2}\right),\\ & =\mathcal{O}(1).\end{array}$$

Combining the two bounds yields:

$${\mathbb{E}}_{a\sim P(a)}[\Vert a{\Vert }^i]=\mathcal{O}(m^{\frac{i}{2}}),$$

the expected value of the norm of a to the i-th power is big O of m to the i over two

as required. $\square$

Using this result, we now bound the whole vector $v=\frac{1}{\sqrt{r}}{\sum }_{i=1}^r\text{vec}(a_i{b}_i^{\top })$v equals one over the square root of r times the sum from i equals one to r of the vectorization of a i times b i transpose

Lemma 3. Let $i\ge 1$. Under Assumption 6:

$${\mathbb{E}}_{v\sim P(v)}\left[\Vert v{\Vert }^i\right]=\mathcal{O}\left((rmn{)}^{\frac{i}{2}}\right)$$

the expected value of the norm of v to the i-th power is big O of the quantity r times m times n, all to the i over two

Proof. For any vectors $a,b$a and b:

$$\begin{array}{rl}\Vert \text{vec}(a{b}^{\top })\Vert & =\sqrt{\sum _{j=1}^m\sum _{k=1}^n(a_j{b}_k{)}^2}=\sqrt{\sum _{j=1}^m{a}_j^2\sum _{k=1}^n{b}_k^2}=\Vert a\Vert \Vert b\Vert ,\\ ⟹\Vert \text{vec}(a{b}^{\top }){\Vert }^i& =\Vert a{\Vert }^i\Vert b{\Vert }^i.\end{array}$$

Applying Lemma 2 under Assumption 6 for each summand of $v=\frac{1}{\sqrt{r}}{\sum }_{l=1}^r\text{vec}(a_l{b}_l^{\top })$v:

$$\begin{array}{rl}{\mathbb{E}}_{v\sim P(v)}\left[\Vert \text{vec}(a_l{b}_l^{\top }){\Vert }^i\right]& ={\mathbb{E}}_{a_l\sim P(a_l)}\left[\Vert a_l{\Vert }^i\right]{\mathbb{E}}_{b_l\sim p(b_l)}\left[\Vert b_l{\Vert }^i\right],\\ & =\mathcal{O}\left((mn{)}^{\frac{i}{2}}\right).\end{array}$$

Applying the triangle inequality:

$$\begin{array}{rl}{\mathbb{E}}_{v\sim P(v)}\left[\Vert v{\Vert }^i\right]& ={\mathbb{E}}_{v\sim P(v)}\left[{\Vert \frac{1}{\sqrt{r}}\sum _{l=1}^r\text{vec}(a_l{b}_l^{\top })\Vert }^i\right],\\ & \le {\mathbb{E}}_{v\sim P(v)}\left[{\left(\frac{1}{\sqrt{r}}\sum _{l=1}^r\Vert \text{vec}(a_l{b}_l^{\top })\Vert \right)}^i\right],\\ & =r^{\frac{i}{2}}{\mathbb{E}}_{v\sim P(v)}\left[{\left(\frac{1}{r}\sum _{l=1}^r\Vert \text{vec}(a_l{b}_l^{\top })\Vert \right)}^i\right].\end{array}$$

Now, as $i\ge 1$i is greater than or equal to one, we can apply Jensen’s inequality:

$${\left(\frac{1}{r}\sum _{l=1}^r\Vert \text{vec}(a_l{b}_l^{\top })\Vert \right)}^i\le \frac{1}{r}\sum _{l=1}^r\Vert \text{vec}(a_l{b}_l^{\top }){\Vert }^i,$$

yielding:

$${\mathbb{E}}_{v\sim P(v)}[\Vert v{\Vert }^i]\le r^{(\frac{i}{2}-1)}\sum _{l=1}^r{\mathbb{E}}_{v\sim P(v)}[\Vert \text{vec}(a_l{b}_l^{\top }){\Vert }^i]=r^{\frac{i}{2}}\mathcal{O}\left((mn{)}^{\frac{i}{2}}\right)=\mathcal{O}\left((rmn{)}^{\frac{i}{2}}\right).\square$$

Our proof borrows techniques used to prove linearisation of the ES update in Section C.1 by bounding the tail probability of any polynomial under the low-rank distribution outside of the ball $B_{\rho }(\mu )$B rho of mu. To apply the concentration inequality that would generalise Lemma 1, we show that $v$v has an exponentially decaying tail:

Lemma 4 (Exponential Tail Bound). Let $r<\infty$r be finite and Assumption 6 hold. Then all elements of $v$v are sub-exponential and for $\sqrt{d}{\sigma }_d=o(1)$the square root of d times sigma d is little o of one there exists some constant $C>0$C greater than zero such that:

$$\mathbb{P}(\Vert {\sigma }_{d}v\Vert \ge \rho )\le 2d\exp \left(-C\frac{\rho }{\sqrt{d}{\sigma }_d}\right).$$

Proof. In matrix form:

$$E=\frac{1}{\sqrt{r}}\sum _{i=1}^r{a}_i{b}_i^{\top }.$$

The elements of $E$E are thus:

$$E_{j,k}=\frac{1}{\sqrt{r}}\sum _{i=1}^r{a}_{ij}{b}_{ik}.$$

As $a_{ij}$a i j and $b_{ik}$b i k are independent sub-Gaussian random variables with zero mean, it follows from Vershynin (2018, Lemma 2.8.6)Vershynin that their product $a_{ij}{b}_{ik}$a i j times b i k is a zero-mean sub-exponential variable with a uniform norm $\Vert a_{ij}{b}_{ik}{\Vert }_{{\psi }_1}<\infty$psi one norm is finite. Finally, a finite sum of sub-exponential variables is sub-exponential (Wainwright, 2019, Eq. (2.18))according to Wainwright with a uniform norm, so all elements of $E$E and hence $v=\text{vec}(E)$v equals vector E are sub-exponential and zero-mean with a uniform ${\psi }_1$-norm $K<\infty$psi one norm K is finite.

We now bound $\mathbb{P}(\Vert {\sigma }_{d}v\Vert \ge \rho )=\mathbb{P}(\Vert v\Vert \ge \frac{\rho }{{\sigma }_d})$. For the vector $v$v, it follows for $t\ge 0$:

$$\Vert v\Vert \ge t⟹\underset{j}{\max }|v_j|\ge \frac{t}{\sqrt{d}}.$$

This is easily proven via the contrapositive: if ${\max }_j|v_j|<\frac{t}{\sqrt{d}}$ then

$$\Vert v{\Vert }^2=\sum _{j=1}^d{v}_j^2<d\frac{t^2}{d}=t^2,$$

implying $\Vert v\Vert <t$. This means for $t\ge 0$:

$$\begin{array}{rl}\mathbb{P}(\Vert v\Vert \ge t)& \le \mathbb{P}\left(\underset{j}{\max }|v_j|\ge \frac{t}{\sqrt{d}}\right),\\ & \le \sum _{j=1}^d\mathbb{P}\left(|v_j|\ge \frac{t}{\sqrt{d}}\right).\end{array}\text{\hspace{1em}(26)}$$

As $v_j$v j is a sub-exponential variable with finite uniform sub-exponential norm, by definition (Vershynin, 2018, Proposition 2.8.1)according to work by Vershynin in 2018 there exists a finite $K$K such that for all $j$j:

$$\mathbb{P}\left(|v_j|\ge \frac{t}{\sqrt{d}}\right)\le 2\exp \left(-\frac{t}{\sqrt{d}K}\right).$$

Applying to Eq. (26) yields:

$$\mathbb{P}(\Vert v\Vert \ge t)\le 2d\exp \left(-\frac{t}{\sqrt{d}K}\right).$$

Now, using $t=\frac{\rho }{{\sigma }_d}$t equals rho over sigma d and $C=\frac{1}{K}$C equals one over K yields:

$$\mathbb{P}(\Vert {\sigma }_{d}v\Vert \ge \rho )\le 2d\exp \left(-C\frac{\rho }{\sqrt{d}{\sigma }_d}\right).$$

We now use these results to assemble into our key polynomial tail bound:

Lemma 5 (EGGROLL Polynomial Tail Bounds). Let Assumption 6 hold. Let $g(x)$ be polynomial bounded as:

$$\Vert g(x)\Vert \le C(1+\Vert x{\Vert }^p),$$

for some finite polynomial of order $p$ and constant $C>0$. Consider the ball $B_{\rho }(\mu ):=\{x^{\prime }\mid \Vert x^{\prime }-\mu \Vert <\rho \}$. Let $\{\mu +{\sigma }_{d}v\notin B_{\rho }(\mu )\}=\{\Vert {\sigma }_{d}v\Vert \ge \rho \}$ denote the event that a mutation lies outside the ball. Assume ${\sigma }_d=o(d^{-1/2})$. Then for some constant $K>0$ independent of $d$:

$$\Vert {\mathbb{E}}_{v\sim P(v)}[g(\mu +{\sigma }_{d}v)1(A_d)]\Vert =\mathcal{O}\left(\sqrt{d}\exp \left(-K\frac{\rho }{\sqrt{d}{\sigma }_d}\right)\right),$$

and in particular the right-hand side is $o(1)$ as $d\to \infty$.

Proof. Let

$$A_d:=\{\mu +{\sigma }_{d}v\notin B_{\rho }(\mu )\}$$

and denote $\mathbb{P}(A_d):={\mathbb{E}}_{v\sim P(v)}[1(A_d)]$. Our proof proceeds as in Lemma 1 to obtain:

$$\Vert {\mathbb{E}}_{v\sim P(v)}[g(\mu +{\sigma }_{d}v)1(A_d)]\Vert \le C^{\prime }\mathbb{P}(A_d)+C^{\prime \prime }{\sigma }_d^p{\mathbb{E}}_{v\sim P(v)}[\Vert v{\Vert }^{p}1(A_d)].$$

where $C^{\prime }=C(1+2^{p-1}\Vert \mu {\Vert }^p)$ and $C^{\prime \prime }=C{2}^{p-1}$ are constant in $d$. Applying the Cauchy–Schwarz inequality to the second expectation gives:

$${\mathbb{E}}_{v\sim P(v)}[\Vert v{\Vert }^{p}1(A_d)]\le \sqrt{{\mathbb{E}}_{v\sim P(v)}[\Vert v{\Vert }^{2p}]}\cdot \sqrt{\mathbb{P}(A_d)}.$$

Applying Lemma 3 with fixed $r$ and $d=mn$:

$$\sqrt{{\mathbb{E}}_{v\sim P(v)}[\Vert v{\Vert }^{2p}]}=\mathcal{O}\left(d^{\frac{p}{2}}\right).$$

Now, $\mathbb{P}(A_d)=\mathbb{P}(\Vert {\sigma }_{d}v\Vert \ge \rho )$. From Lemma 4, there exists some $K>0$ such that:

$$\begin{array}{rl}\mathbb{P}(\Vert {\sigma }_{d}v\Vert \ge \rho )& \le 2d\exp \left(-K\frac{\rho }{\sqrt{d}{\sigma }_d}\right),\\ ⟹\sqrt{\mathbb{P}(A_d)}& =\mathcal{O}\left(\sqrt{d}\exp \left(-K\frac{\rho }{\sqrt{d}{\sigma }_d}\right)\right),\end{array}$$

where we have absorbed the factor of $\frac{1}{2}$ into $K$, hence:

$${\mathbb{E}}_{v\sim P(v)}[\Vert v{\Vert }^{p}1(A_d)]=\mathcal{O}\left(d^{\frac{p+1}{2}}\exp \left(-K\frac{\rho }{\sqrt{d}{\sigma }_d}\right)\right).$$

Now, as $\sqrt{d}{\sigma }_d=o(1)$the square root of d times sigma sub d is little o of one, ${\sigma }_d^p{d}^{\frac{p}{2}}=o(1)$sigma sub d to the p times d to the p over two is little o of one, hence:

$${\sigma }_d^p{\mathbb{E}}_{v\sim P(v)}[\Vert v{\Vert }^{p}1(A_d)]=\mathcal{O}\left(\sqrt{d}\exp \left(-K\frac{\rho }{\sqrt{d}{\sigma }_d}\right)\right).$$

Applying our bounds yields our desired result:

$$\Vert {\mathbb{E}}_{v\sim P(v)}[g(\mu +{\sigma }_{d}v)1(A_d)]\Vert =\mathcal{O}\left(\sqrt{d}\exp \left(-K\frac{\rho }{\sqrt{d}{\sigma }_d}\right)\right)=o(1).$$

where the $o(1)$little o of one bound follows from the fact that the exponential factor dominates $\sqrt{d}$the square root of d and $\sqrt{d}{\sigma }_d=o(1)$the square root of d times sigma sub d is little o of one. $\square$

Theorem 3 (EGGROLL Convergence to Linearity). Let Assumptions 3, 4, 5 and 6 hold and ${\sigma }_d=o(d^{-1/2})$sigma sub d is little o of d to the minus one half and $L_d({\sigma }_{d}d{)}^2=o(1)$L sub d times the quantity sigma sub d times d, squared, is little o of one. Then there exists some $K>0$K greater than zero such that:

$$\Vert \text{vec}({\hat{g}}_{\text{LR}})-\nabla f(\mu )\Vert =\mathcal{O}(L_d({\sigma }_{d}d{)}^2)+\mathcal{O}\left(\frac{\sqrt{d}}{{\sigma }_d^2}\exp \left(-K\frac{\rho }{\sqrt{d}{\sigma }_d}\right)\right)=o(1),$$ $$\Vert \text{vec}({\hat{g}}_{\text{LR}})-{\nabla }_{\mu }J(\theta )\Vert =\mathcal{O}\left({\sigma }_d\sqrt{d}\cdot \left(1+L_d{\sigma }_d{d}^{\frac{3}{2}}\right)\right)=o(1).$$

almost surely with respect to the distribution over $\mu$.

Proof. We start with the definition of the vectorised EGGROLL update:

$$\begin{array}{rl}\text{vec}({\hat{g}}_{\text{LR}})-\nabla f(\mu )& =\frac{1}{{\sigma }_d}{\mathbb{E}}_{v\sim P(v)}[v\cdot f(\mu +{\sigma }_{d}v)]-\nabla f(\mu ),\\ & =\frac{1}{{\sigma }_d}{\mathbb{E}}_{v\sim P(v)}[v\cdot f(\mu +{\sigma }_{d}v)]-\underset{=0}{\underbrace{\frac{1}{{\sigma }_d}{\mathbb{E}}_{v\sim P(v)}[v]}}\cdot f(\mu )-\underset{I_d}{\underbrace{{\mathbb{E}}_{v\sim P(v)}[v{v}^{\top }]}}\nabla f(\mu )\\ & +\underset{=0}{\underbrace{\frac{1}{2{\sigma }_d}{\mathbb{E}}_{v\sim P(v)}[{\sigma }_d^{2}v{v}^{\top }{\nabla }^{2}f(\mu )v]}},\\ & =\frac{1}{{\sigma }_d}{\mathbb{E}}_{v\sim P(v)}\left[v\cdot \underset{:=T_d(v)}{\underbrace{\left(f(\mu +{\sigma }_{d}v)-f(\mu )-{\sigma }_d{v}^{\top }\nabla f(\mu )+\frac{{\sigma }_d^2}{2}{v}^{\top }{\nabla }^{2}f(\mu )v\right)}}\right],\\ & =\frac{1}{{\sigma }_d}{\mathbb{E}}_{v\sim P(v)}[v\cdot T_d(v)],\end{array}$$

where we have used the fact that the expectation of an odd function under a symmetric, zero mean distribution is always zero, and $P(v)$ satisfies this under Assumption 6, hence ${\mathbb{E}}_{v\sim P(v)}[v{v}^{\top }{\nabla }^{2}f(\mu )v]=0$, and ${\mathbb{E}}_{v\sim P(v)}[v{v}^{\top }]=I_d$ from Lemma 6. Consider the ball $B_{\rho }(\mu ):=\{x^{\prime }\mid \Vert x^{\prime }-\mu \Vert <\rho \}$. We now split the integral into two regions, the first within the ball and the second outside:

$$\begin{array}{rl}\frac{1}{{\sigma }_d}{\mathbb{E}}_{v\sim P(v)}[v\cdot (f(\mu +{\sigma }_{d}v)-f(\mu ))]& =\underset{:=I_{\text{loc}}}{\underbrace{\frac{1}{{\sigma }_d}{\mathbb{E}}_{v\sim P(v)}[v\cdot T_d(v)1(\Vert {\sigma }_{d}v\Vert <\rho )]}}\\ & +\underset{:=I_{\text{tail}}}{\underbrace{\frac{1}{{\sigma }_d}{\mathbb{E}}_{v\sim P(v)}[v\cdot T_d(v)1(\Vert {\sigma }_{d}v\Vert \ge \rho )]}}.\end{array}$$

Consider the region inside the ball:

$$\begin{array}{rl}\Vert I_{\text{loc}}\Vert & =\frac{1}{{\sigma }_d}\Vert {\mathbb{E}}_{v\sim P(v)}[v\cdot T_d(v)1(\Vert {\sigma }_{d}v\Vert <\rho )]\Vert ,\\ & \le \frac{1}{{\sigma }_d}{\mathbb{E}}_{v\sim P(v)}[\Vert v\Vert |T_d(v)|1(\Vert {\sigma }_{d}v\Vert <\rho )].\end{array}\text{\hspace{1em}(27)}$$

Within this region, $f(\mu +{\sigma }_{d}v)$f of mu plus sigma d v is $C^2$C two continuous under Assumption 5. We can thus write $f(\mu +{\sigma }_{d}v)$f of mu plus sigma d v using a first-order Taylor expansion about $\mu$mu with a Hessian (second order derivative) remainder within the ball:

$$\begin{array}{rl}f(\mu +{\sigma }_{d}v)& =f(\mu )+{\sigma }_d\nabla f(\mu {)}^{\top }v+{\sigma }_d^2{v}^{\top }\left({\int }_0^1(t-1){\nabla }^{2}f(\mu +t{\sigma }_{d}v)dt\right)v,\\ ⟹T_d(v)& ={\sigma }_d^2{v}^{\top }\left({\int }_0^1(t-1){\nabla }^{2}f(\mu +t{\sigma }_{d}v)dt\right)v+\frac{{\sigma }_d^2}{2}{v}^{\top }{\nabla }^{2}f(\mu )v,\\ & ={\sigma }_d^2{v}^{\top }\left({\int }_0^1(t-1)({\nabla }^{2}f(\mu +t{\sigma }_{d}v)-{\nabla }^{2}f(\mu ))dt\right)v.\end{array}$$

Applying the Lipschitz bound on the Hessian from Assumption 5:

$$\begin{array}{rl}|T_d(v)|& \le {\sigma }_d^2\Vert v{\Vert }^2\Vert {\int }_0^1(t-1)({\nabla }^{2}f(\mu +t{\sigma }_{d}v)-{\nabla }^{2}f(\mu ))dt\Vert ,\\ & \le {\sigma }_d^2\Vert v{\Vert }^2{\int }_0^1(t-1)\Vert {\nabla }^{2}f(\mu +t{\sigma }_{d}v)-{\nabla }^{2}f(\mu )\Vert dt,\\ & \le {\sigma }_d^2\Vert v{\Vert }^2{\int }_0^1(t-1)L_d\Vert t{\sigma }_{d}v\Vert dt,\\ & ={\sigma }_d^3\Vert v{\Vert }^3{L}_d\left|{\int }_0^1(t-1)tdt\right|,\\ & =\frac{L_d}{6}{\sigma }_d^3\Vert v{\Vert }^3.\end{array}$$

Using this to bound Eq. (27):

$$\begin{array}{rl}\Vert I_{\text{loc}}\Vert & \le \frac{L_d}{6}{\sigma }_d^2{\mathbb{E}}_{v\sim P(v)}\left[\Vert v{\Vert }^{4}1(\Vert {\sigma }_{d}v\Vert <\rho )\right],\\ & \le \frac{L_d}{6}{\sigma }_d^2{\mathbb{E}}_{v\sim P(v)}\left[\Vert v{\Vert }^4\right].\end{array}$$

Now, (for fixed $r$r) we apply the identity ${\mathbb{E}}_{v\sim P(v)}\left[\Vert v{\Vert }^4\right]=\mathcal{O}((mn{)}^2)$the expectation of the norm of v to the fourth is order m n squared with $mn=d$m n equals d from Lemma 3:

$$\Vert I_{\text{loc}}\Vert =\mathcal{O}(L_d({\sigma }_{d}d{)}^2).$$

We now bound the tail region outside the ball:

$$\begin{array}{rl}{I}_{\text{tail}}& =\frac{1}{{\sigma }_d}{\mathbb{E}}_{v\sim P(v)}\left[v\cdot T_d(v)1(\Vert {\sigma }_{d}v\Vert \ge \rho )\right],\\ & \le \frac{1}{{\sigma }_d}{\mathbb{E}}_{v\sim P(v)}\left[\Vert v\Vert |T_d(v)|1(\Vert {\sigma }_{d}v\Vert \ge \rho )\right],\\ & =\frac{1}{{\sigma }_d^2}{\mathbb{E}}_{v\sim P(v)}\left[\Vert {\sigma }_{d}v\Vert |T_d(v)|1(\Vert {\sigma }_{d}v\Vert \ge \rho )\right].\end{array}$$

Now under Assumptions 3, 4 and 5, $f(\mu +{\sigma }_{d}v)$f of mu plus sigma d v is polynomial bounded, $\Vert \nabla f(\mu )\Vert =\mathcal{O}(1)$the norm of the gradient of f is order one and $\Vert {\nabla }^{2}f(\mu )\Vert$the norm of the Hessian of f is polynomial bounded hence there exists some finite constant $C>0$C greater than zero and finite polynomial order $p$p such that:

$$\Vert {\sigma }_{d}v\Vert |T_d(v)|\le C(1+\Vert \mu +{\sigma }_{d}v{\Vert }^p).$$

We thus apply Lemma 5:

$$\begin{array}{rl}\frac{1}{{\sigma }_d^2}{\mathbb{E}}_{v\sim P(v)}\left[\Vert {\sigma }_{d}v\Vert |T_d(v)|1(\Vert {\sigma }_{d}v\Vert \ge \rho )\right]& =\mathcal{O}\left(\frac{\sqrt{d}}{{\sigma }_d^2}\exp \left(-K\frac{\rho }{\sqrt{d}{\sigma }_d}\right)\right),\\ & =\mathcal{O}\left(\frac{d\sqrt{d}}{d{\sigma }_d^2}\exp \left(-K\frac{\rho }{\sqrt{d}{\sigma }_d}\right)\right).\end{array}$$

Now, as ${\sigma }_d\sqrt{d}=o(1)$sigma d times the square root of d is little o of one, the exponential term dominates the prefactor $\frac{d\sqrt{d}}{d{\sigma }_d^2}$d times the square root of d, over d sigma d squared, we conclude:

$$\frac{1}{{\sigma }_d^2}{\mathbb{E}}_{v\sim P(v)}\left[\Vert {\sigma }_{d}v\Vert |T_d(v)|1(\Vert {\sigma }_{d}v\Vert \ge \rho )\right]=o(1)$$

Our final result follows from:

$$\begin{array}{rl}\Vert \text{vec}({\hat{g}}_{\text{LR}})-{\nabla }_{\mu }J(\theta )\Vert & =\Vert \text{vec}({\hat{g}}_{\text{LR}})-\nabla f(\mu )+\nabla f(\mu )-{\nabla }_{\mu }J(\theta )\Vert ,\\ & \le \Vert \text{vec}({\hat{g}}_{\text{LR}})-\nabla f(\mu )\Vert +\Vert \nabla f(\mu )-{\nabla }_{\mu }J(\theta )\Vert .\end{array}$$

We have already shown $\Vert \text{vec}({\hat{g}}_{\text{LR}})-\nabla f(\mu )\Vert =o(1)$the norm of vec g hat L R minus the gradient of f of mu is little o of one and under the assumptions for this theorem, Theorem 1 holds and so $\Vert \nabla f(\mu )-{\nabla }_{\mu }J(\theta )\Vert =o(1)$the norm of the gradient of f of mu minus the gradient with respect to mu of J of theta is little o of one. $\square$

D Asymptotic Rank Analysis

For convenience, we work with random vectors in our analysis. We analyse the vector $v^r=\text{vec}(E^r)$v r, which is the vectorization of E r, which is the vectorisation of the low-rank matrix $E^r$E r. We denote $v=\text{vec}(E)$v, the vectorization of E, which is the vectorisation of the full rank matrix $E$E. Note $v\sim \mathcal{N}(0,I_d)$v is distributed as a multivariate normal with mean zero and identity covariance which we denote as $P(v)$P of v. We write $v^r$v r as a standardised sum of $r$r independent, zero-mean random vectors. Let

$$u_i=\text{vec}(a_i{b}_i^{\top }),\text{\hspace{1em}(28)}$$

u i equals the vectorization of a i times b i transpose

where recall $a_i$a i and $b_i$b i are the $i$th column vectors of $A$ and $B$ so:

$$v^r=\frac{1}{\sqrt{r}}\sum _{i=1}^r{u}_i.$$

v r equals one over the square root of r, times the sum from i equals one to r of u i

Denoting the covariance matrix of $p(u)$ as ${\Sigma }_u$sigma u, the central limit theorem proves that the distribution of $v^r$v r converges in distribution to a zero-mean Gaussian $\mathcal{N}(0,{\Sigma }_r)$normal distribution with covariance sigma r. In Lemma 6, we derive the covariance matrix for ${\Sigma }_u$sigma u, which we prove is the identity. Our analysis uses an Edgeworth expansion (Bhattacharya & Ranga Rao, 1976)by Bhattacharya and Ranga Rao to characterise precisely the rate at which $P(v^r)$P of v r converges to the limiting Gaussian distribution. In Lemma 7, we make an Edgeworth expansion of $P(v^r)$P of v r to show that it is dominated by $\mathcal{O}(r^{-1})$order r to the minus one terms and higher. These are then used to prove Lemma 8, which allows us to bound the integral of the remainder of the Edgeworth expansion, thereby characterising how fast $P(v^r)$P of v r converges to the limiting Gaussian distribution.

Lemma 6. Let Assumption 1 hold and $u_i$u i be defined in Eq. (28). Then the variable $u_i$u i has identity covariance matrix:

$${\Sigma }_u:={\mathbb{E}}_{u_i\sim p(u_i)}[u_i{u}_i^{\top }]=I_d,$$

has finite 4th-order absolute moments:

$${\mathbb{E}}_{u_i\sim p(u_i)}[\Vert u_i{\Vert }^4]<\infty ,$$

and the vector $v^r=\text{vec}(E^r)$v r is zero-mean and has identity covariance matrix:

$${\Sigma }_v:={\mathbb{E}}_{v^r\sim P(v^r)}[v^r{v}^{r\top }]=I_d$$

Proof. Under the vec operator, the vector $u_i$u i can be written element wise as:

$$u_i=[a_1{b}_1,a_2{b}_1,\dots a_m{b}_1,a_1{b}_2,\dots a_m{b}_n{]}^{\top }.$$

We note that all elements in the vector $u_i$u i have zero mean, and so the covariance matrix is the expectation of the outer product:

$${\Sigma }_u={\mathbb{E}}_{u_i\sim p(u_i)}[u_i{u}_i^{\top }].$$

The diagonal elements of ${\Sigma }_u$sigma u are:

$${\mathbb{E}}_{a_i,b_j}[(a_i{b}_j{)}^2]={\mathbb{E}}_{a_i}[a_i^2]{\mathbb{E}}_{b_j}[b_j^2]=1.\text{\hspace{1em}(29)}$$

the expected value of a i b j squared equals the expectation of a i squared times the expectation of b j squared, which equals one
As all elements of $a,b$a and b and $ϵ$epsilon are zero-mean, off-diagonal elements are zero:

$${\mathbb{E}}_{a_i,b_j,a_k,b_l}[a_i{b}_j{a}_k{b}_l]=0i\ne k\text{ or }j\ne l.\text{\hspace{1em}(30)}$$

Using Eqs. (29) and (30), our first result follows:

$${\Sigma }_u=I_d.$$

Now, as $u_i$u i is a vector of elements which are sums and products of variables which all have finite 4th order moments from Assumption 1, it immediately follows that $u$u has finite 4th order absolute moments.

For our final result, we can write $v^r$v r as sum of independent variables:

$$v^r=\sum _{i=1}^r\left(r^{-\frac{1}{2}}{u}_i\right)=\sum _{i=1}^r{x}_i,$$

where $x_i:=\frac{1}{\sqrt{r}}{u}_i$x i is defined as one over the square root of r times u i. As $v^r$v r is a sum of zero-mean vectors, it is also zero-mean. We use the fact that the covariance of $r$r i.i.d. random variables is equal to the sum of the individual covariances, hence

$$\begin{array}{rl}{\mathbb{E}}_{v^r}[v^r{v}^{r\top }]& =r{\mathbb{E}}_{x_i}[x_i{x}_i^{\top }],\\ & =r{\mathbb{E}}_{u_i}\left[\frac{1}{r}{u}_i{u}_i^{\top }\right],\\ & ={\mathbb{E}}_{u_i}[u_i{u}_i^{\top }],\\ & =I_d,\end{array}$$

as required. $\square$

Using Lemma 6Lemma six, we see the asymptotic Gaussian density of $v^r$v r is a standard normal:

$$g(v^r)=\frac{1}{\sqrt{(2\pi {)}^d}}\exp \left(-\frac{\Vert v^r{\Vert }^2}{2}\right)\text{\hspace{1em}(31)}$$

g of v r is one over the square root of two pi to the d power, times the exponential of negative norm v r squared over two

which is the density of $P(v)$P of v, where recall $v=\text{vec}(E)$v equals the vectorization of E, is the vectorisation of the full rank matrix $E$E.

Although $P(v^r)$P of v r does not have a density in the usual sense for low-rank $r$r, we can still approximate it with a distribution $\hat{p}(v^r)$p hat of v r by making a Taylor series expansion of its characteristic function, which always exists regardless of whether $P(v^r)$P of v r has a well-defined density or not. We now derive the 4th order Edgeworth expansion for $P(v^r)$P of v r. Our proof reveals that 3rd order cumulants control all terms in the expansion that decay at rate $\mathcal{O}\left(r^{-\frac{1}{2}}\right)$big O of r to the minus one half. As 3rd order cumulants are all zero due to symmetry in Assumption 1, the overall decay rate is controlled by $\mathcal{O}(r^{-1})$big O of r to the minus one terms associated with 4th order cumulants. It is for this reason that we obtain a faster convergence rate than the standard central limit theorem.

Lemma 7. Let Assumption 1 hold and let $v^r=\text{vec}(E^r)$v r equals the vectorization of E r and $u_i$u i be defined in Eq. (28)equation twenty-eight. Let $g(v^r)$g of v r denote the limiting Gaussian density in Eq. (31). Then, the 2nd order Edgeworth expansion of $v^r$v r is a distribution $\hat{P}(v^r)$P hat of v r defined by the approximate density:

$$\hat{p}(v^r)=g(v^r)+\frac{1}{4!r}g(v^r)\sum _{i,j,k,l}{\kappa }_{i,j,k,l}^4{H}_{i,j,k,l}(v^r),$$

where:

$$H_{i,j,k,l}(v^r):=\exp \left(\frac{\Vert v^r{\Vert }^2}{2}\right)\frac{{\partial }^4}{\partial v_i^r\partial v_j^r\partial v_k^r\partial v_l^r}\exp \left(-\frac{\Vert v^r{\Vert }^2}{2}\right)$$

is a 4th order Hermite polynomial associated with $g(v^r)$g of v r (Laplace, 1811; Hall, 1992; Withers, 2000).

Proof. We denote the characteristic function of $P(u_i)$P of u i as:

$${\varphi }_U(\omega )=\int \exp (-i{\omega }^{\top }u)dP(u),$$

and the characteristic function of $P(v^r)$P of v r as:

$${\varphi }_r(\omega )=\int \exp (-i{\omega }^{\top }u)dP(v^r).$$

Recall $v^r=\frac{1}{\sqrt{r}}{\sum }_{i=1}^r{u}_i$v r equals one over the square root of r times the sum of u i is the sum of $r$r i.i.d. copies of $\frac{1}{\sqrt{r}}{u}_i$one over the square root of r times u i. Using the scaling property of the Fourier transform, the characteristic function of $\frac{1}{\sqrt{r}}{u}_i$one over the square root of r times u i is ${\varphi }_U\left(\frac{\omega }{\sqrt{r}}\right)$phi U of omega over the square root of r. The distribution of a sum of $r$r independent random

variables is given by the $r$-fold convolution of the individual distributions. As convolution in the spatial domain corresponds to multiplication in the frequency domain, the characteristic function of $v^r$v r is (Bhattacharya & Ranga Rao, 1976)Bhattacharya and Ranga Rao, 1976:

$${\phi }_r(\omega )={\left({\phi }_U\left(\frac{\omega }{\sqrt{r}}\right)\right)}^r$$

Taking logarithms yields the log-characteristic function:

$$\begin{array}{rl}\log {\phi }_r(\omega )& =r\log \left({\phi }_U\left(\frac{\omega }{\sqrt{r}}\right)\right),\\ & =r{K}_U\left(\frac{\omega }{\sqrt{r}}\right),\end{array}$$

where $K_U(\omega ):=\log {\phi }_U(\omega )$K U of omega is defined as the log of phi U of omega. The cumulants are defined by

$${\kappa }_{i_1,\dots ,i_n}^{(n)}:=i^{-n}{\frac{{\partial }^n{K}_U(\omega )}{\partial {\omega }_{i_1}\cdots \partial {\omega }_{i_n}}|}_{\omega =0}.$$

The Edgeworth expansion proceeds by a Taylor expansion of $r{K}_U\left(\frac{\omega }{\sqrt{r}}\right)$r times K U of omega over the square root of r about $\omega =0$omega equals zero. A 4th order expansion yields:

$$\begin{array}{rl}r{K}_U\left(\frac{\omega }{\sqrt{r}}\right)\approx r{K}_U(0)& +\sqrt{r}\sum _i{\omega }_i{\kappa }_i^1+\frac{1}{2!}\sum _{i,j}{\omega }_i{\omega }_j{\kappa }_{i,j}^2\\ & +\frac{1}{3!\sqrt{r}}\sum _{i,j,k}{\omega }_i{\omega }_j{\omega }_k{\kappa }_{i,j,k}^3+\frac{1}{4!r}\sum _{i,j,k,l}{\omega }_i{\omega }_j{\omega }_k{\omega }_l{\kappa }_{i,j,k,l}^4,\end{array}$$

where $K_U(0)=0$. Under Assumption 8, $u_i$ is symmetric, hence all odd-order cumulants vanish: ${\kappa }^1={\kappa }^3=0$. The second-order cumulant satisfies

$$\sum _{i,j}{\omega }_i{\omega }_j{\kappa }_{i,j}^2=-{\omega }^{\top }{\Sigma }_u\omega ,$$

and from Lemma 6 we have ${\Sigma }_u=I$. Substituting yields:

$$r{K}_U\left(\frac{\omega }{\sqrt{r}}\right)\approx -\frac{\Vert \omega {\Vert }^2}{2}+\frac{1}{4!r}\sum _{i,j,k,l}{\omega }_i{\omega }_j{\omega }_k{\omega }_l{\kappa }_{i,j,k,l}^4.$$

Exponentiating and expanding the exponential to first-order in $1/r$ gives:

$$\begin{array}{rl}{\phi }_r(\omega )& =\exp \left(r{K}_U\left(\frac{\omega }{\sqrt{r}}\right)\right),\\ & \approx \exp \left(-\frac{\Vert \omega {\Vert }^2}{2}\right)\left(1+\frac{1}{4!r}\sum _{i,j,k,l}{\omega }_i{\omega }_j{\omega }_k{\omega }_l{\kappa }_{i,j,k,l}^4\right).\end{array}$$

Taking the inverse Fourier transform (with the convention ${\mathcal{F}}^{-1}(f)(v)=(2\pi {)}^{-d}\int e^{i{\omega }^{\top }v}f(\omega )d\omega$) yields:

$$\hat{p}(v^r)=g(v^r)+\frac{1}{4!r}\sum _{i,j,k,l}{\kappa }_{i,j,k,l}^4\frac{{\partial }^4}{\partial v_i^r\partial v_j^r\partial v_k^r\partial v_l^r}g(v^r),$$

and using the identity $H_{i,j,k,l}(v^r)=g(v^r{)}^{-1}\frac{{\partial }^4}{\partial v_i^r\partial v_j^r\partial v_k^r\partial v_l^r}g(v^r)$, we recover the stated Edgeworth density.
□

We now apply key results from Bhattacharya & Ranga Rao (1976)Bhattacharya and Ranga Rao, 1976 to bound the difference in expectation between the low-rank distribution and the Edgeworth approximation as well as the difference in expectation between the true ES Gaussian distribution and the Edgeworth approximation.

Lemma 8. Let $f(v):=f(M=\mu +\sigma \text{mat}(v))$f of v be defined as f of M equals mu plus sigma times mat of v, let $P(v)=\mathcal{N}(0,I_d)$P of v follow a standard normal distribution, $P(v^r)$P of v r be the distribution of $v^r$v r and $\hat{P}(v^r)$P hat of v r be the 2nd order Edgeworth expansion of $P(v^r)$P of v r. Let Assumptions 1 and 7 hold and let $v^r=\text{vec}(E^r)$v r equal the vectorization of E r and $u_i$u i be defined in Eq. (28). Then:

$$\begin{array}{rl}\Vert {\mathbb{E}}_{v^r\sim P(v^r)}[v^r\cdot f(v^r)]-{\mathbb{E}}_{v^r\sim \hat{P}(v^r)}[v^r\cdot f(v^r)]\Vert & =\mathcal{O}(r^{-1}),\\ \Vert {\mathbb{E}}_{v\sim P(v)}[v\cdot f(v)]-{\mathbb{E}}_{v\sim \hat{P}(v)}[v\cdot f(v)]\Vert & =\mathcal{O}(r^{-1}).\end{array}$$

Proof. From Lemma 7, we have shown that the Edgeworth expansion for $P(v^r)$P of v r is controlled by 4th order cumulants and higher, that is;

$$\hat{p}(v^r)=g(v^r)+\frac{1}{4!r}g(v^r)\sum _{i,j,k,l}{\kappa }_{i,j,k,l}^4{H}_{i,j,k,l}(v^r).\text{\hspace{1em}(32)}$$

We show that the three assumptions needed to apply Bhattacharya & Ranga Rao (1976, Theorem 20.1)Bhattacharya and Ranga Rao, Theorem 20.1 to obtain our result using Eq. (32) hold. Firstly, the boundedness assumption of the integrand holds:

$$\underset{v^r}{\sup }\frac{\Vert f(v^r)v^r\Vert }{1+\Vert v^r\Vert }\le \underset{v^r}{\sup }|f(v^r)|<\infty .$$

Secondly, the sampling regularity assumption that $u_i$u i (as defined in Eq. (28)) is zero-mean i.i.d. (satisfied under Assumption 1) with finite 4th order moments (satisfied from Lemma 6) holds. Let ${\phi }_U(\omega )$phi U of omega denote the characteristic function of $p(u)$p of u, then the final assumption we need to verify is the Cramer condition: ${\mathrm{limsup}}_{\Vert \omega \Vert \to \infty }{\phi }_U(\omega )<1$, which is satisfied from the Riemann-Lebesgue lemma Folland (1999, Theorem 8.22)Folland, 1999 because $p_0(\cdot )$p zero is absolutely continuous under Assumption 1 and hence $|{\phi }_U(\omega )|\to 0$ as $\Vert \omega \Vert \to \infty$. Our first result thus follows from applying Bhattacharya & Ranga Rao (1976, Theorem 20.1)Bhattacharya and Ranga Rao:

$$\Vert {\mathbb{E}}_{v^r\sim P(v^r)}[v^r\cdot f(v^r)]-{\mathbb{E}}_{v^r\sim \hat{P}(v^r)}[v^r\cdot f(v^r)]\Vert =\mathcal{O}(r^{-1}).$$

We now derive our second result.

$$\begin{array}{rl}{\mathbb{E}}_{v\sim \hat{P}(v)}[v\cdot f(v)]& =\int v\cdot f(v)g(v)\left(1+\frac{1}{4!r}\sum _{i,j,k,l}{\kappa }_{i,j,k,l}^4{H}_{i,j,k,l}(v)\right)dv,\\ & ={\mathbb{E}}_{v\sim P(v)}[v\cdot f(v)]-\int v\cdot f(v)g(v)\frac{1}{4!r}\sum _{i,j,k,l}{\kappa }_{i,j,k,l}^4{H}_{i,j,k,l}(v)dv,\end{array}$$

hence

$$\begin{array}{rl}\Vert {\mathbb{E}}_{v\sim P(v)}[v\cdot f(v)]-{\mathbb{E}}_{v\sim \hat{P}(v)}[v\cdot f(v)]\Vert & =\frac{1}{r}\Vert \int v\cdot f(v)\frac{1}{4!}\sum _{i,j,k,l}{\kappa }_{i,j,k,l}^4{H}_{i,j,k,l}(v)g(v)dv\Vert ,\\ & \le \frac{1}{r}\int \Vert v\Vert \cdot |f(v)|\frac{1}{4!}\sum _{i,j,k,l}|{\kappa }_{i,j,k,l}^4{H}_{i,j,k,l}(v)|g(v)dv.\end{array}$$

Now by definition, $H_{i,j,k,l}(v)$H i j k l of v is a 4th order Hermite polynomial and under Assumption 7, $|f(v)|$the absolute value of f of v is bounded, hence $\Vert v\Vert \cdot |f(v)|\frac{1}{4!r}{\sum }_{i,j,k,l}|{\kappa }_{i,j,k,l}^4{H}_{i,j,k,l}(v)|$ has polynomial growth of order 5 and is bounded by:

$$\Vert v\Vert \cdot |f(v)|\frac{1}{4!}\sum _{i,j,k,l}|{\kappa }_{i,j,k,l}^4{H}_{i,j,k,l}(v)|\le C(1+\Vert v{\Vert }^5)$$

for some finite $C>0$C greater than zero. As the expectation of a finite order polynomial under $\mathcal{N}(0,I_d)$a standard normal distribution is bounded, it thus follows:

$$\Vert {\mathbb{E}}_{v\sim P(v)}[v\cdot f(v)]-{\mathbb{E}}_{v\sim \hat{P}(v)}[v\cdot f(v)]\Vert \le \frac{1}{r}\int C(1+\Vert v{\Vert }^5)g(v)dv=\mathcal{O}(r^{-1}),$$

as required.

$\square$

Using Lemma 8Lemma eight, we have all ingredients needed derive our main about the convergence result, which follows after some simple algebra on the norm:

Theorem 4. Let Assumptions 1 and 7 hold, then:

$$\Vert {\nabla }_{\mu }J(\theta )-{\hat{g}}_{LR}^r{\Vert }_F=\mathcal{O}(r^{-1}).$$

Proof. We start by converting the Frobenius norm to vector form using Eq. (13):

$$\begin{array}{rl}\Vert {\nabla }_{\mu }J(\theta )-g_{LR}^r{\Vert }_F& =\Vert \frac{1}{\sigma }(\text{vec}({\mathbb{E}}_E[E\cdot f(W=M+\sigma E)])-\text{vec}({\mathbb{E}}_{E^r}[E^r\cdot f(W=M+\sigma E^r)]))\Vert ,\\ & =\Vert \frac{1}{\sigma }({\mathbb{E}}_E[\text{vec}(E)f(W=M+\sigma E)]-{\mathbb{E}}_{E^r}[\text{vec}(E^r)f(W=M+\sigma E^r)])\Vert ,\\ & =\Vert \frac{1}{\sigma }({\mathbb{E}}_v[vf(v)]-{\mathbb{E}}_{v^r}[v^{r}f(v^r)])\Vert ,\end{array}$$

where $f(v):=f(M=\mu +\sigma \text{mat}(v))$ and $v=\text{vec}(E)$ is the vectorisation of variable $E$, which is distributed as $v\sim P(v):=\mathcal{N}(0,I_d)$. Let $\hat{P}(v)$ be the distribution for the 2nd order Edgeworth expansion, which we derived in Lemma 7Lemma seven. Since $\hat{P}(v^r)$ and $\hat{P}(v)$ are identified as the same Edgeworth-expanded distribution on ${\mathbb{R}}^d$, we may equivalently write:

$${\mathbb{E}}_{v^r\sim \hat{P}(v^r)}[v^{r}f(v^r)]={\mathbb{E}}_{v\sim \hat{P}(v)}[v^{r}f(v)],$$

hence:

$$\begin{array}{rl}{\mathbb{E}}_v[vf(v)]-{\mathbb{E}}_{v^r}[v^{r}f(v^r)]& ={\mathbb{E}}_v[vf(v)]-{\mathbb{E}}_{v\sim \hat{P}(v)}[vf(v)]+{\mathbb{E}}_{v^r\sim \hat{P}(v^r)}[v^{r}f(v^r)]-{\mathbb{E}}_{v^r}[v^{r}f(v^r)],\\ ⟹\Vert {\nabla }_{\mu }J(\theta )-{\hat{g}}_{LR}^r{\Vert }_F& \le \frac{1}{\sigma }\Vert {\mathbb{E}}_v[vf(v)]-{\mathbb{E}}_{v\sim \hat{P}(v)}[vf(v)]\Vert \\ & +\frac{1}{\sigma }\Vert {\mathbb{E}}_{v^r\sim \hat{P}(v^r)}[v^{r}f(v^r)]-{\mathbb{E}}_{v^r}[v^{r}f(v^r)]\Vert .\end{array}$$

Applying Lemma 8Lemma eight to each bound yields our desired result:

$$\Vert {\nabla }_{\mu }J(\theta )-{\hat{g}}_{LR}^r{\Vert }_F=\mathcal{O}(r^{-1}).$$

$\square$

D.1 Mean Field Score Function Approximator

We will use $n$th order Bessel functions of the second kind $K_n(z)$K n of z (Basset, 1888; Macdonald, 1899; Watson, 1944) , which are conveniently represented by the integral equations:

$$K_n(z)={\int }_0^{\infty }\exp (-z\cosh \theta )\cosh (n\theta )d\theta .$$

Bessel functions are the solutions to systems of differential equations that occur naturally in phenomena where there is strong radial symmetry, typically involving the propagation of spherical waves from points like the ripples formed from water droplets (Whitham, 1999) . For our setting, Bessel functions describe the

probability density of the product of rotationally invariant random variables, whose solution is analogous to the interference pattern of two spherical wave propagators.

Using the representation, we find the derivative of the zeroth order function takes the recursive form:

$$\frac{d{K}_0(z)}{dz}=-{\int }_0^{\infty }\exp (-z\cosh \theta )\cosh (\theta )d\theta =-K_1(z).\text{\hspace{1em}(33)}$$

More generally, the derivative of the $n$th order Bessel function is Watson (1944, Section 3.71, Eq. 4)Watson, 1944:

$$\frac{d{K}_n(z)}{dz}=\frac{n}{z}{K}_n(z)-K_{n+1}(z).\text{\hspace{1em}(34)}$$

D.2 Derivation of Mean-field Approximators

To derive a mean-field approximation, we assume that the elements of $A$A and $B$B are drawn independently from the set of generalised Gaussian distributions (GGDs):

Assumption 8. Assume each element $a_{i,j}\sim \mathcal{GG}(s,p)$a i j distributed as G G of s and p and $b_{i,j}\sim \mathcal{GG}(s,p)$b i j distributed as G G of s and p of $A$A and $B$B is independently distributed according to the zero-mean generalised Gaussian distribution $\mathcal{GG}(s,p)$ with density:

$$\mathcal{GG}(x|s,p)=\frac{p}{2s\Gamma \left(\frac{1}{p}\right)}\exp \left(-{\left|\frac{x}{s}\right|}^p\right),$$

where $0<s<\infty$s is between zero and infinity is the scale parameter, $p>0$p is greater than zero the shape parameter and $\Gamma (\cdot )$gamma is the gamma function.

We observe common distributions emerge from the set of GGDs including the Laplace for $p=1$p equals 1, the Gaussian for $p=2$p equals 2 and the uniform over $[-s,+s]$ in the limit $p\to \infty$as p approaches infinity.

If we make the assumption that all elements of $E$E are independent (this is true as $r$r grows) then we can write $p(E)\approx \hat{p}(E):={\prod }_{i=1}^m{\prod }_{j=1}^{n}p(E_{i,j})$ as the product of the marginal distributions. Under this approximation, the score function can be defined element-wise as:

$$[{\nabla }_E\log p(E){]}_{i,j}\approx \hat{S}(E_{i,j}):={\partial }_{E_{i,j}}\log p(E_{i,j}).$$

Using this approximation we apply the score function $\hat{S}(\cdot )$S hat element-wise to the matrix $E$E:

$$g_{\text{LR}}\approx {\hat{g}}_{\text{MF}}:=-\frac{1}{\sigma }{\mathbb{E}}_{E\sim p(E)}\left[f(W=M+\sigma E)\cdot \hat{S}\odot (E)\right].$$

For $r=1$r equals 1, $\hat{S}(\cdot )$S hat has a convenient analytic form for all members of the set of GGDs:

Theorem 5. Let Assumption 8 hold and $r=1$r equals 1. Then the distribution over marginals $p(E_{i,j})$ is:

$$p(E_{i,j})=\frac{p}{{\left(s\Gamma \left(\frac{1}{p}\right)\right)}^2}{K}_0\left(\frac{2|E_{i,j}{|}^{\frac{p}{2}}}{s^p}\right),\text{\hspace{1em}(35)}$$

where $K_0(\cdot )$K sub zero is the zeroth-order modified Bessel function of the second kind and the marginal score function is defined element-wise as:

$$\hat{S}(E_{i,j})=-\frac{K_1\left(\frac{2|E_{i,j}{|}^{\frac{p}{2}}}{s^p}\right)}{K_0\left(\frac{2|E_{i,j}{|}^{\frac{p}{2}}}{s^p}\right)}\cdot \frac{p|E_{i,j}{|}^{\frac{p}{2}-1}\text{sign}(E_{i,j})}{s^p}.$$

Proof. For $r=1$, we denote the elements of vector $A$ as $a_i$ and elements of vector $B$ as $b_j$, then the elements of matrix $E=A{B}^{\top }$ are: $E_{i,j}=a_i{b}_j$. We now derive the distribution of the unnormalised variables: $E_{i,j}$

using the formula for the distribution of the product of two independent random variables (Rohatgi, 1976; Grimmett & Stirzaker, 1993)Rohatgi 1976, and Grimmett and Stirzaker 1993:

$$\begin{array}{rl}p(E_{i,j})& ={\int }_{-\infty }^{\infty }p(a_i)p\left(b_j=\frac{E_{i,j}}{a_i}\right)\frac{1}{|a_i|}d{a}_i,\\ & ={\left(\frac{p}{2s\Gamma \left(\frac{1}{p}\right)}\right)}^2{\int }_{-\infty }^{\infty }\exp \left(-{\left|\frac{a_i}{s}\right|}^p\right)\exp \left(-{\left|\frac{E_{i,j}}{a_{i}s}\right|}^p\right)\frac{1}{|a_i|}d{a}_i,\\ & =2{\left(\frac{p}{2s\Gamma \left(\frac{1}{p}\right)}\right)}^2{\int }_0^{\infty }\exp \left(-{\left|\frac{a_i}{s}\right|}^p\right)\exp \left(-{\left|\frac{E_{i,j}}{a_{i}s}\right|}^p\right)\frac{1}{|a_i|}d{a}_i,\end{array}$$

where we have used symmetry of the integrand about $0$zero to derive the final line. Now, making the substitution $x={\left(\frac{a_i}{s}\right)}^p$x equals a sub i over s, to the power of p, we have:

$$\frac{d{a}_i}{dx}=\frac{s{x}^{\frac{1}{p}-1}}{p},a_i=s{x}^{\frac{1}{p}}$$

hence:

$$p(E_{i,j})=\frac{p}{{\left(s\Gamma \left(\frac{1}{p}\right)\right)}^2}\frac{1}{2}{\int }_0^{\infty }\exp \left(-x-\frac{1}{x}\frac{|E_{i,j}{|}^p}{s^{2p}}\right)\frac{1}{x}dx.$$

Now, we use the identity (Temme, 1996, Theorem 9.42)from Temme 1996:

$$K_0(z)=\frac{1}{2}{\int }_0^{\infty }\exp \left(-x-\frac{z^2}{4x}\right)\frac{1}{x}dx,$$

with $z=\frac{2|E_{i,j}{|}^{\frac{p}{2}}}{s^p}$z equals two times the absolute value of E sub i j to the power of p over two, all over s to the p to yield:

$$p(E_{i,j})=\frac{p}{{\left(s\Gamma \left(\frac{1}{p}\right)\right)}^2}{K}_0\left(\frac{2|E_{i,j}{|}^{\frac{p}{2}}}{s^p}\right),$$

as required for Eq. (35). Now we derive the marginal score function by applying the chain rule:

$$\begin{array}{rl}{\partial }_{E_{i,j}}\log p(E_{i,j})& ={\partial }_{E_{i,j}}\log K_0\left(\frac{2|E_{i,j}{|}^{\frac{p}{2}}}{s^p}\right),\\ & ={\partial }_z\log K_0\left(z=\frac{2|E_{i,j}{|}^{\frac{p}{2}}}{s^p}\right){\partial }_{E_{i,j}}\frac{2|E_{i,j}{|}^{\frac{p}{2}}}{s^p},\\ & ={\partial }_z\log K_0\left(z=\frac{2|E_{i,j}{|}^{\frac{p}{2}}}{s^p}\right)\frac{p|E_{i,j}{|}^{\frac{p}{2}-1}\text{sign}(E_{i,j})}{s^p},\\ & =\frac{{\partial }_z{K}_0\left(z=\frac{2|E_{i,j}{|}^{\frac{p}{2}}}{s^p}\right)}{K_0\left(z=\frac{2|E_{i,j}{|}^{\frac{p}{2}}}{s^p}\right)}\cdot \frac{p|E_{i,j}{|}^{\frac{p}{2}-1}\text{sign}(E_{i,j})}{s^p},\\ & =-\frac{K_1\left(\frac{2|E_{i,j}{|}^{\frac{p}{2}}}{s^p}\right)}{K_0\left(\frac{2|E_{i,j}{|}^{\frac{p}{2}}}{s^p}\right)}\cdot \frac{p|E_{i,j}{|}^{\frac{p}{2}-1}\text{sign}(E_{i,j})}{s^p},\end{array}$$

where we have used the identity ${\partial }_z{K}_0(x)=-K_1(x)$ from Eq. (33). $\square$

For $r>1$r greater than one we can derive $\hat{S}(\cdot )$S hat for the Gaussian sampling case:

Theorem 6. Let Assumption 8 hold and $p=2$p equals two. Then the distribution over marginals $p(E_{i,j})$p of E i comma j is:

$$p(E_{i,j})=\frac{2\sqrt{r}|\sqrt{r}{E}_{i,j}{|}^{\frac{r-1}{2}}}{s^{r+1}\sqrt{\pi }\Gamma \left(\frac{r}{2}\right)}\cdot K_{\frac{r-1}{2}}\left(\frac{2|\sqrt{r}{E}_{i,j}|}{s^2}\right).$$

and the score function is (for $E_{i,j}\ne 0$E i j not equal to zero):

$$\hat{S}(E_{i,j})=\frac{r-1}{E_{i,j}}-\frac{2\sqrt{r}\text{sign}(E_{i,j})}{s^2}\frac{K_{\frac{r+1}{2}}\left(\frac{2|\sqrt{r}{E}_{i,j}|}{s^2}\right)}{K_{\frac{r-1}{2}}\left(\frac{2|\sqrt{r}{E}_{i,j}|}{s^2}\right)}.$$

Proof. Each element $E_{i,j}$E i j is the sum of $r$r independent variables $u_{i,j,l}:=a_{i,l}{b}_{j,l}$u i j l, defined as a i l times b j l distributed according to Eq. (35) with $p=2$p equals two:

$$E_{i,j}=\frac{1}{\sqrt{r}}\sum _{l=1}^r{a}_{i,l}{b}_{j,l}=\frac{1}{\sqrt{r}}\sum _{l=1}^r{u}_{i,j,l}.$$

Let $Z_{i,j}=\sqrt{r}{E}_{i,j}$Z i j equals the square root of r times E i j, hence:

$$Z_{i,j}=\sum _{l=1}^r{u}_{i,j,l}.$$

We first find the density $p(Z_{i,j})$p of Z i j. From Eq. (35), the distribution of each $u_{i,j,l}$u i j l is:

$$p(u_{i,j,l})=\frac{2}{s^2\pi }{K}_0\left(\frac{2|u_{i,j,l}|}{s^2}\right)$$

We use the fact that the PDF of a sum of $r$r independent random variables (i.e. $Z_{i,j}$Z i j) is given by the $r$r-fold convolution of the individual PDFs. As convolution in the spatial domain is equal to multiplication in the frequency domain, the PDF $p(Z_{i,j})$p of Z i j follows by taking Fourier transform of $p(u_{i,j,l})$p of u i j l, taking the power $r$r and then taking the inverse Fourier transform:

$$p(Z_{i,j})={\left(\frac{2}{s^2\pi }\right)}^r{\mathcal{F}}^{-1}\left[{\left(\mathcal{F}\left[K_0\left(\frac{2|\cdot |}{s^2}\right)\right]\right)}^r\right](Z_{i,j}),$$

where recall from Section A with $d=1$d equals one, $\mathcal{F}[f](\omega ):=\int f(x)\exp (-i\omega x)dx$F of f of omega is the integral of f of x times e to the minus i omega x denotes the Fourier transform and ${\mathcal{F}}^{-1}[\tilde{f}](x):=\frac{1}{2\pi }\int \tilde{f}(\omega )\exp (i\omega x)d\omega$the inverse Fourier transform. Taking the Fourier transform of the Bessel function:

$$\begin{array}{rl}\mathcal{F}\left[K_0\left(\frac{2|\cdot |}{s^2}\right)\right](\omega )& =\int \exp (-i\omega x)K_0\left(\frac{2|x|}{s^2}\right)dx,\\ & =\int \cos (\omega x)K_0\left(\frac{2|x|}{s^2}\right)dx-i\int \sin (\omega x)K_0\left(\frac{2|x|}{s^2}\right)dx,\\ & =\int \cos (\omega x)K_0\left(\frac{2|x|}{s^2}\right)dx,\\ & =2{\int }_0^{\infty }\cos (\omega x)K_0\left(\frac{2x}{s^2}\right)dx,\end{array}\text{\hspace{1em}(36)}$$

where we have used the fact that $K_0\left(\frac{2|x|}{s^2}\right)$K zero is an even function of $x$x and so its integral with $\sin (\omega x)$sine of omega x in the second line is zero. Using a standard result, we can evaluate the integral in Eq. (36) Gradshteĭn et al. (2015, 6.671 Integral 14)Gradshteĭn and colleagues:

$$\mathcal{F}\left[K_0\left(\frac{2|\cdot |}{s^2}\right)\right](\omega )=\frac{\pi }{\sqrt{{\omega }^2+{\left(\frac{2}{s^2}\right)}^2}},$$

hence:

$$\begin{array}{rl}p(Z_{i,j})& ={\left(\frac{2}{s^2\pi }\right)}^r{\mathcal{F}}^{-1}\left[\frac{{\pi }^r}{{\left({\omega }^2+{\left(\frac{2}{s^2}\right)}^2\right)}^{\frac{r}{2}}}\right](Z_{i,j}),\\ & ={\left(\frac{2}{s^2}\right)}^r{\mathcal{F}}^{-1}\left[{\left({\omega }^2+{\left(\frac{2}{s^2}\right)}^2\right)}^{-\frac{r}{2}}\right](Z_{i,j}),\\ & ={\left(\frac{2}{s^2}\right)}^r\frac{1}{2\pi }\int \exp (i\omega Z_{i,j}){\left({\omega }^2+{\left(\frac{2}{s^2}\right)}^2\right)}^{-\frac{r}{2}}d\omega ,\\ & ={\left(\frac{2}{s^2}\right)}^r\frac{1}{2\pi }(\int \cos (\omega Z_{i,j}){\left({\omega }^2+{\left(\frac{2}{s^2}\right)}^2\right)}^{-\frac{r}{2}}d\omega \\ & +i\int \sin (\omega Z_{i,j}){\left({\omega }^2+{\left(\frac{2}{s^2}\right)}^2\right)}^{-\frac{r}{2}}d\omega ),\\ & ={\left(\frac{2}{s^2}\right)}^r\frac{1}{2\pi }\int \cos (\omega Z_{i,j}){\left({\omega }^2+{\left(\frac{2}{s^2}\right)}^2\right)}^{-\frac{r}{2}}d\omega ,\\ & ={\left(\frac{2}{s^2}\right)}^r\frac{1}{\pi }{\int }_0^{\infty }\cos (\omega Z_{i,j}){\left({\omega }^2+{\left(\frac{2}{s^2}\right)}^2\right)}^{-\frac{r}{2}}d\omega ,\end{array}\text{\hspace{1em}(37)}$$

where we have used the fact that the integrand is an even function and so its integral with $\sin (\omega Z_{i,j})$sine of omega Z i j is zero to derive the penultimate line. To evaluate the integral in Eq. (37) we apply Gradshteǐn et al. (2015, 3.771 Integral 2)Gradshtein and colleagues:

$$\begin{array}{rl}p(Z_{i,j})& ={\left(\frac{2}{s^2}\right)}^r\cdot \frac{1}{\sqrt{\pi }\Gamma \left(\frac{r}{2}\right)}{\left(\frac{s^2|Z_{i,j}|}{4}\right)}^{\frac{r-1}{2}}\cdot K_{\frac{r-1}{2}}\left(\frac{2|Z_{i,j}|}{s^2}\right),\\ & =\frac{2|Z_{i,j}{|}^{\frac{r-1}{2}}}{s^{r+1}\sqrt{\pi }\Gamma \left(\frac{r}{2}\right)}\cdot K_{\frac{r-1}{2}}\left(\frac{2|Z_{i,j}|}{s^2}\right).\end{array}$$

Using the transformation of variables $E_{i,j}=\frac{1}{\sqrt{r}}{Z}_{i,j}$E i j equals one over the square root of r times Z i j yields our desired results:

$$\begin{array}{rl}p(E_{i,j})& =\sqrt{r}p(Z_{i,j}=\sqrt{r}{E}_{i,j}),\\ & =\frac{2\sqrt{r}|\sqrt{r}{E}_{i,j}{|}^{\frac{r-1}{2}}}{s^{r+1}\sqrt{\pi }\Gamma \left(\frac{r}{2}\right)}\cdot K_{\frac{r-1}{2}}\left(\frac{2|\sqrt{r}{E}_{i,j}|}{s^2}\right).\end{array}$$

Now, we derive the score function:

$$\begin{array}{rl}{\partial }_{E_{i,j}}\log p(E_{i,j})& =\frac{r-1}{2}\cdot {\partial }_{E_{i,j}}\log |\sqrt{r}{E}_{i,j}|+{\partial }_{E_{i,j}}\log K_{\frac{r-1}{2}}\left(\frac{2|\sqrt{r}{E}_{i,j}|}{s^2}\right),\\ & =\frac{r-1}{2E_{i,j}}+\frac{\frac{2{\partial }_{E_{i,j}}|\sqrt{r}{E}_{i,j}|}{s^2}{\partial }_x{K}_{\frac{r-1}{2}}\left(x=\frac{2|\sqrt{r}{E}_{i,j}|}{s^2}\right)}{K_{\frac{r-1}{2}}\left(\frac{2|\sqrt{r}{E}_{i,j}|}{s^2}\right)},\\ & =\frac{r-1}{2E_{i,j}}+\frac{\frac{2\sqrt{r}\text{sign}(E_{i,j})}{s^2}{\partial }_x{K}_{\frac{r-1}{2}}\left(x=\frac{2|\sqrt{r}{E}_{i,j}|}{s^2}\right)}{K_{\frac{r-1}{2}}\left(\frac{2|\sqrt{r}{E}_{i,j}|}{s^2}\right)},\end{array}$$

Now, from Eq. (34)34 for $E_{i,j}\ne 0$E i j not equal to zero:

$$\begin{array}{rl}\frac{{\partial }_x{K}_{\frac{r-1}{2}}(x)}{K_{\frac{r-1}{2}}(x)}& =\frac{\frac{r-1}{2x}{K}_{\frac{r-1}{2}}(x)-K_{\frac{r+1}{2}}(x)}{K_{\frac{r-1}{2}}(x)},\\ & =\frac{r-1}{2x}-\frac{K_{\frac{r+1}{2}}(x)}{K_{\frac{r-1}{2}}(x)},\\ ⟹{\partial }_{E_{i,j}}\log p(E_{i,j})& =\frac{r-1}{2E_{i,j}}+\frac{(r-1)\text{sign}(E_{i,j})}{2|E_{i,j}|}-\frac{2\sqrt{r}\text{sign}(E_{i,j})}{s^2}\frac{K_{\frac{r+1}{2}}\left(\frac{2|\sqrt{r}{E}_{i,j}|}{s^2}\right)}{K_{\frac{r-1}{2}}\left(\frac{2|\sqrt{r}{E}_{i,j}|}{s^2}\right)},\\ & =\frac{r-1}{2E_{i,j}}+\frac{(r-1)}{2E_{i,j}}-\frac{2\sqrt{r}\text{sign}(E_{i,j})}{s^2}\frac{K_{\frac{r+1}{2}}\left(\frac{2|\sqrt{r}{E}_{i,j}|}{s^2}\right)}{K_{\frac{r-1}{2}}\left(\frac{2|\sqrt{r}{E}_{i,j}|}{s^2}\right)},\\ & =\frac{r-1}{E_{i,j}}-\frac{2\sqrt{r}\text{sign}(E_{i,j})}{s^2}\frac{K_{\frac{r+1}{2}}\left(\frac{2|\sqrt{r}{E}_{i,j}|}{s^2}\right)}{K_{\frac{r-1}{2}}\left(\frac{2|\sqrt{r}{E}_{i,j}|}{s^2}\right)},\end{array}$$

as required. $\square$

E EGGROLL Speed

All timings were done on a single GPU on a GH200 (equivalent to a single H100) for a linear model with dimension 8192 in bfloat16, allowing a maximum batch size of 1024. For the graph in Figure 2a, we pre-generate the noises instead of integrating the noise generation into the forward pass.

Figure 7: Relative speed of EGGROLL, when including jax noise regeneration.

In Figure 7, we consider the impact of regenerating noises on-the-fly using jax PRNG. The darker area and value in parenthesis for EGGROLL and OpenES indicate the speed when regenerating noises on-the-fly, while the full bar indicates the speed when the noises are already generated.

We regenerate noises on the fly in our primary jax codebase, but pre-generating the EGGROLL perturbations beforehand is also a practical possibility since low-rank perturbations only require a small amount of memory, proportional to the square root of the size of the original parameter matrices.

F Arithmetic Intensity Analysis

In this section, we derive the arithmetic intensity of standard batched inference, Gaussian matrix ES, and EGGROLL. We calculate arithmetic intensity as the number of operations divided by the total number of bytes read from or written to. For context, for the (b)float16 datatype on an H100 GPU, there are approximately 1000 teraFLOPS of compute (without sparsity) and 3.35 TB/s of GPU memory bandwidth, meaning that the roofline threshold is approximately 300 ops/byte, defined as the minimum for computation needed for it to be the bottleneck instead of memory movement.

In the following subsections, we are considering a single linear layer with mean parameter $M\in {\mathbb{R}}^{d_{out}\times d_{in}}$M in R dimensions d out by d in and a batch of inputs $u\in {\mathbb{R}}^{B\times d_{in}}$u in R dimensions B by d in. All operations occur with a precision of $s$s bytes per element.

F.1 Arithmetic Intensity of Standard Batched Inference

In standard batched inference, we wish to simply calculate $u{M}^T$u times M transpose. The total bytes read as input are $B\times d_{in}\times s$B times d in times s (for $u$u) and $d_{out}\times d_{in}\times s$ (for $M$M), and the total bytes written as output are $B\times d_{out}\times s$. The total number of operations are $B\times d_{in}\times d_{out}\times 2$ since matrix multiplication requires both multiplications and additions for each element of $u$u across all of $d_{out}$d out. Therefore, the arithmetic intensity is:

$$\frac{B\times d_{in}\times d_{out}\times 2}{B\times d_{in}\times s+B\times d_{out}\times s+d_{out}\times d_{in}\times s}.$$

When $s=2$s equals 2 (for (b)float16) and $d_{out}=d_{in}=m$d out and d in both equal m, the arithmetic intensity simplifies to

$$\frac{Bm}{2B+m}.$$

The batch size needed to achieve a desired arithmetic intensity of $A$A is derived as follows:

$$\begin{array}{rl}Bm& =2AB+Am\\ Bm-2AB& =Am\\ B& =\frac{Am}{m-2A}\end{array}$$

Therefore, achieving an arithmetic intensity of 300 ops/byte with $m=8192$m equals 8192 requires a minimum batch size of 324.

F.2 Arithmetic Intensity of Gaussian Matrix ES

In Gaussian matrix ES, we assume access to pre-generated perturbations of shape ${\mathbb{R}}^{B\times d_{out}\times d_{in}}$of dimensions B by d out by d in. The total bytes read as input are $B\times d_{in}\times s$ (for $u$u) and $B\times d_{out}\times d_{in}\times s$ (for $M$M), and the total bytes written as output are $B\times d_{out}\times s$. Otherwise, the total number of operations is identical to standard batched inference, giving us an arithmetic intensity of

$$\frac{B\times d_{in}\times d_{out}\times 2}{B\times d_{in}\times s+B\times d_{out}\times s+B\times d_{out}\times d_{in}\times s}=\frac{d_{in}\times d_{out}\times 2}{d_{in}\times s+d_{out}\times s+d_{out}\times d_{in}\times s}.$$

When $s=2$s equals 2 (for (b)float16) and $d_{out}=d_{in}=m$d out and d in both equal m, the arithmetic intensity simplifies to

$$\frac{m}{2+m}.$$

This means that arithmetic intensity is always strictly less than 1, regardless of batch size or dimensionality. The common way to increase arithmetic intensity is to bring it closer to standard batched inference, reusing the same perturbation across multiple inputs. For instance, when $m=8192$m equals 8192, achieving an arithmetic intensity of 300 ops/byte requires that each perturbation is reused at least 324 times, and smaller values of $m$m need to be reused even more often.

F.3 Arithmetic Intensity of EGGROLL

For EGGROLL, we assume access to the pre-generated decomposed perturbations $A\in {\mathbb{R}}^{B\times d_{\text{out}}\times r}$A in R dimensions B by d out by r and $B\in {\mathbb{R}}^{B\times d_{\text{in}}\times r}$B in R dimensions B by d in by r. Therefore, the bytes read as pure input are $B\times d_{\text{in}}\times s+B\times (d_{\text{in}}+d_{\text{out}})\times r\times s+d_{\text{out}}\times d_{\text{in}}\times s$ and the bytes written as pure output are $B\times d_{\text{out}}\times s$. However, the efficient low-rank perturbation calculation requires writing and reading an intermediate matrix of shape $B\times r$, so the total bytes read are

$$(B\times d_{\text{in}}+B\times (d_{\text{in}}+d_{\text{out}}+2)\times r+d_{\text{out}}\times d_{\text{in}}+B\times d_{\text{out}})\times s.$$

The total number of operations includes the amount for standard batch inference, $B\times d_{\text{in}}\times d_{\text{out}}\times 2$, along with the rank-$r$r perturbations, $B\times (d_{\text{in}}+d_{\text{out}})\times r\times 2$, and the final sum between the main calculation and perturbation $B\times d_{\text{out}}$. Therefore, the arithmetic intensity is

$$\frac{B\times d_{\text{in}}\times d_{\text{out}}\times 2+B\times (d_{\text{in}}+d_{\text{out}})\times r\times 2+B\times d_{\text{out}}}{(B\times d_{\text{in}}+B\times (d_{\text{in}}+d_{\text{out}}+2)\times r+d_{\text{out}}\times d_{\text{in}}+B\times d_{\text{out}})\times s}.$$

When $s=2$s equals two (for (b)float16) and $d_{\text{out}}=d_{\text{in}}=m$d out and d in both equal m, the arithmetic intensity simplifies to

$$\begin{array}{rl}& \frac{Bm+2Br+\frac{B}{2}}{B+Br(2+\frac{2}{m})+m+B}\\ =& \frac{m+2r+\frac{1}{2}}{2+r(2+\frac{2}{m})+\frac{m}{B}}.\end{array}$$

The batch size needed to achieve a desired arithmetic intensity of $A$A is derived as follows:

$$\begin{array}{rl}2A+rA(2+\frac{2}{m})+\frac{Am}{B}& =m+2r+\frac{1}{2}\\ \frac{Am}{B}& =m+2r+\frac{1}{2}-2A-rA(2+\frac{2}{m})\\ B& =\frac{Am}{m-2A+2r+\frac{1}{2}-rA(2+\frac{2}{m})}\end{array}$$

Note that the only difference with the critical batch size of standard batched inference is the additional $2r+\frac{1}{2}-rA(2+\frac{2}{m})$ in the denominator. Therefore, achieving an arithmetic intensity of 300 ops/byte with $m=8192$ and $r=1$ requires a minimum batch size of 352, compared to 324 for standard batched inference. This means that EGGROLL can saturate compute with unique perturbations per input, unlike Gaussian matrix ES.

Note that there is an overhead of $Bm(4r+1)$B m times, 4 r plus 1 flops relative to standard batched inference, resulting in an additional compute rate of $\frac{Bm(4r+1)}{2B{m}^2}=\frac{4r+1}{2m}$, which is effectively negligible for large enough matrices.

G EGG Architecture

In the following section, we detail the design of our EGG model, which follows the high-level structure of modern pre-layernorm decoder-only language models, but replaces self-attention with a modified minGRU and standard layernorms with a custom variant to enable pure integer training. See Algorithm 2 for an overview of the forward pass of the EGG architecture.

Algorithm 2 EGG forward pass

Require: Input token $t\in {\mathbb{U}}_8$, input state $s\in {\mathbb{I}}_8^{l\times D}$, network parameters $\theta$
Ensure: Output vector $y\in {\mathbb{I}}_8^D$ and output state $s^{\prime }\in {\mathbb{I}}_8^{l\times D}$
$s^{\prime }\leftarrow {\mathbb{I}}_8^{l\times D}$ initialised to 0
$y\leftarrow \text{EMBED}({\theta }_{\text{emb}},t)$
for $i\in \{0,\dots ,l-1\}$ do
$y^{\prime },s_i^{\prime }\leftarrow \text{GRU}({\theta }_{\text{gru},i},\text{LN}({\theta }_{\text{ln}1,i},y),s_i)$
$y\leftarrow {\mathbb{I}}_8({\mathbb{I}}_{32}(y^{\prime })+{\mathbb{I}}_{32}(y))$
$y^{\prime }\leftarrow \text{MLP}({\theta }_{\text{mlp},i},\text{LN}({\theta }_{\text{ln}2,i},y))$
$y\leftarrow {\mathbb{I}}_8({\mathbb{I}}_{32}(y^{\prime })+{\mathbb{I}}_{32}(y))$
end for
$y\leftarrow \text{LN}({\theta }_{\text{lnout},i},y)@{\theta }_{\text{head}}^T$

G.1 Motivation

Since EGGROLL does not rely on gradients, we can explicitly design a language model architecture to be efficient and hardware-friendly at inference time. In particular, we design EGG under the following constraints to emphasise the flexibility of EGGROLL:

Pure Integer Training: On H100 systems, int8 is the fastest datatype and int8 matrix multiplication with int32 accumulation is the fastest tensor core operation. Furthermore, integer datatypes are much simpler to implement in hardware, providing massive energy savings for high-throughput systems (Horowitz, 2014). Therefore, we keep all weights in int8 and all activations in integer formats, never casting to floating point at any point during training. This stands in contrast to the standard approach for language model quantisation through “quantisation aware training” with backpropagation, where floating point activations are still necessary (Wang et al., 2023).

Nonlinear RNN: Modern language models use sequence-parallel architectures like Transformers and SSMs, since they enable stable gradients without backpropagation through time. However, most of these architectures cannot handle simple state tracking (Merrill et al., 2024), whereas classic recurrent networks like LSTMs and GRUs can do so with a single layer. Since EGGROLL does not require backpropagation through time, we can train on unbounded sequence lengths (Li et al., 2023) with nonlinear RNNs of broader complexity classes. Specifically, we develop a variant of the minGRU model (Heck & Salem, 2017) that performs all operations in integer formats.

Removal of all Activation Functions: Inspired by Foerster (2017)prior work, we remove all activation functions, like the rectified linear unit and hyperbolic tangent, due to the nonlinearity present in the int8 datatype. Specifically, the saturated addition of int8 values provides sufficient nonlinearity due to the implicit clipping of values to the int8 dynamic range, which evolution strategies can exploit.

G.2 Notation and Operations

We use the constant $l\in {\mathbb{Z}}^{+}$l in the set of positive integers to denote the number of layers of the model and $D=4^d$D equals four to the d as the hidden dimension of the model, where $d\in {\mathbb{Z}}^{+}$d in the set of positive integers.

We use ${\mathbb{I}}_n$I sub n to denote an $n$-bit signed integer and ${\mathbb{U}}_n$U sub n to denote an $n$-bit unsigned integer. We denote casting vector $\vec{u}$ to format ${\mathbb{I}}_n$ as ${\mathbb{I}}_n(\vec{u})$I sub n of u, which implicitly includes clipping to the bounds of the datatype. To ensure symmetry between positive and negative values of each datatype, we consider the value $-2^{n-1}$minus two to the n minus one to be invalid for datatype ${\mathbb{I}}_n$; for instance, for 8-bit signed integers we only allows value from -127 to 127.

We use the following operations:

• $\vec{u}@M$ indicating scaled vector-matrix multiplication of ${\mathbb{I}}_8^n\times {\mathbb{I}}_8^{n,m}\to {\mathbb{I}}_8^m$u at M, indicating scaled vector-matrix multiplication from i 8 n and i 8 n by m to i 8 m, corresponding to int8 tensor core multiplication with int32 accumulation and scaling. The details of this operation are described in Section G.4.
• $a\cdot b$a dot b indicates dot product with int32 accumulation, ${\mathbb{I}}_8^n\times {\mathbb{I}}_8^n\to {\mathbb{I}}_{32}$from i 8 n and i 8 n to i 32, and $a\odot b$ indicates the Hadamard (elementwise) product.
• Standard integer operations: $+$ for addition, $-$ for subtraction, and $\odot$ for element-wise multiplication.
• $|u|$ indicates taking the element-wise absolute value of $u$u, ${\mathbb{I}}^n\to {\mathbb{I}}^n$.
• $\text{sign}(u)$ indicates taking the element-wise sign of $u$u, giving 1 for positive values, -1 for negative values, and 0 for zero.
• $\text{sum}(u)$ indicates taking the sum of all elements in $u$u (casting to ${\mathbb{I}}_{32}$ to prevent overflow): ${\mathbb{I}}^n\to {\mathbb{I}}_{32}$.
• $u\gg n$ indicates an elementwise bitwise right shift by $n$n, which is typically equivalent to $2^{-n}u$. Similarly, $u\ll n$ indicates a bitwise left shift by $n$, which is typically equivalent to $2^{n}u$.
• Square-bracket indexing. For instance $M[i,j]$ extracts the element at index $i$ in axis 0 and index $j$ in axis 1, following the zero-based indexing convention.

G.3 Parameter Initialisation

The standard initialisation for matrix parameters in our model is rounding 16 times a sample from the standard normal, and casting to ${\mathbb{I}}_8$. This can be precomputed on a CPU since this is only done once at the start of training.

The egg model has the following parameters (where an additional subscript of $i$i indicates that there is a version of this parameter for each layer of the model):

• ${\theta }_{\text{emb}}\in {\mathbb{I}}_8^{256\times D}$theta emb in i 8 of size 256 by D, following standard initialisation.
• ${\theta }_{\text{head}}\in {\mathbb{I}}_8^{256\times D}$theta head in i 8 of size 256 by D, following standard initialisation.
• ${\theta }_{\text{lnout}}\in {\mathbb{I}}_8^D$theta l n out in i 8 D, initialised to 16 for each element.
• ${\theta }_{\text{ln}1,i},{\theta }_{\text{ln}2,i}\in {\mathbb{I}}_8^D$theta l n 1 i and theta l n 2 i in i 8 D, initialised to 16 for each element
• ${\theta }_{\text{mlp},i,1}\in {\mathbb{I}}_8^{4D\times D}$ and ${\theta }_{\text{mlp},i,2}\in {\mathbb{I}}_8^{D\times 4D}$theta m l p i 1 of size 4 D by D and theta m l p i 2 of size D by 4 D, both in i 8, following standard initialisation.
• ${\theta }_{\text{GRU},i,[\text{Wf},\text{Uf},\text{Wh},\text{Uh}]}\in {\mathbb{I}}_8^{D\times D}$theta G R U i for the weights in i 8 size D by D, following standard initialisation.
• ${\theta }_{\text{GRU},i,[\text{bfm bh}]}\in {\mathbb{I}}_8^D$theta G R U i for the biases in i 8 D, initialised to 0 for each element.

In total there are $513D+l(4D+12D^2)$513 D plus l times the quantity 4 D plus 12 D squared parameters in the model.

G.4 Matrix Multiplication

Tensor cores in GPUs are able to calculate fast vector-matrix multiplications with int32 accumulation as $uM\in {\mathbb{I}}_{32}^m$u M in i 32 m where $u\in {\mathbb{I}}_8^n$u in i 8 n and $M\in {\mathbb{I}}_8^{n\times m}$M in i 8 n by m. For our purposes, we define $u@M$ as a scaled multiplication:

$$u@M:={\mathbb{I}}_8\left(\frac{uM}{16\sqrt{n}}\right).$$

u at M is defined as the i 8 casting of u M divided by 16 times the square root of n

Note that when $n=4^d$n equals 4 to the d, the division operation just becomes a right-shift by $4+d$, which is fast to calculate.

We choose this scaled matrix multiplication because we initialise $M$ to 16 times standard normal samples for each element, so dividing by $16\sqrt{n}$ preserves the magnitude of $u$ for the output. In particular, if all elements of $u$ and $M$ are drawn from independently from the standard normal distribution multiplied by 16, the central limit theorem tells us that the expected value per element of the output will be $256\sqrt{n}$, so dividing by $16\sqrt{n}$ preserves the standard deviation of 16.

G.5 Embedding

Our embedding function takes as input an embedding matrix ${\theta }_{\text{emb}}\in {\mathbb{I}}_8^{256\times D}$theta emb in the set of 8 bit integers of shape 256 by D and an input token $t\in {\mathbb{U}}_8$t, an 8 bit unsigned integer, and simply outputs the vector corresponding to that token: ${\theta }_{\text{emb}}[t]\in {\mathbb{I}}_8^D$theta emb of t, in the set of 8 bit integers of length D.

G.6 Layer Normalisation (LN)

Our layer normalisation operation involves multiplying our input $u\in {\mathbb{I}}_8^D$u with a weight ${\theta }_{\text{ln}}\in {\mathbb{I}}_8^D$theta l n before dividing by the mean absolute value of $u$u.

We decide to divide by the mean absolute value of the input instead of the more common root-mean-squared since square roots are expensive on integers. Note that the $L1$ norm after dividing the input by the mean absolute value (when using real numbers) is $D$ instead of 1, which we intentionally choose to preserve more bits of information given the limited range of ${\mathbb{I}}_8$the 8 bit integer range.

We calculate the mean absolute value of input $u$u as:

$$u_{\text{mav}}={\mathbb{I}}_8(\text{sum}(|u|)\gg (2d)),$$

u m a v equals the 8 bit integer conversion of the sum of the absolute values of u, bit shifted right by 2 d

Note that we can safely cast the mean absolute value to an ${\mathbb{I}}_8$8 bit integer without overflow given the properties of the mean of a set, though we lose precision due to truncating the fractional component.

The output of layernorm is calculated as:

$$\text{DIVIDE}({\mathbb{I}}_{16}(u)\odot {\mathbb{I}}_{16}({\theta }_{\text{ln}}),u_{\text{mav}}).$$

divide the element-wise product of 16 bit versions of u and theta l n by u m a v

Since division is an expensive operation, we precompute it using a lookup table. Note that the product of two ${\mathbb{I}}_8$8 bit integer values will always remain in the dynamic range of ${\mathbb{I}}_{16}$16 bit integers, so our lookup table will be of shape $2^{16}\times 2^8$.

G.7 MLP

Each MLP block consists of two weight parameters: ${\theta }_1\in {\mathbb{I}}_8^{4D\times D}$ and ${\theta }_2\in {\mathbb{I}}_8^{D\times 4D}$theta one of shape 4 D by D, and theta two of shape D by 4 D. Given an input $u\in {\mathbb{I}}_8^D$u, we calculate the output as:

$$(u@{\theta }_1^T)@{\theta }_2^T.$$

u matmul theta one transpose, matmul theta two transpose

Note that we do not use an activation function, because the @ operation is already nonlinear due to the saturated conversion from ${\mathbb{I}}_{32}$ to ${\mathbb{I}}_8$

G.8 GRU

Each GRU block accepts an input vector and state $u,s\in {\mathbb{I}}_8^D$u and s consists of 6 weight parameters: ${\theta }_{Wf},{\theta }_{Uf},{\theta }_{Wh},{\theta }_{Uh}\in {\mathbb{I}}_8^{D\times D}$ and ${\theta }_{bf},{\theta }_{bh}\in {\mathbb{I}}_8^D$theta W f, theta U f, theta W h, and theta U h of shape D by D, and bias terms theta b f and theta b h.

Using these weight matrices, we calculate the following vectors:

$$\begin{array}{rl}f& =\sigma ({\mathbb{I}}_8({\mathbb{I}}_{32}(u@{\theta }_{Wf}^T)+{\mathbb{I}}_{32}(s@{\theta }_{Uf}^T)+{\mathbb{I}}_{32}({\theta }_{bf}))),\\ \hat{f}& ={\mathbb{I}}_8((({\mathbb{I}}_{32}(f)+127)\odot {\mathbb{I}}_{32}(s))\gg 8),\\ \hat{h}& =\varphi ({\mathbb{I}}_8({\mathbb{I}}_{32}(u@{\theta }_{Wh}^T)+{\mathbb{I}}_{32}(\hat{f}@{\theta }_{Uh}^T)+{\mathbb{I}}_{32}({\theta }_{bh}))),\\ h& =s+{\mathbb{I}}_8((({\mathbb{I}}_{32}(f)+127)\odot ({\mathbb{I}}_{32}(\hat{h})-{\mathbb{I}}_{32}(s)))\gg 8),\end{array}$$

where $h$ is the output and the new hidden state. In the typical GRU, $\sigma$ stands for the sigmoid function while $\varphi$ stands for the hyperbolic tangent, but we find that setting these as identity operations is sufficient due to the nonlinearity already present in the clipped addition. One can view this clipped addition operation as scaled and shifted version of the “hard” tanh and sigmoid operators.

To explain why we perform these operations, we can analyse this relative to the original GRU. The $f$ vector for the standard GRU has all elements between 0 and 1 due to the sigmoid, but our elements are between -127 and 127. Therefore, to calculate $\hat{f}$ (which is typically just $f\odot s$), we first add 127 to $f$, getting the range between 0 and 254 before multiplying by $s$ before bit-shifting right by 8 again to bring our values back to the ${\mathbb{I}}_8$ dynamic range. We apply similar logic to calculate the final $h$, which is typically just $h=s+f\odot (\hat{h}-s)$ but needs to be rescaled to keep the int8 dynamic range.

G.9 Fitness Calculation in Integer Types

The “fitness” used in language model pretraining is the log-likelihood of correctly generating the next token, treating the outputs of the language model as logits (unnormalised log probabilities). If $t^{\prime }\in {\mathbb{U}}_8$t prime in the set of unsigned 8 bit integers is the next token to predict and $y\in {\mathbb{I}}_8^{256}$y in the set of 256 dimensional signed 8 bit integer vectors are the logits, we can calculate the log likelihood as follows:

$$\begin{array}{rl}{y}^{\prime }& ={\mathbb{I}}_{32}(y)+128,\\ o& =y^{\prime }[t^{\prime }]-\text{LOG2}[\text{sum}(\text{EXP2}[y^{\prime }])],\end{array}$$

y prime equals y cast to 32 bit integers plus 128, and o equals y prime at index t prime minus the LOG 2 of the sum of EXP 2 of y prime

where $o$o is the loss for one token. We implement EXP2 and LOG2 as lookup tables, where

$$\begin{array}{rl}\text{EXP2}[i]& =16\times 2^{i/16},\\ \text{LOG2}[i]& =16\times {\log }_2(i/16).\end{array}$$

Note that each element in EXP2 for any ${\mathbb{U}}_8$unsigned 8 bit input requires at most 20 bits, so the sum of exponents across all possible choices is at most 28 bits, meaning we have to precompute LOG2 for $2^{28}$2 to the 28th values.

H EGG Pretraining with Integer EGGROLL

The core ideas of EGGROLL still apply in this integer-based training setting, but we have to make some modifications to ensure it only uses integer operations.

H.1 Adding EGGROLL Perturbations

For parameter $\theta \in {\mathbb{I}}_8^{m\times n}$theta in I 8 of dimensions m by n that represents a matrix multiplication, we first sample rank-1 perturbation vectors for each index in the batch: $A\in {\mathbb{I}}_8^m$A and $B\in {\mathbb{I}}_8^n$B. We sample these vectors from the standard random normal multiplied by 16 and rounded to the nearest ${\mathbb{I}}_8$8 bit integer (clipping if necessary). To prevent the use of floating-point arithmetic on the accelerator, we pre-generate a large matrix of these random values, randomly indexing into it to get the perturbation vectors.

Given an input $u\in {\mathbb{I}}_8^n$u, instead of calculating $u@{\theta }^T$u times theta transpose, we calculate

$${\mathbb{I}}_8\left(\frac{u{\theta }^T+((u\cdot B){\mathbb{I}}_{32}(A)\gg (4+\hat{\sigma }))}{16\sqrt{n}}\right).$$

The value of $\hat{\sigma }$sigma hat is a hyperparameter, related to the $\sigma$sigma in the main paper as $\sigma =2^{-\hat{\sigma }}$sigma equals 2 to the power of negative sigma hat. Note that the batched forward pass remains efficient since it still simply performs a batched vector-vector dot product in int8 (with int32 accumulate) and a batched vector-scalar product in int32.

We apply this same logic to the embedding matrix, since we can interpret $\theta [t]$theta of t as $\text{one\_hot}(t)\theta$a one hot encoding of t times theta and still apply our rank-1 updates in that context. In practice, this means replacing $u\cdot B$u dot B with $B[t]$B of t.

H.2 Fitness Shaping

We employ a simple fitness shaping scheme based on antithetical pairs. Specifically, given raw fitnesses $s^{+},s^{-}$s plus and s minus, for the positive and negative sample of the antithetical pair respectively, the transformed fitness for the noise is:

$$\text{sign}(s^{+}-s^{-}),$$

Note that the only possible values for the fitness after shaping are $\{-1,0,1\}$.

H.3 Parameter Update

For parameter $\theta \in {\mathbb{I}}_8^{m\times n}$theta that represents a matrix multiplication (or embedding vector), suppose the sampled batch of rank-1 perturbation vectors are $A\in {\mathbb{I}}_8^{N\times m}$A and $B\in {\mathbb{I}}_8^{N\times n}$B, and let the fitnesses after shaping be $F\in {\mathbb{I}}_8^N$F. Then we calculate an intermediate value $E\in {\mathbb{I}}_{32}^{m\times n}$E as:

$$E=(\text{diag}(F)A{)}^{T}B.$$

We use $E$E to determine if each element of $\theta$theta should be increased or decreased. In particular, when the absolute value of $E$ is above a pre-specified threshold we move $\theta$ by one discrete bin in the direction of the sign of $E$. Since there are only 255 unique values for each element in ${\mathbb{I}}_8$the set of 8 bit integers, restricting updates to single bins improves stability without compromising the ability for a parameter to get to any other value with relatively few updates. In particular, we have a real-valued hyperparameter, $\alpha \in (0,1)$alpha between zero and one such that the threshold equals

$${\mathbb{I}}_{32}\left(16\times {\Phi }^{-1}\left(\frac{1-\alpha }{2}\right)\right)\times 16\sqrt{N},$$

where $\Phi$phi is the normal cumulative distribution function. Note that this threshold can be precalculated on a CPU. We observe that $\alpha$ approximately equals the fraction of parameters that are updated at each step.

We currently do not incorporate any momentum or other optimiser states, but this remains critical future work to improve the speed of convergence for pure integer training.

Across model sizes and population size, we find that setting $\hat{\sigma }$sigma hat to 4 and letting $\alpha$ decay over training steps as $\frac{1}{.015t+1}$ gives consistently strong results.

I EGG Ablations

In our main experiments, we use a fixed data batch size of 16 sequences for population sizes 2 and powers of 4 ranging from 4 to $4^{10}=1048576$four to the tenth power, or one million forty eight thousand five hundred seventy six. In this section, we vary the batch size by powers of 4, ranging from 4 to $4^5=1024$four to the fifth power, or one thousand twenty four, while varying population size by powers of 4 from 16 to 1048576. When the batch size, $b$b is greater than half of the population size, $N$N, we give each antithetical pair $\frac{2b}{N}$ sequences, functionally giving a cleaner fitness signal to each member of the population. This also means that the number of parallel “inferences” required is $\max (2b,N)$.

Figure 8: Test loss curves when varying data batch size and population size.

In Fig. 8, we observe that the final test loss for each population size is relatively constant beyond a specific data batch size threshold. At the top right of the figure, we observe a decrease in loss for small population sizes after $b>\frac{N}{2}$b is greater than N over two, which is an artifact of the increased compute usage necessary to use the full data batch. Ignoring this artifact, the minimum batch size for near-optimal performance at a given population size $N$ appears to be $\frac{N}{4^6}$N over four to the sixth. We see that large population sizes need larger data batches for improved performance, since a batch size of 4 results in nearly identical performance for population sizes $4^9=262144$ and $4^{10}=1048576$, but this diverges as data batch size increases.

J Distributed EGGROLL Framework

To facilitate the large-scale experiments, where we scale population sizes beyond 1M, we develop a lightweight distributed training framework designed to minimise network overhead.

J.1 Base-3 Fitness Packing and Bandwidth Efficiency

A key bottleneck in distributed training is the communication of gradients or results. We address this via a custom base-3 packing scheme for fitness vectors. Since workers evaluate perturbations in antithetic pairs, the raw signal is discretised into ternary values $\{-1,0,+1\}$. These are mapped to $\{0,1,2\}$ and packed five at a time into a single byte:

$$\text{byte}=\sum _{i=0}^4{v}_i\cdot 3^i$$

This yields an effective bitrate of 1.6 bits per value (near the ${\log }_{2}3\approx 1.585$log base two of three, approximately one point five eight five theoretical limit). Consequently, the network payload per chunk is approximately $52+\text{chunk\_size}/10$ bytes, rendering bandwidth usage independent of model size.

J.2 System Architecture

The system employs a Coordinator-Worker topology. The Coordinator maintains the global state and assigns population chunks to Workers. Workers calculate fitness on GPU, apply signal shaping (chunk mean filtering, adaptive thresholding), and return only the packed ternary fitness, minimising traffic significantly compared to standard gradient transmission.

K Fine-tuning of Integer Quantised Models

K.1 Quantisation Procedure

To maximise population throughput and reduce device memory during EGGROLL fine-tuning, we represent the large matrix-multiplication parameters of RWKV in an int8 weight format while keeping non-matmul parameters (e.g., small biases / bookkeeping tensors) in floating point, bf16. Following Jacob et al. (2017)Jacob and colleagues, for each weight matrix $W\in {\mathbb{R}}^{d_{in}\times d_{out}}$W of dimensions d in by d out, we use symmetric per-channel int8 quantisation with an absmax scale. For each output channel we first compute:

$$s_i=\max \left(\frac{{\max }_j|W_{i,j}|}{127},ϵ\right),$$

where $ϵ$epsilon is some small scalar. Then, we store each $s_i$s i in bf16, and quantise weights as

$$Q_{i,j}=\text{clip}\left(\text{round}\left(\frac{W_{i,j}}{s_i}\right),-127,127\right)\in {\mathbb{I}}_8.$$

Every matrix parameter is stored as a dictionary containing the quantised weight matrix $Q$, the scale parameters per channel $\{s_i\}\forall i\in 1,\dots ,d_{out}$ and an input scale factor $s_x$ in bf16 precision. At runtime, the forward pass is computed by scaling the input vector by $s_x$ and the quantised matrix $Q$ with the scales per channel, $[s_1,\dots ,s_{d_{out}}]$,

$$x_{n+1}=(x_n\odot s_x{)}^T(W\odot [s_1,\dots ,s_{d_{out}}]).$$

K.2 Integrating integer-quantised EGGROLL with Adam

EGGROLL performs black-box (ES) optimisation directly over the parameter representation used in the forward pass, including integer quantised weights. We integrate this with the Adam optimiser (Kingma & Ba, 2014)of Kingma and Ba by maintaining Adam’s moment estimates in bf16, while enforcing that all quantised tensors remain on the int8 lattice.

ES gradients. EGGROLL estimates gradients via antithetic ES perturbations and score-weighted averaging.
This yields a bf16 gradient estimate for: (i) floating-point parameters (when present), (ii) quantised matrix
parameters via a low-rank perturbation pathway, and (iii) scale parameters $\{s_i\}\forall i\in 1,\dots ,d_{\text{out}}$s i for all i from 1 to d out and $s_x$s x via
explicit scale perturbations. We then pass these gradients to Adam (Optax), which produces an update tensor $u$u
for each parameter leaf.

Adam updates for int8 tensors (discretised). For integer parameters (notably int8), Adam produces a
real-valued proposal $u$u (stored in bf16). Since the parameter itself must remain int8, we convert this proposal
into a sparse unit-step update using a normalised thresholding rule. Let $Q\in {\mathbb{Z}}_8^{m\times n}$Q be an int8 tensor and
$u\in {\mathbb{R}}^{m\times n}$u be Adam’s proposed update. We compute a per-tensor z-score normalisation

$$z=\frac{u-\mu (u)}{\sigma (u)+10^{-8}},$$

then apply a threshold $\tau$tau to form the integer step

$$\Delta =\text{sign}(z)\cdot 1\{|z|\ge \tau \}\in \{-1,0,+1{\}}^{m\times n}.$$

Finally we update by unit increments and clip to the valid int8 range:

$$Q\leftarrow \text{clip}(Q+\Delta ,-127,127).$$

Intuitively, Adam supplies a magnitude- and history-aware proposal, while the discretisation enforces the
integer constraint and yields a stable, sparse update pattern (only entries with sufficiently large normalised
updates are modified).

Memory considerations. We store Adam’s optimiser state (moments) in bf16 for all array-valued leaves
to reduce memory footprint, while keeping scalar bookkeeping in full precision. This keeps the dominant
memory cost of optimisation close to that of the parameters themselves, which is particularly important when
fine-tuning large models with large ES populations.

Model distillation. We distil a non-quantised model into the quantised RWKV-7 model by matching the
two distributions in teacher forced examples from GSM8k. More specifically, the fitness for a given set of
parameters, ${\mu }_i$mu i, is computed as follows:

$$f_{{\mu }_i}(x_{1:T})=\sum _{t=1}^T\text{KL}\left(p_t||q_t(\cdot ;{\mu }_i)\right),$$

the fitness is the sum over time steps t of the K L divergence between the distribution of the non-quantised model p t and the distribution of the quantised model q t

where $x_{1:T}$x 1 through T is a subsequence of tokens taken from the solutions of GSM8K and $\text{KL}(p_t||q_t(\cdot ;{\mu }_i))$ is the
Kullback-Leibler divergence between the distribution of the non-quantised model, $p_t$p t, and the distribution of
the quantised model $q_t$q t over the vocabulary at token $t$t.

L Fine-tuning Pretrained Transformer LLMs with Verifiable Rewards

This section describes compares EGGROLL to standard RL from Verifiable Rewards (RLVR). We first describe
our experimental results, before including details of the infrastructure used to run these experiments.

L.1 Results

Here we demonstrate that EGGROLL can be used to fine-tune pre-trained LLMs on verifiable rewards.
We use the vLLM library Kwon et al. (2023)Kwon and colleagues for efficient inference. More infrastructure detail is given in
Section L.2.

We first fine-tune the Qwen3-4B-Base model Yang et al. (2025)Yang and colleagues on the DeepScaleR Agentica Organization et al. (2025)Agentica Organization and colleagues, a dataset of 40k maths questions. As in standard RLVR, the model generates a chain-of-thought (CoT)
followed by a final answer. Fitness is then simply calculated by extracting the final answer and comparing it to
the ground truth answer Shao et al. (2025)Shao and colleagues. We evaluate performance on MATH500 Hendrycks et al. (2021)Hendrycks and colleagues,
OlympiadBench He et al. (2024)He and colleagues, AIME24 Balunović et al. (2026)Balunovic and colleagues, AMC, and MinervaMath Lewkowycz et al. (2022)Lewkowycz and colleagues. Training curves are shown in Figure 9. Here we see that fine-tuning with EGGROLL significantly

improves performance over the base model. In Section L.1 we show final accuracies with EGGROLL and with the equivalent RL experiment. The RL values are taken from Liu et al. (2025)Liu and colleagues, and we match all the relevant shared hyperparameters and setup, such as maximum response length and prompt phrasing. We see that EGGROLL is able to match the RL optimisation with very minimal hyperparameter tuning, a LoRA rank of 1 and a moderately small population size of 2048. Full hyperparameter details are given in Table 3.

Figure 9: Training curves for fine-tuning Qwen3-4B-Base on the DeepScaleR math dataset. Similar to RL from Verifiable Rewards (RLVR), we see that optimising with EGGROLL is able to improve chain-of-thought reasoning performance on a range of math benchmarks.

	MATH500	OlympiadBench	AIME24	AMC	MinervaMath	Average
Qwen3-4B-Base	50.2	24.4	10.0	33.7	21.7	28.0
+EGGROLL	75.8	37.3	13.3	49.4	31.3	41.4
+RL	67.4	33.5	16.7	49.4	40.1	41.4

Table 1: Final test accuracies when training on the DeepScaleR dataset to optimise verifiable rewards with EGGROLL and RL. We see that EGGROLL significantly boosts performance from the base model and is able to match the equivalent RL experiment.

Since EGGROLL can be used to optimise non-differentiable objectives we next try optimising for pass@kpass at k. While zero-shot (pass@1pass at one) is differentiable, the pass@kpass at k objective is not as it depends on multiple samples from the model. This means it cannot be optimised easily with RL. In Figure 10 we fine-tune the Qwen3-1.7B model on the DeepScaleR dataset with a population size of 256, LoRA rank 1, and $K=4$K equals four. We see that EGGROLL successfully optimises both the pass@1pass at one (differentiable) and pass@kpass at k (non-differentiable) objectives. In Figure 10 (right) we plot the number of distinct answers in 4 samples from the model. We see then when optimising for pass@kpass at k the answer diversity sampled by the model increases over training, whereas when optimising for zero-shot (pass@1pass at one) the model collapses towards a single final answer.

L.2 Training Infrastructure for Large-Scale Transformer LLMs

EGGROLL facilitates the fine-tuning of transformer-based LLMs at scale. We achieve this by repurposing the vLLM inference engine, leveraging its high-throughput kernel implementations and native support for multi-LoRA serving. The system utilises vLLM’s native Tensor Parallelism (TP) to shard the model weights across the GPUs within a node, while cross-node parallelisation is employed for the concurrent evaluation of the LoRA population.

To render ES-based optimisation feasible and efficient across a wide range of model sizes, we implement several critical systems-level optimisations:

Custom WorkerExtension and Sharding-Aware Updates By implementing a custom WorkerExtension, we effectively convert the vLLM inference engine into a training-capable runtime. This extension allows the optimisation logic to reside within the GPU process space, enabling direct,

Figure 10: Using EGGROLL to optimise non-differentiable objectives. Left: Fitness curves comparing training with pass@1 (differentiable) versus pass@k (non-differentiable), where $K=4$K equals 4. Right: The mean number of unique final answers generated per 4-sample set. We observe that when optimizing for pass@k increases answer diversity, whereas optimizing for zero-shot accuracy (pass@1) reduces it.

in-place manipulation of the model’s weights. A significant complexity of this integration is vLLM’s internal tensor parallelism, which frequently fuses weights (e.g. combining q_proj, k_proj, and v_proj into a single qkv_proj tensor). Our update mechanism is explicitly “sharding-aware”; it constructs a dictionary which maps individual LoRA updates to the specific fused slices held by each local GPU rank. This ensures that the global ES update is mathematically consistent across all distributed shards.

Layer-wise Memory Management To prevent out-of-memory (OOM) errors during the update phase, the WorkerExtension performs the ES weight application in a streaming, layer-wise fashion. By processing one layer at a time and clearing temporary buffers, the memory overhead of the update remains independent of the total model depth. This allows for the fine-tuning of models of very different sizes with a VRAM footprint barely exceeding that of standard inference.

Direct GPU-to-GPU Weight Synchronization After computing the ES update on the primary rank, we broadcast the updated parameters to all model instances using NCCL via PyNcclCommunicator. This approach bypasses CPU-based communication and instead uses hardware interconnects to transfer weights directly between GPUs, preventing synchronization from becoming a bottleneck when scaling to more nodes.

Meta-Device Blueprinting To initialise models that exceed the physical RAM of the control node, we employ Meta-Device Initialisation. Using accelerate’s init_empty_weights, we instantiate a “meta” version of the model to derive the weight shapes and sharding requirements for the LoRA adapters. This allows the system to generate a complete parameter blueprint for models of arbitrary size without ever allocating the full weight tensors in system memory.

vLLM Engine Settings Throughout the different experiments with vLLM, we use the following engine settings. These generally allow for high throughput across model sizes (e.g. at least 800 tokens/second), but we haven’t performed hyperparameter sweeps, so potentially faster, more memory-efficient settings may be used for improved results.

Parameter	Value
Tensor parallel size	2, 4
Data type	auto
Enable prefix caching	True
Enforce eager execution	True
Enable LoRA	True
Max LoRAs	$\lceil \text{population\_size}/\text{num\_engines}\rceil$
GPU memory utilisation	0.90
Max number of sequences	384
Max model length	$\max (1024,512+\text{max\_tokens})$
Max batched tokens	$\text{prompt\_batch\_size}\times 1024$
Load format	auto

Table 2: vLLM engine configuration parameters to allow for high throughput EGGROLL training on large-scale transformer LLMs.

Parameter	Value
Population size	256, 2048
Sigma	0.001
Learning Rate	0.001
Max Response Length	4096
Temperature	0.0, 0.7
Samples Per Prompt	1, 4
Pass at K	True, False
LoRA Rank	1
LoRA Reuse Steps	4

Table 3: Hyperparameters for the verifiable reward transformer fine-tuning experiments in Section L.1.

M Fine-tuning Time Series Foundation Model: High-Frequency Trading

The preceding experiments demonstrate the effectiveness of EGGROLL on natural language reasoning tasks. We now investigate whether EGGROLL can effectively fine-tune pretrained foundation models on a fundamentally different data modality: structured time series. We focus on high-frequency trading (HFT) for two reasons. First, HFT generates data at an unprecedented scale. The S&P 500 constituents alone produced approximately 3.8 trillion tokens of order flow data between 2016 and 2021, comparable to the largest natural language corpora. Second, the domain presents a well-defined downstream task (order execution) with a natural reward signal: the realised profit and loss, also known as PnL, making it amenable to fine-tuning via evolution strategies.

Order execution takes place in limit order books (LOBs), which are the mechanism upon which modern financial exchanges operate (Gould et al., 2013; Bouchaud et al., 2018)as discussed by Gould and colleagues in 2013 and Bouchaud and colleagues in 2018. They allow market participants to submit limit orders that specify the details of intended transactions. Specifically, each limit order contains the order type, direction, price, and quantity. The continuous stream of these orders is known as the order flow. LOBs aggregate the limit orders that have not been matched yet. Unlike natural language, where tokens are purely symbolic, order flow messages comprise both categorical values (e.g., order type, direction) and numerical values (e.g., price, quantity) in which magnitude carries semantic meaning. This structure provides a distinct test of EGGROLL’s ability to fine-tune foundation models on time series sequential data.

A central objective in this context is order execution, which consists of buying or selling a specified quantity of an asset within a given time window. The goal is to maximise profit by transacting at favourable prices. In prior reinforcement learning approaches to this problem, the action space is usually simplified (Frey et al., 2023; Mohl et al., 2025; Ning et al., 2021)as seen in work by Frey and colleagues, Mohl and colleagues, and Ning and colleagues. In contrast, we aim to give the model full flexibility in choosing

Figure 11: Training curves for order execution with EGGROLL. Left: Mean PnL over training epochs for the baseline ($\sigma =0$, orange dashed) and EGGROLL ($\sigma =0.01$, blue solid). Right: PnL standard deviation over training epochs. Shaded regions indicate the interquartile range across runs.

limit orders, i.e., to freely choose the order type, direction, price, and quantity. We achieve this by tokenising the limit order book messages and providing the model with a token-level action space.

Foundation models have recently been used to generate synthetic order flow (Nagy et al., 2023; Li et al., 2025)as shown by Nagy and colleagues, as well as Li and colleagues and have been shown to replicate realistic market behaviour (Nagy et al., 2025)in more recent work through next-token prediction. We therefore first pretrain a foundation model on tokenised limit order book messages, and then fine-tune it using EGGROLL for the order execution task. The pretraining follows the approach of Nagy et al. (2023)Nagy and colleagues: we employ an S5 model architecture (Smith et al., 2023)by Smith and colleagues that generates next-token probabilities, with cross-entropy as the training loss. The pretraining is conducted on the LOBSTER data set (Huang & Polak, 2011)by Huang and Polak for the Google stock (GOOG) in 2022, which contains around 25B tokens.

Subsequently, we fine-tune the model using EGGROLL. The training parameters are summarised in Table 4. The task is to execute a sell order of $Q=30$Q equals thirty shares within a horizon of $T=10$T equals ten steps. In each episode, the LOB is initialised based on a LOB snapshot followed by 10 warm-up background messages. In each step, the population members generate their messages, which are then followed by 50 real background data messages. The orders are executed using the Jax-LOB (Frey et al., 2023)simulator by Frey and colleagues. We perform the fine-tuning on a fixed time window for GOOG in January 2023. Following (Galim et al., 2025)Galim and colleagues, we apply LoRA with rank 4 on all projection matrices while freezing SSM parameters and layer norms. Performance is evaluated using PnL based on the executed prices and the initial mid price. Specifically, for a sell task of total quantity $Q$Q, the PnL is computed as

$$\sum _{i=1}^N{q}_i{p}_i-Q{P}_{\text{mid}}^{\text{init}},$$

the sum of q i times p i from i equals one to N, minus Q times the initial mid price

where $q_i$q i and $p_i$p i denote the quantity and price of the $i$-th executed trade and $P_{\text{mid}}^{\text{init}}$P mid init is the mid-price at the beginning of the execution window. If the agent does not execute the entire quantity by the end of the episode, an automatic market order is submitted selling the remaining quantity. To improve robustness to outliers, fitness is defined as a rank-based transformation of the PnL. Specifically, for a population of size $M$M, the PnL values

$$\mathcal{P}=\{{\text{PnL}}_1,\dots ,{\text{PnL}}_M\},$$

P, defined as the set of PnL values one through M

are mapped to the interval $[-0.5,0.5]$, where $\text{rank}({\text{PnL}}_i)\in \{0,\dots ,M-1\}$the rank of PnL i, which is between zero and M minus one denotes the rank of the $i$-th individual’s PnL:

$$F_i=\frac{1}{2}-\frac{\text{rank}({\text{PnL}}_i)}{M-1},$$

F i equals one half minus the rank of PnL i over M minus one

Training curves over 6,500 epochs are shown in Figure 11. The baseline policy ($\sigma =0$)with sigma equals zero, corresponding to the pretrained model, achieves a mean PnL of approximately 4,700. In contrast, EGGROLL fine-tuning ($\sigma =0.01$)with sigma equals point zero one improves the mean PnL to around 12,000, corresponding to a roughly 155% improvement over the baseline. The right panel of Figure 11 depicts the PnL standard deviation during fine-tuning: it initially increases to around 3,100 during the first 2,500 epochs, which corresponds to an exploration phase where the population tries out diverse strategies, before decreasing to approximately 400 by the end of training, indicating that the population concentrates around a high-performing policy.

Hyperparameter	Value
Model	LOBS5-360M
Parallel generations per GPU	2,048
Total parallel generations	65,536
LoRA rank	4
Sigma	0.01
Learning rate $\eta$	0.001
Epochs	6,500

Table 4: Model and EGGROLL fine-tuning settings for high-frequency trading.

Figure 12: Training curves and wall clock times for cooperative Multi Particle Environments. Hyperparameter optimisation yielded equal batch sizes for all algorithms on the same environment. All EGGROLL runs used rank 1 perturbations. Shaded regions are standard errors of mean values.

N Experimental Details

N.1 Multi Agent Reinforcement Learning Experiments

Table 5: Hyperparameter Ranges Used in MPE Sweeps for EGGROLL and OpenES

Hyperparameter	Values
activation	pqn, tanh
pop_size	128, 512, 1024, 2048, 4096
learning_rate	0.01, 0.05, 0.1, 0.5
lr_decay	0.3, 0.7, 1.0
sigma	0.1, 0.2, 0.3, 0.4, 0.5
rank_transform	true, false

Table 6: Hyperparameter Ranges Used in MPE Sweeps for IPPO

Hyperparameter	Values
activation	relu, tanh
pop_size	128, 512, 1024, 2048, 4096
learning_rate	5e-5, 1e-4, 2.5e-4, 1e-3
entropy_coef	0.001, 0.005, 0.01

Table 7: MPE Simple Spread v3

Hyperparameter	eggroll	open_es	ippo
activation	tanh	tanh	tanh
deterministic_policy	true	true	false
learning_rate	0.01	0.01	0.001
lr_decay	0.7	0.7	linear
layer_size	64	64	64
n_layers	3	3	3
pop_size	128	128	128
optimizer	adamw	adamw	adam
rank	1	1	-
rank_transform	false	false	-
sigma	0.5	0.5	-
n_minibatches	-	-	4
update_epochs	-	-	4
gamma	-	-	0.99
gae_lambda	-	-	0.95
epsilon_clip	-	-	0.2
entropy_coef	-	-	0.01
value_coef	-	-	0.5
max_grad_norm	-	-	0.5

Table 8: MPE Simple Speaker Listener v4

Hyperparameter	eggroll	open_es	ippo
activation	tanh	tanh	relu
deterministic_policy	true	true	false
learning_rate	0.01	0.01	0.001
lr_decay	0.7	0.3	linear
layer_size	64	64	64
n_layers	3	3	64
pop_size	512	512	512
optimizer	adamw	adamw	adam
rank	1	1	-
rank_transform	true	true	-
sigma	0.5	0.5	-
n_minibatches	-	-	4
update_epochs	-	-	4
gamma	-	-	0.99
gae_lambda	-	-	0.95
epsilon_clip	-	-	0.2
entropy_coef	-	-	0.005
value_coef	-	-	0.5
max_grad_norm	-	-	0.5

Table 9: MPE Simple Reference v3

Hyperparameter	eggroll	open_es	ippo
activation	pqn	tanh	relu
deterministic_policy	true	true	false
learning_rate	0.01	0.01	0.001
lr_decay	0.3	0.3	linear
layer_size	64	64	64
n_layers	3	3	3
pop_size	4096	4096	4096
optimizer	adamw	adamw	adam
rank	1	1	-
rank_transform	false	true	-
sigma	0.1	0.3	-
n_minibatches	-	-	4
update_epochs	-	-	4
gamma	-	-	0.99
gae_lambda	-	-	0.95
epsilon_clip	-	-	0.2
entropy_coef	-	-	0.01
value_coef	-	-	0.5
max_grad_norm	-	-	0.5

We train on three cooperative Multi Particle Environments (MPEs) (Lowe et al., 2017)introduced by Lowe and colleagues implemented in JaxMARL (Rutherford et al., 2024)by Rutherford and colleagues with feed-forward networks of width 64 and depth 3, performing Bayesian hyperparameter optimisation for each environment and algorithm. All runs were executed on NVIDIA A100-SXM4-40GB GPUs. We find that the optimal batch size is consistent across algorithms on the same environment. Figure 12 shows that EGGROLL with rank 1 trains up to 2.4 times faster than OpenES for large batch sizes while staying competitive in performance.

Figure 13: (a) Comparison of our finetuned RWKV 7G 7 billion parameter model using 8 GPUS with the results reported by Qiu et al. (2025)Qiu and colleagues on similarly sized Qwen models. (b) Performance of our finetuned RWKV 7G 14 billion parameter model on hard reasoning tasks using 32 GPUs for 12 hours. The model was trained using the DeepScaleR dataset and the best checkpoint was chosen by evaluating on AIME24. Due to the size of the model we were not able to run similar baseline experiments using GRPO.

N.2 Reasoning Fine-tuning Experiments: Countdown

We ran a Bayesian hyper-parameter sweep (Snoek et al., 2012)by Snoek and colleagues for both GRPO and EGGROLL and used the best set found to run the experiments in figure 4b. For GRPO we swept over sampling temperature and learning rate, whereas for EGGROLL we swept over the standard deviation of the ES sampling $\sigma$sigma and the learning rate scale. The best hyper-parameters found are detailed on tables 10 (EGGROLL) and 11 (GRPO). All of the experiments run in 8 hours on a NVIDIA H200 GPU.

Hyperparameter	Value
Model	RWKV 7g1.5B
Optimiser	Gradient descent
ES standard deviation $\sigma$	$7\times 10^{-4}$
Rank $r$	1
Learning-rate scale ${\eta }_{\text{scale}}$	0.125
Population size	256
Parallel generations per GPU	1536
Prompts per epoch	6
Generation / thinking length	1000 tokens
Train / val temperature	0 / 0
Parallel validations	128

Table 10: Key hyperparameters for EGGROLL training on Countdown with FastRWKV-7g1.5B.

We also run an experiment where we increase the number of GPUs to 8 and use a bigger model, RWKV 7g7B, on the Countdown task, allowing for stronger final performance. Notably, we compare to the results reported by Qiu et al. (2025)Qiu and colleagues on Countdown. Figure 13a shows that starting from our significantly weaker model (RWKV 7g7B v.s. Qwen 2.5-7B), we are able to train to a higher validation accuracy (72.9%), v.s. the ones reported for training with GRPO (52.8%) and Open ES (66.8%). Qiu et al. (2025)Qiu and colleagues do not report the wall clock time or the hardware used for their experiments which makes it difficult to establish a fair comparison.

Hyperparameter	Value
Model	RWKV 7g1.5B
Optimiser	Radam
Learning rate $\eta$	$3\times 10^{-6}$
Generations per prompt $G$	8
Parallel generations per GPU	64
Prompts per epoch	8
Generation length	1000 tokens
Number of minibatches	4
PPO clip parameter ${ϵ}_{\text{clip}}$	0.2
Train / val temperature	1 / 0
Parallel validations	128

Table 11: Key hyperparameters for GRPO training on Countdown with AssociativeScanRWKV-7g1.5B.

N.3 Reasoning Fine-tuning Experiments: GSM8K

We used the hyper-parameters found for Countdown as a starting point and reduced the learning rates for both GRPO and EGGROLL using linear search until we found the best performing one on the validation set. Our experiments for GSM8K run on 8 NVIDIA H200 GPUS for 8 hours each. We also increase the standard deviation, $\sigma$sigma, parameter for ES (from $7\times 10^{-4}$ to $2\times 10^{-3}$)from seven times ten to the minus four to two times ten to the minus three as the significantly bigger population sizes (8096 v.s. 512) allow for much more stable training and aggressive exploration.

Hyperparameter	Value
Model	RWKV 7g7B
ES standard deviation $\sigma$	$2\times 10^{-3}$
Rank $r$	1
Learning-rate scale ${\eta }_{\text{scale}}$	0.06
Generations per prompt $G$	512
Parallel generations per GPU	1024
Total parallel generations	8192
Prompts per epoch	16
Generation length	1000 tokens
Noise reuse factor	1
Freeze non-LoRA params	True
Train / val temperature	0 / 0
Parallel validations	128

Table 12: Key hyperparameters for multi-GPU EGGROLL training on GSM8K with FastRWKV-7g7B.

N.4 Reinforcement Learning Experiments

Next, we compare the performance of EGGROLL against standard OpenES as implemented in Salimans et al. (2017)Salimans and colleagues on reinforcement learning tasks. Given the small network sizes, we can use OpenES at this scale, but as network sizes increase, the use of vanilla OpenES becomes computationally infeasible. We use the standard formulation of simply optimising for the final return in the environment. For both EGGROLL and OpenES, we perform hyperparameter optimisation (HPO) separately for each environment. For each algorithm–environment pair, we define plausible ranges for all key hyperparameters based on prior work and preliminary experiments. We then perform 20 random search trials, where each trial corresponds to a single training run with a randomly sampled hyperparameter configuration. Each configuration is evaluated based on the final return achieved by the mean policy parameters at the end of training. After all trials, we select the configuration that yields the highest final return. Using this best configuration, we then run 10 independent seeds to evaluate performance and report the mean and standard error of the mean across these seeds.

Figure 14: Comparison of reinforcement learning results: Mean returns for each environment and algorithm across 10 random seeds. The returns are evaluated using the mean of the parameters. HPO was conducted for each algorithm/environment pair. The shaded region is the standard error of the mean.

Hyperparameter	Value
Model	RWKV 7g7B
Learning rate $\eta$	$1\times 10^{-6}$
Generations per prompt $G$	8
Parallel generations per GPU	32
Total parallel generations	256
Prompts per epoch	32
Generation length	1000 tokens
Number of minibatches	16
Number of workers (processes)	8
PPO clip parameter ${ϵ}_{\text{clip}}$	0.2
Train / val temperature	1 / 0
Parallel validations	128

Table 13: Key hyperparameters for multi-GPU GRPO training on GSM8K with AssociativeScanRWKV-7g7B.

Hyperparameter	Value
Model	RWKV 7g7B
Optimiser	EGGROLL (Quantised))
ES standard deviation $\sigma$	0.4
Rank $r$	1
Learning-rate scale ${\eta }_{\text{scale}}$	$3\times 10^{-7}$
Population size	8192
Parallel generations per GPU	256
Prompts per epoch	1
Generation / thinking length	256 tokens
Train / val temperature	0 / 0
Parallel validations	128

Table 14: Key hyperparameters for quantised EGGROLL training on GSM8K (teacher-forced) with RWKV-7g7B.

We use policy networks with 3 layers of 256 neurons and a range of environments that demonstrate different capabilities. We evaluate across the Navix (Pignatelli et al., 2024)Pignatelli and colleagues, Craftax (Matthews et al., 2024)Matthews and colleagues, Brax (Freeman et al., 2021)Freeman and colleagues, Kinetix (Matthews et al., 2025)Matthews and colleagues, and Jumanji (Bonnet et al., 2024)Bonnet and colleagues suites of environments. We evaluate 16 environments in total. We choose environments that are not trivial or impossible for PPO to solve, according to the original papers. We also choose environments that belong to different categories (e.g., environment size in Kinetix or categories in Jumanji).

We show a subsample of the evaluated environments in Fig. 4a with the remaining results and hyperparameter details in Appendix N.4. Our findings show that EGGROLL is competitive with Open ES on 7/16 environments, underperforms on 2/16, and outperforms on 7/16. This does not take into account the speed-ups when compared to using OpenES with full-rank updates (see Figure 15). We postulate that the reason for this performance increase is that the large networks are difficult to optimise for OpenES and lend themselves well to low-rank updates.

We present here the hyperparameter ranges we used for hyperparameter optimisation, as well as all hyperparameter settings for all the experiments. All RL experiments were run on an NVIDIA L40S GPU. For PPO, we use the same methodology to tune the hyperparameters as we did for OpenES and EGGROLL as described in Section 6.2. We report the ranges and the final hyperparameters here. We train PPO agents using Rejax (Liesen et al., 2024)Liesen and colleagues. We use the activation function from Gallici et al. (2025)Gallici and colleagues in our experiments, which we refer to as the “pqn” activation function in our hyperparameter tables.

Figure 15: Comparison of reinforcement learning results: Mean and standard deviation of training time. Note that some of the timing difference is due to the differences in episode lengths, which is why the total time for EGGROLL sometimes appears longer than OpenES despite EGGROLL being faster on a per-timestep basis.

Table 15: Hyperparameter Ranges for EGGROLL and OpenES

Hyperparameter	Values
pop_size	512, 1024, 2048, 4096
n_parallel_evaluations	1, 4, 8
rank	1, 2, 4
optimizer	adamw, sgd, adam
learning_rate	1e-3, 1e-2, 1e-1
lr_decay	0.995, 0.999, 0.9995, 1.0
sigma	0.05, 0.2, 0.5
sigma_decay	0.995, 0.999, 0.9995, 1.0
rank_transform	true, false
deterministic_policy	true, false

Table 16: Hyperparameter Ranges for PPO

Hyperparameter	Values
clip_eps	0.1, 0.2, 0.3
ent_coef	0, 0.0001, 0.001
gae_lambda	0.9, 0.95, 0.98
gamma	0.95, 0.99, 0.995, 0.999
learning_rate	0.0001, 0.0003, 0.001
max_grad_norm	0.5, 1, 2
layer_size	256
n_layers	3
normalize_observations	true
normalize_rewards	false
num_envs	64, 128, 256
num_epochs	4, 8, 16
num_minibatches	16, 32, 64
num_steps	64, 128, 256
reward_normalization_discount	0.99
skip_initial_evaluation	false
vf_coef	0.5, 0.75, 1

Table 17: CartPole-v1

Hyperparameter	eggroll	open_es
activation	pqn	pqn
deterministic_policy	false	true
learning_rate	0.1	0.1
lr_decay	0.9995	0.9995
layer_size	256	256
n_layers	3	3
n_parallel_evaluations	1	4
pop_size	2048	512
optimizer	sgd	adamw
rank	4	/
rank_transform	false	true
sigma	0.2	0.5
sigma_decay	0.999	0.9995

Table 18: Pendulum-v1

Hyperparameter	eggroll	open_es
activation	pqn	pqn
deterministic_policy	false	true
learning_rate	0.01	0.01
lr_decay	0.995	0.995
layer_size	256	256
n_layers	3	3
n_parallel_evaluations	1	4
pop_size	4096	4096
optimizer	adam	adamw
rank	4	/
rank_transform	false	false
sigma	0.05	0.05
sigma_decay	0.995	1

Table 19: brax/ant

Hyperparameter	eggroll	open_es
activation	pqn	pqn
deterministic_policy	false	false
learning_rate	0.01	0.1
lr_decay	0.9995	0.995
layer_size	256	256
n_layers	3	3
n_parallel_evaluations	1	8
pop_size	2048	512
optimizer	adam	adam
rank	1	/
rank_transform	false	false
sigma	0.05	0.05
sigma_decay	0.9995	0.9995

Table 20: brax/humanoid

Hyperparameter	eggroll	open_es
activation	pqn	pqn
deterministic_policy	true	false
learning_rate	0.1	0.1
lr_decay	1	0.995
layer_size	256	256
n_layers	3	3
n_parallel_evaluations	8	8
pop_size	4096	1024
optimizer	adam	sgd
rank	1	/
rank_transform	true	true
sigma	0.2	0.2
sigma_decay	0.9995	0.995
Tables 16 through 20 provide the hyperparameter details for the experiments. Table 16 lists the hyperparameter ranges for PPO. Tables 17 through 20 detail specific hyperparameter settings for eggroll and open e s across the environments: CartPole-v1, Pendulum-v1, brax ant, and brax humanoid.

Table 21: brax/inverted_double_pendulum

Hyperparameter	eggroll	open_es
activation	pqn	pqn
deterministic_policy	true	true
learning_rate	0.1	0.1
lr_decay	1	0.995
layer_size	256	256
n_layers	3	3
n_parallel_evaluations	1	1
pop_size	2048	4096
optimizer	adam	adam
rank	2	/
rank_transform	true	true
sigma	0.5	0.05
sigma_decay	0.995	1

Table 22: craftax/Craftax-Classic-Symbolic-AutoReset-v1

Hyperparameter	eggroll	open_es
activation	pqn	pqn
deterministic_policy	false	false
learning_rate	0.01	0.001
lr_decay	0.995	0.995
layer_size	256	256
n_layers	3	3
n_parallel_evaluations	4	8
pop_size	2048	4096
optimizer	sgd	adamw
rank	1	/
rank_transform	false	false
sigma	0.05	0.05
sigma_decay	1	0.995

Table 23: craftax/Craftax-Symbolic-AutoReset-v1

Hyperparameter	eggroll	open_es
activation	pqn	pqn
deterministic_policy	false	false
learning_rate	0.01	0.1
lr_decay	0.999	0.995
layer_size	256	256
n_layers	3	3
n_parallel_evaluations	1	4
pop_size	512	1024
optimizer	sgd	adam
rank	4	/
rank_transform	true	false
sigma	0.05	0.5
sigma_decay	0.999	1

Table 24: jumanji/Game2048-v1

Hyperparameter	eggroll	open_es
activation	pqn	pqn
deterministic_policy	false	true
learning_rate	0.1	0.01
lr_decay	1	0.999
layer_size	256	256
n_layers	3	3
n_parallel_evaluations	4	4
pop_size	1024	1024
optimizer	adamw	adamw
rank	1	/
rank_transform	false	true
sigma	0.5	0.05
sigma_decay	0.9995	0.9995

Table 25: jumanji/Knapsack-v1

Hyperparameter	eggroll	open_es
activation	pqn	pqn
deterministic_policy	false	false
learning_rate	0.1	0.01
lr_decay	0.999	1
layer_size	256	256
n_layers	3	3
n_parallel_evaluations	4	1
pop_size	1024	2048
optimizer	sgd	adamw
rank	4	/
rank_transform	true	true
sigma	0.05	0.5
sigma_decay	1	0.995

Table 26: jumanji/Snake-v1

Hyperparameter	eggroll	open_es
activation	pqn	pqn
deterministic_policy	false	false
learning_rate	0.001	0.001
lr_decay	0.9995	1
layer_size	256	256
n_layers	3	3
n_parallel_evaluations	8	1
pop_size	4096	2048
optimizer	adam	sgd
rank	1	/
rank_transform	true	false
sigma	0.05	0.2
sigma_decay	0.9995	1

Table 27: kinetix/l/hard_pinball

Hyperparameter	eggroll	open_es
activation	pqn	pqn
deterministic_policy	true	true
learning_rate	0.01	0.01
lr_decay	0.995	1
layer_size	256	256
n_layers	3	3
n_parallel_evaluations	8	1
pop_size	2048	512
optimizer	sgd	sgd
rank	4	/
rank_transform	true	true
sigma	0.05	0.5
sigma_decay	0.999	0.9995

Table 28: kinetix/m/h17_thrustcontrol_left

Hyperparameter	eggroll	open_es
activation	pqn	pqn
deterministic_policy	false	false
learning_rate	0.1	0.001
lr_decay	0.9995	1
layer_size	256	256
n_layers	3	3
n_parallel_evaluations	4	1
pop_size	512	1024
optimizer	sgd	adam
rank	4	/
rank_transform	true	true
sigma	0.5	0.5
sigma_decay	1	0.999

Table 29: kinetix/s/h1_thrust_over_ball

Hyperparameter	eggroll	open_es
activation	pqn	pqn
deterministic_policy	false	false
learning_rate	0.1	0.01
lr_decay	0.995	0.995
layer_size	256	256
n_layers	3	3
n_parallel_evaluations	1	1
pop_size	512	2048
optimizer	adamw	sgd
rank	1	/
rank_transform	true	true
sigma	0.5	0.05
sigma_decay	0.9995	1

Table 30: navix/Navix-DoorKey-8x8-v0

Hyperparameter	eggroll	open_es
activation	pqn	pqn
deterministic_policy	false	false
learning_rate	0.01	0.01
lr_decay	0.9995	1
layer_size	256	256
n_layers	3	3
n_parallel_evaluations	1	8
pop_size	1024	2048
optimizer	adamw	adam
rank	1	/
rank_transform	false	true
sigma	0.05	0.05
sigma_decay	1	1

Table 31: navix/Navix-Dynamic-Obstacles-6x6-Random-v0

Hyperparameter	eggroll	open_es
activation	pqn	pqn
deterministic_policy	false	false
learning_rate	0.01	0.01
lr_decay	0.999	1
layer_size	256	256
n_layers	3	3
n_parallel_evaluations	4	1
pop_size	512	4096
optimizer	adam	adam
rank	2	/
rank_transform	false	false
sigma	0.05	0.2
sigma_decay	1	0.995

Table 32: navix/Navix-FourRooms-v0

Hyperparameter	eggroll	open_es
activation	pqn	pqn
deterministic_policy	false	false
learning_rate	0.01	0.001
lr_decay	0.999	0.9995
layer_size	256	256
n_layers	3	3
n_parallel_evaluations	4	4
pop_size	2048	2048
optimizer	sgd	adam
rank	4	/
rank_transform	true	false
sigma	0.05	0.05
sigma_decay	0.9995	0.9995

Table 33: PPO Hyperparameters (Set 1)

Hyperparameter	CartPole	Pendulum	Ant	Humanoid	IDP	CraftaxClassic	CraftaxSymbolic	Game2048
activation	pqn	pqn	pqn	pqn	pqn	pqn	pqn	pqn
clip_eps	0.2	0.1	0.2	0.3	0.1	0.2	0.2	0.3
ent_coef	0.0001	0.001	0	0.0001	0.0001	0.0001	0	0.001
gae_lambda	0.9	0.95	0.95	0.9	0.98	0.98	0.9	0.9
gamma	0.995	0.999	0.995	0.95	0.99	0.95	0.95	0.99
learning_rate	0.0003	0.0003	0.0003	0.0001	0.001	0.001	0.0003	0.0003
max_grad_norm	0.5	1	0.5	2	2	2	2	2
layer_size	256	256	256	256	256	256	256	256
n_layers	3	3	3	3	3	3	3	3
normalize_obs	true	true	true	true	true	true	true	true
normalize_rew	false	false	false	false	false	false	false	false
num_envs	256	256	64	256	64	128	256	64
num_epochs	4	16	8	4	4	4	4	8
num_minibatches	32	16	32	64	64	32	32	16
num_steps	128	256	128	64	128	128	64	64
rew_norm_discount	0.99	0.99	0.99	0.99	0.99	0.99	0.99	0.99
skip_initial_eval	false	false	false	false	false	false	false	false
vf_coef	0.5	1	1	0.75	1	0.5	0.75	0.75

Table 34: PPO Hyperparameters (Set 2)

Hyperparameter	Knapsack	Snake	HardPinball	ThrustLeft	ThrustBall	DoorKey	DynamicObs	FourRooms
activation	pqn	pqn	pqn	pqn	pqn	pqn	pqn	pqn
clip_eps	0.1	0.3	0.1	0.2	0.2	0.1	0.1	0.1
ent_coef	0.0001	0.001	0.0001	0.0001	0.0001	0.0001	0.001	0.001
gae_lambda	0.9	0.95	0.9	0.9	0.95	0.98	0.98	0.9
gamma	0.99	0.999	0.99	0.995	0.999	0.95	0.999	0.99
learning_rate	0.0001	0.0001	0.0001	0.0001	0.0001	0.0003	0.001	0.001
max_grad_norm	0.5	0.5	1	2	0.5	0.5	1	1
layer_size	256	256	256	256	256	256	256	256
n_layers	3	3	3	3	3	3	3	3
normalize_obs	true	true	true	true	true	true	true	true
normalize_rew	false	false	false	false	false	false	false	false
num_envs	256	128	256	256	64	64	128	256
num_epochs	4	4	16	16	16	16	4	8
num_minibatches	64	16	16	32	16	64	16	32
num_steps	128	128	64	128	64	256	128	256
rew_norm_discount	0.99	0.99	0.99	0.99	0.99	0.99	0.99	0.99
skip_initial_eval	false	false	false	false	false	false	false	false
vf_coef	0.75	0.75	0.5	0.5	0.5	0.75	0.5	0.75