HiST: Spatial Gene Expression Prediction from Histology

Official PyTorch implementation of HiST: A Hierarchical Sparse Transformer for Cross-Modal Spatial Transcriptomics Modeling (ICML 2026).

HiST predicts spatial gene expression from co-registered H&E histology patches. It models measured spatial-transcriptomics (ST) sites as a sparse field on a compact two-dimensional lattice, builds multiscale context with a sparse encoder-decoder, and avoids materializing background tokens.

Paper · OpenReview · Preprocessed Data (Google Drive) · ST preprocessing · SPAN

Highlights

• Lattice-indexed sparse modeling. Image features stay aligned with the observed tissue footprint instead of being padded into a dense slide-level feature tensor.
• Dyadic spatial hierarchy. Coordinate-indexed downsampling expands the receptive field rapidly; occupancy-aware pixel shuffle restores the original measured support in the decoder.
• Local and slide-level context. Sparse regular/shifted-window attention models local geometry, while configurable global tokens provide a low-bandwidth slide-level conditioning path.
• Coordinate-only augmentation. Reflection, shear, rotation, and spot dropout remain available with precomputed patch embeddings, so patch images do not need to be re-encoded every epoch.
• Four patch-encoder regimes. Use precomputed embeddings, a frozen image encoder, parameter-efficient LoRA adaptation, or full end-to-end fine-tuning.

Table of contents

• Architecture at a glance
• Installation
• Quick start
• Related repositories
• Data preparation
• Training and outputs
• Advanced model interfaces
• Implementation guide
• Configuration reference
• Repository map
• Citation and acknowledgements
• License

Architecture at a glance

The default slide_configs: "111-21" uses three encoder resolutions and two decoder resolutions. It projects 1024-dimensional UNI features to width 512, coarsens the sparse lattice twice, and recovers the original measured support with same-scale skip features. Three global tokens carry slide-level context, and a final linear head predicts log-expression at every input location. See configs/model/hist.yaml for the released settings.

HiST overview. Patch embeddings are arranged on a sparse lattice, modeled across spatial scales, and recovered at measured ST sites.

Scope

HiST consumes aligned ST locations and H&E patches or patch embeddings. It does not open raw whole-slide images, perform registration, or choose new prediction locations. Those operations belong in the preprocessing pipeline. Predictions are produced at the supplied measured locations.

Installation

HiST requires Python 3.9 or later. A GPU is recommended but not required by the training entrypoint.

git clone https://github.com/wwyi1828/HiST.git
cd HiST

python -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip
pip install -e .

The model uses PyTorch and DGL. If the default resolver does not select builds compatible with your CUDA runtime, install the matching PyTorch and DGL builds first, then run pip install -e ..

For raw-patch modes with the UNI encoder, obtain the weights under the UNI terms of use and set:

export UNI_WEIGHTS=/path/to/uni_weights.bin

UNI_WEIGHTS is not needed in the default precomputed mode.

Quick start

With preprocessed features under HIST_DATA_ROOT, train the default precomputed-UNI model on one fold:

export HIST_DATA_ROOT=/path/to/processed_datasets

python -m tasks.gene_prediction.main \
  dataset=NCBI_brain \
  test_fold=1 \
  logging.wandb.enabled=false

HiST uses Hydra, so any setting can be overridden with a dotted command-line key. The sections below describe the input format, training modes, outputs, and supported architecture options.

Related repositories

Repository	Role
PatchPreprocess	Prepare aligned gene-expression targets, patch features, and coordinates from HEST-style data for HiST.
SPAN	Related sparse hierarchical framework for WSI classification, segmentation, and survival analysis.
Coordinate-aligned CLAM	WSI patch extraction for the SPAN preprocessing route; the HEST-style HiST route starts from existing paired ST data and patches.

Data preparation

HiST consumes paired spatial-transcriptomics data and H&E patch features prepared in the HEST format.

• Preprocessed data: Download prepared datasets from Google Drive and jump to Quick start.
• Custom data: Process new cohorts with the companion PatchPreprocess repository.

Using PatchPreprocess

HEST-style paired ST data and H&E patches
  -> PatchPreprocess gene/morphology preprocessing
  -> HiST training and evaluation

To prepare your own data, follow the gene and morphology preprocessing guide.

The unified preprocessing path expects HEST-style inputs:

/path/to/HEST1K/
├── st/<sample_id>.h5ad
├── patches/<sample_id>.h5
└── metadata/<sample_id>.json

The unified entrypoint, pipeline/extract_molmor_feats_unified.py, harmonizes genes, matches expression and morphology rows by barcode, derives the coordinate representations below, and writes aligned H5 files. Follow PatchPreprocess for its environment and current CLI, and use one preprocessing revision and gene vocabulary across all folds in an experiment.

