Port of μP to JAX (Haiku specifically)

[davisyoshida]

μP is a method for reparameterizing NNs so that hyperparameter optima stay fixed as you scale them up. I wrote a port to JAX here.

As an example, I trained a bunch of transformers on PTB, with and without μP. As you can see, μP causes the optimal learning rate to stay fixed as you scale the models up:

(In the paper, the authors do this experiment but for LMs with 6.7B parameters).

This is specifically for Haiku, but if people are interesting in making versions for FLAX etc., it might be helpful to take a look at this. I had to change the design substantially from the authors' original repo, since a lot of things which are easy in Torch (via mutability) aren't in JAX, and vice versa. I'm also interest in any suggestions on improving the design or increasing usability.

[YouJiacheng]

Great Job!
But Why not implement muP in optax?

[YouJiacheng]

Yes. I mean that you can implement it in optax, and provide a initializer(per array or initialize the whole pytree at once.)

[YouJiacheng]

Besides, you can make the first update of muP optimizer correctly initialize the parameters, though it is a bit tricky.

[davisyoshida]

Unfortunately you can't figure out what transformations you need to do from only the tensor shape. You do need to know a little bit about how it's being used to know what dimension is fan-in/fan-out. You could maybe guess from the parameter names, but I preferred to create a custom haiku parameter creator so you can directly check which module is requesting the parameter.

[YouJiacheng]

Yes, but that will substantially couple muP library with a NN library.
I preferr to ask user to provide fan-in axis information by a pytree.
And muP library author/community can implement a automatic fan-in axis assigner for each popular NN libraries.
This make muP library portable and extensible.

[davisyoshida]

You're right, that's a good idea. I'll take a shot at decoupling them.