# GIVT: Generative Infinite-Vocabulary Transformers

Michael Tschannen1 Cian Eastwood1,∗ Fabian Mentzer2,∘
1Google DeepMind 2Google Research

###### Abstract

We introduce generative infinite-vocabulary transformers (GIVT) which generate
vector sequences with real-valued entries, instead of discrete tokens from a
finite vocabulary. To this end, we propose two surprisingly simple
modifications to decoder-only transformers: 1) at the input, we replace the
finite-vocabulary lookup table with a linear projection of the input vectors;
and 2) at the output, we replace the logits prediction (usually mapped to a
categorical distribution) with the parameters of a multivariate Gaussian
mixture model. Inspired by the image-generation paradigm of VQ-GAN and
MaskGIT, where transformers are used to model the discrete latent sequences of
a VQ-VAE, we use GIVT to model the unquantized real-valued latent sequences of
a VAE. When applying GIVT to class-conditional image generation with iterative
masked modeling, we show competitive results with MaskGIT, while our approach
outperforms both VQ-GAN and MaskGIT when using it for causal modeling.
Finally, we obtain competitive results outside of image generation when
applying our approach to panoptic segmentation and depth estimation with a
VAE-based variant of the UViM framework.††∗Work done as a Student Researcher
at Google DeepMind.
∘Significant technical contributions.

##  1 Introduction

After becoming the dominant architecture in natural language processing
shortly after their introduction, Transformers [62] have also recently become
very popular in computer vision [16, 54, 34]. Dosovitskiy et al. [16] showed
that by splitting images into sequences of patches, linearly embedding those
patches, and then feeding the resulting sequence of features to a transformer
encoder leads to powerful image classifiers that outperform CNN-based
architectures at large model and data scale. This strategy is now standard for
many discriminative vision task including classification [16], detection [34],
and segmentation [54].

Figure 1:  We compare the standard discrete-token generative transformer
(left) to our continuous, infinite-vocabulary variant (GIVT, right), using the
same decoder-only architecture. At the input, GIVT linearly embeds a sequence
of _real-valued vectors_ instead of discrete tokens via lookup. At the output,
GIVT predicts the parameters of a continuous distribution over real-valued
vectors rather than a categorical distribution over a finite vocabulary.
Figure 2: Selected samples from the GIVT-Causal for 10 ImageNet classes (130,
130, 138, 144, 933, 145, 360, 207, 829, 248).

However, mapping an image to a sequence of (unquantized) feature vectors is
not suitable for transformer-based image generation since generative
transformers decoders consume and predict discrete tokens from some fixed,
finite vocabulary. Such a structure naturally fits natural language, for which
decoder-only models enable powerful sequential generative modeling and
efficient training via teacher forcing [62, 44]. To harness these capabilities
for images, recent works [18, 6, 7, 33, 40] have employed a two-stage approach
which first learns a Vector-Quantized Variational Autoencoder (VQ-VAE) [61] to
map images to a sequence of discrete tokens, and then learns a transformer
decoder to model the latent discrete-token distribution. Another advantage of
such VQ-VAE-based image tokenization is that it enables interleaved multimodal
generative models, simply by concatenating the vocabularies of the different
modalities including text and images [1, 26, 2].

While this two-stage approach was successfully applied to image and multimodal
generation, it has several issues. The vocabulary size in VQ-VAE determines
how much information can be stored in the latent code sequence, and how much
of the visual modeling is done by the VQ-VAE decoder. Having a small
vocabulary can make the latent modeling easier but also makes the latent code
less informative, which prevents control of the low-level details in image
generation, and impacts quality when using the tokens for dense prediction
[30, 36] or low-level discriminative tasks [1, 26]. To mitigate this issue,
the vocabulary size can be increased, which, however, can lead to low
vocabulary utilization [40] so that high-fidelity VQ-VAE setups rely on a
range of advanced techniques, such as entropy losses [6] or codebook-splitting
[30]. Furthermore, large vocabularies lead to correspondingly large embedding
matrices consuming a lot of memory, which can be an issue particularly in
multimodal contexts with vocabularies combined across modalities.

To avoid these issues, we propose a modification to decoder-only transformers
which removes the need for discrete tokens and therefore fixed, finite
vocabularies. Specifically, we propose—to our knowledge for the first time—a
generative transformer decoder which operates on sequences of real-valued
vectors. As real-valued vectors can be seen as an infinite vocabulary, we call
this a Generative Infinite-Vocabulary Transformer (GIVT). We make two minimal
changes compared to the standard transformer decoder architecture [62, 44],
see Fig. 1: 1) at the input, rather than using a sequence of discrete tokens
to look up a finite vocabulary of embeddings, we linearly embed a sequence of
real-valued vectors; and 2) at the output, rather than predicting the
parameters of a categorical distribution over a finite vocabulary (via
logits), we predict the parameters of a continuous distribution over real-
valued vectors. We train this model in the same way as standard transformer
decoders: with a causal attention mask and teacher forcing [62], and
alternatively also explore fast progressive masked-bidirectional-modelling as
in MaskGIT [12, 6, 7].

While GIVT can, in principle, be applied to any sequence of feature vectors,
it can be challenging to directly model sequences that are prohibitively long
or follow a complex distribution (_e.g_., the sequence of RGB pixels obtained
by flattening a high-resolution image). Therefore, similar to the two-stage
approach with VQ-VAEs and analogous the two-stage approach of latent-diffusion
models [47, 43], we first learn a lower-dimensional latent space with a
Gaussian-prior VAE, and then model it with GIVT.

Additionally, we port various inference schemes from the sequence-modeling
literature (_e.g_., temperature sampling, classifier-free guidance).
Surprisingly, this leads to a model which is comparable to or better than VQ-
based approaches, by purely relying on real-valued tokens.

Our main contributions can be summarized as follows:

  1. 1.

We show that GIVT matches or outperforms the standard discrete-token
transformer decoder on both image generation and dense prediction tasks like
semantic segmentation and depth estimation (via UViM [30]).

  2. 2.

We derive variants of standard sampling techniques for the continuous case,
such as temperature sampling, beam search, and classifier-free guidance (CFG,
[23]), and demonstrate their effectiveness.

  3. 3.

We study the relationship between the degree of VAE latent-space
regularization (via KL-term weighting [27, 22]) and the resulting properties
of GIVT.

We emphasize that the VAE and GIVT training only relies on standard techniques
from the deep-learning toolbox, and not the advanced training techniques of
the VQ-VAE literature like auxiliary losses [61, 6] on the latent
representation, codebook reinitialization [32], or dedicated optimization
algorithms [30, 25].

Figure 3:  GIVT-Causal training and inference. Left: During training, we
sample a sequence of real-valued latent vectors from the VAE encoder, and
train GIVT via teacher forcing. Right: During inference, we sample a sequence
of vectors (left-to-right) and feed it to the VAE decoder. We note that we
also explore MaskGIT-like GIVT models not shown here. _No component uses a
quantizer._

##  2 Related work

VQ-VAE for visual data tokenization  Following the success of pixel-space
autoregressive modeling [60, 50, 37, 8] for image generation, moving the
autorgressive modeling to the latent space of VQ-VAEs [61, 46] emerged as a
more efficient alternative. The use of GANs and perceptual losses for VQ-VAE
training as well as modern causal [18, 67, 63] and masked [6, 7, 33]
transformers for latent modeling led to substantial quality improvements.
Another active area leveraging VQ-VAEs is interleaved multimodal generative
modeling of images and text [1, 26, 2]. Further, VQ-VAEs are a popular choice
to tokenize the label space of dense prediction vision tasks [30, 36].
Finally, some language-inspired techniques for self-supervised learning from
images rely on VQ-VAE representations [3, 65, 33].

Discretized mixtures of distributions  replace the dense prediction of the
logits of a categorical distribution with a continuous mixture model which is
subsequently discretized. This approach was proposed in [50] for pixel-space
autoregressive modeling, to reduce the number of model parameters and to
improve learning efficiency, and is also popular in neural compression [38,
10, 39].

