Autoregressive Image Generation with Linear Complexity: A Spatial-Aware Decay Perspective

TL;DR

This paper introduces LASAD, a novel linear attention mechanism that preserves spatial relationships in images, enabling efficient autoregressive image generation with high quality and reduced computational complexity.

Contribution

The paper proposes LASAD, a spatial-aware decay attention mechanism that maintains 2D spatial relationships, improving linear attention's effectiveness in image generation.

Findings

LASADGen achieves state-of-the-art performance on ImageNet.

LASAD reduces computational complexity to linear scale.

The method effectively captures long-range dependencies in images.

Abstract

Autoregressive (AR) models have garnered significant attention in image generation for their ability to effectively capture both local and global structures within visual data. However, prevalent AR models predominantly rely on the transformer architectures, which are beset by quadratic computational complexity concerning input sequence length and substantial memory overhead due to the necessity of maintaining key-value caches. Although linear attention mechanisms have successfully reduced this burden in language models, our initial experiments reveal that they significantly degrade image generation quality because of their inability to capture critical long-range dependencies in visual data. We propose Linear Attention with Spatial-Aware Decay (LASAD), a novel attention mechanism that explicitly preserves genuine 2D spatial relationships within the flattened image sequences by computing position-dependent decay factors based on true 2D spatial location rather than 1D sequence positions. Based on this mechanism, we present LASADGen, an autoregressive image generator that enables selective attention to relevant spatial contexts with linear complexity. Experiments on ImageNet show LASADGen achieves state-of-the-art image generation performance and computational efficiency, bridging the gap between linear attention's efficiency and spatial understanding needed for high-quality generation.