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.

Full-Diagram Deep Dives

Three of the strategies above don't reduce to a single hand-drawn LEGO panel — they combine multiple embeddings per token (scPRINT-2) or move the token itself up a level, from gene to whole cell (STATE, STACK). Full walkthroughs below.

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)
    bins  = np.clip(bins, 0, N_BINS - 1)              # guard the top-rank edge case
    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:

# tested against geneformer>=1.0 — the TranscriptomeTokenizer constructor
# signature has changed across major versions; pin your version and check
# the CHANGELOG before upgrading an existing tokenization pipeline.
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.

Strategy C: single-nucleotide tokenization (Evo/HyenaDNA-style). For DNA rather than expression data, the same "hand-roll it first" approach looks like this — no gene vocabulary, no binning, just one token per base:

# Vocab = 5 (4 bases + N for ambiguous/unknown); no merge operations at all.
nuc_vocab = {"A": 0, "T": 1, "G": 2, "C": 3, "N": 4}

def nucleotide_tokenize(seq: str) -> list[int]:
    return [nuc_vocab.get(c.upper(), 4) for c in seq]

seq = "ATGCATGCN..."   # in practice: a 1Mb+ window read from a reference/FASTA
tokens = nucleotide_tokenize(seq)
# -> every character becomes a token; zero OOV by construction, but the
#    context length equals len(seq) exactly — at genome scale this is only
#    tractable with a linear-time backend (Hyena/Mamba), not full attention.

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.