Tokenization Strategies

The Problem: Biology is Messy, AI Needs Order

Depending on what we want the model to learn, we use different strategies to sort and present these "bricks".

Visual Guide to Core Tokenization Concepts

Before diving into the details, let's understand the core concepts visually using our LEGO analogy.

The Basics: Identity & Count

A token is a single unit of information for the model. In single-cell data, we usually need two pieces of information combined:

1. What is it? (The Gene Identity, e.g., "TP53" or a red brick).
2. How much? (The Expression Value, e.g., "50 counts" or a stack of 50 bricks).

The challenge is how to combine these two very different types of information into a single vector representing the token.

Strategy: Rank-Based Tokenization

The Problem: Sometimes we have "technical noise" (batch effects). One experiment might yield giant stacks of bricks (high sequencing depth), and another yields tiny stacks, even for the same biological cell type.

The Solution: Ignore the exact height. Just line them up from tallest to shortest. As long as the relative order is preserved (Red is taller than Blue), the resulting sequence of tokens is the same, making the model immune to batch differences in sequencing depth.

Used in: iSEEEK, Geneformer

Strategy: Expression Binning

The Problem: Standard language models work best with a fixed dictionary of words (categorical data). Actual expression counts are continuous numbers (1, 2, 50, 1000...).

The Solution: We create buckets (bins) for different ranges of heights. Instead of saying "Height 45", we throw it into the "Medium Height Bucket". Now, the token isn't a number, it's a category: [GeneID] + [MediumBin].

Used in: scGPT, scBERT

Strategy: Read-Depth-Aware (RDA) Tokenization

The Problem: Different cells have different sequencing depths—some have 10,000 total counts, others only 1,000. This makes expression values hard to compare directly.

