Issue with Beta NfNet

[simomagi]

I think that the implementation of nfnet does not consider the update of beta. In the paper https://arxiv.org/pdf/2101.08692.pdf , the authors update beta after each residual block as:

$\beta_i = sqrt (expectedvariance_i) 
expectedvariance_{i+1} = expectedvariance_{i}  + \alpha^{2}$ , for i=1..n,
 where n is the number of the residual blocks and expectedvariance_0 = 1

In the implementation, it seems to me that beta is equal to the initialization everywhere. Is this correct?

[vballoli]

Yeah, it'll be fixed soon in addition to the training scripts. Thanks a lot for noticing this and bringing this up. Appreciate the effort!

[simomagi]

Nice! Furthermore i think that the output of the activation should be divided by beta and not multiplied. I am working to fix the code. If you'd like i will make a pull request.