Perceiver IO: A General Architecture for Structured Inputs & Outputs

1 INTRODUCTION . Humans have a remarkable ability to take in data from many sources , integrate it seamlessly , and deploy it in the service of a range of goals . Most machine learning research focuses on building bespoke systems to handle the stereotyped inputs and outputs associated with a single task . This is true even for models that handle multiple modalities . A typical approach independently processes each input with a modality specific architecture ( for example using a 2D ResNet ( He et al. , 2016 ) for vision and a Transformer ( Vaswani et al. , 2017 ) for language ) , integrates them afterwards using a third fusion network , and reads out the result in a task-specific manner . The complexity of systems like this can grow dramatically as the inputs or outputs grow more diverse ( e.g . Abramson et al . ( 2020 ) ; Vinyals et al . ( 2019 ) ; Ramesh et al . ( 2021 ) ) , and the structure of a task ’ s inputs and outputs may place strong constraints on how data is processed , making adaptation to new settings difficult . Is the development of problem-specific models for each new set of inputs and outputs unavoidable ? Life would be drastically simpler if a single neural network architecture could handle a wide variety of both input modalities and output tasks . In this work , we propose such an architecture , with the ultimate goal of building a network that can easily integrate and transform arbitrary information for arbitrary tasks . Our starting point is the Perceiver ( Jaegle et al. , 2021 ) , an architecture which has demonstrated a remarkable ability to handle data from many modalities with no changes to the network architecture . The Perceiver uses attention to map inputs of a wide range of modalities to a fixed-size latent space that is further processed by a deep , fully attentional network . This process decouples the bulk of the network ’ s processing from the size and modality-specific details of the input , allowing it to scale to large and multimodal data . But the Perceiver can only handle simple output spaces like classification . Much of the complexity of real-world tasks comes from the variety , size , and structure of their outputs , and in this regard the original Perceiver can ’ t be considered general purpose . In this work , we develop a mechanism for decoding structured outputs – language , optical flow fields , audiovisual sequences , symbolic unordered sets , etc . – directly from the Perceiver latent space , which allows the model to handle a host of new domains without sacrificing the benefits of deep , domain-agnostic processing . To do this , we produce each output by attending to the latent array using an output query that specifies the semantics of that particular output . For example if we wanted the model to predict optical flow on i one particular pixel we could compose a query from the pixel ’ s xy coordinates plus an optical flow task embedding : the model would then attend using the query and produce a single flow vector . As a result , our architecture can produce many outputs , each with arbitrary shape and structure , and yet the latent features in our architecture remain agnostic to the shape and structure of the outputs . Perceiver IO does this using a fully attentional read-process-write architecture : inputs are encoded ( read ) to a latent space , the latent representation is refined ( process ) via many layers of processing , and the latent space is decoded ( write ) to produce outputs . This approach inherits the best features of both Transformers – which leverage domain agnostic primitives for nonlocal processing of inputs – and the encoder-decoder architectures ( e.g . Ronneberger et al . ( 2015 ) ; Newell et al . ( 2016 ) ) that are in widespread use in high-bandwidth domains such as computer vision or multimodal processing . This approach allows us to decouple the size of elements used for the bulk of the computation ( the latent ) from the size of the input and output spaces , while making minimal assumptions about the spatial or locality structure of the input and output . Perceiver IO ’ s decoding procedure uses an attention mechanism to map from latents to arbitrarily sized and structured outputs using a querying system that can flexibly specify the semantics needed for outputs on a wide range of domains , including dense and multitask settings . This decoder allows Perceiver IO to serve as a drop-in replacement for a wide range of specialist networks currently in use on a set of challenging domains , while improving performance on tasks like classification that could be handled by the Perceiver . The proposed architecture can be applied with unprecedented levels of generality . Perceiver IO can replace the Transformers used in BERT ( Devlin et al. , 2019 ) and AlphaStar ( Vinyals et al. , 2019 ) . At the same time , Perceiver IO produces state-of-the-art results on the Sintel optical flow benchmark ( Butler et al. , 2012 ) and good results on ImageNet image classification ( Deng et al. , 2009 ) . Perceiver IO produces compelling results even when handling highly diverse multimodal data , such as on joint { video , audio , label } auto-encoding in Kinetics ( Smaira et al. , 2020 ) and joint audio-video classification on AudioSet ( Gemmeke et al. , 2017 ) . Perceiver IO allows us to simplify pipelines and remove domain-specific assumptions : we process language without using tokenizers , fine-tune on multiple classification tasks simultaneously and without the need for [ CLS ] tokens ( Sec . 4.1 ) , estimate optical flow without relying on explicit architectural features for multiscale correspondence ( Sec . 4.2 ) , learn joint representations of video , audio , and labels without separate network trunks ( Sec . 4.3 ) , and perform image classification with no information about the 2D structure of images ( Sec . 4.4 ) . 2 RELATED WORK . Neural network research has long sought architectures that can handle large , arbitrarily structured inputs and outputs . Autoencoding ( Hinton & Zemel , 1994 ) was among the first attempts to build representation which could encode and reproduce high-dimensional inputs like images . As hardware grew more powerful , neural nets led to breakthroughs in image understanding ( Krizhevsky et al. , 2012 ; Zeiler & Fergus , 2014 ; Szegedy et al. , 2015 ) and interest intensified : autoregressive models that could process and complete samples of handwriting were developed ( Graves , 2013 ) , and new convolutional network designs led to good results in structured output spaces like semantic segmentation ( Farabet et al. , 2012 ; Long et al. , 2015 ; Ronneberger et al. , 2015 ) , pose estimation ( Toshev & Szegedy , 2014 ) , ii detection ( Sermanet et al. , 2014 ) , captioning ( You et al. , 2016 ) , and optical flow ( Fischer et al. , 2015 ) . At the same time , natural language applications research has made extensive progressive in capturing the structured nature of language , typically via autoregressive models ( Brown et al. , 2020 ; Collobert et al. , 2011 ; Radford et al. , 2019 ; Sutskever et al. , 2014 ; Vaswani et al. , 2017 ) or context prediction ( Devlin et al. , 2019 ; Pennington et al. , 2014 ; Mikolov et al. , 2013 ) . Similar to our work , several groups have proposed to solve tasks in multiple domains ( e.g . ( Akbari et al. , 2021 ; Alayrac et al. , 2020 ; Kaiser et al. , 2017 ) ) , but typically across a fixed and predefined set of modalities by means of domain-specific networks . Although single-task specialist networks remain dominant in vision , multi-task learning has become popular ( Doersch & Zisserman , 2017 ; Kokkinos , 2017 ; Misra et al. , 2016 ; Zamir et al. , 2018 ) and individual models achieve generality in a restricted domain : e.g . Mask-RCNN ( He et al. , 2017 ) handles object detection , segmentation , and pose estimation . In language , training or evaluation on multiple tasks has also become common ( Collobert & Weston , 2008 ; Devlin et al. , 2019 ; Luong et al. , 2016 ; Liu et al. , 2019 ; Raffel et al. , 2020a ) . Several groups have demonstrated that Transformers ( originally designed for language ) can be used or adapted to non-language tasks ( e.g . Chen et al . ( 2020 ) ; Lu et al . ( 2021 ) ) , but the limited scalability of Transformers limits their usefulness as general-purpose architectures . Several groups have proposed to use attention to manipulate the size of arrays or to introduce bottlenecks in processing . Set Transformers and related work ( Lee et al. , 2019 ; Goyal et al. , 2021 ) use a learned query ( “ inducing points ” ) to induce local bottlenecks by mapping a set back and forth from a set with fewer elements and learned decoder queries ( “ seed vectors ” ) to map to outputs ( “ pooling by multiheaded attention ” ) . Each layer of these networks has complexity linear in the input size , while Perceivers use a deep latent network with complexity independent of the input and output . Our work uses attention over inputs and outputs of different sizes in part to produce an efficient attention architecture , and several other efficient attention architectures have been proposed , largely for language or small-scale problems ( e.g . Xiong et al . ( 2021 ) ; Wang et al . ( 2020 ) ; Tay et al . ( 2021a ) ; Beltagy et al . ( 2020 ) and see ( Tay et al. , 2021b ) ) . The focus of our work is developing an architecture that is efficient and also performs well in many settings with a wide range of inputs and outputs . Several works use attention to process latent spaces that interface with input/output data using task- or domain-specific architectures ( Carion et al. , 2020 ; Locatello et al. , 2020 ; Wang et al. , 2021 ) , and cross-attention itself is widely used to produce outputs in of a different size or structure from inputs ( Dai et al. , 2019 ; Desai & Johnson , 2021 ; Miech et al. , 2021 ; Vaswani et al. , 2017 ; Raffel et al. , 2020b ; Santoro et al. , 2018 ; Hudson & Zitnick , 2021 ; Ma et al. , 2021 ) . Perceiver IO builds on this body of work to produce a general purpose architecture that can be easily and widely applied . 3 THE PERCEIVER IO ARCHITECTURE . The Perceiver IO architecture builds on the Perceiver ( Jaegle et al. , 2021 ) , which achieved its crossdomain generality by assuming that its input is a simple 2D byte array : a set of elements ( which might be pixels or patches in vision , characters or words in language , or some form of embedding , learned or otherwise ) , each described by a feature vector . The model then encodes information about the input array using a smaller number of latent feature vectors , using Transformer-style attention , followed by iterative processing and a final aggregation down to a category label . Rather than output a single category , Perceiver IO aims to have the same level of generality with respect to its outputs as the Perceiver has with respect to its inputs : that is , it should produce arbitrary output arrays . We can predict each element of the output array using another attention module by querying the latent array using a query feature vector unique to the desired output element . In other words , we define a query array with the same number of elements as the desired output . The queries may be hand-designed , learned embeddings , or a simple function of the input . They attend to the latents to yield an output array of the desired shape .

Summary: This paper proposed perceiver IO, which is a general architecture for structured input & output. Perceiver IO is based on the perceiver architecture which scales linearly with the size of inputs and outputs and augments with a flexible querying mechanism similar to NeRF. Different from prior work, perceiver IO directly operates in the raw input space -- UTF-8 bytes for language, xy coordinates in optical flow, raw audio, etc. The output can be arbitrary in size and different structures, and with non-autoregressive decoding, the model can handle large input and output sizes. The proposed model is tested on variety of tasks, including language modeling, optical flow, video audio class autoencoding, image classification, and starcraft II, and achieves superior performance.