ICLR-2022 Perceiver IO：A General Architecture for Structured Inputs & Outputs

2022-01-29 · ICLR· Vision · vision-language

Paper: Perceiver IO: A General Architecture for Structured Inputs & Outputs

Code: https://github.com/lucidrains/perceiver-pytorch

Perceiver IO: A Unified Multimodal Architecture

Abstract

Perceiver IO is proposed as a general model that can handle data in diverse modalities and scales linearly with input and output size. A flexible query mechanism supports outputs of various sizes and semantics, removing the need for task-specific architectural engineering. The same architecture achieves strong results on natural language processing, visual understanding, multi-task learning, and multimodal understanding.

Introduction

Most machine learning research focuses on building bespoke systems for rigid input–output patterns tied to a single task. As inputs or outputs grow more diverse, the complexity of such systems can increase sharply, and the structure of a task’s inputs and outputs can strongly constrain how data are processed, making adaptation to new settings difficult. Must we inevitably develop problem-specific models for every new combination of inputs and outputs?

Perceiver uses attention to map inputs from diverse modalities into a fixed-size latent space, which is further processed by a deep, fully attention-based network. This decouples most of the network’s computation from input size and modality-specific details, allowing it to scale to large and multimodal data.

However, Perceiver can only handle simple output spaces such as classification. In this work, the authors develop a mechanism to decode structured outputs—language, optical flow fields, audiovisual sequences, unordered sets of symbols, and more—directly from the Perceiver latent space, enabling the model to tackle many new domains without giving up the benefits of deep, domain-agnostic processing.

To do so, each output is generated by attending over the latent array with an output query whose semantics specify that particular output.

For example, to predict optical flow at a specific pixel, a query can be formed from the pixel’s $(x,y)$ coordinates plus a task embedding for flow.
The model then uses the query in attention and produces a single flow vector. The architecture can produce many outputs, each with arbitrary shape and structure, while latent features in the model remain agnostic to output shape and structure.

Perceiver IO achieves this with a fully attention-based read–process–write architecture: inputs are encoded (read) into a latent space, latent representations are refined through multiple processing layers (process), and the latent space is decoded (write) to produce outputs. By combining the strengths of Transformers—domain-agnostic nonlocal processing of inputs and an encoder–decoder design—this approach decouples element-wise input size from the size of the input space and reduces complexity.

The Perceiver IO Architecture

Encoding, Processing, And Decoding

Each module applies a global query–key–value (QKV) attention operation followed by a multilayer perceptron (MLP). In Transformer architectures, the MLP is typically applied independently to each element along the index dimension.

Both the encoder and decoder take two input matrices: the first feeds the module’s key and value networks, and the second feeds the query network. The module output has the same index dimension as the query input (i.e., the same number of elements), which is why encoder and decoder modules can produce outputs of different sizes.

Why not use a Transformer directly?

Transformers scale poorly in both compute and memory. They deploy attention throughout the full stack, generating queries and keys from the entire input at every layer, so each layer incurs quadratic time and memory in sequence length. For long inputs such as images, training without preprocessing is often infeasible.

By contrast, Perceiver IO uses attention non-uniformly: first to map inputs into a latent space, then to process within that space, and finally to map back to the output space.

The architecture therefore avoids quadratic dependence on input or output size: encoder and decoder attention scale linearly with input and output size respectively, while latent attention is independent of both.

It also demands less compute and memory, so Perceiver IO can scale to larger inputs and outputs. Whereas Transformers are typically used when inputs and outputs have at most a few thousand dimensions, this model shows strong results on data with hundreds of thousands of input and output dimensions.

Decoding The Latent Representation with a Query Array

The second step decodes representation vectors in the latent space. Given a latent representation of size $N \times D$, the goal is to produce an output matrix of size $O \times E$. Query information should therefore reflect the downstream task and capture any structure required in the outputs, which may include spatial locations in images or token positions in sequences.

Queries are constructed by combining (concatenating or adding) a set of vectors into a query vector that contains all information relevant to one of the $O$ desired outputs.

For tasks with simple outputs such as classification, these queries can be reused across examples and learned from scratch. For outputs with spatial or sequential structure, learned positional encodings or Fourier features are included to represent the position to be decoded. For multi-task or multimodal outputs, a single learned query per task or modality lets the network distinguish one task or modality from another, analogous to how positional encodings let attention distinguish one position from another.

To assess the generality of Perceiver IO, the authors evaluate it on tasks spanning multiple domains and data types, including language understanding (masked language modeling and downstream fine-tuning), visual understanding (optical flow and image classification), symbolic game representations (StarCraft II), and multimodal and multi-task settings.

Experiments

See the original paper for additional experiments.

Yang Tan (谭扬)