JAE joint embedding explorer

Where the data and the method come from, then the settings and counts that the saved analysis records.

Data source. NeurIPS 2021 Open Problems BMMC CITE-seq (GEO GSE194122). Detailed dataset provenance and usage conditions are on the Dataset page.

Notebook lineage

Author: JAE — Qiao Liu & Wanwen Zeng, winners of the NeurIPS 2021 Joint Embedding task (team Amateur); reimplemented in PyTorch for this demo.
Source / context: PC5251 teaching demo: a walk through one joint-embedding method, then controlled experiments that change one part of its training at a time.
Notebook: jae_walkthrough.ipynb

What the saved analysis records

The data, settings and counts behind the panels on this page. The right-hand column names the analysis step that produced each value.

Property	Recorded value	Produced by
Dataset	NeurIPS 2021 Open Problems BMMC CITE-seq · GEO GSE194122	dataset record
Cells	66,175 bone-marrow mononuclear cells	stage 1 · preprocessing
Measured per cell	13,953 genes (GEX) and 134 surface proteins (ADT), read simultaneously	stage 1 · preprocessing
Labels kept out of the input	45 cell types · 9 batches (site × donor) · cell-cycle phase	stage 1 · preprocessing
Cell-cycle scoring	log1p(.X) then zero-centre scale, matching upstream JAE · 43 S genes and 54 G2M genes used	stage 1 · preprocessing
Agreement with the dataset's own phase call	31.32% of cells	stage 1 · preprocessing
Transformation	TruncatedSVD to 100 components on GEX (25.32% of variance); ADT kept at 134 and not reduced	stage 2 · transformation
Negative share of the model input	21.47%, which the model's output layer (a ReLU, never negative) cannot reproduce	stage 2 · transformation
Latent layout	cell type 0–45 · batch 45–54 · cell cycle 54–56 · free 56–61	stage 3 · modelling
Training	batch 256 · learning rate 0.001 · max 150 epochs, stop after 10 without improvement · 61 epochs run	stage 3 · training
Loss weights	0.7 / 0.2 / 0.05 / 0.05 · batch term = uniform	stage 3 · training
Variants trained	14 configurations, each with its own 2-D projection and its own six measures	stage 3 · training
Evaluation set	66,175 cells; silhouette scores on a stratified subsample of 8,000 cells	stage 4 · evaluation
Held-out split	6,618 cells (10%) chosen at random among cells (seed 42), not a held-out donor	stage 4 · evaluation
Analysis saved	2026-07-30 17:57:45	stage 4 · saving

Counts in this table are for the full dataset. The matrix picture in the Design matrix section draws on a display sample of 5,004 cells × 160 columns.

All embeddings, metrics and training curves on this page come from the saved analysis.

How the six metrics are computed here, compared with the scib benchmarking package

NMI: scib tunes the resolution of Louvain clustering to maximise NMI (scib.cl.opt_louvain). Louvain is not available here, so KMeans is run over a range of k and the best NMI is kept. Both procedures tune on the metric itself; KMeans assumes compact, convex clusters and Louvain does not.
cell-type ASW: Same formula as scib, (ASW + 1)/2, but computed on a stratified subsample rather than on all cells, because silhouette scoring grows with the square of the cell count.
batch ASW: Same per-cell-type 1 − |silhouette| construction as scib, on the same subsample. Groups with fewer than 10 cells or a single batch are skipped, as in scib.
graph connectivity: Same definition as scib (for each cell type, the share of its cells in the largest connected part of a nearest-neighbour graph), with the graph built on the subsample.
cell-cycle conservation: scib compares, batch by batch, how much variance the S and G2M scores explain before and after integration. Here this is done with PCA and linear regression, and the “before” reference is the 234-value model input rather than scib's unintegrated data.
trajectory conservation: scib uses diffusion pseudotime (scanpy). Here it is replaced by the graph distance from the root cell on a nearest-neighbour graph, compared with the dataset's reference pseudotime by Spearman correlation.

What a row and a column are here

Cells × measured features (genes + proteins)

66,175 × 13,953 + 134

One row is one cell. The columns form two blocks, one per kind of measurement (genes, proteins), processed separately until Stage 2 joins them.

66,175

Cells

13,953

Genes (GEX)

134

Proteins (ADT)

Cell types in y

Batches in y

How a row is built

