From stateful code to purified JAX:�how to build your neural net framework
Sabrina J. Mielke • @sjmielke • 2021-07-01
JAX as “numpy + jit()”
“Tired of Topic Models? Clusters of Pretrained Word Embeddings Make for Fast and Good Topics too!”
Sia, Dalmia, Mielke. EMNLP 2020.
Speed and coolness!
JAX as “numpy + jit()”
omg Jeff Dean!
be still my beating heart
JAX as “numpy + grad()”
Reversible jump MCMC requires the Jacobian determinant of a diffeomorphism:
jit() + vmap()
= speeeeeed
Tape-based autograd, e.g. PyTorch:
Pure transformation-based autograd: JAX
But how to “train” a big neural network?
How can weights and gradients be stateful if our Tensors can’t be mutated?
Basis for most of this talk:�https://sjmielke.com/jax-purify.htm�...a blogpost from March 2020!
Is it still relevant?
Good news: it never was*! :)�(* or it will never not be)
What will we accomplish?
No SotA models or frameworks.
No “idiomatic” elegance.
No batching (vmap), distributing (pmap), and XLA-compiling (jit).
A “module” in PyTorch
modules own their weights...
...and use them when building
the execution graph/tape
Using a PyTorch module: mutation!
mutate gradients and weights!
...off to the notebook!
pure-ification
Let the user write stateful methods and the framework convert such functions into actually pure functions that we can call grad() on!
The purification scheme, visually
This is what I want it to look like:
not quite self.w…
but close enough.
use as before :)
Let’s implement this first...
...off to the notebook!
ours
dm-haiku
randomness handling!
ours
flax.linen
More details and complications:�https://sjmielke.com/jax-purify.htm
Follow me on Twitter @sjmielke :)