Slot Attention

Adaptive Slot Attention allows for a dynamic number of slots by generating a differentiable one-hot mask to decide whether to use a slot. This is implemented using MultiHeadSlotAttention wrapped in an AdaptiveSlotWrapper.

Workflow:

1. Define a MultiHeadSlotAttention module.
2. Wrap it with AdaptiveSlotWrapper, specifying a temperature (Gumbel-softmax temperature).
3. The wrapper returns both the slots and a keep_slots tensor.

Loss Minimization: To minimize the number of slots used for a scene (as suggested in the paper), you can add an auxiliary loss calculated as the sum of the keep_slots tensor to your main loss function.

import torch
from slot_attention import MultiHeadSlotAttention, AdaptiveSlotWrapper

# 1. Define slot attention
slot_attn = MultiHeadSlotAttention(
    dim = 512,
    num_slots = 5,
    iters = 3,
)

# 2. Wrap the slot attention
adaptive_slots = AdaptiveSlotWrapper(
    slot_attn,
    temperature = 0.5 # gumbel softmax temperature
)

inputs = torch.randn(2, 1024, 512)

# 3. Forward pass returns slots and the mask
slots, keep_slots = adaptive_slots(inputs) # (2, 5, 512), (2, 5)

# 4. Auxiliary loss to minimize number of slots used
keep_aux_loss = keep_slots.sum()  # add this to your main loss with some weight

The MetaSlotAttention class wraps an existing slot attention module to enable meta-learning via vector quantization. It quantizes the generated slots to a shared codebook and performs a deduplication process: it identifies duplicate slots (based on their codebook indices), merges them by averaging their values, and packs the unique slots to the left of the sequence while masking the duplicates.

Initialization

MetaSlotAttention requires an existing slot_attention module (which must have a .dim attribute) and a codebook_size.

Returns

The forward method returns a tuple of (dedup_slots, mask, vq_loss):

• dedup_slots: The slots after quantization, merging, and deduplication. Duplicate slots are zeroed out.
• mask: A boolean mask indicating which slots in dedup_slots are unique (True) and which are duplicates/padded (False).
• vq_loss: The vector quantization loss used for training the codebook.

from slot_attention.meta_slot_attention import MetaSlotAttention

# Assuming 'my_slot_attention' is a pre-defined Module with a .dim attribute
meta_slot_attn = MetaSlotAttention(
    slot_attention = my_slot_attention,
    codebook_size = 1024,
    vq_decay = 0.9,
    kmeans_init = True
)

# inputs: [batch, sequence, dim]
# slot_attn_kwargs: arguments passed to the underlying slot_attention module
dedup_slots, mask, vq_loss = meta_slot_attn(inputs, **slot_attn_kwargs)

SlotAttentionExperimental is a module that implements an experimental variant of slot attention using gated residual connections and iterative updates. It learns a set of latent 'slots' by interacting with input features through weighted attention mechanisms.

Key Parameters:

• num_slots: The number of latent slots to learn.
• dim: The dimensionality of the input and slot features.
• iters: The number of iterative update steps (default is 3).
• eps: Small epsilon value for numerical stability (default is 1e-8).
• hidden_dim: The dimensionality of the hidden layer in the internal FeedForward networks (default is 128).

Behavior:

1. Slot Initialization: Slots are initialized using a reparameterization trick from learned slots_mu and slots_logsigma parameters.
2. Iterative Updates: For a specified number of iters, the module performs:
   • A gated residual update of slots using WeightedAttention against the inputs.
   • A gated residual update of slots using a FeedForward network.
   • A gated residual update of inputs using WeightedAttention against the slots.
   • A gated residual update of inputs using a FeedForward network.
3. Returns: The module returns a tuple containing the updated slots and the updated inputs.

Note: You can optionally pass a num_slots argument to the forward method to override the default self.num_slots defined during initialization.

from slot_attention.slot_attention_experimental import SlotAttentionExperimental
import torch

# Configuration
num_slots = 8
dim = 256

# Initialize model
model = SlotAttentionExperimental(num_slots=num_slots, dim=dim)

# Dummy input: (batch, num_features, dim)
b, n, d = 4, 64, 256
inputs = torch.randn(b, n, d)

# Forward pass
slots, updated_inputs = model(inputs)

print(f"Slots shape: {slots.shape}")           # Expected: (4, 8, 256)
print(f"Inputs shape: {updated_inputs.shape}") # Expected: (4, 64, 256)

The SlotAttention class implements the slot attention mechanism, which uses an iterative process to compete for features in an input sequence. It uses a learned Gaussian prior for the slots and updates them via a GRU cell based on attention weights.

Parameters

• num_slots (int): The number of slots to be used.
• dim (int): The dimensionality of the input and slot features.
• iters (int, optional): The number of iterations for the competitive attention process. Defaults to 3.
• eps (float, optional): A small epsilon value for numerical stability in L1 normalization. Defaults to 1e-8.
• hidden_dim (int, optional): The dimensionality of the hidden layer in the MLP applied after the GRU update. Defaults to 128.

Forward Pass

• inputs (torch.Tensor): The input tensor of shape (batch, sequence_length, dim).
• num_slots (int, optional): If provided, overrides the num_slots defined during initialization. Defaults to self.num_slots.

