Scaling Laws for Generative Mixed-Modal Language Models

Introduction

Generative language models have been developed for a wide range of data modalities, including natural language text Brown et al. (2020), code (Chen et al., 2021; Fried et al., 2022), images (Ramesh et al., 2021; Yasunaga et al., 2022), and molecules or proteins (Chilingaryan et al., 2022; Hsu et al., 2022). Recent work has also introduced unified models (Aghajanyan et al., 2022; Reed et al., 2022; Wang et al., 2022; Zellers et al., 2022) that can simultaneously model multiple modalities. One advantage of generative modeling in these cases is that the models scale well in practice; adding data, compute, or parameters typically improves model quality. These scaling trends have been carefully studied for uni-modal models (Kaplan et al., 2020; Hoffmann et al., 2022) and some recent work focuses on pairs of modalities (Droppo & Elibol, 2021; Henighan et al., 2020). However, the scaling behavior of larger number of modalities remains largely unstudied.

We present an extensive empirical study of scaling laws for mixed-modal generative language models over tokens. We assume that every modality can be represented as a sequence of tokens (e.g. VQ-VAEs for images (Esser et al., 2020) or HuBERT for speech (Hsu et al., 2021)). With this assumption, we can train a single discrete language model to represent data with arbitrary subsets of modalities presented in arbitrary orders. Such mixed-modal models are very general, but it is an open question the extent to which scale alone will be enough to overcome the inherent competition that comes as we add more modalities to a single model.

Through extensive experimentation, including over 250 individual experiments with seven modalities and model sizes ranging from 8 million to 30 billion, we have identified a scaling law that reflects the contributions of individual modalities and an additional term that captures the interaction between modalities (whether it be one of competition or synergy). We develop mixed-modal scaling laws that directly model competition between modalities and correctly predict data and model regimes where competition between modalities during training progresses into synergy. Specifically, we showed that our scaling laws correctly predicted the compute regime (30B model size, 45B token size), where we saw the complete reduction of modality competition for the Speech and Text modalities.

We also report a number of new empirical phenomena that arise during the training of mixed-modal models, including the tendency for the models to prioritize the optimization of a single modality at different stages of training. Our findings demonstrate that these phenomena can be primarily explained through the scaling law of interaction within the mixed-modal model. Additionally, we present new insights and guidelines for how to set key hyperparameters based on the terms of our scaling laws when optimal uni-modal hyper-parameters are known.

We develop neural scaling laws for mixed-modalities models that include text, speech, images, code, and their numerous couplings.

We discover a set of scaling laws describing the competition between arbitrary modalities.

We provide a simple recipe for selecting hyper-parameters in a multi-modal setting when optimal uni-modal hyper-parameters are known.

We uncover correlations between the scaling laws parameters we propose and various training phenomena, including training stability, optimal batch size, and coordinate ascent-like behavior in the optimization process across different modalities.

Related Work

Neural scaling laws quantify the relationship between model size, dataset size, compute budget, and performance, when training neural networks. This concept was introduced by Hestness et al. (2017), who observed a power law relationship and later scaled to much larger models by Kaplan et al. (2020).

Hoffmann et al. (2022) developed a unified formula for scaling laws, and provided recipes for compute-optimal training by adding data-dependent scaling terms unlike previous power law parameterizations. Other researchers have applied these principles to specific tasks and different parameterization of Transformers. Clark et al. (2022) examined the application of neural scaling laws to Mixture of Experts (MoE) models. Dettmers et al. (2022); Dettmers & Zettlemoyer (2022) studied the relationship between scaling laws and lower precision, which refers to using lower-precision data types, such as 16-bit floating point numbers, in neural networks. Gordon et al. (2021) and Ghorbani et al. (2021) applied these principles to Neural Machine Translation (NMT).

