Mixture of Tokenizers — multiple parallel embedding pathways with learned routing.

Uses N separate tokenizer embedding pathways, each with a different vocabulary size and embedding dimension, combined via learned soft routing weights. This allows the model to dynamically select the best tokenization granularity for each position.

Architecture

Input [batch, seq_len]
      |
+-- Tokenizer 1: embedding(vocab_1, embed_1) -> dense(hidden_size) --+
+-- Tokenizer 2: embedding(vocab_2, embed_2) -> dense(hidden_size) --+
+-- ...                                                              +
+-- Tokenizer N: embedding(vocab_N, embed_N) -> dense(hidden_size) --+
      |
Router: shared_embed -> dense(N) -> softmax -> weights [batch, seq_len, N]
      |
Weighted sum -> [batch, seq_len, hidden_size]
      |
Transformer blocks -> final norm -> last timestep
      |
[batch, hidden_size]

Usage

model = MixtureOfTokenizers.build(
  hidden_size: 256,
  num_tokenizers: 4,
  tokenizer_vocab_sizes: [256, 512, 1024, 2048],
  tokenizer_embed_dims: [32, 64, 128, 256]
)

References

• "Mixture-of-Tokenizers" (Pham et al., 2024) — multi-granularity tokenization