Starting point

Idea

Use evolution strategies to adapt frozen LLMs through low-rank adapters, but search in a DoRA-like weight decomposition instead of plain LoRA coordinates. The central hypothesis is that separating magnitude from direction gives ES a better-conditioned search space for long-context reward optimisation.

The intended paper is a delta paper for ICLR: not a new foundation model method, but a focused empirical and theoretical study connecting EGGROLL-style low-rank ES with weight-decomposed adapter training.

Reasoning

Evolution Strategies at the Hyperscale argues that ES can be made much more efficient at large scale by structuring perturbations as low-rank matrices, and gives theory for high-dimensional ES behaviour. DoRA argues that LoRA misses part of the learning capacity of full fine-tuning because it does not explicitly separate weight magnitude from direction; its fix is to learn magnitude while applying low-rank adaptation to direction.

For LLM agents, especially long-context agents, the reward signal is often sparse, delayed, noisy, or attached to non-differentiable tool-use trajectories. That makes gradient-free optimisation attractive, but only if the search space is not wasteful. A DoRA-like decomposition may give ES a more useful geometry than direct LoRA-space perturbations: radial changes in magnitude and angular changes in direction can be explored separately rather than entangled in the same low-rank coordinates.

The first experiment should be deliberately small. It is a smoke test for the JAX-to-PyTorch ES translation and for the experimental comparison itself: same ES implementation, same objective, same budget, with plain LoRA-style perturbations compared against DoRA-style decomposed perturbations.

Baseline implementation

The current baseline candidate is the submitted PyTorch port of EGGROLL, provisionally called D-HyperscaleES. The project now has a local pre-XP validation harness at src/validate_eggroll_port.py, configured by experiments/eggroll_port_validation.yaml. This harness does not run an LLM and does not replace an official xp run. It checks the mechanics that should be invariant between the JAX and PyTorch implementations:

• EggRollLinear.forward applies the same low-rank perturbation as JAX EggRoll.do_mm.
• EggRoll.do_updates has the expected update sign and scaling.
• EggRoll.convert_fitnesses matches the standard EGGROLL reward normalisation.

The implementation should be used initially only as noiser=eggroll with freeze_nonlora=True. The open_es and eggrollbs variants in the submitted port are not faithful enough to use as baselines yet: open_es samples LoRA-structured perturbations during generation while updating as if it had sampled full-rank noise, and eggrollbs does not actually make its baseline directions zero-noise during generation.

Other known differences from the JAX reference are acceptable for the first baseline but should be documented in any experiment note. The PyTorch random stream is deterministic but not exactly JAX’s fold_in plus jax.random.normal; the HuggingFace layer targeting is an approximation of the JAX es_map; freeze_nonlora=False is not implemented; and the submitted set_epoch_noise caches all layer noise for the generation batch, which may matter at larger model, rank, or population size.

The code ownership boundary is explicit:

• the submitted PyTorch EGGROLL port remains under src/dhyperscalees;
• the project-specific decomposed candidate lives under src/decomposed_eggroll;
• the original matrix smoke runner remains src/run_matrix_learning.py;
• the comparison runner is src/run_eggroll_vs_decomposed_matrix_learning.py.

This separation is intentional. It should stay easy to review which behaviour comes from the submitted baseline and which behaviour comes from the new decomposed variant.

The decomposed forward path should also keep the EGGROLL computational idea. For a noisy direction

$$V_i=V+{\alpha }_i{A}_i{B}_i^{\top },$$

the implementation computes

$$x{V}_i^{\top }=x{V}^{\top }+{\alpha }_i(x{B}_i)A_i^{\top }$$

without constructing the dense perturbation $A_i{B}_i^{\top }$. The DoRA row norm is also computable from factors:

$$\Vert V_o+{\alpha }_i{A}_{i,o}{B}_i^{\top }{\Vert }^2=\Vert V_o{\Vert }^2+2{\alpha }_i{A}_{i,o}(V_o{B}_i)+{\alpha }_i^2{A}_{i,o}(B_i^{\top }{B}_i)A_{i,o}^{\top }.$$

This is the path used by src/decomposed_eggroll/dora_eggroll.py; the current matrix implementation no longer needs to materialise the dense direction perturbation during the forward pass.

The synthetic empirical ladder has now been run and consolidated in Empirical sketch proofs, matrix tests, and LLM tests. The sequence cleared these gates:

• plain EGGROLL learned the controlled matrix objective;
• full decomposed EGGROLL beat plain EGGROLL robustly on the matrix seed/rank sweep, while direction-only and magnitude-only ablations failed;
• both methods ran through causal-LM module patching on a random Llama-style model;
• both methods stayed finite and improved answer-token NLL on a small pretrained causal LM;
• the pretrained seed/rank sweep was stable and mildly favorable to decomposed EGGROLL, but not decisive.

The next gate should no longer be another logprob smoke test. It should be a tiny real-reward benchmark: generated arithmetic or GSM8K-style prompts where the model samples an answer and ES receives only a scalar exact-answer reward, optionally with minimal parseability shaping if the fully sparse reward is all zero. This is the first test that touches the actual paper claim about reward-only adaptation.

Baseline equations

For a linear weight matrix $W\in {\mathbb{R}}^{d_{\mathrm{out}}\times d_{\mathrm{in}}}$ and rank $r$, EGGROLL draws a seed-determined Gaussian matrix for each antithetic pair:

$${\Xi }_{e,k}=\left[\begin{array}{c}{B}_{e,k}\\ A_{e,k}\end{array}\right],B_{e,k}\in {\mathbb{R}}^{d_{\mathrm{in}}\times r},A_{e,k}\in {\mathbb{R}}^{d_{\mathrm{out}}\times r}.$$