VAEs with learned priors  A rich body of literature studies improving VAEs
with learned priors: Inverse autoregressive flows emerged as a popular choice
[28, 9]. Other approaches are based on normalizing flows [58] or a mixture of
variational posteriors with pseudo-inputs [55]. For VAEs with discrete (non-
VQ) latents, learned priors based on Restricted Boltzmann Machines were
studied [59, 48].

Time-series modeling with Transformers  A variety of works has recently
explored transformers for time-series modeling/forecasting. Those works either
use a regression loss [69, 42, 31, 20, 11], quantile forecasting [17, 35], or
resort to discretizing/binning the data [45]. Somewhat related, [41, 64]
regress continuous speech features from discrete tokens. None of these models
predict a continuous distribution like GIVT that allows for autoregressive
generation.

##  3 Generative infinite-vocabulary transformers

As mentioned in Sec. 1, our method is conceptually similar to recent works
that train decoder-only transformer models on the discrete codes of VQ-VAEs
[18, 66, 6, 7], with the crucial difference being that we do not quantize
(_i.e_., do not use VQ). We now describe the components of our method.

###  3.1 VAE training

We first train a _continuous_ $\beta$-VAE [22] with Gaussian encoder and prior
as originally proposed by [27]. Given an input image $x$, the encoder $E$
predicts mean $\mu$, and variance $\sigma$ of a multivariate normal
distribution with diagonal covariance matrix, and samples a representation $z$
from $\mathcal{N}(\mu,\sigma)$ using the reparameterization trick [27]. The
VAE decoder then maps the latent sequence back to an image. Since we use a
Gaussian encoder distribution, the KL-term in the evidence lower bound (ELBO)
[27] can be computed in closed form as described in [27, Sec. F.1]. As for the
reconstruction/likelihood term in the ELBO, we rely on a mixture of MSE,
perceptual loss and GAN loss for image generation following [6, 18], or the
categorical cross-entropy for dense prediction tasks [30]. Our encoder
spatially-downsamples $x$, whereby we obtain a $z$ with spatial dimensions
$h\times w$ and feature dimension $d$, with $h{=}\lceil H/16\rceil,w{=}\lceil
W/16\rceil$, given a $H{\times}W$ input $x$.

The hyperparameter $\beta$ multiplying the KL-term controls how strongly $z$
is regularized. As we shall see in Sec. 5, this regularization of the VAE is
important to be able to model the resulting (true) latent distribution $p(z)$
well.

| Model | Inference | Steps |  FID$\downarrow$ |  Precision$\uparrow$ |  Recall$\uparrow$
---|---|---|---|---|---|---
GANs |  BigGAN-deep [5] |  |  | $\numprint{6.95}$ | $\numprint{0.87}$ | $\numprint{0.28}$
|  StyleGAN-XL [51] |  |  | $\numprint{2.30}$ | $\numprint{0.78}$ | $\numprint{0.53}$
Diffusion |  ADM [13] |  | 250 | $\numprint{10.94}$ | $\numprint{0.69}$ | $\numprint{0.63}$
Models |  ADM-G [13] |  CG $=1.0$ | 250 | $\numprint{4.59}$ | $\numprint{0.82}$ | $\numprint{0.52}$
|  DiT-XL/2 [43] |  | 250 | $\numprint{9.62}$ | $\numprint{0.67}$ | $\numprint{0.67}$
|  DiT-XL/2-G [43] |  CFG $=1.50$ | 250 | $\numprint{2.27}$ | $\numprint{0.83}$ | $\numprint{0.57}$
Sequence |  VQGAN [18] |  Top-k $=$ Mixed |  256⋆ | $\numprint{17.04}$ |  |
Models |  VQGAN [18] |  Top-k $=600$, CG $=0.05$ |  256⋆ | $\numprint{5.20}$ |  |
|  MaskGIT [6] | $t_{C}=4.5$ | 16 |  $\numprint{4.916}$† |  $\numprint{0.836}$† |  $\numprint{0.489}$†
|  MaskGIT [40] |  $t_{C}=4.5$, CFG $=0.1$ | 16 | $\numprint{4.509}$ | $\numprint{0.860}$ | $\numprint{0.465}$
|  GIVT-Causal _(Ours)_ |  |  256⋆ | $\numprint{12.6431}$ | $\numprint{0.7148}$ | $\numprint{0.5813}$
|  GIVT-Causal _(Ours)_ |  Beam $=8/1$ |  256⋆ | $\numprint{6.9137}$ | $\numprint{0.7506}$ | $\numprint{0.5597}$
|  GIVT-Causal _(Ours)_ | $t=0.95$ |  256⋆ | $\numprint{9.1759}$ | $\numprint{0.7332}$ | $\numprint{0.5792}$
|  GIVT-Causal _(Ours)_ |  $t=0.95$, DB-CFG $=0.5$ |  256⋆ | $\numprint{4.2914}$ | $\numprint{0.8265}$ | $\numprint{0.5032}$
|  GIVT-MaskGIT _(Ours)_ |  | 16 | $\numprint{7.5157}$ | $\numprint{0.7941}$ | $\numprint{0.5097}$
|  GIVT-MaskGIT _(Ours)_ |  $t_{C}=16$, DB-CFG $=0.2$ | 16 | $\numprint{5.5739}$ | $\numprint{0.8560}$ | $\numprint{0.4521}$
Table 1:  Results on class-conditional $256{\times}256$ ImageNet, where GIVT-
Causal models outperform their quantization-based counterparts. We report FID
as well as precision and recall (where available). We use the standard ADM
evaluation suite, where FID is calculated w.r.t. the training set. _CG_ :
Classifier guidance acceptance rate or scale, _CFG $=w$_: Classifier-free
guidance with weight $w$ [23], _DB-CFG $=w$_: Our distribution based CFG
variant (Sec. 3.3.2), _Top-k_ : Top-k sampling [19] (“mixed” refers to
multiple $k$), _Beam $=B/F$_: Beam search with $B$ beams, $F$ fans, _$t$_ :
Temperature sampling by scaling the predicted $\sigma$ of our models with $t$,
_$t_{C}$_ : Choice temperature for MaskGIT. _Steps_ shows inference steps.
Additional comments: †Numbers obtained by us from public code, ⋆Inference uses
activation caching.

###  3.2 GIVT training

We next train a GIVT to predict $p(z)$ or $p(z|c)$ (when a conditioning signal
$c$ is available, _e.g_., in class-conditional generation). We flatten $z$
into a $hw$-length sequence of $d$-dimensional _real-valued_ vectors (or “soft
tokens”). Note how this differs from the standard VQ-VAE-based setup, where
the latent transformer decoder models a $hw$-length sequence of integers
denoting codebook indices. Apart from this change, we only perform two small
changes to the standard transformer decoder-only architecture (see Fig. 1): We
replace the embedding lookup tables at the input with a single linear layer to
project from $d$ to the transformer hidden dimension. At the output, we do not
predict a categorical distribution, and instead let the transformer predict
the parameters of a continuous distribution.

Assuming channel-wise independence, we can model each of the $d$ entries with
a $k$-mixture GMM. The GIVT model then predicts $3dk$ parameters ($dk$ means,
scales, and mixture probabilities) that we interpret as $d$ scalar GMMs.
Experimentally, we found it beneficial to normalize the mixture probabilities
with a softmax activation, and the variance parameters with softplus.

We use the standard cross-entropy loss (which is equivalent to minimizing
negative log-likelihood) on the distribution $\tilde{p}$ predicted by GIVT,
and minimize
$\mathcal{L}_{\text{T}}=\sum_{c}\mathbb{E}_{z}\left[-\log\tilde{p}(z|c)\right]$,
assuming the the classes or conditioning signal $c$ uniformly distributed. We
train two types of GIVT models, as described next.