One row is one bone-marrow cell. Its genes and proteins were measured in the same droplet, which is what lets one embedding combine them.
The columns are two blocks side by side: 13,953 gene columns, then 134 protein columns. Both are normalised the same way; only the gene block is later reduced.
An entry is the normalised abundance of that molecule in that cell — not a count, and not a probability.
The 45 cell types, the 9 batches and the cell-cycle phase are labels y. They are never columns of X, and they reach the model only through the Stage 3 loss.

The picture below shows a sample of the matrix. It draws from 5,004 cells sampled across all cell types and 160 columns (26 marker genes and all 134 proteins); the counts above are for the full dataset.

Stage 1 · Data & preprocessing

One sample = one bone-marrow cell, measured two ways at once

CITE-seq reads two signals from the same cell: the mRNA it is currently transcribing (GEX, 13,953 genes) and the proteins sitting on its surface (ADT, 134 proteins, each detected by an antibody). Both describe the same cell, and a joint embedding combines them into one description.

Input X

one sample: one bone-marrow cell
GEX: 13,953 gene counts
ADT: 134 protein counts
shape of X: 66,175 × 14,087

Labels y (not inputs)

cell type: 45 types, annotated by hand
batch: 9 site × donor
cell cycle: continuous S / G2M scores
pseudotime: erythroid only, 10,373 cells

Task type

The goal is a useful description of each cell, not a prediction → unsupervised representation learning.
But JAE is semi-supervised: the labels never enter the input, only the loss function (the score training minimises).

The design-matrix bridge: what a row, a column and a single entry each are

Before any transformation, get the shape of X clear. In the figure below, every row = one cell, every column = one gene or one protein, and the intersection x_ij = the normalised abundance of that molecule in that cell. Click the heatmap to pick a cell (row) and a column, or use the menus and the Random cell button; the X[i,:] and X[:,j] panels follow.

Pick a column: Filter cells by lineage:

X: a display slice of 66,175 × 14,087

Drawn here: a 150 row × 160 column display slice, sorted by cell type. The full matrix is 66,175 × 14,087. The three colour bands at the right edge are the labels y (cell type coloured by lineage, batch, S score); they are not columns of X. Clicking the heatmap changes the selected row and column at once.

X[i,:] → one cell, in its native units

Selected cell AGTTCCCTCTTCCTAA-1-s1d1 · true type pDC · batch s1d1 · its label was never seen during training. 9 of its 12 most strongly expressed genes are housekeeping genes, which most cells use heavily (mitochondrial, ribosomal, MALAT1 and similar). The most strongly expressed genes are not the most informative, which is why the model needs the SVD and the labels to find what separates cell types.

X[:,j] → one column across all cells

Column j = CD3E (gene mRNA). Each dot is one cell of the display sample, in its lineage's row; the short vertical bar is that lineage's mean, and the long line marks the selected cell. The further apart the lineage means, the better this column tells cell types apart, which is what lets the model recover cell type from X.

Where is y? Cell type, batch and cell-cycle phase are labels. They are drawn as colour bands to the right of the matrix, not as columns of X. They take no part in the SVD that follows; they enter the model only in Stage 3, through the loss function.

Stage 2 · Transformation

The two modalities take different routes, then get concatenated

The gene block has nearly fourteen thousand columns and is extremely sparse (most entries are zero), so TruncatedSVD reduces it to 100 dimensions: 100 combined columns that keep its strongest patterns. The protein block has only 134 columns, each one a known protein, so it is left alone. Whether to reduce is decided for each kind of measurement separately.

GEX counts (66,175 × 13,953)
   │ L1 normalise ×1e4 → log10(1+x)
   └─ TruncatedSVD(100) ───────────┐
                                   ├─→ concat → X_model (234 dims) → JAE
ADT counts (66,175 × 134)       │
   │ L1 normalise ×1e4 → log10(1+x) │
   └─ no reduction, keep 134 ──────┘

y (cell type / batch / cycle) bypasses this whole path.
It appears only in the Stage 3 loss function.

The SVD is fitted on all the cells. The upstream script 1_pca_pretrain.py pools both competition phases before running it, so the cells later used for evaluation also shape the transformation. This is called a transductive setup; the competition rules allowed it.

How much variance each of the 100 components explains

Bars: the share of gene-expression variance each component explains; line: the running total. Hover a bar to read it. All 100 components together explain 25.3%. Single-cell gene data are extremely sparse, and most of the remaining variance is sequencing noise, so here the reduction removes noise rather than compressing without loss.

Stage 3 · Model & training

