Scalable Adaptive Computation for Iterative Generation (2212.11972v2)

Abstract: Natural data is redundant yet predominant architectures tile computation uniformly across their input and output space. We propose the Recurrent Interface Networks (RINs), an attention-based architecture that decouples its core computation from the dimensionality of the data, enabling adaptive computation for more scalable generation of high-dimensional data. RINs focus the bulk of computation (i.e. global self-attention) on a set of latent tokens, using cross-attention to read and write (i.e. route) information between latent and data tokens. Stacking RIN blocks allows bottom-up (data to latent) and top-down (latent to data) feedback, leading to deeper and more expressive routing. While this routing introduces challenges, this is less problematic in recurrent computation settings where the task (and routing problem) changes gradually, such as iterative generation with diffusion models. We show how to leverage recurrence by conditioning the latent tokens at each forward pass of the reverse diffusion process with those from prior computation, i.e. latent self-conditioning. RINs yield state-of-the-art pixel diffusion models for image and video generation, scaling to 1024X1024 images without cascades or guidance, while being domain-agnostic and up to 10X more efficient than 2D and 3D U-Nets.

• The paper introduces RINs, a novel architecture that decouples core computations from input data, enhancing scalability in generative models.
• It employs a dual-token system where latent tokens concentrate intensive computation through global self-attention, combined with dynamic cross-attention.
• Experiments demonstrate up to tenfold efficiency gains and superior FID and Inception Scores compared to traditional models in image and video generation.

Scalable Adaptive Computation for Iterative Generation: A Review

The paper presents an architecture called Recurrent Interface Networks (RINs), designed to optimize the computation of generative models for high-dimensional data such as images and videos. RINs differentiate themselves from conventional models by decoupling core computation from data dimensionality, which facilitates adaptive computation and improved scalability. This approach mitigates the inefficiencies seen in prevalent architectures that uniformly allocate computation across input and output spaces.

Summary of Methodology

RINs leverage attention mechanisms for processing information differentially based on task requirements. The architecture employs two categories of tokens: interface tokens, which directly relate to the input data, and latent tokens, which undergo the majority of computational processing. The bulk of computation, specifically global self-attention, occurs on the latent tokens, while cross-attention dynamically routes information between interface and latent tokens. The separation and selective focus on latent tokens allow RINs to efficiently handle large-scale data sets, making them notably more efficient than 2D and 3D U-Nets used in state-of-the-art diffusion models.

A notable aspect of the architecture is its ability to integrate bottom-up and top-down feedback loops through stacked RIN blocks. These loops enhance routing expressiveness and computational depth but also introduce challenges associated with recurrent computations. To address this, the paper introduces latent self-conditioning, where latent tokens are conditioned based on previous iterations in the reverse diffusion process. This method effectively forms a deeper network without requiring expansive latents, leading to significant efficiency gains.

Experimental Results

The architecture is tested using diffusion models for image and video generation tasks. RINs demonstrate state-of-the-art performance, scaling to $1024 \times 1024$ images without necessitating cascading models or guidance and outperforming existing models by up to tenfold in efficiency. The empirical results underscore the architecture's domain-agnostic nature, providing versatility across different generative tasks while ensuring optimal allocation of computational resources.

The experiments indicate that RINs deliver better FID and Inception Scores compared to traditional and convolution-based models, supporting their effective computation allocation and scalable design. Interestingly, despite having fewer inductive biases than convolutional architectures, RINs manage to maintain competitive performance even on smaller dataset tasks such as CIFAR-10, highlighting their adaptability and robustness.

Implications and Future Directions

The proposed architecture extends the theoretical understanding of adaptive computation in high-dimensional generative modeling. Practically, it suggests alternative strategies for computational allocation in AI models, potentially reducing the need for computationally intensive techniques such as cascades or guidance that other models typically require.

The conception of latent self-conditioning presents new opportunities for optimization, suggesting avenues for future research into adaptive and scaled routing within recurrent computations. This self-conditioning mechanism enhances depth and expressiveness without compromising on computational efficiency, presenting a valuable aspect for further exploration.

RINs offer a notable contribution to generative modeling in AI, providing a scalable architecture with flexible computation allocation strategies. As we advance in the realms of generative AI, there is scope for integrating RINs with latent diffusion strategies and adaptive guidance techniques to further enhance model performance and efficiency. The findings advocate for a reconsideration of traditional evolutionary architectures in favor of models that allow better adaptability and resource efficiency tailored to data dynamism.