Temporal AI model predicts drivers of cell state trajectories across human aging

The researchers developed MaxToki, a temporal AI model trained on nearly one trillion gene tokens across the human lifespan, which successfully predicts long-term cell state trajectories and identifies experimentally verified targets for reversing age-related functional decline.

The Gist

Imagine you have a massive library containing the "diaries" of nearly every cell in the human body. These diaries record what genes are active, what the cell is doing, and how it changes over time. For a long time, scientists could only read these diaries one page at a time. They could see what a cell looked like at age 20, or at age 60, but they couldn't easily predict the journey between those ages or figure out what specific events caused a cell to age faster or slower.

Enter MaxToki, a new artificial intelligence model described in this paper. Think of MaxToki not just as a reader, but as a time-traveling storyteller who has read almost a trillion pages of these cellular diaries.

Here is how it works, broken down into simple concepts:

1. The "Time-Traveling Storyteller" (The Model)

Most AI models today are like snapshots. They look at a photo of a cell and say, "This is a heart cell." But cells aren't static photos; they are movies. They grow, change, and age.

MaxToki is different because it was trained on trajectories. Instead of just looking at one cell, it learned to watch a whole movie of a cell's life from birth to old age.

• The Analogy: Imagine learning to drive. A standard AI might show you a picture of a car and ask, "What is this?" MaxToki, however, has watched millions of videos of cars driving from a stoplight to a highway, learning how the speed, gears, and steering change over time. Because of this, it can predict what happens next in the movie, even if it's never seen that specific car before.

2. Learning the "Language" of Cells

Cells speak a language made of genes. MaxToki learned this language by reading a massive library called Genecorpus-175M, which contains about 175 million single-cell snapshots from healthy and diseased tissues.

• The Trick: Instead of memorizing exact numbers (which can be messy due to technical errors), MaxToki learned the ranking of genes. It learned that if Gene A is the "most important" in a cell, and Gene B is second, that relationship matters more than the exact volume of their activity. This is like learning that a sentence makes sense because of the order of the words, not just the spelling.

3. The Two-Stage Training

MaxToki didn't learn everything at once. It had a two-step education:

• Stage 1 (The Vocabulary Lesson): It learned to generate single cells. It practiced writing a "page" of a cell's diary so perfectly that if you gave it a few words, it could finish the sentence.
• Stage 2 (The Plot Lesson): It learned to connect those pages into a story. It was shown a sequence of cells (e.g., a young heart cell, then an older one) and asked to predict:
  1. How much time passed? (e.g., "It took 10 years to get from this state to that state.")
  2. What does the future look like? (e.g., "If we wait 20 years, what will this cell look like?")

4. Predicting the Future and Rewinding Time

Once trained, MaxToki could do some magic tricks:

• Filling in the Blanks: If you showed it a young cell and an old cell, it could accurately guess what the cell looked like in the middle (the "intervening" years).
• Seeing the Unseen: It was tested on cell types and ages it had never seen before. It still worked! It used "in-context learning," meaning it looked at the pattern of the story it was given and applied the rules of aging it had learned to the new character.
• Detecting "Fast-Forward" Buttons: The model could tell if a disease was making cells age faster than normal. For example, it looked at lung cells from heavy smokers and said, "These cells look 5 years older than they should be." It did the same for Alzheimer's patients, detecting that their brain cells were aging faster than healthy controls.

5. Finding the "Villains" and "Heroes" of Aging

The most exciting part is that MaxToki didn't just predict the future; it suggested how to change it.

• The Simulation: Scientists asked MaxToki: "What if we turned off this specific gene? Would the cell age faster or slower?"
• The Discovery: The AI identified several genes that, when overactive, act like a "fast-forward" button for aging in heart cells.
• The Proof: The team took these predictions to the lab. They overexpressed (turned up) these "villain" genes in human heart cells grown in a dish and in mice.
  • Result: The cells aged faster, became stiff, and the mice developed heart failure.
  • Conclusion: The AI was right. It found real biological drivers of aging.

Why This Matters

Think of aging as a long, winding road. For decades, we've been trying to fix potholes on this road by guessing. MaxToki is like a GPS that can see the whole road ahead. It can tell us:

1. Where the road is going to break down (predicting disease).
2. Which specific rocks (genes) are causing the potholes.
3. If we remove those rocks, can we smooth out the road and make the journey longer and healthier?

This paper shows that we can now use AI to not just observe aging, but to simulate interventions in a computer before we ever touch a patient. It's a powerful new tool for finding cures for heart disease, dementia, and other age-related conditions by programming our cells to take a healthier path through time.

Technical Summary

1. Problem Statement

Current foundational AI models for biology (e.g., Geneformer, scGPT) typically treat cell states as static snapshots. They lack the ability to model temporal dynamics, meaning they cannot predict how cellular responses unfold over time or how perturbations affect long-term cell state trajectories (such as those seen in development, aging, or disease progression).

• Limitation: Existing models cannot forecast future cell states, infer past states, or predict the time required to transition between states.
• Challenge: Aging occurs over decades, making longitudinal sampling of the same individual in clinically inaccessible tissues (like the heart or brain) impossible. There is a need for a model that can learn shared aging trajectories from population-based data to predict interventions that reverse or slow age-related decline.

2. Methodology: MaxToki

The authors developed MaxToki, a temporal AI model designed to generate past, intervening, and future cell states across dynamic trajectories.

A. Data Curation

