Step size scaling for larger models

[Howuhh]

Hi! I tried scaling MLP sizes a bit, and almost always they start to learn worse or slower. I noticed that for larger widths (e.g. 1024+) step size is always stays at almost zero (<1e-8), while all other metrics are of similar magnitude compared to 128 hidden (like td error). The only change is that l1 norm, which is currently a global sum over all parameters. For larger models this sum grows quickly and reduces the learning rate.

Is this expected or normal? It feels weird to me. Larger models do not necessary should have smaller step sizes.

I tried to replace global sum with the average, but from initial experiments it does not work...

[mohmdelsayed]

Hi @Howuhh, this is a great question. Your observation makes sense, since the L1 scales with the width of each layer, leading to a very small step size. Performance scaling with the network width is something we are actively working on and will be included in a new paper revision, along with other exciting new things, so please stay tuned!

I’ll leave this issue open and follow up once the new revision is released.

[Howuhh]

Thanks @mohmdelsayed!

Actually, we tried the same thing for few weeks now, with the goal to advance sota on dm-hard (from simba-v2) as much as possible under some constraints, like only 2M transitions. We tried wide variety of empirical tricks from the literature (bronet, simba, simba-v2 and many other) but none of them worked. However, I only know notices the issue with learning rate, so maybe those results are not as credible due to that.

Two most interesting observations so far:

1. Inspired by the work https://arxiv.org/abs/2504.01961, I measured cos sim between consecutive gradients / gradient interference. For actor I don't notices anything strange. But for critic, cos sim always near 1 with high interference (with and without sparsity), which is usually a bad sign in off-policy learning (like simba). From test on multiple envs, I also notices that I can remove traces for critic and results stays mostly the same, while for actor traces are critical and without them training diverges. Given high similarity of gradients it kinda make sense (?), that if all steps are same, then traces do not distribute credit assignment for truly important transitions. However, I tried to orthogonalize grads via g_t - proj_g_prev g_t and it breaks training. Not sure for now...

Example from dm_control/cartpole-balance-v0:

2. I measured Deviation from Isometry (DfI), which is a common metric in continual learning (e.g. https://arxiv.org/abs/2602.08040), basically, how orthonormal our weight matrix is) during training. I was curious to find out what effect sparse init has. So, mostly, it reduces DfI for the actor, which is positively correlated with the final return, and without sparse init DfI grows indefinitely. As a simple test, I switched the init from default (in my case, lecun normal in flax) to orthogonal init with scale 1. This resulted in a much larger drop in DfI, stable training and higher final return. After that, sparse init for the actor almost don't do anything (but still seems useful for critic, not sure). I planned to explore some strategies to orthogonalize weights during training via the newton-shulz, at small scale first.

Example from dm_control/cartpole-balance-v0:

I run those from my rewrite on jax: https://github.com/Howuhh/streaming-drl-jax

P.S. May be you need some working hands or help with the experiments for the paper? I am glad to help and I have access to some GPU's. And I love the work you do, although I am a total noob in theory!

[mohmdelsayed]

Hi @Howuhh,
Thanks for sharing your observations! Once we release our paper revision, I’ll let you know. We can discuss ways to improve sample efficiency in streaming deep RL. I’d love to see streaming deep RL methods achieve state-of-the-art performance, and it would be amazing to have your help with that. Let’s stay in touch and continue this discussion over email.