GIVT-Causal  Here, $z$ is flattened into a sequence line-by-line, and the GIVT
predicts every $d$-dimensional vector in the sequence conditioned on all
previous vectors, masking the self-attention layers to be causal [18, 62].
This training strategy is also called teacher forcing and is analogous to the
latent modeling in VQ-GAN [18]. For class-conditional image generation we
prepend a [CLS] vector to the input sequence, _i.e_., a learned vector for
each class $c$.

GIVT-MaskGIT  Here, as in MaskGIT [6], we mask the input randomly during
training and then gradually uncover the masked tokens during inference. The
only changes compared to [6] are related to our real-valued tokens: since we
have infinitely many tokens, there is no obvious choice to define a special
mask token (when using VQ, one can just extend the vocabulary to contain
special tokens, such as [MASK]). Instead, given $z$ and a mask $M$ indicating
for every spatial location whether it is masked, we first replace the
locations in $z$ corresponding to $M$ with zeros (to remove information), and
then embed it with a single dense layer, as above. Additionally, we
_concatenate_ one of two learned special vectors in the feature dimension, a
[MASK] vector for masked locations, and a [UNMASK] vector otherwise (we half
the dimension of the embedded inputs and special tokens s.t. the final hidden
dimension remains unchanged).

###  3.3 Inference

Given a VAE and GIVT trained as above, we run GIVT during inference either
causally (see Fig. 3) or as in MaskGIT [6]. We now investigate the various
inference schemes that have been developed for discrete transformers, and
derive their continuous counterparts.

####  3.3.1 Temperature Sampling, Nucleus Sampling,
Beam Search

In causal sequence models for text (see [24] for an overview and discussion)
and VQ-GAN-based approaches, it is common to tune the sampling algorithm. We
start with temperature sampling, which for discrete models adapts the softmax
temperature of the categorical distributions predicted at each decoding step.
For GIVT, we instead scale the standard deviation of the predicted Gaussian
distribution(s) and call this strategy “variance scaling”. As we will see in
Sec. 4, this simple change can have a significant impact on sample quality.

Nucleus sampling [24] proposes to collect the largest logits such that its
cumulative probability after normalization exceeds a threshold (for example
0.8), and to sample from this reduced-support distribution. In GIVT, for a
single Gaussian prediction per output dimension, this can be approximated by
truncating the predicted distributions per dimension (thereby choosing a
higher-density support). We observed that this has a similar effect to
variance scaling and therefore do not pursue this strategy.

We also consider beam search, which is the same for GIVT as it is for discrete
transformer decoders. For every sample, we maintain $B$ beams, and at every
step we sample a number of candidates for every beam (we call these “fans”
here). We then compute the cumulative log probability for all beams and fans
up to the current sampling step, and select the $B$ beams with the highest
cumulative log probability. Finally, there is no analogous concept for top-k
sampling [19] in GIVT, because it predicts continuous distributions.

Figure 4: VAE ablation: We vary the KL weight $\beta$ and number of channels
$d$ when training the VAE. Dashed lines show the sampling FIDs obtained with a
Base-size GIVT-Causal. As we increase $\beta$, the sampling FID improves, but
the reconstruction FID increases, lower bounding the best possible sampling
FID. For the main results we use $\beta=0.00005$ and $d=16$.

####  3.3.2 Distribution-Based Classifier Free Guidance

In the diffusion literature, classifier-free guidance (CFG) [23] has been
employed with great success. Concretely, conditional diffusion models are
trained with an additional null class $\emptyset$ to learn the unconditional
data distribution. Then, during inference, the conditional log density is
“moved away” from the unconditional one: given a guidance weight $w$, the
updated (diffusion) score estimate is is obtained as

$\tilde{\epsilon}(z,c)=(1+w)\epsilon(z,c)-w\epsilon(z,\emptyset),$ (1)

where $\epsilon$ estimates the gradient of the log density of the data
distribution, $\epsilon(z,c)\propto\nabla_{z}\log p(z|c)$ (see [23, Sec. 2]).
From this, we now derive a CFG variant for our GIVT, since _we directly
predict a density_. We term this approach “Density-Based CFG” (DB-CFG). Eq. 1
can be written as

$\displaystyle\tilde{\epsilon}(z,c)$ $\displaystyle\propto(1+w)\nabla_{z}\log
p(z|c)-w\nabla_{z}\log p(z|\emptyset)$
$\displaystyle\propto\nabla_{z}\log\left(p(z|c)^{1+w}p(z|\emptyset)^{-w}\right),$

_i.e_., $\tilde{\epsilon}$ estimates the log of the density
$p_{\text{CFG}}(z|c)\propto p(z|c)^{1+w}p(z|\emptyset)^{-w}$ (see Fig. 5).
Thus, we want to adapt our models to sample from $p_{\text{CFG}}$. To this
end, we follow [23] and train GIVT with an additional null class $\emptyset$.
During inference, we run GIVT twice at every step, once conditional on the
actual label $c$ and once conditional on $\emptyset$. To implement classifier-
free guidance, we then have to sample from an unnormalized version of
$p_{\text{CFG}}(z)$ derived from the two GIVT predictions. To this end, we
turn to rejection sampling, which requires: 1) an unnormalized density; 2) a
good proposal distribution $p^{\prime}$, that is close to the true target
distribution; and 3) a scaling factor $K$ to bound the likelihood between
$p^{\prime}$ and the unnormalized target density.

Figure 5:  Visualization of our _Density-Based Classifier-Free Guidance (DB-
CFG)_. We show the conditional and unconditional PDFs predicted by GIVT, and
the resulting CFG PDF for two values of $w$. Note how the CFG distributions
become more peaked. We use rejection sampling to sample from $p_{\text{CFG}}$
(see Sec. 3.3.1).

In our setup, the distributions we mix are all Gaussian. We find empirically
that the unconditional distributions (_i.e_., distributions predicted using
the $\emptyset$ label) tend to have larger variance than the conditional ones
(as visualized in Fig. 5). It is thus sensible to pick sample proposals from
$\mathcal{N}(\mu_{c},2\sigma_{c})$, where $\mu_{c},\sigma_{c}$ are the
parameters predicted by GIVT when given the label $c$. We empirically find
that drawing 1000 samples is enough to find at least one valid sample 99.9% of
the time. For the remaining ${<}0.1$%, fall back to sampling from
$\mathcal{N}(\mu_{c},\sigma_{c})$.

We emphasize that the overhead of DB-CFG is small: it requires two forward
passes (per inference step) instead of one to predict the conditional and
unconditional distribution. We then draw 1000 samples from those in parallel
on an accelerator, which is very fast. We refer to App. 3.3.2 for Python code.

VAE Sample |  Step 1 |  Step 4 |  Step 8 |  Step 16
---|---|---|---|---

Figure 6:  _First Column_ : Two images obtained when sampling from the VAE
prior, resulting in a soup of low-level image textures. _Remaining Columns_ :
Visualizing the output of GIVT-MaskGIT, for two ImageNet classes (947, 94),
after 1, 4, 8, 16 inference steps. As expected, the samples start to become
more coherent as we perform more inference steps.

##  4 Experiments

We explore various tasks to show the generality of GIVT.

###  4.1 Image generation

We use ImageNet and explore class-conditional generation (where we condition
our GIVT on class labels) for 256px and 512px, as well as _un_ conditional
generation for 256px.

VAE  We closely follow the setup of MaskGIT [6]. We use their VAE
architecture, built of ResBlocks (as detailed in App. A; encoder and decoder
have a combined 53.5M parameters), remove the VQ layer and related losses, and
replace it with a linear layer predicting $\mu,\sigma$ (Sec. 3.1). We use the
same weights for reconstruction, perceptual, and GAN-loss, as well as
identical optimizer parameters, as in [6, 40]; we only vary the latent
dimension $d$ and weight $\beta$ of the KL-term. By default, we set the token
dimension to 16 (_i.e_., we predict 16 means and variances per token) and
$\beta=5\cdot 10^{-5}$.

