Typical cosine similarity range

[sebgao]

Hi authors,

Thanks for the excellent work and I am very interested about this paper! I am trying to reproduce it, especially the training part. I am curious about the typical range of the cosine similarity loss, since in my experiment the alignment metric peaks above 0.9 during early iterations, then remains on a long plateau, which seems very strange. Could you share what range the alignment loss typically lies in and whether this behavior is expected?

[hila-chefer]

Hi @sebgao, thanks for your interest!
in practice, we employ the 1s complementary of cosine similarity for the feature loss (i.e., lambda x, y : torch.mean(1 - torch.nn.CosineSimilarity(dim=-1)(x, y)), mostly for loss visualization purposes, gradient is invariant to this).
The logged loss for the ImageNet experiment (first 500k steps are logged) is attached below, I wouldn't say it looks like the loss is plateauing, which suggests there could be some issue with the implementation? what coefficient did you use for the feature loss?

[sebgao]

Thank you for your valued reply! My coefficient for feature alignment loss follows the paper and is set to 0.8. I am investigating my implementation further but my loss curve looks like your provided one till 100K iters! By the way, is there any EMA decay warmup from 0.0 to 0.9999 in your implementation?

[hila-chefer]

happy to help! the similarity until step 100k is a good sign :-) we don't employ a warmup for EMA, but that could possibly help boost performance to some extent.

[Shenyi-Z]

Hi @hila-chefer and @sebgao, thank you so much for the great work and discussion!

I've been trying to reproduce that curve and ran into one detail at step 0 that I can't quite reconcile with the released model code, so I wanted to ask in case I'm misreading something.

If I trace through initialize_weights, the last linear of every block'sadaLN_modulation is zero-initialized:

# Self-Flow/src/model.py L351-L353
for block in self.blocks:
    nn.init.constant_(block.adaLN_modulation[-1].weight, 0)
    nn.init.constant_(block.adaLN_modulation[-1].bias, 0)

So at step 0 we get gate_msa = gate_mlp = 0 and every block reduces to the residual identity (block(x, c) = x). Together withEMA = deepcopy(model), the raw features at any layer of student and teacher are just x_embedder(x_τ) + pos_embed andx_embedder(x_τ_min) + pos_embed. With R_M = 0.25, 75% of the tokens have x_τ[i] = x_τ_min[i], so the raw cos(s_l, t_k) is already very close to 1, which lines up beautifully with the start of your figure (1 - cos ≈ 0.005).

The part I'm having trouble with is that Eq. 6 places h_θ = SimpleHead on the student side, and both of its linears seem to receive xavier from the global _basic_init:

# Self-Flow/src/model.py L329-L336
def initialize_weights(self):
    def _basic_init(module):
        if isinstance(module, nn.Linear):
            torch.nn.init.xavier_uniform_(module.weight)
            if module.bias is not None:
                nn.init.constant_(module.bias, 0)
    self.apply(_basic_init)

If my understanding is correct, no matter how close s_l is to t_k, a random nonlinear projection in D = 1152 should send s_l to an essentially random direction relative to t_k. The expected cosine would then be ≈ 0, so the loss should start at1 - cos ≈ 1 rather than ≈ 0.

Could you point me to what I'm missing? For instance, is there an additional initialization on the projector (e.g. zero-init on the last linear plus a residual to approximate identity at step 0), or does the teacher side also pass through the same projector (BYOL-style), or something else entirely? I just want to make sure I'm not overlooking a detail of the training-side implementation.

Thanks again, and apologies if I missed something obvious!

[sebgao]

Hi @Shenyi-Z, in my case the 1.0 - cos metric actually starts from ~1.0 at step 0 and drops below 0.1 within a few iterations.

[Shenyi-Z]

Hi @Shenyi-Z, in my case the 1.0 - cos metric actually starts from ~1.0 at step 0 and drops below 0.1 within a few iterations.

Thanks @sebgao ! I observed similar results in my own experiments, so this makes perfect sense.