The 61 numbers are split into blocks, each with its own job

JAE is an autoencoder: it compresses each cell's 234 input values into 61 numbers, the latent vector, and learns to rebuild the input from them. Each stretch of those 61 numbers has a job. JAE's cell-type "classifier" is not a separate network: it simply reads a slice of the latent vector (2_ae_pretrain_cite.py:272-274 in the upstream code). The chart shows the selected cell's 61 numbers:

JAE in one sentence: the first 45 numbers should peak at the cell's type, the middle 9 should be flat (no batch stands out), the next 2 should track the cell cycle, and the last 5 have no target, left free for the network to use. Hover a bar to read its value.

Four losses, with the weights from the upstream code

Loss	Latent block it acts on	Weight
MSE(decoder(z), x)	all 61 dims	0.70
CE(z[0:45], cell type)	first 45 dims	0.20
CE(z[45:54], uniform)	middle 9 dims	0.05
MSE(z[54:56], S/G2M scores)	next 2 dims	0.05

Training minimises the weighted sum of these four terms. MSE (mean squared error) scores how well the decoder rebuilds the input x from the latent vector z; CE (cross-entropy) scores a block against a target. The third line is the key one. It does not ask the model to get the batch wrong, which would still require knowing it. It asks for equal probability over all 9 batches, removing batch information from those 9 numbers. This is a simpler alternative to adversarial batch correction, which trains a second network to detect the batch.

Training: cell-type accuracy rises, batch accuracy falls to chance

Solid lines: how often the cell-type block names the right type, and the batch block the right batch, after each epoch (one pass over the training cells). After 61 epochs: cell-type accuracy 86.5%, batch accuracy 7.1% (chance is 11.1%). The dashed line is the reconstruction loss, scaled to the chart height. Cell-type accuracy climbing while batch accuracy falls to chance shows the two label losses doing their separate jobs. Hover the chart to read each epoch.

The upstream output layers can only give values ≥ 0. In the upstream code, the last layer of both the encoder and the decoder is a ReLU, which turns negative values into zero. As a result (1) the decoder output is always ≥ 0, while 21% of the model input is negative (SVD components can be negative), so that part cannot be rebuilt; and (2) the two cell-cycle numbers are also ≥ 0, while roughly 70% of the true S/G2M scores are negative. This demo keeps both ReLUs; Experiment 5 (the Lab · Output activations section) trains the jae_linear_out variant with both removed.

Stage 4 · Evaluation & output

Scoring the embedding, then tracing a point back to a real cell

Six metrics for the full JAE embedding

Metric	All cells	Held-out only	Rewards
NMI agreement between clusters and cell type	0.7775	0.7559	biology kept
cell type ASW are cell types well separated	0.6407	0.6240	biology kept
cell cycle cons. did the cell-cycle signal survive	0.8778	0.7897	biology kept
trajectory cons. did the erythroid trajectory survive	0.9011	0.9380	biology kept
batch ASW are batches mixed within a cell type	0.8379	0.8254	batch removed
graph conn. does each cell type stay in one piece	0.9316	0.9344	batch removed
Mean of the six	0.8278	0.8113

Higher is better for every metric. Four reward keeping biological structure, two reward removing batch effects, and the last row is their plain average. The first column scores all 66,175 cells; the second only the 6,618 held-out cells, whose labels the model never saw.

Biological conservation ↔ batch removal is a trade-off

Each dot is one of the 14 trained variants. The horizontal axis is the mean of the four biological-conservation metrics; the vertical axis is the mean of the two batch-removal metrics. Top right is best, but no variant leads on both at once: the harder a variant removes batch effects, the more biological structure tends to go with them. The highlighted dot is the variant chosen in the Variant menu below; click a dot to choose it. Experiment 1 (the Lab · Batch-loss strength section) moves along this trade-off.

How the six metrics are computed. They follow the definitions of the scib benchmarking package, reimplemented with numpy and scikit-learn. The main differences: NMI runs KMeans over a range of k instead of tuning Louvain resolution; trajectory conservation uses distances along a nearest-neighbour graph instead of diffusion pseudotime; and the silhouette-based metrics use a stratified subsample of 8,000 cells. All 6 metrics differ from scib in implementation; the Source evidence section lists each difference.

The embedding: 61 numbers drawn as a 2-D map (UMAP)

Colour by: Variant:

TBNKMyeloidErythroidProgenitorOther

