Next Patch Prediction for Autoregressive Visual Generation

Introduction

Autoregressive models, foundational to large language models (LLMs) (Vaswani et al., 2017; Devlin et al., 2018; Radford et al., 2018; Raffel et al., 2020; Radford et al., 2019; Brown et al., 2020; Zhang et al., 2022), generate content through the prediction of subsequent tokens in a sequence. This Next Token Prediction paradigm enables LLMs to excel in a variety of natural language processing tasks, exhibiting human-like conversational abilities (Ouyang et al., 2022; OpenAI, 2022; 2023; Google, 2023; Anthropic, 2023; Workshop et al., 2022; Touvron et al., 2023a; b; Bai et al., 2023a; Yang et al., 2023; Team, 2023; Bi et al., 2024) and demonstrating remarkable scalability (Kaplan et al., 2020; Henighan et al., 2020; Hoffmann et al., 2022; Wei et al., 2022; Alabdulmohsin et al., 2022; Chowdhery et al., 2023; Anil et al., 2023). Such advancements illustrate the potential for achieving general-purpose artificial intelligence systems.

Inspired by the success of autoregressive models in the language domain, their applications for image generation are significantly explored. Notable approaches, including VQVAE (Van Den Oord et al., 2017; Razavi et al., 2019), VQGAN (Esser et al., 2021; Lee et al., 2022), DALL-E (Ramesh et al., 2021), and Parti (Yu et al., 2021; 2022), introduce image tokenizers that convert continuous images into discrete tokens, employing autoregressive models to sequentially generate these tokens, thereby achieving image generation. In parallel, diffusion models (Song & Ermon, 2019; Ho et al., 2020; Song et al., 2020; Dhariwal & Nichol, 2021; Lu et al., 2022a; Ho et al., 2022a; Ho & Salimans, 2022; Rombach et al., 2022; Ramesh et al., 2022; Saharia et al., 2022; Rombach et al., 2022; Peebles & Xie, 2023; Podell et al., 2023; Xue et al., 2023; Chen et al., 2023a; b; Betker et al., 2023; Li et al., 2024a; Esser et al., 2024) emerge as a distinct and rapidly evolving paradigm in image generation. With wide open-source community support, diffusion models dominate the visual generation tasks. However, the fundamental differences in the underlying methodologies of autoregressive and diffusion models pose significant challenges for developing a unified framework that integrates both language and vision tasks. More recently, a pioneering work LlamaGen (Sun et al., 2024) achieves the next token prediction paradigm for image generation with a vanilla autoregressive model, Llama, bringing the field one step closer to building a unified model between language and vision.

In this work, we follow LlamaGen (Sun et al., 2024) for autoregressive image generation and rethink the modeling of the NTP paradigm in the following aspects.

a). The NTP paradigm, widely successful in large language models, leverages the high information density of text tokens. However, this property does not directly transfer to image data, where information density is naturally low compared to language data. Our key insight is to process image tokens into high information density units (i.e., patches), potentially leading to better performance in autoregressive image generation.

b). Transformer-based autoregressive models incur significant computational costs during training, with the total cost approximately scaling as $C\approx 6WN$, where $W$ represents the number of model parameters and $N$ denotes the input sequence length. While maintaining the model architecture, we could manage to reduce the input length of image tokens, improving training efficiency.

c). Unlike language data, the image modality inherently exhibits a hierarchical property in understanding and generation. This suggests that the training process for autoregressive image generation could be modeled in a multi-scale, coarse-to-fine manner, potentially enhancing generation quality.

Building on these insights, we introduce Next Patch Prediction (NPP), a simple yet effective method for autoregressive visual generation. Specifically, the input image tokens are grouped and aggregated into patch tokens of higher information density through an intra-patch average operation. With the resulting patch tokens as a shorter input sequence, the autoregressive model is trained to predict the next patch, thus significantly reducing the computational cost. We further propose a multi-scale patch grouping strategy, modeling patches in a coarse-to-fine manner, which enhances training efficiency and generation quality. As shown in Figure 1, experiments on a diverse range of models from 100M to 1.4B parameters demonstrate that the next patch prediction paradigm could reduce the training cost to $\sim 0.6\times$ while improving image generation quality by up to 1.0 FID score on the ImageNet benchmark. We highlight that our method retains the original autoregressive model architecture without introducing additional trainable parameters or specifically designing a custom image tokenizer, ensuring flexibility for seamless adaptation to various autoregressive models addressing visual generation tasks.

Next Patch Prediction

We introduce the Next Patch Prediction paradigm in Figure 2. The input image is initially encoded into image token indexes, which are then mapped to token embeddings of sequence length $N$. Considering the naturally low information density of image data, our key idea is to aggregate multiple tokens into groups of units containing higher information density. Specifically, we group tokens into non-overlapping patches and generate a sequence of patch embeddings with length $\frac{N}{K}$, where $K$ is the number of tokens associated with each patch. To avoid introducing extra parameters during this compression process, we simply adopt an intra-patch average operation to compute the patch embeddings. Formally, given the embedding function $E$, for the $i$-th patch $p_{i}$ associated with $K$ image tokens $x_{k}^{i}$ in the input sequence, the patch embedding is formulated as,

