Change default initializers for RBM to be scale invariant
Created by: PhilipVinc
This PR does 2 thins: changes the default initialzier of DenseSymm
and RBMSymm
to have a variance 1/(features * n_sites)
, and also change the default initializer of Dense
to lecun_normal
(so a rescaled normal distrbution with variance 1/n_sites
) instead of the fixed value of 0.01
we had previously. This also affect RBM
.
With regards to RBM/Dense
, I think a more appropriate initializer should be variance 1/(alpha*n_sites)
, but flax reccomends lecun_normal
so i'm not sure what would be best...
before merging this i'd like everyone to run a few tests on this branch, for example for code that uses the default initializers of Dense and DenseSymm.