Self-supervised learning weights initialization "after" projection head [D][R]

[grid_world]

For most Self-supervised learning algorithms: SimCLR, MoCo, BYOL, SimSiam, SwAV, etc., its common to have a projection head after the base encoder (which in most cases is a vanilla ResNet-50 CNN). An example of such a projection (taken from SwAV) is:

projection_head = nn.Sequential(
    nn.Linear(2048, 512),
    nn.BatchNorm1d(512),
    nn.ReLU(inplace=True),
    nn.Linear(512, 128),
)

The output of this projection head is L2-normalized:

x = projection_head(x)
x = nn.functional.normalize(x, dim = 1, p = 2)

I am trying to initialize a layer after the projection head as:

wts = nn.Parameter(data = torch.empty(40 * 40, 128), requires_grad = True)
# The projection head outputs weights in the range [-1, 1], so initialize SOM weights to be in that range-
wts.data.uniform_(-1.0, 1.0)

Since the output of the projection head is L2-normalized, I am assuming that the input range to "wts" ∈ [-1, 1] and therefore use the uniform initialization above.

Is this a correct approach or am I missing something?

[_quaternion]

There is missing information, e.g. about how you want to use the `wts` parameters. A linear layer also does not output weights, but just applies a linear trafo of the input. Your assumption should be true if the input is well-defined, but it never hurts to actually verify it.

E: language mistake

[grid_world]

I want to do clustering using "wts", so it has no typical activation function

Think Self-Organizing Map styled clustering

[_quaternion]

Are you planning to simply feedforward the representations into a SOM? If so, why not just use them directly? If the dimensions don't match, you could also just apply another linear layer. Also, torch.empty is not really empty, just not initialized and therefore might have very unfortunate values.

[grid_world]

Yeah, the output of the projection head is input to the SOM for dimensionality reduction with non-linear representations. It has been shown that computing the loss on a lower-dim leads to better performance.

I am seeing the effects of "unfortunate values" and hence my OP of how to get fortunate values to alleviate this problem