Affine transformation estimation improves visual self-supervised learning

TL;DR

This work augments multi-view self-supervised learning with a model-agnostic affine transformation estimation module that adds a predictive geometric loss $l_{ ext{affine}}$ for affine parameters $\boldsymbol{\phi}=[\theta, t_x, t_y, \sigma, s_x, s_y]$. By balancing invariance to standard SSL augmentations with equivariance to affine changes, the approach improves downstream accuracy and accelerates convergence across SimCLR, BYOL, and Barlow Twins on Tiny ImageNet-pretrained encoders evaluated on CIFAR-10/100 and Caltech101. Extensive ablations show that a single-view, vector-difference aggregation often yields the best trade-off between performance and compute, while the four affine components collectively offer the most benefit. The findings suggest that injecting a geometric supervision signal provides complementary information to the conventional ID-based SSL loss, enabling more expressive representations with modest parameter overhead (~4%).

Abstract

The standard approach to modern self-supervised learning is to generate random views through data augmentations and minimise a loss computed from the representations of these views. This inherently encourages invariance to the transformations that comprise the data augmentation function. In this work, we show that adding a module to constrain the representations to be predictive of an affine transformation improves the performance and efficiency of the learning process. The module is agnostic to the base self-supervised model and manifests in the form of an additional loss term that encourages an aggregation of the encoder representations to be predictive of an affine transformation applied to the input images. We perform experiments in various modern self-supervised models and see a performance improvement in all cases. Further, we perform an ablation study on the components of the affine transformation to understand which of them is affecting performance the most, as well as on key architectural design decisions.

Figures (7)

• Figure 1: SimCLR simclr architecture.
• Figure 2: Architecture diagram depicting our method. Our contributions are enclosed in the purple box. The remainder of the diagram illustrates a typical, modern self-supervised method. Note that we only depict the case of one random view being used in our module because the process is exactly the same if the other random view is included.
• Figure 3: Downstream performance for SimCLR with and without the affine module during training.
• Figure 4: Downstream performance for BYOL with and without the affine module during training.
• Figure 5: Downstream performance for Barlow Twins with and without the affine module during training.