JetFormer: An autoregressive generative model of raw images and text
Authors: Michael Tschannen
André Susano Pinto
Alexander Kolesnikov
Venue: ICLR 2025 Poster (2025)
Keywords: generative pretraining
multimodal models
vision-language model
text-to-image
normalizing flow
generative model
Primary area: generative models
TL;DR: We present JetFormer, an autoregressive generative model of raw images and text. The key distinction from previous models is that JetFormer does require pre-trained image tokenizers and scales to high-resolution images.

Abstract: Removing modeling constraints and unifying architectures across domains has been a key driver of the recent progress in training large multimodal models. However, most of these models still rely on many separately trained components such as modality-specific encoders and decoders. In this work, we further streamline joint generative modeling of images and text. We propose an autoregressive decoder-only transformer---JetFormer---which is trained to directly maximize the likelihood of raw data, without relying on any separately pretrained components, and can understand and generate both text and images. Specifically, we leverage a normalizing flow model to obtain a soft-token image representation that is jointly trained with an autoregressive multimodal transformer. The normalizing flow model serves as both an image encoder for perception tasks and an image decoder for image generation tasks during inference. JetFormer achieves text-to-image generation quality competitive with recent VQVAE- and VAE-based baselines. These baselines rely on pretrained image autoencoders, which are trained with a complex mixture of losses, including perceptual ones. At the same time, JetFormer demonstrates robust image understanding capabilities. To the best of our knowledge, JetFormer is the first model that is capable of generating high-fidelity images and producing strong log-likelihood bounds.

Published as a conference paper at ICLR 2025
JETFORMER: AN AUTOREGRESSIVE GENERATIVE MODEL OF RAW IMAGES AND TEXT
Michael Tschannen; André Susano Pinto; Alexander Kolesnikov
Google DeepMind
(tschannen, andresp, akelesnikov
ABSTRACT
Removing modeling constraints and unifying architectures across domains has been a key driver of the recent progress in training large multimodal models. However, most of these models still rely on many separately trained components such as mooaalty-speccffcc encoders and decoders. this work, we further stream- line joint generative modeling of images and text. We propose an autoregressive decoder-only Sennformer--eeFoomee----
is trained to directly maximize the likelihood of raw data, without relying on any separately pretrained components, and can understand and generate both text and images. Specifically, we lever- age a normalizing flow model to obtain a soft-token image representation that is jointly trained with autoregressive multimodal transformer. The normalizing flow model serves as both an image encoder for perception tasks and an image decoder for image generation tasks during inference. JetFormer achieves text-to- image generation quality competitive with recent VQVAE- and VAE-based base- lines. These baselines rely pretrained image autoencoders, which are trained with a complex mixture of losses, including perceptual ones. At the same time, JetFormer demonstrates robust image understanding capabilities. To the best of our knowledge, JetFormer is the first model that is capable of generating high- fidelity images and producing strong likeliiooooo bounds.
INTRODUCTION
The "Bitter lesson" (Sutton 2019) has been the prime force behind the recent progress in machine learning and artificial intelligence research. It suggests that general-purpose methods which effec- tively leverage large amounts of compute and data prevail over specialized techniques designed by domain experts. Arguably the most prominent examples in this context are transformer decoder-only models trained for next-token prediction (Vaswani et al.) 2017 Radford et al. 2018), that outperform task-specific NLP systems, and transformer encoders in computer vision (Dosovitskiy et al.) 2021
Strudel et al. 2021 Li et al. 2022), that achieve better quality than CNN-based models.
This trend is also visible in the current pursuit of extending L.MSs both understand and generate multiple modalities such as text and images with a single model. A powerful paradigm in the lit-
erature (Aghajanyan et al. 2022 Kim et al., 2023 AAhaaanyyn et al. 2023 You al. (023) is to model the image tokens using discrete tokens obtained via (VQIVAEs (van den Oord et al.) 2017
Esser al. 2020 Ramesh al. 202T). One limitation of these approaches that the conversion from image into tokens and vire-versa is performed by separate, frozen, codaaiity-specfic and lossy encoder (and decoder) trained ahead of time. As a result, this image encoder may be agnostic to the actual task at hand and limit the performance of the resulting model (Dong et al. 2023 Pan
jet al. 2024. Xu et al. 2024).
To obtain general architecture that can generate multiple modalities but does not have (limiting) components pretrained ahead of time, we develop new generative model: the JetFormer. It can be trained from scratch and optimized end-to-end for the log- likeeiill of raw training data. We demonstrate this for text and pixels. To this end, we combine normalizing flow (Dinh et al. 2016 Kingma & Dhariwal 2018 for computing soft-token image representation with decoder-only transformer (Vaswani et al. 2017) and a soft-token Gaussian mixture loss (Tschannen et al. 2024)
"Equal contribution.

---

Published as a conference paper at ICLR 2025
Targets:
Loss:
d.thhh.
Flow Invertible (1., lossless) and differentiable
Transformer
Inputs:
+R0S>
<BOK>
Figure 1: Visualization of JetFormer training with teacher-forcing. The ground truth consists of text tokens and images. The images are converted into soft tokens via normalizing flow. The sequence of tokens is processed by an auto-regressive transformer and its outputs parameterize either a discrete distribution or Gaussian mixture depending whether the target a discrete or soft token.
The key insight behind the JetFormer model, is that a powerful normalizing flow (which we call "jet", hence the model name) be used to encode images into latent representation suitable for autoregressive modeling. Intuitively, raw image patches encoded as pixels have very complex structure, which makes direct autoregression futile: date, there no convincing demonstration of this. At the same time, the flow model is lossless and can be trained together with the (multimodal) autoregressive model end-to-end. At inference time an image decoder is readily available, since our flow model is invertible in closed form.
Although we only optimize for log-likelihood, is worth noting that doing so naively does not guar- antee being able to generate images with global coherence. Similarly the vast majority of work on high-fidelity image generation (Esser et al. 2020 Dhariwal Nichol| 202T Ho & Salimans) 2022), guide the model to focus on the high-level information. To this end, explore two approaches.
First, we introduce a novel technique that is based on image augmentation during training. The main idea is to add Gaussian noise during training, but reduce it through the course of training. Intuitively, this pushes the model towards prioritizing high-level information early on; see Section [3.2. for more details. Even though noise curriculum during training is inspired by diffusion models (Sohl- Dickstein et al. (015), it is very different on the technical level, and the resulting model does not perform progressive image denoising inference time.
Second, we explore two ways of managing redundancy natural images. JetFormer readily allows excluding a subset of redundant dimensions from the autoregressive model. As an alternative, explore PCA for reducing image dimensionality.
We conduct experiments ImageNet class-conditional image generation and web-scale mul- timodal generation, thereby demonstrating the JetFormer works and scales to both text-to-image generation and vision- language understanding with a single model.
In summary, our contributions are:
We present JetFormer, a generative model composed of a transformer and a normalizing flow that can be trained from scratch to model text and raw pixels jointly, end-to-end.
We show that image augmentation based on a noise curriculum can significantly boost image generation quality of such likelihood-based models.
We demonstrate our proposed end-to-end model is competitive with less flexible techniques when trained web-scale data, and can generate both images and text.
2

---

Published as a conference paper at ICLR 2025
Selected class-conditional 256x256 samples from the JetFormer-L trained on Imagenet. The samples produced with the CFG sampler, where CFG strength is equal to 4.0.
Figure 2:
RELATED WORK
Generating natural images autoregressively as a sequence of discrete- (sub- biib-)pixels was ex- tensively explored in the literature using CNNs (Van den Oord al., [2016bla] Salimans et al. 2016) or transformers (Parmar et al.; 2018 Chen et al. 2020). While achieving excellent results log-likel,hood, these models
computationally expensive and do not scale well to high image resolutions. A related family of models are normalizing flows (Dinh et al.] 2014 2016 Kingma & Dhariwal 2018 Ho et 2019), invertible models which are trained map image pixels a sim- plc prior by maximizing log- lieliikl. These models scale better but achieve lower likelihood than autoregressive models and empirically fail generate nigh-fidelity images, even for low resolutions.
More recently, compressing the high- diggesiiinsii image pixel space to lower- dimensional se- quence of discrete tokens via a pretrained, frozen VQ-VAE (van den Oord et al., [201 Razavi al.] 2019), and then modeling the compressed sequence with transformer decoder has emerged as scalable technique for high-fidelity image generation (Esser et al,, 2020. Yu et al.) [022a Ramesh al. 2021: Yu et al. 2022cc Gafni et al. 2022). To enable semantic compression, VQ-VAEs typi- cally rely on perceptual and GAN losses. Moreover, VQ-VAEEbased representations are common in
the context of dense prediction tasks (Kolesnikov al. 2022 Lu et al. 2022 Mizrahi et al. 2024 Mentzer et al. 2024). in particular when modeling multiple modalities jointly. GIVT (Tschannen
al. 2024) showed that combining autoencoder and autoregressive transformer can be ap- sequences, by directly modeling feature vectors in the latent space of a VAF, without any quantization. Somewhat related, (Weiss al. 2021 Nachmani al.
plied to continuous-valun
2023 Meng
et al. 2024) explored soft tokens for speech synthesis.
VQ-VAEs are also becoming popular in the context of Vision-Language Models (VLMs). Such models are typically either trained from scratch (Radford et al.) 2021 Jia et al. 2021 Wang et al., 2022b Yu et al., 2022b) on web-scale data or constructed by combining and tuning pretrained vision encoder and a pretrained language model (Alayrac et al., 2022 Chen et al. 2023b Wang al. 2022a Li al., 2023; Huang et al.) 2023 Liu et al. | 2024 Beyer et al.) 2024). and can solve broad range of task which can be cast as text output. To enable pixel outputs for such models simple way extend the text vocabulary with VQ-VAF
VAF tokens (Aghajanyan et al. 2022 Kim al. 2023
Aghajanyan et al. 2023 You et al. 2023 Pan et al. (202).. Other works (Dong al. 2023 Ge ct al.
2024 Zhou et al. 2024) combine VLMs with (latent) diffusion models (Sohl-Dicksten) et al. 2015
Dhariwal & Nichol, Rombach et al. 2022 Saharia et al., 2022 Ramesh et al.) (222) to enable image generation capabilities. JetFormer is related to this category of models, but unlike previous models does rely any pretrained (VQ-)VAB vision encoder/decoder..
METHOD
Modeling natural images with autoregressive transformers poses many obstacles. Doing in pixel space is a viable approach (Parmar et al.] 2018; Chen et al. 2020), but it quickly becomes com- putationally prohibitive. Even image of size 256x256x3 would require predicting/sampling