Variant jae_full, loss weights [0.7, 0.2, 0.05, 0.05]. Each point is one cell (5,004 shown, sampled across cell types), placed by a UMAP of that variant's 61 numbers. Use Colour by and Variant to change the view; the lineage filter in the Design matrix section applies here too. Hover a point to identify it, and click it to select that cell here and in the Design matrix section. Maps of different variants are rotated or reflected to line up (orthogonal Procrustes); their shapes are unchanged.

Traceability: which cell is this point?

barcode: AGTTCCCTCTTCCTAA-1-s1d1
true type: pDC
lineage: Myeloid
batch: s1d1
S / G2M: -0.54 / -0.74
erythroid pseudotime: not applicable
strongest genes: MALAT1, MT-CO1, B2M, MT-CO3, MT-CO2
training label: held out (never trained on)

Every point on the map is one real cell: here are its barcode, batch, annotated cell type and strongest genes, and the bars show its 134 surface-protein values (the three largest are labelled; hover a bar to read it).

Lab: what each part of training changes

Six controlled experiments, each with one factor changed at a time: the batch-loss weight, the label losses, the batch target, the input measurements, the output layers and the random seed. Every variant shares the same data, the same train/validation split, the same architecture and the same random seed; only the factor named in each experiment differs. All 14 variants were trained in advance.

Each experiment shows its variants' 2-D maps of the same 5,004 cells, drawn on shared axes, then their metrics. Use “Colour every embedding below by” to recolour all maps; hover a point to identify the cell.

Colour every embedding below by: s1d1s1d2s1d3s2d1s2d4s2d5s3d1s3d6s3d7

Experiment 1

Raising the batch-loss weight from 0 to 0.5: is biological structure removed too?

The only thing that changes is λ_batch, the weight on the batch loss (the third line of the loss table). The paper uses 0.05. Start with the maps coloured by batch: if batch correction is working, the colours within each group of cells should be thoroughly mixed.

λ_batch = 0

[0.7, 0.2, 0.00, 0.05]

λ_batch = 0.05

[0.7, 0.2, 0.05, 0.05] ← paper

λ_batch = 0.2

[0.7, 0.2, 0.20, 0.05]

λ_batch = 0.5

[0.7, 0.2, 0.50, 0.05]

Variant	NMI	type ASW	cycle cons.	traj. cons.	batch ASW	graph conn.	bio mean	batch mean	mean of 6	cell-type acc.
λ_batch = 0 jae_batch0	0.792	0.641	0.891	0.928	0.850	0.944	0.813	0.897	0.841	84.8%
λ_batch = 0.05 jae_full	0.777	0.641	0.878	0.901	0.838	0.932	0.799	0.885	0.828	86.5%
λ_batch = 0.2 jae_batch20	0.785	0.643	0.873	0.757	0.821	0.945	0.765	0.883	0.804	84.9%
λ_batch = 0.5 jae_batch50	0.786	0.640	0.886	0.930	0.845	0.954	0.811	0.899	0.840	86.9%

Going from λ_batch=0 to 0.5, the batch-removal mean moves from 0.897 to 0.899 while the biological-conservation mean moves from 0.813 to 0.811. The paper's 0.05 sits at 0.885 / 0.799. Compare these gaps with the seed-to-seed spread in Experiment 6.

Experiment 2

Which loss is the latent structure actually coming from?

Starting from "no training at all", add the supervision back one term at a time. These four variants answer the question: can an autoencoder grow cell-type structure on its own, or does it need labels to push it there?

① No training at all

TruncatedSVD to the same 61 dims

② Plain autoencoder

[1.0, 0.0, 0.00, 0.00]

③ + cell type

[0.7, 0.3, 0.00, 0.00]

④ Full JAE

[0.7, 0.2, 0.05, 0.05]

Variant	NMI	type ASW	cycle cons.	traj. cons.	batch ASW	graph conn.	bio mean	batch mean	mean of 6	cell-type acc.
① No training at all svd61	0.675	0.544	0.780	0.897	0.874	0.915	0.724	0.894	0.781	—
② Plain autoencoder ae_only	0.511	0.486	0.671	0.742	0.855	0.813	0.603	0.834	0.680	1.1%
③ + cell type ae_celltype	0.779	0.648	0.865	0.919	0.838	0.930	0.803	0.884	0.830	84.3%
④ Full JAE jae_full	0.777	0.641	0.878	0.901	0.838	0.932	0.799	0.885	0.828	86.5%