For thread $i$, let $k_i=\lfloor i/2\rfloor$ and let

$$s_i=\{\begin{array}{ll}+1& i\text{ even},\\ -1& i\text{ odd}.\end{array}$$

The perturbed forward pass for input $x_i$ is

$$y_i=x_i{W}^{\top }+\frac{\sigma s_i}{\sqrt{r}}{x}_i{B}_{e,k_i}{A}_{e,k_i}^{\top }.$$

This is the key computational trick. EGGROLL does not need to construct the dense perturbation

$$\Delta W_i=\frac{\sigma s_i}{\sqrt{r}}{A}_{e,k_i}{B}_{e,k_i}^{\top }\in {\mathbb{R}}^{d_{\mathrm{out}}\times d_{\mathrm{in}}}$$

before applying it. By associativity,

$$x_i\Delta W_i^{\top }=\frac{\sigma s_i}{\sqrt{r}}{x}_i{B}_{e,k_i}{A}_{e,k_i}^{\top }.$$

The term $x_i{B}_{e,k_i}$ first projects the input onto $r$ random directions in input space, producing only $r$ coefficients. Multiplication by $A_{e,k_i}^{\top }$ then turns those coefficients back into an output-space correction. For rank $1$, this is especially simple: with $B=b$ and $A=a$,

$$x_i\Delta W_i^{\top }=\sigma s_i(x_{i}b)a^{\top }.$$

So the perturbation is determined by two vectors, $b\in {\mathbb{R}}^{d_{\mathrm{in}}}$ and $a\in {\mathbb{R}}^{d_{\mathrm{out}}}$, not by a full matrix. For rank $r$, it is a sum of $r$ such vector outer products. The perturbation-side cost is therefore proportional to $(d_{\mathrm{in}}+d_{\mathrm{out}})r$ rather than $d_{\mathrm{in}}{d}_{\mathrm{out}}$, and the implementation can regenerate the factors from seeds instead of storing a dense perturbation for every sampled thread.

This does not remove the base model multiplication $x_i{W}^{\top }$; the frozen or trainable base weight is still a dense matrix. What EGGROLL avoids is the extra dense matrix multiplication that would come from explicitly sampling and applying a full ES perturbation of $W$.

With raw rewards $R_i$ and population size $N$, the ungrouped normalised fitness is

$$z_i=\frac{R_i-\mathrm{mean}(R)}{\sqrt{\mathrm{var}(R)+ϵ}}.$$

When several generations share the same prompt, the grouped version subtracts a within-prompt mean but keeps the global standard deviation:

$$z_{q,g}=\frac{R_{q,g}-\frac{1}{G}{\sum }_{h=1}^G{R}_{q,h}}{\sqrt{\mathrm{var}(R)+ϵ}},$$

where $q$ indexes prompts and $g$ indexes generations within a prompt group of size $G$.

The corresponding SGD-as-ascent update for a LoRA-structured EGGROLL layer is

$$\Delta W=\eta \frac{\sigma }{\sqrt{N}\sqrt{r}}\sum _{i=1}^N{z}_i{s}_i{A}_{e,k_i}{B}_{e,k_i}^{\top }.$$

The PyTorch implementation sets the gradient to the negative of this ES estimate so that the standard PyTorch optimiser step, which subtracts gradients, performs ascent on reward.

Decomposed ES target

The proposed method should keep the same ES estimator but change the searched coordinates. A DoRA-like layer writes the effective weight as a magnitude times a normalised direction. For a channel-wise magnitude vector $m$ and directional matrix $V$,

$$W(m,V)=m\odot \frac{V}{\Vert V{\Vert }_c},$$

where $\Vert \cdot {\Vert }_c$ denotes a per-output-channel norm and $\odot$ broadcasts $m$ over input channels.

A simple decomposed perturbation would sample low-rank directional noise and optional magnitude noise:

$$V_i=V+\frac{{\sigma }_V{s}_i}{\sqrt{r}}{A}_{e,k_i}{B}_{e,k_i}^{\top },m_i=m+{\sigma }_m{s}_i{ϵ}_{e,k_i}.$$

The noisy weight used for generation would then be

$$W_i=W(m_i,V_i).$$

The first real comparison should match compute and reward budget between this decomposed parameterisation and the plain EGGROLL baseline above. If decomposed ES only wins because it has more effective perturbation scale, more trainable scalar degrees of freedom, or a different normalisation, the comparison is not meaningful.

What would falsify this

• The PyTorch ES translation fails to reproduce expected behaviour on a simple smoke-test objective.
• Decomposed ES is consistently less stable or less sample-efficient than plain LoRA-space ES under matched compute.
• Any observed gain disappears once the comparison controls for trainable parameter count, perturbation scale, reward normalisation, and wall-clock budget.
• The method only helps on synthetic objectives and does not transfer to small long-context reward tasks.

Open questions

• What is the smallest exact-answer reward benchmark that is not completely flat under a small pretrained model and a matched ES budget?
• If fully sparse correctness is all zero, what shaped but still black-box reward is acceptable as the next bridge?
• What decomposition granularity is best for ES: per-weight matrix magnitude, per-output-channel magnitude, or a coarser adapter-level scale?
• Can the theory show lower variance or better conditioning in decomposed coordinates without overclaiming convergence for LLM training?
• Which ICLR framing is strongest: reward-only long-context adaptation, black-box agent optimisation, or scalable ES for parameter-efficient tuning?