• Genecorpus-175M: A pretraining corpus of ~175 million single-cell transcriptomes from healthy and diseased human tissues. Malignant cells and immortalized lines were excluded to avoid mutational noise.
• Genecorpus-Aging-22M: A specialized corpus of ~22 million cells from ~3,800 donors spanning birth to the 90s. This was used to construct ~100 million simulated aging trajectories stratified by cell type and sex.
• Short Trajectories: Partial reprogramming data (OSKM) from T cells, fibroblasts, and endothelial cells to model rapid state transitions.

B. Architecture and Training Strategy

MaxToki is a Transformer decoder model trained in two stages:

1. Stage 1 (Pretraining): The model learns to generate single-cell transcriptomes using an autoregressive objective.
   • Input Representation: Cells are encoded as rank value encodings (genes ranked by relative expression within a cell, scaled by corpus-wide expression). This deprioritizes housekeeping genes and emphasizes transcription factors while being robust to technical artifacts.
   • Scale: Two variants were trained: a 217M parameter model and a 1B parameter model.
   • Optimization: Utilized NVIDIA BioNeMo, FlashAttention-2, and mixed-precision (bf16) training to handle massive context sizes and GPU memory constraints.
2. Stage 2 (Temporal Training): The context window was extended to 16,384 tokens using RoPE (Rotary Positional Embeddings) scaling.
   • Prompt Structure: The model receives a context trajectory (2–3 cells + time tokens between them) and a query (either a target cell or a target time).
   • Dual Tasks:
     • Task 1 (Timelapse Prediction): Given a context and a query cell, predict the numerical time elapsed (using Mean Squared Error (MSE) loss with continuous numerical tokenization).
     • Task 2 (Cell Generation): Given a context and a time duration, generate the resulting cell transcriptome (using Cross-Entropy (CE) loss).

C. In Silico Perturbation Strategy

To identify drivers of aging, the authors developed a strategy where specific genes in the query cell are "inhibited" (removed from the rank encoding) or overexpressed. The model then predicts the new timelapse.

• If the predicted timelapse is shorter, the perturbation is deemed rejuvenating.
• If the predicted timelapse is longer, the perturbation is deemed pro-aging.

3. Key Contributions

1. First Temporal Foundation Model for Biology: MaxToki is the first model capable of learning continuous cell state trajectories across the human lifespan, moving beyond static snapshot analysis.
2. Continuous Numerical Tokenization: The model treats time as a continuous numerical variable rather than categorical bins, allowing it to generalize to unseen ages and interpolate between states.
3. Self-Supervised Interpretability: The model learned to prioritize transcription factors (TFs) without explicit labeling, identifying them as key drivers of state transitions through attention mechanism analysis.
4. Experimental Validation Pipeline: The study bridges in silico prediction with rigorous in vitro and in vivo validation, confirming that AI-predicted aging drivers cause functional decline in real biological systems.

4. Key Results

A. Predictive Performance

• Generalization: MaxToki successfully predicted timelapses for held-out ages (e.g., ages 21, 31, 41 not seen in training) and held-out cell types (e.g., specific microglia subtypes).
• Accuracy: The model achieved a correlation of 0.77 for held-out ages and 0.85 for held-out cell types, significantly outperforming linear regression baselines and simple age-matching heuristics.
• Fidelity: Generated cells were indistinguishable from ground truth in embedding spaces, correctly classified by external cell-type classifiers, and maintained single-cell resolution (not averaging into doublets).

B. Disease and Aging Inference

• Age Acceleration: The model inferred significant age acceleration in disease states it had never seen during training:
  • Lung Fibroblasts (Pulmonary Fibrosis): Inferred 15 years of age acceleration.
  • Lung Epithelial Cells (Heavy Smoking): Inferred 5 years of age acceleration.
  • Microglia (Alzheimer's Disease): Inferred 3 years of age acceleration.
  • Resilience: Microglia from patients with Alzheimer's resilience (neuropathology present but no cognitive decline) did not show age acceleration, suggesting the model captures disease-specific vs. general aging signals.

C. Experimental Validation of Pro-Aging Drivers

The model predicted novel pro-aging targets in cardiac cell types. Top candidates were validated experimentally:

• ZBTB16: Inhibition in human cardiac endothelial cells caused increased senescence and cell death.
• P4HA1 and RASGEF1B:
  • In Vitro: Overexpression in human iPSC-derived cardiomyocytes and primary cardiac fibroblasts induced gene networks associated with inflammation, mitochondrial dysfunction, and senescence.
  • Functional Decline: Perturbed cardiomyocytes showed slowed calcium cycle kinetics and increased rhythm irregularities.
  • In Vivo: AAV9-mediated overexpression of P4HA1 and RASGEF1B in adult mice caused a significant decline in systolic function (ejection fraction and global longitudinal strain) within 6 weeks.

5. Significance

• Accelerated Discovery: MaxToki provides a framework to screen millions of potential interventions in silico before costly long-term clinical trials, specifically for aging and age-related diseases.
• Mechanistic Insight: The model successfully identified known aging pathways (mTOR, AGE-RAGE, SASP) and novel drivers (P4HA1, RASGEF1B), validating its biological relevance.
• Therapeutic Programming: By understanding the "drivers" of cell state trajectories, this approach enables the programming of therapeutic trajectories to reverse aging or induce resilience in specific tissues.
• Scalability: The use of continuous time tokenization and self-supervised learning on population data offers a scalable solution for studying diseases where longitudinal human data is unavailable.

In conclusion, MaxToki represents a paradigm shift from static genomic modeling to temporal trajectory modeling, offering a powerful tool to decode the mechanisms of aging and discover interventions to program therapeutic cellular trajectories.