EGGROLL: Low-Rank Evolution Strategies for Hyperscale Optimization of Billion-Parameter Models

Forum topic · ✨步子哥 · 2025-11-23

Summary

EGGROLL (Evolution Guided General Optimization via Low-rank Learning) is a black-box optimization algorithm that replaces the full-rank perturbations of traditional evolution strategies (ES) with low-rank matrix products. Instead of sampling an m x n perturbation matrix costing O(mn) memory and compute, EGGROLL generates two small random matrices A (m x r) and B (n x r) with r << min(m, n), forming perturbations of the form (1/sqrt(r)) * A B^T with complexity O(r(m+n)). Deterministic random number generators reconstruct noise on the fly, decoupling memory from population size, while batched low-rank adapters let all population members share a single forward pass. Theory shows low-rank updates converge to full-rank ES updates at an O(1/r) rate, and population averaging yields high effective-rank updates. Experiments demonstrate up to 100x training throughput gains, parity with full-rank ES on reinforcement learning benchmarks, competitive performance with GRPO for LLM reasoning fine-tuning, and stable pretraining of purely integer-type recurrent language models. The official JAX implementation is open-sourced as a research preview in the HyperscaleES GitHub repository, with a simplified single-file version in the nano-egg repo.

EGGROLL: Low-Rank Evolution Strategies for Hyperscale Optimization

EGGROLL (Evolution Guided General Optimization via Low-rank Learning) is a gradient-free black-box optimization algorithm designed to scale evolution strategies (ES) to models with billions of parameters. This post summarizes the algorithm's design, theory, applications, comparisons, and open-source implementation.

Key points

Problem: Traditional ES perturbs every weight matrix with a full-rank noise matrix, costing O(mn) memory and compute per population member — infeasible for billion-parameter models.
Solution: EGGROLL replaces full-rank perturbations E ∈ R^{m×n} with low-rank products ε = (1/√r) A B^T, where A ∈ R^{m×r}, B ∈ R^{n×r}, and r << min(m, n). The 1/√r factor keeps variance matched to the full-rank case.
Complexity: Storage and forward-pass cost drop from O(mn) to O(r(m+n)). With m = n = 1000 and r = 10, this is roughly a 50x reduction.
Engineering tricks: Deterministic random number generators rebuild perturbations on demand from seeds (near-zero memory, decoupled from population size N), and batched low-rank adapters let the whole population share one base forward pass: y = Wx + (σ/√r) A (B^T x).
Theory: Low-rank updates provably converge to full-rank ES updates at an O(1/r) rate; averaging over a sufficiently large population yields high effective-rank updates. Even r = 1 works without significant performance loss.

Applications and results

Reinforcement learning: On tabula-rasa RL benchmarks (e.g., MuJoCo continuous control), EGGROLL matches full-rank ES in final reward and convergence while delivering up to 100x training throughput gains — approaching pure batch-inference speed.
LLM fine-tuning: For improving reasoning, EGGROLL's performance is competitive with GRPO (Group Relative Policy Optimization), offering a gradient-free alternative suited to sparse, result-only reward settings.
Integer models: EGGROLL enables stable pretraining of nonlinear recurrent language models operating purely on integer datatypes, where gradients are undefined and backpropagation-based methods fail.

Comparisons

| Feature | Traditional ES | EGGROLL | | :--- | :--- | :--- | | Perturbation | Full-rank E ∈ R^{m×n} | Low-rank (1/√r) AB^T | | Storage | O(mn) | O(r(m+n)), ~0 via on-the-fly noise | | Forward-pass cost | O(mn) | O(r(m+n)) | | Throughput | Limited | Up to 100x improvement |

vs. LoRA: LoRA is parameter-efficient fine-tuning that trains small bypass adapters while freezing base weights; EGGROLL updates all parameters using low-rank perturbations, and its population-averaged updates can achieve high effective rank — potentially advantageous on tasks needing complex reasoning.
vs. gradient-based methods (Adam/PPO): EGGROLL needs no differentiable objective, avoids credit assignment, is robust to hyperparameters, and mitigates reward hacking; gradient methods remain preferable when a differentiable loss exists (e.g., supervised fine-tuning).

Open-source implementation

Official JAX implementation: ESHyperscale/HyperscaleES on GitHub (research preview). Core code lives in src/hyperscalees/, with an end-to-end example in tests/end_to_end_test.py and LLM experiment entry point python -m llm_experiments.general_do_evolution. A Dockerfile is included.
A simplified single-file implementation, nano-egg, demonstrates pure int8 pretraining of language models with EGGROLL.
Requirements: Python ≥ 3.10, JAX, NumPy; optional vllm and tensorboard for accelerated LLM fine-tuning experiments.
Caveat: The code is a research preview — expect instability, evolving interfaces, and the need for custom engineering (distributed orchestration, hyperparameter tuning of N, r, σ, monitoring) before production use. Recommended practices include phased deployment, cloud infrastructure, containerization, and active monitoring.