Exploring the Role of the Bottleneck in Slot-Based Models Through Covariance Regularization

Object-centric representations hold the potential to significantly improve the generalization capabilities of computer vision models through their ability to factorize and represent a scene as the composition of objects. Greff et al. (2020) Goyal and Bengio (2020) van Steenkiste et al. (2019) Despite growing interest Locatello et al. (2020) Singh et al. (2021) Singh et al. (2022) Kipf et al. (2021) Elsayed et al. (2022), most slot-based models still struggle on scenes with complex textures and on real world images. Karazija et al. (2021) Papa et al. (2022) Recent work has shown that strong encoders are likely needed to produce latent features that are clusterable by Slot Attention from these complex scenes. Seitzer et al. (2022) Biza et al. (2023) However, these stronger encoders loosen the architecture’s information bottleneck and reduce the incentive for object separation between slots. This tradeoff between a stronger encoder and object separation is an artifact of a reliance on architectural approaches to changing the strength of the bottleneck.

In this paper, we attempt to overcome this tradeoff through the use of a loss function-based approach to constricting the bottleneck. Specifically, we adapt the projection head and variance and covariance losses from VICReg Bardes et al. (2021) to slot-based methods. By placing additional restrictions on the slot features, the reconstruction task is made more difficult and the bottleneck is constricted.

Our original hypothesis was that our proposed loss would improve instance segmentation performance and produce slot representations that improve downstream performance. While we do not achieve these results in this work, we confirm the intuitions about the role of the bottleneck in slot-based architectures underlying these predictions and present some interesting preliminary results.

Slot Attention Locatello et al. (2020) is a convolutional autoencoder that iteratively applies a modified attention mechanism to latent vectors in order to obtain a permutation-invariant set of object-specific representations, called slots. The slot attention module relies on a variation of dot-product attention that treats slots as queries that compete to "explain away" the encoder output. While initial iterations of this architecture fail on datasets with complex textures Karazija et al. (2021), recent work demonstrates that a stronger backbone allows this architecture to achieve a more competitive performance. Biza et al. (2023)

DINOSAUR Seitzer et al. (2022) is a slot-based architecture that relies on a frozen DINO ViT Wysoczańska et al. (2022) as its encoder and trains using a DINO-feature reconstruction loss. The feature reconstruction objective is justified through an experiment where Slot Attention and DINOSAUR with a MLP decoder are trained on COCO using a frozen ViT encoder with an image reconstruction and DINO-feature reconstruction objective, respectively. While Slot Attention learns a degenerate stripe pattern, the feature reconstruction objective learns instance segmentation. Despite its feature reconstruction objective, DINOSAUR still suffers from training instabilities, preventing the end-to-end training of the architecture. Seitzer et al. (2022)

VICReg Bardes et al. (2021) is a self-supervised, non-contrastive representation learning method with a joint embedding architecture that relies only on its loss function to prevent collapse. This is a simpler architecture than other non-contrastive methods such as BYOL Grill et al. (2020) or SimSiam Chen and He (2020) which frequently rely on momentum encoders, stop gradients, quantization, or batch normalization in order to avoid collapse.

There have been a number of attempts to incorporate contrastive losses into slot-based models for video Kipf et al. (2019) Racah and Chandar (2020). Unlike these works, we apply a non-contrastive objective to slot representations as a sort of regularization term and are not restricted to training on video.

Slot Attention produces a representation of an input image using a convolutional encoder and then iteratively applies the slot attention module in order to obtain a permutation-invariant set of $K$ slots, each with $D_{slots}$ dimensions. Locatello et al. (2020) We propose the addition of a MLP projection head $h_{\phi}:\mathbb{R}^{D_{slots}}\to\mathbb{R}^{D_{proj}}$ where $D_{proj}\gg D_{slots}$ (see Figure 1). In this projection space, we will compute the variance and covariance losses:

where $z_{j}$ is the $j^{th}$ dimension of the projection space, $\gamma$ is a hyperparameter representing the desired variance for each feature, and $\epsilon$ is a small scalar included for numerical stability. Covariance is calculated over the features of the projection space, giving a covariance matrix of size $D_{proj}\times D_{proj}$. The total loss function for our method is:

where $X$ is an input image, $\hat{X}$ is the reconstruction, $Z$ is the projection of the slot representations for $X$ through $h_{\phi}$, and $\beta$ is a hyperparameter that weights the variance and covariance losses and dictates the strength of the information bottleneck.

Empirically we find that the variance loss term is not necessary for our task. While VICReg Bardes et al. (2021) finds that the variance loss term helps avoid representation collapse, we find that our primary reconstruction objective is sufficient to prevent collapse and that the variance loss can be removed. We call this loss using only the reconstruction and covariance terms covLoss.

In addition to covLoss, we test another loss called cosineLoss:

where $cosemb(z_{i},z_{j})=\max(0,\cos(z_{i},z_{j})-0.2)$. This loss is similar to the InfoNCE loss van den Oord et al. (2019), and the intuition is to encourage the model to make embeddings of slots far from each other. Note that unlike covLoss, this loss is applied directly to the slot feature vectors. Attempts to apply this loss to the projection of the slot vectors diverges. Intuitively, since $D_{proj}$ can be quite large compared to $D_{slots}$, the projection head can learn to project similar slot feature vectors to constant orthogonal spaces, that can pose a local optima for the cosineLoss and hinder the effect we seek to achieve.

The primary contribution of this project is a loss-based approach to constricting the bottleneck of slot-based models. These losses are proposed in hopes of allowing the use of the stronger encoders needed to produce cluster-able features from real-world images while still achieving object separation. While the results we show in this report are not perfect, they validate our intuitions about the role of the bottleneck in producing object separation. From our experiments, it appears that the covLoss and cosineLoss losses are an insufficient bottleneck to produce strong object separation results, although we were unable to thoroughly explore different hyperparameter combinations and design decisions due to constraints on our computational resources.

Other approaches to constricting the bottleneck include experimenting with different, weaker decoder architectures. Recent work like SLATE Singh et al. (2021) proposes an alternative to the spatial broadcast decoders used throughout this project. However, these architectural approaches make it necessary to adjust the decoder architecture for each dataset or downstream task which is burdensome and can be non-trivial. Another source of future work is the exploration of objective functions besides reconstruction. While DINOSAUR is a first step towards moving beyond image reconstruction objectives, it still suffers from training stability issues that prevent the entire architecture from being trained end-to-end. The intersection of slot-based models and self-supervised learning is practically unexplored and a promising direction for future research.