PaperlayerView exact revision on arXiv

Source-linked AI summary

MEGABYTE: Predicting Million-byte Sequences with Multiscale Transformers

Lili Yu, Dániel Simig, Colin Flaherty, Armen Aghajanyan, Luke Zettlemoyer, Mike Lewis

arXiv:2305.07185v2cs.LG

TL;DR

Long-sequence autoregressive transformers face costly attention and feedforward computation, limiting their application to million-byte data. MEGABYTE addresses this with global modeling across patches and local modeling within patches. The resulting architecture supports competitive byte-level language modeling, ImageNet density estimation, and raw-audio modeling, though the experiments remain much smaller than those for state-of-the-art language models.

  • Problem

    Transformers scale poorly to million-byte sequences because self-attention and large per-position feedforward networks are expensive.

  • Method

    MEGABYTE segments byte sequences into patches, uses a large global autoregressive transformer across patches, and uses smaller local models to predict bytes within patches.

  • Results

    MEGABYTE outperforms existing byte-level models across tasks and modalities, while producing competitive language-modeling results with subword models.

  • Takeaways & Limitations

    The results support tokenization-free autoregressive modeling of sequences exceeding one million tokens across language, image, and audio tasks.

  • Takeaways & Limitations

    The experiments are far below the scale of state-of-the-art language models, so scaling MEGABYTE to larger models and datasets remains future work.

Abstract

from arXiv · show

Autoregressive transformers are spectacular models for short sequences but scale poorly to long sequences such as high-resolution images, podcasts, code, or books. We proposed 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.

1. Introduction

MEGABYTE models million-byte sequences by combining a global transformer over patches with smaller local models within patches. This design reduces attention cost, enlarges per-patch feedforward capacity, and improves decoding parallelism while supporting byte-level modeling.

  • Architecture: MEGABYTE segments byte sequences into fixed-size patches and combines a large global transformer over patches with smaller local autoregressive models within patches.The patch embedder encodes bytes, the global module models patch representations, and the local module predicts bytes within each patch.
  • Efficiency: MEGABYTE reduces self-attention cost to O(N^(4/3)) through a two-level decomposition of long sequences.Optimal patch sizes produce two shorter sequences whose attention remains tractable for long inputs.
  • Efficiency: With patch size P, MEGABYTE can use a feedforward layer with mP parameters once instead of a baseline layer with m parameters P times.This reallocates computation from repeated per-position feedforward layers to larger per-patch layers.
  • Experiments: Experiments report competitive byte-level language modeling with subword models, state-of-the-art ImageNet density-estimation perplexities, and raw-audio modeling.The comparisons use fixed compute and data budgets across models.
  • Limitations and extensions: MEGABYTE’s patch representation is not translation invariant, so byte sequences can receive different representations at different positions within a patch.The authors explore causal convolutional layers to provide translation-invariant context before chunking.

3. Efficiency Analysis

MEGABYTE reduces long-sequence computation by splitting sequences into global patch-level and local within-patch processing. This lowers attention and feedforward costs while improving generation parallelism.

  • Attention: MEGABYTE decomposes sequences into global and local models, giving overall complexity O(T^2/P^2 + TP).The global model processes T/P patches, while the local model processes P-length patches.
  • Feedforward Layers: Over 98% of FLOPS in large Transformers are spent on position-wise feedforward layers, which MEGABYTE applies per patch instead of per byte.With patch size P, MEGABYTE can use a feedforward layer with mP parameters once for the same cost as m parameters used P times.
  • Combined Analysis: MEGABYTE with P = 8 uses fewer FLOPS than Transformers and Linear Transformers across model sizes from 660M to 173B and sequence lengths up to 1M tokens.The compared MEGABYTE global/local model sizes are 452M/151M, 5.8B/604M, and 170B/3.2B parameters.
  • Parallelism in Decoding: MEGABYTE can reduce generation inference costs by nearly a factor of P when the global model has many more layers than the local model.Generating each patch requires O(L_global + P · L_local) serial operations versus O(P · L_global + P · L_local) for a Transformer.

4. Experimental setup