---

Published as a conference paper at ICLR 2025
nearly 200 000 tokens. Alternativlyy, modeling images at the patch-level to tame the computational complexity creates its own challenges. Each patch is a sample from a complex distribution, and producing all of its dimensions in single forward pass fails to model pixel interactions.
Currently, the most common approach to overcome these issues is to leverage a standalone image encoder/tokenizer (and decoder/detokenizzer model, which encodes image sequence of (usu- ally discrete) tokens. Such an image encoder performs semantic compression and thus reduces the computational burden. However, this type of approach has significant shorteomings: one needs to tolerate precision loss due to compression and to commit to the image encoder that was trained ahead of time and may not be suitable for the modeling task hand.
In this paper we overcome both of these issues and propose JetFormer, a generative autoregressive decoder-only model able to model both text and images, and directly learn from the raw training data. The JetFormer model is trained end-to-end, without relying on any lossy modality- encoders/tokenizzrs. Our model built on two key insights.
First, we use continuous (also called "soft") tokens to represent images. As shown in GIVT (Tschan- nen et al.) 2024), transformer decoders can generate the soft-token-sequence produced by VAE encoder for high-fidelity image generation. Specifically, GIVT replaces the categorical prediction head with Gaussian Mixture Model (GMM) that models log-likelihood of soft image embeddings.
Second, instead of VAE (that needs to be trained ahead of time), we use a normalizing flow model to learn a soft-token image representation suitable for autoregressive modeling. As flow models are lossless by design, they won't suffer from representation collapse and can be trained simultaneously with the transformer model without auxiliary losses, eliminating the necessity to use pretrained encoders and enabling full end-to-end learning of autoregressive transformers from raw images.
Naturally, the above two insights can be combined with the standard transformer for text, forming a simple and unified multimodal model, capable of learning from image and text data.
3.1 MODELING IMAGES IN PIXEL SPACE WITH SOFT TOKENS AND NORMALIZING FLOWS
As outlined above, we model an image x using a normalizing flow model (Dinh et al. 2014 2016 Kingma & Dhariwal 2018) f(x) that losslessly maps an image into a sequence of embeddings { J, which we also call "soft tokens". Note that the flow preserves the total number of the input dimensions. These embeddings are then modeled by the deep autoregressive model p, where the outputs are modeled with GMM, as proposed in GIVT. We then maximize the image log-likelihood lower bound I::
L(x) logp(f(x)) + log det (Of(z)
){
where
Oxt
f(x) = [2], 22, ,2, sn] and p(2) II
(2)
Note the log- determiiennn term arises from the normalizing flow model, as a part of the data log- likelihood, see Dinh et al. (2016). Further, to ensure correctness, we apply image dequantization, as outlined in Theis et al. (2015). This amounts adding uniform noise the input images I, s.t. T+u, where u ~ U [0, IT. This guarantees that we optimize lower bound on the discrete image log-likelihood. For clarity, we point out that both p and f both have learnable parameters which are optimized with gradient-based method (training via teacher forcing is illustrated in Figure
In simple words, JetFormer for images an autoregressive model, where inputs and targets are produced by the flow model, which re-encodes input images. Due to the end-to-end nature of the objective, the flow model is incentivized to learn a sequence of embeddings, that makes autore- gressive modeling effective as possible. During inference, the autoregressive model generates sequence of soft tokens, which then need to be decoded into an image using the inverse of the flow.
3.2 IMPROVING MODEL QUALITY FOR NATURAL IMAGES
While JetFormer works out of the box, we have found several modeling enhancements which greatly improve the quality of the generated images. In particular, factoring out latent dimensions, the use of massifier-rree-uuidae during sampling, and novel noise curriculum.
4

---

Published as a conference paper at ICLR 2025
Factoring out redundant dimensions Natural images are redundant, intrinsically low- dimensional signals with low-frequency components dominating the spectrum. The design of Jet- Former enables a simple and effective way to leverage this observation and improve model quality, while also reducing computational burden.
The key observations is that not all output dimensions of the invertible flow need be further pro- cessed by the autoregressive model. We can model a subset of dimensions (i.e. a subset of channels) with Gaussian distribution, PN, and the remaining ones with the autoregressive transformer:
(f(()) du!
L(z) = log p(2) + log (x) + log det
where [2, = f(x)
Intuitively, we expect redundant dimensions be "factored out" غ, as they do not require further heavy processing. We verify our intuition empirically in the experimental section and in Figure /c
As a strong baseline to the above approach, we also consider a more direct approach handle redundancy in images. Before feeding x to the flow model, we reshape it into a sequence of flattened patches and apply a learnable, invertible linear map W along the channel dimension. We want this map to learn separate important dimensions of the flattened patches from redundant ones. To this end, we feed the first channels its output =W T the normalizing flow and model the remaining channels with a Gaussian distribution. Intuitively, given the simplicity of the transform W applied before factoring out part of the sequence. minimizing the NLL while training will ensure that the part of the sequence is modelled by JetFormer, whereas the low-level noise will be mapped to the Gaussian prior. This similar to the reasoning behind probabilistic PCA (Tipping & Bishop 1999). Indeed, observe that the transform learned by the model is close to applying PCA to image patches (see Figure 6d), and we observe that when initializing W with PCA and freezing it we obtain similar results. See Appendix B for formal definition of this approach.
Classifier-free guidance Following common practice in the diffusion and autoregressive image modeling literature, we employ classifier-free guidance (CFG) (Ho & Salimans 2022) which was previously shown to substantially improve sample quality. We reuse the distribution-bas based variant implemented via rejection sampling for GMMs from (Tschannen et al. (024)) without modification.
3.2.1 1 RGB NOISE CURRICULUM DURING TRAINING
It is common to explicitly factorize data into semantically meaningful parts to improve image quality. One approach is to model RGB pixels as sequences of increasing color depth and/or resolution (Kolesnikov & Lampert| 2017 Menick & Kalchbrenner 2019: Nash et al. (2021). Similarly, adding noise to RGB pixels closely related to reducing the color depth and effective resolution. This has led to the interpretation of diffusion models, where denoisers are trained at different noise levels according to a predefined noise schedule, as learning a hicrarchical representation in pixel space induced by the noise schedule (Kingma & Gao 2023 Dieleman 2024).
Building on this intuition, we alter the training procedure by introducing "noise curriculum": additive Gaussian pixel noise during JetFormer training. The noise is strongest in the beginning of the training and decays towards zero. We use cosine decay schedule for the noise standard deviation.
In the beginning of the training, when strong (high-variance) noise is added to the image, JetFormer learns to model coarse image information (see Figure 66)) As training progresses, the model grad- ually learns finer level of detail, while "remembering" previously learned patterns. In the end of the training, JetFormer uses the correct distribution for training. Intuitively, this scheme prioritizes modeling of high-level image structure without sacrificing overall performance.
Importantly, unlike in diffusion, the noise curriculum merely acts as a data augmentation during training. The model is conditioned on the noise magnitude and, at inference, image is not gradually denoised, but generated autoregressively in the latent space of normalizing flow.
For an integer-valued RGB image 1, the noisy image obtained as [1 + (0, where the noise scale 0F as a function of the training progresses E [0, 1] follows the cosine schedule
1+cos(tbr) The shape of the noise schedule is visualized in Figure
at on
5

---

Published as a conference paper at ICLR 2025
Negatva log- aggelkalhhood Icalld lina)
with curriculam
-60
10 Noise standard devitation (dashed
noss curicauum
175
-0
-43
125
baseline samples
100
.20
50
25
80%
samples
Training progress
Training progress
Figure 3: Training with noise curriculum on RGB input images. The left plot demonstrates negative log- ongllikelhhood progression (solid lines) and the noise strength schedule (dashed lines). Importantly, both NLL curves land very similar points, despite visual quality measured by FID (middle plot) and the typical samples (right plot) being very different: observe that noise curriculum guides the JetFormer towards modeling high-level image structure.
3.3 JOINT GENERATIVE MODELING OF PIXELS AND TEXT
We explore JetFormer for multimodal generative modeling, where the model can perform both dis- criminative and generative tasks on all modalities, focusing on images and text. Sophisticated mod- els of this class are trained on interleaved sequences of images and text (Aghajanyan et al.) [022] Kim et al. 2023 Aghajanyan al. 2023). often with post-training refinement, which enables few- shot image-to-ted togge-to-textt text (e.g. captioning) and text-to-image (e.g. image editing) capabilities. Here, we follow (Kim al. 2023 You et al. (023) and consider image-text pairs from the web a proxy for more complex, interleaved setups, and do not involve a post-training step. While conceptually sim- ple, this allows to explore vision-language understanding tasks such captioning and VQA, as well as text-to-ine text-to-image image generation. Extending the image generation approach discussed in the previ- ous section to this setup is straight-forward We simply extend the transformer backbone generating soft tokens to modeling language tokens produced by standard language tokenizer with separate prediction head and softmax.
We train sequences of both image tokens followed by text tokens, and vice versa, only applying loss to the second part (modality) of the sequence. We use the respective negative log-likelhhodd, and weight to balance the two. We observed that applying the loss to the full sequence leads to worse results, possibly because predicting an image prefix effectively means unconditionally modeling web images, which is very challenging. We expect this to change for interleaved sequences, which may provide a stronger conditioning signal.
For text-to-image generation, the text prefix acts conditioning, and image generation is per- formed as described in Section 3.1 For image-to-text generation the normalizing flow acts as an image encoder. The model uses the same soft token space for generation and understanding.
EXPERIMENTS
Architecture We rely the simple, transformer-bused design from (Kolesnikov al. 2024) for the normalizing flow model. This design uses affine coupling layers (predicting an element-wise scale and shift) consisting of spatial and channel-wise splitting functions split the activations in two parts, and a stack ViT blocks (Dosovitskiy et al.] (02T)) applied to half of the activations to infer the affine transform. Here, we only use channel-wise splitting we found in preliminary experiments that spatial splitting did not improve modeling quality. We set the depth to 32 coupling blocks, each consisting of stack of 4 The (input and output) feature shape for this model is H W
VIT blocks of width 512, and with 8 attention heads.
pp 3pp)) when feeding the full image, and (H) W. d) when using a dense invertible map or PCA to reduce dimensionality prior to applying p flow. For image size = W 256 and patch size = 16 this amounts to 256 768, after flattening spatial dimensions and with dimensionality reduction to d = 128 to 256x 128.

---

Published as a conference paper at ICLR 2025
Table 1: Comparison of JetFormer trained for 500 epochs on ImageNe1256 with baselines from the literature. For large enough scale JetFormer approaches models without components pretrained in an extra step.
extra step
FID
Precision
Recall
NLL
BigGAN-deeep (Brock et al.
2018
6.95
0.87
0.28
ADM-G (Dhariwal & Nichol
2021
4.59
0.82
0.52
LDM-4G (Rombach al.
2022)
VAE
3.60
0.87
0.48
VQGAN
[Esser al. 2020)
VQ-VAE VAE
5.20
ViT-VQGAN YYu et al. 2022a)
"Q-VAE
3.04
GIVT-Causal (Tschannen et 2024)
VAE
3.35
0.84
0.53
JetFormer-B
7.25
0.72
0.44
3.06
JetFormer-L
6.64
0.69
0.56
3.05
For the latent autoregressive decoder-only backbone, we rely on the Gemma architecture (Gemma Team et al., 2024). We consider 3 different model shapes amounting to 350M, 750M and 1.3B parameters which are largely inspired by previous decoder-only models for image generation (Esser) (ct al. 2020) (see Appendix A). For the GMM prediction head, unless explicitly noted, we set the number of mixtures k 1024, use multivariate Gaussians of dimension d 128 with diagonal covariance. For text we use a sentencepiece tokenizer with vocabulary size 32k trained on the English portion of the C4 corpus made available by Raffel al. (2020). We set the maximum number of text tokens to 64, but do not explicitly model padding tokens (i.e. during training we mask out attention elements corresponding to padding tokens and adapt the RoPE positions to skip them). When training for class-condite consiisondiiinall image generation, we use a prefix of 16 learned tokens per class (instead of single one), as we observed that this improves sample quality.
Training recipe We train the latent decoder-only backbone on concatenated sequences of discrete text and soft tokens (flow latents) via teacher forcing for NLL according to the respective distribution of the tokens (categorical for text tokens and GMM for soft tokens). Whether the sequence is an image-text or text-image sequence is sampled at random for every example. Normalizing flow (or learnable invertible patch embedding) is trained along with the maasformer--ackkoor end-to-end. This does not require any dedicated techniques thanks to the NLL for the GMM being differentiable evaluated w.r.t. When
both w.r. to the GMM parameters as well as the soft tokens/features training for captioning with image prefix (i.e. for image- sequences, where normalizing flow serves as a vision encoder), we stop the gradient at the flow output during pretraining. We did not observe improved performance when passing the gradients.
We use the Adam optimizer with learning rate 10-3, decoupled weight decay of 10-4, B2 parameter 0.95 and clip the gradient norms 1.0. We set the batch size to 4k. We also apply dropout with probability 0.1 at the output of the self-attention and MLP blocks, which we found to improve image sample quality. For both class conditional image generation and text-to-image generation we drop the conditioning with 10% probablilty and replace it with a learned [NOLABEL] token for CFG. Unless explicitly stated, we apply the RGB noise schedule the input images described in Section 3.1 with initial noise standard deviation To = 64 (for pixel values in the range [0, 255)). We decay to for ImageNet setup and to for multimodal setup. Inspired by the VAE-based GIVT, where for every example latent is sampled from the VAE encoder (i.e. approximate posterior) which might have regularizing effect, we add Gaussian noise with standard deviation 0.3 the flow latents. We normalize the image NLL to bits per dimension is common in the image generative modeling literature and apply weight of 0.0025 to the text NLL such that the loss magnitude per token is roughly identical for both modalities.
Training data For training class-conditional image generation models we use ImageNetlk (Rus- sakovsky et al. 2015). For multimodal generation we rely on the image-text pairs from WebLI data set (Chen et al. 2023b). In both cases, resize the images so that the shorter side is 256 pix- els while preserving the aspect ratio and extract a 256x256 central crop. Besides the RGB noise described earlier we do not apply any augmentation, except for random left-right flipping on Im-
7

---

Published as a conference paper at ICLR 2025
Table 2: Summary of the main JetFormer (trained for 100 epochs) ablations performed on class- conditional ImageNet256 generation. We demonstrate relative importance of various components and highlight the interesting interplay between PCA preprocessing and the noise curriculum.
FID
Precision
Recall
NLL
JetFormer-B
7.84
0.75
0.39
3.14
normalizing flow
117.76
0.17
0.32
6.84
no noise curriculum
44.71
0,45
0.28
3.05
factored-out dimensions no end-to-end training
17.29
0.65
0.26
3.13
11.16
0.68
0.33
3.08
learned inv. projection
10.19
0.73
0.32
4.78
no GMM (Gaussian loss only)
9.46
0.77
0.30
3.14
single class token
8.85
0.73
0.37
3.14
PCA preproc. + JetFormer-B
8.79
0.77
0.35
PCA preproe. JetFormer-B (no noise cur.)
13.16
0.71
0.31
kk On ImageNetlk, we train for 100 epochs in ablations, and 500 epochs otherwise. On WebLI, we train for 1B examples seen per modality, so 1B examples for sext-to-image only models, and 2B total for models that also support image-to-text (understanding) tasks.
Evaluations and metrics Following the diffusion literature, we use the ADM FID evaluation suite (Dhariwal & Nichol (021) (with 50k reference samples) and precision/recall (Sajjadi al., 2018) to assess image sample quality on ImageNetlk. For sext-to-image, we adopt the common MS-COCO FID-30K, generating images for captions from 30k randomly sampled COCO validation images and evaluating FID against reference statistics from the full COCO validation set. We report this metric both zero-shot and after fine-tuning on the COCO training set. As image understanding tasks we consider ImageNet zero-shot classification and fine-tune for image captioning (reporting CIDEr score) and visual question answering (VQA, measuring VQAv2 accuracy).
4.1 CLASS-CONDITIONAL IMAGE GENERATION
Table [ shows the sample quality of JetFormer trained on ImageNet with image resolution of 256 256 along with some baselines from the literature. Model samples are show Figure 2 and Appendix(F Despite being explicit NI.I. model and not using advanced image encoders, our Jet- Former model is competitive with those baselines. Interestingly, JetFormer has high recall of 0.56. We hypothesize that this is consequence of our model being explicit log- lekeeekl model and thus not suffering from the mode collapse.
Table 2 shows ablations of key design choices of JetFormer:
Removing the normalizing flow results a catastrophic loss in performance. This confirms our intuition that re-encoding image pixels with a flow model is essential.
Omitting the noise curriculum further results in much worse results.
We also confirm that not factoring out dimensions after the flow model leads to quality loss.
First training the normalizing flow minimize NLLI w.r.t. Gaussian prior, and then training the autoregressive transformer on the latent representation of the frozen flow model (factoring out redundant dimensions the frozen latent space with a learnable, invertible linear map) results in lower sample quality than end-to-end training.
Factoring out dimensions with a learnable linear mapping before the flow (as described in See. 3.2 leads to better results, but underperforms post-flow factoring out.
Reducing the number GMM components from 1024 to for the soft-token loss result in relatively mild performance drop and significant drop in recall, indicating more mixtures enable a better coverage of the distribution.
Finally, doing class-oondiiioning with a single class token (instead of 16 tokens) leads to slight performance drop. likely due the weaker conditioning signal.
Interestingly, we observe that modeling images after the PCA transform leads to somewhat worse results. However, in the presence of PCA, the noise curriculum becomes less important. It highlights
8

---

Published as a conference paper at ICLR 2025
Table 3: Comparison with baselines from the literature for text-to-image generation (0-shot COCO FID-30k, 256x256). We included autoregressive (first group) and diffusion models (second group) which use raw image-text data (without e.g. re-captioning) and do not rely on pretrained text en- coders. Models with image understanding capabilities are marked with T&I. FII (ft.) is obtained when fine-tuning on the COCO training set, and NLL is measured on the COCO validation set after fine-tuning. JetFormer is the only method which does not rely on any extra step.
extra step
#param.
FID
FID (ft.)
NLL (ft.)
DALL-E Ramesh 2021
VQ-VAE
12B
27.50
Vog (Ding et al. 2021)
VQ-VAE
4B
27.10
2ie (Ding al.
2022
VQ-VAE
6B
24.00
17.50
ARGVLT (T&I) (Kim al.
2023
VQ-VAE
0,45B
16.93
MAGVLT (T&I) (Kim et al.
2023
VQ-VAE
0.45B
12.08
Make-A-Scene (Ganni et al.
2022
VQ-VAE
4B
11.84
7.55
LDM-KL-8-G (Rombach al.
2022
VAE
1,45B
12.63
GLIDE (Nichol et al. 2022)
Super-res.
6B
12.24
DALL-E-2 (Ramesh et al.
2022
Super-res.
5.2B
10.39
JetFormer-L (T&I)
2.75B
20.86
13.70
3.86
JetFormer-I.
2.75B
18.63
13.07
3.85
importance of prioritizing the high-level information in images, if visual quality is the end goal. It also shows that manual preprocessing may reduce the necessity for various modeling tricks, yet it also shows that full end-to-end modeling prevails when done right.
4.2 EFFECT OF THE NOISE CURRICULUM
We have experimented with different levels of initial noise, and generally find that for 256x256 images the initial noise with standard deviation of 64 annealed to 0 during the training to be the best for the ImageNet setup. The initial noise levels of 128, 64, 32 and 0 result in FID 8.62, 7.84, 8.59 and 44.71 respectively. For the multimodal setup we observe that annealing the noise standard deviation to 3 (thus leaving small level of noise) improves FID. More analysis of the interplay between NLL and FID is presented in Appendix C
In Figure 3 we demonstrate the effect of noise curriculum with the initial standard deviation of 64. First, we observe that, in comparison to the noise-free baseline, the final NLL is not affected dramatically. However, as would be expected, initial stages of training have much worse NLLs due to strong noise. Second, being enabled. This is also evident by the final samples from the noise-free model and model with noise. The latter has much more pronounced emphasis on the high-level image structure.
demonstrate that FID drastically improves with the noise curriculum
4.3 MULTIMODAL GENERATION AND UNDERSTANDING
Figure 4 shows multimodal JetFormer models of different sizes, trained for text-to-image genera- tion only (T21) and both T2I and image-to-text generation (T&I). Zero-shot samples are show in Appendix B In both cases, increasing the model size improves quality. Training models for mixed T&I generation leads reduction in T2I generation quality, but also equips the model with vision- language understanding capabilities.
Tables and compare T2I and T&I JetFormer models with models from the literature. Sim- ilar prior work trained in generative fashion on image-text pairs relying on pretrained VQ-VAE (ARGVLTT MAGVIT (Kim et al.) 2023)) achieves better performance on T2I generation, but lags in understanding performance, which highlights the benefits of having access to unquantized end-to- end learned features for understanding. JetFormer also approaches the understanding performance of image-text models pretrained for understanding such as CLIP and CapPa, although these models have a significantly smaller size.
Finally, we present T&I baseline using pretrained, frozen VAE instead of an end-to-end trained normlizing flow in Appendix C and find that JetFormer clearly outperforms this baseline in T2I generation and all I2T tasks.

---

Published as a conference paper at ICLR 2025
(0-shot)
ImageNet (0-shot)
(fine-tuned)
VOAV2 (fine-tuned]
120
70
30
typo
52
69
TG
T21
50
116
68
121
40
114
67
46
112
20
44
2.0 110
paramstars
paremaleers
Figure 4: Zero-shot image sample quality, zero-shot ImageNet classification accuracy (by scoring class labels, without prompts), and captioning and VQA performance as a function of model size for JetFormer variants. Increasing the model size improves all the metrics. T&I models generally perform worse than unidirectional models which are trained either for T2I or 12T generation.
Table 4: Comparison with baselines pretrained raw image-text data and fine-tuned for image captioning (COCO captions; CIDEr) and VQA (VQAv2; test-dev acc.). Models marked with T&I were jointly pretrained for captioning and T21 generation. "The numbers are from (Tschannen al. 2023 for models pretrained on 900M image-text pairs, which most closely matches our setup.
extra step
COCO cap.
VQAv2
CapPa L/14 (Tschannen et al.
2023
118.7
68.6
CLIP L/14 (Radford al. 202
118.2
67.9
ARGVLT (T&I) (Kim et 2023
VQ-VAE VQ-VAE
94.7
MAGVITT Large (T&I) (Kim et al.
2023
110.7
65.7
JetFormer-B (I2T)
118.7
67.2
JetFormer-L (T&I)
119.8
70.0
5 CONCLUSION
In this paper we introduce JetFormer, a novel class of generative models that combines normalizing flows and autoregressive models with soft tokens. To the best of our knowledge. is the first image model, which is capable to synthesize high-resolutions images and provide an explicit (and competitive) NL.L bounds for the raw images. JetFormer is fully end-to-end trainable model (with no components pretrained ahead of time), which means that it can be fully tailored towards the task at hand, without being limited by the external and frozen components. Being able to compute NLL is also important feature of our model. NLL is tangible score closely related to compression capabilities, and can be used to compare various generative models across different modeling classes or for hill-cling fill-climiing. Also, by measuring NLL score one can ensure the absence of the mode collapse, because mode collapse will lead to deterioration NLL on the hold-out data.
We note that JetFormer in its current form also has some limitations. The visual quality of its samples lags behind state-of-the-ant diffusion models that leverage pretrained latent representations. Additionally, the full end-to-end nature of JetFormer also comes with increased computational quirements. However, given JetFormer's 's simple design, we believe that it can be scaled well so that the benefits of end-to-end training can come to full fruition.
Reproducibility Statement We provide detailed information about the training recipe, the archi- tecture, hyper-parameters and the training data in Section 4 and Appendix A
Ethics Statement This paper describes system for understanding and generation image-text data, with focus on characterizing and exploring its performance on academic data sets. It fits into the broader class of large multimodal models trained on data from the web, and the same ethical implications as for those prior works apply here. In particular, when releasing models publicly, extensive measures to de-biasing the data should be taken, and models should be safety-tuned and red-teamed prior release. Content-filtes to further improve safety. We refer to the corresponding papers for a more in-depth discussion, for example (Radford et al. 2021 Chen al. 2023at Po et al. 2024).
deploying such
be added to the inference pipeline
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENTS
We thank Lucas Beyer for detailed review and feedback on the final manuscript. We further thank Joan Puigeerver for timely infrastructure contributions.
REFERENCES
Armen Aghajanyan, Yao Huang, Candace Ross, Vladimir Karpukhin, Hu Xu, Naman Goyal, Dmytro Okhonko, Mandar Joshi, Gargi Ghosh, Mike Lewis, and Luke Zettlemoyer. CM3: A causal masked multimodal model of the internet. arXiv:2201.07520. 2022.
Armen Aghajanyan, Lili Yu, Alexis Conneau, Wei-Ning Hsu. Karen Hambardzumyan, Susan Zhang, Stephen Roller, Naman Goyal, Omer Levy, and Luke Zettlemoyer. Scaling laws for gen- erative mixed-modal language models. In ICML, 2023.
Jean-Baptiste Alayrac. Jeff Donahue, Pauline Luc, Antoine Miech, lain Barr, Yana Hasson, Karel Lenc, Arthur Mensch, Katie Millican, Malcolm Reynolds, Roman Ring. Eliza Rutherford, Serkan Cabi, Tengda Han, Zhitao Gong, Sina Samangooei, Marianne Monteiro, Jacob Menick, Sebastian Borgeaud, Andrew Brock, Aida Nematzadeh, Sahand Sharifzadeh, Mikolaj Binkowski, Ricardo Barreira, Oriol Vinyals, Andrew Zisserman, and Karen Simonyan. Flamingo: A visual language model for few-shot learning. In NeurlPS, 2022.
Lucas Beyer, Andreas Steiner, André Susano Pinto, Alexander Kolesnikov, Xiao Wang, Daniel Salz, Maxim Neumann, Ibrahim Alabdulmohsin, Michael Tschannen, Emanuele Bugliarello, et al. PaliGemma: A versatile 3B VLM for transfer. arXhv:2407.07726. 2024.
Andrew Brock, Jeff Donahue, and Karen Simonyan. Large scale GAN training for high fidelity natural image synthesis. In ICLR, 2018
Mark Chen, Alee Radford. Rewon Child, Jeffrey Wu, Heewoo Jun, David Luan, and Ilya Sutskever. Generative pretraining from pixels. In /CMI., 2020.
Xi Chen, Josip Djolonga, Piotr Padlewski, Basil Mustafa, Soravit Changpinyo, Jialin Wu, Car- los Riquelme Ruiz, Sebastian Goodman, Xiao Wang, Yi Tay, Siamak Shakeri, Mostafa De- hghani, Daniel Salz, Mario Lucic, Michael Tschannen, Arsha Nagrani, Hexiang Hu, Mandar Joshi, Bo Pang. Ceslce Montgomery, Paulina Pictrzyk, Marvin Ritter, AJ Piergiovanni, Matthias Minderer, Filip Pavetic, Austin Waters, Gang Li, Ibrahim Alabdulmohsin, Lucas Beyer, Julien Amelot, Kenton Lee, Andreas Peter Steiner, Yang Li, Daniel Keysers, Anurag Arnab, Yuanzhong Xu, Keran Rong, Alexander Kolesnikov, Mojtaba Seyedhosseini, Anelia Angelova, Xiaohua Zhai, Neil Houlsby, and Radu Soricut. PaLI-X: On scaling up multilingual vision and language model. arXiv:2305. 18565, 2023a.
Xi Chen, Xiao Wang, Soravit Changpinyo, AJ Piergiovanni, Piotr Padlewski, Daniel Salz, Sebastian Goodman, Adam Grycner, Basil Mustafa, Lucas Beyer, et al. PaLI: A jointly-scaled multilingual language-image model. ICLR, 2023b.
Prafulla Dhariwal and Alexander Nichol. Diffusion models beat GANs on image synthesis. NeurIPS, 2021.
Sander Dieleman. Diffusion is a Blog post, 2024. html
//nnerr.
Ming Ding, Zhuoyi Yang, Wenyi Hong, Wendi Zheng, Chang Zhou, Da Yin, Junyang Lin, Xu Zou, Zhou Shao, Hongxia Yang, et al. Cogview: Mastering sot-to-image generation via transformers. NeurlPS, 2021.
Ming Ding, Wendi Zheng, Wenyi Hong, and Jie Tang. Cogview2: Faster and better text-to-image generation via hierarchical transformers. NeurlPS, 2022.
Laurent Dinh, David Krueger, and Yoshua Bengio. NICE: Non-linear independent components estimation. arXiv:/410.855 2014.
11

---

Published as a conference paper at ICLR 2025
Laurent Dinh, Jascha arXiv:160S.08803, 2016.
and Samy Bengio. Density estimation using real NVP.
Runpei Dong, Chunrui Han, Yuang Peng, Zekun Qi, Zheng Ge, Jinrong Yang, Liang Zhao, Jianjian Sun, Hongyu Zhou, Haoran Wei, Xiangwen Kong, Xiangyu Zhang, Kaisheng Ma, and Li Yi. Dreamllm: Synergistic multimodal comprehension and creation. In ICLR, 2023.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszko- reit, and Neil Houlsby, An image is worth 16x16 words: Transformers for image recognition at scale. ICLR, 2021.
Patrick Esser, Robin Rombach, and Björn Ommer. Taming transformers for high-resolution image synthesis. In CVPR, 2020.
Oran Gafni, Adam Polyak, Oron Ashual, Shelly Sheynin, Devi Parikh, and Yaniv Taigman. Make- a-scene: Scene-base sexttto-image generation with human priors. In ECCV, 2022.
Yuying Ge, Sijie Zhao, Jinguo Zhu, Yixiao Ge, Kun Yi, Lin Song, Chen Li, Xiaohan Ding, and Ying Shan. Seed-x: Multimodal models with unified multi-granularity comprehension and generation. arXiv:2404. 14396, 2024.
Gemma Team, Thomas Mesnard, Cassidy Hardin, Robert Dadashi, Surya Bhupatiraju, Shreya Pathak, Laurent Sifre, Morgane Rivière, Mihir Sanjay Kale, Juliette Love, et al. Gemma: Open models based on gemini research and technology. arXiv:2403. 08295, 2024.
Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. arXiv:2207. /2598, 2022.
Jonathan Ho, Xi Chen, Aravind Srinivas, Yan Duan, and Pieter Abbeel. Flow++: Improving flow- based generative models with variational dequantization and architecture design. In /CML, 2019.
Shaohan Huang, Li Dong, Wenhui Wang, Yaru Hao, Saksham Singhal, Shuming Ma, Tengchao Lv, Lei Cui, Owais Khan Mohammed, Qiang Liu, Kriti Aggarwal, Zewen Chi, Johan Bjorck, Vishrav Chaudhary, Subhojit Som, Xia Song, and Furu Wei. Language is not all you need: Aligning perception with language models. arXiv:2302. 2023.
Chao Jia, Yinfei Yang, Ye Xia, Yi-Ting Chen, Zarana Parekh, Hieu Pham, Quoc Le, Yun-Hsuan Sung, Zhen Li, and Tom Duerig. Scaling up visual and vision-language representation learning with noisy supervision. In ICML, 2021.
Sungwoong Kim, Daejin Jo. Donghoon Lee, and Jongmin Kim. MAGVLT: Masked generative vision-and-language transformer. In CVPR, 2023.
Diederik Kingma and Ruiqi Gao. Understanding diffusion objectives as the ELBO with simple data augmentation. NeurIPS, 2023.
Durk P Kingma and Prafulla Dhariwal. Glow: Generative flow with invertible 1x1 convolutions. NeurIPS, 2018.
Alexander Kolesnikov and Christoph Lampert. PixelCNN models with auxiliary variables for natural image modeling. In ICML, 2017.
Alexander Kolesnikov, André Susano Pinto, Lucas Beyer, Xiaohua Zhai, Jeremiah Harmsen, and Neil Houlsby. UViM: A unified modeling approach for vision with learned guiding codes.
NeurlPS, 2022.
Alexander Kolesnikov, André Susano Pinto, and Michael Tschannen. Jet: modern transformer- based normalizing flow. arXiv:2412.15129 2024.
Junnan Li, Dongxu Li, Silvio Savarese, and Steven Hoi. BLIP-2: Bootstrapping language-image pre-training with frozen image encoders and large language models. arXiv:2301.1259// 20233.
Yanghao Li, Hanzi Mao, Ross Girshick, and Kaiming He. Exploring plain vision transformer back- bones for object detection. In ECCV, 2022.
12

---

Published as a conference paper at ICLR 2025
Tsung- Yi Lin, Michael Maire, Serge Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and C Lawrence Zitnick. Microsoft COCO: Common objects in context. In (CCV,2014
Haotian Liu, Chunyuan Li, Qingyang Wu, and Yong Jae Lee. Visual instruction tuning. NeurIPS, 2024.
Jiasen Lu, Christopher Clark, Rowan Zellers, Roozbeh Mottaghi, and Aniruddha Kembhavi. Unified-IO: A unified model for vision, language, and multi-modal tasks. In ICLR, 2022.
Lingwei Meng, Long Zhou, Shujic Liu, Sanyuan Chen, Bing Han, Shujie Hu, Yanqing Liu, Jinyu Li, Sheng Zhao, Xixin Wu, et al. Autoregressive speech synthesis without vector quantization. 00Xiv:2407.08551,, 2024.
Jacob Menick and Nal Kalchbrenner. Generating high fidelity images with subscale pixel networks and multidimensional upscaling. ICLR, 2019.
Fabian Mentzer, David Minnen, Eirikur Agustsson, and Michael Tschannen. Finite scalar quantiza- tion: VQ-VAE VAE made simple. ICLR, 2024.
David Mizrahi, Roman Bachmann, Oguzhan Kar, Teresa Yeo, Mingfei Gao, Afshin Dehghan, and Amir Zamir. 4M: Massively multimodal masked modeling. NeurlPS, 2024.
Eliya Nachmani, Alon Levkovitch, Roy Hirsch, Julian Salazar, Chulayuth Asawaroengchai, Soroosh Mariooryad, Ehud Rivlin, RJ Skerry-Ryan, and Michelle Tadmor Ramanovich. Spo- ken question answering and speech continuation using spectrogram-powered Ilm. arXiv preprint arXiv:2305. 15255, 2023.
Charlie Nash, Jacob Menick, Sander Dieleman, and Peter Battaglia. Generating images with sparse representations. In ICML, 2021.
Alexander Quinn Nichol, Prafulla Dhariwal, Aditya Ramesh, Pranav Shyam, Pamela Mishkin, Bob Mcgrew, Ilya Sutskever, and Mark Chen. GLIDE: Towards photorealistic image generation and editing with text-guided diffusion models. In ICML, 2022.
Kaihang Pan, Siliang Tang, Juncheng Li, Zhaoyu Fan, Wei Chow, Shuicheng Yan, Tat-Seng Chua, Yueting Zhuang, and Hanwang Zhang. Auto-encoding morph-tokens for multimodal 1lm. arXiv:2405.01926. 2024.
Niki Parmar, Ashish Vaswani, Jakob Uszkoreit, Lukasz Kaiser, Noam Shazeer, Alexander Ku, and Dustin Tran. Image transformer. In ICML, 2018.
Ryan Po, Wang Yifan, Vladislav Golyanik, Kfir Aberman, Jonathan T Barron, Amit Bermano, Eric Chan, Tali Dekel, Aleksander Holynski, Angjoo Kanazawa, et al. State of the art on diffusion models for visual computing. In Computer Graphics Forum, 2024.
Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever. Improving language under- standing by generative pre-training. 2018.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, al. Learning transferable visual models from natural language supervision. In ICML, 2021.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. JMLR, 2020.
Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Ilya Sutskever. Zero-shot text-to-image generation. In ICML, 2021.
Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text- conditional image generation with clip latents. arXiv:2204.0612, 2022.
Ali Razavi, Aaron Van den Oord, and Oriol Vinyals. Generating diverse high-fidelity images with VAE-2.. NewrIPS, 2019.
13

---

Published as a conference paper at ICLR 2025
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution image synthesis with latent diffusion models. In CVPR, 2022.
Olga Russakovsky, Jia Deng. Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, al. Imagenet large scale visual recognition challenge. IJC,, 2015.
Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily L Denton, Kamyar Ghasemipour, Raphael Gontijo Lopes, Burcu Karagol Ayan, Tim Salimans, et al. Photorealistic text-to-image diffusion models with deep language understanding. NeurIPS, 2022.
Mehdi SM Sajjadi, Olivier Bachem, Mario Lucic, Olivier Bousquet, and Sylvain Gelly, Assessing generative models via precision and recall. NeurIPS, 2018.
Tim Salimans, Andrej Karpathy, Xi Chen, and Diederik P Kingma. PixeICNN++: Improving the PixeICNN with discretized logistic mixture likelihood and other modifications. In ICLR, 2016.
Amanprect Singh, Vivek Natarjan, Meet Shah, Yu Jiang, Xinlei Chen, Devi Parikh, and Marcus Rohrbach. Towards VQA models that can read. CVPR, 2019.
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In ICML, 2015.
Robin Strudel, Ricardo Garcia, Ivan Laptev, and Cordelia Schmid. Segmenter: Transformer for semantic segmentation. In CVPR, 2021.
Richard Sutton. The bitter lesson. Incomplete Ideas (blog). 2019.
Lucas Theis, Aäron van den Oord, and Matthias Bethge. A note on the evaluation of generative models. arXiv:1511.01844 2015.
Michael E Tipping and Christopher Bishop. Probabilistic principal component analysis. Journal of the Royal Statistical Society Series B: Statistical Methodology. 1999.
Michael Tschannen, Manoj Kumar, Andreas Steiner, Xiaohua Zhai, Neil Houlsby, and Lucas Beyer. Image captioners are scalable vision learners too. In NewrIPS, 2023.
Michael Tschannen, Cian Eastwood, and Fabian Mentzer. GIVT: Generative infinite-vocabulary transformers. In ECCV, 2024.
Aaron Van den Oord, Nal Kalchbrenner, Lasse Espeholt, Oriol Vinyals, Alex Graves, et al. Condi- tional image generation with pixelenn decoders. NeurlPS, 2016a.
Aäron Van den Oord, Nal Kalchbrenner, and Koray Kavukcuoglu. Pixel recurrent neural networks. In ICML, 2016b.
Aäron van den Oord, Oriol Vinyals, and Koray Kavukcuoglu. Neural discrete representation learn- ing. NewrIPS, 2017.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. NeurIPS, 2017.
Jianfeng Wang, Zhengyuan Yang, Xiaowei Hu, Linjie Li, Kevin Lin, Zhe Gan, Zicheng Liu, Ce Liu, and Lijuan Wang. GIT: A generative image-to-text transformer for vision and language. TMLR, 2022a.
Zirui Wang. Jiahui Yu, Adams Wei Yu, Zihang Dai, Yulia Tsvetkov, and Yuan Cao. SLM Simple visual language model pretraining with weak supervision. ICLR, 2022b.
RonJ J Weiss, RJ Skenry- -yann Eric Battenberg, Soroosh Mariooryad, and Diederik P Kingma. Wave- Tacotron: Spectrogram-frre end-to-end text-to-speech synthesis. In ICASSP, 2021.
Yifan Xu. Xiaoshan Yang. Yaguang Song, and Changsheng Xu. Libra: Building decoupled vision system on large language models. arXiv:2405. 10140, 2024.
14

---

Published as a conference paper at ICLR 2025
Haoxuan You, Mandy Guo, Zhecan Wang, Kai- Chang, Jason Baldridge, and Jiahui Yu. CoBIT: A contrastive bi-directional image-text generation model. arXiv:2333..3345, 2023.
Jiahui Yu, Xin Li, Jing Yu Koh, Han Zhang, Ruoming Pang, James Qin, Alexander Ku, Yuanzhong Xu, Jason Baldridge, and Yonghui Wu. Vector-quantized image modeling with improved VQ- GAN. ICLR, 2022a.
Jiahui Yu, Zirui Wang, Vijay Vasudevan, Legg Yeung, Mojtaba Seyedhosseini, and Yonghui Wu. Coca: Contrastive captioners are image-text foundation models. TMLR, 2022b.
Jiahui Yu, Yuanzhong Xu, Jing Yu Koh, Thang Luong, Gunjan Baid, Zirui Wang, Vijay Vasude- van, Alexander Ku, Yinfei Yang, Bureu Karagol Ayan, Benton C. Hutchinson, Wei Han, Zarana Parekh, Xin Li, Han Zhang, Jason Baldridge, and Yonghui Wu. Scaling autoregressive models for content-rich text-to-image generation. 7MLR, 2022c.
Chunting Zhou, Lili Yu, Arun Babu, Kushal Tirumala, Michihiro Yasunaga, Leonid Shamis, Jacob Kahn, Xuezhe Ma. Luke Zettlemoyer, and Omer Levy. Transfusion: Predict the next token and diffuse images with one multi-modal model. arXiv:2408. /039, 2024.
15

---

Published as a conference paper at ICLR 2025
APPENDIX - SUPPLEMENTARY MATERIAL
A ARCHITECTURES AND HYPERPARAMETERS
Table 5 shows the shapes of the autoregressive transformer and normalizing flow components for different JetFormer variants. One can see that the dense layer predicting the GMM parameters are relatively large. To save memory, we store this layer in the bfloate 6 format (instead of float 32 as the other parameters) which does not impact the quality. Also note that at inference time applying this layer can be greatly sped by first sampling the mixture and then only inferring the mean and covariance that mixture.
We rely on the design from (Kolesnikov et al. 2024) for the normalizing flow, which uses stack affine coupling blocks (Dinh et al., 2016, Eqn. 7, 8). For each block the relation between input sequence y/1 and output 11+1
Vi+ 1 yi
V1 = 1
9i+1 = (w: + b.:))) @ ((a((T)))
+
Y. = Vi+ 1 (a.(( (T(++11).
Here, and are obtained by splitting yi into equally-shaped halves along the channel dimension according to a randomly initialized partition (and M/+1 results from merging M and fi+1 according to this partition). The scale a, and shift b, are inferred by two separate heads from shared VIT g:-
The normalizing flow component has about 450M to 650M parameters, depending the variant. This is more than typical CNN-based (AEs for image generation in the literature (Esser al.
2020 Rombach et al.) 2022), but comparable to common ViT-based VQ- VQ-VAEs (Yu et al. 2022a
Table 6 shows the hyper- parameters used to fine-tune JetFormer. Overall observed dropout to have a significant impact on the image sample quality. When reporting fine-tuning-basee metrics report test-dev accuracy from the
report the median of 10 runs, except in VQAv2 where evaluation server.
Table Architecture details for different JetFormer variants.
Section for discussion.
JetFormer-B
JetFormer-M
JetFormer-L
Autoregressive transformer
Depth
24
48
Width
1024
1536
1536
MLP hidden dim.
4096
6144
6144
Num. heads
16
16
16
Num. KV heads
1
Head dim.
64
96
96
Vocab. size
32k
32k
32k
Num. mixtures
1024
1024
1024
Num. GMM param.
269M
404M
404M
Num. vocab, param.
33M
49M
49M
Num. backbone param.
389M
847M
1.65B
Normalizing flow
Num. coupling blocks
32
32
32
Block width
512
512
512
Block depth
4
Block MLP hidden dim.
2048
2048
2048
Block num. heads.
8
8 650M
Num. param.
447M
447M
Total num. param.
1.38B
1.75B
2.75B
B FACTORING OUT REDUNDANT DIMENSIONS: PCA-INSPIRED VARIANT
Recall that for this variant, before feeding x to the flow model, we reshape it into a sequence of flattened patches and apply learnable, invertible linear map W along the channel dimension.
16

---

Published as a conference paper at ICLR 2025
Table Hyper-parameter details for fine-tuning tasks.
Batch size
Learning rate
Weight decay
Task
Model Size
Epochs
Dropout
Sampler
COCO (FID-300)
B/M/L
10
256
le-4
le-5
0.1
CFG
VQAv2
B/M/L
10
128
3e-5
3e-6
0.0
Greedy
COCO caption
B/M
10
128
3e-5
3e-6
0.0
Greedy
COCO caption
256
3e-5
3e-6
0.0
Greedy
For this variant, the loss can be written as:
[(x) = loglogg(ff (()( + (x)) + log det (ff()) 0iT
+ (HW/p²) log det W
where PN is the Gaussian distribution. xWT is decomposed into two slices as [ i| =WT of shapes HW/p² x d and HW/p² x (3p² - d), respectively, where the original image has shape H W and is split it into p x patches which then flattened.
Intuitively, factoring out dimensions via W rather than at the flow output limits the complexity and hence potentially leads to worse results (see Table 2
When compting W via PCA to obtain the bottom two rows in Table we sample 000 images from the ImageNet training set, split them into patches, and compute PCA using the built-in implementa- tion from scikit-learn without whitening.
C ADDITIONAL ABLATIONS
VAE baseline We train VAE following the design of VQGAN (Rombach al. 2022) producing a sequence of 256 128-dimensional tokens like the normalizing flow in JetFormer. We remove the GAN and perceptual losses ensure fair comparison with JetFormer (which solely maximizes the data likelihood, without relying on GAN or perceptual losses). We then train the JetFormer-I BB decoder-only model on this VAË representation, following the JetFormer-B T&I training recipe. The results in the Table 7 below show that JetFormer outperforms the VAE baseline by solid margin (in particular in T21 generation).
Table 7: Comparison of JetFormer-B B trained for T&I with a 2-stage VAE-based variant.
COCO-FID30K
INet O-shot
COCO Cap
VQAv2
VAE + Transformer (2 stages)
95.4
37.6
103.9
64.0
VAE + Trans. stages, noise curr.)
87.6
40.0
106.7
64.9
JetFormer (end-to-end)
30.6
42.7
113.9
65.4
Effeet of noise curriculum on NLL Intuitively, the noise curriculum biases the training process towards those high-likelihood solutions with high perceptual quality, at the expense of a slight in- crease in NLL (see Sec. B.2.1 for more discussion). Table g shows that longer training with noise curriculum eliminates the gap in NLL compared to training without curriculum.
Table 8: Sample quality and NLL as a function of noise curriculum and training duration.
noise curr.
#epochs
FID-50k
NLL
No
100
44.71
3,05
Yes
100
7.84
3.14
Yes
500
7.25
3.06
Yes
1000
7.10
3.04
17

---

Published as a conference paper at ICLR 2025
Effect of final noise scale when fine-tuning for T21 generation We further visualize the interplay
between the final noise scale and NLL when fine-tuning T2I models on COCO in Fig.
17
3.8
3,7
(dashed)
FID (solld)
3.6
15
3.5
Jetformer-B Jetformer-L
13
Nolse scale
Figure 5: FID and NLL results obtained when fine-tuning JetFormer for T2I generation on COCO with different values for the final noise scale. JetFormer pretrained with a curriculum decaying the noise scale from 64 to 3 and the noise scale is then decayed during fine-tuning from 3 to one of {0, 1.3]. We observe that decaying the noise scale to 0 leads the best NLL but worst FID. The two model sizes have similar NLL at noise scale 3 but different FIDs, possibly because NLL is bounded by the noise scale but larger model scale helps to generalize better.
18

---

Published as a conference paper at ICLR 2025
D IMAGE AND LATENT MANIPULATION
In Figure we visualize the effect of different image or latent manipulations. Note that these images indicative of the final image generation quality and their purpose is illustrate certain quali- tative properties. We use four images: full 256 x56 view of a plane from ImageNet (Russakovsky al. 2015), a 64x64 x4 crop of an image of bird from ImageNet, x6x64 crop of scene from MS-COCO Lin et al.] 2014) to analyze out-of-distribution generalization (vs. ImageNet), and a 128 128 crop of an image with text from TextVQA (Singh al. 2019), visualize the effect of significantly different domain (vs. natural images). On those images we show:
(a) Ground truth.
(b) Additive RGB noise N(0, I) with scale og = 64. This corresponds to the images at the begin- ning of the noise curriculum. Although this significantly affects the perception of the zoomed in view of the bird or details of the clouds, has smaller impact on the more zoomed out features such the plane x56x255 or the large text letters.
(c) In Section 3.2 we discussed factoring out redundant dimensions. Here we use a model trained on ImageNet where the dimensions are factored out after the normalizing flow to map the images to the latent space. After we keep the 128 autoregressive dimensions fixed but resample the 640 dimensions modelled by the Gaussian prior. Despite these corresponding to 640 768 latent dimensions of each patch, they have little impact the overall image structure. It shows the model successfully maps important information to the 128 dimensions modeled with the autoregressive prior which are unmodified in this manipulation.
(d) Zeroing out the gaussian latents after a learnable projection. Overall the model learned to map the relevant dimensions into the 128 per patch modelled further by normalizing flow.
(e) PCA reconstruction of patches when using 128 dimensions. PCA is from ImageNet images.
(f) VAE encoding and decoding. We use VAE from (Tschannen et al. 2024) which was pre- trained on ImageNet with perceptual losses, overall produces sharp images but looses signif- icant detail and is not able to preserve details in outto--distribution images, egg. with text.
1819.
1819
1819.
1819.
1819.
1810.
REROS del currilo
And
del caemita
del cheritt
frt
mode podforac
mode an
aree as padiondd
tahare
beyetilla
bayrtlll
beyrtill
bayetilla
Lase pwc
(a)
(b)
(c)
(d)
(e)
())
Figure 6: Visualization of image and latent manipulations. Note these images are not indicative of the final image generation quality. Best viewed in digital format.
19

---

Published as a conference paper at ICLR 2025
E ZERO-SHOT MS-COCO SAMPLES
body waber bear
dde
plate broccal sercsalandoe cance ota 1atle
tco siiting nea the window loccng suter bricculdhpg
large einphart feeding bret uning In muni
fhe tydnat gutting cater ento simett
antier white dothing attempt na
cleck trave It miles soy.
Figure 7: Selected zero-shot samples from JetFormer-L. (T2I) for captions from MS-COCO.
20

---

Published as a conference paper at ICLR 2025
F IMAGENET SAMPLES
Figure 8: Random, non-cherry picked samples from JetFormer-L for selected ImageNet classes: flamingo, bustard, golden retriever, hummingbird, bell pepper, volcano and ring-tailed lemur.
21

---

Published as a conference paper at ICLR 2025
Narminn 13
Figure 9: Random, non-cherry picked samples from JetFormer-L for random ImageNet classes.
22


Reviews:
[{"review_id": "xAuTV7RAzO", "reviewer": "Reviewer_9Pqq", "summary": "This paper proposes JetFormer, an end-to-end autoregressive transformer designed for multimodal generation of text and images. JetFormer utilizes a normalizing flow to encode images into soft tokens, eliminating the need for pretrained encoders and enabling direct learning from raw data. The model’s unified structure for text and image processing enhances its multimodal adaptability, with additional techniques like a noise curriculum to improve image quality. JetFormer achieves competitive performance with state-of-the-art VAE-based models in tasks such as text-to-image generation and image captioning.", "questions": "(suggestion) To validate the advantages of JetFormer’s end-to-end learned image embeddings, additional experiments are needed to compare its performance against models using well-established VAE-based embeddings. Such comparative analysis would clarify whether JetFormer’s learned embeddings offer benefits over pre-trained embeddings in terms of image generation quality and task-specific adaptability.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. JetFormer processes text and images within a single autoregressive transformer, removing the need for modality-specific encoders and supporting seamless multimodal generation tasks.\n2. JetFormer’s fully end-to-end training from raw data enables it to potentially learn task-specific representations, enhancing adaptability without relying on pre-trained embeddings.", "weaknesses": "1. The claim that pre-trained modality-specific encoders (e.g., VQ-VAE) limit performance due to task-agnostic design lacks sufficient quantitative or qualitative evidence to substantiate this limitation convincingly.\n2. While end-to-end training in JetFormer may offer the advantage of learning task-specific latent representations, this approach could also introduce challenges in terms of training stability and computational cost. The paper does not provide a comparative analysis between end-to-end and two-stage methods (e.g., VAE-based embeddings followed by generative model training), leaving it unclear whether the benefits of end-to-end learning outweigh the potential costs and stability issues.\n3. JetFormer’s approach of combining text and image modalities into a unified embedding space lacks clear justification. Without evidence of improved multi-task performance or multimodal synergy, it is unclear if a shared embedding space is necessary or beneficial for this model’s architecture.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "vUOjGggINv", "reviewer": "Reviewer_jKRf", "summary": "The paper titled \"JetFormer: An autoregressive generative model of raw images and text\" presents a novel system designed for the understanding and generation of image-text data, with a particular focus on its performance on academic datasets. The model, JetFormer, is part of a broader category of large multimodal models trained on web data, and it shares similar ethical considerations, such as the need for de-biasing, safety-tuning, and the use of content filters to ensure safe deployment.\n\nJetFormer is characterized by its end-to-end training approach, which, while computationally demanding, offers the potential for scalability and improved performance. However, the model currently faces limitations in visual quality compared to state-of-the-art diffusion models that utilize pretrained latent representations. Additionally, the computational requirements are significant, especially when dealing with high-resolution images, which necessitates efficient handling of image data, such as modeling at the patch level.\n\nThe paper also discusses a PCA-inspired variant to reduce redundant dimensions, which involves reshaping input data into sequences of flattened patches and applying a learnable, invertible linear map. This approach aims to manage the complexity of the data and improve the model's efficiency.\n\nOverall, the contributions of the paper include the development of JetFormer, an exploration of its capabilities and limitations, and the proposal of methods to enhance its performance and safety in practical applications.", "questions": "Please refer to the weaknesses.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Originality: The paper introduces innovative techniques to enhance the quality of image generation models, particularly through the use of a novel noise curriculum and the factoring out of latent dimensions. The approach of factoring out redundant dimensions using a PCA-inspired variant before applying a flow model is a creative solution to improve model efficiency and performance. This originality is further demonstrated by the introduction of classifier-free guidance during sampling, which adds a unique aspect to the model's design.\n\nQuality: The quality of the research is evident in the thorough experimentation and evaluation conducted. The paper employs robust evaluation metrics, such as the ADM FID evaluation suite and precision/recall metrics, to assess image sample quality on ImageNet1k. Additionally, the use of MS-COCO FID-30k for text-to-image tasks and the reporting of CIDEr scores for image captioning and VQAv2 accuracy for visual question answering further demonstrate the comprehensive nature of the quality assessment. The results indicate that the proposed enhancements significantly improve the quality of generated images.\n\nClarity: The paper is well-structured and clearly articulates the methodologies and findings. The explanation of the process of factoring out redundant dimensions and the subsequent impact on model performance is detailed and easy to follow. The use of visual aids or diagrams, if included, would further enhance the clarity of the presentation, but the textual explanation alone provides a solid understanding of the research contributions.\n\nSignificance: The significance of the paper lies in its potential impact on the field of image generation and understanding. By addressing the computational challenges associated with modeling images at the patch level and proposing solutions that reduce computational burden while maintaining or improving image quality, the research offers valuable insights and tools for future developments in the field. The ability to improve model quality for natural images and the implications for tasks such as zero-shot classification, image captioning, and visual question answering highlight the broad applicability and importance of the work.", "weaknesses": "Model Performance and Design Choices: The paper highlights several design choices that impact model performance, such as the use of dropout and PCA transforms. It notes that modeling images after PCA leads to worse results and that omitting noise curriculum results in significantly poorer outcomes. These observations suggest that the model's performance is sensitive to specific preprocessing and design choices, which may limit its robustness and generalizability.\n\nFactoring Dimensions: The paper discusses the impact of factoring out dimensions at different stages of the model. It finds that not factoring out dimensions after a flow model leads to quality loss, and while factoring out dimensions with a learnable linear mapping before the flow improves results, it still underperforms compared to post-flow factoring. This indicates a potential area for improvement in the model's architecture to optimize performance.\n\nLack of Comprehensive Evaluation: While the paper reports metrics based on the median of 10 runs and test-dev accuracy for VQAv2, it does not provide a thorough evaluation across diverse datasets or scenarios. A more extensive evaluation would help in understanding the model's strengths and limitations better.\n\nAddressing these weaknesses would enhance the paper's contribution to the field and provide a clearer understanding of the model's capabilities and limitations.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "edA2ZZ8U7X", "reviewer": "Reviewer_Lv3i", "summary": "The paper proposes an end-to-end autoregressive model that can jointly understand and generate both image and text. The model consists of two components. A normalizing flow is applied to convert images into soft tokens for autoregressive modeling. The main backbone is a decoder-only transformer that predicts the next soft token. The paper explores two effective approaches to help model convergence. One is to introduce an evolving noise level during the training schedule. Another one is to dropout the redundant dimensions within the soft token when modeling the distribution auto-regressively. Experiments are conducted to evaluate the proposed method both quantitatively and qualitatively.", "questions": "1. About information leakage. Since the image tokenizer is also trainable, it seems possible for the image tokenizer to attend to the next input token that the auto-regressive transformer should predict, i.e. the information of the next token is leaked into the input tokens, causing a trivial solution. Just wonder why this won't happen in Jetformer?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "1. The paper is well-writen and is easy to follow.\n2. The proposed method is innovative, showing a promising approach towards end-to-end image-text joint probabilistic modeling. It shows that it is possible to directly optimize NLL for potentially any modality.\n3. The paper identifies two crucial training techniques to help the model converge better, which are crucial for the method's effectiveness.\n4. Experiments are conducted sufficiently. It also shows the proposed method can potentially benefit from scaling up the model size.", "weaknesses": "1. In figure 3 of the paper, it seems that the noise curriculum makes the final NLL get slightly higher than without using the noise curriculum, but experimentally, the model clearly benefits from the noise curriculum training technique. Does this indicate that NLL may not be the best objective for image-text generation task? A more theoretical explanation is helpful for the readers to understand this phenomenon.\n2. The detailed architecture of the normalizing flow model is not clearly explained. Better make it clear by showing diagrams of the model blocks and calculations.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "zaVnGOybOu", "reviewer": "Reviewer_NBCt", "summary": "This paper propose a generative model for generating text and images, it propose to use a normalizing flow model to simultaneously encode and decode images. The flow model generates soft token image representation that is jointly trained with an autoregressive decoder only transformer architecture. \n\nThe main contribution of the paper is to first explore the usage of normalizing flow models as image tokenizers that can achieve an end-to-end training without the need to use a pretrained autoencoder for images.", "questions": "- What do you think it is the main practical advantage of removing the pretrained image tokenizer for autoregressive image modeling?\n\n- Recent papers as Transfusion [3] shows strong results by training with mixed modal corpus of data, a discussion on how this approach compare to your method would be beneficial. \n\n- Recent developments in image generation and deep learning guided compression shows that using a pretrained autoencoder to compress and reduce image dimensionality is really effective, moreover decoupling image tokenization from the generative training strategy can have some benefits in terms of training dynamics. What is your perspective on these considerations? \n\n\n[3] Transfusion: Predict the Next Token and Diffuse Images with One Multi-Modal Model. Zhou et al", "limitations": "", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- The paper is well written and easy to follow.\n\n- The exploration of a normalizing flow model in vision language pretraining setting is an original and unexplored idea.\n\n- The proposed method eliminate the needs to rely on pretrained components, by training the whole system end-to-end, the use of soft-tokens with normalizing flow is needed to avoid the mode collapse typical of VAEs.", "weaknesses": "- The central claims of the paper are not supported by experimental results: The model performance on both text-to-image taskt is worse with respect to older and smaller models based on pretrained VQ-VAE (e.g MAGVLT [1]), while the improvement on understanding tasks is slight, training for image understanding greatly reduces the performance on image generation. \n\n- The paper does not compare with recent baselines such as VAR models [2] for simple settings like ImageNet256 where VAR achieve 1.92 FID and Jetformer has 6.64 using the same training data and similar model size.\n\n- Optimizing image tokens during training can add instability to the training dynamics, this is especially true in complex large scale multimodal settings. \n\n- The main comparison to validate the usage of this modeling approach with respect to use a VQ-VAE in a controlled setting would be to compare it against a baseline, trained in the exact same setting  by substituting the normalizing flow with a VQ-VAE.\n\n- How the approach scales is not clearly analized in the paper by just showing to models scales with relative low difference in parameters count\n\n[1]: MAGVLT: Masked Generative Vision-and-Language Transformer\n\n[2]: Visual Autoregressive Modeling: Scalable Image Generation via Next-Scale Prediction, Tian et al.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)