First, the accuracy of reading cell type off an argmax over the first 45 dimensions: the plain autoencoder manages only 1.1% (nothing asked those dimensions to represent type, so it should be near the 2.2% chance level), rising to 84.3% once type supervision is added, and 86.5% for the full JAE. No surprise there.

The surprise is the NMI column (agreement between clusters and cell types): 0.675 → 0.511 → 0.779 → 0.777. The plain autoencoder (0.511) is worse than not training at all (0.675). A 61-dimensional bottleneck optimised purely for reconstruction does not organise cell-type structure better — it actually scrambles structure the SVD already had (remember the output ReLU truncates half the directions; see experiment 5). Only once type supervision is added does NMI jump to 0.779.

So the conclusion here is not "labels are the icing on the cake". It is: in this architecture, the readable cell-type structure in the latent space comes mainly from the labels, not from the reconstruction loss.

Experiment 3 · Mechanism check

The key control: swap "uniform distribution" for "classify correctly"

This is the most important experiment on the page. JAE's batch term pushes the 9 batch logits toward a uniform distribution. The natural question is: what if it were an ordinary "get the batch right" term instead? Intuitively, getting it right means batch information is still sitting in the latent space — exactly what we did not want. The two variants are identical apart from that one line of the loss.

Push toward uniform (what JAE does)

CE(z[45:54], uniform)

Classify the batch correctly (control)

CE(z[45:54], true batch)

Variant	NMI	type ASW	cycle cons.	traj. cons.	batch ASW	graph conn.	bio mean	batch mean	mean of 6	cell-type acc.
Push toward uniform (what JAE does) jae_full	0.777	0.641	0.878	0.901	0.838	0.932	0.799	0.885	0.828	86.5%
Classify the batch correctly (control) jae_batch_supervised	0.608	0.524	0.919	0.704	0.666	0.806	0.689	0.736	0.704	86.0%

Batch classification accuracy: 7.1% with the uniform target versus 97.4% with the classify-correctly target (chance is 11.1%). Batch-removal score: 0.885 versus 0.736. That confirms the mechanism: it is the target distribution of the loss, not the mere presence of a batch term, that does the batch correction. Switch to classifying correctly and those 9 dimensions become a high-fidelity encoder of batch identity — precisely what we were trying to avoid.

Experiment 4

What does the joint embedding actually buy over a single modality?

Zero out either the 100 gene (GEX) dimensions or the 134 protein (ADT) dimensions. Architecture, parameter count and training configuration all stay exactly the same; the model simply cannot see one of the two measurements. This measures what combining the two adds.

GEX only

134 protein dims zeroed

ADT only

100 GEX SVD dims zeroed

GEX + ADT jointly

all 234 dims visible

Variant	NMI	type ASW	cycle cons.	traj. cons.	batch ASW	graph conn.	bio mean	batch mean	mean of 6	cell-type acc.
GEX only jae_gex_only	0.761	0.645	0.840	0.937	0.791	0.855	0.796	0.823	0.805	81.2%
ADT only jae_adt_only	0.697	0.555	0.867	0.680	0.851	0.884	0.700	0.867	0.755	80.0%
GEX + ADT jointly jae_full	0.777	0.641	0.878	0.901	0.838	0.932	0.799	0.885	0.828	86.5%

Mean of six: GEX only 0.805, ADT only 0.755, joint 0.828. But look per metric rather than at the total: cell-type ASW favours GEX, while cycle conservation favours ADT. And trajectory conservation is actually highest for GEX alone (0.937) rather than the joint model (0.901) — adding a modality does not improve every metric monotonically. The two modalities carry information that does not fully overlap, which is both the point of "joint" and its cost.

Experiment 5 · Found by reading the source

What do those two ReLUs cost?

Upstream puts a ReLU on the last layer of both the encoder and the decoder. The consequences: the decoder output is always ≥ 0, but 21% of the model input is negative (SVD components are signed), so that part is structurally unreconstructable; and the two cell-cycle dimensions are also pinned at ≥ 0, while about 70% of the true S/G2M scores are negative. jae_linear_out removes both ReLUs and changes nothing else.

As upstream (ReLU output)

encoder/decoder last layer = ReLU

Both ReLUs removed

encoder/decoder last layer = linear

