PaperlayerView exact revision on arXiv

Source-linked AI summary

Scalable Adaptive Computation for Iterative Generation

Allan Jabri, David Fleet, Ting Chen

arXiv:2212.11972v2cs.LGcs.CVcs.NE

TL;DR

RINs address the limited scalability of architectures that allocate computation uniformly across redundant, high-dimensional data. They route information through compact latent tokens and reuse latent state during iterative diffusion generation. The resulting models outperform leading image and video diffusion approaches while reducing inference computation by up to 10×.

  • Problem

    Existing architectures allocate computation uniformly, while scalable adaptive computation for large input and output spaces remains unresolved.

  • Method

    RINs route information between input-linked interface tokens and compact latent tokens, using latent self-conditioning to reuse prior latent context during recurrent diffusion computation.

  • Results

    RINs outperform leading image and video diffusion models from 64×64 to 1024×1024 and reduce video inference-step FLOPs by 10×.

  • Takeaways & Limitations

    RINs provide a domain-agnostic architecture that is particularly effective for iterative generation, where persistent context can amortize routing challenges.

  • Takeaways & Limitations

    The reported advantages are more pronounced in recurrent settings where inputs change gradually and persistent context can be propagated.

Abstract

from arXiv · show

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.

1. Introduction

RINs address the difficulty of scaling generative models by routing information through compact latent tokens instead of applying global computation uniformly across high-dimensional data. In recurrent diffusion settings, latent self-conditioning supports expressive routing, yielding improved efficiency and performance for image and video generation.

  • Motivation: Natural data is often redundant and unevenly structured, but prevailing architectures allocate computation uniformly across input data.This creates a scalability challenge for high-dimensional generation, where computation could instead focus on complex image regions or changing video content.
  • Architecture: RINs partition hidden units into interface tokens tied to the input and compact latent tokens decoupled from data dimensionality.Cross-attention routes information from interface tokens to latents and back, while most computation operates on the latents.
  • Architecture: Stacked RIN blocks alternate reading, global latent computation, and writing, allowing bottom-up and top-down context to inform subsequent routing.This produces deeper and more expressive information routing without quadratic pairwise interactions across interface tokens.
  • Recurrent computation: Latent self-conditioning reuses latents from previous diffusion iterations instead of reinitializing them, providing persistent context without backpropagation through time.The approach addresses routing challenges when the task and routing problem change gradually across recurrent computation steps.
  • Results: 10× fewer FLOPs per inference step accompany RINs’ reported advantages over leading diffusion approaches for both image and video generation.For ImageNet, RINs outperform leading models from 64×64 through 1024×1024; on Kinetics600 video prediction, they surpass leading approaches while reducing per-step FLOPs by 10×.
  • Contributions: The paper proposes RINs as a domain-agnostic architecture for adaptive computation and scalable generation of high-dimensional data.The authors specifically identify recurrent computation settings as favorable for amortizing routing difficulty and report gains over U-Net diffusion models.

2. Method

RINs separate data-facing interface tokens from compact latent tokens, routing information between them so high-capacity computation can focus on latents. Their recurrent latent self-conditioning carries context across diffusion iterations, addressing cold-start routing while supporting adaptive generation.

  • 2. Method: RINs partition hidden units into an input-sized interface and input-independent latents that form a compact computational representation.Interface units connect locally to the input, while the bulk of computation operates on latents decoupled from input dimensionality.
  • 2. Method: Each RIN block reads interface information into latents, computes with latent self-attention and token-wise MLPs, then writes latent information back to the interface.Read and write use cross-attention; stacked blocks enable repeated bottom-up and top-down updates.
  • 2. Method: Unlike U-Nets and standard Transformers, RINs avoid fixed global downsampling or pairwise token attention, reserving high-capacity processing for selectively read latent information.Uniform interface computation remains lightweight, while latent processing adapts to input structure and content.
  • 2.1. Background: Iterative Generation with Diffusion: Diffusion generation repeatedly denoises a noisy state, making its gradual refinement a natural recurrent setting for RINs.The denoising network receives the noisy image, time step, and optional conditioning, then predicts noise for the next reverse transition.
  • 2. Method: RINs initialize image or video interfaces by tokenizing inputs with positional encodings, while latents begin as learned embeddings and can include conditioning tokens.Images use patch embeddings and videos use 3-D patches; class labels and diffusion time steps are concatenated to the latent set.
  • 2. Method: Latent self-conditioning uses prior latent activations to mitigate cold-start routing in recurrent settings where inputs and routing problems change gradually.At inference, previous-step latents initialize the current step; training estimates latents with a zero-conditioned pass and stops gradients through those estimates.

3. Experiments

Experiments evaluate RINs on image generation and video prediction, including implementation choices, adaptive computation, comparisons with leading methods, and ablations. The results show strong quality and efficiency, while routing and latent self-conditioning materially affect performance.

  • 3.3. Comparison to SOTA: RINs attain superior ImageNet generation quality with fewer parameters and FLOPs, using 256 latents even for 1024×1024 images with 16,384 tokens.At 1024×1024, this configuration is reported as 2× more efficient than the 256×256 ADM U-Net while operating at 4× higher resolution.
  • 3.3. Comparison to SOTA: RINs surpass leading Kinetics-600 video-prediction approaches while reducing computation by up to 10× per step without guidance.The video model uses 3D patches and otherwise shares the architecture used for 256×256 image generation.
  • 3.4. Ablations: Latent self-conditioning correlates with improved sample quality, and a training rate of 0.9 is used for the best reported results.A rate of 0 means self-conditioning is absent during training and inference; positive rates use prior latent context, always at inference.
  • 3.4. Ablations: More frequent read-write operations delay performance saturation and improve sample quality, while RINs handle up to 4096 interface tokens with reasonable performance.The ablations support iterative routing and tokenization flexibility across patch sizes.
  • 3.1. Implementation Details: A sigmoid noise schedule with appropriate temperature is more stable than cosine during training, especially for larger images, while cosine remains sufficient for sampling.The default sigmoid temperature is 0.9; noise-schedule choice has less impact during sampling.
  • 3.5. Visualizing Adaptive Computation: During generation, read attention becomes sparser and more data-specific, concentrating latent computation on regions with higher information.In video prediction, attention also favors information not directly copied from conditioning frames, such as object motion and panning.

