SOMA: Scaling Zero-Order Pretraining through Model Sharding

SOMA reduces interference in zeroth-order updates by training recurrent experts on separate losses. It improves over tested monolithic ZO controls, not over BPTT. Its main tradeoff is independent optimization versus jointly learned representations.

Source

Status And Credibility

Stanford-authored preprint with a detailed appendix, controlled independent-versus-summed-loss experiments, and explicit negative BPTT comparisons. No conference acceptance was verified. The homepage’s earlier work date is not the arXiv submission date. The paper says training/evaluation code and checkpoints are released, but this ingest did not locate a paper-matched public training repository or checkpoint collection in the paper, author homepage, or public author repository listing. The browser demo is verified, not the entire claimed reproduction package.

The demo has evolved: its current default is a 1B research model with a learned SetAttn-EMA combiner, alongside the original 617K tiny backprop-free model. Do not treat that later demo’s default as the paper’s fixed-router probability-average system or as evidence that its entire training was ZO.

Mechanism

  • Cluster a byte-level FineWeb-Edu corpus; initialize expert bodies from a seed model.
  • Each expert is an LSTM/MLP predictor with its own recurrent state and local objective, trained with SPSA-style weight probes. Experts do not jointly optimize the routed ensemble loss.
  • A fixed TF-IDF/SVD/centroid router selects up to four experts from observed context; their next-byte probabilities are averaged.
  • The main small-expert campaign shares a frozen tied embedding/decoder learned during the seed stage. The seed decoder receives an exact local cross-entropy gradient, without BPTT through the recurrent body. This is not an entirely derivative-free training pipeline.
  • Recipe distinction: the wide fixed-total-size controls with retain and update separate tied heads. The width-32 campaign uses one shared frozen head. Do not describe every reported SOMA result as frozen-head training.
  • SOMA uses Adam-style moments for estimated body gradients. No BPTT does not mean no optimizer-state memory.

For a fixed separable objective, equal-sized parameter blocks, dense Rademacher probes and small perturbations, the paper derives the leading relative-variance ratio

This removes cross-expert perturbation noise for that separable problem; it does not prove that partitioning an arbitrary coupled world model preserves its optimum, or that routed inference loss improves by the same factor.

Evidence And Accounting

ComparisonPrimary evidenceBoundary
Modest sharding versus monolithic ZOTable 4, roughly 8.44M parameters and 150 aggregate GPU-hours: gets 1.76 test nats/byte, monolithic SPSA 2.00–2.11, reproduced EGGROLL 2.21.Compares measured systems; EGGROLL is a reproduced FP32 GRU recipe, not the original integer model.
Controlled local-loss mechanismSame architecture, starting states, data, probes and update counts: independent losses improve test loss by 0.035 nats/byte after 1,000 updates, across three seeds.Stronger mechanism isolation than comparing different architectures; not a universal loss guarantee.
BPTT comparisonTable 5, PDF p.28: at 100 GPU-hours, SOMA 1.8257, BPTT 1.0535, in ensemble validation nats/byte.Values interpolate recorded curves; hardware and recipes differ; both exclude initialization. The paper explicitly says these recipes do not show a training advantage over backpropagation.
Inference throughput versus : 2.36M versus 257k byte-tokens/s, including routing, at test losses 1.68 versus 1.71.The better endpoint used 59.9× aggregate training compute. This is an inference advantage, not a matched-training-budget win; BPE-token throughput is not directly comparable.
ContextTraining windows are 1,024 bytes; evaluation routes on 256 observed bytes and scores 768 targets, resetting state between windows.Does not establish reliable memory across ten thousand steps or unbounded-context learning.

The approximately 80k estimated GPU-hours describe the broader research campaign, not one matched-budget checkpoint. Aggregate GPU-hours, idealized parallel elapsed time, actual elapsed time, arithmetic operations, initialization and preprocessing MUST remain distinct in comparisons.

Social Claim Audit

The author’s October thread argues that small independent experts are the only scalable path for ZO pretraining and that accumulated probes must span parameter space. Those are stronger claims than the paper establishes. The paper itself acknowledges other ZO pretraining routes, including EGGROLL. Accumulated directions at different weights do not reconstruct one stationary gradient, and spanning all directions is neither necessary nor sufficient for successful optimization.

The root post criticizes ZO speedrunning generally; it does not explicitly name Noise or Dust. In the surrounding Noise discussion, the relevant concern is extrapolating quick threshold attainment into large-scale convergence. Its shared scaling caveat is legitimate, but Dust probes activations with token-level signals, not simply every weight against one global scalar. The dimensionality argument must be evaluated for the actual estimator.

Audit Of The Supplied Analysis

The user-provided second-agent analysis correctly identified local expert losses, loss of joint representation learning, the BPTT memory advantage, and the approximate 1.83 versus 1.05 comparison. Corrections preserved here: the numbers are validation nats/byte, not an unspecified error or test score; GPU budgets are not identical hardware/recipe controls; shared versus trainable heads differ by experiment; and 1,024-byte reset windows do not test very long temporal credit. The hypothetical 10,000-step example is a mechanism illustration, not an experiment in SOMA.

Foundation TSFM Relevance

Agenda slotVerdictEvidenceMissing pieces
Streaming state, long context, and constant updatesadjacentRecurrent experts can be trained without temporal backward graphs.Long-memory retention, continuous state carry and numeric trajectories.
Native multivariate encodingwarningIndependent experts isolate optimization noise.Cross-channel coupling and shared latent-state coordination may be lost.
Dynamic compute allocationadjacentRouting bounds active experts.End-to-end cost-quality curves including seed, routing and all training work.

Open Questions

  • Can a small ZO-trained recurrent memory module preserve rare early events at matched GPU-hours and peak memory versus full/checkpointed/truncated BPTT?
  • Can experts exchange enough predictive information for coupled multivariate dynamics without reintroducing shared-loss noise?
  • Which gains survive validation-only selection, multiple seeds, longer state-carry windows and a fixed total-compute budget?