A PyTorch implementation of slot attention

Check out our implementation at https://github.com/untitled-ai/slot_attention.

Also, we’re hiring! You can see our open roles, and email us at jobs@imbue.com if you find a role for you.

Outputs of our slot attention model, demonstrating the model's ability to divide objects, or parts of objects, into slots.

Why we chose this paper

Here’s a question for you, intelligent reader: how many distinct objects do you see in the upper-left image?

This is not a hard question for us humans. However, unsupervised machine learning models have a surprisingly difficult time answering this: they struggle to understand which pixels constitute an object by human standards. And while there’s been success in training supervised models to complete this task ¹, it’s largely unsolved on the unsupervised side ².

You can see the issue here. It’s important that machines actually understand what objects are, because countless elementary tasks that humans can do require this basic capability.

Enter Object-Centric Learning with Slot Attention. This nifty technique introduces a novel way to extract objects from the original input without any training or supervision. After an image is processed (for example, through a CNN) into feature representations, a Slot Attention module forces “slots” to compete over these features, using an iterative attention mechanism similar to the attention mechanism popular in other domains of machine learning ³. At each iteration, slots compete to explain parts of input features, and then update their representation.

Figure 3 in the paper visualizes the attention masks at each iteration, illustrating how each slot's representation updates after each iteration.

It’s exciting that unsupervised models can now determine objects from complex scenes, and can even do so without using depth, motion, or other helpful signals.

What’s next?

While slot attention is a step in the right direction toward unsupervised object understanding, there are plenty of interesting open questions:

• How can this work on real data with many objects and noisy backgrounds?
• Our team works heavily with video and objects in motion. How well would it work if we determined objects in an image based on what they’re doing - for example, whether they’re moving, or static?
• Instead of reconstruction loss, what would happen if we use a contrastive loss instead? (See this related work by Sindy Löwe)

By reimplementing this paper and sharing our code, we hope we can allow others to help answer these questions, and more.

Note: We did not fully reproduce the result of the paper due to time and GPU constraints, so drop a ticket on our Github page if you have any issues. If you have thoughts on the model, where you think it’d be useful, etc., please shoot us a note at josh@imbue.com. We’re happy to chat!

Sources

1. Object-Centric Learning with Slot Attention, https://arxiv.org/abs/2006.15055.

2. Slot Attention in Tensorflow, https://github.com/google-research/google-research/tree/master/slot_attention.