Variant	NMI	type ASW	cycle cons.	traj. cons.	batch ASW	graph conn.	bio mean	batch mean	mean of 6	cell-type acc.
As upstream (ReLU output) jae_full	0.777	0.641	0.878	0.901	0.838	0.932	0.799	0.885	0.828	86.5%
Both ReLUs removed jae_linear_out	0.808	0.617	0.686	0.680	0.838	0.950	0.698	0.894	0.763	91.2%

Mean of six: 0.828 (as upstream) versus 0.763 (ReLUs removed), a difference of -0.064. Cycle conservation alone: 0.878 → 0.686. A question raised by reading the source code is settled by a controlled experiment.

Experiment 6 · Seed-to-seed spread

How much does the score move if you only change the random seed?

This measures how much the metrics move when nothing changes but the random seed: the yardstick for the differences in Experiments 1–5. The three variants differ only in torch.manual_seed.

Seed 0

torch.manual_seed(0)

Seed 1

torch.manual_seed(1)

Seed 2

torch.manual_seed(2)

Variant	NMI	type ASW	cycle cons.	traj. cons.	batch ASW	graph conn.	bio mean	batch mean	mean of 6	cell-type acc.
Seed 0 jae_full	0.777	0.641	0.878	0.901	0.838	0.932	0.799	0.885	0.828	86.5%
Seed 1 jae_full_seed1	0.779	0.643	0.818	0.938	0.839	0.909	0.795	0.874	0.821	84.9%
Seed 2 jae_full_seed2	0.758	0.653	0.874	0.855	0.810	0.953	0.785	0.882	0.817	77.7%

The three seeds land between 0.817 and 0.828 on the mean of six, a spread of 0.010. Half-ranges per metric: NMI ±0.010, type ASW ±0.006, cycle cons. ±0.030, traj. cons. ±0.041, batch ASW ±0.014, graph conn. ±0.022. Differences in the mean of six smaller than about 0.010 are within seed-to-seed variation.

All results

14 variants × six metrics

Evaluate on:

Variant	NMI	type ASW	cycle cons.	traj. cons.	batch ASW	graph conn.	bio mean	batch mean	mean of 6	cell-type acc.
No training · SVD to 61 dims svd61	0.675	0.544	0.780	0.897	0.874	0.915	0.724	0.894	0.781	—
Plain autoencoder ae_only	0.511	0.486	0.671	0.742	0.855	0.813	0.603	0.834	0.680	1.1%
AE + cell type ae_celltype	0.779	0.648	0.865	0.919	0.838	0.930	0.803	0.884	0.830	84.3%
JAE λ_batch=0 jae_batch0	0.792	0.641	0.891	0.928	0.850	0.944	0.813	0.897	0.841	84.8%
Full JAE (paper config) jae_full	0.777	0.641	0.878	0.901	0.838	0.932	0.799	0.885	0.828	86.5%
JAE λ_batch=0.2 jae_batch20	0.785	0.643	0.873	0.757	0.821	0.945	0.765	0.883	0.804	84.9%
JAE λ_batch=0.5 jae_batch50	0.786	0.640	0.886	0.930	0.845	0.954	0.811	0.899	0.840	86.9%
Batch term classifies correctly jae_batch_supervised	0.608	0.524	0.919	0.704	0.666	0.806	0.689	0.736	0.704	86.0%
GEX only jae_gex_only	0.761	0.645	0.840	0.937	0.791	0.855	0.796	0.823	0.805	81.2%
ADT only jae_adt_only	0.697	0.555	0.867	0.680	0.851	0.884	0.700	0.867	0.755	80.0%
Output ReLUs removed jae_linear_out	0.808	0.617	0.686	0.680	0.838	0.950	0.698	0.894	0.763	91.2%
Full JAE · seed 1 jae_full_seed1	0.779	0.643	0.818	0.938	0.839	0.909	0.795	0.874	0.821	84.9%
Full JAE · seed 2 jae_full_seed2	0.758	0.653	0.874	0.855	0.810	0.953	0.785	0.882	0.817	77.7%
Upstream hyperparameters, batch=32 lr=1e-4 jae_full_faithful	0.769	0.624	0.657	0.927	0.846	0.951	0.744	0.898	0.796	82.9%

Scored on all 66,175 cells, including those whose labels the model trained on. Click “Held-out 10% only” to score only the cells it never saw.

What every variant shares. The data, the train/validation split (seed 42), the architecture, the optimiser, batch size, learning rate, the early-stopping rule, the UMAP parameters and seed, and the metric implementations and their sampling seed.
The one exception is jae_full_faithful, which uses the upstream batch size and learning rate (batch 32, learning rate 1e-4) and so changes two factors at once; it is not part of any experiment above.

