Implement Barlow Twins

[OlivierDehaene]

🌟 New SSL approach addition

Approach description

Implement Barlow Twins (arxiv link).

Pseudocode:

# f: encoder network
# lambda: weight on the off-diagonal terms
# N: batch size
# D: dimensionality of the representation
#
# mm: matrix-matrix multiplication
# off_diagonal: off-diagonal elements of a matrix
# eye: identity matrix

for x in loader: # load a batch with N samples
    # two randomly augmented versions of x
    y_a, y_b = augment(x)

    # compute representations
    z_a = f(y_a) # NxD
    z_b = f(y_b) # NxD

    # normalize repr. along the batch dimension
    z_a_norm = (z_a - z_a.mean(0)) / z_a.std(0) # NxD
    z_b_norm = (z_b - z_b.mean(0)) / z_b.std(0) # NxD

    # cross-correlation matrix
    c = mm(z_a_norm.T, z_b_norm) / N # DxD

    # loss
    c_diff = (c - eye(D)).pow(2) # DxD
    # multiply off-diagonal elems of c_diff by lambda
    off_diagonal(c_diff).mul_(lambda)
    loss = c_diff.sum()

    # optimization step
    loss.backward()
    optimizer.step()

Open source status

The model implementation is not yet available. However, it will be open sourced at: https://github.com/facebookresearch/barlowtwins

• the model implementation is available
• the model weights are available
• who are the authors: Jure Zbontar, Li Jing, Ishan Misra, Yann LeCun, Stéphane Deny

[OlivierDehaene]

I would be glad to work on this if you want.

[prigoyal]

EDIT: We are checking with the authors if they planned to provide an implementation themselves! I'll update here if authors plan to do it in which case I'd recommend to wait. We will communicate the plan.

Hi @OlivierDehaene , the proposal looks great and we would love to have this in VISSL. Go ahead :)

Also take a look at projects/ folder where you can add a README etc to ensure authors are properly credited for their work in VISSL.

[IgorSusmelj]

You could use the code here: https://github.com/IgorSusmelj/barlowtwins
We used the SimSiam model and just swapped the loss. We're still tuning some parameters since we're running the code on CIFAR10 and the paper does not mention experiments on this dataset.

[OlivierDehaene]

@prigoyal sorry I only just saw your edit.

I started an implementation here.

[jzbontar]

Hi Oliver,

paper author here. Thanks for volunteering to implement Barlow Twins in VISSL! Please go ahead with the implementation.
I can review your code and, potentially, run it on our cluster to verify the accuracy.

Also, the official code should be out next week. We are making final tweaks to the codebase.

Best,
Jure

[OlivierDehaene]

Hello Jure,

Cool! I think I'll wait for the code to come out before continuing as I'm not really sure about the value of some hyper-parameters.

Nice paper by the way :)

Cheers

[prigoyal]

@jzbontar , it it possible to share the hyperparams setup in @OlivierDehaene 's PR? :)

[jzbontar]

@OlivierDehaene,

I also think that waiting for us to release to code would make sense so that the VISSL implementation and our official implementation match. Like I said, you can expect us to release the code sometime next week.

Out of curiosity, which hyper-parameters are you unsure about?

[OlivierDehaene]

Out of curiosity, which hyper-parameters are you unsure about?

@jzbontar I only have questions regarding the optimization hyper-parameters:

• Do you adjust the learning rate as a function of the minibatch size like in MoCo and SimCLR?
• Do you use mixed precision? If so and with Apex, using which optimization level?

The first can be inferred from the "We follow the optimization protocol described in BYOL" sentence but I would prefer to be sure :)

[jzbontar]

• This is how we set the learning rate:

def adjust_learning_rate(args, optimizer, loader, step):
    max_steps = args.epochs * len(loader)
    warmup_steps = 10 * len(loader)
    base_lr = args.learning_rate * args.batch_size / 256
    if step < warmup_steps:
        lr = base_lr * step / warmup_steps
    else:
        step -= warmup_steps
        max_steps -= warmup_steps
        q = 0.5 * (1 + math.cos(math.pi * step / max_steps))
        end_lr = base_lr * 0.001
        lr = base_lr * q + end_lr * (1 - q)
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr
    return lr

• We use torch AMP and train in half precision.

[jzbontar]

We released the code for Barlow Twins.