CortexJEPA

CortexJEPA is a joint-embedding predictive model for cortical spatial transcriptomics. The model treats a tissue section as a set of local cell/spot neighborhoods, where each cell is described by sparse gene expression and by its spatial position in the cortex. Instead of learning from cells as independent expression vectors, CortexJEPA learns to predict masked cell representations from their surrounding spatial context. The resulting encoder captures both molecular state and local cortical organization.

This repository contains the inference package and released checkpoints used for marmoset and macaque cortical analyses. The released models generate cell-level embeddings, predict cortical layer labels, regress the marmoset PrAl axis, and infer a 16-target set of cortical multimodal attributes.

Model Overview

CortexJEPA follows the JEPA idea at the scale of spatial transcriptomics: the model predicts representations of hidden target cells, rather than reconstructing raw gene counts. A local spatial region is divided into context cells and masked target cells. The context encoder observes the visible cells, the target encoder produces the representation target for hidden cells, and a predictor learns to infer those target representations from context.

The model has five main components:

• Sparse gene encoder. Nonzero genes in each cell are mapped through learned gene embeddings, weighted by expression values, pooled into one cell token, and projected to the model hidden space.
• Spatial Transformer encoder. Cell tokens are processed by a 12-layer Transformer with 768 hidden dimensions and 12 attention heads.
• 3D positional encoding. Cell coordinates are injected through rotary positional encoding, allowing attention to depend on cortical geometry.
• JEPA predictor. A 6-layer, 384-dimensional predictor uses context cell representations and target coordinates to predict masked target-cell representations.
• Downstream heads. Fine-tuned checkpoints attach classification or regression heads to the encoder for cortical layer, PrAl, and multimodal prediction.

The default inference window contains up to 2048 cells/spots per local region. For layer prediction, the released heads use seven labels: L1, L2, L2-Al, L3, L4, L5, and L6.

Released Models

Folder	Species	Task	Output
models/marmoset/pretrained	marmoset	JEPA encoder	768-dimensional cell embeddings
models/marmoset/layer	marmoset	cortical layer classification	L1 to L6, including L2-Al
models/marmoset/pral	marmoset	PrAl regression	one scalar per cell
models/marmoset/multimodal	marmoset	multimodal cortical regression	16 scalar targets per cell
models/macaque/pretrained	macaque	JEPA encoder	768-dimensional cell embeddings
models/macaque/layer	macaque	cortical layer classification	L1 to L6, including L2-Al

The marmoset models use a 34,805-gene vocabulary. The macaque models use a 16,266-gene vocabulary. Cross-species inference is supported through the bundled one-to-one ortholog table for human, mouse, macaque, and marmoset gene symbols; genes without a mapped target-vocabulary ortholog are omitted.

The 16 marmoset multimodal targets are: PrAl, All-Neuron-Density, L3-Neuron-Density, L5-Neuron-Density, L2-Neuron-Density, L6-Neuron-Density, L4-Layer-Thickness, HighVisA-Network, LatV1-Network, L3-Layer-Thickness, PV/(PV+SST), FG1, MedV1-Network, L4-Neuron-Density, Myelinmap, and FG2.

Installation

Clone the repository with Git LFS so the checkpoint weights are downloaded:

git lfs install
git clone https://huggingface.co/jiangshuan6/CortexJEPA
cd CortexJEPA
pip install -e ".[hub]"

For a local checkout that already contains the weights:

pip install -e .

Input Data

The high-level API accepts an anndata.AnnData object or an .h5ad path. Input data should contain:

• expression values in .X, .raw.X, or a selected .layers[...];
• spatial coordinates in adata.obsm["spatial"] or coordinate columns in adata.obs;
• a tissue grouping column, normally adata.obs["segment"].

If segment is missing, the package can write a default segmentation:

import anndata as ad
from cortexjepa import add_default_segment_to_h5ad

adata = ad.read_h5ad("sample.h5ad")
adata = add_default_segment_to_h5ad(
    adata,
    segment_column="segment",
    n_segments=80,
)

Basic Usage

Load the released marmoset and macaque model bundles from the cloned repository:

from cortexjepa import MacaqueCortexJEPA, MarmosetCortexJEPA

marmoset = MarmosetCortexJEPA.from_pretrained(
    "models/marmoset",
    device="cuda",
)
macaque = MacaqueCortexJEPA.from_pretrained(
    "models/macaque",
    device="cuda",
    load_pral=False,
)

Generate CortexJEPA embeddings:

adata = marmoset.embed(
    adata,
    expression_key="X",
    batch_size=1,
)

emb = adata.obsm["X_cortexjepa_marmoset"]

Predict cortical layer:

adata = marmoset.predict_layer(
    adata,
    expression_key="X",
    batch_size=1,
)

layer = adata.obs["cortexjepa_marmoset_layer"]

Run cross-species layer prediction by converting input genes through the ortholog table before inference:

adata = macaque.predict_layer(
    adata,
    expression_key="X",
    source_species="human",
    batch_size=1,
)

layer = adata.obs["cortexjepa_macaque_layer"]

Predict the marmoset PrAl axis:

adata = marmoset.predict_pral(
    adata,
    expression_key="X",
    batch_size=1,
)

pral = adata.obs["cortexjepa_marmoset_pral"]

Run the marmoset multimodal model:

model = MarmosetCortexJEPA.from_pretrained(
    "models/marmoset",
    device="cuda",
    load_pretrained=False,
    load_layer=False,
    load_pral=False,
    load_multimodal=True,
)

adata = model.predict_multimodal(
    adata,
    expression_key="X",
    batch_size=1,
)

Multimodal outputs are written to adata.obs with keys such as cortexjepa_marmoset_multimodal_PrAl, cortexjepa_marmoset_multimodal_Myelinmap, and cortexjepa_marmoset_multimodal_prediction_count.

Run multiple marmoset outputs together:

result = marmoset.predict(
    adata,
    expression_key="X",
    batch_size=1,
)

adata = result.adata
adata.write_h5ad("sample.cortexjepa_marmoset.h5ad")

Default marmoset keys are:

adata.obsm["X_cortexjepa_marmoset"]
adata.obs["X_cortexjepa_marmoset_count"]
adata.obs["cortexjepa_marmoset_layer"]
adata.obs["cortexjepa_marmoset_layer_prediction_count"]
adata.obs["cortexjepa_marmoset_pral"]
adata.obs["cortexjepa_marmoset_pral_prediction_count"]

Notes

• CortexJEPA predictions are model-derived annotations. They should be interpreted together with experimental metadata, tissue quality, and independent biological validation.
• Cross-species predictions depend on one-to-one ortholog coverage and the selected target model vocabulary.
• Spatial coordinates and segmentation affect local region construction, so consistent preprocessing is important when comparing samples.
• The repository is intended for research use.

Citation

If you use CortexJEPA, please cite the CortexJEPA paper. The final citation metadata can be added here after publication.