GIVT  For GIVT-Causal, we follow the original transformer decoder architecture
[62] in decoder-only mode, but remove biases from attention layers, MLP
blocks, and LayerNorms, and replace ReLU by GELU as is common practice. For
GIVT-MaskGIT, we simply remove the attention mask during training and feed
masked inputs instead of shifted ones. Unless otherwise stated, we use the
BERT-Large configuration [12]. We configure our GIVT models to predict a
single Gaussian distribution by default, but also explore predicting a mixture
of 4 distributions per dimension, and predicting a multivariate Gaussian
modeling the full covariance matrix of the tokens (instead of assuming a
diagonal covariance matrix). For the conditional generation experiments, we
use a learned embedding which we prepend to the embedded token sequence. To
train GIVT, we use the AdaFactor optimizer from [68] with a cosine schedule
(300 epochs; with linear warmup of 30 epochs), set the learning rate and
weight decay to $10^{-3}$ and $10^{-4}$, respectively, the optimizer
$\beta_{2}$ parameter to $0.95$, the dropout probability to $0.2$ for GIVT-
causal and $0.4$ for GIVT-MaskGIT, and the batch size to $8192$. We use the
same data augmentation as during VAE training (see [6, 40]), and sample from
the VAE encoder distribution for every batch (an additional source of
randomness besides data augmentation).

We implement GIVT in JAX [4] and use the JAX version of TensorFlow Probability
[14] to implement the predicted distributions and compute the log-
probabilities. We train GIVT-Casual and GIVT-MaskGIT models for $256\times
256$ ImageNet, and GIVT-MaskGIT additionally on $512\times 512$. We note that
our VAE is trained on $256\times 256$ images, and we use it also for
$512\times 512$ experiments (like [6]).

Figure 7:  _First two plots:_ Showing the impact of DB-CFG (Sec. 3.3.2) on
sampling FID of our class-conditional $256\times 256$ models. First, we show
our GIVT-Causal, with three sampling temperatures $t$, and see a compounding
effect of $t$ and DB-CFG. Then, we show our GIVT-MaskGIT, using different
inference schedules, and observe that apart from a linear schedule, all
explored schedules achieve similar quality. _Right-most plot:_ Average
standard deviation predicted by GIVT-MaskGIT: when using CFG, we perform
inference both with class labels (conditional) and without (unconditional). We
plot the average predicted scale over 24 samples for class 94.

GIVT-MaskGIT inference  We fix the number of inference steps to 16 (as in
[6]), and explore various inference schedules which define how many tokens are
uncovered in each inference step. We employ the cosine schedule from the
paper, as well as various power schedules (_i.e_. letting $r=i/16$ at step
$i$, the fraction of masked tokens is given by $\cos(\pi/2r)$ for the cosine
schedule, and by $1{-}r^{\alpha}$ for the pow:$\alpha$ schedule). As in [6],
we sort tokens by likelihood at each step and sample using a “choice
temperature” $t_{C}$.

Exploring the VAE latent space  To better understand the interplay between the
feature dimension, the strength of the KL regularization $\beta$, the
reconstruction quality of the VAE, and the sampling quality of the GIVT, we
train VAEs with latent dimension in $\\{4,8,16,32\\}$ and $\beta$ in
$\\{2.5\cdot 10^{-5},5\cdot 10^{-5},10^{-4},2\cdot 10^{-4}\\}$ using the VAE-
training setup described at the beginning of this section. For each of the
resulting VAEs we train a causal GIVT with the smaller BERT-Base [12]
dimensions.

Evaluation  For the VAEs we report “reconstruction FID”, the FID obtained when
reconstructing the 50k ImageNet validation images. For our GIVT variants and
baselines, we report the sampling FID [21] when sampling a balanced set of 50k
images covering all ImageNet classes. In both cases, we rely on the well-
established ADM TensorFlow Suite [13], which uses the entire ImageNet training
set as a reference. Furthermore, we also report Precision and Recall [49].
Finally, we evaluate the representation learning capabilities by training a
linear classifier on an average-pooled intermediate representation of
unconditional GIVT-Causal as in prior work [8, 66] (see App. D for details).

###  4.2 Panoptic segmentation and depth estimation

We build on the UViM framework [30], which uses a VQ-VAE to compress the label
space of computer-vision dense-prediction tasks, and an encoder-decoder
transformer taking the RGB image as an input and predicting the associated
dense labels as discrete codes in the VQ-VAE latent space. Here, we replace
the VQ-VAE with a $\beta$-VAE and use a GIVT encoder-decoder to model the
continuous latent code. For the VAE, we use the same transformer-based
autoencoder architecture (6-layer encoder and 12-layer decoder) and cross-
entropy loss as [30]. We set $d=16$ and KL weight $\beta=2.5\cdot 10^{-4}$ for
panoptic segmentation and $\beta=2\cdot 10^{-4}$ for depth estimation. To
build an encoder-decoder GIVT model the same way as in [30], we employ the
causal variant described for ImageNet generation and insert a cross-attention
layer after each self-attention layer. Following [30], we use the
ImageNet-21k-pretrained ViT-L/16 from [53] as the encoder, set the image
resolution to 512px, and adopt the preprocessing and optimization hyper-
parameters from [30]. We use the UViM variant without encoder and decoder
context [30]. Finally, we consider variance scaling and beam search, selecting
the best parameters on a held-out subset of the training set as in [30].

##  5 Results

###  5.1 Image generation

VAE latent space  In Fig. 4 we show how varying the weight $\beta$ of the KL
term affects (i) the VAE reconstruction FID and (ii) the sampling FID of a
Base-size GIVT-Causal trained on the corresponding latent sequence. For (i),
increasing $\beta$ leads to worse reconstruction FID since the VAE can store
less information in the latent. It shifts more of the modeling effort to the
VAE decoder, so that the decoder becomes gradually more generative, which
affects sampling quality (see [56, Sec., 7] for an extensive discussion).

For (ii), we see the opposite trend: increasing $\beta$ leads to decreased
(better) sampling FID for GIVT models trained on the latent sequence.
Arguably, this is because the VAE latent sequence more closely follows the
Gaussian prior, and hence becomes easier for the GIVT to model. Note, however,
that reconstruction FID lower-bounds the sampling FID, so one should chose a
VAE with reconstruction FID smaller than the desired sampling FID. For this
reason, we chose $\beta=5\cdot 10^{-5}$ for the VAE used in our main models,
which obtains a reconstruction FID of 3.66 (note that our best GIVT-Causal
model has sampling FID 4.29, see Table 1). Finally, we see that the latent
dimension $d$ has quite a small impact on reconstruction and sampling FID.

Sampling FID  In Table 1 we show the sampling FID for three class-conditional
model classes on $256\times 256$ ImageNet: GANs, diffusion-based approaches,
and sequence modeling approaches. One can observe that our GIVT-Causal models
outperform the quantization-based VQ-GAN in terms of FID, even though VQ-GAN
has more than $4\times$ the model parameters. We study the impact of our DB-
CFG for both our models in Fig. 7. On the left, for causal models, we see the
impact of temperature scaling, and obtain the best results with $t=0.95$. We
did not observe benefits from temperature scaling for our GIVT-MaskGIT models,
possibly due to the high-confidence selection-scheme already favoring high-
certainty tokens. In the middle of Fig. 7, we instead show the impact of
different inference schedules on GIVT-MaskGIT. We find that the cosine
schedule used in the original VQ-VAE-based MaskGIT yields the best FID (by a
small margin). Moving on to $512\times 512$ ImageNet, we show Sampling FID of
available models in Table 2. We see that our GIVT-MaskGIT scales to this
resolution, and obtains an FID of 5.81. Finally, we present _un_ conditional
results in App. C. This task is considerably harder, but GIVT outperforms the
diffusion-based ADM [13] by a solid margin.

