Probabilistic Slot Attention: learning identifiable object representations

In brief

• Standard slot attention binds image features into object slots, but lacks mathematical guarantees that representations are identifiable across training runs.
• Probabilistic Slot Attention (PSA) reformulates slot competition as inference within a hierarchical Gaussian Mixture Model (GMM).
• The paper establishes theoretical identifiability conditions, proving that slots recover true object properties up to permutation and linear scaling.

This retrospective explainer is dated to the December 2024 NeurIPS conference. It discusses Identifiable object-centric representation learning via probabilistic slot attention by Avinash Kori, Francesco Locatello, Ainkaran Santhirasekaram, Francesca Toni, Ben Glocker and Fabio De Sousa Ribeiro.

Suppose two people sort a toy box containing red cubes, blue cylinders, and yellow spheres into three storage bins. The first person might put all cubes in bin 1, cylinders in bin 2, and spheres in bin 3. The second person might group them by color, putting red items in bin 1 and yellow items in bin 2. Both groupings successfully partition the toy box, but they have learned incompatible definitions of what makes an “object.”

This is the core challenge of representation identifiability in object-centric deep learning. Standard Slot Attention architectures use iterative competitive routing to assign visual tokens to latent vectors (“slots”). While they produce striking visual segmentation masks on multi-object benchmarks like CLEVR, there is no mathematical guarantee that the internal coordinate systems learned by different slots will be consistent or unique.

Probabilistic Slot Attention (PSA) solves this problem by defining an explicit generative distribution over object slots, establishing formal conditions under which unsupervised object representations are guaranteed to be identifiable.

The problem with standard deterministic slots

In standard Slot Attention, $K$ slots $s_1, \dots, s_K \in \mathbb{R}^D$ are sampled from a single static Gaussian distribution $\mathcal{N}(\mu, \sigma^2 I)$ and updated over $T$ iterations using cross-attention:

\[A_{i, j} = \frac{\exp\left(\frac{1}{\sqrt{D}} k_i^T q_j\right)}{\sum_{l=1}^K \exp\left(\frac{1}{\sqrt{D}} k_i^T q_l\right)}, \quad v_j = \sum_{i=1}^N A_{i, j} v_i,\]

where $k_i$ and $v_i$ are key and value projections of the image feature tokens, and $q_j$ are query projections of slot $s_j$. A recurrent update (such as a GRU) then produces the next slot state: $s_j^{(t+1)} = \mathrm{GRU}(s_j^{(t)}, v_j)$.

Because the initial slots are sampled from an isotropic Gaussian with zero prior knowledge of object categories, the resulting slot assignments are prone to:

1. Permutation ambiguity: Slots arbitrarily swap positions across scenes.
2. Coordinate rotation: Latent dimensions can be freely rotated without altering reconstruction loss.
3. Slot collapse / under-clustering: Two slots may merge or split identical object instances inconsistently.

Hierarchical GMMs for identifiable binding

PSA replaces the isotropic Gaussian prior with a hierarchical generative process. An image $x$ is composed of $K$ distinct object entities generated by an underlying mixture of $M$ canonical object components:

\[p(s) = \sum_{m=1}^M \pi_m \, \mathcal{N}\left(s; \boldsymbol{\mu}_m, \boldsymbol{\Sigma}_m\right).\]

Each individual image generates a local GMM reflecting the specific objects present in that environment. Across the entire training dataset, these local GMMs aggregate into an Aggregate GMM, forming a structured latent space of recurring object properties.

Figure from the paper. Overview of Probabilistic Slot Attention (PSA). An encoder processes environment images $\mathbf{x}_1, \mathbf{x}_2, \mathbf{x}_3$ into local GMMs $\mathbf{z}_1, \mathbf{z}_2, \mathbf{z}_3$. These aggregate into a global GMM space from which slots $\mathbf{s}_1, \dots, \mathbf{s}_5$ are sampled to reconstruct image $\tilde{\mathbf{x}}$. Source: Avinash Kori et al.. Extracted from the source manuscript; click to enlarge.

Identifiability guarantees

The paper proves that under this generative formulation, the true latent object parameters $(\boldsymbol{\mu}_m, \boldsymbol{\Sigma}_m)$ are identifiable up to permutation and affine transformation. The proof builds on non-linear independent component analysis (ICA) with conditional mixture distributions:

\[p(\mathbf{z} \mid \mathbf{u}) = \prod_{j=1}^K p(z_j \mid \mathbf{u}),\]

where $\mathbf{u}$ represents auxiliary environmental conditioning (such as multi-scene contexts). By enforcing that the aggregate GMM components exhibit linearly independent characteristic functions, PSA eliminates the unidentifiable rotation subspace that plagues standard autoencoders.

Downstream advantages

1. Zero-Shot Transfer: Because slot properties represent identifiable physical factors (such as shape, material, and spatial coordinates), slots trained on simple synthetic environments transfer cleanly to novel combinations without retraining.
2. Consistent Visual Reasoning: When slots are fed into downstream relational or symbolic reasoning models, each slot index reliably refers to the same semantic concept, eliminating label-switching errors.
3. Principled Model Selection: The mixture log-likelihood provides an explicit criterion to determine the optimal number of slots $K$ and mixture components $M$ directly from data.

Paper and sources

• Avinash Kori, Francesco Locatello, Ainkaran Santhirasekaram, Francesca Toni, Ben Glocker and Fabio De Sousa Ribeiro. Identifiable object-centric representation learning via probabilistic slot attention. NeurIPS 2024 (Advances in Neural Information Processing Systems 37, pages 93300–93335).