Understanding per-head orthogonalization lr

[eric-tc-wong]

Hello,

Thank you for implementing this feature. I did some quick tests but don't see much improvement from orthogonalizing heads independently (needs verification). One change is in the adjust_lr of q/k/v, which I think is scaled smaller by 1/√H when using num_heads=12. I don't quite understand the implications. Allowing heads to orthogonalize independently probably generates larger updates, but could reducing lr by 3.46x in my 12-head model be over-compensating for that? Is this something I have to tune?

Would appreciate your insights. Thanks!

[alint77]

It also gives worse loss in my testing, and in modded-nanogpt as well (they use NS on 2heads together instead of one iirc).

But the main point of per-head-NS is to reduce the computation cost, since doing NS on n small matrices is cheaper than a big one with the same size. Noah Amsel (@NoahAmsel) has done some benchmarking on this and measured the speedups.

[eric-tc-wong]

Thank you for the confirmation! I hope this means my setup is likely correct.
I also started reading on the computation speedup a moment ago and will need to check whether it's worth optimizing for the system I'm on.

My main interest in this feature comes from the statement: "Orthogonalizing the fused matrix blends information across heads — often not what you want." Orthogonalization of the full matrix may help reduce overlapping features across heads. Allowing a mixture of overlapping and non-overlapping heads might be best according to Zhang et. la. https://arxiv.org/html/2605.08933v1
In this paper, they also mentioned the attention head shape scaling, so assume they did some testing.

For head-wise and group-wise Muon variants, splitting a Q/K/V matrix changes the effective matrix shape on which the polar-factor update is applied. We therefore apply a shape-aware learning-rate transfer

Kimi K3 seems to apply per-head orthogonalization, claiming it helps stability. I'm worried this might become important when scaling to larger production model, but it'll be expensive to do A B comparison on large models.

[JohnLangford]

Per-head loss makes sense in my opinion, but we haven't been able to make it work empirically yet. This might be a matter of further testing to figure out the right way to do things.

[eric-tc-wong]

Thanks, John. That is my opinion as well.
I'm also currently running some experiments and I haven't been able to show the stability benefits of per-head orthogonalization yet. Per-head orthogonalization does lead to higher loss, which corroborate results from Zhang et. la. https://arxiv.org/html/2605.08933v1 . However, my main concern is stability at production scale so I'm willing to trade some performance.
A confounding factor might be the weights decaying base_lr *weight decay, while the updates are scaled by 1/sqrt(group_heads).

[JohnLangford]

Sounds good.