Additionally, Henighan et al. (2020) and Droppo & Elibol (2021) examined the application of neural scaling laws to generative language models in different modalities, including image generation and acoustic models. Cherti et al. (2022) also examined multi-modal training but did not specifically focus on generative models. To our knowledge, we are the first to investigate the phenomenon of interactions, competition, and interference between multiple modalities during training and provide a recipe for optimal mixed-modal training.

Interestingly, similar competition and scaling phenomenon have been observed for multi-lingual models. Conneau et al. (2019) observed a “curse of multilinguality,” where training in multiple languages can lead to interference between languages, resulting in decreased performance. Goyal et al. (2021) and Shaham et al. (2022) demonstrated that this interference could occur even on models much smaller than the available training data, but scaling up the model size can improve synergy and alleviate interference. These findings align with our findings in the mixed-modal scenario, suggesting that similar principles apply when training on multiple modalities.

Definitions

Modalities are traditionally distinguished by the data source, domain, or sensor affinity. For example, the code domain is typically seen as distinct from text due to the different data involved (e.g., GitHub vs. CommonCrawl). This also applies to auditory or visual modalities, which are captured with different sensors. Yet the decisions are not always clear, for example different languages are often all within the domain of the text. Given that we are studying neural scaling laws across modalities, we aim to have an empirically testable modality definition.

We define $\sigma-$membership of a set of samples, $D_{\alpha}$ through the following membership function.

We empirically define modality by comparing the perplexity of one data set to another. Suppose the perplexity of the secondary data set over the probability distribution of the primary set is greater than $\sigma$ times the mean perplexity of the primary set. In that case, we consider them to be distinct modalities. This definition distinguishes modalities by source, domain, sensor affinity, and language. We use the standard definition of perplexity ($\operatorname{ppl}$). Using this definition with $\sigma=3$, we decided to select seven modalities that we describe in detail below: Text, Image, Image-Text, Speech, Speech-Text, Code, Molecules.

Additionally, we define source modality as the type of token the sample contains, which within our setting will be; Text, Speech, or Image.

2 Uni-Modal Scaling Laws

We selected the Hoffmann et al. (2022) parameterization of scaling laws due to its precise representation of data factors and its additive nature, which allows for easy extension to multiple modalities. This parameterization (Equation 2) describes the loss based on the number of model parameters ($N$) and the number of tokens ($|D|$) through three constituent parts: the minimal achievable loss ($E$), the functional approximation error ($\frac{A_{j}}{N^{\alpha_{j}}}$), and the optimization or convergence error ($\frac{B_{j}}{|D_{j}|^{\beta_{j}}}$). These three factors are captured through seven learned parameters, providing a precise description of the loss.

It is well established that the upper bounds for $\beta$ and $\alpha$ are both $\frac{1}{2}$, which provides a clear understanding of how well transformers coupled with gradient descent algorithms scale in relation to the optimal scaling for each modality (Hoffmann et al., 2022).

Empirical Setting

For our text corpus, we use the same data as was used in OPT Zhang et al. (2022) for a total of 180B tokens. This dataset is primarily in English, although it contains other languages, as no explicit language filtering was done.

For all images, we convert them to discrete tokens using the Make-A-Scene visual tokenizer (Gafni et al., 2022), which gives 1024 tokens from an 8192 vocabulary per image. We select a custom subset of 600 million images across Schuhmann et al. (2022), and a custom image-text dataset scraped from Common Crawl. We remove all NSFW images and images that contain watermarks. Our Image dataset only contains the image and not the caption for a total of 614 billion tokens.

We utilize the Image dataset described above but align it with captions available from the image for a total of 690 Billion tokens. We call this our Image-Text dataset.

We used a combination of custom web-mined speech data and unlabeled speech in several public datasets. The web-mined speech dataset contains only unlabeled data in the form of long podcasts or news. We follow a series of preprocessing steps to improve the data quality and remove music and sensitive speech data. We also use a LangID model to select English-only speech. Our public data collection covers various speech styles and content topics, including LibriSpeech (Read-Books), CommonVoice in Read-Wiki, VoxPopuli from the Parliament domain, and Spotify Podcast and People’s Speech as web speech. Thanks to this combination, our Speech dataset offers a rich diversity.

