Direct Coloring for Self-Supervised Enhanced Feature Decoupling

Fig. 1 shows a schematic diagram of our proposed framework. Our method is conceptually simple and is categorized as a non-contrastive approach, in the sense that it does not require negative views. Similar to some prior work such as [6, 39, 11], this framework is also constructed from a pair of symmetric networks, where the base design consists of encoders followed by projectors. The basic procedure starts with a random data augmentation function $\tau$ generating two augmented views for each sample image from a batch of samples $X$, namely $X_{1}$ and $X_{2}$, to be fed to the pair of symmetric networks. However, our framework distinguishes itself from others by further architectural modification as well as its loss function. The specific core idea is based on enhancing feature decoupling by inducing meaningful correlation (coloring) between features followed by eliminating unnecessary redundancy. We will empirically assess how this oriented coloring also diversifies the useful features. We view feature decoupling as the main role of data augmentation, and thus facilitates learning useful sparse features (as opposed to less useful dense features) [35, 36]. However, we must emphasize that only proper augmentation leads to desired feature decoupling, otherwise the learning would not capture plausible general features. Hence, we carefully engineer a coloring transform that allows for effective feature decoupling relying on proper augmentation. This coloring is shortly followed by a decorrelation (aka whitening) process allowing for dimensional collapse avoidance. Assuming the proper coloring (relying on proper data augmentation), the loss function is given as follows:

where $\mathcal{L_{C}}$ and $\mathcal{L_{W}}$ are coloring loss and whitening loss, respectively. Here we have:

where, $\lambda$ and $\alpha$ are weighting factors. Also $C_{ij}$ and $W_{ij}$ are elements of cross-correlation matrices computed for the coloring and whitening processes, respectively (see Fig. 1), while $E_{ij}$, is the target cross-correlation matrix used for the desired coloring. In other words, matrix $C$ is the cross-correlation matrix computed between the two output vectors of the coloring projector heads, whereas matrix $W$ is the cross-correlation matrix computed between the outputs of the final projector heads (whitening projector heads). Matrix $E$ is the desired colored cross-correlation matrix, computed as will be explained in the next section. We have:

where $z^{(1)}$ and $z^{(2)}$ are the normalized outputs of the coloring heads respectively for views $x_{1}$ and $x_{2}$, $m$ is the batch size. Similarly, $W{{}_{ij}}$ values are computed from the normalized outputs of the whitening heads.

A certain layer (close to the output layer) of each network is connected to a projector head intended for coloring the features at this level. The output layer is also connected to a projector head for the decorrelation process, whitening the embedding space. The coloring transform is performed on the cross-correlation matrix of the outputs of the projectors. The goal is to guide the network to target the desired features, by inducing controlled useful correlation in the cross-correlation matrix. The redundancy would be reduced in the next stage (the whitening process), by setting the cross-correlation matrix to an identity matrix. In the case of whitening, the diagonal elements are set to 1, encouraging similarity in representation, while the off-diagonal elements are to be close to zero, reducing redundancy in the feature representation [39].

Consider the standard practice of coloring transform under the Gaussian model for a multivariate Gaussian signal. Here the premise is that the signal (vector of random variables) at hand has a non-identity covariance matrix and the goal is to transform the signal to a new signal with the desired covariance matrix. To this end, first we need to perform the whitening transform, which transforms the multivariate Gaussian signal to a signal with an identity covariance matrix, similar to the ideal white noise. Whitening transform here is essentially a decorrelation by change of basis, and then scaling the principal axes to unit length. After that, the transformed signal will undergo a coloring process in which the data is scaled in desired directions toward desired variances and then rotated. This results in projecting the covariance matrix of the signal to that of an arbitrary multivariate Gaussian signal, hence the covariance matrix of the colored signal will be the desired covariance matrix [16].

However, in SSL literature in general, the purpose of whitening is mainly decorrelation, as was done in [11, 39]. We also follow this idea, aiming at instantiating the colored signal as the one with the desired cross-correlation. To this end, we deviate from the general procedure of coloring in which coloring is preceded by whitening, and introduce direct coloring relying on some Bayesian prior (as discussed in the next section) which is compatible with gradient based learning, while conceptually easier to perform. We perform direct coloring by projecting the (estimated) cross-correlation matrix to the desired cross-correlation via a gradient-based learning process minimizing the following functional:

where $E_{ij}$ are the elements of the desired cross-correlation matrix, whereas $C_{ij}$ are elements of the cross-correlation matrix computed from two embedding vectors.

