Unlocking Slot Attention by Changing Optimal Transport Costs

TL;DR

This paper introduces MESH, a novel cross-attention module that enhances slot attention by integrating optimal transport techniques, enabling better handling of dynamic object counts in videos and improving performance on object-centric benchmarks.

Contribution

It establishes a connection between slot attention and optimal transport, and proposes MESH, a new method that combines unregularized and regularized optimal transport for improved object modeling.

Findings

MESH significantly outperforms standard slot attention on multiple benchmarks.

The method effectively handles videos with a dynamic number of objects.

MESH improves tie-breaking in object-centric modeling.

Abstract

Slot attention is a powerful method for object-centric modeling in images and videos. However, its set-equivariance limits its ability to handle videos with a dynamic number of objects because it cannot break ties. To overcome this limitation, we first establish a connection between slot attention and optimal transport. Based on this new perspective we propose MESH (Minimize Entropy of Sinkhorn): a cross-attention module that combines the tiebreaking properties of unregularized optimal transport with the speed of regularized optimal transport. We evaluate slot attention using MESH on multiple object-centric learning benchmarks and find significant improvements over slot attention in every setting.