Many public datasets also come with text aligned with speech. We take ASR and TTS data from Multilingual LibraSpeech and VoxPopuli and form the Speech-Text dataset.

We use the InCoder data (Fried et al., 2022).

We utilize the Simplified Molecular Input Line Entry System (SMILES, where the chemical’s structure is serialized into a string of symbols) representation from the Zinc dataset prepared by Chilingaryan et al. (2022).

2 Tokenization

Our mixed-modal generative models use a unified tokenization over all the modalities mentioned. This tokenizer processes data from all modalities into discrete tokens, which can be processed jointly by our model and trained with a single loss.

We use a Vector Quantized Variational autoencoders (VQGAN Esser et al. (2020)) model to tokenize image data into discrete tokens. The VQGAN model compresses each image into a grid of image tokens, where an encoder encodes each token into a vector. This process reduces the context size of the transformer by a factor of $3*X^{2}$, where $X$is the spatial reduction rate, or patch size, and 3 is the number of image channels. Online clustering is then performed, mapping each vector to the nearest entry of a learned codebook. We use a variant of the VQGAN from Gafni et al. (2022), which has a spatial reduction of 8 and a codebook size of 8192. This model is trained with extra perceptual losses to specific image regions, such as faces and salient objects, which improves the fidelity of the generated images. To be most effective in the language model stage, the visual tokenizer needs to effectively represent a image, and the correlated decoder needs to reconstruct the generated image tokens into high quality image data. We benchmark various image pretokenizers for those properties in Appendix A.3.1.

We use a Hidden-Unit BERT (HuBERT) Hsu et al. (2021) model for tokenizing our speech data. HuBERT is a self-supervised learning (SSL) model. It is trained to predict a masked subset of the speech signal using a mask language model objective, and has been found to be effective in learning a combined acoustic and language model over the continuous speech inputs. An offline clustering step to then used to generate discrete units. We use the Base HuBERT model in our work (model and training details see the appendix A.3.2). The final HuBERT units are generated through K-means clustering of the third iteration feature at the last layer, with a codebook size of 2000. Our HuBERT model encodes audio at 50Hz, and we compress a 16kHz audio by about 120 times, while effectively retaining specch information (Analysis see A.3.2).

Finally, we randomly sample 10 million sentences from all the data sets mentioned above and train a BPE model, where image and speech tokens take up a single token. We do an additional digit splitting for a vocab size of $2^{16}$ (Sennrich et al., 2016).

3 Model Architecture

We study the family of decoder-only models described in GPT-3 Brown et al. (2020) and OPT Zhang et al. (2022). We limit ourselves to training up to 6.7 billion-parameter models for all our uni-modal and bi-modal scaling laws and train up to 30B parameters to measure the generalizability of our scaling laws. For completeness, we present model architecture and their respective sizes in Table 2. We use learned positional encodings across all model architectures.

4 Causal Masking Objective

Instead of the traditional left-to-right causal language modeling objective, we use the causal masked objective from Aghajanyan et al. (2022). This provides a form of bidirectional context for sequence infilling, and also supports more aggressive generalization. For example, causally masked models trained only on data with text followed by images can still flip the ordering to generate images from text, since they were not strictly trained to predict tokens left to right. Recent work also shows that this masking does not hurt language modeling performance or the generative capacity of the models (Fried et al., 2022; Bavarian et al., 2022). We provide additional support for this claim in § A.2.

5 Training Procedure