We can identify two advantages of direct coloring:

1. 1.

   The process is computationally faster as it does not require us to initially project the elements of $C$ to the identity matrix before setting them to desired values.

2. 2.

   It does not assume a multivariate Gaussian, and in fact, the nature of the desired cross-correlation depends heavily on the effect of the argumentation.

Proper or standard augmentation is an integral part of any SSL framework as it decouples features properly, allowing for learing useful features. We build upon this to further enhance the feature decoupling. To this end, the desired colored cross-correlation matrix is computed from augmented data by further decorrelating the dense features while keeping the correlation between sparse features. Specifically, a pair of variational autoencoders (VAE) are trained separately on pairs of augmented views $x_{1}$ and $x_{2}$ for sample image $x$ under standard augmentation protocols. Then $E$, the cross-correlation between the normalized latent vectors from the VAEs are computed as the desired cross-correlation. See Fig. 1 (Right). The specific settings and scenarios are further discussed under experiments and ablation study. Elements of $E$ are computed along the whole dataset samples ($n$) as follows:

where $z^{(1)}$ and $z^{(2)}$ are the latent vectors of the top and bottom VAEs in Fig. 1 (Right).

Here we want to demonstrate that the proposed loss function based on a Bayesian prior (desired target cross-correlation) is a solution to the Maximum A Posteriori (MAP) estimation, specifically considering the prior and likelihood components in this work. The MAP estimation aims to find the most probable model parameters given both the data and some prior knowledge i.e., it combines a prior probability ${p(\Theta)}$ and a likelihood ${p(X|\Theta)}$ to find the best model parameters ${\Theta}$:

here we have a prior in the form of the target colored cross-correlation matrix $\mathbf{E}$ from VAEs, which specifies the desired correlation structure between the augmented views. Hence the prior is Gaussian distribution with mean ${E}$ and a certain variance $\sigma^{2}$ as follows:

Moreover, the likelihood is composed of two terms: the coloring likelihood $p_{color}(X|\Theta)$ and the whitening likelihood $p_{whiten}(X|\Theta)$, which capture the respective effects of the coloring and whitening processes on the learned representations. These likelihoods are defined as follows:

Thus, we have:

Since maximizing the product of probabilities is equivalent to minimizing the negative logarithm of the product, the MAP objective can be expressed in terms of our loss function:

Comparing this expression to our loss function $\mathcal{L}$, we can see that it aligns with the terms in Equation 12 and 13,i.e., joint optimization of whitening and Bayesian coloring terms in our loss function reflects the MAP estimation with a Gaussian prior (see supplementary for more analysis).

The main part of the experiments is performed on ImageNet dataset [10], under linear evaluation on ImageNet as well as transfer learning on smaller datasets for classification task. However, we also assess the approach on detection and segmentation tasks with different datasets. The experimental results are mainly obtained by building on Solo-Learn [9], a recently developed open access library of visual SSL approaches. Solo-Learn [9] provides an implementation of the existing baselines that we compared against in this section.

Datasets: In this study we make use of ImageNet [10], CIFAR10/100 [20], Tiny ImageNet [21], as well as VOC0712 [12] and COCO [22]. In our ablation study we use a separate dataset, ImageNet-100.

Baselines: For comparison, we contrast our framework to different classes of the recent baselines, including contrastive, non-contrastive, clustering-based, and whitening (aka redundancy reduction) baselines, as well as baselines primarily based on vision transformers. These baselines include SimCLR [6], BYOL [13], SimSiam [7], SwAV [4], Barlow-Twins (BT) [39], Whitening-MSE (W-MSE) with $d=4$ [11], and DINO [5].

Standard evaluation is performed under linear and transfer learning settings following the details of former baselines [11, 39, 13]. For linear evaluation, the common procedure is to remove the projector heads and train a linear classifier placed on top of one of the fixed encoders under the supervised setting over the source data (used for pre-training). For transfer learning purposes (classification, detection, and segmentation), the same standard procedure is performed except that the supervised training and evaluation is performed on the target data. Following [39], we also perform linear and transfer learning evaluation, except here we have four heads to remove, both whitening and coloring projector heads. In linear evaluation a whitening head is replaced with the linear classifier (a fully connected layer followed by a softmax) for the evaluation process. The learning rate for linear and transfer learning evaluation consistently starts with $10^{-3}$ and exponentially decays to $10^{-6}$ for some 500 epochs.