Samples  Fig. 2 shows 10 samples from GIVT-Causal, and App. E shows 10 samples
for the GIVT-MaskGIT. We can see that the model produces high-fidelity,
coherent samples. To study sample diversity, we sample multiple images for a
fixed class in App. E.

In Fig. 6, one can see two samples from our VAE (obtained by decoding latents
sampled from the prior), which show a soup of image textures. We then show
different steps of the GIVT-MaskGIT inference, and observe similar behavior as
in the VQ-based model ([6, Fig. 2]).

Model | Inference | Steps |  FID$\downarrow$
---|---|---|---
ADM [13] |  | 250 | $\numprint{23.2}$
ADM-G [13] |  CG $=1.0$ | 250 | $\numprint{7.72}$
DiT-XL/2 [43] |  | 250 | $\numprint{12.03}$
DiT-XL/2-G [43] |  CFG $=1.5$ | 250 | $\numprint{3.04}$
MaskGIT [6] | $t_{C}=4.5$ | 16 |  $\numprint{7.8007}$†
GIVT-MaskGIT _(Ours)_ |  | 16 | $\numprint{6.3577}$
GIVT-MaskGIT _(Ours)_ |  DB-CFG $=0.2$ | 16 | $\numprint{5.81271300238609}$
Table 2:  Results on class-conditional $512{\times}512$ ImageNet. We use the
standard ADM evaluation suite for metrics, where FID is calculated w.r.t. the
training set. Our masked GIVT-MaskGIT obtains competitive FID scores with only
16 inference steps. †Numbers obtained by us from public code.

Ablations and visualizations  We did not observe clear benefits from using
more than one mixture component per latent dimension in our image generation
experiments (both for causal and MaskGIT-based GIVTs). This is perhaps
unsurprising since our VAE’s latent codes are regularized towards a unimodal
Gaussian prior. Using multiple mixtures might be more beneficial when, _e.g_.,
training GIVT on less-regularized latent spaces such as those of vision
encoders pretrained in supervised or contrastive settings.

Further, predicting the full covariance matrix per latent vector, rather than
assuming a diagonal covariance, only led to very modest gains of about 3% for
MaskGIT-based GIVT. On the other hand, a full covariance matrix makes DB-CFG
less tractable than the diagonal covariance (because the high dimensional
multivariate distribution has more regions of low density).

In Fig. 7 (right), we visualize the predicted standard deviations as a
function of the GIVT-MaskGIT inference step. The standard deviation gradually
decreases, meaning that the predictions later in the sampling process become
more certain. Furthermore, the unconditional predictions generally have a
higher standard deviation, as expected.

Representation learning  Table 3 shows the ImageNet linear probing accuracy of
unconditional GIVT-Causal and generative models from the literature (we chose
the model variants closest in terms of model size and compute). GIVT-Causal
matches VIM+ViT (ViT-VQ-GAN) [66] which has more than $2\times$ the model
parameters and $4\times$ the sequence length (and hence FLOPs). GIVT-Causal is
only outperformed by MAGE [33], whose latent encoder-decoder architecture is
better suited for representation learning than decoder-only models. An
investigation of the probing accuracy as function of the layer index can be
found in App. D.

###  5.2 Panoptic segmentation and depth estimation

Table 4 compares the performance of a GIVT-based UViM variant with a VQ-VAE-
based baseline (both without encoder/decoder context) on COCO Panoptic 2017
[29] and NYU Depth v2 [52]. We report the panoptic quality metric (PQ) [29]
and RMSE, respectively, and find that our GIVT-based model outperforms the
baseline by a small margin in panoptic segmentation and performs slightly
worse in depth estimation. In App. E we show visual results.

Model | Type | #Tok. | #Param. |  Acc.$\uparrow$
---|---|---|---|---
BigBiGAN [15] |  |  | 344M | 61.3
iGPT-L [8] | dec.-only | 1024 | 1362M | 60.3
VIM+CNN [66] | dec.-only | 1024 | 650M | 61.8
VIM+ViT [66] | dec.-only | 1024 | 650M | 65.1
MAGE ViT-L [33] | enc.-dec. | 256 | 404M | 78.9
GIVT-Causal _(Ours)_ | dec.-only | 256 | 304M | 65.1
Table 3:  ImageNet linear probing accuracy of unconditional GIVT-Causal and generative models from the literature. GIVT-Causal matches VIM+ViT (ViT-VQ-GAN) [66] which has more than $2\times$ the model parameters and $4\times$ the sequence length (and hence FLOPs). _Type_ : (Latent) generative model type. _#Param._ : Number of parameters of the full (latent) generative model.  |  COCO Pan. (PQ$\uparrow$) |  NYU Depth v2 (RMSE$\downarrow$)
---|---|---
| recon. | inference |  recon. | inference
UViM [30] | 66.0 | 39.0 |  0.183 | 0.459
GIVT (ours) | 71.0 | 40.2 |  0.195 | 0.474
Table 4: UViM based on GIVT-Causal and VQ-VAE evaluated on panoptic
segmentation (COCO Panoptic 2017) and depth estimation (NYU Depth v2). We
report the panoptic quality (PQ) and RMSE for the VAE/VQ-VAE reconstructions
of the validation set label maps (recon.) and the inference metrics on the
actual dense prediction tasks (inference). GIVT obtains metrics comparable to
the VQ-based UViM.

##  6 Conclusion

In this paper, we proposed two small modifications to standard transformer
decoder-only models to make them suitable to generating real-valued vectors.
To our knowledge, this is the first decoder-only model amenable to generating
sequences of real-valued vectors. In the context of image generation with VQ-
GAN or Mask-GIT, this side-steps training difficulties such as low codebook
usage in VQ-VAEs and corresponding mitigations like entropy losses or
codebook-splitting algorithms, by enabling the use of standard VAEs which are
much easier to train. Furthermore, our method avoids large embedding matrices
because the feature representations can directly be consumed and predicted by
our GIVT model. Surprisingly, our simple, quantization-free approach
outperforms its VQ-based equivalent when using causal transformers for class-
conditional image generation, and obtains strong performance in dense
prediction tasks when applied to UViM. We hope that future work explores
applications of GIVT to other modalities such as audio and time-series
modeling.

## Acknowledgments

We would like to thank André Susano Pinto, Neil Houlsby, Eirikur Agustsson,
and Lucas Theis for inspiring discussions and helpful feedback on this
project. We also thank Han Zhang for support with the VAE training code.