What the analysis shows

Each finding is followed by the numbers it rests on.

The joint embedding keeps cell-type structure while mixing batches

Four of the six metrics reward keeping biological structure and two reward removing batch effects. The full JAE model (jae_full) scores:

NMI 0.777 · cell-type ASW 0.641 · cell-cycle conservation 0.878 · trajectory conservation 0.901 · batch ASW 0.838 · graph connectivity 0.932 · mean of six 0.828

This is one point on the trade-off between the two, not a maximum of both.

Keeping biology and removing batch effects pull against each other

The trade-off chart in the Evaluation section plots every trained variant by the mean of its four biological metrics against the mean of its two batch metrics.

jae_full — biological conservation 0.799 · batch removal 0.885 · 14 variants on the same axes

Experiment 1 in the lab moves along that trade-off by changing the batch-loss weight λ_batch.

The batch loss works by flattening, not by predicting

The loss asks the batch block for equal probability over every batch, rather than for the right batch. Experiment 3 in the lab changes exactly that one line.

batch term = uniform · loss weights 0.7 / 0.2 / 0.05 / 0.05 · latent block 45–54 of 61

It is a simpler alternative to adversarial batch correction: no second network is trained to detect the batch. The loss constrains only those 9 numbers; the Insight page tests how often a cell's nearest neighbours, over all 61 numbers, share its batch.

Changing only the random seed moves the mean of six by about 0.01

Experiment 6 trains the same configuration three times, changing only the random seed. That spread is the bar any difference on this page has to clear.

jae_full 0.828 · jae_full_seed1 0.821 · jae_full_seed2 0.817

Differences smaller than this are within seed-to-seed variation.

Method & reproducibility

How to reproduce the analysis

data: GEO GSE194122 / openproblems S3
SVD seed: 123, 100 components
train / val split: 90/10, seed 42
display sample: 5,004 cells, stratified, seed 7
cycle genes matched: 43 S + 54 G2M / 97
analysis saved: 2026-07-30 17:57:45

python3 tools/precompute_features.py   # normalise + SVD + cycle scores
python3 tools/train_jae.py             # all variants
python3 tools/precompute_metrics.py    # the six metrics
python3 tools/precompute_page.py       # pack the JSON payload
python3 tools/build.py                 # inject into templates -> index.html

TensorFlow, anndata and scanpy were not available, so the model is reimplemented layer for layer in PyTorch, the .h5ad files are read directly with h5py, and scanpy's cell-cycle scoring is rewritten in numpy. The whole pipeline runs on a CPU; no GPU is needed.

Glossary

Every term this page uses in a specific sense, with what it means here — on this dataset, in this experiment — rather than in general.

CITE-seq: a protocol that reads two signals from the same cell at once: the mRNA it is transcribing, and the proteins on its surface. That simultaneity is the entire reason a "joint" embedding is possible.
GEX and ADT: as used here: GEX is the gene-expression half of each cell's measurement (thousands of genes, very sparse); ADT is the antibody-derived-tag half (134 surface proteins, dense and individually interpretable). They take different routes through Stage 2 for exactly that reason.
joint embedding: as used here: one list of numbers per cell (here 61) that combines both measurements, rather than one list per measurement. The task is to produce a representation, not a prediction.
the design matrix X, and the labels y: as used here: one row of X is one cell; one column is one gene or one protein. Cell type, batch and cell-cycle phase are labels, not columns — they never enter the input and reach the model only through the loss function in Stage 3.
TruncatedSVD: a linear dimensionality reduction that does not centre the data first, which is what makes it usable on a sparse matrix. Here it takes GEX from thousands of genes down to 100 components; ADT is left alone.
transductive: as used here: the SVD is fitted on all the cells, including the ones later used for evaluation. Information from the evaluation cells therefore helps build the transformation. The competition permitted it.
autoencoder, and the latent vector: a network trained to reproduce its own input through a narrow middle layer. That middle layer's output — here 61 numbers per cell — is the embedding everything else on this page is computed from.
the latent layout: as used here: JAE's "classifier" is not a network, it is a slicing of those 61 numbers. Named ranges own named jobs — a cell-type block, a batch block, a cell-cycle pair, and a free remainder. The Transformation and modelling section draws the layout for the selected cell.
the uniform batch loss: as used here: the loss does not ask the model to get the batch wrong — that would still mean batch identity had been learned, merely inverted. It asks for equal probability over every batch, which removes batch information from those 9 numbers. It is a simpler alternative to adversarial batch correction: no second network is trained to detect the batch.
λ_batch: the weight on that batch term. Experiment 1 in the lab tries 0, 0.05, 0.2 and 0.5; the upstream method uses 0.05.
NMI: normalised mutual information: how well a clustering of the embedding agrees with the annotated cell types, from 0 (no agreement) to 1 (perfect agreement). Here the clustering is KMeans, run over a range of k, keeping the best NMI.
ASW (average silhouette width): as used here, two different things under one name: cell-type ASW rewards cells of one type sitting together, and batch ASW rewards batches being mixed within a cell type. One wants structure, the other wants structure destroyed — which is why they pull against each other.
graph connectivity: whether the cells of one type form a single connected neighbourhood rather than several islands. It penalises an embedding that shatters a cell type across the picture.
cell-cycle conservation: how much of the cell-cycle signal that was present before the embedding survives it. It is a conservation score: higher means less of that biology was scrubbed away.
trajectory conservation: whether an ordering that exists in the original data — here the erythroid maturation pseudotime — is still recoverable from the embedding.
the six metrics, taken together: four of them score biological conservation and two score batch removal. The trade-off chart in the Evaluation section plots the mean of the first four against the mean of the last two: top right is best, and no variant on this page reaches it.
held-out 10%: as used here: a random 10% of cells, at a recorded split seed — not a held-out donor. Every batch appears on both sides of that split.