All models were trained using the metaseqhttps://github.com/facebookresearch/metaseq code base, which includes an implementation of causal masking Zhang et al. (2022). The training used the PyTorch framework Paszke et al. (2019), with fairscale to improve memory efficiency through fully sharded model and optimizer states Baines et al. (2021). The training also uses Megatron-LM Tensor Parallelism Shoeybi et al. (2019) to support large model runs, and we use bf16 Kalamkar et al. (2019) to improve training stability. Given the large volume of data, we performed a single epoch of training, using each training document once. The batch size per GPU was determined based on the total world size of the experiment, the level of model parallelism, and the total target batch size in terms of the number of tokens. To ensure stable training, we applied gradient clipping with a maximum norm of 1.0 and used the Adam optimizer with $\beta_{1}=0.9$, $\beta_{2}=0.98$ Kingma & Ba (2015). We used the built-in polynomial decay learning rate scheduler in MetaSeq with 500 warmup updates and the end learning rate set to 10% of the peak learning rate.

We tracked all experiments using the Aim experiment tracker (Arakelyan et al., 2020). To ensure consistent training strategies across our experiments, we implemented a model restart policy using the Aim experiment tracker and callbacks. Specifically, if training perplexities do not decrease after 500 million tokens, the training run is restarted with a reduced learning rate with a factor of 0.8 of the current time step. This policy helps remove variance in the scaling laws due to differences in training procedures and allows us to scale up the number of asynchronous experiments significantly.

All experiments were conducted in a two-month time frame with a cluster of 768 80GB A100 GPUs. The majority of experiments used 64 GPUs at a time.

Scaling Laws

We first aim to discover scaling laws for each of the individual modalities we listed above. We train seven different model sizes, from 8 million to 6.7 billion, on seven different modalities on three different dataset sizes (5B, 10B, 100B).

In Figure 1, we share the training curves for all modalities and model sizes for the largest data size (100B tokens), and the final performance of all models in Figure 2. Overall, we see that scaling dynamics are fundamentally different across modalities, scale, and dataset size (which further reinforces our selection of dataset-size-dependent parameterization of scaling laws).

For each modality, we fit the seven parameters from Equation 2, following the procedure in Hoffmann et al. (2022). Specifically, we minimize

We then set $A_{j}=e^{a_{j}}$, $B_{j}=e^{b_{j}}$, $E_{j}=e^{e_{j}}$. In order to identify the optimal minima, we followed the method outlined by Hoffmann et al. (2022) and employed the L-BGFS algorithm on the same grid of initialization values. Our only deviation was using a higher value for the Huber loss parameter $\sigma$, which was necessary for generalization to held-out data in our multi-modal setting. The optimal values obtained were not located on the boundaries of the initialization grid.

The scaling laws for each modality are presented in Table 1. The parameters for each modality vary significantly. Some modalities, such as Code and Molecules, demonstrate more efficient use of the power of scale compared to others, such as Image. Our coefficients for Text are similar to those reported by Chinchilla, although it should be noted that we used a different dataset for our analysis. This accounts for any differences in the results.

2 Bi-Modal Scaling Laws

We also estimate scaling laws for training on two modalities: $\mathcal{L}\left(N,D_{i},D_{j}\right)$, where $N$ represents the model size, and $D_{i}$ and $D_{j}$ represent the two datasets being used. In the case where $D_{i}$ and $D_{j}$ are completely independent and have no mutual information between them, we expect the minimal achievable loss to be the average of the two monomodal scaling laws, given by $0.5*\left[\mathcal{L}\left(\infty,D_{i}\right)+\mathcal{L}\left(\infty,D_{j}\right)\right]$. This is because we are averaging over the loss and subsampling both test datasets equally ($|D_{i}|=|D_{j}|$). On the other hand, if there is some form of mutual information present between $D_{i}$ and $D_{j}$, we can expect the loss to be reduced by some maximal factor $C_{i,j}$. When considering finite model size and data regimes, there will be competition between the function approximation and optimization processes, which can be modeled using the same form as in Equation 2. We present our scaling law for mixed modal models in Equation 4.

An additional benefit to this parameterization is the additive or linear nature, which allows us to extend our parameterization to $n$-modal scaling laws.

