Implementation-aligned method#
This page describes the algorithm implemented by the current public code. It is also the reference for terminology used in the revised manuscript.
1. UOT coupling between adjacent slices#
For cells in two adjacent observed slices, DeepSpatial builds a hybrid cost from normalized spatial distance, gene-expression cosine distance, and a cell-type mismatch penalty:
$$ C_{ij} = \alpha d^{x}{ij} + (1-\alpha)d^{g}{ij} + \lambda_{\mathrm{type}}d^{c}_{ij}. $$
The first two distance matrices are independently divided by their maximum. The cell-type term is zero up to a small numerical epsilon for a matched label and one for a mismatched label. The coupling is obtained with entropy-regularized unbalanced Sinkhorn transport. Sampled endpoint pairs are drawn from the normalized coupling with replacement.
DeepSpatial.setup_data exposes all solver controls: alpha_spatial,
cell_type_penalty, uot_reg, uot_tau, uot_num_itermax, and
uot_stop_thr. These values are stored in data_config when a checkpoint is
saved.
2. Joint continuous flow matching#
Spatial coordinates $x$, gene expression $g$, and one-hot cell-type vectors $c$ follow the same linear probability path. For each modality $m\in{x,g,c}$,
$$ m_t=(1-t)m_0+t m_1, \qquad u_m=m_1-m_0. $$
Cell type is therefore a continuous relaxation, not a discrete categorical process during training. The network predicts three velocity fields and all three are optimized with MSE velocity matching. The implementation first takes the per-sample mean-squared error over each modality. This averaging provides the single feature-dimension normalization used by the total objective:
$$ \mathcal L_m = \mathbb E!\left[\frac{1}{d_m}\lVert v_m-u_m\rVert_2^2\right], \qquad \mathcal L=\mathcal L_x+\lambda_g\mathcal L_g+\lambda_c\mathcal L_c. $$
The defaults are $\lambda_g=0.1$ and $\lambda_c=10$. There is no
cross-entropy objective in the final joint flow. During inference, the
continuous cell-type state is integrated together with $x$ and $g$ and is
converted to a label by argmax only when the trajectory output is assembled.
The standard pipeline expects cell-type annotations because they contribute to the UOT cost, the joint flow state, and the reconstructed labels.
3. GiT architecture#
GiT uses the following modality-specific projections:
spatial coordinates:
Linear(2, H), producing one spatial token;gene expression: non-overlapping gene patches, each projected by a two-layer MLP (
Linear–SiLU–Linear);continuous one-hot cell-type state:
Linear(K, H), used as global conditioning.
Sinusoidal timestep embeddings are used for flow time, physical depth, and the
slice gap. Their conditioning vector modulates adaLN-Zero Transformer blocks.
Separate output heads predict spatial, gene-expression, and cell-type
velocities. The implementation does not use a two-layer spatial MLP or
nn.Embedding for cell type.
4. Explicit bidirectional ODE generation#
Forward trajectories integrate physical flow time from $t=0$ to $t=1$. For a terminal cell sampled from the upper slice, reverse generation starts at physical time $t=1$. The numerical solver still advances an auxiliary time $u:0\rightarrow1$, using the exact reverse-time reparameterization
$$ q(u)=s(1-u), \qquad \frac{dq}{du}=-v(q(u),1-u). $$
The original endpoint ordering $(z_0,z_1)$ and forward-oriented $\Delta z=z_1-z_0$ are retained. The implementation does not swap endpoints, negate $\Delta z$, run a forward solve, or flip a completed trajectory to emulate reverse integration.
For a requested target depth $t$, the parent direction is sampled as
$$ P(\mathrm{forward}\mid t)=1-t, \qquad P(\mathrm{backward}\mid t)=t. $$
Thus direction selection is depth dependent rather than
Bernoulli(0.5).
5. Density-aware depths and gene sparsification#
Let the endpoint slices contain $N_0$ and $N_1$ cells. Virtual depths are sampled with density proportional to the linear interpolation of endpoint cell counts,
$$ p(t)=\frac{2[(1-t)N_0+tN_1]}{N_0+N_1}. $$
The inverse CDF is used directly; when $N_0=N_1$, this reduces to uniform depth sampling. The number of virtual cells in a gap is
$$ N_{\mathrm{virtual}}= \max!\left(1,\left\lfloor \frac{N_0+N_1}{2}\frac{|z_1-z_0|}{\mathrm{thickness}} \right\rfloor\right). $$
After ODE generation, each virtual cell retains the top-$k$ predicted genes, where $k$ is the number of non-zero genes in its sampled parent cell. This is the final gene-sparsification step.
There is no oversampled raw population, spatial subset optimization, or
post-hoc spatial pruning. Every generated virtual cell is retained. The
corresponding private helper is therefore named _generate_virtual_cells.
6. API defaults and reported benchmark settings#
Public API defaults provide a general starting point. The reported Fig. 2 benchmark explicitly overrides the parameters shown below; reproducibility scripts should set them rather than relying on defaults.
Parameter |
Public API default |
Reported Fig. 2 setting |
|---|---|---|
|
50,000 |
total number of cells in the observed input slices |
|
50.0 |
50.0 |
|
0.8 |
0.1 |
|
0.05 |
1.0 |
|
50 |
50 |
|
$10^{-2}$ |
$10^{-2}$ |
training epochs |
100 |
10 |
All checkpoint metadata includes the data, model, and training configuration so that a saved run can be audited without inferring values from later defaults.