Question about weight_decay interaction with fraction

Hello,

Thanks for making this wonderful work available.

I have a question regarding the interaction between weight_decay and ortho_fraction in Dion3. weight_decay is applied to the whole weight matrix whereas scatter_add_ is applied on the top_k. configs/dion3_160m.yaml and configs/normuon_160m.yaml kept weight_decay constant, but I want to understand whether there's a need to scale weight_decay for consistency between Dion3 with ortho_fraction<1 and NorMuon. Would the weight_decay need to be smaller when using a small ortho_fraction since rows are decayed more often than they are selected?

I intended to do hyperparameter tuning, but the recommendation is to use a larger ortho_fraction for smaller model, and a larger one for large models. Would the optimal hyperparameters shift with the change in ortho_fraction and what is your recommendation?

Start by comparing the weight_decay and ortho_fraction settings in configs/dion3_160m.yaml and configs/normuon_160m.yaml, then trace how Dion3 applies decay and scatter_add_ to the selected rows. Done means establishing whether scaling is needed and documenting how ortho_fraction changes hyperparameter recommendations.