Stanford's SOMA Slashes AI Training Noise by Sharding Models Into Expert Shards
A new architecture called SOMA shards models into independent experts, slashing the gradient variance that has blocked zero-order methods from scaling to pretraining.
- Stanford paper SOMA shards a model into independent LSTM experts trained with zero-order SPSA.
- Separable per-expert losses provably cut relative gradient variance to roughly 1/N of a shared-loss estimator.
- At 8.44M params and 150 GPU-hours, SOMA N=2 hits 1.76 nats/byte versus 2.00-2.21 for monolithic baselines.
- Experts exchange no gradients, activations or optimizer state, enabling disaggregated training across GPUs.
- SOMA N=256 with top-4 routing delivers 9.19x inference throughput at lower loss than N=8.
- Code and checkpoints promised but not yet released at time of preprint.
SOMA shards models to scale forward-only training
A Stanford preprint proposes an architectural fix for a central problem in zero-order optimization: gradient estimates become noisier as the number of jointly perturbed parameters grows. Scaling Zero-Order Pretraining through Model Sharding (paper) introduces SOMA, which divides a model into independently trained experts so each estimator handles fewer parameters. The design could make larger models practical on forward-only hardware, with non-differentiable objectives, or in clusters constrained by interconnect bandwidth.
Random directions hit a dimension wall
Zero-order optimization updates weights from loss evaluations without computing derivatives. Its common workhorse, SPSA, perturbs every selected parameter along a random direction, compares the losses from positive and negative perturbations, and averages estimates from multiple directions.
With M jointly perturbed parameters and K perturbation directions, the leading relative gradient variance scales approximately as M/K. Maintaining a comparable estimate as the model grows therefore requires proportionally more perturbations, each adding forward evaluations.
Methods including MeZO, HiZOO, and EGGROLL reduce that cost through optimizer changes such as low-rank perturbations. SOMA instead limits the dimension seen by each estimator through model sharding.
Separate losses contain the noise
SOMA partitions the model into N expert submodels, implemented as LSTMs in the paper. A predetermined router assigns each data cluster to an expert, and SPSA trains every expert against its own scalar loss. Experts exchange no gradients, activations, or optimizer state during training.
This story is for Pro members
You've reached the end of the free preview. Upgrade to AlphaSignal Pro to read the full article - and everything else behind the paywall.
Comments
