Low NextCell gene-token Jaccard (~5%) on zero-shot OPC→oligodendrocyte trajectory — expected for pretrained 1B, or data-prep / inference misconfiguration?

Hi maxToki team!

Thank you for the work, the code is very clear and robust, and was fairly easy to test it myself.

Summary

I'm running zero-shot NextCell inference with the pretrained MaxToki-1B checkpoint on a human OPC → COPs → MFOL → Oligodendrocyte trajectory. Inference completes successfully on 20,000 queries, but evaluation shows very low overlap between predicted and ground-truth gene-token lists (mean Jaccard ~4.8%), and trajectory visualizations suggest predictions collapse to a generic expression profile rather than matching query-specific next states.

I'd appreciate guidance on whether this is expected without trajectory-specific fine-tuning, or if something in our data prep / inference / evaluation pipeline looks wrong. Also were this kind of trajecotry were used in your original model training? I'm trying to understand if this is related to model limitaion, training data, batch effect etc.

---

Goal

Evaluate whether pretrained MaxToki-1B can predict the next cell state along an oligodendrocyte differentiation trajectory using the documented data_prep → assemble-queries → predict --generate-next-cell workflow.

Zero-shot NextCell inference with MaxToki-1B (theodoris-lab/MaxToki) on a human OPC → COPs → MFOL → Oligodendrocyte trajectory completes on 20k queries, but gene-token overlap with ground truth is very low (mean Jaccard 4.8%, median 3.2%). Predictions look collapsed/generic in UMAP and do not align with the reference trajectory in PHATE.

Is this expected without trajectory-specific fine-tuning, or is something wrong in our setup?

---

Setup

Data: Human snRNA-seq; rule-based stages OPC=0, COPs=1, MFOL=2, Oligo=3; grouped by donor (time_group); 100 cells per (donor, cell_type).

Pipeline: README workflow — tokenize → assemble-paragraphs --is-eval → assemble-queries → predict --generate-next-cell.

Key params:

• task-ratio 0.0 (100% NextCell)
• min/max-timepoints 3/4, max-timepoint 3
• time-group-columns time_group
• V1 resources: token_dictionary_v1.json, gene_median_dictionary_v1.json, ensembl_mapping_dict_v1.json

Inference: MaxToki-1B-bionemo, seq-length 16384, top-k 50, temperature 1.0, max-tokens-to-generate 4096, bf16-mixed, time dict at maxtoki_traj_time_dictionary.json.

---

Token / time dictionary

• Tokenize uses base token_dictionary_v1.json (~20k Ensembl gene tokens).
• Paragraphs/queries/inference use auto-generated maxtoki_traj_time_dictionary.json: base dict + appended <boq>, <eoq>, and relative timesteps -3 … +3 (max-timepoint=3, negative timelapses allowed).
• obs["time"] = absolute stage (0–3); question tokens = relative timelapses between cells.
• Gene symbols in h5ad are mapped to Ensembl via ensembl_mapping_dict_v1.json during tokenization.

Concern: README examples use max-timepoint 730; we used 3. Since time tokens are appended to the vocab, could our timestep token IDs be misaligned with the pretrained checkpoint? (Inference log: original vocab_size: 20284, padded vocab_size: 20352.)

---

Results

Metric	Value
Queries	20,000
Mean / median Jaccard	4.8% / 3.2%
Avg generated tokens/query	~353

---

Questions

1. Were this kind of trajectories were presented to the model during training?
2. Is low zero-shot NextCell performance expected on a new trajectory (human OPC→oligo) without second-stage temporal fine-tuning on this dataset?
3. Are our time / max-timepoint settings reasonable? We used discrete stage indices 0–3 with max-timepoint=3. Should we use a different time encoding (e.g. pseudotime, larger max-timepoint)?
4. Is task-ratio=0.0 + assemble-queries the correct eval setup for NextCell inference, or is there a recommended query format / held-out scheme from the paper?
5. Recommended inference hyperparameters for NextCell (top-k, temperature, greedy vs sampling)?
6. Recommended evaluation metrics beyond gene-token Jaccard? (e.g. rank correlation, marker-gene recovery, benchmarks from the paper)
7. Checkpoint path: we used the HuggingFace MaxToki-1B-bionemo weights directly. Is import_hf required before inference for best results?
8. Any common pitfalls with V1 ensembl_mapping_dict / gene symbol mapping that could artificially deflate Jaccard?

---

Thanks for any pointers — especially on whether we should expect reasonable zero-shot NextCell on this setup, or whether trajectory-specific fine-tuning is required before meaningful evaluation.

Start with the README's tokenize → assemble-paragraphs --is-eval → assemble-queries → predict --generate-next-cell workflow, checking the V1 dictionaries, relative time tokens, and Hugging Face checkpoint conversion against the documented setup. Reproduce the reported evaluation and compare generated-token counts, Jaccard, and trajectory visualizations; done means identifying whether the result is expected zero-shot behavior or a concrete data, vocabulary, checkpoint, or inference-configuration problem.