With the resulting patch tokens forming a shorter input sequence, the autoregressive model is trained to predict the next patch. However, due to the absence of an explicit ground truth (GT) index for a patch token, Cross-Entropy (CE) Loss cannot be directly applied. To maintain the model architecture without introducing additional parameters, we propose a patch CE loss that supervises the model using the associated $K$ image token GT indexes $Index_{k}^{i}$ in the next patch $p_{i}$. Specifically, given the next patch predictions as $Pred_{i}$, and recalling the path sequence length $\frac{N}{K}$, the loss function is formulated as:

Moreover, we propose a multi-scale, coarse-to-fine patch grouping strategy that leverages the natural hierarchical structure of image data. Specifically, the grouping function begins with a large kernel size, resulting in large patches and a short patch sequence length, allowing the autoregressive model to capture coarse representations. As training progresses, the patch size is gradually reduced to $1\times 1$, enabling the model to learn finer details, making the next token prediction a special case of the next patch prediction paradigm. To balance training efficiency and model performance, we introduce a segment scheduling factor $\lambda$ and set the number of patch levels. As illustrated in Figure 3, during the total training steps $T$, each segment is represented as $\lambda^{L}T-\lambda^{(L-1)}T$ with a patch size (PS) of $2^{L}\times 2^{L}$, where $L$ denotes the current patch level. The computational cost is reduced by a factor of $\frac{1}{PS}$ due to the shorter sequence length at each level.

Experiments

Benchmark. We build the Next Patch Prediction based on LlamaGen (Sun et al., 2024) and evaluate it on the class-conditional image generation task using the standard ImageNet1K benchmark (Deng et al., 2009).

Model Architecture. For the image encoder, we adopt the same VQGAN tokenizer trained by LlamaGen on ImageNet1K. The tokenizer has a vocabulary size of 16,384 and downsamples the input image at a fixed ratio of $16\times 16$. For the autoregressive model, we adopt the same setting as LlamaGen. Note that our method does NOT introduce any extra trainable parameters and thereby can be easily extended to other autoregressive models or scaling up to similar tasks.

Training & Inference Settings. All the model are trained for 300 epochs following the same setting of LlamaGen (Sun et al., 2024): base learning rate of $1\times 10^{-4}$ per 256 batchsize, AdamW optimizer with $\beta_{1}=0.9,\beta_{2}=0.95$, weight decay $=0.05$, gradient clipping set to 1.0. To enable smooth transfer between different patch size segments, we set learning rate warmup for the first 1 epoch and linearly decay to $1\times 10^{-5}$ for the last $1/5$ number of epochs in each segment. The dropout ratio in the autoregressive model backbone is set to 0.1. We also set the class token embedding dropout ratio to 0.1 for classifier-free guidance. For inference, we strictly adopt the same inference setting as LlamaGen (Sun et al., 2024).

Baseline Methods. We choose baseline methods from popular image generation models, including GAN (Brock et al., 2018; Kang et al., 2023; Sauer et al., 2022), Diffusion models (Dhariwal & Nichol, 2021; Ho et al., 2022b; Rombach et al., 2022; Peebles & Xie, 2023), masked-prediction models (Chang et al., 2022) and autoregressive models (Esser et al., 2021; Yu et al., 2021; Lee et al., 2022; Tian et al., 2024). As our method is built upon LlamaGen (Sun et al., 2024), we take it as a strong baseline and mainly compare our method with it.

2 Main results

We compare our method with various baseline works on class-conditional ImageNet 256×256 benchmark and show the results in Table 1. Our method achieves state-of-the-art performance on a diverse model size from 100M to 3B parameters compared to baseline methods. For LlamaGen-B and LlamaGen-L, our method improves the FID score up to 1.0 with a higher inception score. For LlamaGen-XL and LlamaGen-XXL, our methods trained on 256$\times$256 resolution even outperform the baseline work trained on 384$\times$384 resolution. In general, the proposed method outperforms the baseline work LlamaGen by improving the image generation quality up to 1.0 FID scores with a lower training cost. We also provide the training cost study in Section 3.6.

3 Effect of Patch Size

We study the effect of different patch sizes and present the results in Table 2. In this experiment, we disable the multi-scale grouping strategy and set the segment scheduling factor $\lambda=1/2$ with the patch size adopted in the first $1/2$ number of training epochs ablated. We find that with a larger patch size (PS), the learned knowledge cannot be smoothly transferred to the case with $PS=1\times 1$, leading to a performance drop.

4 Effect of Segment Scheduling Factor λ\lambda

We provide a study on the effect of different segment scheduling factors as shown in Table 3. In this study, we disable the multi-scale grouping strategy and set $PS=2\times 2$. We show that larger $\lambda$ results in lower training computational cost but with slight performance degradation.

5 Effect of Multi-scale Grouping Strategy

We provide a study on the effect of the multi-scale grouping strategy as shown in Table 4. In this study, we set $\lambda=1/2$ and compare different numbers of patch levels $\#L$. Experiments show this strategy makes a trade-off between training computational cost and image generation quality.

