Benchmarking

Code Availability

The benchmarking, evaluation, and experiment-configuration code in this repository is fully public and reproduces every comparison reported in the paper. The core Guided Tokenization implementation — k-mer selection and motif-preserving tokenization — is currently under IP review and is available from the corresponding author upon reasonable request.

Where included under results/, the token-selection outputs (the actual selected k-mer/motif lists used in each experiment) are provided as data files, so every reported comparison can be reproduced exactly as run. The selection algorithm that produced those lists — i.e. how the k-mers were chosen, not which k-mers were chosen — is the IP-reviewed part described above. In short: the "what" is public and reproducible; the "how it was chosen" is available from the corresponding author upon reasonable request.

Experiment infrastructure for the Guided Tokenization (GT) paper revision. Built on top of the guided_tokenizer library and examples/ at the repo root — see those for the underlying GT tokenizer/vocab-expansion API this code wraps.

Guided Tokenization

Guided Tokenization (GT) is a supervised, task-specific vocabulary-adaptation strategy for fine-tuning genomic language models. Standard tokenizers (BPE, k-mer) can fragment biologically meaningful subsequences — e.g. splitting a TATA box promoter motif across multiple tokens — which loses signal that's relevant to the downstream classification task. GT identifies task-relevant motifs from labeled training data (via class-specific k-mer frequency and/or gradient-based attribution), adds them to the tokenizer's vocabulary, and initializes their embeddings from the mean of their constituent subword embeddings so they inherit the pretrained model's existing representations. At tokenization time, these motifs are preserved as single tokens rather than being split apart by the base BPE merges.

This repository contains the benchmarking, training, and evaluation harness used to produce every reported comparison in the paper. The token-selection methodology itself (how motifs are scored and selected) is described in the paper's Methods and is under IP review — see Code Availability above.

Layout

configs/
  models/          per-model config (HF path, conda env, architecture)
  datasets/        per-dataset config (source, path, num_classes, ...)
  experiments/      concrete experiment configs (model + dataset + tokenization + training)
  sweep.yaml        list of experiment config paths to run as a batch
src/
  data/             dataset loading (HF/local/GUE/GBM) + preprocessing
  tokenization/      GT wrapper, k-mer extraction, random/freq baselines, vocab expansion
  models/           model loading dispatch across the 4 model families
  training/         trainer, metrics, timing/memory callbacks
  analysis/         motif coverage, FLOPs estimate, results aggregation
scripts/
  download_data.py               pull GUE/GBM/resLens into /data/gt_benchmarks
  generate_gue_dataset_configs.py  materialize one dataset config per GUE subtask
  prepare_fasta.py                convert a dataset's train split into per-class FASTA (for KMC)
  extract_kmers_kmc.py            KMC-based class-specific unique k-mer extraction (CPU, offline)
  run_kmer_pipeline.py            batch driver: prepare_fasta + extract_kmers_kmc over many dataset configs
  select_tokens.py                rank KMC output into augment/prioritize token lists (CPU-only modes)
  precompute_guided_tokens.py     in-process class-specific k-mer extraction (pure-Python fallback; prefer the KMC path above for real runs)
  precompute_weighted_tokens.py   offline gradient/embedding-norm token scoring — needs GPU + a loaded model
  generate_sweep_configs.py       materialize the full experiment matrix from the constants at its top
  run_experiment.py               run ONE experiment from a config file
  run_sweep.py                    expand sweep.yaml into a (config, conda_env) job manifest
  launch.sh                       execute the manifest: one conda env at a time, parallel across GPUs
results/            gitignored experiment outputs

Running one experiment

conda activate mgx_seqlens   # must match the model config's conda_env
python benchmarking/scripts/run_experiment.py --config benchmarking/configs/experiments/main/example_seqlens_bpe_reslens_short_seed0.yaml

Running the full sweep

# 1. Materialize every experiment config from the matrix (main + ablation + gt_variants)
python benchmarking/scripts/generate_sweep_configs.py

# 2. Precompute motifs for every (task, strategy) that needs them — do this
#    BEFORE run_sweep, since k-mer extraction and gradient scoring are slow
#    offline steps not meant to run per-seed.
python benchmarking/scripts/precompute_guided_tokens.py --dataset-config benchmarking/configs/datasets/reslens_amr_short.yaml --out benchmarking/results/motifs/reslens_amr_short_guided.txt
python benchmarking/scripts/precompute_weighted_tokens.py --model-config benchmarking/configs/models/seqlens_89m.yaml --dataset-config benchmarking/configs/datasets/reslens_amr_short.yaml --out benchmarking/results/motifs/reslens_amr_short_weighted.txt

# 3. Expand the sweep into a per-conda-env job manifest
python benchmarking/scripts/run_sweep.py --sweep benchmarking/configs/sweep.yaml --manifest-out /tmp/jobs.tsv

# 4. Launch — processes one conda env block at a time, parallel across GPUs within a block
bash benchmarking/scripts/launch.sh /tmp/jobs.tsv

# 5. Aggregate results into summary tables
python benchmarking/src/analysis/collect_results.py --results-root benchmarking/results --out-dir benchmarking/results/summary

