Slight update, more fancy initialization of DyT weights (instead of having them be ones) seems to help a lot in my case (although it's still not as good as just

Slight update, more fancy initialization of DyT weights (instead of having them be ones) seems to help a lot in my case (although it's still not as good as just using RMSNorm). Do something like this on the very first training step (`x` is the input to the layer): y = x.to(torch.float32) y = y * torch.rsqrt(y.pow(2).mean(-1, keepdim=True) + 1e-6) z = torch.tanh(self.alpha * x) scale = (y / (z + 1e-6)).mean(dim = -2).flatten() self.weight.detach().copy_(scale) This basically tries to initialize the weights so that the output of DyT is closer to what RMSNorm would have outputted, and it seems to help.