Reasoning Across Space

Code for Reasoning Across Space: Tiny Recursive Models for Spatial Omics (1st FoGen Workshop, ICML 2026).

Tiny Recursive Models (TRMs) have emerged as a promising alternative paradigm to scaling models. Instead of scaling along the depth axis, they apply one small network to its own output repeatedly, carrying a latent scratchpad alongside a current answer and refining both. Effective depth grows with the number of recursions while the parameter count stays fixed (Jolicoeur-Martineau, 2025).

While they excel at discrete tasks like Sudoku and the ARC-AGI benchmark, their utility on real-world biological data was unclear. Those benchmarks are self-verifiable; biological readouts are noisy, partially annotated and continuous.

Spatial omics is a hard test. Models there are usually local, passing information around a neighbourhood graph or splitting the slide into patches, and are typically fit on the section they are asked to explain. Sections are few and annotations scarce, which is where trading parameters for recursions should pay.

Two tasks, one 896K-parameter model:

Task	Data	Predicts	Scored by
SDI — spatial domain identification	DLPFC, 12 Visium sections	one of 7 cortical layers per spot	ARI
H2G — histology to gene expression	HER2ST, 36 ST sections	302 gene values per spot	PCC-k, RVD

Neither uses test-time training. SDI holds out a whole slice, excluded from gene selection, PCA and training; H2G holds out a whole slide.

Getting started

Python >= 3.11.

uv venv && source .venv/bin/activate
uv pip install -e .          # [cache] to rebuild HER2ST features, [download] for DLPFC

With data in place (see Data below):

python -m spatial_trm.sdi.train model=trm holdout_slice=151670
python -m spatial_trm.h2g.train model=trm

Each run writes best.pt, last.pt, results.json, eval_history.{json,png} and per-slice predictions to output_dir, plus its git hash, parameter count and active SDPA backends.

Both entrypoints are Hydra apps:

python -m spatial_trm.h2g.train model=trm_act_shallow training=pcc_halt \
    model.per_class_halt=true model.halt_prob_threshold=0.4

Data

DLPFC (~2 GB). Obtain the 12 spatialLIBD sections in Space Ranger layout, one directory per slice holding {sid}_filtered_feature_bc_matrix.h5, tissue_positions_list.csv and {sid}_truth.txt, then convert:

python -m spatial_trm.sdi.data.download --raw_dir data/dlpfc_raw --skip-download

The Google Drive folder behind the no---skip-download path carries only the *_truth.txt annotations; the expression matrices come from spatialLIBD.

HER2ST (~0.4 GB raw, ~0.7 GB cached). Caching needs a GPU and Prov-GigaPath weights: accept the license for prov-gigapath/prov-gigapath and huggingface-cli login first.

python -m spatial_trm.h2g.data.fetch --out data/her2st
python -m spatial_trm.h2g.data.build_cache --data data/her2st \
    --splits data/splits/her2st_b1_holdout.csv \
    --out cache/stem_gigapath_her2st --holdout B1 --backbone gigapath

Breast cancer (zero-shot target). Build one projection per LOSO fold — a shared projection would evaluate checkpoints in the wrong basis, and the tooling refuses it.

python -m spatial_trm.sdi.data.download_breast_cancer --raw_dir data/breast_cancer_raw
python -m spatial_trm.sdi.data.preprocess_breast_cancer \
    --raw_dir data/breast_cancer_raw \
    --dlpfc_processed_dir data/dlpfc_processed_loso_151670 \
    --out_dir data/breast_cancer_processed_loso_151670 --label_type fine

Per-fold DLPFC features are built by scripts/sdi_loso.sh.

See docs/REPRODUCING.md for the runs behind the paper's tables and figures.

Repo structure

src/spatial_trm/
  models/     TRMSpatial + the three non-recursive baselines, shared by both tasks
  core/       training loop, ACT halt objectives, batching, evaluation helpers
  sdi/        DLPFC: data, cross-entropy losses, ARI evaluation, train.py
  h2g/        HER2ST: data, MSE losses, PCC/RVD evaluation, train.py
  analysis/   recursion ablations, LOSO tables, zero-shot transfer, halt-step analysis
configs/      Hydra configs, one tree per task: {sdi,h2g}/{model,training,data}
scripts/      end-to-end drivers
docs/         paper.pdf, figures, reproduction notes

The tasks share the model, recursion schedules, EMA, checkpoint selection and halt objectives. A task implements core.task.Task and supplies only its loss, halt statistic and metrics, so sdi/ and h2g/ stay thin.

Halt objectives

Each supervises the halt logit with BCE against a soft target from a quality statistic on a per-step randomised spot holdout, disjoint from the prediction loss spots.

Recipe	Target	Task
loss_proxy	sigmoid((tau - loss_holdout) * s)	SDI (held-out CE), H2G (held-out MSE)
ce_halt	sigmoid((tau - CE) * s), optionally per class	SDI
mse_halt / rvd_halt	sigmoid((tau - stat) * s)	H2G
pcc_halt	sigmoid((PCC - tau) * s), note the sign	H2G
random_k	fire at a randomly drawn depth	SDI

model.per_class_halt=true emits one logit per class or gene and averages them for the slide-level decision.

Results

Figures and numbers below are from the paper; see docs/paper.pdf for the full tables, ablations and appendix.

Depth from repetition, saturating around three blocks. The outer-step boundary is read-only, so N_outer and N_ZH enter only through their product: equal products give a bit-identical forward pass.

Latent recursion carries real weight — dropping N_ZL from 6 to 1 costs up to 0.25 ARI. On a frozen model, 1 latent step is near-random and 6 recover sharp laminae.

Recursion beats matched depth. A Universal Transformer reusing one block 21 times has the same effective depth and reaches 0.483 mean ARI against TRM's 0.559. The gain is the factorisation into latent refinement, prediction refinement and an outer loop, not weight sharing.

The latent transfers across tissue. DLPFC checkpoints run on breast cancer without retraining. The recursive latent gains more from added spatial context than a non-recursive baseline (+0.15 vs +0.115 ARI), so it encodes something complementary to spatial position rather than redundant with it.

Regression saturates earlier. With frozen histology features the deep-recursion update carries the gains; at the default config every N_ZL from 1 to 8 gives the same PCC to four decimals.

Predicted expression still clusters into regions a pathologist recognises:

ACT needs adapting for continuous targets. Classical ACT assumes a solved/not-solved signal. Supervising the halt head per gene and moving the threshold inside the converged halt-probability distribution gives input-adaptive halting: easy slides stop at the floor, hard slides run to K_max, ~42% less compute at matched PCC.

Slide-level halt recipes collapse to one extreme, but their BCE term still cuts RVD 2-3x as a variance-aware regulariser.

Citation

This work (docs/paper.pdf):

@inproceedings{rajwade2026reasoning,
  title     = {Reasoning Across Space: Tiny Recursive Models for Spatial Omics},
  author    = {Rajwade, Dhruva Abhijit and Rapsomaniki, Marianna},
  booktitle = {1st FoGen Workshop, ICML},
  year      = {2026},
}

The Tiny Recursive Model this builds on:

@article{jolicoeurmartineau2025less,
  title   = {Less is More: Recursive Reasoning with Tiny Networks},
  author  = {Jolicoeur-Martineau, Alexia},
  journal = {arXiv preprint arXiv:2510.04871},
  year    = {2025},
  url     = {https://arxiv.org/abs/2510.04871},
}

MIT License — see LICENSE.