ueaj @_ueaj
ueaj @_ueaj · 18h
New blog! You can just keep increasing the amount of heads in your model with no diminishing returns on ICL up to atleast 4x. For reference that would make the o_proj head dimensions in this experiment 16k x 2k.
Additionally, if you perform a truncated SVD on full rank master weights to train MLA instead of training them as two separate matrices, you can recover most of the ICL capability but with less memorization.
I think MLA specialized optimizers are a direction worth exploring and are very underserved rn. Unfortunately I have more important projects to attend to and I've burned like 300$ on compute for this already. I would highly recommend someone trying to scale this up and see how well we can do.
[Chart: "arc/output_loss" — lines for mla-rat-l512-r64-b0p9-arc1-200m, mla-base-l512-r64-b0p9-arc1-200m, quad-kv-b0p9-arc1-200m, single-kv-b0p9-arc1-200m, x-axis train/tokens_seen 20M-180M, y-axis loss 0-8, all converging to similar low loss; tooltip shows single-kv-b0p9-arc1-200m: 0.74451, quad-kv-b0p9-arc1-200m: 0.66843 at 199,884,800 tokens]
ueaj @_ueaj · Apr 15
Replying to @_ueaj
Also something like MLA should be trained like QAT but instead of converting a high precision matrix to low precision you convert a full rank master weights into low rank latent projections...
Note from Claude Sonnet 5
A technical thread on Multi-head Latent Attention (MLA) architecture experiments — scaling attention heads without diminishing ICL returns, and using truncated SVD on full-rank weights (analogous to quantization-aware training) to recover in-context-learning capability with less memorization. Directly relevant to the brain_graph_1 "KDA-MLA attention at the ends only" architecture plan noted in project memory.
machine learningattention mechanismsmlamulti-head latent attentionin-context learningtwitterueajbrain_graph_1architecture research