The Solution: Keep the continuous expression values (don't discretize!), but add special "depth tokens" that tell the model: "This cell was supposed to have T=10,000 counts, but we only sampled S=1,000." The model learns to mentally "scale up" the values during pretraining.

Key Benefit: Enables expression enhancement—the model can predict what gene expression would look like at higher sequencing depth, effectively denoising sparse data.

Used in: scFoundation, AIDO.Cell

Strategy: Genome-Coordinate Tokenization

The Problem: In chromatin data (ATAC-seq), we don't have predefined "genes". We just have regions on the genome that are "open" (accessible). These regions change depending on the cell type.

The Solution: Imagine the genome as a giant LEGO baseplate ruler. We don't define the brick type; we define where the bricks are placed. The token isn't a name, it's a set of coordinates: Chromosome number, Start position on the ruler, and End position.

Used in: ChromFound

Strategy: Cell2Sentence (C2S)

The Idea: Instead of stacking bricks to show expression level, lay them out in a horizontal line. High expression? Repeat that brick many times. Low expression? Just one or two repeats.

Why It Matters: This transforms cell data into text-like "sentences" that standard large language models (GPT, LLaMA) can understand. You can even ask natural language questions about cells!

Used in: C2S-Scale (27B params)

Strategy: Amino Acid Tokenization

The Idea: Proteins are chains of just 20 different building blocks (amino acids). Each amino acid is one token - like having 20 different LEGO brick types that snap together in a chain.

Why It Works: Just like language models learn word patterns ("the" often follows "in"), protein models learn amino acid patterns that determine protein function and structure.

Used in: ESM3, ProGen2

Strategy: Codon-Level Tokenization

The Idea: Instead of tokenizing single DNA letters (A, T, G, C), group them into triplets called codons. There are 64 possible triplet combinations (4^3).

Why It Matters: Multiple codons can code for the same amino acid (synonymous codons), but the choice affects translation speed and mRNA stability. "Silent" mutations can still cause disease!

Used in: CodonFM, EnCodon

Strategy: BPE + IUPAC Encoding for Diploid Genomes

The Problem: Reference genome models ignore individual genetic variation. How do you encode personalized genomes with heterozygous variants (where you inherited different alleles from each parent)?

The Solution: Use IUPAC ambiguity codes to represent heterozygous sites directly in the sequence (e.g., Y for C/T, R for A/G). Then apply Byte-Pair Encoding (BPE) to learn variable-length subword tokens that capture regulatory motifs.

Why It Works: This enables native diploid genome modeling without separate haplotype processing. The model learns the biological significance of both homozygous and heterozygous positions.

Used in: VariantFormer

Strategy: K-mer Tokenization

The Idea: Instead of reading DNA one letter at a time (A, T, G, C), slide a window of K letters along the sequence. Each window position becomes one token. For 6-mers, there are 4^6 = 4,096 possible tokens.

Why It Works: It's like reading words instead of individual letters. "ATGCAT" carries more biological meaning than six separate letters, capturing local motifs and regulatory elements.

Used in: Nucleotide Transformer, DNABERT

Strategy: Macrogene Tokenization

The Idea: Instead of using individual genes as tokens, group genes from ALL species into "macrogenes" based on protein sequence similarity (via ESM2). This creates a universal vocabulary that works across species.

Why It Works: Genes with similar protein functions cluster together regardless of species, enabling cross-species integration WITHOUT requiring one-to-one homolog mappings. Human CD4 and mouse Cd4 end up in the same macrogene!

Used in: SATURN

Key Papers Implementing These Strategies

Comparisons & Trade-offs

Every granularity level buys something and pays for something:

The table below is the same trade-off, one strategy at a time:

Tokenization Strategy Comparison

Context Length vs Model Scale Trade-offs

Different tokenization strategies and architectures dictate the maximum sequence length (context window) a model can handle, which impacts the biological scope it can capture.

Choosing the Right Tokenization Strategy

⚖️ Side-by-Side: Encoding Expression — Rank Ordering vs Value Binning

Geneformer and scGPT both turn a single cell's expression vector into a token sequence a transformer can read, but they diverge on the fundamental question: should a gene's expression level be encoded implicitly by its rank among the other genes (Geneformer), or explicitly as a discrete value token paired with the gene (scGPT)?

Up a Level: Cells as Tokens — STACK vs STATE

The two models above tokenize genes within a cell. STACK and STATE (both from the Arc Institute) move up a granularity — the token is the whole cell, and attention runs over a set or table of cells. They diverge on how a set of cells becomes a perturbation prediction: STACK learns in context from unlabeled example cells at inference, while STATE trains an explicit embedding-plus-transition pipeline.

🛠️ Hands-On Practice

The steps below walk through tokenizing a single-cell dataset for a transformer foundation model — from an AnnData count matrix to two model-ready token streams: a rank-ordered sequence (Geneformer-style) and a (gene, value-bin) sequence (scGPT-style). The core logic is shown from scratch in NumPy so the mechanics are transparent, then mapped to the real library call.

Environment & packages

Install a minimal tokenization stack. scanpy/anndata hold the matrix; the model-specific tokenizers (Geneformer's TranscriptomeTokenizer, scGPT's GeneVocab) emit the final token IDs and handle the gene vocabulary, special tokens, and truncation.

# conda / mamba recommended
conda create -n tok python=3.10 -y
conda activate tok

pip install scanpy anndata numpy datasets
# model-specific tokenizers:
pip install geneformer          # rank-value encoding (expects Ensembl IDs)
# pip install scgpt             # value binning (GeneVocab tokenizer)

Hardware. Tokenization is CPU-bound and cheap — a laptop handles 100k+ cells. A GPU is only needed for the downstream model forward pass, not for building the token streams.

Data structures & formats

• AnnData — input: adata.X (raw counts), adata.var["ensembl_id"] (Geneformer keys on Ensembl IDs, not symbols), adata.obs["n_counts"] (total counts per cell)
• Token dictionary — a fixed gene → integer ID vocabulary; genes absent from it are out-of-vocabulary and silently dropped
• Rank-encoded sequence — an ordered list of gene token IDs per cell (no value tokens), truncated to the model context length (2,048 for Geneformer)
• Value-binned sequence — aligned arrays of gene token IDs and discrete value-bin IDs (e.g. 51 bins) per cell
• Special tokens — <cls> (cell-level embedding), <pad> (padding to a fixed length), <mask> (masked-token pretraining objective)
• Output — Geneformer emits a HuggingFace .dataset (Arrow) of token IDs; scGPT works from tensors of (gene_ids, values)

Minimal code walkthrough

Encode a toy count matrix with both strategies to expose the mechanics, then map to the real Geneformer tokenizer.

import numpy as np

# Toy: 3 cells x 6 genes (raw UMI counts). In practice: adata = sc.read_h5ad(...)
genes  = np.array(["CD3D", "MS4A1", "GAPDH", "NKG7", "LYZ", "ACTB"])
counts = np.array([[12,  0, 40, 3,  0, 30],
                   [ 0, 25, 35, 1, 18, 22],
                   [ 5,  0, 50, 9,  2, 28]], dtype=float)

# ---------- Strategy A: rank-value encoding (Geneformer-style) ----------
# 1. Normalize each gene by its NON-ZERO median across the corpus. This
#    down-weights housekeeping genes that are uniformly high (GAPDH, ACTB).
nz          = np.where(counts > 0, counts, np.nan)
gene_median = np.nanmedian(nz, axis=0)
norm        = counts / gene_median

# 2. Per cell, rank genes by normalized expression high -> low; drop zeros.
def rank_encode(cell):
    idx = np.where(cell > 0)[0]
    return genes[idx[np.argsort(-cell[idx])]].tolist()

for i in range(counts.shape[0]):
    print(f"cell {i} rank tokens:", rank_encode(norm[i]))
# -> ordered gene tokens only; magnitude lives in POSITION, no value token.

# ---------- Strategy B: value binning (scGPT-style) ----------
N_BINS = 5   # scGPT uses ~51 in practice
def value_bin(cell):
    idx   = np.where(cell > 0)[0]
    vals  = np.log1p(cell[idx])
    ranks = vals.argsort().argsort()                 # per-cell quantile rank
    bins  = np.floor(ranks / len(idx) * N_BINS).astype(int)
    return list(zip(genes[idx].tolist(), bins.tolist()))

for i in range(counts.shape[0]):
    print(f"cell {i} (gene, bin):", value_bin(counts[i]))
# -> aligned (gene token, value-bin token) pairs; magnitude is explicit.

In production you would not hand-roll this — the model ships a tokenizer that also handles the gene vocabulary, special tokens, and truncation:

from geneformer import TranscriptomeTokenizer

# adata must carry raw counts, adata.var["ensembl_id"], adata.obs["n_counts"].
tk = TranscriptomeTokenizer({"cell_type": "cell_type"}, nproc=4)
tk.tokenize_data(
    "data_dir/",        # folder of .h5ad / .loom files
    "output_dir/",
    "tokenized",        # output prefix
    file_format="h5ad",
)
# -> output_dir/tokenized.dataset : rank-ordered token IDs per cell,
#    median-normalized internally and truncated to 2,048 genes.

Common pitfalls & tips

• Gene ID system mismatch. Geneformer's vocabulary is keyed on Ensembl IDs, not gene symbols — map symbols → Ensembl before tokenizing, or most genes become OOV and vanish.
• Wrong normalization state. Geneformer expects raw counts and median-normalizes internally; scGPT expects normalized/log1p input. Double-normalizing (feeding already-log-normed data to Geneformer) corrupts the ranking.
• Check vocabulary coverage. Any gene not in the token dictionary is dropped silently; a low-overlap panel (targeted or non-human data) can lose most of its signal before the model sees a single token. Log the fraction of counts retained.
• Sequence-length truncation drops the tail. Rank encoding keeps only the top ~2,048 genes; lowly expressed but biologically important genes fall off the end, and sparse cells yield short, padding-heavy sequences.
• Value bins are per-cell and relative. "Bin 50" in one cell is not the same absolute expression as "bin 50" in another — binning is computed within each cell, so depth and batch effects shift the mapping; don't compare raw bin IDs across cells.
• Don't forget special tokens. The <cls> position carries the cell-level embedding and <pad> must be masked in attention; omitting or mis-placing them silently degrades every downstream readout.