Robust and Controllable Object-Centric Learning through Energy-based Models

TL;DR

This paper introduces extours, a simple, general energy-based model using Transformers for learning object-centric representations that are robust, controllable, and effective in segmentation and generalization tasks.

Contribution

It presents a novel energy-based approach for object-centric learning that is permutation-invariant, easily integrable, and improves robustness and compositional generalization.

Findings

Achieves high-quality object segmentation

Demonstrates robustness against distribution shifts

Enables systematic scene re-composition

Abstract

Humans are remarkably good at understanding and reasoning about complex visual scenes. The capability to decompose low-level observations into discrete objects allows us to build a grounded abstract representation and identify the compositional structure of the world. Accordingly, it is a crucial step for machine learning models to be capable of inferring objects and their properties from visual scenes without explicit supervision. However, existing works on object-centric representation learning either rely on tailor-made neural network modules or strong probabilistic assumptions in the underlying generative and inference processes. In this work, we present \ours, a conceptually simple and general approach to learning object-centric representations through an energy-based model. By forming a permutation-invariant energy function using vanilla attention blocks readily available in Transformers, we can infer object-centric latent variables via gradient-based MCMC methods where permutation equivariance is automatically guaranteed. We show that \ours can be easily integrated into existing architectures and can effectively extract high-quality object-centric representations, leading to better segmentation accuracy and competitive downstream task performance. Further, empirical evaluations show that \ours's learned representations are robust against distribution shift. Finally, we demonstrate the effectiveness of \ours in systematic compositional generalization, by re-composing learned energy functions for novel scene generation and manipulation.