Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

GQ-VAE: A gated quantized VAE for learning variable length tokens

Theo Datta, Kayla Huang, Sham Kakade, David Brandfonbrener

TL;DR

GQ-VAE tackles tokenization by learning discrete, variable-length tokens that can serve as a drop-in tokenizer for large language models. It introduces an encoder, a vector quantizer, a gate, and a decoder to form tokens whose boundaries are learned and adjustable, guided by reconstruction, length, compression, and codebook objectives. The approach achieves compression competitive with BPE and improves downstream language modeling when matched for compression, while producing a more uniform token distribution in the tail, suggesting better learnability for rare tokens; it remains a proof-of-concept evaluated on a toy dataset. Overall, GQ-VAE demonstrates a promising direction toward learned tokenization that minimizes architectural changes to LLMs and opens avenues for applying discrete, variable-length compression in diverse domains.

Abstract

While most frontier models still use deterministic frequency-based tokenization algorithms such as byte-pair encoding (BPE), there has been significant recent work to design learned neural tokenizers. However, these schemes generally add to underlying language model complexity and force large changes to architecture, making them hard to implement at large scales. To overcome these challenges, we propose the gated quantized variational autoencoder (GQ-VAE), a novel architecture that can be independently pre-trained to serve as a drop-in replacement for existing tokenizers. The key innovation of the architecture is to learn to encode variable-length discrete tokens. GQ-VAE improves compression and language modeling performance over a standard VQ-VAE tokenizer, and approaches the compression rate and language modeling performance of BPE. Interestingly, if we use BPE with a smaller vocabulary, such that the compression is equivalent between GQ-VAE and BPE, we find that GQ-VAE improves downstream language model learning. We conclude with a discussion of several exciting avenues for future work. Code can be found at https://github.com/Theo-Datta-115/gq-vae.

GQ-VAE: A gated quantized VAE for learning variable length tokens

TL;DR

GQ-VAE tackles tokenization by learning discrete, variable-length tokens that can serve as a drop-in tokenizer for large language models. It introduces an encoder, a vector quantizer, a gate, and a decoder to form tokens whose boundaries are learned and adjustable, guided by reconstruction, length, compression, and codebook objectives. The approach achieves compression competitive with BPE and improves downstream language modeling when matched for compression, while producing a more uniform token distribution in the tail, suggesting better learnability for rare tokens; it remains a proof-of-concept evaluated on a toy dataset. Overall, GQ-VAE demonstrates a promising direction toward learned tokenization that minimizes architectural changes to LLMs and opens avenues for applying discrete, variable-length compression in diverse domains.

Abstract

While most frontier models still use deterministic frequency-based tokenization algorithms such as byte-pair encoding (BPE), there has been significant recent work to design learned neural tokenizers. However, these schemes generally add to underlying language model complexity and force large changes to architecture, making them hard to implement at large scales. To overcome these challenges, we propose the gated quantized variational autoencoder (GQ-VAE), a novel architecture that can be independently pre-trained to serve as a drop-in replacement for existing tokenizers. The key innovation of the architecture is to learn to encode variable-length discrete tokens. GQ-VAE improves compression and language modeling performance over a standard VQ-VAE tokenizer, and approaches the compression rate and language modeling performance of BPE. Interestingly, if we use BPE with a smaller vocabulary, such that the compression is equivalent between GQ-VAE and BPE, we find that GQ-VAE improves downstream language model learning. We conclude with a discussion of several exciting avenues for future work. Code can be found at https://github.com/Theo-Datta-115/gq-vae.
Paper Structure (15 sections, 7 equations, 7 figures)

This paper contains 15 sections, 7 equations, 7 figures.

Table of Contents

  1. Introduction
  2. GQ-VAE Architecture
  3. Architecture and inference procedure
  4. Training
  5. Reconstruction.
  6. Compression.
  7. Length decoding.
  8. Vector quantization.
  9. Putting it together
  10. Implementation Details
  11. Results
  12. Compression and reconstruction
  13. Language Modeling
  14. Limitations
  15. Conclusion

Figures (7)

  • Figure 1: GQ-VAE Architecture. D=Decoder, Q=Quantizer. Encoder and Gater are transformers. The decoder head is illustrated in (Figure \ref{['fig:decoder']}).
  • Figure 2: Decoder Head
  • Figure 3: Compression and Vocabulary Size. VQ-VAE represents a fixed token length GQ-VAE, and higher vocab results continue the linear trend but are omitted for scale. GQ-VAE vocab sizes are distilled to represent only the number of unique tokens, while the codebook sizes in the model are [2k, 5k, 10k, 20k, 50k].
  • Figure 4: Compression for fixed-length token models. These models are all trained on the same hyperparameter setup as the 3.85 bits/byte variable length GQ-VAE baseline.
  • Figure 5: Language Modeling with GQ-VAE, BPE, and VQ-VAE (fixed token length GQ-VAE). Models are trained on all of tinystories, so the lower-compression GQ-VAE trains for more steps on the same data. The exact "used" vocabulary sizes for these models are BPE=10000, GQ-VAE=10314, VQ-VAE=13201.
  • ...and 2 more figures