## References

  * Aghajanyan et al. [2022] Armen Aghajanyan, Po-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.
  * Aghajanyan et al. [2023] 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 generative mixed-modal language models.  In _ICML_ , 2023.
  * Bao et al. [2021] Hangbo Bao, Li Dong, Songhao Piao, and Furu Wei.  BEiT: BERT pre-training of image transformers.  In _ICLR_ , 2021.
  * Bradbury et al. [2018] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang.  JAX: composable transformations of Python+NumPy programs, 2018\.
  * Brock et al. [2018] Andrew Brock, Jeff Donahue, and Karen Simonyan.  Large scale gan training for high fidelity natural image synthesis.  In _ICLR_ , 2018.
  * Chang et al. [2022] Huiwen Chang, Han Zhang, Lu Jiang, Ce Liu, and William T Freeman.  MaskGIT: Masked generative image transformer.  In _CVPR_ , pages 11315–11325, 2022.
  * Chang et al. [2023] Huiwen Chang, Han Zhang, Jarred Barber, AJ Maschinot, José Lezama, Lu Jiang, Ming Yang, Kevin P. Murphy, William T. Freeman, Michael Rubinstein, Yuanzhen Li, and Dilip Krishnan.  Muse: Text-to-image generation via masked generative transformers.  In _ICML_ , 2023.
  * Chen et al. [2020] Mark Chen, Alec Radford, Rewon Child, Jeffrey Wu, Heewoo Jun, David Luan, and Ilya Sutskever.  Generative pretraining from pixels.  In _ICML_ , pages 1691–1703, 2020.
  * Chen et al. [2016] Xi Chen, Diederik P Kingma, Tim Salimans, Yan Duan, Prafulla Dhariwal, John Schulman, Ilya Sutskever, and Pieter Abbeel.  Variational lossy autoencoder.  In _ICLR_ , 2016.
  * Cheng et al. [2020] Zhengxue Cheng, Heming Sun, Masaru Takeuchi, and Jiro Katto.  Learned image compression with discretized gaussian mixture likelihoods and attention modules.  In _CVPR_ , pages 7939–7948, 2020.
  * Das et al. [2023] Abhimanyu Das, Weihao Kong, Rajat Sen, and Yichen Zhou.  A decoder-only foundation model for time-series forecasting.  _arXiv:2310.10688_ , 2023.
  * Devlin et al. [2019] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova.  Bert: Pre-training of deep bidirectional transformers for language understanding.  _NAACL-HLT_ , 2019.
  * Dhariwal and Nichol [2021] Prafulla Dhariwal and Alexander Nichol.  Diffusion models beat GANs on image synthesis.  _NeurIPS_ , pages 8780–8794, 2021.
  * Dillon et al. [2017] Joshua V Dillon, Ian Langmore, Dustin Tran, Eugene Brevdo, Srinivas Vasudevan, Dave Moore, Brian Patton, Alex Alemi, Matt Hoffman, and Rif A Saurous.  Tensorflow distributions.  _arXiv:1711.10604_ , 2017.
  * Donahue and Simonyan [2019] Jeff Donahue and Karen Simonyan.  Large scale adversarial representation learning.  In _NeurIPS_ , 2019.
  * Dosovitskiy et al. [2021] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, 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.
  * Eisenach et al. [2020] Carson Eisenach, Yagna Patel, and Dhruv Madeka.  MQTransformer: Multi-horizon forecasts with context dependent and feedback-aware attention.  _arXiv:2009.14799_ , 2020.
  * Esser et al. [2020] Patrick Esser, Robin Rombach, and Björn Ommer.  Taming transformers for high-resolution image synthesis.  In _CVPR_ , pages 12868–12878, 2020.
  * Fan et al. [2018] Angela Fan, Mike Lewis, and Yann Dauphin.  Hierarchical neural story generation.  In _ACL_ , pages 889–898, 2018.
  * Garza and Mergenthaler-Canseco [2023] Azul Garza and Max Mergenthaler-Canseco.  TimeGPT-1.  _arXiv:2310.03589_ , 2023.
  * Heusel et al. [2017] Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter.  GANs trained by a two time-scale update rule converge to a local nash equilibrium.  _NeurIPS_ , 2017.
  * Higgins et al. [2016] Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner.  Beta-VAE: Learning basic visual concepts with a constrained variational framework.  In _ICLR_ , 2016.
  * Ho and Salimans [2022] Jonathan Ho and Tim Salimans.  Classifier-free diffusion guidance.  _arXiv:2207.12598_ , 2022.
  * Holtzman et al. [2019] Ari Holtzman, Jan Buys, Li Du, Maxwell Forbes, and Yejin Choi.  The curious case of neural text degeneration.  In _ICLR_ , 2019.
  * Huh et al. [2023] Minyoung Huh, Brian Cheung, Pulkit Agrawal, and Phillip Isola.  Straightening out the straight-through estimator: Overcoming optimization challenges in vector quantized networks.  In _ICML_ , 2023.
  * Kim et al. [2023] Sungwoong Kim, Daejin Jo, Donghoon Lee, and Jongmin Kim.  MAGVLT: Masked generative vision-and-language transformer.  In _CVPR_ , pages 23338–23348, 2023.
  * Kingma and Welling [2013] Diederik P Kingma and Max Welling.  Auto-encoding variational bayes.  _arXiv:1312.6114_ , 2013.
  * Kingma et al. [2016] Durk P Kingma, Tim Salimans, Rafal Jozefowicz, Xi Chen, Ilya Sutskever, and Max Welling.  Improved variational inference with inverse autoregressive flow.  _NeurIPS_ , 2016.
  * Kirillov et al. [2019] Alexander Kirillov, Kaiming He, Ross Girshick, Carsten Rother, and Piotr Dollár.  Panoptic segmentation.  In _CVPR_ , pages 9404–9413, 2019.
  * Kolesnikov et al. [2022] 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.  _NeurIPS_ , pages 26295–26308, 2022.
  * Kunz et al. [2023] Manuel Kunz, Stefan Birr, Mones Raslan, L. Ma, Zhuguo Li, Adele Gouttes, Mateusz Koren, Tofigh Naghibi, Johannes Stephan, Mariia Bulycheva, Matthias Grzeschik, Armin Keki’c, Michael Narodovitch, Kashif Rasul, Julian Sieber, and Tim Januschowski.  Deep learning based forecasting: a case study from the online fashion industry.  _arXiv:2305.14406_ , 2023.
  * Łańcucki et al. [2020] Adrian Łańcucki, Jan Chorowski, Guillaume Sanchez, Ricard Marxer, Nanxin Chen, Hans JGA Dolfing, Sameer Khurana, Tanel Alumäe, and Antoine Laurent.  Robust training of vector quantized bottleneck models.  In _IJCNN_ , pages 1–7, 2020.
  * Li et al. [2023] Tianhong Li, Huiwen Chang, Shlok Mishra, Han Zhang, Dina Katabi, and Dilip Krishnan.  Mage: Masked generative encoder to unify representation learning and image synthesis.  In _CVPR_ , pages 2142–2152, 2023.
  * Li et al. [2022] Yanghao Li, Hanzi Mao, Ross Girshick, and Kaiming He.  Exploring plain vision transformer backbones for object detection.  In _ECCV_ , pages 280–296, 2022.
  * Lim et al. [2021] Bryan Lim, Sercan Ö Arık, Nicolas Loeff, and Tomas Pfister.  Temporal fusion transformers for interpretable multi-horizon time series forecasting.  _International Journal of Forecasting_ , pages 1748–1764, 2021.
  * Lu et al. [2022] 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.
  * Menick and Kalchbrenner [2018] Jacob Menick and Nal Kalchbrenner.  Generating high fidelity images with subscale pixel networks and multidimensional upscaling.  _arXiv:1812.01608_ , 2018.
  * Mentzer et al. [2020] Fabian Mentzer, Luc Van Gool, and Michael Tschannen.  Learning better lossless compression using lossy compression.  In _CVPR_ , pages 6638–6647, 2020.
  * Mentzer et al. [2023a] Fabian Mentzer, Eirikur Agustsson, and Michael Tschannen.  M2T: Masking transformers twice for faster decoding.  _arXiv:2304.07313_ , 2023a.
  * Mentzer et al. [2023b] Fabian Mentzer, David Minnen, Eirikur Agustsson, and Michael Tschannen.  Finite scalar quantization: VQ-VAE made simple.  _arXiv:2309.15505_ , 2023b.
  * Nachmani et al. [2023] Eliya Nachmani, Alon Levkovitch, Julian Salazar, Chulayutsh Asawaroengchai, Soroosh Mariooryad, RJ Skerry-Ryan, and Michelle Tadmor Ramanovich.  Lms with a voice: Spoken language modeling beyond speech tokens.  _arXiv:2305.15255_ , 2023.
  * Nie et al. [2022] Yuqi Nie, Nam H Nguyen, Phanwadee Sinthong, and Jayant Kalagnanam.  A time series is worth 64 words: Long-term forecasting with transformers.  In _ICLR_ , 2022.
  * Peebles and Xie [2022] William Peebles and Saining Xie.  Scalable diffusion models with transformers.  _arXiv:2212.09748_ , 2022.
  * Radford et al. [2018] Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever.  Improving language understanding by generative pre-training.  2018\.
  * Rasul et al. [2023] Kashif Rasul, Arjun Ashok, Andrew Robert Williams, Arian Khorasani, George Adamopoulos, Rishika Bhagwatkar, Marin Bilovs, Hena Ghonia, Nadhir Hassen, Anderson Schneider, Sahil Garg, Alexandre Drouin, Nicolas Chapados, Yuriy Nevmyvaka, and Irina Rish.  Lag-Llama: Towards foundation models for time series forecasting.  _arXiv:2310.08278_ , 2023.
  * Razavi et al. [2019] Ali Razavi, Aaron Van den Oord, and Oriol Vinyals.  Generating diverse high-fidelity images with VQ-VAE-2.  _NeurIPS_ , 2019.
  * Rombach et al. [2022] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer.  High-resolution image synthesis with latent diffusion models.  In _CVPR_ , 2022.
  * Sadeghi et al. [2019] Hossein Sadeghi, Evgeny Andriyash, Walter Vinci, Lorenzo Buffoni, and Mohammad H Amin.  PixelVAE++: Improved pixelvae with discrete prior.  _arXiv:1908.09948_ , 2019.
  * Sajjadi et al. [2018] Mehdi SM Sajjadi, Olivier Bachem, Mario Lucic, Olivier Bousquet, and Sylvain Gelly.  Assessing generative models via precision and recall.  _NeurIPS_ , 2018.
  * Salimans et al. [2016] Tim Salimans, Andrej Karpathy, Xi Chen, and Diederik P Kingma.  PixelCNN++: Improving the PixelCNN with discretized logistic mixture likelihood and other modifications.  In _ICLR_ , 2016.
  * Sauer et al. [2022] Axel Sauer, Katja Schwarz, and Andreas Geiger.  StyleGAN-XL: Scaling StyleGAN to large diverse datasets.  In _SIGGRAPH_ , 2022.
  * Silberman et al. [2012] Nathan Silberman, Derek Hoiem, Pushmeet Kohli, and Rob Fergus.  Indoor segmentation and support inference from RGBD images.  In _ECCV_ , pages 746–760, 2012.
  * Steiner et al. [2021] Andreas Steiner, Alexander Kolesnikov, Xiaohua Zhai, Ross Wightman, Jakob Uszkoreit, and Lucas Beyer.  How to train your ViT? data, augmentation, and regularization in vision transformers.  _TMLR_ , 2021.
  * Strudel et al. [2021] Robin Strudel, Ricardo Garcia, Ivan Laptev, and Cordelia Schmid.  Segmenter: Transformer for semantic segmentation.  In _CVPR_ , pages 7262–7272, 2021.
  * Tomczak and Welling [2018] Jakub Tomczak and Max Welling.  Vae with a vampprior.  In _AISTATS_ , pages 1214–1223, 2018.
  * Tschannen et al. [2018] Michael Tschannen, Olivier Bachem, and Mario Lucic.  Recent advances in autoencoder-based representation learning.  _arXiv:1812.05069_ , 2018.
  * Tschannen et al. [2023] Michael Tschannen, Manoj Kumar, Andreas Steiner, Xiaohua Zhai, Neil Houlsby, and Lucas Beyer.  Image captioners are scalable vision learners too.  In _NeurIPS_ , 2023.
  * Vahdat and Kautz [2020] Arash Vahdat and Jan Kautz.  NVAE: A deep hierarchical variational autoencoder.  _NeurIPS_ , pages 19667–19679, 2020.
  * Vahdat et al. [2018] Arash Vahdat, Evgeny Andriyash, and William Macready.  Dvae#: Discrete variational autoencoders with relaxed boltzmann priors.  _NeurIPS_ , 2018.
  * Van Den Oord et al. [2016] Aäron Van Den Oord, Nal Kalchbrenner, and Koray Kavukcuoglu.  Pixel recurrent neural networks.  In _ICML_ , pages 1747–1756, 2016.
  * van den Oord et al. [2017] Aäron van den Oord, Oriol Vinyals, and Koray Kavukcuoglu.  Neural discrete representation learning.  _NeurIPS_ , 2017.
  * Vaswani et al. [2017] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  _NeurIPS_ , 2017.
  * Villegas et al. [2022] Ruben Villegas, Mohammad Babaeizadeh, Pieter-Jan Kindermans, Hernan Moraldo, Han Zhang, Mohammad Taghi Saffar, Santiago Castro, Julius Kunze, and Dumitru Erhan.  Phenaki: Variable length video generation from open domain textual descriptions.  In _ICLR_ , 2022.
  * Wang et al. [2023] Jiaming Wang, Zhihao Du, Qian Chen, Yunfei Chu, Zhifu Gao, Zerui Li, Kai Hu, Xiaohuan Zhou, Jin Xu, Ziyang Ma, Wen Wang, Siqi Zheng, Chang Zhou, Zhijie Yan, and Shiliang Zhang.  LauraGPT: Listen, attend, understand, and regenerate audio with GPT.  _arXiv:2310.04673_ , 2023.
  * Wang et al. [2022] Rui Wang, Dongdong Chen, Zuxuan Wu, Yinpeng Chen, Xiyang Dai, Mengchen Liu, Yu-Gang Jiang, Luowei Zhou, and Lu Yuan.  Bevt: Bert pretraining of video transformers.  In _CVPR_ , pages 14733–14743, 2022.
  * Yu et al. [2022a] 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 VQGAN.  _ICLR_ , 2022a.
  * Yu et al. [2022b] Jiahui Yu, Yuanzhong Xu, Jing Yu Koh, Thang Luong, Gunjan Baid, Zirui Wang, Vijay Vasudevan, Alexander Ku, Yinfei Yang, Burcu 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.  _TMLR_ , 2022b.
  * Zhai et al. [2022] Xiaohua Zhai, Alexander Kolesnikov, Neil Houlsby, and Lucas Beyer.  Scaling vision transformers.  In _CVPR_ , pages 12104–12113, 2022.
  * Zhou et al. [2021] Haoyi Zhou, Shanghang Zhang, Jieqi Peng, Shuai Zhang, Jianxin Li, Hui Xiong, and Wancai Zhang.  Informer: Beyond efficient transformer for long sequence time-series forecasting.  In _AAAI_ , pages 11106–11115, 2021.

