CrIBo : Self-Supervised Learning via Cross-Image Object-Level Bootstrapping

1. Semantically Coherent Image Regions

The first step towards cross-image object-level bootstrapping is to identify semantically coherent image regions within the images. To that end, we rely on an online joint-space clustering algorithm from [Stegmuller_2023_CVPR].

Joint-space clustering.

$${\mathbf{c}}_i^k=\sum _{n\in [N]}{\mathbf{Q}}_i^{\ast }(n,k)\cdot {\mathbf{z}}_i^n,i\in \{1,2\}$$

2. Cross-Image Object Matchings

In the section on semantically coherent image regions, we detail the methodology for identifying objects within the views of a given image and establish the strategy for their alignment from one view to another. Following this step is the critical phase of aligning such objects across images. We rely on a pair of memory banks ${\mathcal{M}}_1$ and ${\mathcal{M}}_2$, populated with object representations extracted from the first and second views during previous iterations. Thus, the memory banks ${\mathcal{M}}_i$ can be queried with an object representation ${\mathbf{c}}_i^k$ to retrieve its nearest neighbor:

$$\text{NN}({\mathbf{c}}_i^k,{\mathcal{M}}_i)\triangleq \underset{\mathbf{c}\in {\mathcal{M}}_i}{\arg max}\frac{{\mathbf{c}}^{\mathrm{\top }}{\mathbf{c}}_i^k}{\Vert \mathbf{c}{\Vert }_2\Vert {\mathbf{c}}_i^k{\Vert }_2}$$

Here, we consider $\text{NN}({\mathbf{c}}_1^k,{\mathcal{M}}_1)$ and ${\mathbf{c}}_1^k$ as cross-image object-level matches.

Cycle-consistent matchings.

Figure: Illustration of cycle-consistent matchings. Such matchings are invariant to data augmentations and reciprocal, loosely speaking.

At the beginning of the pretraining phase, when the encoder is randomly initialized, relying on the nearest neighbor operator to identify cross-image object-level matches is not well grounded. Indeed, this straightforward approach assumes that the representations of similar objects are close in the embedding space, an assumption that does not hold at the initialization. Therefore, enforcing consistency between such pairs of semantically distinct objects should be avoided. To that end, we design a semantic closeness condition, which is only fulfilled when the matches are established based on high-level features, thus enabling the discarding of invalid positive pairs that do not share a semantic similarity.

At a high level, we consider a pair of objects as a valid positive pair if they are the nearest neighbors of each other. In practice, the criterion is implemented as follows. Starting from an object representation ${\mathbf{c}}_1^k$ in the batch of first views, its nearest neighbor $\mathtt{\text{nn}}({\mathbf{c}}_1^k,{\mathcal{M}}_1)$ in the queue is found using the nearest neighbor operator. To ensure that the criterion is invariant to data augmentation, we then switch to the queue of second views via the $\mathtt{\text{jump}}(\cdot )$ operator:

$$\mathtt{\text{jump}}({\mathbf{c}}_1^k)\triangleq {\mathbf{c}}_2^k,\mathtt{\text{jump}}({\mathbf{c}}_2^k)\triangleq {\mathbf{c}}_1^k$$

From this step, we obtain $\mathtt{\text{jump}}(\mathtt{\text{nn}}({\mathbf{c}}_1^k,{\mathcal{M}}_1))$ and retrieve its NN in the batch of second views $\mathtt{\text{nn}}(\mathtt{\text{jump}}(\mathtt{\text{NN}}({\mathbf{c}}_1^k,{\mathcal{M}}_1)),{\mathcal{B}}_2)$. The final step consists of switching back to the batch of first views and verifying if the resulting object is the same as the one we started from. More formally:

