lucidrains/slot-attention: A PyTorch Block for Object-Centric Clustering

The gap this fills between the paper and your encoder output

Slot Attention as described by Locatello et al. is an iterative attention mechanism that binds a set of learned slots to the parts of an input. The paper is object-centric learning: given a feature map, produce a small number of vectors where each vector tends to correspond to one object. Reimplementing the normalization and the iteration order correctly is fiddly, and the official Google Research code lives inside a larger research tree rather than a package you can pin. This repository packages the module as a Python distribution named slot_attention, so the unit of reuse is a torch.nn.Module you drop after your encoder. The intended user is someone who already has a CNN or transformer producing a sequence of feature vectors and wants the slot iteration without copying paper code. It is not aimed at someone who wants a finished object discovery system. The README shows no dataset loader, no training script, and no checkpoint.

What the forward pass actually does with your tensor

The mechanism is visible in the signature. SlotAttention takes num_slots, dim and iters at construction, and the README's example passes a tensor of shape (2, 1024, 512): batch two, 1024 input positions, 512 feature channels. The output is (2, 5, 512), so the 1024 positions are compressed into five slot vectors of the same width. That is the whole contract. The module does not care whether the 1024 positions came from a convolutional feature map flattened or from a sequence of patches. iters defaults to 3, which is the number of attention iterations run before the slots are returned. The README states that after training the network is reported to generalize to a slightly different number of slots, and that you can override it per call with num_slots in forward, which returns (2, 8, 512) in the given example. That per-call override is the interesting detail: the slot count is not baked into the weights, so the same trained module can be asked for more or fewer clusters at inference. The README does not explain how the override interacts with the internal parameters, and I cannot confirm from the supplied material whether the extra slots are initialized the same way as the original ones.

Adaptive slots and the auxiliary loss you have to wire yourself

The second half of the README covers MultiHeadSlotAttention combined with AdaptiveSlotWrapper, citing the Adaptive Slot Attention paper on dynamic slot number. The wrapper takes a temperature argument, documented as the gumbel softmax temperature, shown as 0.5. Calling the wrapped module returns two tensors: the slots, shape (2, 5, 512), and keep_slots, shape (2, 5). The second is the differentiable one-hot mask indicating whether each slot is used. The README is explicit that the auxiliary loss for minimizing the number of slots used in a scene is keep_slots.sum(), to be added to the main loss with some weight. That weight is not given. The README does not suggest a value, a schedule, or a range. If you use AdaptiveSlotWrapper, tuning that coefficient is on you, and the documentation gives you no starting point. Note also that the adaptive path uses MultiHeadSlotAttention rather than the plain SlotAttention class from the first example; the README does not say whether AdaptiveSlotWrapper works with the single-head module.

Install and the smallest working call

Installation is a single command: pip install slot_attention. The usage example is four lines of setup plus one call. You construct SlotAttention with num_slots, dim and iters, then feed a tensor of shape (batch, positions, dim). There is no config file, no environment variable, and no CLI, so there is nothing else to configure. The package name on PyPI uses an underscore while the distribution is referred to as slot_attention throughout, which is worth keeping straight if you pin it in a requirements file. The README does not state a minimum PyTorch version, a Python version, or whether CUDA is required. It also does not mention mixed precision, torch.compile, or any device placement guidance. For a module this small that is probably fine, but if you are integrating into an existing training stack you will be discovering those constraints by reading the source rather than the documentation.

Where this package stops and your work begins

The clearest limitation is scope. There is no decoder in the README, and Slot Attention in the original work is trained with a reconstruction objective through a decoder that maps slots back to the input space. Without that, the slots have no pressure to correspond to objects; the module will happily return five vectors that are a lossy summary of your features and nothing more. The README also gives no training loop, no optimizer settings, no learning rate, and no dataset. Anyone expecting to pip install and get object masks will be disappointed. A second limitation is the reported slot-count generalization. The README says the network is reported to be able to generalize to slightly different numbers of slots, and the word slightly is doing real work there. There is no number attached, no evaluation, and no statement about what happens at double the training slot count. Treat the per-call num_slots override as an experiment to run, not a guarantee. Third, the README's own update note points to the official Google Research release, which means this package is a convenience reimplementation with a citation list rather than the reference implementation. If you are reproducing paper numbers, that distinction matters.

How it compares to using the official research code

The alternative is the official repository linked in the README under google-research. The difference in approach is packaging versus completeness. The official tree sits inside a larger research codebase with the training pipeline, configs and evaluation that the paper used, which is what you want if you are reproducing results or need the decoder and data setup. This package is the opposite trade: a small, pip-installable module with a stable import path, at the cost of everything around the module. Choosing between them is not about which is better in the abstract. If your job is to add object-centric slots to an existing PyTorch model, the official tree drags in a research environment you do not need. If your job is to reproduce the paper, this package gives you the attention block and leaves you to build the rest, and you will likely end up reading the official code anyway to get the objective right. A third option, writing the module yourself from the paper, is more work than it sounds because the iteration and normalization details are where implementations diverge.

Version drift, maintenance signals and the MIT licence

The listed releases are 1.2.1a, 1.2.2 and 1.4.0, the last dated 2024-08-20. The repository's last push is dated 2026-06-06, so there is activity after the most recent release, but the material does not say what that activity contains or whether a newer version has been published. That gap is the thing to check before you pin: confirm on PyPI that the version you install contains MultiHeadSlotAttention and AdaptiveSlotWrapper, since those appear in the README and it is not stated which release introduced them. The citation list runs through 2026, including MetaSlot and OrthoRF, which suggests the README tracks follow-up work even if the code has not been re-released alongside it. On licensing, the repository carries MIT. That is permissive and places few obligations on how you use the module in your own codebase, but it says nothing about the licences of the papers or the official implementation you might combine it with, and it is not legal advice. If you are shipping a product that depends on this, the practical maintenance cost is low because the surface area is one module, but the cost of the surrounding work (decoder, loss weighting, training data) is entirely yours and is not something the package reduces.

Editorial conclusion

Adopt this if you already have a PyTorch encoder and want the Slot Attention iteration as a ready-made module rather than transcribing the paper yourself; the install is one pip command and the forward signature is small. Do not adopt it if you need a complete object-discovery pipeline, pretrained weights, or a training loop, because none of those are in the repository material. Before committing, verify the released version on PyPI matches the API in the README, since the README shows AdaptiveSlotWrapper and MultiHeadSlotAttention while the listed releases stop at 1.4.0 from August 2024 and the last push to master is dated 2026-06-06.

Official sources

1. Issues
2. License: MIT
3. lucidrains/slot-attention on GitHub
4. README
5. Releases