##  Appendix A Training details

###  A.1 Image generation

For the image generation experiments on ImageNet, we adapt the CNN-based VQ-
GAN tokenizer from MaskGIT (see vqgan_tokenizer.py). We replace the VQ layer
and related losses with a Gaussian reparametrization layer [27], and we use
the hyper parameters given in Table 5 for the CNN. For our transformers, we
use the standard pre-LN setup (see, _e.g_., UViM vtt.py for a concrete
implementation). We use the BERT-Large [12] configuration for the main models
(24 layers, 16 attention heads, MLPs with hidden dimension 4096, embedding
dimension 1024).

embedding_dim | 32
---|---
filters | 128
num_res_blocks | 2
channel_multipliers | (1, 1, 2, 2, 4)
conv_downsample | False
activation_fn | “swish”
norm_type | “GN”
Table 5: Hyperparameters of the ImageNet CNN-based VAE encoder/decoder. Note
that the 32 features of the embedding are split into 16 means and 16 scales,
so our actual representation has $d=16$ channels after reparametrization.

ImageNet preprocessing  We preprocess the ImageNet data as follows, following
prior work [6, 40]:

  * •

Decode JPEGs

  * •

Random crop such that between 80% and 100% of the source image are retained

  * •

Resize to target resolution ($256{\times}256$ or $512{\times}512$) using
bicubic filters with anti-aliasing

  * •

Randomly flip the image horizontally with probability 0.5

###  A.2 Panoptic Segmentation and Depth Estimation

For our UViM panoptic segmentation and depth estimation experiments, we adopt
the public UViM GitHub repo and only replace the VQ layer in stage I, adapt
the transformer decoder in stage II, and modify the losses accordingly.

