Loss implementation about barlowtwins

[laura-wang]

Thanks for your quick reponse! I've tried, they produce the same results. Thanks a lot!

[jzbontar]

Hi Laura,

the psudocode in the paper does in fact show a different way of computing the Barlow Twins loss compared to the method we use in main.py. However, both ways of computing the loss should produce the same result.

Consider the following code snippet:

import torch

def off_diagonal(x):
    n, m = x.shape
    assert n == m
    return x.flatten()[:-1].view(n - 1, n + 1)[:, 1:]

c = torch.rand(3, 3)
lambd = 0.1

# compute loss as in the psuedocode
c_diff = (c - torch.eye(3)).pow(2)
off_diagonal(c_diff).mul_(lambd)
loss = c_diff.sum()
print('from pseudocode:', loss.item())

# compute loss as in main.py
on_diag = torch.diagonal(c).add_(-1).pow_(2).sum()
off_diag = off_diagonal(c).pow_(2).sum()
loss = on_diag + lambd * off_diag
print('from main.py:', loss.item())