Directory layout

Set the root once:

export HIST_DATA_ROOT=/path/to/processed_datasets

A precomputed-feature dataset should look like:

$HIST_DATA_ROOT/
└── ALLVisium/
    └── NCBI_brain/
        ├── gene/
        │   ├── slide_001.h5
        │   └── slide_002.h5
        └── imge_UNI/
            ├── slide_001.h5
            └── slide_002.h5

Raw-patch training uses imge_RAW/ instead of imge_UNI/. Dataset locations are declared in configs/dataset_config/; a path can also be overridden directly with Hydra.

H5 input format

The gene and image files for each slide share the same file stem and must have identical spot ordering, where row i refers to the same barcode or measured location:

File	Key	Status	Shape / dtype	Purpose
gene/<slide>.h5	mol_feats	Required	[N, G], floating point	Non-negative expression in linear space. HiST applies log1p internally; do not store already-log-transformed targets.
	cords	Required	[N, 2], integer	Canonical non-negative, unique lattice coordinates used for evaluation and as the fallback training coordinates.
	float_cords	Required for coordinate augmentation	[N, 2], float32	Continuous near-lattice coordinates transformed before collision-free integer mapping. Without this key, training silently falls back to cords.
	union_gene_names	Required by the default export path	[G], strings	Gene names in exactly the same order as columns of mol_feats.
	hvg_indices / heg_indices	Required when selected by gene_type	one-dimensional integer arrays	Ordered column indices used to select HVGs or HEGs.
	orig_cords	Optional	[N, 2], floating point	Original pixel/physical coordinates for provenance and visualization.
imge_<BACKBONE>/<slide>.h5	mor_feats	Required in precomputed mode	[N, C], floating point	Patch embeddings whose width matches the configured backbone input, for example 1024 for UNI.
imge_RAW/<slide>.h5	mor_feats	Required in raw-patch modes	[N, H, W, 3]	Raw RGB H&E patches on a [0, 255] scale for frozen, lora, or full.

If you add custom spot-level arrays, extend dataset.py so they follow the same filtering and sorting operations as the declared fields.

Coordinate system: continuous geometry, integer operators

The raw ST coordinates do not need to be integers. The coordinates passed to the sparse model do: they must be non-negative, unique, compact integer lattice indices. A typical conversion from pixel coordinates p is:

step_px     = target_lattice_unit_um / source_mpp_um_per_px
float_cords = (p - min(p, axis=0)) / step_px
cords       = collision_free_integer_map(float_cords)

The companion preprocessing code computes this mapping from the source image resolution and the requested physical patch/lattice unit. It retains the original and continuous coordinates because integerization is an approximation, not a lossless coordinate transform. HiST's runtime mapper rounds near-lattice points, resolves collisions by assigning an unused nearby grid site, and translates the result so that both coordinate minima are zero.

Integer dtype alone is not sufficient. Do not pass unnormalized WSI pixel coordinates directly: large gaps create an unnecessarily large lookup canvas, and duplicate coordinates overwrite the one-site/one-token correspondence. The current attention rulebook casts lattice coordinates to signed 16-bit integers before computing the canvas size, so each coordinate maximum must be at most 32766. In practice the memory cost of the temporary canvas usually demands a much tighter range.

The integer lattice defines hierarchical parent-child relationships, attention-window membership, and local relative positions. For irregular or subcellular assays, the lattice unit and collision policy are modeling choices.

Before training, ensure that expression values are finite and non-negative, union_gene_names matches the expression width, cords contains unique non-negative integer pairs with max <= 32766, and float_cords is present when coordinate augmentation is enabled.

Registering datasets and groups

Dataset discovery is a Hydra registry rather than a hard-coded path list. Add a file under configs/dataset_config/, for example:

# configs/dataset_config/my_study.yaml
# @package _global_
datasets:
  cohort_a:
    path: "${oc.env:HIST_DATA_ROOT,./data}/cohort_a"
    gene_dir: "gene"
    gene_type: "hvg"
    num_genes: 2000

Field	Description
path	Dataset root containing the gene and image-feature folders.
gene_dir	Gene-file directory, relative to path or absolute. Defaults to gene.
gene_type	hvg, heg, or null for all columns. The corresponding ranking key must exist when hvg or heg is selected.
num_genes	Positive prefix length of the selected ranking, or null for the complete ranking/all columns.

Select it with dataset_config=my_study dataset=cohort_a. Registry files may also define named dataset_groups combining multiple cohorts. Datasets in a group must share identical gene ordering, morphology width, and coordinate conventions.