Linear evaluation: Linear evaluation on classification task is performed on ImageNet, Tiny ImageNet, CIFAR10, and CIFAR100. The evaluation consists of 500 epochs of supervised training (on labeled data) and then testing. Note that the encoder is fixed and only linear classifier undergoes training.

Transfer learning: Transfer learning consisted of pre-training on ImageNet dataset and evaluation on other datasets, as presented in the results. Accordingly, on classification task, transfer learning was performed with CIFAR10 and CIFAR100, whereas on detection and segmentation task it is performed on VOC0712 and COCO respectively. Similar to the case of linear evaluation, the supervised training is performed for 500 epochs before testing the performance.

In this section the results for different datasets, tasks, and learning paradigms are presented. Following is respectively the classification results with ImageNet, CIFAR1/100, and Tiny ImageNet, as well as object detection results with VOC0712 and segmentation results with COCO.

We assess the effect of the presence, and location of the coloring heads. Given the standard setting presented for the framework, the coloring heads are connected to layer 16 of ResNet18, with hidden layer dimension of 2048. Under this setting, the baseline performance on classification task is a top-1 accuracy of $80.93\%$. Removing the coloring heads and corresponding optimization terms from the loss function, the framework is reduced to BT, and the accuracy descends to $79.69\%$. However, with the current output size of 2048, there is a $1.24\%$ improvement with the coloring heads. Next we assess the performance with coloring head connected to different layers.

Coloring at layer 10: When the coloring heads are connected to layer number 10, the top-1 accuracy descends to $79.97\%$ (from $80.93\%$). We suspect that this drop is mainly because layers 11-16 are no longer under coloring process. Here, the desired coloring cross-correlation is computed from a VAE with both encoder and decoder of size 10, same as the location of coloring heads.

Coloring at layer 17: When the coloring heads are connected to the same layer as the whitening heads, layer number 17, the performance drops to $78.81\%$. Roughly speaking, coloring could be seen as the opposite of whitening. Accordingly we hypothesize that what actually happens in case of parallel whitening and coloring heads, is that the result follows the superposition principle, meaning that coloring is cancelled out as the weight factor $\lambda$ in Equation 1, gives less weight to coloring.

The experiments with a range of different projector dimensionality indicate the sensitivity of the method to this factor, as the top-1 accuracy with output size 512, 1024, 2048, and 4096 are respectively $74.31\%$, $77.15\%$, $80.93\%$, and $81.66\%$. This is a similar behavior to that of Barlow-Twins.

We briefly assess the case in which coloring matrix is computed from a pair of autoencoder, instead of VAEs. The performance degrades to $80.21\%$. The slightly better performance of VAEs might be due to the fact that the latent space of VAE is regularized, which allows for more robust cross-correlation in terms of variance of the elements.

We consider an architecturally simpler instantiation of the direct coloring in which, the cross-correlation is replaced with auto-correlation, as shown in Fig. 3. The idea turned out to be effective. As shown in Fig. 3, the the framework consists of only one network, in which the encoder undergoes the coloring before passing its output to whitening heads. Specifically, layer 16 of the encoder is connected to a coloring head where the auto-correlation of its output is becoming colored. Similarly in case of whitening, the auto-correlation of the output is becoming whitened. The desired coloring matrix is also an auto-correlation matrix computed from the latent space of one VAE. Amazingly, while the computational complexity reduces by half, the performance only slightly drops to $80.64\%$. Corresponding loss function of coloring and whitening process as well as the total loss is presented in the supplementary material.

Let’s consider direct coloring a rather general technique for any SSL framework. Here, in terms of loss function, we have two terms, first term being the original loss of the framework to be ungraded denoted as $f$, plus a second term corresponding to the direct coloring, hence we have $\mathcal{L}=f+\lambda g$. In essence, direct coloring in this setting would be seen as a constrained optimization problem in terms of Lagrange multipliers. The first term alone, $f$, would be solved by a trivial solution, complete collapse. However, this is subjected to the second term, $g$, as the coloring constraint. We observe that direct coloring would substantially reduce the chance of complete collapse, because to find the solution to this constrained optimization problem, the optimizer looks for points where the gradient vector of $f$ and $g$ are parallel to each other. Since complete collapse is only a solution to $f$ alone and certainly not a solution to $g$, if one chooses a proper $\lambda$, one can substantially avoid the complete collapse. Further theoretical and empirical evidence on this as well as empirical evidence on other baselines with direct coloring technique is presented in supplementary .