Jet: A Modern Transformer-Based Normalizing Flow
Authors: Alexander Kolesnikov
André Susano Pinto
Michael Tschannen
Venue: Accepted by TMLR ()

Abstract: In the past, normalizing generative flows have emerged as a promising class of generative models for natural images. This type of model has many modeling advantages: the ability to efficiently compute log-likelihood of the input data, fast generation, and simple overall structure. Normalizing flows remained a topic of active research but later fell out of favor, as visual quality of the samples was not competitive with other model classes, such as GANs, VQ-VAE-based approaches or diffusion models. In this paper we revisit the design of coupling-based normalizing flow models by carefully ablating prior design choices and using computational blocks based on the Vision Transformer architecture, not convolutional neural networks. As a result, we achieve a much simpler architecture that matches existing normalizing flow models and improves over them when paired with pretraining. While the overall visual quality is still behind the current state-of-the-art models, we argue that strong normalizing flow models can help advancing the research frontier by serving as building components of more powerful generative models.

Published in Transactions on Machine Learning Research (03/2025)
Jet: A Modern Transformer-Based y Flow
Alexander Kolesnikov* Google DeepMind
akolesndkonn@ggenle.coom
André Susano Pinto' Google DeedMind
undresp@googlle
Michacl Tschannen' Google Deed
tschamnen httpoeneeenooogge.ccoomo com
Reviewed on OpenReview:
https: // openreview. net/ forum? id= jdumaks 7ZY
Abstract
In the past, normalizing generative flows have emerged as promising class of generative models for natural images. This type of model has many modeling advantages: the ability to efficiently compute log-likelihood of the input data, fast generation, and simple overall structure. Normalizing flows remained a topic of active research but later fell out of favor, as visual quality of the samples was not competitive with other model classes, such as GANs, VQ-VAE-basee based approaches or diffusion models. In this paper we revisit the design of coupling-ban based normalizing flow models by carcfully ablating prior design choices and using computational blocks based on the Vision Transformer architecture, not convolutional neural networks. As result, achieve much simpler architecture that matches existing normalizing flow models and improves over them when paired with pretraining. While the overall visual quality is still behind the current state-of-tho-are models, we argue that strong normalixing flow models can help advancing the research frontier by serving as building components more powerful generative models.
1 Introduction
In this paper we explicitly do not attempt to devise the new state-of-the art image modeling approach or propose a new paradigm. Instead, we revisit the long known but recently neglected class of models for generalive modeling: coupling- based normalizing flows. Normalizing flows have important capabilities that make them a useful tool for modern generative modeling.
At high-level, normalizing flow model learns bijective (and thus invertible) mapping 9 from the input space to the latent space, where the latent space follows simple distribution, e.g. Gaussian distribution. A complex bijective transformation can be constructed by stacking multiple coupling blocks, which are bijective and invertible in closed form by design and are parametrized by deep neural networks.
Normalizing flow models can be directly trained by computing data log-likelihood in the simple (e.g. Gaus- sian) latent space after the learnable and diffcrentiable mapping g is applicd transform training examples. For data generation, the inverse transformation g is readily available, which can be used map easy-to- sample Gaussian latent space to the samples from the target distribution. The two explicit, differentiable and lossless mappings and gー1 can be used as building blocks for more complex generative systems. For example, normalizing flows are used as a critical component of more complex systems: recent examples include CChen et al. 2016 Kingma et al. 2016 Tschannen et al.j 20244) and (Tschannen el al. 2024b) that leverage normalizing flows to facilitate image modeling with VAEs and autoregressive transformers respectively. This motivates us to revisit the normalizing flow model class.
"Equal contribution.

---

Published in Transactions on Machine Learning Research (03/2025)
Image palches
Mergo
dimentoos
Repeat times
Figure 1: Overview of the Jet model. The dashed box contains coupling layer computing an affine transform from one half of the input dimensions (patches or features) and then applying it to the other half of the input dimensions. The full model is obtained by stacking N such invertible coupling layers.
In the mid-to-late 2010s normalizing flow models were a topic of active research. NICE DDnn et al. 2014) was the early normalizing flow model for images. It introduced the main building block behind normalizing flow models: an additive coupling layer. RealNVP (Dinh et al., 2017)) improved over NICE by introducing multiscale architeclure and affine coupling layers that additionally perform scaling transformation. Subse- quently, in Glow [Kingma & Dhariwal 2018), authors introduce two more components:
invertible dense from using dequantization flow trick and generalized variant of the affine coupling block.
layer and specialized activation normalization layer. Finally, Flow++ (Ho et al. 2019) shows improvements
In this paper we mainly concentrate on revisiting optimal design for the normalizing flow models. We focus on both performance and simplicity of the final model. We built on top of the prior literature and put all components under careful scrutiny. Our final model has a radically simpler design and only relies on the plain affine coupling blocks parametrized by the Vision Transformer model. Our key contributions can be summarized as follows:
We use Vision Transformer building blocks instead of convolutional neural networks, which leads to significant performance improvement.
We drastically simplify the overall architecture by eliminating many components from prior models:
No multiscale components and early factored-out channels
No invertible dense layers
No "activation normalization" layers
No dequantization flow
No generalized coupling transformation
We show the proposed architccture matches SOTA results in terms of negative log-likelihood common image benchmarks when trained on the same data-limited datasets.
We demonstrate that pretraining Jet model on large corpus of natural images is highly effective and simple, and that it surpasses prior state-of-the-arl results.
2 Method
In the section we introduce the Jet model. We first introduce the architecture, then describe training procedure and important implementation details.
2

---