test_fold is a one-based value in 1..4. Four-fold splits are generated from slide files at runtime; each cohort should contain at least four slides.

Training and outputs

Inspecting configurations and sweeps

Inspect the fully composed and resolved configuration without opening data or starting training:

python -m tasks.gene_prediction.main \
  --cfg job --resolve \
  dataset=NCBI_brain \
  model.pos_emb=alibi

Use --cfg job --resolve to inspect or archive the fully resolved configuration for reproducibility.

Hydra multirun can execute all four runtime folds:

python -m tasks.gene_prediction.main --multirun \
  dataset=NCBI_brain \
  test_fold=1,2,3,4 \
  logging.wandb.enabled=false

For sweeps, assign each condition a unique, filename-safe logging.results.save_path such as hist_rope; the automatic model slug does not encode every override and can otherwise collide.

Patch-encoder modes

HiST always trains the sparse spatial encoder-decoder and gene head. The mode controls how morphology features are obtained and whether the image encoder is also updated:

model.train_mode	Input folder	Additional trainable parameters	Pixel augmentation	Typical use
precomputed	imge_<patch_backbone>/	None in the patch encoder (identity)	No	Default and fastest; reuse embeddings while retaining coordinate augmentation.
frozen	imge_RAW/	None; image-encoder weights stay frozen	Optional	Re-encode raw patches and experiment with pixel-space augmentation.
lora	imge_RAW/	LoRA adapters on the attention Q/K projections	Optional	Adapt the image encoder with substantially fewer trainable parameters.
full	imge_RAW/	The complete image encoder (full means trainable)	Optional	End-to-end fine-tuning with the highest memory and compute cost.

Raw-patch modes instantiate the UNI ViT-L/16 encoder. Ensure UNI_WEIGHTS points to your local pretrained weights before training. Pixel augmentation is enabled with training.patch_aug=true, independently of coordinate augmentation.

Example LoRA run:

python -m tasks.gene_prediction.main \
  dataset=NCBI_brain test_fold=1 \
  patch_backbone=UNI model.train_mode=lora \
  model.lora_rank=8 model.lora_alpha=16 \
  model.lora_lr_mult=5.0 training.patch_aug=true \
  logging.wandb.enabled=false

model.lora_rank, model.lora_alpha, and model.lora_lr_mult control Q/K adapter capacity, scaling, and learning rate.

Output files

By default, aggregate metrics are written to:

results/gene_prediction/hist.json

Final-epoch per-slide predictions and slide tokens are written under:

results/gene_prediction/predictions/hist/<model_config>/<run>_fold<k>/
results/gene_prediction/slide_token/hist/<model_config>/<run>_fold<k>/

Each prediction H5 contains:

Key	Shape	Meaning
predictions	[N, G]	Non-negative predictions in linear expression space after expm1 and ReLU.
targets	[N, G]	Linear-space targets from mol_feats.
coords	[N, 2]	Sorted integer lattice coordinates used for evaluation.
gene_names	[G]	Byte strings aligned with prediction columns.

Rows are saved in coordinate-sorted order; use the exported coords to map predictions back to the original tissue spots. Slide file stems should be unique across datasets in a group.

Each slide-token H5 contains:

Key	Shape	Meaning
slide_token	[2T, D] by default	The final encoder global tokens followed by the decoder pre-output global tokens. T is the number of valid token initializers; the released model uses T=3, D=512.
token_init_types	[T]	Initializer labels such as max, mean, and fix1e-4.

Setting logging.results.save_predictions=false suppresses both per-slide prediction files and slide-token files. The aggregate JSON remains enabled and stores each metric as a four-element fold array under:

<run>/<model_config>/spot_metrics/<metric>
<run>/<model_config>/slide_avg_metrics/<metric>

Here, spot_metrics is a pooled metric over all flattened spot–gene entries, while slide_avg_metrics computes the metric separately per slide and takes the unweighted mean across slides. Metrics and per-slide predictions are evaluated and saved at the final epoch.

The training objective is mean squared error in log1p expression space. Final evaluation reports coefficient of determination (R2) and Pearson correlation coefficient (PCC) on linear-space values, matching the paper. Their order or a subset can be selected through training.eval_metrics.

Advanced model interfaces

The options below are the supported Hydra controls for architecture experiments.

Topology DSL

model.slide_configs is a compact depth-per-resolution string:

111-21
^^^ ^^
enc dec