We selected seven different pairs: Image-Text|Code, Image-Text|Speech-Text, Image-Text|Text, Speech|Text, Code|Text, Molecule|Code, and Speech|Code. While other couplings are available, we cannot do an exhaustive sweep due to computational constraints. We selected these pairs to maximize variety. For example, while Code|Text is known to perform well, Image-Text|Code may not offer as much benefit.

We create a dataset for each coupling and token target count where each subdataset contributes 50% of the tokens. We train using the same hyper-parameters as the uni-modal trainings and fit the scaling laws per modality coupling using the same procedure and optimization process (Equation 3).

We present the empirical results in Figure 3.

2.2 Breaking The Competition Barrier

Given these laws, we can now make predictions about what scale will be required to overcome modal competition and achieve synergy from training on each pair of modalities. By modality competition, we refer to the empirical phenomena of two modalities performing worse than if we trained two individual models on the same number of per-modality tokens. By synergy, we mean the inverse. We can define the notion of synergy formally through our scaling laws. If

we are reducing the loss beyond the independent modeling of the modalities and therefore are synergistic; otherwise, we say the modalities are in competition. When both sides of the inequality are equal, we call this the competition barrier for the two modalities. We present our extrapolated scaling laws with the predicted competition barrier in Figure 4.

We can then find the compute-optimal model size and token count that breaks the competition barrier by minimizing a compute cost over the competition barrier. We select the approximation from Kaplan et al. (2020).

For the Speech|Text coupling, the predicted compute optimal parameters are $N=28.35\textrm{B}$ and $D=45.12\textrm{B}$

To test this hypothesis, we select the closest architecture available from Zhang et al. (2022), which is the 30B parameterization and 50B tokens, slightly above the predicted data regime, to cover any error in our approximation. We train three models a, 350M, 2.7B, and 30B models on either Speech, Text, or Speech|Text. We plot the ratio of the average of the Speech and Text models perplexity per timestep by Speech|Text perplexity, the competition barrier and predictions from our scaling laws in Figure 5. As we see, the prediction does hold, and we achieve a model that crosses the competition barrier. Further scaling is likely to further improve the synergy, but we leave this exploration to future work.

Emergent Phenomena

We observed a number of emergent behaviors during training, many of which can be predicted from the modality-specific constants in our scaling laws. We briefly document these behaviors here; each is potentially worthy of study in future work.

Intermittent Coordinate Ascent Like Training: Different source modalities in a multi-modal setting are optimized at different paces, with some modalities even pausing their training progression for a significant amount of steps.

When looking at average perplexity over the dataset, the training dynamics are always consistently smooth and somewhat monotonically decreasing (Figure 1). But looking at the sub-perplexities of the modalities shows a different picture; certain modalities flatten out during training (see left figure in Figure 6).

In Figure 7, we plot the percent of the submodality that exhibits flatness, where flatness is defined as an area of the training curves where loss does not decrease (we do not count the warm-up period of optimization as part of this percentage).

Rate of Phenomena 1 Diminishes Past A Certain Scale: The rate of intermittent coordinate ascent-like training is correlated with scale ($N$) and $\alpha_{i,j}$.

Most of this intermittent coordinate ascent-like training can be reduced by simply increasing the model size. Intuitively, this makes sense as the increased functional approximation space should give the models enough capacity to simultanouesly optimize all of the modalities (Figure 7). Additionally, we discover that the empirically found $\alpha_{i,j}$, which describes the functional approximation cost across two modalities, is highly correlated with the uni-modal optimization flatness in the training regime. We found no correlation between $\beta_{i,j}$ and optimization flatness.

Optimal Batch Size for Modalities $i$ and $j$ is Correlated with $\beta_{i,j}$

We fixed the batch size to 1M tokens, but the question of the optimal batch for each modality and modality coupling remains. We train four versions of all over a subset of models, overall modalities, and selected couplings of modalities with batch sizes 1M, 2M, 4M, and 8M over 5B tokens, with the exception of modalities that contain Text for which we add 0.5M batch size experiments. We use the same training regime as mentioned in § 4.5. We present our results in Figure 8. Additionally for the bi-modal coupling experiments we plot $\log$ of the ratio between the optimal batch size and the sum of the optimal batch sizes for the sub-datasets against the $\beta_{i,j}$ of the discovered scaling laws in § 5.2. We found no correlation between $\alpha_{i,j}$ and optimal batch-size.