$$\mathcal{O}({\mathbf{c}}_i^k)\triangleq \{\begin{array}{ll}true& \text{if}{\mathbf{c}}_i^k=\mathtt{\text{jump}}(\mathtt{\text{nn}}(\mathtt{\text{jump}}(\mathtt{\text{nn}}({\mathbf{c}}_i^k,{\mathcal{M}}_i)),{\mathcal{B}}_j))\\ false& \text{otherwise}\end{array}i,j\in \{1,2\},i\ne j$$

3. Self-Supervised Training Objectives

CrIBo’s final training objective involves three distinct types of self-supervision: cross-view object-level, cross-image object-level, and global cross-image. All levels of self-supervision are based on the self-distillation mechanism, which employs a student-teacher pair of siamese networks $f_s$ and $f_t$, respectively. Given two augmented views ${\tilde{\mathit{x}}}_1$ and ${\tilde{\mathit{x}}}_2$, global representations and object representations can be obtained for each view and each network within the siamese pair.

Cross-View Object-Level Self-Supervision

The object representations from the teacher and student are fed to their respective head $h_s$ and $h_t$ to output probability mass functions whose sharpness can be modulated via temperature parameters (${\tau }_s$ and ${\tau }_t$). For the teacher, this translates to:

$${\mathbf{p}}_{i,t}^k=\mathtt{\text{softmax}}({\text{head}}_t({\mathbf{c}}_{i,t}^k)/{\tau }_t),i\in \{1,2\}$$

The student's projections are obtained analogously. The cross-view object-level loss is expressed as the averaged cross-entropy over all pairs of object representations and making the loss symmetric with respect to the student/teacher:

$${\mathcal{L}}_{cv}^o=\frac{1}{K}\sum _{k\in [K]}(H({\mathbf{p}}_{1,t}^k,{\mathbf{p}}_{2,s}^k)+H({\mathbf{p}}_{2,t}^k,{\mathbf{p}}_{1,s}^k))$$

where $H(\mathbf{a},\mathbf{b})=-\sum _{l=1}^L{\mathbf{a}}_l\log ({\mathbf{b}}_l)$ denotes the cross-entropy.

Cross-Image Object-Level Self-Supervision

Once all object representations in the batch have been retrieved, the cross-image object-level loss is enforced in a manner analogous to its cross-view counterpart. The nearest neighbors are first fed to the projection head as follows:

$${\tilde{\mathbf{p}}}_{i,t}^k=\mathtt{\text{softmax}}(h_t(\mathtt{\text{nn}}({\mathbf{c}}_{i,t}^k))/{\tau }_t),i\in \{1,2\}$$

Note that these projections are only defined for the teacher since the $\mathtt{\text{nn}}(\cdot )$ operator is not differentiable. The formulation of the cross-image object-level loss aligns with that in the cross-view loss up to the filtering of invalid positive pairs:

$${\mathcal{L}}_{ci}^o=\frac{1}{Z_1}\sum _{k\in [K]}{\mathbb{1}}_{\{\mathcal{O}({\mathbf{c}}_1^k)\}}H({\tilde{p}}_{1,t}^k,p_{2,s})+\frac{1}{Z_2}\sum _{k\in [K]}{\mathbb{1}}_{\{\mathcal{O}({\mathbf{c}}_2^k)\}}H({\tilde{p}}_{2,t}^k,p_{1,s})$$

where ${\mathbb{1}}_{\{\cdot \}}$ denotes the indicator function and where $Z_1$ and $Z_2$ are normalization constants, i.e., $\sum _{k\in [K]}{\mathbb{1}}_{\{\mathcal{O}({\mathbf{c}}_1^k)\}}$ and $\sum _{k\in [K]}{\mathbb{1}}_{\{\mathcal{O}({\mathbf{c}}_2^k)\}}$, respectively.

Global Cross-View Self-Supervision

Importantly, both the clustering and the bootstrapping steps are contingent on the quality of the representations. To avoid potential instabilities, we incorporate a global cross-view loss into our framework, offering meaningful self-supervision at any stage of the pretraining. This loss is analogous to the one from the cross-view object-level loss except that it is applied to the global representation (${\mathbf{z}}_{1,s}$, ${\mathbf{z}}_{2,s}$, ${\mathbf{z}}_{1,t}$, and ${\mathbf{z}}_{2,t}$), which are fed to a dedicated head $\bar{h}$:

The global cross-view projections are defined as follows:

$${\bar{p}}_{i,t}=\mathtt{\text{softmax}}({\bar{h}}_t({\mathbf{z}}_{i,t})/{\bar{\tau }}_t),i\in \{1,2\}$$

where ${\bar{\tau }}_s$ and ${\bar{\tau }}_t$ are the corresponding temperature parameters for the $\bar{L}$-dimensional output distribution. The global cross-view loss is defined as:

$${\mathcal{L}}_{cv}^g=H({\bar{p}}_{1,t},{\bar{p}}_{2,s})+H({\bar{p}}_{2,t},{\bar{p}}_{1,s})$$

where $H(\mathbf{a},\mathbf{b})=-\sum _{l=1}^{\bar{L}}{\mathbf{a}}_l\log ({\mathbf{b}}_l)$.

Finally, we obtain the overall training objective of CrIBo as a sum of its individual components:

$${\mathcal{L}}_{\text{tot}}={\mathcal{L}}_{cv}^g+{\mathcal{L}}_{cv}^o+{\mathcal{L}}_{ci}^o$$

The student's parameters are optimized to minimize the above loss whereas the parameters of the teacher are updated via an exponential moving average of the student's weights.