Published in Transactions on Machine Learning Research (03/2025)
4.3
3.90
VIT
VIT depth
4.2
3.85
Nogative log-likelihood
CNN
Negalive log-likelihood
4.1
3.80
2
3.75
4.0
3.70
3.65
3.8
3.7
3.55
3.6
3.50
10'
10
10
10
10' TPU Core-hours
10
TPU Core-hours
(a)
(b)
Figure 2: Effect of different architecture design choices the validation NLL (in bits per dimension), function of training compute. Figure 2a Results on ImageNet-1k 64 x 64 for CNN vs VIT blocks (the marker size is proportional to the model parameter count). VIT blocks clearly outperform CNN blocks for a given training compute budget. Figure 25 Results on ImageNet-21k 32 x 32 for different ViT depths. Increasing the block depth leads to improved results up to depth
2.1
Jet model
The Jet model has a very simple high-level structure. First, the input image is split into K flat patches. flattened into a sequence of vectors, and then we repeatedly apply affine coupling layers (Dinh et al, 2017) We illustrate the Jet model architccture Figure 1
The input to a coupling layer is x E Rx2d where K is the total number of patches and 2d is the number of dimensions within each patch after flattening. We then apply dimension splitting transformation. producing 21,42 E kxd The first part, is not modified, yielding V1- The sccond part, is modified with element-wise affine transformation V/2 (x, + b) where s, b E Kxd Scaling factor s and bias are functions of Φ1- We use a single Vision Transformer (ViT) backbone, f Rxxd - Kx2d which maps Dj to scale and bias, effectively implementing the functions s(x1) and 6(x1). Finally, we merge y1 and V2 by applying the inverse of the splitting transformation that was applied x, and arrive at the output
y Kx2d
To stabilize training additionally apply sigmoid transformation, a,, that squashes the scale coefficient to (0, 1) range and then multiply it by fixed constant m to extend the range to (0, m). In practice we always set m = 2, which allows to amplify or suppress the intermediate activations, while keeping activations in a reasonable range when stacking many coupling layers.
Overall, single coupling block can be formalized follows:
V1 =
1/2 = (x2 + b(x1)) (a(21)) m
As we show below in Section 2.3 this transformation can be inverted for arbitrary functions b(x1) and s(x1).
2.2 Training and optimization objective.
We now recall the fundamentals of the normalizing flow model training. The key assumption behind the normalizing flows that the target density p(x) can be modeled as a simple distribution, (x(2) (e.g. standard Gaussian) after using bijective transformation g: + transform the original data. Applying the "change of variable" identity from the basic probability calculus obtain tractable model for the
3

---

Published in Transactions on Machine Learning Research (03/2025)
probability density function:
0g(x)) 0x
p(x) = Px(((()) det
We implement function with the Jet model described above. The determinant of the derivative, 0g(a) dr is equal to the product of determinant of the derivatives of individual coupling layers. And for each coupling block the determinant computation is trivial and equal to the product of all scaling factors in the affine transform. See (Dinh et al., 2017) for details on the derivation. In practice we maximize data log-likelihood, the optimization objective is
2 Σ 19(x); 1og((2)) + ΣΣ llo (s}) logm)
where we use index i to iterate over all dimensions of an array inside the sum, and upper index & to iterate over total number of blocks. The first sum term arises from taking the logarithm of the standard Gaussian density, applied to every output dimension. The second sum term the log determinant, which is computed as sum of logarithms of all scale values across all layers, plus logarithm of the fixed m multiplier.
In practice normalize the above objective by the total number of dimensions and we keep all constant terms. Even though it does not affect optimization, by doing this we obtain log-likelihood estimate, which can be interpreted as "bit-per-dimension" (bpd) (assuming logarithm base of 2). For example, if the objective value is 3.1, it means that the average uncertainty per input dimension is 3.1 bits. A uniformly random image model should yield approximately bpd. This is useful when comparing models across different modeling classes and for various correctness checks.
Note that correctly model discrete densities (e.g. images where pixels are discrete-valued) we employ standard dequantization procedure described in (Theis al. 2015). In practice, it mcans that add random [0, 1] uniforın noise to inpul images, which consist of discrete values in the range [0, 255]. We perform this dequantization procedure both during training and evaluation.
2.3 Inverse transformation and image generation
It is easy to obtain the inverse of the above transformation in closed form. Notably, the VIT function that computes bias and scale terms does not need to be inverted:
T₁ 1/1
1/2
b(x1)
T2 o (s(xr))
The computational complexity of computing the inverse exactly the same as computing the normalizing flow itself. New images can be sampled by first sampling the target density (i.e. Gaussian noise) and then applying the inverse transformation.
2.4 Initialization
Careful initialization is essential for training deep normalizing flow model with large number of coupling blocks. We employ a simple yet very effective initialization scheme. The final lincar projection of the ViT f is initialized with zero weights. As a result, predicted bias values, b(x1) are 0. The scale values are equal to (0) = 0.5. When we set m = 2, then the all scaling factors become equal to 1.
As result, the Jet model behaves identity function at initialization. Empirically, we find this sufficient to ensure stable optimization at the beginning of training. Consequently, we do not need add " "ActNorm" layer that is commonly used in the normalizing flow literature to achieve similar effect.
4

---