Data

Large datasets live on fast storage at /data/gt_benchmarks/, symlinked into the repo as benchmarking/../data (i.e. guided_tokenizer2/data). scripts/download_data.py pulls resLens/GBM/GUE from HuggingFace. 16S data (~150GB) must be scp'd manually from the university HPC into /data/gt_benchmarks/16s/.

CPU preprocessing pipeline: FASTA -> KMC -> token selection

Requires kmc/kmc_dump/kmc_tools (KMC3, install via conda install -n mgx_seqlens -c bioconda -c conda-forge kmc) and the kneed Python package (pip install kneed, only needed for the GPU-based weighted-token step later).

# 1. Convert a task's train split to one FASTA file per class
python benchmarking/scripts/prepare_fasta.py --dataset-config benchmarking/configs/datasets/cds_vs_noncds.yaml

# 2. KMC: find k-mers present in one class and absent from every other
#    (chained N-1 subtract ops per class for N-class tasks — see the
#    script's docstring). k=4..12 by default.
python benchmarking/scripts/extract_kmers_kmc.py --task-name cds_vs_noncds --k-min 4 --k-max 12

# Or run both steps for many dataset configs at once:
python benchmarking/scripts/run_kmer_pipeline.py \
    --dataset-configs benchmarking/configs/datasets/*.yaml benchmarking/configs/datasets/gue/*.yaml \
    --k-min 4 --k-max 12

# 3. Rank into final token lists (augment = new vocab tokens; prioritize =
#    overlap with an existing tokenizer's vocab). Weighted/gradient tokens
#    are NOT produced here — see precompute_weighted_tokens.py, GPU-only.
python benchmarking/scripts/select_tokens.py \
    --unique-kmers-tsv /data/gt_benchmarks/cds_vs_noncds/kmers/unique_kmers.tsv \
    --mode augment prioritize --top-n 100 \
    --tokenizer omicseye/seqLens_4096_512_89M-at-base-multi \
    --out-dir /data/gt_benchmarks/cds_vs_noncds/tokens

Two important corrections made while wiring this up (verified against the actual downloaded data, not assumed from the paper plan):

• resLens is 13-class (Sequence/Class columns; 12 drug-class labels
  • non_ARG), not binary — the original scaffolding guessed label/2 classes from the task name alone. See configs/datasets/reslens_amr_*.yaml.
• GBM repo ids on the Hub carry a Genomic_Benchmarks_ prefix (katarinagresova/Genomic_Benchmarks_demo_coding_vs_intergenomic_seqs), not the bare short names used in the paper plan.
• KMC counts canonical (double-stranded) k-mers by default, folding a k-mer and its reverse complement into one counter — wrong for tokenization, where only the literal forward-strand substring matters. extract_kmers_kmc.py passes -b to disable this.
• GUE's EPI_* tasks (enhancer-promoter interaction) are sequence-PAIR classification (enhancer,promoter,label columns), which the current single-sequence loader doesn't support — excluded from GUE_TASKS/GUE_UNSUPPORTED_TASKS in src/data/benchmark_tasks.py rather than silently mishandled.
• 16S (4,288 classes) is excluded from the batch k-mer pipeline: the chained-subtraction approach is O(N) per class / O(N^2) total, which is fine for the largest task here (GUE virus_species_40, 25 classes) but would need a different algorithm (e.g. one global union pass) at 4,288 classes — not attempted until the 16S data itself is downloaded.
• Cleanup-ordering race in the first batch run: extract_kmers_kmc.py originally deleted a class's own "base" k-mer db as soon as that class's own chain-subtraction finished. Since all classes subtract in parallel and finish at different times, a fast class could delete its base db while a slower class's chain still needed to read it, causing kmc_tools: Cannot open file ...base.kmc_pre errors. This silently dropped rows (each per-class failure was caught and logged to stderr, not fatal) rather than corrupting output outright, but it meant undercounting for any multi-class task where the race window was wide enough to hit — confirmed in the first run for the 5 highest-class-count tasks (gue_virus_species_40 25 classes, gue_fungi_species_20 20, reslens_amr_short/reslens_amr_long 13, gue_virus_covid 9); binary and low-class tasks were unaffected in practice (narrower window) but not immune in principle. Fixed by moving cleanup to a separate phase (barrier after ALL classes finish subtracting, before ANY class's files are deleted); all 5 affected tasks were re-extracted from scratch with the fix and verified.

Known gaps in the underlying guided_tokenizer library

The library implements motif-aware tokenization and mean-subword initialization, but does NOT implement class-specific k-mer selection, gradient-attribution weighted-token selection, or a "combined" variant — those are implemented from scratch here in src/tokenization/ (see kmer_extraction.py, weighted_tokens.py). gt_wrapper.py also works around a latent OOV-detection bug in guided_tokenizer.model_utils.augment_model (it checks convert_tokens_to_ids(m) is None, which real HF fast tokenizers never return for OOV tokens) by reimplementing the mean-subword init correctly in src/tokenization/vocab_expansion.py.