Rate of Deteriorating Training Dynamics is Correlated with $\alpha_{i,j}$ and $N$

The stability of training can be captured by looking at the total count of gradient norm spikes throughout the lifetime of the training. A large number of gradient spikes can indicate a poor training setting, from selecting the wrong learning rate or batch size to having low-quality data. Additionally, larger models tend to be harder to stabilize, reflecting in a larger amount of gradient spikes. We hypothesize that lower values of $\alpha_{i,j}$, reflecting higher competition between modalities, will correlate with more gradient norm spikes. We present the empirical correlation between $\log(N)/\alpha_{i,j}$ and # of Gradient Norm Spikes in Figure 9. We see a highly predictive relationship between model size ($N$) and the rate of mixed-modal competition ($a_{i,j}$ to the stability of the training run. We found no correlation between $\beta{i,j}$ and the # of gradient norm spikes.

Conclusion

We have provided extensive experimentation and analysis into the scaling properties of mixed-modal generative models. By developing a scaling law that reflects the contributions of individual modalities and the interaction between them, we have gained a deeper understanding of scaling mixed-modal models and the training dynamics of these models. Our findings also include a set of empirical phenomena observed during the training process and training dynamics that can be primarily explained through various interaction terms in our newly proposed scaling law. Additionally, we have developed guidelines for selecting critical hyper-parameters based on our scaling law, providing a valuable tool for practitioners in the field. Overall, our research has advanced the knowledge and understanding of mixed-modal generative models and will help develop unified models that can handle multiple modalities simultaneously.

Acknowledgements

We thank Hrant Khachatrian and Hrayr Harutyunyan for their discussions about the exact formulation of the mixed-modal scaling laws. We also thank Adam Polyak and Oran Gafni for training the Make-A-Scene tokenizer used in this work, and Rich James for editing the paper.

References

Appendix A Appendix

All models are trained with pre-norm and using ReLU activation. We apply a dropout of 0.1 throughout, but we do not apply any dropout to embeddings. We also use weight decay of 0.1. To initialize the weights, we use a variant based on Megatron-LM codebase, which involves using a normal distribution with a mean of zero and a standard deviation of 0.006. We truncate this normal distribution within two standard deviations and observed substantial gain in both training stability and performance.

A.2 Causal Masked vs. Causal Objective

We measure the impact of the choice of objective by conducting an additional scaling law on our Speech and Text datasets on the standard (causal) language modeling objective. Everything is kept constant except for the objective, including the training procedures. We present the empirically fit scaling law parameters in Table 3.

Note that both objectives optimize the joint probability of tokens; therefore, if there was a significant difference in our perplexity, we should expect to see it reflected in a difference in scaling law parameters. Instead, we see that the scaling laws seem to be close to identical, with whatever minor differences within the error of our approximation.

A.3 Tokenization

Modeling long-range dependencies with raw pixel input of an image (for example, total sequence length for a 256-pixel image in RGB form is 196608) is non-trivial, especially with transformers, which in their vanilla form scale poorly with sequence length. Recently, Vector Quantized Variational autoencoders (VQ-VAE, or Discrete-VAE) have been proposed, which learn discrete image representations, allowing a later generative model to generate images in the discrete latent space, just like a standard language model. VQ-VAE reduces the context size of a transformer by a factor of $3*X^{2}$ ($X$ is the spatial reduction rate, and 3 is the number of image channels), where information loss is unavoidable. VQ-VAE is trained to optimize the evidence lower bound of distribution of data. Esser et al. (2020) introduced VQGAN, which improves upon VQVAE by introducing an adversarial loss produced by a discriminator, reconstructing images with much higher quality. Recently, Gafni et al. (2022) trained a new image tokenizer with a better training objective focusing on faces or objects, which is adopted for this work and denoted as $\text{VQGAN}_{\text{MAS}}$. To be most effective in the later language model stage, the image tokenizer must represent an image effectively. The correlated decoder must reconstruct the generated image tokens into high-quality image data. We benchmark the following pre-trained tokenizers on these properties:

