zeroth-order optimization

Francois Chaubard @FrancoisChauba1 · 11h
for everyone new to zero order, zoo is the future and has amazing properties over BP.

but requires a change in architecture as there is no "gradient flow" you need to worry about. you just jump. rendering most previous architectures over engineered and previous mental models unhelpful.

but the issue w ZO will always be dimensionality that you cant solve wo changing the architecture. i tried the last 2 years wo much success on anything except fine tuning. could never pretrain.

here is the key issue:
if you want to train an 8B parameter model, over lets say 100k steps, w 64 perturbations, then you have only ever checked 6 million halfspaces, less than 0.1%. so 99.9% of the directions in the model have never been explored ever. thus the gradient noise is way too high esp at later stages of training.

initial training will look good, bc initially, all directions give finite differences in loss. but later training stalls abruptly, bc in late training almost all directions are degenerate so every probe produces near 0 finite difference. (this is also why speedrunning zoo is not a good idea.. wont tell you much..)

but split those 8B parameters into independently trained 10K experts, SOMA-style, and w just 10 perturbations per step for 1K steps, each expert gets 10k probes, enough for those directions to span all directions in design space. finite differences remain high always.

estimating a useful update in 10K dimensions isnt too hard and even better each expert training is independent loop, so no comms btwn gpus!

this enables pretraining at scale w zoo. and nothing else does. not hyperbole. ive trained everything else. you need to shard the model.

this is how the brain is organized too btw. all sensory information is sent to the thalamus (except olfactory) which routes those signals to tiny cortical columns in your neocortex that are massively inhibited and train independently from each other and we ensemble the results at test time. [cut off]

zeroth-order optimizationpretrainingbackpropagationmixture of expertssomaneuroscienceml architecture