4. Related Work

RINs relate to adaptive and recurrent architectures while targeting scalable pixel diffusion for images and videos. Their distinguishing use of latent self-conditioning propagates routing context across recurrent computation.

  • RINs use recurrence to propagate routing context through deep computation graphs and amortize iterative routing costs.
  • RINs decouple computation from input structure using latent units, while allocating compute non-uniformly with fixed, differentiable computation graphs.
  • Unlike prior diffusion Transformers that lack strong large-input results, RINs provide a domain-agnostic path toward scalable image and video generation.
  • Latent self-conditioning distinguishes RINs from related self-conditioning methods by conditioning on hidden activations rather than data predictions.

5. Conclusion

The conclusion presents RINs as a domain-agnostic architecture that separates interface and latent tokens to scale adaptive computation for high-dimensional generation. It identifies latent self-conditioning and combinations with complementary diffusion techniques as future directions.

  • RINs partition hidden units into interface and latent tokens, decoupling core computation from data layout.
  • Persistent context in recurrent computation amortizes latent construction while retaining efficient training.
  • RINs perform efficiently for image and video generation, although the conclusion frames broader modality scaling as an intended scope.
  • Future work includes understanding latent self-conditioning and combining RINs with guidance and latent diffusion.

A. Architecture Implementation Pseudo-code

The pseudo-code initializes image interfaces from patch tokens and processes them with latent tokens through repeated attention and feed-forward blocks. Latent self-conditioning optionally incorporates prior latents before computation, and readout reconstructs image outputs.

  • Algorithm 3 changes interface initialization for other tasks or modalities while retaining the RIN computation blocks.
  • Each core block reads interface tokens into latents with cross-attention, applies latent self-attention and feed-forward computation, then writes back to interfaces.
  • Image interfaces are created by strided convolutional patchification, layer normalization, and positional embeddings.
  • Latent self-conditioning adds a transformed previous-latent state to the current latent initialization when prior latents are available.
  • RIN computation repeats the core block for a configurable number of blocks and layers per block.
  • The readout linearly projects interface tokens and rearranges them from patches into the output image.

B. More Details of Training / Sampling Algorithms, and Noise schedules

The appendix specifies continuous-time noise schedules and DDIM/DDPM sampling updates used with RIN diffusion models. It compares cosine and sigmoid schedules, finding the sigmoid schedule more stable for training larger images while cosine remains adequate for sampling.

  • Noise schedules: Continuous-time noise scheduling is represented by a function γ(t), with cosine and sigmoid implementations provided.
  • Sampling algorithms: The appendix presents DDIM and DDPM as the diffusion updating rules used for x_t estimation.
  • Sampling algorithms: DDIM estimates the next sample from the current noisy input, predicted clean sample, and noise schedule values γnow and γnext.
  • Sampling algorithms: DDPM updates combine the predicted clean sample, reconstructed noise, schedule ratio, and fresh Gaussian noise to produce x_next.
  • Noise schedules: The sigmoid schedule with appropriate temperature is more stable than cosine during training, particularly for larger images.
  • Noise schedules: For sampling, schedule choice has less impact, so the default cosine schedule can suffice.

C. Hyper-parameters and Other Training Details

This section presents the model and training hyper-parameters, along with the hardware and optimizer settings used for training.

  • Most models are trained on 32 TPUv3 chips with a batch size of 1024, while larger resolutions use 64 TPUv3 or 256 TPUv4 chips.The 512×512 and 1024×1024 models use the larger configurations, respectively.
  • All models use the LAMB optimizer.
  • Model hyper-parameters are provided in Table C.1.
  • Training hyper-parameters are provided in Table C.2.

D. Sample Visualizations

The visualizations show class-conditional image samples, high-resolution generation, multimodal video prediction, and adaptive computation focused on motion regions.

  • Selected ImageNet 256×256 samples cover classes including space shuttle, arctic fox, go-kart, macaw, white wolf, lion, siberian husky, lorikeet, giant panda, and cockatoo.
  • RINs generate uncurated class-conditional ImageNet samples at 768×768 and 1024×1024 resolutions using classifier-free guidance.The 768×768 visualization states that the model uses input scaling and operates directly on pixels without a multiscale cascade.
  • Kinetics-600 video prediction exhibits multiple possible futures, including varying movements, trajectories, outcomes, and camera motions.Examples are shown at 16×64×64 with conditioning frames from the test set.
  • Adaptive-computation visualizations are provided for ImageNet 256×256 samples and Kinetics-600 video prediction.
  • For Kinetics-600, read attention and computation focus on regions of motion rather than regions that could be generated by copying conditioning frames.The attention visualization subsamples videos by 2× in time and includes read attention on conditioning frames.
Loading 2212.11972v2…