Returns a tensor of shape (batch, num_slots, dim) representing the updated slots.

import torch
from slot_attention import SlotAttention

b, n, d = 2, 16, 64
inputs = torch.randn(b, n, d)

# Initialize SlotAttention
model = SlotAttention(num_slots=8, dim=d, iters=3)

# Forward pass
slots = model(inputs)

print(slots.shape)  # torch.Size([2, 8, 64])

The AdaptiveSlotWrapper is a PyTorch module that wraps an existing SlotAttention or MultiHeadSlotAttention instance. It adds a mechanism to predict whether each slot should be 'kept' or 'discarded' using a differentiable one-hot mask via the Gumbel-Softmax trick (straight-through estimator).

This is useful for implementing auxiliary losses that penalize the use of unnecessary slots, allowing the model to dynamically decide the number of active slots for a given scene.

Key features:

• Wraps SlotAttention or MultiHeadSlotAttention.
• Uses a linear layer (pred_keep_slot) to predict slot usage logits.
• Employs Gumbel-Softmax to provide a differentiable 'hard' mask for training while maintaining gradients via the straight-through estimator.
• Returns both the generated slots and a keep_slots mask (a float tensor of shape [batch, num_slots] containing values in {0., 1.}).

from slot_attention.adaptive_slot_wrapper import AdaptiveSlotWrapper
from slot_attention.slot_attention import SlotAttention
import torch

# Initialize base slot attention
slot_attn = SlotAttention(
dim = 256, num_slots = 10)

# Wrap it to enable adaptive slot selection
wrapper = AdaptiveSlotWrapper(
    slot_attn = slot_attn,
    temperature = 1.0
)

# Forward pass
x = torch.randn(8, 16, 256)  # [batch, sequence, dim]
slots, keep_mask = wrapper(x)

# slots: [batch, num_slots, dim]
# keep_mask: [batch, num_slots] (values are 0.0 or 1.0)

The forward method performs the iterative slot attention process.

Arguments

• inputs (torch.Tensor): The input features of shape (batch, n, dim).
• num_slots (int, optional): Overrides the num_slots defined during initialization. If None, the default self.num_slots is used.

Returns

• slots (torch.Tensor): The updated slots of shape (batch, num_slots, dim).

Implementation Details

1. Slot Initialization: Slots are sampled from a learned Gaussian distribution defined by slots_mu and slots_logsigma.
2. Iterative Updates: For the specified number of iters, the model:
   • Computes attention between slots (queries) and inputs (keys/values).
   • Applies L1 normalization to the attention weights.
   • Updates slots using a GRUCell based on the attended values.
   • Applies a residual MLP block to the slots.

import torch
from slot_attention.multi_head_slot_attention import MultiHeadSlotAttention

b, n, d = 2, 16, 256
num_slots = 8

model = MultiHeadSlotAttention(num_slots=num_slots, dim=d)
inputs = torch.randn(b, n, d)

# Perform slot attention
slots = model(inputs)

# slots shape: (2, 8, 256)
print(slots.shape)

The MultiHeadSlotAttention class implements a multi-head version of the slot attention mechanism. It uses a learned Gaussian distribution (mean slots_mu and log-variance slots_logsigma) to initialize slots, which are then iteratively updated via attention and a GRU cell.

Parameters

• num_slots (int): The number of slots to be used.
• dim (int): The input and output feature dimension.
• heads (int, default=4): The number of attention heads.
• dim_head (int, default=64): The dimension of each individual head.
• iters (int, default=3): The number of iterative update steps.
• eps (float, default=1e-8): Small epsilon value for L1 normalization.
• hidden_dim (int, default=128): The hidden dimension used in the internal MLP.

from slot_attention.multi_head_slot_attention import MultiHeadSlotAttention
import torch

# Example initialization
model = MultiHeadSlotAttention(
    num_slots=8,
    dim=256,
    heads=4,
    dim_head=64,
    iters=3
)

The OrthoSlotWrapper is a PyTorch Module wrapper designed to add an orthogonality loss to an existing slot attention module. It encourages the learned slot representations to be orthogonal to each other by penalizing the squared pairwise inner products of centered slot representations (excluding the diagonal).

When you wrap a slot attention module with OrthoSlotWrapper, the forward pass returns both the original output of the slot attention module and a new ortho_loss scalar. If the original module returns a tuple, the ortho_loss is appended to that tuple.

import torch
from slot_attention import SlotAttention # Assuming SlotAttention is the base module
from slot_attention.ortho_slot_wrapper import OrthoSlotWrapper

# 1. Initialize your base slot attention module
slot_attention = SlotAttention(
    dim = 256,
    num_slots = 10,
    # ... other parameters
)

# 2. Wrap it with OrthoSlotWrapper
wrapped_slot_attention = OrthoSlotWrapper(slot_attention)

# 3. Forward pass
# slots will be the original output, ortho_loss is the penalty for non-orthogonality
slots, ortho_loss = wrapped_slot_attention(inputs)

# 4. Use ortho_loss in your total loss calculation
total_loss = reconstruction_loss + ortho_loss
total_loss.backward()