6 Training Cost Study

We provide a study on the training cost as shown in Table 4. We compare baseline methods with different settings of our method, reporting the average computation cost (GFLOPs). Our method achieves the best balance of model performance and training efficiency.

7 Generation Results

In Figure 4, we present some of the generations results on ImageNet 256$\times$256 benchmark.

Related Works

Generative adversarial networks (GANs) (Goodfellow et al., 2014; Brock et al., 2018; Karras et al., 2019; Kang et al., 2023) are the pioneering method for visual generation in the deep learning era, focusing on learning to generate realistic images through adversarial training. Inspired by language model architectures, BERT-style models (Chang et al., 2022; 2023; Yu et al., 2023a; b) emerge, using masked-prediction techniques to learn to predict missing parts of images, much like how BERT predicts masked words in text. Autoregressive models (Esser et al., 2021; Ramesh et al., 2021; Yu et al., 2022), inspired by GPT, predict the next token in a sequence. These methods often involve an image tokenization step (Kingma & Welling, 2013; Van Den Oord et al., 2017), converting pixel space into a more semantically meaningful representation. Diffusion models (Ho et al., 2020; Song & Ermon, 2019; Song et al., 2020; Dhariwal & Nichol, 2021; Rombach et al., 2022; Lin et al., 2024a) introduce a novel approach, treating visual generation as a reverse diffusion process, where images are gradually denoised from Gaussian noise through a series of steps.

2 Multimodal Foundation Models

Recent advancements in vision-and-language models (Liu et al., 2024; Zhu et al., 2023; Dai et al., 2024; Peng et al., 2023; Zhang et al., 2023; Pagnoni et al., 2024; Ma et al., 2024a; Lin et al., 2023; 2024b; Jin et al., 2024) have demonstrated impressive capabilities in visual understanding through visual instruction tuning. However, unifying the understanding and generation tasks in multimodal models is still being explored. Most existing approaches (Sun et al., 2023b; a; Dong et al., 2024; Ge et al., 2023; Zhou et al., 2024; Xiao et al., 2024; Xie et al., 2024; Li et al., 2024b; Gu et al., 2024; Tang et al., 2024; Ma et al., 2024b; Tong et al., 2024; Li et al., 2024c; Deng et al., 2024) focus on integrating diffusion models with other existing pre-trained models, rather than adopting a unified next-token prediction paradigm. These methods often require complex designs to link two distinct training paradigms, which makes scaling up more challenging. Some works (Yu et al., 2024b; Qu et al., 2024; Zhu et al., 2024; Han et al., 2024; Shi et al., 2024; Luo et al., 2024; Chen et al., 2024) focus on image tokenizer for better compression and reconstruction of image data, which is crucial for the image generation quality. Some pioneering efforts (Lu et al., 2022b; 2023; Bai et al., 2023b; Team et al., 2023; Team, 2024; Wu et al., 2024; Wang et al., 2024; Sun et al., 2024; Tian et al., 2024; Pang et al., 2024; Yu et al., 2024a; He et al., 2024) explore incorporating image generation into large language models using an autoregressive approach, achieving promising results. However, most of them do not focus on the autoregressive image generation method or directly adopt the next token prediction paradigm. More recently, some efforts Shao et al. (2024); Pagnoni et al. (2024); The et al. (2024) for LLM study novel sequence modeling methods, proposing to aggregate text tokens into units (i.e., patch or concept) with higher information density. In this work, we draw inspiration from them and propose a novel next patch prediction paradigm for autoregressive visual generation without introducing additional trainable parameters or specifically designing a custom image tokenizer, ensuring flexibility for seamless adaptation to various autoregressive image generation tasks.

Limitation and Future Works

Currently, due to limited computational resources, we are unable to provide experimental results on text-to-image tasks or develop a unified vision-language model. However, our method retains the original autoregressive model architecture without introducing additional trainable parameters or specifically designing a custom image tokenizer. We believe it can be seamlessly adapted to various autoregressive models addressing similar tasks, offering a baseline framework for future exploration. For future works, exploring the Next Patch Prediction paradigm for autoregressive video generation is an interesting direction, where the input sequence length is significantly longer, resulting in higher computational costs. By leveraging our method, we anticipate improved efficiency and scalability in handling the extended sequences inherent in video data. Another potential research direction is to explore more effective methods for obtaining patch embeddings. In this work, we adopt a simple intra-patch averaging operation, which is straightforward but may be limited in its capacity to capture fine-grained details. Designing more powerful and expressive patch extraction methods could further improve training efficiency and enhance generation quality.

Conclusion

In this work, we introduce a novel Next Patch Prediction paradigm that improves autoregressive image generation by grouping and aggregating image tokens into high-density patch tokens. This approach reduces the computational cost to $\sim 0.6\times$ while improving image generation quality by up to 1.0 FID score on the ImageNet benchmark. Our method retains the original autoregressive model architecture without introducing additional trainable parameters or custom image tokenizers, thereby making the next patch prediction paradigm seamlessly adapted to various autoregressive models addressing image generation tasks.

References