Published in Transactions on Machine Learning Research (03/2025)
2.5 Dimension splitting
We explore various oplions for channel splitting. One option is to perform channel-wise split, by splitting the channels of each image patch into the two equal groups (Dinh et al. 2014) Kingma & Dhariwal 2018). The splitting is random within each coupling layer and fixed ahead of time (independently for cach layer) This is a simple-to aaplement-strategy that ensures diverse channel mixing.
We also implement various splitting strategies facilitate spatial mixing. To this end, we explore 3 strategies that respect images' 2D topology: row-wise alternating patch splitting, column-wise alternating patch- splitting, and the "checkerboard" splitting (Dinh et al. 2017). See Section 3 for the empirical investigation of performance for various design choices of splitting operations.
Splitting implementation details. The natural way to perform channel splitting is to use array indexing operations. However, indexing can be slower than matrix multiplication on modern accelerators. Thus, we implement channel splitting as matrix multiplication with precomputed 0 and (frozen) weight matrices. A nice by-product of this approach is that dimension merging (inverse of splitting) can be trivially implemented as matrix multiplications with the transposed weight matrices.
Numerical precision considerations. Matrix multiplications on modern accelerators often implicitly uses half-precision for multiplying individual values, while accumulating the result using full float32 precision. We note that such loss in precision may lead to numerical issues, as, for example, the uniform [o, 1] de- quantization noise loses most of its entropy and leads to overly optimistic log-liiel estimates. To avoid this, we enable full precision matrix multiplication mode when aplittnng/merging the dimensions. Note that we still use fast default settings when computing the ViT function inside coupling layers. So the overall efficiency drop from enabling full precision for aplitting/mergingg is negligible overall.
3 Experiments
Throughout the paper, we keep our experimental setup simple and unified. Additionally, the code is available in the big_ visvisson codebasd] For the ViT architecture, we follow the original paper Dosovitskiy et al. (2021). Our VIT models inside the coupling layers do not. have initial patchification or final pooling and linear projections. Unless stated otherwise, we set the patch size such that the total mumber of patches is equal to 256.
For the optimizer we use AdamW (Loshchilov al. 2017 We set second momentum 32 parameter to 0.95 to stabilize training. We cosine learning rate decay schedule.
Datasets. We perform experiments on three datasets: Imagenet-1k, 1k, Imagenet-214, 221 and CIFAR-10, across two input resolutions: 32 x 32 and 64 x 64 (except for C-F) To downsample Imagenet- images we follow the standard protocol (Van den Oord et al. 2016b to ensure a correct comparison to the prior art (i.e. we use the preprocessed data provided by (Van den Oord et al. 2016b) where available). To downsample Imagenet-21L. images we use TensorFlow resize operation with method set to AREA For CIFAR-10 we use the original dataset resolution. Importantly, to make sure our results are comparable to the literature, we do not perform any data augmentations. For comparison with Denseflow (rriié et al. 2021 which uses different protocol see Appendix A.2
3.1
Main results
We conduct extensive sweep that. includes Jet. models trained across varying computational capacity, data size (ImageNet-Ik Ik and ImageNet-21k) and resolutions (32x x 32 and 64 64). For the model capacity sweep we explore the following configurations (approximately spanning 2 orders of magnitude in compute intensity):
number of coupling layers in (16,32,66 32,64) for ImageNet- and in {32, 64, 128) for ImageNet-21k,
the depth of ViT used inside the coupling layer in (1,2,3, 6,
5

---

Published in Transactions on Machine Learning Research (03/2025)
ImageNet- Ik 32x32
ImageNet- IK 64x64
ImageNet-21k 32x32
ImageNet-2IK 64x64
Train
3.75
4.2
3.9
3,70
3.8
T1N
3.65
3.7
3.60
3.2
3.8
3.6
3.55
3.6
3.50
3.1
10' 10
10'
10'
10
10'
TPU Core-hours
Figure NLL as a function of training compute obtained when training Jet architectures with a range of architecture pyper-paramaters, for different data sets. The size of each marker proportional to the number of parameters in the model configuration. Overall observe normalizing flow models benefit from scale, yet ImageNet-lk models start to overfit. When increasing the amount of data to ImageNet-221 size, we observe little overfitting and strong scaling trends.
Table 1: Negative loo-lokoliilleee likelihood of Jet models and the state-of thettrttttrtttte coupling- based normalizing flow model Flow++ (Ho et al. 2019) from the literature. Jet (21k) indicates the result obtained when using a Jet model pretrained on ImageNet21k. Overall, we observe that all tasks benefit from this pretraining and that Jet matches outperforms +low+++.
(a) ImageNet-lk 64 x 64
(b) ImageNet-1k 32 32
(c) CIFAR-1O 32 x 32
Model
Result (NLL4)
Model
Result (NLLJ)
Model
Result (NLL))
+low++ 3.69
+low++ 3.86
Flow
3.08
Jet
3.656
Jel
3.931
Jet (121k)
3.018
Jet (121k)
3.580
Jet (121k)
3.857
the ViT embedding dimension in {256, 512, 768] for ImageNet-lk and in {512, 768, 1024] for ImageNet-225 The number of heads tied to the embedding dimension and equal to {4, 8, 12] and {8, 12, 16) respectively.
We use a fixed standard learning rate of 3c-4, weight decay of 1c-5 and train for 200 epochs for ImageNet- and for 50 epochs for ImageNet-21k. We additionally investigate transfer learning setup and finetune our best ImageNet-2lk models on ImageNet-Ik Ik and CIFAR-10. Full sweep results are shown in Figure 3 Additionally, we present key numerical results in Table [
Our first observation is that due to the high expressive power of the Jet. model parameterized by VIT model, it tends to quickly overfit on ImageNet-lk, which has only 1.2M examples in total. However, despite this, for 64 x 64 Jet attains state-of-the-art NLL of 3.66 bpd. For 32 x 32 input resolution overfit for large models more severe, however reasonably small models still achieve a competitive result of 3.93 bpd.
It is important note that unlabeled natural images are abundant. Thus, the best way tackle overfitting is to increase the amount of training data, opposed to constraining or regularizing the model. We, therefore, use the Imagenet-22 22k dataset with more than 10x more images than in ImageNet-lk. As expected, overfitting is tamed, with larger models leading to increasingly better results.
The next important question is whether a model trained on a larger and more class-diverse ImageNet-21k dataset transfers to ImageNet-lke We observe that with a very light finetuning (30 epochs, learning rate of 1c-5 and 3e- 6 for higher resolution) we obtain statc-of-the-art results on ImageNet- attaining 3.58 and 3.86 bpd on 32 32 and 64 x 64 input resolution, respectively. However, one can argue that ImageNet- 1k is very similar to ImageNet-21k and results are overly optimistic. To address this concern, also transfer the

---

Published in Transactions on Machine Learning Research (03/2025)
3.90
Spatial couplings
Negative log. le-liikhhhood
3.85
row-wise
column-wise
checkerboard
3.80
mix of all
none
3.75
3.70
all spatial
all channel
1:1
2:1
3:1
4:1
5:1
Ratio of channel-wise to spatial-wist -wiee couplings
Figure 4: Ablation of coupling types. Negative log-likelite log-likelhhood on ImageNet-lk 1k 64x64 when varying the ratio of channel-wise to spatial-wise couplings and when using different types of spatial-wisc couplings. Results in table format in Appendix Table 10
ImageNet-21k model to the CIFAR-10 dataset, which has very different type of images and has different downsampling procedure. Again, after finetuning (100 epochs, learning rate of 3e-6) we obtain state-of-the- art performance of 3.02 bpd.
Overall, our results provide strong evidence that the Jet architecture achieves state-of-the art performance in the class of coupling-basel normalizing flow models, and, at the same time, it is very stable and easy to train across wide range of scenarios.
3.2 Ablations
In this section we ablate key design choices for the Jet model. Our default ablation setting is moderately- sized Jet model, trained ImageNet-1k 64 64 for 200 epochs. We set the total depth to 32 coupling layers, the VIT depth to 2 blocks and width to 512 dimensions. In our scaling study this configuration reasonably close to the optimal setup, while being sulficiently fast for the extensive ablations sweeps. We use negative log-likelihook reported in bpd as the main ablation metric.
3.2.1 Coupling types
As presented in Section 2.5 consider 4 different ways to split the channels in coupling block: is channel-wise method (splits the channels into two parts) and three patch-wise wise methods (alternating rows, alternating columns "checkerboard" splitting). Inspired by the prior litcrature, where channel-wiso couplings were more prominent, we explore the following design space: M repeated channel-wise couplings followed by 1 spatial coupling. The spatial coupling is either fixed to one of the methods, or alternates between 3 types. We vary M from Overall, the above choice space gives rise configurations.
Figure presents the quantitative results, suggesting the following insights:
Using interleaved spatial and channel coupling is optimal, where M is optimal, though M E 2..5 achieves very similar performance.
Alternating spatial coupling methods are superior over using any single coupling lypes.
Spatial coupling alone is the worst configuration, while exclusive channel coupling is competitive.
3.2.2 Coupling layers vs ViT depth
Jet has two depth parameters: number of coupling layers and ViT depth within each coupling layer. The total compute is roughly proportional to the product of these two depths. We observe very interesting
7

---

Published in Transactions on Machine Learning Research (03/2025)
Table 2: Ablation of design choices reporting negative log-likelihood (bpd) on ImageNet- 64 64. Table 2a
shows the effect of using masking- modie or pairing-man morin (best) splitting the spatial tokens. Table 2b
shows the effect of using invertible dense layers and/or activation norm. Best results obtained when not using either.
(a)
(b)
Couplings ratio
Masking
Pairing
Activation Norm
Invertible Dense
NLL
all spatial
3.844
3.787
1:1
3.741
3.737
x
3.720
2:1
3.727
3.725
V
X
3.727
3:1
3.723
3.723
X
V
3.741
4:1
3.722
3.719
V
V
3.733
5:1
3.722
3.720
interplay between these parameters, see Figure representing ImageNet-21k 32 32, which is a detailed view of Figure
Specifically, observe that scaling the number of coupling layers, while keeping shallow VII models (e.g. depth 1) results in unfavorable compute-porformance tradc-off. It appears that ViT depth of at least 4 is the necessary condition for the Jel model to stay close to the frontier. For example, a model with 32 coupling layers and ViT depth has roughly the same compute requirements as model with 128 coupling layers and ViT depth 1. However, the former performs much better than the latter for fixed compute.
3.2.3 ViT vs CNNs
To ablate the use VIT instead of CNN block, conduct a similar sweep to our main sweep on ImageNet-
1k 64 x 64 but using a CNN architecture (specifically we use the CNN architecture from (Kolesnikov et al. 2020 This time sweeping the following scttings for the CNN sctup: model depth in {16, 32, 64), CNN block depth in (I.....8), 8), block embedding dimension in {256, 512, 768, 1024, 1536). Block dimension 1536 was not used for model depth 64 due significant memory costs. The results in Figure 2a show that the CNN-based variant lags significantly behind the ViT-based one.
We anticipate that this gap be reduced by using multiscale architectures, commonly done in the literature (Dinh et al. 2017 Kingma & Dhariwal 2018). However, in this paper we strive to simplify design and exclude multiscale archilectures.
3.2.4 Coupling implementation
We investigate two common approaches for implementing coupling layers: masking or pairing. To be concrete, let's assume that we implement coupling layer that splits the input spatially into two groups of patches. In one approach, which we name "masking" mode, we feed the K patches to the VIT block but mask with zeros the ones corresponding to the group. At the output we use only the ones corresponding to the T2 group and ignore the output of the patches corresponding to the æ group. One potential issue with this method that it weakens the residual connections as the tokens from which we predict the output are zero tokens.
Another approach we consider is a "pairing" mode in which we establish pairing between input and output patches (or embeddings). For example when using vertical-stripes pattern, the outputs of the ViT block for patch in the I1 group will predict the scale and bias for patch in the 22 group (e.g. to the patch below). This would make the VIT block processes only K/2 patches.
We experiment with these two implementation types while sweeping the M:1 channelsspatial coupling ratios as in Scction 3.2.T] The results presented in Table [a indicate pairing be superior, though the impact becomes smaller as one increases the number of channel couplings which do not depend on this design decision. We observe that both methods perform very similarly, with the pairing being slightly ahead of

---

Published in Transactions on Machine Learning Research (03/2025)
masking, especially for the scenario when only spatial couplings are used. Thus, for the Jet model, we default to using pairing mode.
3.2.5 Invertible dense layers and activation normalization
Glow Kingmma & Dhariwal 2018) introduces two components to improve the performance of normalizing flows: (1) a learnable, invertible dense layer which replaces the fixed permutation used to split the channels for each coupling; (2) an activation normalization layer with scalar and bias parameters per channel similar to batch normalization.
We use neither of them in Jet, but ablate whether performance could be improved by introducing those in Table 2b As a note, we observed that the use of activation normalization alone to be highly unstable, we sweep additional learning rate and seeds and report the best result found. Overall we found that not using any of those components leads to the best results.
3.2.6 Uniform dequantization vs dequantization flow
+ (Ho et a 2019) introduces a variational dequantization scheme to normalizing flows. Concretely, it proposes to replace the uniform dequantization noise added to the input with an image-conditional, learned noise distribution modeled by another normalizing flow. We ablate this component by training a 64 layer Jet model with 16-layer doquantization flow and compare with an 80-layer base Jet model. The image conditioning of the dequantization flow was implemented by adding cross-altention layers to the ViT-blocks to the input. We observe no significant improvements when using the dequantization flow component.
4 Related work
NICE (Dinh et al. 2014) popularized coupling-basee based normalizing flows with the introduction of the additive coupling layer. RealNVP (Dinh et al., 2017) then increased the flow's expressivity by using affine coupling layers in combination with multiscale architecture, and (Kingma & Dhariwal 2018 Hoogeboom et al. 2019 Sukthanker et al. 2022 proposed additional specialized invertible layers for image modeling. Flow= (Ho el al.) 2019) demonstrated improvements from learning the dequant ization noise distribution along with the flow model.
Another class of likelihood-basee generative models are autorcgressive models which flatten the (sub) pixels of
an image into sequence. Auloregressive modeling is enabled by using CNNs (Van den Oord el. al.) 201 6b
Salimans et al. 2016) or transformers (Parmar et al. 2018 Chen et al. 2020 Kolesnikov & Lampert (2017): Menick & Kalchbrenner (2019) improved performance of autoregressive models with hicrarchical modeling (e.g. over color depth or resolution). While obtaining better results than normalizing flows, autoregressive models are also much slower and do not scale to large resolutions as they require a forward-pass per (ixel.
In the context of normalizing flows, autoregressive dependency patterns between latent variables are a popular approach to improve modeling capabilities of normalizing flows (Kingma et al.] 2016 Papamakarios et al., 2017 Huang et al. 2018). Bhattacharyya et al. (2020) combined autoregressive modeling with a multiscale architecture. Concurrently to this work, Zhai et al. (2024) proposed combination of the transformer-basse autoregressive flow.
5
Conclusion
The Jet model revisits normalizing flows with focus simplicity and performance. While eliminating complex components such as multiscale architectures and invertible layers, Jet matches existing normaliz- ing flow models. The architecture with its regular attention mechanism showed significant improvements when paired with transfer learning and establishes new SOTA results, highlighting the need to study these architectures in data-abundant regime which is common for unsupervised methods.
We see normalizing flows, and Jet in particular, as a useful tool for advancing generative modeling. Due its simple structure and lossless guarantees, it can serve as building block for powerful generative systems.
9

---

Published in Transactions on Machine Learning Research (03/2025)
One recent example is Tschannen al. (2024b), which leverages normalizing flow to enable end-to-end autoregressive modeling of raw high-resolution images. We anticipate more progress in this area and believe that the Jet model will prove itself a powerful normalizing flow component that can be used out-of-the- for a variety of applications.
Acknowledgments
We thank Fabian Mentzer for discussions and feedback on this paper.
References
Apratim Bhattacharyya, Shweta Mahajan, Mario Fritz, Bernt Schiele, and Stefan Roth. Normalizing flows with multi-scale autoregressive priors. In CVPR, 2020.
Mark Chen, Alec Radford, Rewon Child, Jeffrey Wu, Heewoo Jun, David Luan, and Ilya Sutskever. Gener- ative pretraining from pixels. In ICML, 2020.
Chen, Diederik P Kingma, Tim Salimans, Yan Duan, Prafulla Dhariwal, John Schulman, Llya Sutskever, and Picter Abbeel. Variational lossy autoencoder. arXiv:1611.0273) 2016.
Patryk Chrabaszez, Tlya Loshchilov, and Frank Hutter. A downsampled variant of imagenct alternative to the cifar datasets. arXiv:1707.08819 2017.
Laurent Dinh, David Krueger, and Yoshua Bengio. Nice: Non-linear independent components estimation. arXiv:1410.8541 2014.
Laurent Dinh, Jascha Sohl-Diekstein, and Samy Bengio. Density estimation using real NVP. Tn ICLR, 2017.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Un- terthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. ICLR, 2021.
Matej Grcié, Ivan Grubišié, and Siniša Segvić. Densely connected normalizing flows. NeurIPS, 2021.
Jonathan Ho, Xi Chen, Aravind Srinivas, Yan Duan, and Picter Abbeel. Flow+ Improving flow-based generative models with variational dequantization and architecture design. In ICML, 2019.
Emiel Hoogeboom, Rianne Van Den Berg, and Max Welling. Emerging convolutions for generative normal- izing flows. In ICML, 2019.
Chin-Wei Huang, David Krueger, Alexandre Lacoste, and Aaron Courville. Neural autoregressive flows. In ICML, 2018.
Durk P Kingma and Prafulla Dhariwal. Glow: Generative flow with invertible Ixl convolutions. NewrIPS, 2018.
Durk P Kingma, Tim Salimans, Rafal Jozefowicz, Xi Chen, Ilya Sutskever, and Max Welling. Improved variational inference with inverse autoregressive flow. NewrIPS, 2016.
Alexander Kolesnikov and Christoph H Lampert. PixeICNN models with auxiliary variables for natural image modeling. In ICML, 2017.
Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Joan Puigcerver, Jessica Yung, Sylvain Gelly, and Neil Houlsby. Big transfer (BIT): General visual representation learning. In ECCV, 2020.
Ilya Loshchilov, Frank Hutter, et al. Fixing weight decay regularization in adam. arXiv:1711.05100 2017.
Jacob Menick and Nal Kalchbrenner. Generating high fidelity images with subscale pixel networks and multidimensional upscaling. ICLR, 2019.
10

---

Published in Transactions on Machine Learning Research (03/2025)
George Papamakarios, Theo Pavlakon, and lain Murray, Masked autoregressive flow for density estimation. NewrIPS, 2017.
Niki Parmar, Ashish Vaswani, Jakob Uszkoreit, Lukasz Kaiser, Noam Shazeer, Alexander Ku, and Dustin Tran. Image transformer. In ICML, 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.
Rhea Sanjay Sukthanker, Zhiwu Huang, Suryansh Kumar, Radu Timofte, and Luc Van Gool. Generative flows with invertible attentions. In CVPR, 2022.
Lucas Theis, Aäron van den Oord, and Matthias Bethge. A note on the evaluation of generative models. arXiv:111100144, 2015.
Michael Tschannen, Cian Eastwood, and Fabian Mentzer. GIVT: Generative infinite-vcaaullrry transform- ers. In ECCV, 2024a.
Michael Tschannen, André Susano Pinto, and Alexander Kolesnikov. JetFormer: An autoregressive genera- tive model of raw images and text. 10.::2/11.1922.. 2024b.
Aaron Van den Oord, Nal Kalchbrenner, Lasse Espeholt, Oriol Vinyals, and Alex Graves. Conditional image generation with PixeICNN decoders. NewrlPS, 2016a.
Aäron Van den Oord, Nal Kalchbrenner, and Koray Kavukcuoglu. Pixel recurrent neural networks. In ICML, 2016b.
Shuangfei Zhai, Ruixiang Zhang, Preetum Nakkiran, David Berthelot, Jiatao Gu, Huangjie Zheng, Tianrong Chen, Miguel Angel Bautista, Navdeep Jaitly, and Josh Susskind. Normalizing flows are capable generative models. arXiv:2412.00:: 12.06329, 2024.
11

---

Published in Transactions on Machine Learning Research (03/2025)
A Appendix
A.1 Architecture details
Table 3: Architecture details for models used to obtain the main results in Table
ImageNet-Ik 1k
ImageNet-222 22k
32 x 32
64 x 64
32 x 32
64 x 64
Coupling layers
64
64
64
64
VIT depth
6
5
6
6
VIT width
256
512
768
1024
VIT attention heads
8
12
16
A.2 Comparison with other works
The main text of this work reports numbers in the canonical dataset ImageNet 1k x2 and 64x64 (Van den Oord et al. 2016bb and also used the main baseline Flow++ (Ho et al.) 2019). There are other works in the literature which deviate from this protocol rendering number comparisons between works hard. Here we compare our work with Denseflow (Greié al. 2021) which uses different downsampled version of ImageNet (Chaaaszzzz et al. 2017] and uses data augmentations.
The results in Table Siiee@ssoo show that as know in the community the newer downsampled ImageNet is significantly easier with values significantly lower than in Table Additionally it shows once more that ImageNet- 1k is a data constrained setting where data augmentation can play critical role in better numbers but ultimately it lags significantly behind to using transfer learning.
Table 4: Comparison of Jet with Denseflow
ImageNet- lk
32 x 32
64 x 64
Using augmentation from Greié al.
2021
Denseflow
3.63
3.35
Jet
3.638
3.355
No data augmentation Jet
3.656
3.382
Jet (I21k)
3.573
3.300
A.3 Detailed results
Table 5: NLL on ImageNet- 1k 32 x 32 when sweeping architecture hyper parameeer.
depth
16
32
64
ViT dim
256
512
768
256
512
768
256
512
768
ViT depth
1
4.32
4.22
4.18
4.18
4.11
4.08
4.12
4.06
4.04
2
4.11
4.04
4.03
4.03
3.99
4.01
3.99
3.97
4.02
3
4.05
4.00
4.03
3.99
3.97
4.07
3.96
3.98
4.11
4.03
3.99
4.08
3.97
3.99
4.15
3.94
4.03
4.18
4.01
3.99
4.14
3.96
4.01
4.21
3.93
4.07
4.28
4.00
4.00
4.19
3.95
4.04
4.34
3.93
4.12
N/A
12

---

Published in Transactions on Machine Learning Research (03/2025)
Table 6: NLL on ImageNet-Ik 64 x 64 when sweeping architecture hyper-parameters.
Couplings
16
32
64
VIT dim
256
512
768
256
512
768
256
512
768
VIT depth
4.02
3.94
3.90
3.92
3.84
3.81
3.85
3.79
3.77
2
3.86
3.78
3.75
3.78
3.72
3.70
3.74
3.69
3.68
3
3.81
3.73
3.71
3.74
3.69
3.68
3.70
3.66
3.69
3.79
3.71
3.69
3.72
3.67
3.68
3.69
3.66
3.70
3.77
3.70
3.69
3.71
3.66
3.69
3.67
3.66
3.73
3.76
3.69
3.69
3.70
3.66
3.70
3.67
3.66
3.75
Table 7: NLL on ImageNet-1k x 64 when sweeping CNN architecture
Couplings
16
32
64
Block dim
256
512
768
1024
1536
256
512
768
1024
1536
256
512
768
1024
CNN depth
4.17
4.05
4.00
3.96
3.92
4.05
3.95
3.91
3.87
3.83
3.97
3.88
3.84
N/A
4.09
3.97
N/A
3.87
3.83
3.98
3.88
3.83
3.79
3.76
3.90
3.81
3.77
3.75
4.06
3.94
3.88
3.84
3.80
3.95
3.84
3.79
3.77
3.74
3.87
3.78
3.75
3.73
4.04
3.91
3.85
3.82
3.78
3.93
3.83
3.78
3.75
3.74
3.85
3.76
5.97
N/A
4.02
3.90
3.84
3.81
3.77
3.91
3.81
3.77
3.75
3.74
3.84
3.76
3.98
N/A
4.00
3.88
3.83
3.80
3.77
3.90
3.80
3.76
3.74
3.74
3.83
3.89
6.41
3.77
3.98
3.87
3.81
3.78
3.77
3.88
3.79
3.75
3.74
3.75
3.81
3.74
3.87
5.89
Table NLL on ImageNet-21k x
when sweeping architecture hyper-parameters.
Couplings
32
64
128
VIT dim
512
768
1024
512
768
1024
512
768
1024
VIT depth
1
3.79
3.76
3.74
3.74
3.72
3.70
3.71
3.69
3.68
2
3.65
3.63
3.62
3.63
3.61
3.60
3.61
3.60
3.59
3
3.62
3.59
3.58
3.59
3.57
3.56
3.57
3.56
N A
3.59
3.57
3.57
3.57
3.55
N/A
3.55
N/A
N A
5
3.58
3.56
3.56
3.55
3.54
N/A
3.54
N/A
N N/A
6
3.57
3.55
3.54
3.55
3.53
N/A
3.55
N/A
N/A
Table NLL on ImageNet-21k x when sweeping architecture hyper-parameterss.
Couplings VIT dim
32
64
128
512
768
1024
512
768
1024
512
768
1024
VIT depth
1
3.42
3.39
3.37
3.37
3.34
3.34
3.33
3.31
3.29
2
3.29
3.26
3.24
3.26
3.23
3.24
3.23
3.20
3.19
3
3.25
3.22
3.20
3.22
3.19
3.21
3.19
3.16
N/A /
3.23
3.19
3.17
3.19
3.16
3.15
3.16
3.14
N/ A
5
3.21
3.18
3.16
3.18
3.15
N/A
3.15
N/A
N N/A
6
3.20
3.17
3.15
3.17
3.14
3.13
3.14
N/A
N/A
13

---

Published in Transactions on Machine Learning Research (03/2025)
Table 10: NLL on ImageNet-1k 64 x 64 when varying the ratio between channel and spatial couplings.
Couplings ratio
Type of spatial couplings volumn-wist checkerboard mix of all
(channel : spatial)
row-wise
all spatial
3.877
3.882
3.872
3.787
1:1
3.747
3.742
3.741
3.737
2:1
3.735
3.731
3.731
3.725
3:1
3.732
3.729
3.728
3.723
4:1
3.731
3.726
3.724
3.719
5:1
3.728
3.726
3.723
3.720
all channel
3.742
14

---

Published in Transactions on Machine Learning Research (03/2025)
(a) Jet 32 x 32
(b) Jet 64 x 64
(c) Jet (121k) 32 x 32
(d) Jet (121k) 64 x 64
(e) Flow+++ 32 x 32
(I) Flow ++ 64 x 64
Figure 5: Random samples for ImageNet-1k at both 32 x 32 and 64 x 64 resolution. We show samples from Jet when trained from scratch and when finetuning a model pretrained on ImageNet-22 For comparison we also show samples from Flow+ (o et al. 2019
15


Reviews:
[{"review_id": "hva0wKI7QY", "reviewer": "Reviewer_bjSh", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\n* Authors conduct a wide range of ablation studies\n* The paper demonstrate the advantage of pre-training\n\n**Weaknesses**\n* The performance improvements are modest unless pre-training is employed. Jet outperforms Flow++ only with pre-training. For a fair comparison of the pre-trained model, one would need to pre-train Flow++ on Imagenet-21K as well. There also some better results reported on these datasets with normalizing flows (e.g. ANF [1] and VFlow [2]). At the same time, authors repeatedly mention \"state-of-the-art\" performance.\n\n[1] Huang, Chin-Wei, Laurent Dinh, and Aaron Courville. \"Augmented normalizing flows: Bridging the gap between generative flows and latent variable models.\" arXiv preprint arXiv:2002.07101 (2020).\n\n[2] Chen, Jianfei, et al. \"Vflow: More expressive generative flows with variational data augmentation.\" International Conference on Machine Learning. PMLR, 2020.\n\n* Model size and memory consumption are not reported, but form the hyperparameters values I assume it is much more expensive to train and sample from Jet compared to Flow++.\n\n* The paper is not clearly written in several place. Please, see the requested change for the details.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper introduces a model called Jet. It is a normalizing flow model with coupling layers parametrized by vision transformer. The authors aim to simplify the model architecture while improving the performance. Jet does not use many advanced techniques developed for normalizing flows, while demonstrating competitive performance.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "* Tone down state-of-the-art claims.\n* Report number of parameters, memory requirements and sampling time compared to previous works.\n* What is the reasoning behind plain CNN architecture as a baseline and not a model with Self-Attention (as done in Flow++)? According to Figure 1 in Flow++ paper, Attention improves over the standard CNN. I think this would be a better baseline to have.\n\n\n* Writing and missing details:\n   * For some reason non-finetuned result for CIFAR10 are missing. Please, add them.\n   * Result for ablation study in section 3.2.6 are not reported, please fix it.\n   * Move Figure 2 to experimental section.\n   * I think Figure 3 is missing the colour-coding, similarly to Figure 2. It is hard to make any useful conclusions from it currently.\n   * Please, clarify how the dimension splitting is implemented. According to 2.1 patches are flattened so that the input has shape $K \\times 2d$. Is the last dimension in section 2.1 the same as \"spacial dimension\" mentioned in section 2.5? In this case, is spacial splitting performed before the flattening of the patches or after? Does channel-wise splitting split each patch channel-wise or is it split over the \"patch\" dimension?\n   * Section 3.2.4. Paragraph 1 talks about \"K patches\", paragraph 2 talks about \"N/2 patches\". I assume K and N are the same thing. Please, standardize the terminology to avoid confusion.", "broader_impact_concerns": "NA", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "2fDOcFjYdu", "reviewer": "Reviewer_HWBQ", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "__Strengths__\n\nAlthough normalizing flows have recently fallen out of favour for many generative tasks, they have the benefit of producing exact likelihoods, while being fast to sample from. Few work exists that evaluates the use of this specific architectural advancement (ViT) for this class of models. The proposed method is far simpler in terms of model complexity, showing these ViT base architectures can outperform previous methods with a simpler design space. Standard datasets (Imagenet, CIFAR-10) are used for evaluation, while performance is measured using negative log likelihood (NLL), also a standard metric. Comparisons with a state-of-the-art flow model shows a favourable performance for the proposed method. Extensive ablations support the claim of improved performance over a wide range of hyper-parameters.\n\n\n__Weaknesses__\n\nAlthough the claims are mostly supported, the scope of the method is potentially very general while the scope of the experiments is narrow.\n\nThe only metric considered is NLL. While this is a standard metric in this space, NLL is known to not always correlate with image quality for image generation tasks. More recent work on normalizing flows often include FID to further support claims about model performance.\n\nThe major contribution of the paper is the empirical validation of a ViT architecture for normalizing flows.\nThe experiments are limited to toy vision datasets. Although these are standard, they leave questions about the performance of the method beyond a toy setting. This can be important since transformers have quadratic complexity, but the images considered are relatively small and may not reveal scaling properties of the model. Other related factors that may be of interest is the effect of patch size, number of tokens, and different data modalities.\n\nOnly one baseline primarily used for comparison since it is the state-of-the-art. However, it would be instructive to include qualitative comparisons with other generative models to provide the reader a greater context of how the proposed method fits into the space of generative modelling. Often, normalizing flow methods are compared with autoregressive likelihood methods, VAEs, and even diffusion models, since these methods can also provide exact likelihoods or at least approximations for NLL, while also potentially sharing architectural similarities. Although this is the only work to my knowledge that specifically investigates a ViT architecture for normalizing flows, it is not the only to investigate the use of transformers.\n\nThere are also few limitations discussed, making the paper seem one sided, making it difficult for the reader to determine applicability of the model outside of toy settings, or potential directions of future research.\n\nTransfer learning is also investigated to mitigate overfitting, but seems to be an orthogonal to the main contributions of the paper. This approach to prevent overfitting is a general method not specific to normalizing flows, and does not seem to add to the contributions.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper investigates the potential benefits of using new transformer based architectures (ViT) for normalizing flows, a recently neglected class of generative models. The proposed method significantly simplifies the model architecture without the numerous engineered components found in previous normalizing flow methods, only relying on the patch-based attention of the ViT architecture. Evaluations are performed on standard image generation tasks, using a standard metric.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Generally, the style of the paper appears to be of a technical report that covers the results of an extensive hyper-parameter sweep.\nAlthough the proposed approach shows significant improvement across the range of tested hyper-parameters, the scope is very limited.\nThere is limited discussion on limitations of the method, the scope makes it difficult to even imagine how the method may be applied beyond a toy setting, and the presentation/discussion of the results does not convey how the work fits into the greater space of work on generative models.\n\nI recommend adding final evaluations using a metric that is not NLL (FID) to further strengthen claims about model performance. At minimal a discussion, if not new experiments, to investigate limitations and scaling properties of the method would be highly valuable. The inclusion of final quantitative results for other relevant methods along with some discussion would help place the work in the space of other generative modelling work.\n\nThe transfer learning experiments contributes little to the primary claims of the paper, potentially moving these into the appendix would improve the coherence of the paper. Unless there is a way to make these experiments more relevant that I have not thought of.", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "Da3to5zE6C", "reviewer": "Reviewer_9Uwh", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n-----------\nThis paper goes into great detail about the architectural + other ablations, which serves as a powerful artifact for anyone interested in applying recent architectures (e.g. vision transformers) to flow models.\n\nWeaknesses\n---------------\nThe major weakness is that the comparison on imagenet32x2 and 64x64 is a bit suspect. The paper claims to compare on the right dataset of Crabaszcz et al, but the canonical reference is Van der oord et al 2016 (pixel-rnn/cnn paper). This can be fine if the baselines also use the same dataset, but it seems that the flow++ paper uses Van der oord, so this needs to be clarified.\nThere is also a question about the baselines used. In particular, Flow++ is 6 years old, which warrants a deeper look into more recent architectures. For example, the paper \"densely connected normalizing flows\" should satisfy the baseline requirements as a one step architecture and should be compared with (but note that this reference does not use the right/comparable dataset either!).", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper improves upon previous (one step) normalizing flow models by imbuing them with more recent architectural changes.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "See weaknesses.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes a new Transformer-based architecture, Jet, for normalizing flows. The reviewers are in consensus on being a bit on the fence with regards to the significance of the findings of the paper. After having read through the paper, I am of the opinion that the acceptance can be defended:

* the topic is of general interest
* the paper makes a thorough scaling analysis of the investigated architectures (Jet and CNN baseline)
* the paper gives and in depth description and analysis of the model architecture.

The reviewers have some additional suggestions for improvements, some of which the authors have already agreed to do. I encourage the authors to make these because this will improve the usefulness and impact of the paper.
Decision: Accept (poster)