SlowGuess/NeurIPS_2023_OCR huggingface dataset:
MEGABYTE: Modeling Million-byte Sequences with Multiscale Transformers
Lili Yu*
Daniel Simig*
Colin Flaherty*
Armen Aghajanyan
Luke Zettlemoyer Mike Lewis
Autoregressive transformers are spectacular models for short sequences but scale poorly to long sequences such as high-resolution images, podcasts, code, or books. We propose MEGABYTE, a multi-scale decoder architecture that enables end-to-end differentiable modeling of sequences of over one million bytes. MEGABYTE segments sequences into patches and uses a local submodel within patches and a global model between patches. This enables sub-quadratic self-attention, much larger feedforward layers for the same compute, and improved parallelism during decoding—unlocking better performance at reduced cost for both training and generation. Extensive experiments show that MEGABYTE allows byte-level models to perform competitively with subword models on long context language modeling, achieve state-of-the-art density estimation on ImageNet, and model audio from raw files. Together, these results establish the viability of tokenization-free autoregressive sequence modeling at scale.
Sequences of millions of bytes are ubiquitous; for example, music, image, or video files typically consist of multiple megabytes. However, large transformer decoders (LLMs) typically only use several thousand tokens of context (Brown et al., 2020; Zhang et al., 2022a)—both because of the quadratic cost of self-attention but also, more importantly, the cost of large feedforward networks per-position. This severely limits the set of tasks where LLMs can be applied.
We introduce MEGABYTE, a new approach to modeling long byte sequences. First, byte sequences are segmented into fixed-sized patches, loosely analogous to tokens. Our model then consists of three parts: (1) a patch embedder, which simply encodes a patch by losslessly concatenating embeddings of each byte, (2) a global module, a large autoregressive transformer that inputs and outputs patch representations and (3) a local module, a small autoregressive model
Figure 1: Overview of MEGABYTE with patch size  $P = 4$ . A small local model autoregressively predicts each patch byte-by-byte, using the output of a larger global model to condition on previous patches. Global and Local inputs are padded by  $P$  and 1 token respectively to avoid leaking information about future tokens.
that predicts bytes within a patch. Crucially, we observe that for many tasks, most byte predictions
$$
h _ {t} ^ {\text {e m b e d}} = E _ {x _ {t}} ^ {\text {g l o b a l - e m b e d}} + E _ {t} ^ {\text {p o s}} \quad t \in [ 0.. T), E ^ {\text {g l o b a l - e m b e d}} \in \mathbb {R} ^ {V \times D _ {G}},
$$
$$
E ^ {\text {p o s}} \in \mathbb {R} ^ {T \times D _ {G}}, h ^ {\text {e m b e d}} \in \mathbb {R} ^ {T \times D _ {G}}
$$
$$
\begin{array}{l l} h _ {k} ^ {\text {g l o b a l - i n}} & = \left\{ \begin{array}{l l} E ^ {\text {g l o b a l - p a d}}, & \text {i f} k = 0, \\ h _ {((k - 1) \cdot P): (k \cdot P)} ^ {\text {e m b e d}}, & k \in [ 1,.., K), \end{array} \right. \end{array} \qquad E ^ {\text {g l o b a l - p a d}} \in \mathbb {R} ^ {P \times D _ {G}}, K = \frac {T}{P}
$$
$$
h _ {0: K} ^ {\text {g l o b a l - o u t}} = \text {t r a n s f o r m e r} ^ {\text {g l o b a l}} \left(h _ {0: K} ^ {\text {g l o b a l - i n}}\right) \quad h ^ {\text {g l o b a l - o u t}}, h ^ {\text {g l o b a l - i n}} \in \mathbb {R} ^ {K \times P \cdot D _ {G}}
$$
$$
\begin{array}{r l} {h _ {k, p} ^ {\mathrm {l o c a l - i n}}} & {= w ^ {\mathrm {G L}} h _ {k, (p \cdot D _ {G}): ((p + 1) \cdot D _ {G})} ^ {\mathrm {g l o b a l - o u t}} + \left\{ \begin{array}{l l} E _ {x _ {(k \cdot P + p - 1)}} ^ {\mathrm {l o c a l - p a d}}, & \text {i f} p = 0 \\ E _ {x _ {(k \cdot P + p - 1)}} ^ {\mathrm {l o c a l - e m b e d}}, & p \in [ 1,,.., P) \end{array} \right. \qquad E _ {x _ {(k \cdot P + p - 1)}} ^ {\mathrm {l o c a l - p a d}} \in \mathbb {R} ^ {D _ {L}}, w ^ {\mathrm {G L}} \in \mathbb {R} ^ {D _ {G} \times D _ {L}}} \\ & {\qquad E _ {x _ {(k \cdot P + p - 1)}} ^ {\mathrm {l o c a l - e m b e d}} \in \mathbb {R} ^ {V \times D _ {L}}} \end{array}
$$
$$
\begin{array}{l l} h _ {k, 0: P} ^ {\text {l o c a l - o u t}} & = \text {t r a n s f o r m e r} ^ {\text {l o c a l - i n}} (h _ {k, 0: P} ^ {\text {l o c a l - i n}}) \end{array} \quad h _ {k, p} ^ {\text {l o c a l - i n}} \in \mathbb {R} ^ {D _ {L}}, h ^ {\text {l o c a l - o u t}} \in \mathbb {R} ^ {K \times P \cdot D _ {L}}
$$
$$
p (x _ {t} | x _ {0: t}) = \operatorname {s o f t m a x} (E ^ {\mathrm {l o c a l - e m b e d}} h _ {k, p} ^ {\mathrm {l o c a l - o u t}}) _ {x _ {t}} \qquad \qquad t = k \cdot P + p
$$
Figure 2: Summary of MEGABYTE with vocabulary  $V$ , sequence length  $T$ , global and local dimensions  $D_G$  and  $D_L$ , and  $K$  patches of size  $P$ . Transformer layers use masked self attention to not observe future timesteps.
are relatively easy (for example, completing a word given the first few characters), meaning that large networks per-byte are unnecessary, and a much smaller model can be used for intra-patch modelling.
MEGABYTE has three main advantages over Transformers for long sequence modeling:
1. Sub-quadratic self-attention Most work on long sequence models has focused on mitigating the quadratic cost of self-attention. MEGABYTE decomposes long sequences into two shorter sequences, and optimal patch sizes reduces the self-attention cost to  $O\left(N^{\frac{4}{3}}\right)$ , which remains tractable for even long sequences.
2. Per-patch feedforward layers In GPT3-size models, more than  $98\%$  of FLOPS are used in computing position-wise feedforward layers. MEGABYTE uses large feedforward layers per-patch rather than per-position, enabling much larger and more expressive models for the same cost. With patch size  $P$ , where a baseline transformer would use the same feedforward layer with  $m$  parameters  $P$  times, MEGABYTE can use a layer with  $mP$  parameters once for the same cost.
3. Parallelism in Decoding Transformers must perform all computations serially during generation because the input to each timestep is the output from the previous timestep. By reusing the global representation over multiple time steps during local model decoding, MEGABYTE allows greater parallelism during generation. For example, a MEGABYTE model with 1.5B parameters can generate sequences  $40\%$  faster than a standard 350M Transformer, whilst also improving perplexity when trained with the same compute.
Together, these improvements allow us to train much larger and better-performing models for the same compute budget, scale to very long sequences, and improve generation speed during deployment.
MEGABYTE also provides a strong contrast to existing autoregressive models that typically use some form of tokenization, where sequences of bytes are mapped to larger discrete tokens (Sennrich et al., 2015; Ramesh et al., 2021; Hsu et al., 2021). Tokenization complicates pre-processing, multi-modal modelling, and transfer to new domains, while hiding useful structure from the model. It also means that most state-of-the-art models are not truly end to end. The most widely used approaches to tokenization require language-specific heuristics (Radford et al., 2019) or lose information (Ramesh et al., 2021). Replacing tokenization with efficient and performant byte models would therefore have many advantages.
We conduct extensive experiments for both MEGABYTE and strong baselines. We use a fixed compute and data budget across all models to focus our comparisons solely on the model architecture rather than training resources, which are known to benefit all models. We find that MEGABYTE allows byte-level models to perform competitively with subword models on long context language modeling, achieve state-of-the-art perplexities for density estimation on ImageNet, and allow audio modelling from raw audio files. Together, these results establish the viability of tokenization-free autoregressive sequence modeling at scale.
2 MEGABYTE Transformer
MEGABYTE is an autoregressive model for efficiently modeling long input sequences. MEGABYTE is comprised of 3 components: (1) a patch embedder that inputs a discrete sequence, embeds each element, and chunks it into patches of length  $P$  (2) a large global Transformer that contextualizes patch representations by performing self-attention over previous patches, and (3) a smaller local Transformer that inputs a contextualized patch representation from the global model, and autoregressively predict the next patch.
2.2 Components
Patch Embedder with patch size of  $P$  maps a byte sequence  $x_{0..T}$  to a sequence of patch embeddings of length  $K = \frac{T}{P}$  and dimension  $P \cdot D_G$ .
First, each byte is embedded with a lookup table  $E^{\mathrm{global - embed}} \in \mathbb{R}^{V \times D_G}$  to an embedding of size  $D_G$  and positional embeddings are added.
$$
h _ {t} ^ {\text {e m b e d}} = E _ {x _ {t}} ^ {\text {g l o b a l - e m b e d}} + E _ {t} ^ {\text {p o s}} \quad t \in [ 0.. T ] \tag {1}
$$
Then, byte embeddings are reshaped into a sequence of  $K$  patch embeddings with dimension  $P \cdot D_G$ . To allow autoregressive modelling, the patch sequence is padded to start with a trainable patch-sized padding embedding  $(E^{\mathrm{global - pad}} \in \mathbb{R}^{P \times D_G})$ , and the last patch is removed from the input. This sequence is the input to the global model, and is denoted  $h^{\mathrm{global - in}} \in \mathbb{R}^{K \times (P \cdot D_G)}$ .
$$
h _ {k} ^ {\text {g l o b a l - i n}} = \left\{ \begin{array}{l l} E ^ {\text {g l o b a l - p a d}}, & \text {i f} k = 0, \\ h _ {((k - 1) \cdot P): (k \cdot P)} ^ {\text {e m b e d}}, & k \in [ 1,.., K), \end{array} \right. \tag {2}
$$
Global Model is a decoder-only Transformer with dimension  $P \cdot D_G$  that operates on a sequence of  $K$  patches. It incorporates a self-attention mechanism and causal masking to capture dependencies between patches. It inputs a sequence of  $K$  patch representations  $h_{0:K}^{\mathrm{global - in}}$ , and outputs an updated representation  $h_{0:K}^{\mathrm{global - out}}$  by performing self-attention over previous patches.
$$
h _ {0: K} ^ {\text {g l o b a l - o u t}} = \operatorname {t r a n s f o r m e r} ^ {\text {g l o b a l}} \left(h _ {0: K} ^ {\text {g l o b a l - i n}}\right) \tag {3}
$$
The output of the final global layer  $h_{0:K}^{\mathrm{global}}$  contains  $K$  patch representations of dimension  $P \cdot D_G$ . For each of these, we reshape them into sequences of length  $P$  and dimension  $D_G$ , where position  $p$  uses dimensions  $p \cdot D_G$  to  $(p + 1) \cdot D_G$ . Each position is then projected to the dimension of the local model with a matrix  $w^{\mathrm{GL}} \in \mathbb{R}^{D_G \times D_L}$  where  $D_L$  is the local model dimension. We then combine these with byte embeddings of size  $D_L$  for the tokens in the next patch  $E_{x(k \cdot P + p - 1)}^{\mathrm{local - embed}}$ . The local byte embeddings is offset by one with a trainable local padding embedding  $(E^{\mathrm{local - pad}} \in \mathbb{R}^{D_L})$  to allow autoregressive modelling within a patch. This results in a tensor  $h^{\mathrm{local - in}} \in \mathbb{R}^{K \times P \times D_L}$ .
$$
h _ {k, p} ^ {\text {l o c a l - i n}} = w ^ {\mathrm {G L}} h _ {k, (p \cdot D _ {G}): ((p + 1) \cdot D _ {G})} ^ {\text {g l o b a l - o u t}} + E _ {x _ {(k \cdot P + p - 1)}} ^ {\text {l o c a l - e m b e d}} \tag {4}
$$
Local Model is a smaller decoder-only Transformer of dimension  $D_{L}$  that operates on a single patch  $k$  containing  $P$  elements, each of which is the sum of an output from the global model and an embedding of the previous byte in the sequence.  $K$  copies of the local models are run on each patch independently (and in parallel during training), computing a representation  $h^{\mathrm{local - out}}\in \mathbb{R}^{K\times P\cdot D_L}$ .
$$
h _ {k, 0: P} ^ {\text {l o c a l - o u t}} = \operatorname {t r a n s f o r m e r} ^ {\text {l o c a l}} \left(h _ {k, 0: P} ^ {\text {l o c a l - i n}}\right) \tag {5}
$$
Finally, we can compute the probability distribution over the vocabulary at each position. The  $p$ th element of the  $k$ th patch corresponds to element  $t$  of the complete sequence, where  $t = k \cdot P + p$ :
$$
p \left(x _ {t} \mid x _ {0: t}\right) = \operatorname {s o f t m a x} \left(E ^ {\text {l o c a l - e m b e d}} h _ {k, p} ^ {\text {l o c a l - o u t}}\right) _ {x _ {t}} \tag {6}
$$
2.3 Variations and Extensions
Convolutional Patch Encoder: One limitation of patchifying sequences is that it is not translation invariant, and byte sequences may receive a different representation depending on their position in the patch. This may mean, for example, that a model has to relearn the meaning of a word at different offsets. To mitigate this issue, we experimented with augmenting the Patch Embedder with causal convolutional layers, which allow translation-invariant contextual representations of the bytes before they are chunked into patches. We use a stack of convolutional layers, with filter sizes of 3, 5 and 7.
Cross-patch Attention: The Local model uses short sequences for efficiency, and relies on the Global model for long-range information. However, we can increase the context of the Local model with little overhead by allowing it to condition on  $r$  elements from the previous patch. This approach allows the Global model to focus on a longer-range context. Specifically, when computing self-attention in each layer, we concatenate the keys and values with the last  $r$  keys and queries from the previous patch. We use rotary embeddings (Su et al., 2021) to model relative positions between elements in the sequence. This approach is reminiscent of TransformerXL (Dai et al., 2019) but differs by being fully differentiable.
Strided Inference: We observed empirically that the per-token loss within each patch increases towards the end of the patch, as the prediction relies more on the weaker Local model. To alleviate this issue, we propose strided inference, in which we predict the sequence with two forward passes of the full model, whose inputs are offset by  $p/2$  positions from each other. We then combine the first  $p/2$  positions in each patch for our predictions to predict the complete sequence. Similarly to sliding window methods (Press et al., 2020), this approach doubles the cost of inference but improves results.
3 Efficiency Analysis
3.1 Training Efficiency
Attention The cost of attention in a transformer architecture for a sequence of length  $T$  has  $O(T^{2})$  complexity. Much work has been explored reducing this; for example, Sparse Transformers (Child et al., 2019) and Routing Transformers (Roy et al., 2020) show strong results with a complexity  $O(T^{\frac{3}{2}})$ . Many linear attention mechanisms have also been proposed (Katharopoulos et al., 2020; Choromanski et al., 2020), although we are not aware of competitive results on large scale language modeling tasks. As a function of sequence length  $T$  and patch size  $P$ , the Global model has a sequence of length  $\frac{P}{T}$  so uses  $O\left(\frac{T^2}{P^2}\right)$  operations, and the Local model uses  $\frac{P}{T}$  sequences of length  $P$  so uses  $O\left(\frac{TP^2}{P}\right) = O(PT)$  operations. The overall cost of MEGABYTE is therefore in  $O\left(\frac{T^2}{P^2} + TP\right)$ .  $P$  is a hyperparameter that is chosen to create an architecture for sequences of size  $T$ . By setting  $P = T^{\frac{1}{3}}$  the complexity is in  $O(T^{\frac{4}{3}})$ . Using much shorter patches of  $P = T^{\frac{1}{5}}$  would give a complexity of  $O(T^{\frac{8}{5}})$ . The cost is less than the transformer for all non-trivial values of  $P$  such that  $1 < P < T$ .
Feedforward Layers However, attention is not the main cost in large transformers. Instead of increasing the sequence length, transformers are more commonly scaled by increasing the dimension of their latent state  $d$ , and the feedforward network cost dominates the model's overall cost (Kaplan et al., 2020). For example, in the GPT3 architecture, the quadratic self-attention computation accounts for only  $1.4\%$  of FLOPS. Following the approximation of (Kaplan et al., 2020), a forward pass with a large transformer with  $m$  non-embedding parameters on a sequence of length  $T$  uses roughly  $2mT$  FLOPS. MEGABYTE contains two transformers: the Global model uses  $m_g$  parameters on a sequence of length  $\frac{T}{P}$ , and a Local model with  $m_l$  parameters that sees  $\frac{T}{P}$  sequences of length  $P$ , giving an estimate of  $2T(\frac{m_g}{P} + m_l)$  FLOPS. When  $m_g \gg m_l$ , the FLOPS used by MEGABYTE is approximately  $\frac{2Tm_g}{P}$ , allowing a model  $P$  times larger than a transformer with equivalent FLOPS. This analysis holds irrespective of any efficient attention mechanisms used in the transformer.
Combined Analysis To understand efficiency at different sequence lengths and model sizes, we calculate the total FLOPS used by transformers, Linear Transformers and MEGABYTE. For each operation, we use FLOP estimates from (Kaplan et al., 2020), except for attention in Linear Transformers, which we estimate as  $9D$  FLOPS-token $^1$ , where  $D$  is the model embedding dimension. Figure 3 shows that for models of size 660M to 173B and sequence lengths of up to 1M tokens,
This may underestimate the time taken by Linear Transformer decoders, which use a recurrence mechanism that is harder to parallelize on current hardware.
MEGABYTE with  $P = 8$  uses less FLOPS than either transformers or Linear Transformers. Baseline model architectures are based on GPT3, and Megabyte global/local model sizes are 452M/151M, 5.8B/604M, 170B/3.2B respectively.
3.2 Generation Efficiency
Generating long sequences with transformers is slow, because the input to each timestep is the output from the previous timestep, meaning each layer must be computed for each token serially. As running a layer on a single token typically does not saturate the amount of parallelism available within a GPU, for analysis, we model each layer as a constant cost independently of size. Consider a MEGABYTE model with  $L_{\mathrm{global}}$  layers in the Global model and  $L_{\mathrm{local}}$  layers in the Local model and patch size  $P$ , compared with a Transformer architecture with  $L_{\mathrm{local}} + L_{\mathrm{global}}$  layers. Generating each patch with MEGABYTE requires a sequence of  $O(L_{\mathrm{global}} + P \cdot L_{\mathrm{local}})$  serial operations, whereas the Transformer requires  $O(P \cdot L_{\mathrm{global}} + P \cdot L_{\mathrm{local}})$  serial operations. When  $L_{\mathrm{global}} \gg L_{\mathrm{local}}$  (i.e. the Global model has many more layers than the Local model), MEGABYTE can reduce inference costs by a factor close to  $P$ .
Figure 3: Computational cost (FLOPS-token) for different model architectures at different scales. MEGABYTE architectures (here with  $P = 8$ ) use less FLOPS than equivalently sized Transformers and Linear Transformers (Katharopoulos et al., 2020) across a wide range of model sizes and sequence lengths, allowing larger models to be used for the same computational cost.
Controlling for Compute and Data Models show consistent improvements when increasing both data and compute Kaplan et al. (2020); Hoffmann et al. (2022), meaning that one model can outperform another because of an increased training budget instead of an improved architecture. However, in practice, both compute and data are typically limited. We conduct experiments using a fixed compute and data budget across all models to focus comparisons solely on the model architecture rather than training resources. To achieve this, we adjust model hyperparameters (mainly, number of layers) within each architecture so that the forward pass time taken per byte is matched, and then train all models for the same number of bytes.
Comparison Systems We compare MEGABYTE with both a standard decoder-only Transformer and PerceiverAR (Hawthorne et al., 2022). PerceiverAR extends the original transformer with a single cross-attention layer over a much longer context sequence, and is the best performing general purpose autoregressive model we are aware of and achieves state-of-the-art results across several modalities. We implemented both models in the same codebase, and all models share a similar data loader, preprocessing step, and trainer to avoid any artifacts in our compute-controlled experiments.
Training Procedure All models were trained using the Metaseq code base Zhang et al. (2022b). 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). Mixed precision training was used to improve training efficiency at scale Micikevicius et al. (2017). More training details and various model parameters can be found in Section A.1 in the Appendix. To validate our implementation of PerceiverAR, we reproduced their experiments on downsized ImageNet at 64 pixels. By carefully matching hyperparameters, we achieved a bits per byte (bpb) score of 3.53, compared to the reported 3.54 in the original paper.
Inference Methods Several techniques have been proposed for trading off speed for performance during inference with language models, including sliding windows Press et al. (2020) and our strided inference. We only use these methods when comparing with prior published work (Tables 2 and 3).
$^{2}$ https://github.com/facebookresearch/metaseq
Dataset	Total Bytes	bytes/doc	Transformer	PerceiverAR	MEGABYTE
PG-19	10.1GB	411,404	1.057	1.104	1.000
Stories	21.3GB	35,265	1.064	1.070	0.978
Books	79.7GB	509,526	1.097	1.104	1.007
arXiv	91.5GB	58,518	0.816	0.791	0.678
Code	353.7GB	7,461	0.575	0.546	0.411
Table 1: Text dataset sizes and mean document lengths. We also report bpb of various models (Transformer, PerceiverAR, and MEGABYTE) trained with the same compute.
Tokenizer	Vocab	Context Length	Validation	Test
TransformerXL Rae et al. (2019a)	SentPiece	32k	512+1024	45.5	36.3
CompressiveTransformer Rae et al. (2019a)	SentPiece	32k	512+512+2x512	43.4	33.6
PerceiverAR Hawthorne et al. (2022)	SentPiece	32k	2048	45.9	28.9
BlockRecurrent Hutchins et al. (2022)	SentPiece	32k	1024+recurrence	-	26.5
Transformer byte-level (ours)	Bytes	256	2048	81.6	69.4
PerceiverAR byte-level (ours)	Bytes	256	8192	119.1	88.8
MEGABYTE	Bytes	256	8192	42.8	36.4
Table 2: Larger scale experiments on PG19, converting bits-per-byte to word-level perplexities for comparison with prior work. Results below the line are compute-matched. MEGABYTE outperforms other byte models by a wide margin, and gives results competitive with state-of-the-art models trained on subwords.
5 Language Modeling
We evaluated the performance of MEGABYTE on language modeling on a set of 5 diverse datasets emphasizing long-range dependencies: Project Gutenberg (PG-19), Books, Stories, arXiv, and Code.
Datasets We experiment on a range of long form text datasets. The PG-19 dataset Rae et al. (2019b) consists of English-language books written before 1919 and is extracted from the Project Gutenberg online library. The Stories dataset Trinh & Le (2018) is a subset of CommonCrawl data meant to emulate Winograd schemas. Books Gao et al. (2020) is another collection of English-language books. The arXiv dataset contains technical publications written in LATEX from the arXiv online archive. Finally, the Code dataset is a large publicly available dataset of open source code, under Apache, BSD or MIT licenses. More details on dataset sizes and document lengths are shared in Table 1.