BatchEnsemble with more layers and use_bias.
BatchEnsemble with more layers and use_bias.

1. Because the additional cost is negligible, we add BE to layer index 22, the representation_size projection layer, and the head layer.
2. It turned out the BE layers used a setting of no bias. This likely reduces model capacity or trainability so we turn them back on.

I also updated ed.nn.DenseBatchEnsemble:

+ Rename use_ensemble_bias to use_bias, following flax.linen's use_bias.
+ Set default kernel_init to nn.initializers.lecun_normal(), following flax.linen.Dense. This default doesn't affect ViT BatchEnsemble because, like ViT's encoder we change the kernel_init to xavier_uniform(). However, ViT's representation_size Dense layer does use the default kernel init, so BE does too.

PiperOrigin-RevId: 425532701

google/edward2