$ show lucidrains/slot-attention
This repository provides a PyTorch implementation of Slot Attention and its variants for object-centric learning. The core SlotAttention module refines a set of slot representations via iterative attention from slots to inputs, enabling the model to discover and represent objects in a scene. The repository extends this base with several wrappers: MultiHeadSlotAttention for multi-head attention, AdaptiveSlotWrapper for dynamic slot selection via Gumbel softmax, MetaSlotAttention for deduplication using vector quantization, and OrthoSlotWrapper for adding an orthogonality loss to encourage distinct slot representations. The library is designed to be modular and reusable, with each component having minimal dependencies and clear interfaces. The code leverages einops for readable tensor operations and follows modern PyTorch practices. A training script is included to demonstrate the usage on a toy dataset of colored rectangles, showing how the slot attention can be used to discover and deduplicate objects. The project is actively maintained, with recent commits and a healthy number of stars and forks. Dependencies are well-specified and up-to-date, indicating low technical drift. The codebase exhibits strong craftsmanship in its clean abstractions and reusable components, making it a valuable resource for researchers and practitioners working on object-centric representations.
The LLM identified 4 files worth highlighting for reuse:
Implements the core SlotAttention module as described in the paper 'Object-Centric Learning with Slot Attention'. It takes input features and iteratively refines slot representations via attention from slots to inputs, followed by GRU and MLP updates. The component is self-contained, depends only on torch and einops, and is directly reusable in any PyTorch project for object-centric learning tasks.
Extends the base SlotAttention to use multi-head attention, allowing the model to attend to information from different representation subspaces. It maintains the same interface as the base SlotAttention but splits key, value, and query projections into multiple heads. This component demonstrates skill in adapting attention mechanisms and is reusable with minimal modification.
Wraps any slot attention module to predict a differentiable mask for slot usage via a Gumbel softmax over a linear projection of slot representations. Enables dynamic selection of the number of slots used for a given input, as described in the Adaptive Slot Attention paper. The wrapper is highly reusable and adds minimal overhead.
Wraps a slot attention module and applies vector quantization to slot representations to merge duplicate slots. Uses a codebook to quantize slots, merges slots assigned to the same codebook entry by averaging, and returns a mask indicating non-duplicate slots. Demonstrates non-obvious skill in combining vector quantization with attention for deduplication and is reusable with the vector-quantize-pytorch dependency.