The experiments compare MEGABYTE with standard Transformers and PerceiverAR under controlled training conditions. Implementation checks, shared infrastructure, and prior-comparison evaluation procedures support the experimental setup.

  • Controlled comparisons: All models use fixed compute and data budgets so comparisons focus on architecture rather than training resources.
  • Comparison Systems: MEGABYTE is compared with a standard decoder-only Transformer and PerceiverAR, which adds cross-attention over a longer context sequence.
  • Implementation: The models share a codebase, data loader, preprocessing step, and trainer to reduce artifacts in compute-controlled experiments.
  • Validation: PerceiverAR reproduction on 64-pixel ImageNet achieved 3.53 bits per byte versus the original reported 3.54.
  • Evaluation: Sliding-window and strided inference are used only when comparing with prior published work in Tables 3 and 4.

5. Language Modeling

MEGABYTE is evaluated on diverse long-form text datasets under controlled compute and data budgets. It consistently outperforms byte-level baselines and remains competitive with state-of-the-art subword models at larger scale.

  • Datasets: MEGABYTE is evaluated on PG-19, Books, Stories, arXiv, and Code, emphasizing long-range dependencies.The datasets include books, CommonCrawl text, technical publications, and open-source code.
  • Controlled Experiments: Under matched compute and data budgets, MEGABYTE consistently outperforms Transformer and PerceiverAR baselines across the text datasets.The controlled experiments report bits-per-byte for models trained on 80 billion bytes.
  • Scaling Experiment: At 400B training bytes, MEGABYTE outperforms other byte-level models by a wide margin.The larger-scale comparison uses the same model parameters and compute budget as the controlled experiments.
  • Scaling Experiment: MEGABYTE achieves results competitive with state-of-the-art subword models, although the comparison may be confounded by differing compute and tuning.The authors present this as evidence that byte-level modeling may support tokenization-free language models.

6. Image Modeling

MEGABYTE models long image sequences across resolutions using compute-controlled comparisons. It matches state-of-the-art performance at 64×64 and outperforms baselines as resolution increases, including sequences exceeding one million bytes.

  • Experimental Setup: ImageNet experiments cover resolutions from 64×64 to 640×640, with the largest images requiring effective modeling of over 1.2M tokens.The task combines local image patterns with long-range structural context.
  • Comparison with State of the Art: At ImageNet 64×64, MEGABYTE matches PerceiverAR’s state-of-the-art performance while using only half the compute.The MEGABYTE model uses 2.7B global and 350M local parameters and is trained for 1.4T tokens.
  • Scaling to Higher Resolutions: Across image resolutions, vanilla Transformer, PerceiverAR, and MEGABYTE variants are trained with matched GPU hours and data.Baseline models require segmented evaluation for the longest sequences, whereas MEGABYTE uses a single forward pass up to 1.2M tokens.
  • Scaling to Higher Resolutions: MEGABYTE outperforms the baselines at all tested resolutions, demonstrating effective modeling of sequences over 1M bytes.The comparison uses the same compute and data budgets across resolutions.

7. Audio Modeling

MEGABYTE is applied to raw 16-bit audio sequences, addressing the long and high-cardinality nature of audio modeling. It achieves lower bits-per-byte than both PerceiverAR and a vanilla Transformer.

  • Motivation: Audio is a relevant MEGABYTE application because it combines sequential structure with the continuous nature of images.The experiments model raw audio rather than a tokenized representation.
  • Modeling Raw Audio: Raw 16-bit audio requires modeling 65,536 possible values per timestep with a direct softmax.Prior approaches reduce this computational burden through quantization or mixture distributions.
  • Experimental Setup: The audio experiments use 16 kHz, 16-bit audio, 2 terabytes of data, sequence length 524,288, and patch size 32.The dataset contains roughly 18,000 hours of audio.
  • Results: MEGABYTE obtains 3.477 bpb, compared with 3.543 for PerceiverAR and 3.567 for the vanilla Transformer.Lower bits-per-byte indicates the best reported result among these models.

8. Analysis

