Our design choices for the projection and the prediction head are even simpler than the original BYOL [1]}.
We replace the BatchNorm with LayerNorm and the ReLU activations with GeLU activations, making the overall framework free of BatchNorm and consistent with the rest of the transformer blocks.
Our work also incidentally demonstrates that BYOL does not rely on BatchNorm to prevent collapse [2]}.
The projection head and the prediction head have 3 and 2 hidden layers respectively, a hidden dimension of 4096, and an output dimension of 256, following the original design.
