SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Listing 1: Pseudocode of Megabyte model
class MegaByteDecoder: def__init__( self, global_args, local_args, patch_size, ): self_pad  $= 0$  self.batch_size  $=$  patch_size self.globalmodel  $=$  TransformerDecoder(global_args) self.localmodel  $=$  TransformerDecoder(local_args) def forward( self, bytes, ): bytes_global，bytes_local  $=$  self.prepare_input(bytes)
global_bytes_embedding  $\equiv$  self.globalmodel_embedding(bytes_global) global_in  $=$  rearrange( global_bytes_embedding, "b (tp)e->bt(p e)," p= self.batch_size, 1 global_output  $\equiv$  self.globalmodel(global_in) global_output_reshaped  $\equiv$  rearrange( global_output , "b t (p e)->(b t)p e", p= self.batch_size, 1 local_bytes_embedding  $\equiv$  self.localmodel_embedding(bytes_local) local_in  $=$  local_bytes_embedding + global_output_reshaped local_output  $\equiv$  self.localmodel(local_in) batch_size  $\equiv$  bytes_global.shape[0] x  $=$  rearrange(local_output，"（b t）1v ->b(t1)v"，b= batch_size) return x   
def prepare_input(self,bytes): padding_global  $\equiv$  bytes.new(bytes.shape[0],self.batch_size). fill_(self_pad) bytes_global  $\equiv$  torch.cat((padding_global,bytes[(:,:-self. patch_size]),-1) bytes_input  $\equiv$  rearrange(bytes,"b(tp)\~"(btp)p",p= self. patch_size) padding_local  $\equiv$  bytes_input.new(bytes_input.shape[0],1).fill_(self_pad) bytes_local  $\equiv$  torch.cat((padding_local,bytes_input[,:-1]), -1) return bytes_global,bytes_local
C PerceiverAR Implementation
To reproduce PerceiverAR in a compute-controlled setting we extended the standard transformer implementation in metaseq with an additional cross attention layer to compute the latents and match the architecture of PerceiverAR. We trained the model by sampling random spans from each text, matching the procedure used in the PerceiverAR codebase. To be consistent with the original work, we use sliding window evaluation with a stride of num_latents/2 unless otherwise noted. In several cases we used the standard metaseq implementation as opposed to specific techniques reported in the original paper: 1) we used standard attention dropout instead of cross-attention dropout 2) We did not implement chunked attention. We verified our implementation by reproducing the "Standard Ordering" experiments in Table 5 of the Perceiver AR paper. After carefully matching context size, number of latents, the amount of data and training steps used and learning rate, we achieved 3.53 bpb vs 3.54 reported in the original paper.
D.1 Patch scan Implementation
Images have a natural structure, containing a grid of  $n \times n$  pixels each composed of 3 bytes (corresponding to color channels). We explore two ways of converting images to sequences for modeling (see Figure 4). Firstly, raster scan where the pixels are linearized into 3 bytes and concatenated row-by-row. Secondly, patch scan where we create patches of shape  $p \times p \times 3$  bytes
Figure 4: Two ways to model 2D data sequentially. Left, raster scan, by taking bytes row by row and left to right; right, patch scan, where we first split an image into patches, and do raster scan across patches and within a patch.  $(\mathrm{T} = 36, \mathrm{K} = 9, \mathrm{P} = 4)$ .
where  $p = \sqrt{\frac{P}{3}}$ , and then use a raster scan both within and between patches. Unless otherwise specified, MEGABYTE models use patch scan for image data.
D.2 Patch scan vs Raster scan
The patch scan method is inspired by recent works in Vision Transformers (Dosovitskiy et al., 2020), and it is more effective than raster scan for modeling image sequencing. We found it improves both MEGABYTE and Perceiver AR.
(Global) Size	Local Size	context	bpb
MEGABYTE (patch scan)	62M (D=768, L=6)	N/A	8,192 (768 latents)	3.158
MEGABYTE (raster scan)	62M (D=768, L=6)	N/A	8,192 (768 latents)	3.428
Perceiver AR (patch scan)	125M (D=768, L=12)	125M (D=768, L=12)	196,608 (patch size 192)	3.373
Perceiver AR (raster scan)	125M (D=768, L=12)	125M (D=768, L=12)	196,608 (patch size 192)	3.552
Table 14: ImageNet256 performance with patch scan vs raster scan for MEGABYTE and Perceiver AR.
D.3 Longer sequence modeling
For our pg19 scaling experiment, we also use longer context length for MEGABYTE. The results are shown in Table 15. With longer sequence, we didn't observer further improvement, consistent with findings in Hawthorne et al. (2022). We think we will benefit more from longer sequence when we further scale up the model size and data.
context	bpb
MEGABYTE	8,192 (patch size 8)	0.8751
MEGABYTE	16,384 (patch size 8)	0.8787
Table 15: Longer sequence for PG19 dataset. For both experiments, we set global model as 1.3b, local model as 350m, and MEGABYTE patch size as 8.