$\text{VQGAN}(fx,y)$ with different spatial reduction rate $fx$ and diffent vocab size $y$. For example, for a 256px image, 256 tokens will be created with a $\text{VQGAN}(f16)$ tokenizer and 1024 tokens with a $\text{VQGAN}(f8)$ tokenizer.

Our $\text{VQGAN}_{\text{MAS}}256$ and $\text{VQGAN}_{\text{MAS}}512$ use $f8$ and $f16$ spatial reduction, respectively, and have an 8192 vocab size. Our $\text{VQGAN}_{\text{MAS}}256$ is trained with a face-aware loss with the help of a pre-trained face embedding model. Our $\text{VQGAN}_{\text{MAS}}512$ is trained with face+object aware loss with extra downsampling and upsampling layer in the encoder and decoder to reconstruct images with higher resolution. Note, Our $\text{VQGAN}_{\text{MAS}}512$ with 512x512 image input compresses the image to 1024 tokens, thanks to the downsampling layer.

One way to quantify the realism captured by these models is to compute Fréchet Inception Distance (FID) scores of reconstructed images w.r.t. the inputs (R-FIDs). Table 4 shows R-FIDs when reconstructing the whole validation split of the ImageNet dataset. For an image with a 256-pixel resolution, reducing spatial reduction rate or increasing visual vocab size can help achieve lower R-FIDs. Our $\text{VQGAN}_{\text{MAS}}256$ model is superior to its counterpart with the same spatial reduction rate and vocab size, demonstrating the effectiveness of the extra face-aware loss. Interestingly, $\text{VQGAN}_{\text{MAS}}256$ gets a higher R-FID than $\text{VQGAN}_{\text{MAS}}512$, consistent with the result in the original paper. We are also interested in understanding how much the reconstruction process can retain information and if we lose critical image information. We benchmark the representation power via classification accuracy with a pretrained model. We first reconstruct all the images in the ImageNet validation set with different tokenizers, similar to R-FID computation. Then a trained pretrained classifier on ImageNet is used to run inference on the original and reconstructed images. The classification accuracy of original images with 256 or 512-pixel resolution is 81.56 and 82.89, respectively. The accuracy@1 on the reconstructed images and their gap with the raw images are reported in Table 4. Images reconstructed by the $\text{VQGAN}_{\text{MAS}}256$ can best maintain the original information with less than 2 percent degradation in accuracy.

For qualitative comparison, we give all tokenizers except $\text{VQGAN}_{\text{MAS}}512$ an image with 256x256 pixels. $\text{VQGAN}_{\text{MAS}}512$ reconstructs a 512-pixel resolution image and resizes it to 256-pixel resolution for plotting purposes. $\text{VQGAN}~{}(f16)$ a s produce 256 tokens, while $\text{VQGAN}~{}(f8)$ and $\text{VQGAN}_{\text{MAS}}$ models produce 1024 tokens.

We randomly sample two images from ImageNet (top 2 rows in Figure 10). All reconstructed images can maintain vital information about the image and the textures. With a high reduction rate (192), VQGAN with $f16$ spatial reduction can not reproduce every detail of its input but tends to hallucinate parts of it, for example, the eye and the tail of the dog in row 1 and the mirror of the blue car in row 2. By increasing the vocab size, more realistic images can be generated. With a decreased compression rate, the $\text{VQGAN}~{}(f8)$ model and $\text{VQGAN}_{\text{MAS}}$ produce much more realistic reconstructed images. For example, in row 2, the door handle, the clouds, and the mirror’s tree are successfully reconstructed with great detail.