Each character gives the Transformer depth at one resolution. In 111-21, the encoder has three resolutions with one layer each and the decoder has depths two and one. The first encoder stage projects the input; later encoder stages coarsen, decoder stages recover, and the prediction head is appended automatically. Depths are single digits. Gene prediction requires one fewer decoder stage than encoder stages, and the released hierarchy assumes kernel 2 and stride 2.

Slide-calibration tokens and embeddings

model.token_init_types accepts an ordered list of global-token initializers:

Initializer	Behavior
mean, max, std	Pool the current slide's input patch features.
fix<scale>	Create a fixed random buffer, for example fix1e-4.
lrn<scale>	Create a learned parameter, for example lrn1e-4.
numeric value	Create a fixed random buffer with that standard deviation.
[]	Disable global tokens entirely.

The default [max, mean, fix1e-4] creates three tokens. With hybrid attention, they exchange information with local spatial tokens; model.share_qkv=true shares Q/K/V projections between the two paths, while the default uses separate projections.

At each resolution change, model.global_strategy controls the token path:

Value	Global-token update
proj	Independent learned linear projection (default).
identity	Identity when widths match; otherwise a learned width projection.
kernel	Reuse the local spatial operator's feature transform followed by its pooling rule.

Final encoder and decoder tokens are saved through the schema described in Output files.

Operator and attention ablations

Use these options to compare spatial operators, attention, and skip connections:

Interface	Supported values	Notes
model.econvs_type	patchconv (default), dwpatchconv, patchmerge	Sparse encoder coarsening; patchmerge uses canonical four-child slots.
model.dconvs_type	pixelshuffle (default), patchinvs	Sparse decoder recovery intersected with cached active support.
model.trans_type, model.dtrans_type	hybrid (default), hybrid_noshift, swin	hybrid_noshift removes shifted windows but keeps global exchange; swin uses regular/shifted local graphs without local–global edges.
model.pos_emb	rpb (default), alibi, rope, none	Learned relative bias, 2D distance ALiBi, 2D rotary Q/K, or an ablation without positional encoding.
model.connection	concat (default), add, none	Encoder–decoder skip fusion.
model.edge_mode	none (default), count, partial	Optional incomplete-neighborhood normalization for sparse patch convolution.

For patchmerge, both non-none edge modes reduce to count normalization; the count/partial distinction applies to sparse patch convolution.

Implementation guide

This section follows the forward path in source-code order. For supported experiment knobs and commands, use Advanced model interfaces.

1. Input projection and padding

PredictionModel obtains one patch feature per measured location, and the first encoder block projects it to the model width. SPAN_Padder adds zero-padding tokens along boundaries so repeated stride-2 coarsening and recovery have compatible shapes; this padding is removed before prediction. The configured slide-calibration tokens follow a parallel low-bandwidth path and exchange information with local tokens under hybrid attention.

2. Efficient coordinate augmentation

Coordinate augmentation is built in lib/utils/coord_aug.py and src/span/preprocessing/transforms.py. For one slide it:

1. centers float_cords and composes mirror, shear, and rotation into one 2 x 2 matrix;
2. transforms all N coordinates in one matrix multiplication;
3. optionally drops spots and applies the same mask to image features, expression targets, and aligned metadata;
4. translates the transformed support to the origin and maps it to collision-free integer coordinates; and
5. sorts coordinates and every aligned spot-level tensor with the same index.

This changes only an N x 2 coordinate array and reuses the precomputed N x C patch embeddings, so patches are neither re-extracted nor re-encoded. drop.p is the probability of dropping a spot: p=0.75 retains roughly 25% when its multiplier is 1. The optional epoch-dependent probability schedule is experimental with persistent DataLoader workers.

3. Sparse spatial downsampling and upsampling

The current default encoder uses sparse patchconv with kernel 2 and stride 2. The implementation keeps the feature table sparse while using a compact integer ID map to gather only the neighborhoods that contribute to active output sites. Missing child sites contribute zero. Optional count and partial edge modes can compensate for incomplete neighborhoods; the released config uses none.

HiST also provides a dyadic patchmerge path. For each fine coordinate c:

