Sorry for delay! Here is my work on zero order (ZO) optimization.
We achieve SOTA for ZO methods on pretraining controlling for compute and parameters.
ZO methods struggle to improve loss
as model size grows because relative gradient variance increases linearly with the number of perturbed parameters, inhibiting large model training.
All previous methods innovate on the optimizer, but adopt architectures designed for backprop.
Instead, we design an architecture with ZO in mind, that caps the gradient variance as you increase model size.
Introducing SOMA (Sharded Optimization Mixture of Assemblies). SOMA shards the model into tiny experts, and trains each expert independently, on disaggregated GPUs, without communicating gradients, activations or optimizer state during training.
Inspired by the thalamus and cortical columns, each expert specializes on a subset of the train set based on a fixed router.
At isocompute and isoparams, SOMA achieves lower loss vs. all monolithic ZO methods tested (e.g. EGGROLL, more perturbations w Vanilla SPSA, etc).
ZO w SOMA continues to improve at larger model sizes.
Additional inference benefits of this architecture, we can select top-k active experts over N trained experts to reduce inference compute by O(N/k), allowing us to tradeoff inference flops for accuracy wo retraining.