Unsqueeze [CLS] Bottleneck to Learn Rich Representations

TL;DR

UDI tackles information over-compression in self-supervised learning by unsqueezing the CLS bottleneck and introducing a multimodal, multi-level distillation target. It combines context-aligned semantic constraints via self-attention, a shared projector for image- and patch-level semantics, and an extra class token to capture nuisance information, all within a teacher–student ViT framework with Stratified Random Sampling. Empirically, UDI achieves competitive ImageNet-1K results in linear and k-NN evaluations, strong low-shot and transfer performance, and substantial gains on dense prediction tasks such as COCO detection/segmentation and ADE20K segmentation, with ablations validating the benefits of its components. The work demonstrates that preserving nuisances alongside semantics yields richer representations and broader applicability across vision tasks, albeit with a ViT-centric scope and opportunities to extend to other backbones.

Abstract

Distillation-based self-supervised learning typically leads to more compressed representations due to its radical clustering process and the implementation of a sharper target distribution. To overcome this limitation and preserve more information from input, we introduce UDI, conceptualized as Unsqueezed Distillation-based self-supervised learning (SSL). UDI enriches the learned representation by encouraging multimodal prediction distilled from a consolidated profile of local predictions that are derived via stratified sampling. Our evaluations show that UDI not only promotes semantically meaningful representations at instance level, delivering superior or competitive results to state-of-the-art SSL methods in image classification, but also effectively preserves the nuisance of input, which yields significant improvement in dense prediction tasks, including object detection and segmentation. Additionally, UDI performs competitively in low-shot image classification, improving the scalability of joint-embedding pipelines. Various visualizations and ablation studies are presented to further elucidate the mechanisms behind UDI. Our source code is available at https://github.com/ISL-CV/udi.

Figures (9)

• Figure 1: Information compression. UDI retains the most information from $\boldsymbol{X}$ as its attention map of $\texttt{CLS}+$ token (UDI$_+$) closely aligned with the underlying semantics.
• Figure 2: Semantic constraints of SA layer in SRS module.
• Figure 3: Visualization of [cls] attention from the last layer of ViT-S/16 trained with self-distillation based SSLs. DINO300 denotes DINO trained with 300 epochs.
• Figure 4: UDI Framework. UDI is an SSL method based on the joint-embedding strategy with multilevel self-distillation objectives. Specifically, for each image, UDI creates two views with one cropped out from the other, followed by two random augmentations, respectively, for student network $h_\text{s}\circ\text{SRS}_\text{s}\circ f_{\text{s}}$ and teacher network $h_\text{t}\circ\text{SRS}_\text{t}\circ f_{\text{t}}$. UDI employs ViT with an extra class token $\boldsymbol{z}_{\texttt{cls}+}$ for encoder $f$. The dense features from $f$ are then sampled and processed by a Stratified Random Sampling (SRS) module to produce patch-level representations $\widehat{\boldsymbol{Z}}_p$. The class token $\boldsymbol{z}_{\texttt{cls}+}$ is learned to produce multimodal prediction against a target distribution $p^*_t({y}|\boldsymbol{Z}')$ constructed with patch-level predictions ${p}_t(y|\hat{\boldsymbol{z}}')$ and image-level prediction $p_t(y|\boldsymbol{z}'_{\text{cls}})$. The final UDI objective involves maximizing the agreement via cross-entropy loss between (i) $p_t(y|\boldsymbol{z}'_\text{cls})$, $p_s(y|\boldsymbol{z}_{ \text{$\text{cls}$} })$, (ii) $p^*_t(y|\boldsymbol{Z}')$, $p_s(y|\boldsymbol{z}_{\text{cls}+})$, and (iii) ${p}_t(y|\hat{\boldsymbol{z}}'_i)$, ${p}_s(y|\hat{\boldsymbol{z}}_i)$, respectively.
• Figure 5: Clustering process of SSL objectives. (Left) the InfoNCE objective drives an image $\boldsymbol{x}_i$ closer to its augmented views $\boldsymbol{x}_j$ (positive sample) while being far away from other images (negative samples). In the representation space, this leads to the formation of small clusters centered around each image $\boldsymbol{x}$, with their covariance depending on augmentation $\boldsymbol{T}$. (Right) explicit clustering-based SSL, such as DINO and SwAV, partitions the entire dataset into $K$ clusters.