parent = (c0 // 2, c1 // 2)
slot   = 2 * (c0 % 2) + (c1 % 2)

Up to four children are scattered into canonical slots, concatenated, and projected to the next-stage width.

Dyadic encoder-decoder. Each encoder stage groups active child sites into coarse parent tokens; the decoder reverses the hierarchy on the cached active support. The green C illustrates the slide-level global-token stream; the released configuration carries three such tokens.

Before coarsening, the encoder caches the active support at each resolution. Sparse pixel shuffle expands every parent feature into four candidate children, maps them to 2 * parent + offset, and keeps only candidates present in the matching cached support. Recovery therefore returns predictions only at measured locations and remains aligned with same-resolution skip features.

Sparsity note. Feature tensors and attention edges remain sparse, but some operators create an H x W integer ID or occupancy map over the tight coordinate bounding box. Memory is therefore sensitive to coordinate range as well as spot count. Compact lattice coordinates are required for efficient sparse operations.

4. Sparse window and global attention

AttentionBuilder scatters active token IDs onto a tight lattice map, partitions occupied sites into regular and shifted windows, and creates DGL edges only between valid token pairs. Under hybrid attention, a second graph adds local-to-global, global-to-local, and global-to-global interactions. window_size is a half-side parameter, so the released value 3 produces 6 x 6 lattice windows.

TransformerLayer computes Q/K/V projections, scores only graph edges, normalizes incoming edges with DGL edge_softmax, and aggregates values by message passing. It never instantiates a dense N x N attention matrix. Repeated Transformer layers in the same configured stage reuse the stage's constructed graphs.

Local attention supports learned relative positional bias (rpb, the default), 2D distance-based ALiBi, 2D RoPE, or no positional encoding. RPB and ALiBi are applied only to local–local edges. RoPE rotates Q/K on both hybrid paths, using the lattice positions for local tokens and the origin for global tokens.

5. Encoder, decoder, and prediction head

SPAN_Encoder caches local features, global features, and coordinates at every resolution. SPAN_Decoder reverses the hierarchy and combines the recovered features with encoder history using concat, add, or no skip connection. The final linear layer produces one log-expression vector per input location. Predictions are converted back with expm1 and clamped to non-negative values.

Configuration reference

HiST uses Hydra. The full configuration is defined in configs/gene_prediction.yaml and configs/model/hist.yaml; use --cfg job --resolve to inspect every setting. The primary experiment controls are:

Key	Meaning
dataset_config, dataset, dataset_group	Select a dataset registry, one dataset, or a validated harmonized group.
test_fold	One-based runtime fold in 1..4.
patch_backbone, model.train_mode	Select stored features or precomputed, frozen, lora, or full image-encoder training.
model.lora_rank, model.lora_alpha, model.lora_lr_mult	Q/K adapter rank, scaling, and learning-rate multiplier.
model.slide_configs	Transformer depth at each encoder and decoder resolution.
model.econvs_type, model.dconvs_type	Spatial coarsening and recovery operators.
model.trans_type, model.dtrans_type, model.pos_emb	Attention and positional encoding.
model.token_init_types, model.global_strategy, model.share_qkv	Slide-calibration token initialization and resolution path.
model.connection, model.edge_mode	Skip fusion and sparse-neighborhood normalization.
training.lr, training.weight_decay, training.epochs	AdamW and run-length settings.
training.init_bias_with_mean	Initialize gene-head bias from mean training-fold log1p expression.
training.coord_aug, training.patch_aug	Coordinate and raw-patch augmentation.
training.eval_metric_freq, training.eval_metrics	Evaluation interval and final R2/PCC metrics.
logging.wandb.enabled	Enable or disable Weights & Biases.
logging.results.dir, logging.results.save_path, logging.results.save_predictions	Result namespace and per-slide export control.

Repository map

Path	Responsibility
configs/	Hydra task, model, dataset, and logging configuration.
tasks/gene_prediction/	H5 loading, model assembly, training, evaluation, and export.
src/span/	Hierarchical sparse encoder-decoder and configuration builder.
src/span/layers/	Sparse resolution changes, graph attention, Transformers, and positional encoding.
src/span/preprocessing/	Padding, coordinate transforms, integer mapping, and spatial utilities.
lib/backbones/	Raw-patch UNI encoder, pixel augmentation, and LoRA injection.
lib/utils/	Coordinate augmentation, fold splitting, metrics, and result logging.
assets/	README figures exported from the HiST manuscript.

Citation and acknowledgements

If you use HiST, please cite:

@inproceedings{
  wu2026hist,
  title={Hi{ST}: A Hierarchical Sparse Transformer for Cross-Modal Spatial Transcriptomics Modeling},
  author={Weiyi Wu and Xinwen Xu and Xingjian Diao and Siting Li and Zhi Wei and Alma Andersson and Jiang Gui},
  booktitle={Forty-third International Conference on Machine Learning},
  year={2026},
  url={https://openreview.net/forum?id=ptbzlHzmEv}
}

Data preparation is shared with PatchPreprocess and builds on the HEST-1k data ecosystem. Raw-patch runs may also depend on UNI. Please cite these projects and follow their respective data, model-weight, and software terms when used.

License

HiST is released under the MIT License. This license does not replace the separate terms governing external datasets, pretrained weights, or upstream software.