The analysis finds that MEGABYTE benefits from both local and global components, uses context effectively, and gains from strided inference and robust patch-size choices. Its larger global model also improves efficiency and performance under matched compute.

  • Generation speed: 40% faster generation is achieved by a MEGABYTE model with over four times the parameters of the baseline transformer.The comparison uses a 1.3B-parameter global model plus a 218M-parameter local model against a 350M-parameter transformer.
  • Component ablations: Removing either the Local or Global model substantially increases bits-per-byte across arXiv, Librilight-L, and ImageNet, while removing cross-patch local modeling remains competitive.The ablation table reports that both main modules are crucial, whereas the cross-patch modification is comparatively robust.
  • Context use: Later tokens within the context window receive higher likelihoods, indicating effective use of at least 8k bytes of context on PG19.Figure 4 reports this trend for an 8192-context MEGABYTE model.
  • Inference techniques: Strided inference improves performance by combining predictions from two offset forward passes, with gains additive to the standard sliding window.The method uses two passes offset by P/2 tokens and combines the first half of each patch.
  • Hyperparameters: MEGABYTE performs similarly across a wide range of patch sizes, although the optimal size can differ between modalities.This robustness was observed on Image256 and across other modalities.
  • Hyperparameters: A larger Global model uses compute more effectively than a larger Local model, a trend consistent across modalities and patch sizes.Grouping bytes into patches lets the Global model process P times fewer tokens, enabling greater capacity at comparable cost.

9. Related Work

Related work improves long-sequence efficiency through encoder adaptations, decoder chunking, linear or recurrent alternatives, sparse attention, and tokenization. MEGABYTE instead decomposes decoder sequences into two shorter sequences while retaining dense attention and byte-level modeling.

  • Long-sequence research has primarily targeted the quadratic cost of self-attention.
  • Encoder patching and downsampling methods cannot be naively transferred to decoder-only models because they may leak information from future bytes.
  • Decoder efficiency methods include chunking, linear or recurrent alternatives, state space models, and sparse attention.
  • MEGABYTE decomposes the complete sequence into two shorter sequences, producing sub-quadratic attention while preserving a dense-attention decoder approach.
  • Tokenization shortens decoder sequences by mapping multiple bytes to fixed-vocabulary tokens but can require language-specific heuristics and limit out-of-domain performance.

10. Conclusion

The conclusion presents MEGABYTE as a scalable architecture for long byte sequences, with competitive language-modeling results and applicability across modalities. It also notes that the experiments remain much smaller than those of state-of-the-art language models.

  • MEGABYTE models sequences of over 1 million tokens and outperforms existing byte-level models across tasks and modalities.
  • MEGABYTE achieves competitive language-modeling results with subword models, supporting byte-level modeling as an alternative to tokenization.
  • The experimental scale is far below that of state-of-the-art language models, motivating future scaling to larger models and datasets.

A. Appendices

The appendices document training procedures, matched-compute comparisons, model architectures, and pseudocode for preparing inputs and combining global and local modules.

  • Training details: Training uses gradient clipping, Adam, polynomial learning-rate decay, pre-norm ReLU models, dropout, and weight decay.
  • Experimental controls: Models within each architecture are matched by update time and trained for the same number of updates under fixed compute and data budgets.
  • Model specifications: Architecture tables report model size, layers, embedding dimensions, attention heads, batch sizes, learning rates, context lengths, latents, and patch sizes.
  • MEGABYTE implementation: The pseudocode embeds bytes globally by patch, applies a global TransformerDecoder, adds global outputs to local byte embeddings, and applies a local TransformerDecoder.
  • MEGABYTE implementation: The final output reshapes local predictions back into the original batch-by-sequence layout.
  • Input preparation: Input preparation reshapes bytes into patches, prepends padding for local autoregressive inputs, and returns global and local byte sequences.

C. PerceiverAR Implementation

MEGABYTE models image bytes with either raster scan or patch scan, which rasterizes within and between image patches. Patch scan is more effective for image sequencing than raster scan and improves both MEGABYTE and Perceiver AR.

  • Image sequencing: Images can be serialized by raster scanning bytes row by row or by scanning patches and rasterizing within and between them.Patch scan first divides each image into p × p × 3 patches.
  • Image sequencing: Unless otherwise specified, MEGABYTE uses patch scan for image data.
  • Image sequencing: Patch scan is more effective than raster scan for image sequencing and improves both MEGABYTE and Perceiver AR.The comparison is reported for ImageNet256 performance in Table 13.
  • Longer sequences: Longer sequences did not improve PG19 results under the reported model configuration.The experiment used a 1.3b global model and a 350m local model; the authors expect longer sequences may help with larger models and more data.
Loading 2305.07185v2…