Lastly, we reconstruct images from a textbook (row 4 in Figure 10) or screenshots of tables from scientific papers (row 5 in Figure 10). All models struggle to reconstruct the original image, except $\text{VQGAN}_{\text{MAS}}$ models. Figure 10 shows impressive results by $\text{VQGAN}_{\text{MAS}}$ models that all text and numbers are human readable. $\text{VQGAN}_{\text{MAS}}256$ produces sharper edges while $\text{VQGAN}_{\text{MAS}}512$ smooths things out.

From all the above examples, the reduced spatial reduction is effective for better tokenization; however, it results in a longer token sequence. Another way to increase image representation is to increase the pixel numbers of images. We reconstruct images with a size of 512x512 for $\text{VQGAN}_{\text{MAS}}512$ in Figure 11. $\text{VQGAN}(f16)$ produce 1024 tokens, while $\text{VQGAN}(f8)$ and $\text{VQGAN}_{\text{MAS}}$ models produce 4096 tokens. $\text{VQGAN}_{\text{MAS}}256$ in Figure 10 outperform VQGAN (f16) by a big margin. With the same token budget, decreasing spatial reduction is more effective than increasing image pixels.

A.3.2 Details of Speech Tokenization

We use the Base HuBERT model in our work. This model comprises a convolutional encoder and 12 layer Transformer, each with an embedding dimension of 768, a feed-forward layer dimension of 3072, and 12 self-attention heads. Pre-training of the model has been performed on 32 GPUs over three iterations, with 400K updates per iteration. The training data consists of 221K hours of unlabeled speech from multilingual Librispeech (MLS) Pratap et al. (2020), Common Voice (CV) Ardila et al. (2019), and VoxPopuli (VP) Wang et al. (2021) in eight languages (English, Spanish, French, German, Dutch, Italian, Polish, Portuguese). The MFCC/6-th layer feature from iteration 1 and the 9-th layer feature from iteration 2 are used as targets, with codebook sizes of 100/500/1000, respectively, following the methodology outlined in Lee et al. (2021).

A typical 16kHz audio with a bit depth of 16 has a bitrate of 64kbps. HuBERT encodes audio at 50Hz with a codebook size of 2000, resulting in a bitrate of 548bps. The effective compression rate is roughly 117. Our model still effectively retains speech information, as shown in Table 5. We compare the word error rate (WER) of a pretrained automatic speech recognition (ASR) model with original audio or reconstructed audio by HuBERT models. We present results with two HuBERT models, one public Hsu et al. (2021) version (HuBERT public) and one trained by us (HuBERT ours). WER of the original audio on LJSpeech is 2.04, the audio reconstructed by HuBERT public degrades by 0.94, while the audio reconstructed by our HuBERT only degrades it by 0.3. A similar phenomenon is observed on the LibriSpeech dataset, where our HuBERT model improves upon HuBERT public and can effectively reconstruct audio with very little information loss.

Appendix B Credit

Armen Aghajanyan: Proposed the original idea, co-authored the ablation plan, executed all the training runs and scaling law research, and was the primary writing author of the paper.

Lili Yu: Core contributor to mixed-modal evaluations framework, drove the selection of speech/image/text tokenizer, secondary writing author of the paper.

Alexis Conneau: Drove high-level direction, co-authored ablation plan, and collected speech datasets.

Wei-Ning Hsu: Provided day-to-day feedback on speech-language model training, trained speech tokenizer, and tokenized all the speech data used throughout the project.

Karen Hambardzumyan: Helped in the core design of scaling laws and writing.

Susan Zhang, Stephen Roller, Naman Goyal: Provided support for the general training of all the models and the metaseq framework.

Omer Levy: Provided, developed, and distilled the story for this paper. Edited paper as well.

Luke Zettlemoyer: Provided support and feedback throughout the whole lifetime of the project. Provided help writing the paper as well as significant feedback for the paper.