simple ViT implementation

[EIFY]

includes necessary modules (LN, bias, scale, constant posemb) and notebook showing it working on MNIST. Initial tuning shows that momentum w/ dualize not quite as performant as Adam but

1. That was a 10-minute effort
2. Not sure how serious we should take MNIST performance after 1000 steps anyway

More of a demonstration but could be merged.

[EIFY]

I have ported this "dualized ViT" to Big Vision and started training a dualized ViT-S/16 on ImageNet-1k.
TL;DR: So far it severely underperforms the baseline but there are still a few knobs to turn.

Implementation: I made a custom branch of modula that dry-runs dualize recursively and caches the target_norm for each atom. I then export them along with the class name of the atom module and graft the behavior of dualize onto the Big Vision implementation by using optax.partition of the optax nightly (see the Big Vision branch that does this). So far the training loss looks like this:

wandb report

LR=0.05, WD=0.005 or 0.0001 (decoupled from LR), momentum w/ beta=0.95
No warm-up, cosine learning rate decay.
Notably l2_params starts out not too different from the baseline but grows quickly to 50x of that. My knee-jerk reaction to raise WD from 0.0001 to 0.005 however seems to make the model even worse:

I don't have great intuition here. Perhaps it still needs LR warm-up even though the model training is stable without it?

Notable architecture (or rather, just scaling) differences from the baseline:

1. GELU is scaled by 1/1.1289 to keep the derivative <= 1.
2. The dot product of the dot-product attention is scaled by 1/d instead of 1/sqrt(d) where d is the dimension of the attention head (same as μP). The final attention output is then scaled by 1/3.
3. The residual branch is scaled by s = 1/(2 * depth) and the residual connection is normalized as x = (1-s) * x + s * y. This is more aggressive than the 1/sqrt(depth) residual branch scaling for ViT here

Optimizer differences from the "conventional" muon:

1. Due to the different target_norm values we effectively have different LR for each layer.
2. Instead of AdamW fallback we have momentum + L2-normalized update for bias (same as Embed) and Lion-like momentum + sign update for scale (following the derivation of muon and consider the scale of layernorm as the linear operator of diag(scale)).
3. We don't clip the sqrt(fanout / fanin) factor from below as sqrt(max(1, fanout / fanin)). It's what it is.

If necessary it may be interesting to bisect these differences.

[gngdb]

I noticed this too and corrected it in #13. I guess it doesn't really matter for performance?