slot_attn.py

# link to package https://github.com/lucidrains/slot-attention

import torch
from torch import nn

class Residual(nn.Module):
    def __init__(self, fn):
        super().__init__()
        self.fn = fn
    def forward(self, x):
        return x + self.fn(x)

class PreNorm(nn.Module):
    def __init__(self, dim, fn):
        super().__init__()
        self.fn = fn
        self.norm = nn.LayerNorm(dim)
    def forward(self, x):
        x = self.norm(x)
        return self.fn(x)

class SlotAttention(nn.Module):
    def __init__(self, num_slots, dim, iters = 3, eps = 1e-8, mlp_hidden_size = 128):
        super().__init__()
        self.num_slots = num_slots
        self.iters = iters
        self.eps = eps
        self.scale = dim ** -0.5

        self.slots_mu = nn.Parameter(torch.randn(1, 1, dim))
        self.slots_sigma = nn.Parameter(torch.randn(1, 1, dim))

        self.to_q = nn.Linear(dim, dim, bias = False)
        self.to_k = nn.Linear(dim, dim, bias = False)
        self.to_v = nn.Linear(dim, dim, bias = False)

        self.gru = nn.GRU(dim, dim)

        self.mlp = Residual(PreNorm(dim, nn.Sequential(
            nn.Linear(dim, mlp_hidden_size),
            nn.ReLU(inplace = True),
            nn.Linear(mlp_hidden_size, dim)
        )))

        self.norm_input = nn.LayerNorm(dim)
        self.norm_slots = nn.LayerNorm(dim)

    def forward(self, inputs):
        b, n, d, n_s = *inputs.shape, self.num_slots

        mu = self.slots_mu.expand(b, n_s, -1)
        sigma = self.slots_sigma.expand(b, n_s, -1)
        slots = torch.normal(mu, sigma)

        slots_shape = slots.shape

        inputs = self.norm_input(inputs)        
        k, v = self.to_k(inputs), self.to_v(inputs)

        for _ in range(self.iters):
            slots_prev = slots

            slots = self.norm_slots(slots)
            q = self.to_q(slots)

            dots = torch.einsum('bid,bjd->bij', q, k) * self.scale
            attn = dots.softmax(dim=1) + self.eps
            attn = attn / attn.sum(dim=-1, keepdim=True)

            updates = torch.einsum('bjd,bij->bid', v, attn)

            slots, _ = self.gru(
                updates.reshape(1, -1, d),
                slots_prev.reshape(1, -1, d)
            )

            slots = slots.reshape(b, -1, d)
            slots = self.mlp(slots)

        return slots

slot_attn = SlotAttention(num_slots=5, dim=512)
inputs = torch.randn(1, 1024, 512)
slot_attn(inputs)

Thank you for sharing your re-implementation! Two slight inconsistencies w.r.t. our paper that caught my eye: (1) we use ReLU (and different hidden layer sizes) in the feedforward network (mlp) instead of LeakyReLU, and (2) we initialize the slots at random for every individual example, but we learn the parameters of the initialization distribution. Initial values of the slots are taken as i.i.d. samples from a Gaussian distribution.

Another comment: In line 60, you normalize/divide the attention coefficients by their mean, but you should instead divide them by their sum (see Eq. 2 in the paper).

We also use a smaller eps of 1e-8 by default (see Table 11a).

all addressed! (i think)

Thanks!

Maybe I haben overlooked it, but how is this wonderful work licensed?