Source evidencedata source · recorded settings · how the metrics are computed

Notebook lineage

What the saved analysis records

Design matrixwhat one row, column and cell mean

What a row and a column are here

How a row is built

One sample = one bone-marrow cell, measured two ways at once

Input X

Labels y (not inputs)

Task type

The design-matrix bridge: what a row, a column and a single entry each are

X: a display slice of 66,175 × 14,087

X[i,:] → one cell, in its native units

X[:,j] → one column across all cells

Transformation and modellingSVD · joint autoencoder · four weighted losses

The two modalities take different routes, then get concatenated

How much variance each of the 100 components explains

The 61 numbers are split into blocks, each with its own job

Four losses, with the weights from the upstream code

Training: cell-type accuracy rises, batch accuracy falls to chance

Evaluationmean of six: JAE 0.828 · plain autoencoder 0.680 · SVD 0.781

Scoring the embedding, then tracing a point back to a real cell

Six metrics for the full JAE embedding

Biological conservation ↔ batch removal is a trade-off

The embedding: 61 numbers drawn as a 2-D map (UMAP)

Traceability: which cell is this point?

Lab: what each part of training changes

Lab · Batch-loss strengthλ = 0 / 0.05 / 0.2 / 0.5 · no monotonic winner

Raising the batch-loss weight from 0 to 0.5: is biological structure removed too?

λ_batch = 0

λ_batch = 0.05

λ_batch = 0.2

λ_batch = 0.5

Lab · Adding supervisionSVD → autoencoder → cell-type labels → full JAE

Which loss is the latent structure actually coming from?

① No training at all

② Plain autoencoder

③ + cell type

④ Full JAE

Lab · Uniform versus correct batch targetsbatch mixing: 0.597 with uniform target · 0.196 with classification

The key control: swap "uniform distribution" for "classify correctly"

Push toward uniform (what JAE does)

Classify the batch correctly (control)

Lab · One modality or bothmean of six: joint 0.828 · genes 0.805 · proteins 0.755

What does the joint embedding actually buy over a single modality?

GEX only

ADT only

GEX + ADT jointly

Lab · Output activationsoriginal ReLU outputs versus linear outputs

What do those two ReLUs cost?

As upstream (ReLU output)

Both ReLUs removed

Lab · Random seedsmean-of-six spread: 0.0105 across three seeds

How much does the score move if you only change the random seed?

Seed 0

Seed 1

Seed 2

Lab · All comparison results14 variants · six metrics · all-cell and held-out scopes

14 variants × six metrics

Conclusionswhat the analysis shows

What the analysis shows

The joint embedding keeps cell-type structure while mixing batches

Keeping biology and removing batch effects pull against each other

The batch loss works by flattening, not by predicting

Changing only the random seed moves the mean of six by about 0.01

How to reproduce the analysis

Glossaryterms as used in this workflow

Glossary