Learning Decorrelated Representations Efficiently Using Fast Fourier Transform

TL;DR

This paper introduces a fast Fourier Transform-based regularizer for decorrelated representation learning that reduces computational complexity from quadratic to near-linear, maintaining accuracy while improving efficiency.

Contribution

It proposes a novel decorrelating regularizer computed via FFT, significantly speeding up training for high-dimensional embeddings compared to existing methods.

Findings

Achieves comparable downstream task accuracy to existing regularizers.

Reduces training time and memory usage for large embedding dimensions.

Provides an effective technique to avoid undesirable local minima.

Abstract

Barlow Twins and VICReg are self-supervised representation learning models that use regularizers to decorrelate features. Although these models are as effective as conventional representation learning models, their training can be computationally demanding if the dimension d of the projected embeddings is high. As the regularizers are defined in terms of individual elements of a cross-correlation or covariance matrix, computing the loss for n samples takes O(n d^2) time. In this paper, we propose a relaxed decorrelating regularizer that can be computed in O(n d log d) time by Fast Fourier Transform. We also propose an inexpensive technique to mitigate undesirable local minima that develop with the relaxation. The proposed regularizer exhibits accuracy comparable to that of existing regularizers in downstream tasks, whereas their training requires less memory and is faster for large d. The source code is available.