lucidrains/perceiver-pytorch

Implementation of Perceiver, General Perception with Iterative Attention, in Pytorch

perceiver‑pytorch – PyTorch implementation of the Perceiver family

What it is – A pure‑Python library that re‑creates the Perceiver and Perceiver IO architectures (and a small experimental variant) from the papers “Perceiver: General Perception with Iterative Attention” and “Perceiver IO: A General Architecture for Structured Inputs & Outputs”. These models are a type of transformer‑style neural network designed to handle very large, high‑dimensional inputs (images, video, audio, point clouds, etc.) by first projecting them onto a much smaller set of learned latent vectors and then iteratively applying cross‑attention and self‑attention.

Why it matters – The original Perceiver papers showed that a single architecture could process many modalities without needing modality‑specific designs. This repository makes the idea easy to experiment with in PyTorch, letting researchers and developers plug the model into image classification, language modelling, or any custom task.

Quick start

pip install perceiver-pytorch

Image classification example (vanilla Perceiver)

import torch
from perceiver_pytorch import Perceiver

model = Perceiver(
    input_channels=3,          # RGB image
    input_axis=2,              # 2‑D data (height × width)
    num_freq_bands=6,
    max_freq=10.,
    depth=6,
    num_latents=256,
    latent_dim=512,
    cross_heads=1,
    latent_heads=8,
    cross_dim_head=64,
    latent_dim_head=64,
    num_classes=1000,
    attn_dropout=0.,
    ff_dropout=0.,
    weight_tie_layers=False,
    fourier_encode_data=True,
    self_per_cross_attn=2,
)

img = torch.randn(1, 224, 224, 3)   # batch‑size‑1 ImageNet‑size image
logits = model(img)                # → shape (1, 1000)

Flexible‑output Perceiver IO

from perceiver_pytorch import PerceiverIO

model = PerceiverIO(
    dim=32,
    queries_dim=32,
    logits_dim=100,
    depth=6,
    num_latents=256,
    latent_dim=512,
    cross_heads=1,
    latent_heads=8,
    cross_dim_head=64,
    latent_dim_head=64,
    weight_tie_layers=False,
    seq_dropout_prob=0.2,
)

seq = torch.randn(1, 512, 32)          # arbitrary input sequence
queries = torch.randn(128, 32)        # decoder queries
out = model(seq, queries=queries)     # → (1, 128, 100)

Language‑modeling variant (PerceiverLM)

from perceiver_pytorch import PerceiverLM

model = PerceiverLM(
    num_tokens=20000,
    dim=32,
    depth=6,
    max_seq_len=2048,
    num_latents=256,
    latent_dim=512,
    cross_heads=1,
    latent_heads=8,
    cross_dim_head=64,
    latent_dim_head=64,
    weight_tie_layers=False,
)

seq = torch.randint(0, 20000, (1, 512))
mask = torch.ones(1, 512).bool()
logits = model(seq, mask=mask)        # → (1, 512, 20000)

Key features (as described in the README)

Single‑line install via pip.
Modular API: Perceiver, PerceiverIO, and PerceiverLM classes cover classification, flexible‑output tasks, and language modelling respectively.
Fourier positional encoding built‑in (toggle with fourier_encode_data).
Configurable depth, latent size, and attention heads to match the original papers or to experiment with smaller models.
Experimental bottom‑up attention variant available under perceiver_pytorch.experimental.Perceiver (adds an induced‑set attention block similar to Set Transformers).
Citation ready: the repository includes ready‑to‑copy BibTeX entries for the Perceiver, Perceiver IO, and related works.

When to use it

You need a modality‑agnostic backbone that can ingest very large tensors (e.g., high‑resolution images, video frames, point clouds) without exploding memory.
You want to prototype cross‑modal or multimodal research where the same encoder can be reused for different data types.
You are exploring flexible output shapes (e.g., segmentation maps, language generation) – Perceiver IO’s query‑based decoder makes this straightforward.

Limitations / notes

The library provides the model definition only; training loops, data pipelines, and performance optimisations are left to the user.
No pre‑trained weights are shipped; you must train from scratch or load your own checkpoints.
The experimental bottom‑up variant is optional and requires importing from the experimental submodule.

References (from the README)

Perceiver: General Perception with Iterative Attention – arXiv:2103.03206
Perceiver IO: A General Architecture for Structured Inputs & Outputs – arXiv:2107.14795
Additional citations for related attention mechanisms are also listed.

Bottom line – perceiver-pytorch is a faithful, lightweight re‑implementation of the Perceiver family, suitable for anyone wanting to experiment with the architecture in PyTorch without digging into the original research codebases.