SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Controlled Experiments Table 1 lists bpb on each dataset. Each model is trained for 80 billion bytes, and models are scaled to use the same compute budget. We carefully tune hyperparameters for all architectures to best utilize the available compute budget. MEGABYTE consistently outperforms both transformers and PerceiverAR across all datasets. We use the same sets of parameters on all dataset. In all experiments presented in Table 1, transformer has size of 320M with context length of 1024, PerceiverAR has size of 248M with context size of 8192 and latent size of 1024, and MEGABYTE global/local model sizes are 758M/262M with context length of 8192 and patch size of 8.
Scaling Experiment We scale up our training data on PG-19 (Table 2), and compare MEGABYTE with byte baselines, as well as converting all results to word-level perplexities to benchmark with state-of-art token based models. We train a byte-level Transformer, PerceiverAR and MEGABYTE models for 400B bytes and the same compute budget using same model parameters as in the controlled experiments. We find that MEGABYTE outperforms other byte-level models by a wide margin at this scale.3 We also compare with the best previously reported numbers for sub-word models. These results may be confounded by differing amounts of compute and tuning used, but show that MEGABYTE gives results competitive with state-of-the-art models trained on subwords. These results suggest that MEGABYTE may allow future large language models to be tokenization-free.
6 Image Modeling
Sequence Modeling on ImageNet We test MEGABYTE on variants of the autoregressive image generation task on ImageNet (Oord et al., 2016), to measure its ability to efficiently use long context. We test on three different resolutions of images, ranging from  $64 \times 64$  to  $640 \times 640$  pixels – the latter
3The only prior byte-level experiments we are aware of are at a smaller scale in Hutchins et al. (2022), who report results equivalent to test perplexities of 46.5 with a version of the BlockRecurrent transformer, and 49.5 with Memorizing Transformers Wu et al. (2022), compared to 36.4 with our model.
requiring the effective modeling of sequences with over 1.2M tokens. This generation task becomes increasingly challenging as the image's resolution grows: doing well on this task requires the modeling of local patterns (textures, lines, etc.) and long-range context that provides information about the high level structure of the image. Inspired by recent works in Vision Transformers (Dosovitskiy et al., 2020), we model image data patch by patch (more details can be found in Appendix D.1).
Comparison with State of the Art We train a large MEGABYTE model on ImageNet 64x64 with Global and Local models sized 2.7B and 350M parameters, respectively, for 1.4T tokens. We estimate that training this model consumed less than half the GPU hours we would have needed to reproduce the best PerceiverAR model described by (Hawthorne et al., 2022). As shown in Table 2, MEGABYTE matches the state-of-the-art performance of PerceiverAR whilst using only half the compute.
ImageNet64	bpb
Routing Transformer (Roy et al., 2020)	3.43
Combiner (Ren et al., 2021)	3.42
Perceiver AR (Hawthorne et al., 2022)	3.40
MEGABYTE	3.40
Table 3: Bits per byte (bpb) on ImageNet  ${64} \times  {64}$  . MEGABYTE matches the current state-of-the-art while only using half the amount of GPU hours to train.
Context	Image64	Image256	Image640
Total len	12288	196608	1228800
Transformer	1024	3.62	3.801	2.847
Perceiver AR	12000	3.55	3.373	2.345
MEGABYTE	Full	3.52	3.158	2.282
Table 4: Bits per byte (bpb) on ImageNet with different resolutions. All models use the same compute and data. MEGABYTE scales well to sequences of over 1M tokens.
Scaling to higher resolutions We compare three transformer variants (vanilla, PerceiverAR, MEGABYTE) to test scalability to long sequences on increasingly large image resolutions. We use our own implementations of these in the same framework and budget the same amount of GPU hours and data to train each of these model variants.
MEGABYTE is able to handle all sequence lengths with a single forward pass of up to 1.2M tokens. We found neither the standard Transformer nor PerceiverAR could model such long sequences at a reasonable model size, so instead we split images into segments of size 1024 and 12000 respectively. For Megabyte, we set patch size as 12 for Image64 and patch size as 192 for Image256 and Image640 datasets. Model sizes are adjusted to match overall training speeds across models and we do not use any form of sliding window evaluation in this experiment. As seen in Table 4, MEGABYTE outperforms baselines across all resolutions in this compute-controlled setting. The precise settings used for each of the baseline models such as context length and number of latents are summarized in Table 12. Results show that MEGABYTE outperforms the other systems at all resolutions, demonstrating an effective model of sequences of over 1M bytes.
7 Audio Modeling
Audio has aspects of both the sequential structure of text and the continuous nature of images, so is an interesting application for MEGABYTE.
Raw audio is typically stored as a sequence of 16-bit integer values (one per timestep); a softmax layer would need to output 65,536 probabilities per timestep to model all possible values. To address this issue, various techniques have been developed to reduce the memory and computational requirements of the softmax layer. For instance, van den Oord et al. (2016) apply  $\mu$ -law companding transformation and quantizes the input into 256 possible values. Alternatively, van den Oord et al. (2017) model the samples using the discretized mixture of logistics distribution introduced by Salimans et al. (2017). Finally, Kalchbrenner et al. (2018) use a dual softmax technique to produce 8 coarse and 8 fine bits. In our approach, we simplify the audio modeling process by directly reading the bytes (256 possible values) from the audio file and conducting an autoregressive language model on top of that. This greatly streamlines the modeling process, making it easier and more efficient.
Our audio modeling approach focuses on  $16\mathrm{kHz}$ , 16-bit audio, which equates to  $32\mathrm{k}$  bytes per one-second clip. We use an extensive audio dataset consisting of 2 terabytes (roughly 18,000 hours) of audio. We use a sequence length of 524,288, a patch size of 32, and a batch size of 32 to facilitate model training. By utilizing these settings, we can effectively train our model on large volumes of audio data, helping to improve its accuracy and efficacy. Our model obtains bpb of 3.477, much
Global Size	(Local) Size	bpb	Generation Time (s)
Transformer	-	350M	1.064	132
MEGABYTE	1.3B	218M	0.991	93
lower than the results with perceiverAR (3.543) and vanilla transformer model (3.567). More ablation results are presented in Table 6.
8 Analysis
We study different behaviors of MEGABYTE. All experiments in the same group use the same compute.
Generation speed We also compare the text generation speed between MEGABYTE and a transformer. We compare a 350M parameter baseline transformer and a MEGABYTE model with a 1.3B parameter Global model and a 218M parameter local model, trained on PG19 with equal compute. As shown in Table 5, the MEGABYTE model achieves much lower perplexity as expected. However, MEGABYTE also generates a sequence of 8192 tokens  $40\%$  faster than transformer, despite having over 4 times the parameters. This speed up is due to the bulk of the parameters being in the Global model, which only needs to be computed once for every 8 tokens, whereas all the parameters in the baseline model are used on every token.
Model Components In Table 6, we analyze the significance of different components in the MEGABYTE architecture by studying arXiv, Librilight-L and ImageNet256 datasets. Removing Local (w/o local model) or global (w/o global model) model, we observe a substantial increase in bpb on all datasets, showing that both parts are crucial. The performance of the model without the cross-patch local model (w/o cross-patch local model) is competitive, indicating that the architecture is robust to this modification. We observe slight improvement on the Librilight-L and ImageNet256 datasets by augmenting the MEGABYTE model with a CNN encoder (w/ CNN encoder). This suggests that the MEGABYTE architecture can benefit from integrating alternative encoding mechanisms.
Effective Use of Context Long-context models often struggle to benefit from the full context (Sun et al., 2021). Figure 7 shows that later tokens within each context window have a higher likelihood, indicating that MEGABYTE can effectively use at least 8k bytes of context on the PG19 dataset.
Table 5: Comparison of bits per byte (bpb) and generation speed of 8192 bytes of transformer model (with context length 1024) and MEGABYTE with context length 8192 and patch size 8.
Arxiv	Audio	Image256
MEGABYTE	0.6871	3.477	3.158
w/o local model	1.263	5.955	4.768
w/o global model	1.373	3.659	3.181
w/o cross-patch attention	0.6781	3.481	3.259
w/ CNN encoder	0.6871	3.475	3.155
Table 8: An illustration of strided inference with patch size 8. Blue and yellow represents two inferences that are shifted by half patch size. Solid line indicates final probability being taking during strided inference.
Table 7: Average log probability assigned to different positions within the context length by MEGABYTE and by a vanilla transformer model on PG19 test set.
Table 6: Ablation of MEGABYTE model components. Models with the same dataset are trained using the same compute. The hyperparameters are listed in Table 12.
Method	Inference Cost	bpb
Basic Inference	1X	0.9079
w/ Sliding Window	2X	0.8918
w/ Strided Inference	2X	0.8926
w/ Sliding & Strided	4X	0.8751
Table 9: Performance of various inference techniques on the PG19 test set using our best MEGABYTE model.
Strided Inference We find that within a single patch, on average, the MEGABYTE performs worse on later tokens within a patch (see Figure 8). Section 2.3 proposes strided inference as a solution, where two forward passes are performed offset by  $\frac{P}{2}$  tokens, and results from the first half of each patch are combined. Table 9 shows performance improvements from strided inference, which are additive with the standard sliding window.
Patch Size. We experimented with various patch sizes on Image256 dataset and found a wide range of values where MEGABYTE performs similarly. We found similar robustness to patch size choices across all modalities, although the optimal patch size itself can be different across modalities.
Local to Global model Size Ratio. We experimented with different Local/Global model size ratios on PG19 dataset. By grouping bytes into patches, MEGABYTE effectively uses  $P$  times less tokens for the Global model as on the Local model—enabling us to increase the size of the Global model with reduced cost. We find that a given compute budget is spent optimally when the Global model is larger than the Local model, consistently across all modalities and various patch sizes.
Patch	Global Size	Local Size	bpb
48	125M	114M (L=11)	3.178
192	125M	125M (L=12)	3.158
768	125M	83M (L=8)	3.186
Table 10: Effects of patch size on performance on the Image256 dataset. All versions use the same amount of GPU hours and data.
Global Size	Local Size	bpb
350M (D=1024,L=24)	290M (D=1024,L=20)	1.014
760M (D=1536,L=24)	262M (D=1024,L=18)	1.002
1.3B (D=2048,L=24)	218M (D=1024,L=15)	0.991
Table 11: Effects of Local / Global model size on the PG19 dataset. Increasing the capacity of global model improves performance. Models are compute and data matched.
Prior research has explored the possibility of improving the efficiency of Transformers on long sequences, primarily motivated by mitigating the quadratic cost of self-attention.
Efficient Encoder Models Several related techniques to ours have been developed for transformer encoder architectures but cannot be straightforwardly applied to decoders. In particular, patchifying operations have previously been used in image encoder models such as ViT (Dosovitskiy et al., 2020), and down- and up-sampling operations have been used for text encoders (Clark et al., 2022), but such methods cannot be naively applied to decoder-only models without leaking information to future bytes in the same patch. MEGABYTE generalizes these approaches to an efficient decoder model by using a intra-patch transformer to predict each sequence element's likelihood, and offsetting the inputs to the two models to avoid leaking information. Jaegle et al. (2021) use self-attention on a shorter latent sequence also resembles patchification, but this technique cannot easily be applied to decoder architectures without leaking information to future timesteps.
Efficient Decoder models Improving the efficiency of decoder models is harder because of the need to make one prediction per timestep, and not leak information to future timesteps. The most popular approaches can be categorized as (1) chunking sequences into smaller blocks, and propagating information from previous blocks with either recurrence (Dai et al., 2019; Hutchins et al., 2022) or cross-attention (Hawthorne et al., 2022), (2) linear alternatives to attention, which typically involve forms of token-level recurrence (Katharopoulos et al., 2020) or state space models (Gu et al., 2021; Smith et al., 2022; Ma et al., 2022), or (3) sparse approximations of attention (Kitaev et al., 2020; Beltagy et al., 2020; Child et al., 2019; Wu et al., 2022). However, the performance of dense attention means it is typically still chosen for large scale decoders (Touvron et al., 2023; Chowdhery et al., 2022). MEGABYTE takes the alternative approach of decomposing the complete sequence into two shorter sequences, giving sub-quadratic attention. We also note that feedforward networks are the dominant cost in large decoders, not self-attention. Our approach to compressing sequences allows much larger models than would be possible when using large feedforward networks at every timestep.
Tokenization The most common approach to shortening sequence lengths in Transformer decoders is to pre-process the input with a form of tokenization, in which multiple bytes are mapped to a single discrete token from a fixed vocabulary. For text, this can be done losslessly using methods such as BPE (Sennrich et al., 2015) and SentencePiece (Kudo & Richardson, 2018), but these approaches can require language-specific heuristics (Radford et al., 2019), limit out-of-domain performance (Sharami
et al., 2023), and can affect prompting and truncated sampling in unpredictable ways. The amount of high-frequency information in images and audio means that tokenization cannot be performed losslessly, and instead clustering (Hsu et al., 2021) or discrete auto-encoders (Ramesh et al., 2021) are used to compress the inputs, which lose information and likely limit generative model performance. Our patches are analogous to traditional lossless tokens, and the Local model performs the role of mapping a hidden state to a distribution over possible patches.
We introduced MEGABYTE, a scalable architecture for modeling long sequences. MEGABYTE outperforms existing byte-level models across a range of tasks and modalities, allowing large models of sequences of over 1 million tokens. It also gives competitive language modeling results with subword models, which may allow byte-level models to replace tokenization. However, the scale of experiments here is far below those of state-of-the-art language models (Brown et al., 2020), and future work should explore scaling MEGABYTE to much larger models and datasets.
Baines, M., Bhosale, S., Caggiano, V., Goyal, N., Goyal, S., Ott, M., Lefaudeauux, B., Liptchinsky, V., Rabbat, M., Sheiffer, S., Sridhar, A., and Xu, M. FairScale: A general purpose modular PyTorch library for high performance and large scale training. https://github.com/facebookresearch/fairscale, 2021.
Beltagy, I., Peters, M. E., and Cohan, A. Longformer: The long-document transformer. arXiv preprint arXiv:2004.05150, 2020.
Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
Child, R., Gray, S., Radford, A., and Sutskever, I. Generating long sequences with sparse transformers. arXiv preprint arXiv:1904.10509, 2019.
Choromanski, K., Likhosherstov, V., Dohan, D., Song, X., Gane, A., Sarlos, T., Hawkins, P., Davis, J., Mohiuddin, A., Kaiser, L., et al. Rethinking attention with performers. arXiv preprint arXiv:2009.14794, 2020.
Chowdhery, A., Narang, S., Devlin, J., Bosma, M., Mishra, G., Roberts, A., Barham, P., Chung, H. W., Sutton, C., Gehrmann, S., et al. Palm: Scaling language modeling with pathways. arXiv preprint arXiv:2204.02311, 2022.
Clark, J. H., Garrette, D., Turc, I., and Wieting, J. Canine: Pre-training an efficient tokenization-free encoder for language representation. Transactions of the Association for Computational Linguistics, 10:73-91, 2022.
Dai, Z., Yang, Z., Yang, Y., Carbonell, J., Le, Q. V., and Salakhutdinov, R. Transformer-xl: Attentive language models beyond a fixed-length context, 2019. URL https://arxiv.org/abs/1901.02860.
Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.11929, 2020.
Gao, L., Biderman, S., Black, S., Golding, L., Hoppe, T., Foster, C., Phang, J., He, H., Thite, A., Nabeshima, N., Presser, S., and Leahy, C. The pile: An 800gb dataset of diverse text for language modeling, 2020.
Gu, A., Goel, K., and Ré, C. Efficiently modeling long sequences with structured state spaces. arXiv preprint arXiv:2111.00396, 2021.
4For example, whether or not a prompt should end in whitespace depends on details of the subword algorithm.
A.1 Training Details
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 0. 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 Motivation
Why is the local model needed? Many of the efficiency advantages of the MEGABYTE design could be realized with the Global model alone, which would resemble a decoder version of ViT (Dosovitskiy et al., 2020). However, the joint distribution over the patch  $p(x_{t+1}, \dots, x_{t+P} | x_{0:t})$  has an output space of size  $256^P$  so direct modeling is only tractable for very small patches. We could instead factor the joint distribution into conditionally independent distributions  $p(x_{t+1} | x_{0:t}) \dots p(x_{t+P} | x_{0:t})$ , but this would greatly limit the model's expressive power. For example, it would be unable to express a patch distribution such as  $50\%$  cat and  $50\%$  dog, and would instead have to assign probability mass to strings such as  $cag$  and  $dot$ . Instead, our autoregressive Local model conditions on previous characters within the patch, allowing it to only assign probability to the desired strings.
Increasing Parameters for Fixed Compute Transformer models have shown consistent improvements with parameter counts (Kaplan et al., 2020). However, the size of models is limited by their increasing computational cost. MEGABYTE allows larger models for the same cost, both by making self attention sub-quadratic, and by using large feedforward layers across patches rather than individual tokens.
Re-use of Established Components MEGABYTE consists of two transformer models interleaved with shifting, reshaping and a linear projection. This re-use increases the likelihood that the architecture will inherit the desirable scaling properties of transformers.
A.3 Model Details
As discussed in Section 4, we conduct experiments using a fixed compute and data budget across all models to focus our comparisons solely on the model architecture rather than training resources. To achieve this, we adjust model hyperparameters within each architecture so that the time taken for a single update is matched and then train all models for the same number of updates. We list all of model details in Table 12 and Table 13.
Model	#L	dmodel	#H	head
S1	125M	12	768	12	64
S2	350M	24	1024	16	64
S3	760M	24	1536	16	96
S4	1.3B	24	2048	32	64
S5	2.7B	32	2560	32	80
S6	6.7B	32	4096	32	128
Table 12: Common Model architecture details by size. For each model size, we show the number of layers (#L), the embedding size ( $d_{\text{model}}$ ), the number of attention heads (#H), the dimension of each attention head ( $d_{\text{head}}$ ).
Model	(Global) Size	Local Size	BS	LR	Context Length (in bytes)
arXiv
Transformer	320M (D=1024, L=22)	N/A	72	2.00E-04	1,024
Perceiver AR	248M (D=1024, L=17)	N/A	72	2.00E-04	8,192 (1024 latents)
MEGABYTE	758M (D=2048, L=14)	262M (D=1024, L=18)	48	2.00E-04	8,192 (patch size 8)
w/o Local model	2.3B (D=2560, L=20)	N/A	48	1.50E-04	8,192 (patch size 4)
w/o global model	N/A	350M (D=1024, L=24)	192	2.00E-04	8,192 (patch size 8)
w/o cross-patch Local model	921M (D=2048, L=17)	350M (D=1024, L=24)	48	2.00E-04	8,192 (patch size 8)
w/ CNN encoder	704M (D=2048, L=13)	262M (D=1024, L=18)	48	2.00E-04	8,192 (patch size 8)
Image task 64 (Table 2)
MEGABYTE	2.7B (D=2560, L=32)	350M (D=1024, L=24)	2	2.00E-04	12,288 (patch size 12)
Image task 64 (Table 4)
Transformer	760M (D=1536, L=24)	N/A	512	3.00E-04	2,048
Perceiver AR	227M (D=1024, L=16)	N/A	512	3.00E-04	12,288 (1024 latents)
MEGABYTE	1.3B (D=2048, L=24)	1.3B (D=2048, L=24)	256	3.00E-04	12,288 (patch size 12)
Image task 256
Transformer	62M (D=768, L=6)	N/A	1536	2.00E-04	1,024
Perceiver AR	62M (D=768, L=6)	N/A	256	2.00E-04	8,192 (768 latents)
MEGABYTE	125M (D=768, L=12)	125M (D=768, L=12)	16	2.00E-04	196,608 (patch size 192)
w/o local model	2.7B (D=4096, L=32)	N/A	16	2.00E-04	196,608 (patch size 48)
w/o global model	125M (D=768, L=12)	125M (D=768, L=12)	16	2.00E-04	196,608 (patch size 192)
w/o cross-patch Local model	250M	156M (D=768, L=15)	16	2.00E-04	196,608 (patch size 192)
w/ CNN encoder	125M (D=768, L=12)	125M (D=768, L=12)	16	2.00E-04	196,608 (patch size 192)
Image task 640
Transformer	83M (D=768, L=8)	N/A	4800	3.00E-04	1,024
Perceiver AR	62M (D=768, L=6)	N/A	2048	3.00E-04	4,096 (1024 latents)
MEGABYTE	125M (D=768, L=12)	83M (D=768, L=8)	32	3.00E-04	1,228,800 (192 patch size)
audio
Transformer	135M (D=768, L=13)	N/A	2048	2.00E-04	1024
Perceiver AR	62M (D=768, L=6)	N/A	384	2.00E-04	8,192 (1024 latents)
MEGABYTE	350M (D=1024, L=24)	125M (D=768, L=12)	256	2.00E-04	524,288 (32 patch size)
w/o local model	2.7B (D=4096, L=32)	125M (D=768, L=12)	256	2.00E-04	524,288 (32 patch size)
w/o global model	350M (D=1024, L=24)	125M (D=768, L=12)	256	2.00E-04	524,288 (32 patch size)
w/o cross-patch Local model	350M (D=1024, L=24)	146M (D=768, L=14)	256	2.00E-04	524,288 (32 patch size)
w/ CNN encoder	350M (D=1024, L=24)	125M (D=768, L=12)	256	2.00E-04	524,288 (32 patch size)
Table 13: Model architecture details. We report the model size, the embedding size (D), number of layers(L), total batch size (BS), learning rate(LR), and context length. When we vary the number of model layers from the standard amount for the given size (Table 12), we note this accordingly. For PerceiverAR models, we note the number of latents used, and for MEGABYTE models we note the patch sizes.
B Pseudocode