##  Appendix B DB-CFG implementation

We show the JAX implementation of the rejection sampler we use for DB-CFG in
Fig. 9.

##  Appendix C Unconditional image generation

In Table 6 we present FID results for unconditional ImageNet generation.

Model | Inference | FID$\downarrow$
---|---|---
ADM [13] |  | $\numprint{26.2}$
MAGE (ViT-B) | temp $=6.0$ | $\numprint{12.080213022509724}$
MAGE (ViT-L) | temp $=6.0$ | $\numprint{9.9528908191744}$
GIVT-Causal _(Ours)_ |  | $\numprint{24.95}$
GIVT-Causal _(Ours)_ | Beam $=4/1$, $t{=}0.95$ | $\numprint{17.67322949324796}$
Table 6:  Results for unconditional $256\times 256$ ImageNet generation. We
use the standard ADM evaluation suite for metrics, where FID is calculated
w.r.t. the training set. We obtained samples for MAGE [33] using their GitHub
code. Note that the MAGE Base and Large models have considerably more model
parameters than our corresponding models because MAGE has a latent encoder-
decoder model (rather than a decoder-only model).

##  Appendix D Probing intermediate representations

We extract intermediate representations form a GIVT-causal trained for
$256\times 256$ unconditional ImageNet generation by average-pooling the
output of a given layer over the sequence dimension (which results in a
feature vector of dimension 1024 for L-sized models). Thereby, the input image
is first encoded with the VAE encoder and then fed to GIVT as during teacher
forcing (_i.e_. the latent sequence is first right-shifted and padded, and
then fed to GIVT with causal attention mask).

Following prior work [8, 66], we investigate the impact of the layer index on
the linear probing accuracy. To speed up the evaluation we rely on 25-shot
classification with the fast evaluator form [68] and test a range of diverse
data sets. The results presented in Fig. 8 show that the range of layers 10 to
20 lead to high accuracy for all data sets.

We choose layer 18 for the linear probing experiments using the full ImageNet
training set. To select hyperparameters we set aside 1% of the training data
and consider the selection of schedules and hyperparameters described in [57,
App. A]. The results, along with other generative models from the literature,
can be found in Table 3. See Sec. 5.1 for a discussion of these results.

Figure 8: 25-shot linear probing accuracy as a function of the GIVT-Causal
layer trained on unconditional $256\times 256$ ImageNet. The range of layers
10 to 20 lead to high accuracy for all data sets.

##  Appendix E Additional visual examples

In Fig. 10, we show reconstructions from our VAE when feeding it images from
the ImageNet validation set. We see low-level texture variations in the
reconstructions when sampling from the encoder distribution multiple times.
Fig. 11 shows multiple samples for a fixed label to show the diversity of our
model and compare with baselines. In Fig. 12 we show various samples from
MaskGIT for the same classes that we shown in Fig. 2. Fig. 13 shows samples
from our $512\times 512$ GIVT-MaskGIT model. We explore changing labels midway
through sampling (_i.e_., after generating the top half of the image) for
GIVT-Causal in Fig. 14.

For our GIVT-Casual UViM models and VQ baselines, we show visual outputs in
Figs. 15 16.

⬇

import jax

import jax.numpy as jnp

import tensorflow_probability as tfp

tfd = tfp.substrates.jax.distributions

# We assume the ps have shape (b, seq_len, c).

def rejection_sample(

seed, p_cond: tfd.Normal, p_uncond: tfd.Normal, w, max_samples=1_000

):

rng_sample, rng_uni = jax.random.split(seed, 2)

scale_simple = jnp.stack([p_cond.scale, p_uncond.scale], -1).max(-1) * 2

simple = tfd.Normal(loc=p_cond.loc, scale=scale_simple) # Proposal.

def unnormalized_pcfg(x):

return jnp.exp((1 + w) * p_cond.log_prob(x) - w * p_uncond.log_prob(x))

# Find scaling factor by checking for max around the conditional loc.

points = p_cond.loc[None, ...] + jnp.linspace(-10, 10, 1001).reshape(

-1, 1, 1, 1)

fac = jnp.max(unnormalized_pcfg(points) / simple.prob(p_cond.loc), axis=0)

# Shape (max_samples, b, seq_len, c)

xs = simple.sample(seed=rng_sample, sample_shape=(max_samples,))

facq = fac * simple.prob(xs)

ys = jax.random.uniform(rng_uni, shape=facq.shape, minval=0.0, maxval=facq)

p = unnormalized_pcfg(xs)

# Mask is True if the sample is valid.

mask = ys < p

# Now we need to do tricks to get the first element in ‘mask‘ that is

# True, in a jit-able way. We do this by making a shifted mask that is

# False for every element after the first True. Example:

# mask [0, 1, 0, 1, 0, 0, 1, 0]

# cmask [0, 1, 1, 1, 1, 1, 1, 1]

# shifted_cmask [0, 0, 1, 1, 1, 1, 1, 1]

# keep [0, 1, 0, 0, 0, 0, 0, 0] # <- picks the first valid!

cmask = jnp.cumsum(mask, axis=0).astype(jnp.bool_)

shifted_cmask = jnp.pad(

cmask, [(1, 0), (0, 0), (0, 0), (0, 0)], constant_values=False

)[:-1]

assert shifted_cmask.shape == mask.shape

keep = jnp.logical_and(cmask, jnp.logical_not(shifted_cmask))

sample = jnp.where(keep, xs, 0).sum(0) # Grab the first valid sample.

ok = mask.sum(0) > 0 # Fall back to pdf_c if not ok.

return jnp.where(ok, sample, pdf_c.sample(seed=rng_sample))

Figure 9: jax.jit-compatible implementation of the rejection sampler used for
DB-CFG Figure 10:  Samples from the VAE encoder distribution: The first column
is an image from the ImageNet validation set. We encode it with the VAE
encoder, obtain the approximate posterior distribution, and sample from it 4
times. The resulting 4 latent sequences are then decoded to images (last 4
columns). The semantic layout of the input image is well preserved, only the
low-level textures change.

GIVT-Causal _(Ours)_ |  GIVT-MaskGIT _(Ours)_ |  Validation Set
---|---|---

BigGAN [5] (via [6, Fig. 9]) |  VQ-based MaskGIT [6, Fig. 9] |
---|---|---

Figure 11: Various samples for the 980 class to show the diversity of our
samples. For reference, we copy examples from from [6, Fig. 9].  Figure 12:
Selected $256\times 256$ samples from GIVT-MaskGIT, for the same 10 ImageNet
classes as in Fig. 2.  Figure 13: Selected $512\times 512$ samples from the
GIVT-MaskGIT, for the same 10 ImageNet classes as in Fig. 2.  Figure 14:
Changing the label half way through sampling (_i.e_., after the top half is
generated) for GIVT-Causal: The first column uses the same label for the top
and bottom half, the other columns switch to a different labels. Top row
labels: “golden retriever” (207) for the top half; “otter” (360), “gorilla”
(366), “camel” (355) for the bottom half. Bottom row labels: “bird house”
(448) for the top half; “boat house” (449), “light house” (437), “bakery”
(415) for the bottom half. Note that for each row, the top row _latent_ is
always the same, but the overall color balance in the _RGB output_ might be
different because of the VAE decoder (possibly due to GroupNorm layers).

Input |  Ground Truth |  GIVT-Causal UViM |  VQ-based UViM
---|---|---|---

Figure 15:  ‘COCO Panoptic Segmentation’ visual examples from our UViM models.
Panoptic segmentation visual examples from COCO Panoptic 2017 for UViM based
on GIVT-Causal and standard VQ-based UViM.

Input |  Ground Truth |  GIVT-Causal UViM |  VQ-based UViM
---|---|---|---

Figure 16:  Depth estimation visual examples from NYU Depth v2 for UViM based
on GIVT-Causal and standard VQ-based UViM.