Transformer encoders and a fused multimodal transformer#
Experimental. Transformer encoders were not used for the results in the paper. This notebook compares them with the standard MLP encoders on identical inputs and reports whatever comes out; it does not assume that attention helps.
UniVI can encode a modality with a transformer over tokens, where each token is one of the cell’s most prominent features (a highly expressed gene, an accessible peak) carrying its value, its rank, and a learned embedding of which feature it is. It can also encode all modalities jointly with one fused transformer over the concatenated tokens, so that gene tokens and peak tokens attend to each other before the latent posterior is formed.
On 10x Multiome PBMCs this notebook:
builds token-friendly inputs (z-scored RNA, TF-IDF-weighted peaks);
trains five models on the same cells and features: MLP (paper objective,
v1), per-modality transformers (v1), per-modality transformers whose ATAC tokens also know their genomic position (v1), MLP (v2), and MLP experts plus a fused transformer (v2);compares alignment, label transfer, ATAC → RNA prediction, fused-embedding quality and training time;
opens up the fused transformer: a single cell’s attention map, how sharp attention is in each layer, how much of it crosses modalities (overall and per cell type), which tokens the pooled embedding depends on (attention rollout) and whether that adds anything to simply being a token, attention sinks, pair-specific gene–peak links, and whether a gene token attends more to cis peaks (near its TSS) than to other peaks in the same cell.
The last test is a falsifiable check: attention weights are not explanations, but if the fused encoder had learned nothing about genomic proximity, cis peaks would receive no more of a gene’s attention than other peaks.
import sys
if "google.colab" in sys.modules:
%pip install -q "univi[tutorials]>=1.2" "pandas==2.2.3"
import time
from pathlib import Path
import anndata as ad
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import scanpy as sc
import scipy.sparse as sp
import torch
from IPython.display import display
from sklearn.model_selection import cross_val_score
from sklearn.neighbors import KNeighborsClassifier
import univi
import univi.datasets as uds
from univi import ModalityConfig, TrainingConfig, UniVIConfig, UniVIMultiModalVAE, UniVITrainer
from univi.config import TokenizerConfig, TransformerConfig
from univi.evaluation import (cross_modal_predict, encode_adata, encode_fused_adata_pair, evaluate_alignment,
pearson_corr_per_feature)
from univi.interpretability import fused_encode_with_meta_and_attn, top_cross_modal_feature_pairs_from_attn
from univi.preprocessing import RNAPreprocessor, split_by_label
from univi.utils.seed import set_seed
from univi.workflows import load_reference, make_loader, save_reference
device = "cuda" if torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu")
set_seed(0)
dense = lambda x: x.toarray() if sp.issparse(x) else np.asarray(x)
print(f"UniVI {univi.__version__} on {device}")
UniVI 1.2.1 on cuda
C:\Users\ashfo\micromamba\envs\univi-release-050\Lib\site-packages\tqdm\auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
from .autonotebook import tqdm as notebook_tqdm
N_EPOCHS = 300
BATCH_SIZE = 256
N_HVG = 2000
N_PEAKS = 20000 # most frequently open peaks in training cells
ATAC_RECON_WEIGHT = 0.2 # balances the ATAC reconstruction sum against RNA's; tune if needed
N_TOKENS_RNA = 128 # genes per cell given to the transformer (the cell's highest z-scores)
N_TOKENS_ATAC = 128 # peaks per cell (the cell's highest TF-IDF values)
D_TOKEN = 64 # token embedding size before the transformer
TRANSFORMER = dict(d_model=128, num_heads=4, num_layers=2, dim_feedforward=256, dropout=0.1, attn_dropout=0.1)
CIS_WINDOW_BP = 100_000 # gene-peak pairs closer than this (same chromosome) count as cis
N_PERMUTATIONS = 200 # peak-coordinate shuffles for the cis-attention null
N_UNTRAINED_CONTROLS = 5 # untrained copies of the fused model, as a reference for the cis test
MAX_ATTENTION_CELLS = 2000
RUN_POSITION_MODEL = True # also train the transformer whose ATAC tokens carry genomic coordinates
N_DE_GENES = 50 # marker genes per cell type for the attention-rollout test
ATAC_FEATURES = "near_genes" # "near_genes": peaks within PEAK_WINDOW_BP of a selected gene; "most_accessible"
PEAK_WINDOW_BP = 100_000 # window around each selected gene's TSS for "near_genes"
N_PSEUDOBULKS = 200 # pseudobulks (clusters of training cells) for correlation-based peak-gene links
N_BACKGROUND = 400 # background peaks per accessibility decile for the correlation null
LINK_Z = 3.0 # correlation z-score at which a peak-gene pair counts as a data-derived link
MIN_COPRESENT = 5 # cells in which both tokens must be present for an attention pair score
ATTRIBUTION_MODEL = "MLP (v1)" # model whose gradients give the attribution scores
N_ATTRIBUTION_CELLS = 500 # test cells used for gradient attribution
N_PSEUDOBULKS_LINEAGE = 60 # pseudobulks within each lineage for within-lineage reference links
# cell types grouped into lineages for within-lineage links (labels not listed are left out)
LINEAGES = {"T/NK": ["CD4 Naive", "CD4 TCM", "CD4 TEM", "CD8 Naive", "CD8 TEM_1", "CD8 TEM_2", "Treg", "MAIT",
"gdT", "NK"],
"B": ["Naive B", "Intermediate B", "Memory B", "Plasma"],
"myeloid": ["CD14 Mono", "CD16 Mono", "cDC", "pDC"]}
TOP_PAIRS_PER_CELL = 20 # strongest (centered) gene -> peak pairs taken from each cell
MIN_COOCCURRENCE = 20 # a pair must have both tokens present in at least this many cells
OUTPUT_DIR = "univi_outputs"
1. Data and token-friendly inputs#
A top-k tokenizer keeps each cell’s largest values. For RNA we z-score log-normalized expression,
so the kept genes are those most above their average level in that cell. For ATAC, binarized peaks
would tie (every open peak equals 1), so we use Signac-style TF-IDF, log1p(TF × IDF × 1e4) with the
IDF fit on training cells: the kept peaks are the cell’s open peaks that are rarest across cells,
which is where cell identity lives. All models receive exactly these inputs.
Which peaks? With ATAC_FEATURES = "near_genes" (the default), only peaks within PEAK_WINDOW_BP of
a selected gene’s TSS are kept (the most accessible N_PEAKS of them). Every peak is then a candidate
regulatory partner of at least one modeled gene, which is what the cis-attention test and the peak–gene
benchmark in Section 5 need: in a run with the N_PEAKS most accessible peaks genome-wide, only about
four gene tokens per cell had any nearby peak among the peak tokens. The cost is that peaks far from every
selected gene, some of which mark cell identity, are left out; the benchmark in Section 3 shows whether
that matters. "most_accessible" restores the genome-wide selection.
data = uds.pbmc_multiome_10k()
rna, atac = data["rna"], data["atac"]
splits = split_by_label(rna.obs["cell_type"], train_fraction=0.8, val_fraction=0.1, seed=0)
print({k: len(v) for k, v in splits.items()})
def tss_table(var):
need = ["chrom", "chromStart", "chromEnd", "strand"]
if not set(need) <= set(var.columns):
return None
v = var[need].dropna()
v = v[v["strand"].astype(str).isin(["+", "-"])]
tss = np.where(v["strand"].astype(str) == "-", v["chromEnd"], v["chromStart"]).astype(np.int64)
return pd.DataFrame({"chrom": v["chrom"].astype(str).to_numpy(), "tss": tss}, index=v.index)
def peaks_near(peak_var, gene_tss, window):
"""Boolean mask of peaks whose midpoint lies within `window` bp of any TSS on the same chromosome."""
chrom = peak_var["chrom"].astype(str).to_numpy()
mid = ((peak_var["chromStart"].astype(np.int64) + peak_var["chromEnd"].astype(np.int64)) // 2).to_numpy()
out = np.zeros(len(mid), dtype=bool)
for c, g in gene_tss.groupby("chrom"):
t = np.sort(g["tss"].to_numpy())
sel = np.flatnonzero(chrom == c)
i = np.clip(np.searchsorted(t, mid[sel]), 1, len(t)) - 1
nearest = np.minimum(np.abs(mid[sel] - t[i]), np.abs(mid[sel] - t[np.minimum(i + 1, len(t) - 1)]))
out[sel] = nearest <= window
return out
rna_prep = RNAPreprocessor(n_hvg=N_HVG, scale=True).fit(rna[splits["train"]])
train_counts = sp.csr_matrix(atac[splits["train"]].layers["counts"])
freq = np.asarray((train_counts > 0).mean(axis=0)).ravel()
tss = tss_table(rna.var)
has_peak_coords = {"chrom", "chromStart", "chromEnd"} <= set(atac.var.columns)
if ATAC_FEATURES == "near_genes" and tss is not None and has_peak_coords:
genes_sel = [g for g in rna_prep.features_ if g in tss.index]
near = peaks_near(atac.var, tss.loc[genes_sel], PEAK_WINDOW_BP)
cand = np.flatnonzero(near & (freq > 0))
keep = cand[np.argsort(-freq[cand], kind="stable")][:N_PEAKS]
print(f"near_genes: {len(cand):,} open peaks lie within {PEAK_WINDOW_BP // 1000} kb of "
f"{len(genes_sel):,} selected genes; keeping {len(keep):,}")
else:
keep = np.argsort(-freq, kind="stable")[: min(N_PEAKS, int((freq > 0).sum()))]
print(f"most_accessible: keeping {len(keep):,} peaks")
df = np.asarray((train_counts[:, keep] > 0).sum(axis=0)).ravel()
idf = np.log1p(train_counts.shape[0] / (1.0 + df)).astype(np.float32)
def tfidf(adata_counts):
x = sp.csr_matrix(adata_counts.layers["counts"][:, keep], dtype=np.float32)
x.data = np.ones_like(x.data) # binarize, then TF-IDF
tf = sp.diags(1.0 / np.maximum(np.asarray(x.sum(axis=1)).ravel(), 1.0)) @ x
out = sp.csr_matrix(tf.multiply(idf[None, :]))
out.data = np.log1p(out.data * 1e4).astype(np.float32)
return ad.AnnData(out, obs=adata_counts.obs.copy(), var=atac.var.iloc[keep].copy())
parts = {k: {"rna": rna_prep.transform(rna[i]), "atac": tfidf(atac[i])} for k, i in splits.items()}
train, val, test = parts["train"], parts["val"], parts["test"]
print("RNA", train["rna"].shape, "| ATAC", train["atac"].shape)
{'train': 7698, 'val': 956, 'test': 977}
near_genes: 22,829 open peaks lie within 100 kb of 2,000 selected genes; keeping 20,000
RNA (7698, 2000) | ATAC (7698, 20000)
2. Five models on the same inputs#
A top-k token says “this feature has this value and this rank in this cell”. Without an identity
embedding the transformer would not know which gene or peak a token is, so the tokenizers below switch
on a learned feature-ID embedding (use_feature_embedding=True). ATAC tokens of the position model
additionally get a genomic coordinate embedding (use_coord_embedding=True: a chromosome embedding
plus multi-scale sinusoidal features of the peak midpoint, from 1 kb to 100 Mb), and its transformer adds
a learned attention bias for binned genomic distance between peak tokens (use_relpos_bias=True; peaks
on different chromosomes share the farthest bin). These options work since UniVI 1.2.0.
The fused transformer is only trained by the v2 objective (in v1 its posterior never enters the
reconstruction loss), so an MLP model with v2 is included to separate the effect of the fused encoder
from the effect of the objective. recon_normalize_by_dim=False keeps v2 reconstruction sums on the
same scale as v1, so the same beta and gamma are used throughout.
A note on fused encoders. When a fused encoder produces the posterior, the per-modality encoders take no part in reconstruction; the alignment term (each modality’s posterior mean pulled toward the fused one) is their only training signal, and with the paper’s schedule it is annealed in from epoch 75. Before UniVI 1.2.0 the zero-weighted term stayed in the loss, the encoders received zero gradients, and Adam’s weight decay drove their weights to zero before alignment began: in a run of this notebook with UniVI 1.1 every per-modality embedding of the fused model was the same point (FOSCTTM 0.000 from tied distances, Recall@10 0.01, label transfer at chance). UniVI 1.2.0 leaves a zero-weighted term out of the loss, so all five models here use the paper’s schedule. The check after training reports any per-modality encoder whose output barely varies across cells.
atac_chroms = sorted(set(train["atac"].var["chrom"].astype(str))) if "chrom" in train["atac"].var else []
def tokenizer(n_features, n_tokens, position=False):
return TokenizerConfig(mode="topk_channels", n_tokens=n_tokens, channels=("value", "rank"),
n_features=int(n_features), use_feature_embedding=True, feature_emb_dim=D_TOKEN,
use_coord_embedding=position, n_chroms=len(atac_chroms) if position else None)
def modalities(encoder):
n_rna, n_atac = train["rna"].n_vars, train["atac"].n_vars
position = encoder == "transformer+position"
if encoder == "mlp": # MLP experts; the tokenizers are used only by a fused encoder
kw_rna, kw_atac = dict(tokenizer=tokenizer(n_rna, N_TOKENS_RNA)), dict(tokenizer=tokenizer(n_atac, N_TOKENS_ATAC))
else:
tcfg_atac = TransformerConfig(**TRANSFORMER, use_relpos_bias=position, relpos_num_bins=32, relpos_max_dist=1e6)
kw_rna = dict(encoder_type="transformer", transformer=TransformerConfig(**TRANSFORMER),
tokenizer=tokenizer(n_rna, N_TOKENS_RNA))
kw_atac = dict(encoder_type="transformer", transformer=tcfg_atac,
tokenizer=tokenizer(n_atac, N_TOKENS_ATAC, position=position))
return [ModalityConfig("rna", n_rna, [512, 256, 128], [128, 256, 512], likelihood="gaussian", **kw_rna),
ModalityConfig("atac", n_atac, [512, 256, 128], [128, 256, 512], likelihood="gaussian",
recon_weight=ATAC_RECON_WEIGHT, **kw_atac)]
def attach_peak_coordinates(model):
"""Coordinates are data, not weights: attach them after building (or loading) a position model."""
v = getattr(model.encoders["atac"], "vec2tok", None)
if v is None or not v.use_coord_emb:
return
var = train["atac"].var
chrom = np.array(var["chrom"].astype(str).map({c: i for i, c in enumerate(atac_chroms)}), dtype=np.int64)
v.set_feature_coords(torch.from_numpy(chrom), torch.from_numpy(np.array(var["chromStart"], dtype=np.float32)),
torch.from_numpy(np.array(var["chromEnd"], dtype=np.float32)))
common = dict(latent_dim=30, beta=1.25, gamma=4.35, encoder_dropout=0.10, decoder_dropout=0.05,
kl_anneal_start=50, kl_anneal_end=85, align_anneal_start=75, align_anneal_end=110)
specs = {
"MLP (v1)": (UniVIConfig(modalities=modalities("mlp"), **common), dict(loss_mode="v1", v1_recon="avg")),
"transformer (v1)": (UniVIConfig(modalities=modalities("transformer"), **common),
dict(loss_mode="v1", v1_recon="avg")),
"transformer + position (v1)": (UniVIConfig(modalities=modalities("transformer+position"), **common),
dict(loss_mode="v1", v1_recon="avg")),
"MLP (v2)": (UniVIConfig(modalities=modalities("mlp"), **common),
dict(loss_mode="v2", recon_normalize_by_dim=False)),
"fused transformer (v2)": (UniVIConfig(modalities=modalities("mlp"), fused_encoder_type="multimodal_transformer",
fused_transformer=TransformerConfig(**TRANSFORMER), **common),
dict(loss_mode="v2", recon_normalize_by_dim=False)),
}
if not RUN_POSITION_MODEL or not atac_chroms:
specs.pop("transformer + position (v1)")
models, histories, minutes = {}, {}, {}
for name, (cfg, kwargs) in specs.items():
set_seed(0)
model = UniVIMultiModalVAE(cfg, normalize_v1_terms=True, **kwargs)
attach_peak_coordinates(model)
n_params = sum(p.numel() for p in model.parameters())
trainer = UniVITrainer(
model, make_loader(train, batch_size=BATCH_SIZE, shuffle=True, drop_last=True),
make_loader(val, batch_size=512),
TrainingConfig(n_epochs=N_EPOCHS, batch_size=BATCH_SIZE, lr=1e-3, weight_decay=1e-4, device=device,
early_stopping=True, patience=50, best_epoch_warmup=110, log_every=100, grad_clip=5.0))
started = time.time()
histories[name] = trainer.fit()
minutes[name] = (time.time() - started) / 60
models[name] = model
print(f"{name}: {n_params / 1e6:.1f}M parameters, best epoch {trainer.best_epoch}, {minutes[name]:.1f} min")
fused_tok = models["fused transformer (v2)"].fused_encoder.vec2tok
assert all(fused_tok[m].id_emb is not None for m in fused_tok), "feature-ID embeddings missing in fused encoder"
spread = pd.DataFrame({name: {m: float(encode_adata(model, val[m], modality=m, device=device,
latent="modality_mean").std(0).mean()) for m in ("rna", "atac")}
for name, model in models.items()}).T
spread.columns = [f"{m} embedding s.d. across cells" for m in spread.columns]
display(spread.round(4))
collapsed = spread[(spread < 1e-3).any(axis=1)].index.tolist()
if collapsed:
print("WARNING: per-modality encoders with (near-)constant output:", collapsed,
"- their alignment, label-transfer and cross-modal metrics are not meaningful.")
MLP (v1): 23.2M parameters, best epoch 129, 6.5 min
transformer (v1): 13.6M parameters, best epoch 112, 9.4 min
transformer + position (v1): 13.7M parameters, best epoch 135, 12.1 min
MLP (v2): 23.2M parameters, best epoch 293, 8.7 min
fused transformer (v2): 25.0M parameters, best epoch 120, 11.3 min
[2026-09-23 19:31:23,644] [UniVITrainer] [INFO] TrainingConfig:
[2026-09-23 19:31:23,645] [UniVITrainer] [INFO] n_epochs: 300
[2026-09-23 19:31:23,646] [UniVITrainer] [INFO] batch_size: 256
[2026-09-23 19:31:23,646] [UniVITrainer] [INFO] lr: 0.001
[2026-09-23 19:31:23,646] [UniVITrainer] [INFO] weight_decay: 0.0001
[2026-09-23 19:31:23,647] [UniVITrainer] [INFO] device: 'cuda'
[2026-09-23 19:31:23,647] [UniVITrainer] [INFO] log_every: 100
[2026-09-23 19:31:23,647] [UniVITrainer] [INFO] grad_clip: 5.0
[2026-09-23 19:31:23,648] [UniVITrainer] [INFO] num_workers: 0
[2026-09-23 19:31:23,648] [UniVITrainer] [INFO] seed: 0
[2026-09-23 19:31:23,648] [UniVITrainer] [INFO] early_stopping: True
[2026-09-23 19:31:23,649] [UniVITrainer] [INFO] patience: 50
[2026-09-23 19:31:23,649] [UniVITrainer] [INFO] min_delta: 0.0
[2026-09-23 19:31:23,649] [UniVITrainer] [INFO] best_epoch_warmup: 110
Training UniVI: 0%| | 0/300 [00:00<?, ?it/s][2026-09-23 19:31:26,702] [UniVITrainer] [INFO] [Epoch 001] Train loss=1795.5100 (beta=0.000, gamma=0.000)
[2026-09-23 19:31:26,906] [UniVITrainer] [INFO] [Epoch 001] Val loss=1865.1799 (beta=1.250, gamma=4.350)
Training UniVI: 33%|███▎ | 99/300 [03:40<07:03, 2.11s/it, beta=1.250, gamma=2.983, train_loss=1547.6104, val_loss=1660.9100][2026-09-23 19:35:05,630] [UniVITrainer] [INFO] [Epoch 100] Train loss=1547.3652 (beta=1.250, gamma=3.107)
[2026-09-23 19:35:05,816] [UniVITrainer] [INFO] [Epoch 100] Val loss=1661.4268 (beta=1.250, gamma=4.350)
Training UniVI: 36%|███▌ | 108/300 [04:01<06:48, 2.13s/it, beta=1.250, gamma=4.226, train_loss=1544.2032, val_loss=1657.5564][2026-09-23 19:35:25,017] [UniVITrainer] [INFO] [Epoch 109] Best tracking warmup ends next epoch (best_epoch_warmup=110).
Training UniVI: 36%|███▋ | 109/300 [04:03<06:49, 2.14s/it, beta=1.250, gamma=4.350, train_loss=1543.6702, val_loss=1657.1236][2026-09-23 19:35:27,163] [UniVITrainer] [INFO] [Epoch 110] New best val loss: 1657.1236
Training UniVI: 37%|███▋ | 110/300 [04:05<06:47, 2.14s/it, beta=1.250, gamma=4.350, train_loss=1544.6350, val_loss=1655.8083][2026-09-23 19:35:29,240] [UniVITrainer] [INFO] [Epoch 111] New best val loss: 1655.8083
Training UniVI: 37%|███▋ | 111/300 [04:07<06:41, 2.12s/it, beta=1.250, gamma=4.350, train_loss=1543.8619, val_loss=1655.0349][2026-09-23 19:35:31,421] [UniVITrainer] [INFO] [Epoch 112] New best val loss: 1655.0349
Training UniVI: 39%|███▉ | 118/300 [04:22<06:16, 2.07s/it, beta=1.250, gamma=4.350, train_loss=1538.7507, val_loss=1653.7773][2026-09-23 19:35:46,052] [UniVITrainer] [INFO] [Epoch 119] New best val loss: 1653.7773
Training UniVI: 43%|████▎ | 128/300 [04:43<05:58, 2.08s/it, beta=1.250, gamma=4.350, train_loss=1536.6810, val_loss=1653.0760][2026-09-23 19:36:06,862] [UniVITrainer] [INFO] [Epoch 129] New best val loss: 1653.0760
Training UniVI: 59%|█████▉ | 178/300 [06:28<04:25, 2.18s/it, beta=1.250, gamma=4.350, train_loss=1516.0855, val_loss=1656.7333][2026-09-23 19:37:51,990] [UniVITrainer] [INFO] Early stopping at epoch 179 (best val loss=1653.0760, best epoch=129)
Training UniVI: 59%|█████▉ | 178/300 [06:28<04:26, 2.18s/it, beta=1.250, gamma=4.350, train_loss=1516.0855, val_loss=1656.7333]
[2026-09-23 19:37:52,022] [UniVITrainer] [INFO] Restored best model from epoch 129 (val loss=1653.0760)
[2026-09-23 19:37:52,154] [UniVITrainer] [INFO] TrainingConfig:
[2026-09-23 19:37:52,155] [UniVITrainer] [INFO] n_epochs: 300
[2026-09-23 19:37:52,156] [UniVITrainer] [INFO] batch_size: 256
[2026-09-23 19:37:52,157] [UniVITrainer] [INFO] lr: 0.001
[2026-09-23 19:37:52,157] [UniVITrainer] [INFO] weight_decay: 0.0001
[2026-09-23 19:37:52,157] [UniVITrainer] [INFO] device: 'cuda'
[2026-09-23 19:37:52,158] [UniVITrainer] [INFO] log_every: 100
[2026-09-23 19:37:52,158] [UniVITrainer] [INFO] grad_clip: 5.0
[2026-09-23 19:37:52,158] [UniVITrainer] [INFO] num_workers: 0
[2026-09-23 19:37:52,159] [UniVITrainer] [INFO] seed: 0
[2026-09-23 19:37:52,159] [UniVITrainer] [INFO] early_stopping: True
[2026-09-23 19:37:52,160] [UniVITrainer] [INFO] patience: 50
[2026-09-23 19:37:52,160] [UniVITrainer] [INFO] min_delta: 0.0
[2026-09-23 19:37:52,160] [UniVITrainer] [INFO] best_epoch_warmup: 110
Training UniVI: 0%| | 0/300 [00:00<?, ?it/s][2026-09-23 19:37:55,945] [UniVITrainer] [INFO] [Epoch 001] Train loss=1827.3201 (beta=0.000, gamma=0.000)
[2026-09-23 19:37:56,176] [UniVITrainer] [INFO] [Epoch 001] Val loss=1966.5991 (beta=1.250, gamma=4.350)
Training UniVI: 33%|███▎ | 99/300 [05:49<11:09, 3.33s/it, beta=1.250, gamma=2.983, train_loss=1588.2788, val_loss=1668.5833][2026-09-23 19:43:45,172] [UniVITrainer] [INFO] [Epoch 100] Train loss=1588.3190 (beta=1.250, gamma=3.107)
[2026-09-23 19:43:45,389] [UniVITrainer] [INFO] [Epoch 100] Val loss=1660.9097 (beta=1.250, gamma=4.350)
Training UniVI: 36%|███▌ | 108/300 [06:22<10:32, 3.29s/it, beta=1.250, gamma=4.226, train_loss=1585.8687, val_loss=1664.3285][2026-09-23 19:44:15,113] [UniVITrainer] [INFO] [Epoch 109] Best tracking warmup ends next epoch (best_epoch_warmup=110).
Training UniVI: 36%|███▋ | 109/300 [06:26<10:32, 3.31s/it, beta=1.250, gamma=4.350, train_loss=1584.5907, val_loss=1663.9755][2026-09-23 19:44:18,394] [UniVITrainer] [INFO] [Epoch 110] New best val loss: 1663.9755
Training UniVI: 37%|███▋ | 110/300 [06:29<10:27, 3.30s/it, beta=1.250, gamma=4.350, train_loss=1586.7886, val_loss=1660.9883][2026-09-23 19:44:21,745] [UniVITrainer] [INFO] [Epoch 111] New best val loss: 1660.9883
Training UniVI: 37%|███▋ | 111/300 [06:32<10:26, 3.32s/it, beta=1.250, gamma=4.350, train_loss=1584.4013, val_loss=1659.4302][2026-09-23 19:44:25,137] [UniVITrainer] [INFO] [Epoch 112] New best val loss: 1659.4302
Training UniVI: 54%|█████▎ | 161/300 [09:25<08:10, 3.53s/it, beta=1.250, gamma=4.350, train_loss=1573.8785, val_loss=1668.8779][2026-09-23 19:47:17,458] [UniVITrainer] [INFO] Early stopping at epoch 162 (best val loss=1659.4302, best epoch=112)
Training UniVI: 54%|█████▎ | 161/300 [09:25<08:08, 3.51s/it, beta=1.250, gamma=4.350, train_loss=1573.8785, val_loss=1668.8779]
[2026-09-23 19:47:17,475] [UniVITrainer] [INFO] Restored best model from epoch 112 (val loss=1659.4302)
[2026-09-23 19:47:17,562] [UniVITrainer] [INFO] TrainingConfig:
[2026-09-23 19:47:17,563] [UniVITrainer] [INFO] n_epochs: 300
[2026-09-23 19:47:17,563] [UniVITrainer] [INFO] batch_size: 256
[2026-09-23 19:47:17,563] [UniVITrainer] [INFO] lr: 0.001
[2026-09-23 19:47:17,564] [UniVITrainer] [INFO] weight_decay: 0.0001
[2026-09-23 19:47:17,564] [UniVITrainer] [INFO] device: 'cuda'
[2026-09-23 19:47:17,564] [UniVITrainer] [INFO] log_every: 100
[2026-09-23 19:47:17,565] [UniVITrainer] [INFO] grad_clip: 5.0
[2026-09-23 19:47:17,565] [UniVITrainer] [INFO] num_workers: 0
[2026-09-23 19:47:17,565] [UniVITrainer] [INFO] seed: 0
[2026-09-23 19:47:17,566] [UniVITrainer] [INFO] early_stopping: True
[2026-09-23 19:47:17,566] [UniVITrainer] [INFO] patience: 50
[2026-09-23 19:47:17,566] [UniVITrainer] [INFO] min_delta: 0.0
[2026-09-23 19:47:17,567] [UniVITrainer] [INFO] best_epoch_warmup: 110
Training UniVI: 0%| | 0/300 [00:00<?, ?it/s][2026-09-23 19:47:21,102] [UniVITrainer] [INFO] [Epoch 001] Train loss=1826.1058 (beta=0.000, gamma=0.000)
[2026-09-23 19:47:21,347] [UniVITrainer] [INFO] [Epoch 001] Val loss=2039.6260 (beta=1.250, gamma=4.350)
Training UniVI: 33%|███▎ | 99/300 [06:19<13:06, 3.91s/it, beta=1.250, gamma=2.983, train_loss=1591.4143, val_loss=1653.2709][2026-09-23 19:53:40,826] [UniVITrainer] [INFO] [Epoch 100] Train loss=1591.6047 (beta=1.250, gamma=3.107)
[2026-09-23 19:53:41,058] [UniVITrainer] [INFO] [Epoch 100] Val loss=1653.8799 (beta=1.250, gamma=4.350)
Training UniVI: 36%|███▌ | 108/300 [06:57<12:08, 3.79s/it, beta=1.250, gamma=4.226, train_loss=1590.4278, val_loss=1653.6150][2026-09-23 19:54:15,394] [UniVITrainer] [INFO] [Epoch 109] Best tracking warmup ends next epoch (best_epoch_warmup=110).
Training UniVI: 36%|███▋ | 109/300 [07:01<12:05, 3.80s/it, beta=1.250, gamma=4.350, train_loss=1588.6068, val_loss=1650.7373][2026-09-23 19:54:19,251] [UniVITrainer] [INFO] [Epoch 110] New best val loss: 1650.7373
Training UniVI: 45%|████▍ | 134/300 [08:37<10:46, 3.90s/it, beta=1.250, gamma=4.350, train_loss=1584.1715, val_loss=1650.5452][2026-09-23 19:55:55,260] [UniVITrainer] [INFO] [Epoch 135] New best val loss: 1650.5452
Training UniVI: 61%|██████▏ | 184/300 [12:06<08:04, 4.18s/it, beta=1.250, gamma=4.350, train_loss=1575.0167, val_loss=1659.3929][2026-09-23 19:59:23,772] [UniVITrainer] [INFO] Early stopping at epoch 185 (best val loss=1650.5452, best epoch=135)
Training UniVI: 61%|██████▏ | 184/300 [12:06<07:37, 3.95s/it, beta=1.250, gamma=4.350, train_loss=1575.0167, val_loss=1659.3929]
[2026-09-23 19:59:23,793] [UniVITrainer] [INFO] Restored best model from epoch 135 (val loss=1650.5452)
[2026-09-23 19:59:23,914] [UniVITrainer] [INFO] TrainingConfig:
[2026-09-23 19:59:23,915] [UniVITrainer] [INFO] n_epochs: 300
[2026-09-23 19:59:23,916] [UniVITrainer] [INFO] batch_size: 256
[2026-09-23 19:59:23,917] [UniVITrainer] [INFO] lr: 0.001
[2026-09-23 19:59:23,917] [UniVITrainer] [INFO] weight_decay: 0.0001
[2026-09-23 19:59:23,918] [UniVITrainer] [INFO] device: 'cuda'
[2026-09-23 19:59:23,919] [UniVITrainer] [INFO] log_every: 100
[2026-09-23 19:59:23,919] [UniVITrainer] [INFO] grad_clip: 5.0
[2026-09-23 19:59:23,919] [UniVITrainer] [INFO] num_workers: 0
[2026-09-23 19:59:23,920] [UniVITrainer] [INFO] seed: 0
[2026-09-23 19:59:23,920] [UniVITrainer] [INFO] early_stopping: True
[2026-09-23 19:59:23,920] [UniVITrainer] [INFO] patience: 50
[2026-09-23 19:59:23,921] [UniVITrainer] [INFO] min_delta: 0.0
[2026-09-23 19:59:23,921] [UniVITrainer] [INFO] best_epoch_warmup: 110
Training UniVI: 0%| | 0/300 [00:00<?, ?it/s][2026-09-23 19:59:25,479] [UniVITrainer] [INFO] [Epoch 001] Train loss=3570.9123 (beta=0.000, gamma=0.000)
[2026-09-23 19:59:25,628] [UniVITrainer] [INFO] [Epoch 001] Val loss=3602.2584 (beta=1.250, gamma=4.350)
Training UniVI: 33%|███▎ | 99/300 [03:00<06:18, 1.88s/it, beta=1.250, gamma=2.983, train_loss=3003.5118, val_loss=3264.1559][2026-09-23 20:02:26,254] [UniVITrainer] [INFO] [Epoch 100] Train loss=3001.6450 (beta=1.250, gamma=3.107)
[2026-09-23 20:02:26,431] [UniVITrainer] [INFO] [Epoch 100] Val loss=3261.3672 (beta=1.250, gamma=4.350)
Training UniVI: 36%|███▌ | 108/300 [03:19<06:06, 1.91s/it, beta=1.250, gamma=4.226, train_loss=2990.6066, val_loss=3258.0605][2026-09-23 20:02:43,430] [UniVITrainer] [INFO] [Epoch 109] Best tracking warmup ends next epoch (best_epoch_warmup=110).
Training UniVI: 36%|███▋ | 109/300 [03:21<06:01, 1.89s/it, beta=1.250, gamma=4.350, train_loss=2985.5543, val_loss=3256.3660][2026-09-23 20:02:45,389] [UniVITrainer] [INFO] [Epoch 110] New best val loss: 3256.3660
Training UniVI: 40%|████ | 120/300 [03:41<05:28, 1.83s/it, beta=1.250, gamma=4.350, train_loss=2970.8141, val_loss=3256.3530][2026-09-23 20:03:05,423] [UniVITrainer] [INFO] [Epoch 121] New best val loss: 3256.3530
Training UniVI: 41%|████ | 123/300 [03:46<05:23, 1.83s/it, beta=1.250, gamma=4.350, train_loss=2967.9897, val_loss=3255.6876][2026-09-23 20:03:10,899] [UniVITrainer] [INFO] [Epoch 124] New best val loss: 3255.6876
Training UniVI: 43%|████▎ | 128/300 [03:56<05:16, 1.84s/it, beta=1.250, gamma=4.350, train_loss=2964.3844, val_loss=3253.5619][2026-09-23 20:03:20,130] [UniVITrainer] [INFO] [Epoch 129] New best val loss: 3253.5619
Training UniVI: 48%|████▊ | 144/300 [04:23<04:24, 1.70s/it, beta=1.250, gamma=4.350, train_loss=2942.4376, val_loss=3253.2350][2026-09-23 20:03:47,661] [UniVITrainer] [INFO] [Epoch 145] New best val loss: 3253.2350
Training UniVI: 49%|████▉ | 148/300 [04:30<04:12, 1.66s/it, beta=1.250, gamma=4.350, train_loss=2941.2103, val_loss=3253.0055][2026-09-23 20:03:54,283] [UniVITrainer] [INFO] [Epoch 149] New best val loss: 3253.0055
Training UniVI: 50%|████▉ | 149/300 [04:32<04:10, 1.66s/it, beta=1.250, gamma=4.350, train_loss=2936.9656, val_loss=3250.7109][2026-09-23 20:03:56,015] [UniVITrainer] [INFO] [Epoch 150] New best val loss: 3250.7109
Training UniVI: 54%|█████▍ | 163/300 [04:55<03:50, 1.69s/it, beta=1.250, gamma=4.350, train_loss=2922.4480, val_loss=3250.1190][2026-09-23 20:04:19,127] [UniVITrainer] [INFO] [Epoch 164] New best val loss: 3250.1190
Training UniVI: 56%|█████▋ | 169/300 [05:05<03:36, 1.65s/it, beta=1.250, gamma=4.350, train_loss=2922.5593, val_loss=3249.9357][2026-09-23 20:04:28,955] [UniVITrainer] [INFO] [Epoch 170] New best val loss: 3249.9357
Training UniVI: 57%|█████▋ | 171/300 [05:08<03:32, 1.65s/it, beta=1.250, gamma=4.350, train_loss=2917.6276, val_loss=3248.7883][2026-09-23 20:04:32,322] [UniVITrainer] [INFO] [Epoch 172] New best val loss: 3248.7883
Training UniVI: 61%|██████▏ | 184/300 [05:29<03:13, 1.67s/it, beta=1.250, gamma=4.350, train_loss=2908.5117, val_loss=3245.7300][2026-09-23 20:04:53,883] [UniVITrainer] [INFO] [Epoch 185] New best val loss: 3245.7300
Training UniVI: 64%|██████▍ | 193/300 [05:45<03:03, 1.71s/it, beta=1.250, gamma=4.350, train_loss=2900.2144, val_loss=3245.2307][2026-09-23 20:05:09,393] [UniVITrainer] [INFO] [Epoch 194] New best val loss: 3245.2307
Training UniVI: 66%|██████▌ | 198/300 [05:53<02:53, 1.70s/it, beta=1.250, gamma=4.350, train_loss=2898.7581, val_loss=3244.8408][2026-09-23 20:05:17,844] [UniVITrainer] [INFO] [Epoch 199] New best val loss: 3244.8408
Training UniVI: 66%|██████▋ | 199/300 [05:53<02:53, 1.72s/it, beta=1.250, gamma=4.350, train_loss=2898.7581, val_loss=3244.8408][2026-09-23 20:05:19,293] [UniVITrainer] [INFO] [Epoch 200] Train loss=2900.2854 (beta=1.250, gamma=4.350)
[2026-09-23 20:05:19,445] [UniVITrainer] [INFO] [Epoch 200] Val loss=3243.2422 (beta=1.250, gamma=4.350)
Training UniVI: 66%|██████▋ | 199/300 [05:55<02:53, 1.72s/it, beta=1.250, gamma=4.350, train_loss=2900.2854, val_loss=3243.2422][2026-09-23 20:05:19,472] [UniVITrainer] [INFO] [Epoch 200] New best val loss: 3243.2422
Training UniVI: 68%|██████▊ | 203/300 [06:02<02:41, 1.66s/it, beta=1.250, gamma=4.350, train_loss=2895.5856, val_loss=3241.9056][2026-09-23 20:05:26,049] [UniVITrainer] [INFO] [Epoch 204] New best val loss: 3241.9056
Training UniVI: 70%|███████ | 211/300 [06:15<02:29, 1.68s/it, beta=1.250, gamma=4.350, train_loss=2892.6076, val_loss=3241.3938][2026-09-23 20:05:39,481] [UniVITrainer] [INFO] [Epoch 212] New best val loss: 3241.3938
Training UniVI: 74%|███████▎ | 221/300 [06:32<02:13, 1.69s/it, beta=1.250, gamma=4.350, train_loss=2884.9056, val_loss=3241.0387][2026-09-23 20:05:56,677] [UniVITrainer] [INFO] [Epoch 222] New best val loss: 3241.0387
Training UniVI: 75%|███████▍ | 224/300 [06:37<02:04, 1.64s/it, beta=1.250, gamma=4.350, train_loss=2884.7712, val_loss=3241.0040][2026-09-23 20:06:01,600] [UniVITrainer] [INFO] [Epoch 225] New best val loss: 3241.0040
Training UniVI: 77%|███████▋ | 230/300 [06:47<01:57, 1.68s/it, beta=1.250, gamma=4.350, train_loss=2879.6026, val_loss=3240.9216][2026-09-23 20:06:11,757] [UniVITrainer] [INFO] [Epoch 231] New best val loss: 3240.9216
Training UniVI: 77%|███████▋ | 231/300 [06:49<01:56, 1.68s/it, beta=1.250, gamma=4.350, train_loss=2878.2083, val_loss=3238.5565][2026-09-23 20:06:13,419] [UniVITrainer] [INFO] [Epoch 232] New best val loss: 3238.5565
Training UniVI: 79%|███████▉ | 237/300 [06:59<01:42, 1.63s/it, beta=1.250, gamma=4.350, train_loss=2875.0379, val_loss=3235.8761][2026-09-23 20:06:23,213] [UniVITrainer] [INFO] [Epoch 238] New best val loss: 3235.8761
Training UniVI: 80%|████████ | 240/300 [07:04<01:36, 1.60s/it, beta=1.250, gamma=4.350, train_loss=2873.3801, val_loss=3235.6627][2026-09-23 20:06:28,048] [UniVITrainer] [INFO] [Epoch 241] New best val loss: 3235.6627
Training UniVI: 81%|████████▏ | 244/300 [07:10<01:31, 1.64s/it, beta=1.250, gamma=4.350, train_loss=2873.8163, val_loss=3234.9037][2026-09-23 20:06:34,655] [UniVITrainer] [INFO] [Epoch 245] New best val loss: 3234.9037
Training UniVI: 82%|████████▏ | 246/300 [07:14<01:29, 1.66s/it, beta=1.250, gamma=4.350, train_loss=2874.6286, val_loss=3233.5850][2026-09-23 20:06:38,074] [UniVITrainer] [INFO] [Epoch 247] New best val loss: 3233.5850
Training UniVI: 84%|████████▍ | 253/300 [07:26<01:18, 1.68s/it, beta=1.250, gamma=4.350, train_loss=2867.6142, val_loss=3232.4353][2026-09-23 20:06:49,984] [UniVITrainer] [INFO] [Epoch 254] New best val loss: 3232.4353
Training UniVI: 85%|████████▍ | 254/300 [07:27<01:17, 1.69s/it, beta=1.250, gamma=4.350, train_loss=2866.5460, val_loss=3232.2365][2026-09-23 20:06:51,691] [UniVITrainer] [INFO] [Epoch 255] New best val loss: 3232.2365
Training UniVI: 87%|████████▋ | 261/300 [07:39<01:05, 1.68s/it, beta=1.250, gamma=4.350, train_loss=2866.1587, val_loss=3231.1862][2026-09-23 20:07:03,406] [UniVITrainer] [INFO] [Epoch 262] New best val loss: 3231.1862
Training UniVI: 91%|█████████ | 272/300 [07:57<00:46, 1.65s/it, beta=1.250, gamma=4.350, train_loss=2855.9465, val_loss=3229.7094][2026-09-23 20:07:21,654] [UniVITrainer] [INFO] [Epoch 273] New best val loss: 3229.7094
Training UniVI: 91%|█████████ | 273/300 [07:59<00:44, 1.65s/it, beta=1.250, gamma=4.350, train_loss=2860.9191, val_loss=3228.0837][2026-09-23 20:07:23,327] [UniVITrainer] [INFO] [Epoch 274] New best val loss: 3228.0837
Training UniVI: 95%|█████████▌| 285/300 [08:19<00:24, 1.65s/it, beta=1.250, gamma=4.350, train_loss=2852.0476, val_loss=3227.1428][2026-09-23 20:07:43,262] [UniVITrainer] [INFO] [Epoch 286] New best val loss: 3227.1428
Training UniVI: 96%|█████████▌| 288/300 [08:24<00:20, 1.70s/it, beta=1.250, gamma=4.350, train_loss=2853.4259, val_loss=3226.8480][2026-09-23 20:07:48,480] [UniVITrainer] [INFO] [Epoch 289] New best val loss: 3226.8480
Training UniVI: 96%|█████████▋| 289/300 [08:26<00:18, 1.70s/it, beta=1.250, gamma=4.350, train_loss=2848.7160, val_loss=3225.5531][2026-09-23 20:07:50,204] [UniVITrainer] [INFO] [Epoch 290] New best val loss: 3225.5531
Training UniVI: 97%|█████████▋| 292/300 [08:31<00:13, 1.66s/it, beta=1.250, gamma=4.350, train_loss=2849.6184, val_loss=3223.5901][2026-09-23 20:07:55,167] [UniVITrainer] [INFO] [Epoch 293] New best val loss: 3223.5901
Training UniVI: 100%|█████████▉| 299/300 [08:41<00:01, 1.70s/it, beta=1.250, gamma=4.350, train_loss=2847.2311, val_loss=3227.2478][2026-09-23 20:08:06,981] [UniVITrainer] [INFO] [Epoch 300] Train loss=2845.8535 (beta=1.250, gamma=4.350)
[2026-09-23 20:08:07,138] [UniVITrainer] [INFO] [Epoch 300] Val loss=3228.8680 (beta=1.250, gamma=4.350)
Training UniVI: 100%|██████████| 300/300 [08:43<00:00, 1.74s/it, beta=1.250, gamma=4.350, train_loss=2845.8535, val_loss=3228.8680]
[2026-09-23 20:08:07,158] [UniVITrainer] [INFO] Restored best model from epoch 293 (val loss=3223.5901)
[2026-09-23 20:08:07,300] [UniVITrainer] [INFO] TrainingConfig:
[2026-09-23 20:08:07,301] [UniVITrainer] [INFO] n_epochs: 300
[2026-09-23 20:08:07,302] [UniVITrainer] [INFO] batch_size: 256
[2026-09-23 20:08:07,302] [UniVITrainer] [INFO] lr: 0.001
[2026-09-23 20:08:07,302] [UniVITrainer] [INFO] weight_decay: 0.0001
[2026-09-23 20:08:07,303] [UniVITrainer] [INFO] device: 'cuda'
[2026-09-23 20:08:07,303] [UniVITrainer] [INFO] log_every: 100
[2026-09-23 20:08:07,303] [UniVITrainer] [INFO] grad_clip: 5.0
[2026-09-23 20:08:07,304] [UniVITrainer] [INFO] num_workers: 0
[2026-09-23 20:08:07,304] [UniVITrainer] [INFO] seed: 0
[2026-09-23 20:08:07,304] [UniVITrainer] [INFO] early_stopping: True
[2026-09-23 20:08:07,305] [UniVITrainer] [INFO] patience: 50
[2026-09-23 20:08:07,305] [UniVITrainer] [INFO] min_delta: 0.0
[2026-09-23 20:08:07,305] [UniVITrainer] [INFO] best_epoch_warmup: 110
Training UniVI: 0%| | 0/300 [00:00<?, ?it/s][2026-09-23 20:08:10,540] [UniVITrainer] [INFO] [Epoch 001] Train loss=3659.3444 (beta=0.000, gamma=0.000)
[2026-09-23 20:08:10,749] [UniVITrainer] [INFO] [Epoch 001] Val loss=3797.3232 (beta=1.250, gamma=4.350)
Training UniVI: 33%|███▎ | 99/300 [06:25<13:36, 4.06s/it, beta=1.250, gamma=2.983, train_loss=3131.9834, val_loss=3249.9659][2026-09-23 20:14:36,548] [UniVITrainer] [INFO] [Epoch 100] Train loss=3130.8394 (beta=1.250, gamma=3.107)
[2026-09-23 20:14:36,772] [UniVITrainer] [INFO] [Epoch 100] Val loss=3250.9677 (beta=1.250, gamma=4.350)
Training UniVI: 36%|███▌ | 108/300 [07:06<13:03, 4.08s/it, beta=1.250, gamma=4.226, train_loss=3123.3944, val_loss=3246.0032][2026-09-23 20:15:13,485] [UniVITrainer] [INFO] [Epoch 109] Best tracking warmup ends next epoch (best_epoch_warmup=110).
Training UniVI: 36%|███▋ | 109/300 [07:10<12:58, 4.08s/it, beta=1.250, gamma=4.350, train_loss=3120.7520, val_loss=3242.1705][2026-09-23 20:15:17,552] [UniVITrainer] [INFO] [Epoch 110] New best val loss: 3242.1705
Training UniVI: 37%|███▋ | 111/300 [07:18<12:47, 4.06s/it, beta=1.250, gamma=4.350, train_loss=3116.7022, val_loss=3240.7087][2026-09-23 20:15:25,591] [UniVITrainer] [INFO] [Epoch 112] New best val loss: 3240.7087
Training UniVI: 39%|███▊ | 116/300 [07:38<12:23, 4.04s/it, beta=1.250, gamma=4.350, train_loss=3109.0126, val_loss=3240.1353][2026-09-23 20:15:45,939] [UniVITrainer] [INFO] [Epoch 117] New best val loss: 3240.1353
Training UniVI: 40%|███▉ | 119/300 [07:50<12:14, 4.06s/it, beta=1.250, gamma=4.350, train_loss=3109.6978, val_loss=3237.8623][2026-09-23 20:15:58,310] [UniVITrainer] [INFO] [Epoch 120] New best val loss: 3237.8623
Training UniVI: 56%|█████▋ | 169/300 [11:17<09:04, 4.16s/it, beta=1.250, gamma=4.350, train_loss=3065.3826, val_loss=3250.7517][2026-09-23 20:19:24,646] [UniVITrainer] [INFO] Early stopping at epoch 170 (best val loss=3237.8623, best epoch=120)
Training UniVI: 56%|█████▋ | 169/300 [11:17<08:45, 4.01s/it, beta=1.250, gamma=4.350, train_loss=3065.3826, val_loss=3250.7517]
[2026-09-23 20:19:24,682] [UniVITrainer] [INFO] Restored best model from epoch 120 (val loss=3237.8623)
| rna embedding s.d. across cells | atac embedding s.d. across cells | |
|---|---|---|
| MLP (v1) | 0.5211 | 0.4586 |
| transformer (v1) | 0.3129 | 0.3017 |
| transformer + position (v1) | 0.2819 | 0.2620 |
| MLP (v2) | 0.3668 | 0.3114 |
| fused transformer (v2) | 0.3110 | 0.2925 |
fig, ax = plt.subplots(figsize=(5.5, 3.2))
for name, h in histories.items():
ax.plot(h["val_loss"], label=name)
ax.set(xlabel="epoch", ylabel="validation loss", yscale="log", title="Validation loss (full beta, gamma)")
ax.legend(frameon=False, fontsize=8)
plt.show()
Losses of v1 and v2 models are different objectives, so compare curves only within an objective.
3. Benchmark on held-out cells#
FOSCTTM and Recall@10: are a cell’s RNA and ATAC embeddings close to each other?
label transfer (RNA → ATAC, k-NN) and its macro-F1 in the worse direction;
ATAC → RNA: median per-gene Pearson correlation of RNA predicted from accessibility alone;
fused kNN accuracy: 5-fold cross-validated 15-NN cell-type accuracy on the embedding that uses both modalities (the fused transformer’s posterior, or the precision-weighted MoE for the others).
labels = test["rna"].obs["cell_type"].astype(str).to_numpy()
observed = dense(test["rna"].X)
rows, fused_z = {}, {}
for name, model in models.items():
z_r = encode_adata(model, test["rna"], modality="rna", device=device, latent="modality_mean")
z_a = encode_adata(model, test["atac"], modality="atac", device=device, latent="modality_mean")
m = evaluate_alignment(Z1=z_r, Z2=z_a, labels_source=labels, labels_target=labels, recall_ks=(10,))
pred = cross_modal_predict(model, test["atac"], src_mod="atac", tgt_mod="rna", device=device)
fz = encode_fused_adata_pair(model, {"rna": test["rna"], "atac": test["atac"]}, device=device,
write_to_adatas=False)["mu"]
fused_z[name] = fz
folds = int(min(5, pd.Series(labels).value_counts().min())) if len(set(labels)) > 1 else 0
acc = (cross_val_score(KNeighborsClassifier(15), fz, labels, cv=folds).mean() if folds >= 2 else np.nan)
rows[name] = {"FOSCTTM ↓": m["foscttm_mean"], "Recall@10 ↑": m["recall_at_k"]["10"]["mean"],
"label transfer ↑": m["label_transfer_acc"], "macro-F1 (worse dir.) ↑": m["worst_direction_macro_f1"],
"ATAC→RNA median r ↑": float(np.nanmedian(pearson_corr_per_feature(observed, pred))),
"fused kNN acc ↑": acc, "train min": minutes[name]}
bench = pd.DataFrame(rows).T
bench.round(3)
| FOSCTTM ↓ | Recall@10 ↑ | label transfer ↑ | macro-F1 (worse dir.) ↑ | ATAC→RNA median r ↑ | fused kNN acc ↑ | train min | |
|---|---|---|---|---|---|---|---|
| MLP (v1) | 0.047 | 0.355 | 0.874 | 0.674 | 0.260 | 0.871 | 6.473 |
| transformer (v1) | 0.074 | 0.185 | 0.759 | 0.559 | 0.219 | 0.838 | 9.422 |
| transformer + position (v1) | 0.073 | 0.164 | 0.771 | 0.563 | 0.223 | 0.830 | 12.104 |
| MLP (v2) | 0.053 | 0.306 | 0.857 | 0.671 | 0.232 | 0.826 | 8.721 |
| fused transformer (v2) | 0.047 | 0.335 | 0.877 | 0.692 | 0.253 | 0.840 | 11.290 |
This is a single seed on one dataset. Before concluding that one encoder is better, repeat the comparison over a few seeds (differences between seeds can be as large as differences between models).
fig, axes = plt.subplots(1, len(models), figsize=(3.6 * len(models), 3.4))
for ax, (name, fz) in zip(np.atleast_1d(axes), fused_z.items()):
a = ad.AnnData(fz.astype(np.float32),
obs=pd.DataFrame({"cell_type": pd.Categorical(labels)}, index=test["rna"].obs_names))
sc.pp.neighbors(a, n_neighbors=30, use_rep="X")
sc.tl.umap(a, random_state=0)
sc.pl.umap(a, color="cell_type", ax=ax, show=False, legend_loc="none", title=name, frameon=False)
plt.tight_layout()
plt.show()
4. Inside the fused transformer#
fused_encode_with_meta_and_attn runs the fused encoder and returns the attention matrices of every
layer (averaged over heads) together with a token map: which slice of the sequence belongs to which
modality and, for top-k tokenizers, which feature each token was in each cell.
fused = models["fused transformer (v2)"]
fused.eval()
n_cells = min(MAX_ATTENTION_CELLS, test["rna"].n_obs)
cells = np.random.default_rng(0).choice(test["rna"].n_obs, n_cells, replace=False)
X_rna, X_atac = dense(test["rna"].X[cells]), dense(test["atac"].X[cells])
def attention_batches(batch=256, model=None, limit=None):
model = fused if model is None else model
for s in range(0, n_cells if limit is None else min(limit, n_cells), batch):
xb = {"rna": torch.as_tensor(X_rna[s:s + batch], dtype=torch.float32, device=device),
"atac": torch.as_tensor(X_atac[s:s + batch], dtype=torch.float32, device=device)}
_, _, tokmap, attn = fused_encode_with_meta_and_attn(model, xb, return_attn=True)
yield s, tokmap, attn
_, tokmap0, attn0 = next(attention_batches(8))
print("token slices:", tokmap0.slices, "| global CLS:", tokmap0.has_global_cls,
"| layers:", len(attn0), "| attention shape:", tuple(attn0[0].shape))
token slices: {'rna': (0, 128), 'atac': (128, 256)} | global CLS: False | layers: 2 | attention shape: (8, 256, 256)
How much attention crosses modalities?#
For each layer, the average attention a query token of one modality sends to key tokens of each modality (rows sum to 1). Off-diagonal mass is cross-modal mixing.
mods = list(tokmap0.slices)
mass = np.zeros((len(attn0), len(mods), len(mods)))
count = 0
for s, tokmap, attn in attention_batches():
shift = 1 if tokmap.has_global_cls else 0
for li, A in enumerate(attn):
A = A.float().cpu().numpy()
for i, qa in enumerate(mods):
a0, a1 = [v + shift for v in tokmap.slices[qa]]
for j, kb in enumerate(mods):
b0, b1 = [v + shift for v in tokmap.slices[kb]]
mass[li, i, j] += A[:, a0:a1, b0:b1].sum(-1).mean(-1).sum()
count += A.shape[0]
mass /= count
fig, axes = plt.subplots(1, len(attn0), figsize=(2.8 * len(attn0), 2.6))
for li, ax in enumerate(np.atleast_1d(axes)):
im = ax.imshow(mass[li], vmin=0, vmax=1, cmap="magma")
for i in range(len(mods)):
for j in range(len(mods)):
ax.text(j, i, f"{mass[li, i, j]:.2f}", ha="center", va="center", color="w", fontsize=9)
ax.set_xticks(range(len(mods)), [f"→{m}" for m in mods])
ax.set_yticks(range(len(mods)), mods)
ax.set_title(f"layer {li + 1}")
plt.tight_layout()
plt.show()
One cell, every token#
The full last-layer attention matrix of a single test cell: rows are query tokens, columns are key tokens, gene tokens first and peak tokens after the white lines. Color is on a log scale, because a few entries dominate. The most-attended tokens (largest column sums) are labeled.
_, tokmap1, attn1 = next(attention_batches(1))
shift1 = 1 if tokmap1.has_global_cls else 0
A1 = attn1[-1][0].float().cpu().numpy()
names1 = ([test["rna"].var_names[i] for i in tokmap1.meta["rna"]["topk_idx"][0].cpu().numpy()]
+ [test["atac"].var_names[i] for i in tokmap1.meta["atac"]["topk_idx"][0].cpu().numpy()])
names1 = (["CLS"] if shift1 else []) + names1
fig, ax = plt.subplots(figsize=(6.2, 5.4))
im = ax.imshow(np.log10(A1 + 1e-6), cmap="magma", aspect="auto")
cut = tokmap1.slices["rna"][1] + shift1
ax.axhline(cut - 0.5, c="w", lw=0.8)
ax.axvline(cut - 0.5, c="w", lw=0.8)
top_cols = np.argsort(-A1.sum(0))[:8]
ax.set_xticks(top_cols, [names1[j][:18] for j in top_cols], rotation=90, fontsize=6)
ax.set_yticks([cut / 2, cut + (A1.shape[0] - cut) / 2], ["gene tokens", "peak tokens"], fontsize=8)
plt.colorbar(im, ax=ax, label="log10 attention")
ax.set_title(f"Last-layer attention of one {labels[cells[0]]} cell", fontsize=9)
plt.tight_layout()
plt.show()
How sharp is attention?#
For every query token, the entropy of its attention distribution divided by the entropy of uniform attention over all tokens (1 = spread evenly, near 0 = focused on one token). Sharp attention in a layer means that layer routes information between specific tokens; near-uniform attention means it mostly averages.
ent = {li: [] for li in range(len(attn0))}
ent_cell = []
for s_, tokmap, attn in attention_batches():
for li, A in enumerate(attn):
A = A.float().cpu().numpy()
h = -(A * np.log(A + 1e-12)).sum(-1) / np.log(A.shape[-1]) # (cells, queries)
ent[li].append(h.ravel())
if li == len(attn) - 1:
ent_cell.append(h.mean(1))
ent_cell = np.concatenate(ent_cell)
fig, axes = plt.subplots(1, 2, figsize=(9, 3), gridspec_kw={"width_ratios": [1, 1.3]})
for li, h in ent.items():
axes[0].hist(np.concatenate(h), bins=50, histtype="step", density=True, lw=1.4, label=f"layer {li + 1}")
axes[0].set(xlabel="normalized attention entropy per query", ylabel="density")
axes[0].legend(frameon=False, fontsize=8)
per_type = pd.Series(ent_cell, index=labels[cells[:len(ent_cell)]]).groupby(level=0).mean().sort_values()
axes[1].barh(per_type.index, per_type.values, color="0.5")
axes[1].set(xlabel="mean normalized entropy, last layer", xlim=(per_type.min() * 0.98, 1.0))
axes[1].tick_params(axis="y", labelsize=6)
plt.tight_layout()
plt.show()
Which tokens does the fused embedding depend on? Attention rollout#
The fused posterior is computed from the mean of the last layer’s token states. Attention rollout (Abnar & Zuidema, 2020) traces that back through the layers: each layer’s attention is mixed with the identity (for the residual connection), row-normalized, and the layers are multiplied, which gives how much every input token contributes to every output token. Averaging over output tokens gives each input token’s share of the pooled embedding. Rollout is an approximation (it ignores the feed-forward blocks and treats heads as averaged), so read it as a ranking.
For each cell type, genes are scored by their mean rollout share over the type’s cells. The heatmap shows the top genes of each type: “attention markers”.
n_genes = test["rna"].var_names.size
cell_types_t = pd.Index(sorted(set(labels[cells])))
imp_sum = np.zeros((len(cell_types_t), n_genes))
present_n = np.zeros((len(cell_types_t), n_genes))
type_n = np.zeros(len(cell_types_t))
cross_share = []
for s_, tokmap, attn in attention_batches():
shift = 1 if tokmap.has_global_cls else 0
R = None
for A in attn:
A = A.float().cpu().numpy()
A_hat = 0.5 * A + 0.5 * np.eye(A.shape[-1])[None]
A_hat /= A_hat.sum(-1, keepdims=True)
R = A_hat if R is None else A_hat @ R
share = R.mean(1) # (cells, tokens): share of pooled output
a0, a1 = [v + shift for v in tokmap.slices["rna"]]
b0, b1 = [v + shift for v in tokmap.slices["atac"]]
gi = tokmap.meta["rna"]["topk_idx"].cpu().numpy()
t_idx = cell_types_t.get_indexer(labels[cells[s_:s_ + len(gi)]])
for c in range(len(gi)):
imp_sum[t_idx[c], gi[c]] += share[c, a0:a1]
present_n[t_idx[c], gi[c]] += 1
type_n[t_idx[c]] += 1
A_last = attn[-1].float().cpu().numpy()
cross_share.append(np.c_[A_last[:, a0:a1, b0:b1].sum(-1).mean(-1), A_last[:, b0:b1, a0:a1].sum(-1).mean(-1)])
cross_share = np.concatenate(cross_share)
rollout = pd.DataFrame(imp_sum / np.maximum(type_n[:, None], 1), index=cell_types_t, columns=test["rna"].var_names)
token_freq = pd.DataFrame(present_n / np.maximum(type_n[:, None], 1), index=cell_types_t, columns=test["rna"].var_names)
when_present = pd.DataFrame(np.where(present_n > 0, imp_sum / np.maximum(present_n, 1), np.nan),
index=cell_types_t, columns=test["rna"].var_names)
show_types = [t for t in cell_types_t if type_n[cell_types_t.get_loc(t)] >= 10]
top_genes_attn = list(dict.fromkeys(g for t in show_types for g in rollout.loc[t].nlargest(3).index))
M = rollout.loc[show_types, top_genes_attn]
M = M.div(M.max(axis=0) + 1e-12, axis=1)
fig, ax = plt.subplots(figsize=(0.22 * len(top_genes_attn) + 2.5, 0.28 * len(show_types) + 1.6))
im = ax.imshow(M.to_numpy(), cmap="viridis", aspect="auto")
ax.set_xticks(range(len(top_genes_attn)), top_genes_attn, rotation=90, fontsize=6)
ax.set_yticks(range(len(show_types)), show_types, fontsize=7)
plt.colorbar(im, ax=ax, label="rollout share (scaled per gene)")
ax.set_title("Attention markers: genes the fused embedding relies on, per cell type", fontsize=9)
plt.tight_layout()
plt.show()
Does attention add anything to being a token? Tokens are a cell’s most over-expressed genes, so a
cell type’s marker genes become tokens in its cells whatever the transformer learns. The test compares
three gene scores for recovering each type’s N_DE_GENES top differentially expressed genes (t-test on
the test cells, one type against the rest), as AUROC over genes that are tokens in that type:
token frequency: fraction of the type’s cells in which the gene is a token (no attention involved);
rollout share: mean rollout share over the type’s cells (attention plus frequency);
share when present: mean rollout share over cells in which the gene is a token (attention only).
If “share when present” stays near 0.5 while token frequency is high, the model’s attention adds little cell-type information beyond what token selection already carries.
from sklearn.metrics import roc_auc_score
de_adata = ad.AnnData(dense(test["rna"].X[cells]).astype(np.float32), obs=pd.DataFrame(
{"cell_type": pd.Categorical(labels[cells])}, index=test["rna"].obs_names[cells]), var=pd.DataFrame(index=test["rna"].var_names))
import warnings
with warnings.catch_warnings(): # inputs are z-scored, so scanpy's log fold changes are undefined (unused here)
warnings.simplefilter("ignore", RuntimeWarning)
sc.tl.rank_genes_groups(de_adata, "cell_type", groups=show_types, method="t-test", n_genes=N_DE_GENES)
rows = {}
for t in show_types:
de = set(de_adata.uns["rank_genes_groups"]["names"][t])
ok = present_n[cell_types_t.get_loc(t)] > 0
genes_t = test["rna"].var_names[ok]
y = np.array([g in de for g in genes_t])
if y.all() or not y.any():
continue
rows[t] = {"genes scored": int(ok.sum()),
"token frequency": roc_auc_score(y, token_freq.loc[t, genes_t]),
"rollout share": roc_auc_score(y, rollout.loc[t, genes_t]),
"share when present": roc_auc_score(y, when_present.loc[t, genes_t])}
de_test = pd.DataFrame(rows).T
display(de_test.round(3))
print("mean AUROC:", de_test[["token frequency", "rollout share", "share when present"]].mean().round(3).to_dict())
| genes scored | token frequency | rollout share | share when present | |
|---|---|---|---|---|
| CD14 Mono | 1355.0 | 0.765 | 0.743 | 0.333 |
| CD16 Mono | 886.0 | 0.815 | 0.802 | 0.435 |
| CD4 Naive | 1137.0 | 0.933 | 0.964 | 0.555 |
| CD4 TCM | 1165.0 | 0.974 | 0.984 | 0.479 |
| CD4 TEM | 805.0 | 0.976 | 0.983 | 0.610 |
| CD8 Naive | 1157.0 | 0.950 | 0.971 | 0.569 |
| CD8 TEM_1 | 819.0 | 0.954 | 0.971 | 0.613 |
| CD8 TEM_2 | 704.0 | 0.988 | 0.987 | 0.646 |
| Intermediate B | 792.0 | 0.978 | 0.952 | 0.428 |
| MAIT | 543.0 | 0.976 | 0.950 | 0.581 |
| Memory B | 766.0 | 0.997 | 0.978 | 0.438 |
| NK | 878.0 | 0.994 | 0.988 | 0.664 |
| Naive B | 570.0 | 0.966 | 0.939 | 0.469 |
| Treg | 670.0 | 0.948 | 0.966 | 0.617 |
| cDC | 717.0 | 0.783 | 0.750 | 0.446 |
| gdT | 618.0 | 0.984 | 0.971 | 0.646 |
| pDC | 437.0 | 0.899 | 0.889 | 0.487 |
mean AUROC: {'token frequency': 0.934, 'rollout share': 0.929, 'share when present': 0.53}
Cross-modal attention by cell type#
The share of last-layer attention that gene tokens send to peak tokens, and peak tokens to gene tokens, per cell type. Cell types that lean more on the other modality are those where the fused encoder mixes chromatin and expression most.
cs = pd.DataFrame(cross_share, columns=["gene → peak", "peak → gene"], index=labels[cells[:len(cross_share)]])
per_type_cs = cs.groupby(level=0).mean().loc[show_types].sort_values("gene → peak")
ax = per_type_cs.plot.barh(figsize=(5.5, 0.25 * len(per_type_cs) + 1.4), width=0.8)
ax.set(xlabel="share of attention sent to the other modality (last layer)", ylabel="")
ax.tick_params(axis="y", labelsize=7)
ax.legend(frameon=False, fontsize=8)
plt.tight_layout()
plt.show()
Collect the last layer’s gene → peak attention#
For the remaining analyses we keep, for every analyzed cell, the block of last-layer attention from gene tokens (queries) to peak tokens (keys), the feature index behind every token, and a double-centered copy of the block (each gene token’s and each peak token’s mean removed), which keeps only pair-specific preferences.
pv = test["atac"].var
peak_chrom = pv["chrom"].astype(str).to_numpy() if "chrom" in pv else None
peak_mid = ((pv["chromStart"].astype(np.int64) + pv["chromEnd"].astype(np.int64)) // 2).to_numpy() \
if {"chromStart", "chromEnd"} <= set(pv.columns) else None
As, Rs, GIs, PIs = [], [], [], []
for s, tokmap, attn in attention_batches():
shift = 1 if tokmap.has_global_cls else 0
a0, a1 = [v + shift for v in tokmap.slices["rna"]]
b0, b1 = [v + shift for v in tokmap.slices["atac"]]
A = attn[-1][:, a0:a1, b0:b1].float().cpu().numpy() # (cells, gene tokens, peak tokens)
As.append(A)
Rs.append(A - A.mean(2, keepdims=True) - A.mean(1, keepdims=True) + A.mean((1, 2), keepdims=True))
GIs.append(tokmap.meta["rna"]["topk_idx"].cpu().numpy()) # feature index of each token
PIs.append(tokmap.meta["atac"]["topk_idx"].cpu().numpy())
A, R, GI, PI = map(np.concatenate, (As, Rs, GIs, PIs))
del As, Rs
print("attention blocks:", A.shape)
attention blocks: (977, 128, 128)
Attention sinks#
Transformers often route much of their attention to a few “sink” tokens that act as a default target rather than carrying pair-specific information. The table lists the peaks that receive the largest share of all gene → peak attention in the last layer, and in how many analyzed cells each is a token. If a few peaks take a large share, rankings of summed attention (such as the next table) are mostly lists of genes paired with those peaks.
received = np.bincount(PI.ravel(), weights=A.sum(1).ravel(), minlength=test["atac"].n_vars)
present = np.bincount(PI.ravel(), minlength=test["atac"].n_vars) # top-k tokens are distinct within a cell
sinks = pd.DataFrame({"share of gene→peak attention": received / received.sum(),
"cells where it is a token": present}, index=test["atac"].var_names)
sinks = sinks.sort_values("share of gene→peak attention", ascending=False)
print(f"top 1 peak: {sinks.iloc[0, 0]:.1%} of all gene→peak attention; top 10 peaks: {sinks.iloc[:10, 0].sum():.1%} "
f"(a uniform split over tokens would give each peak about {1 / PI.shape[1]:.1%} within a cell)")
sinks.head(10).round(4)
top 1 peak: 0.3% of all gene→peak attention; top 10 peaks: 2.4% (a uniform split over tokens would give each peak about 0.8% within a cell)
| share of gene→peak attention | cells where it is a token | |
|---|---|---|
| peaks | ||
| chr16:4090281-4090802 | 0.0029 | 12 |
| chr12:52603270-52603954 | 0.0027 | 37 |
| chr1:183048585-183049120 | 0.0027 | 15 |
| chr18:55423309-55424046 | 0.0027 | 14 |
| chr18:807940-808918 | 0.0024 | 22 |
| chr2:197802899-197803424 | 0.0024 | 16 |
| chr20:24987390-24989486 | 0.0022 | 22 |
| chr3:134562238-134562874 | 0.0021 | 11 |
| chr18:32803310-32803987 | 0.0021 | 10 |
| chr6:32963824-32964326 | 0.0020 | 15 |
Gene → peak pairs#
top_cross_modal_feature_pairs_from_attn sums, over cells, the raw attention from gene tokens to peak
tokens for each cell’s strongest pairs. It is shown first as the built-in summary; it is dominated by
sinks and by features that are tokens in many cells.
_, tokmap_b, attn_b = next(attention_batches(512))
pairs = top_cross_modal_feature_pairs_from_attn(
attn_b, tokmap_b, mod_a="rna", mod_b="atac",
var_names_by_mod={"rna": test["rna"].var_names, "atac": test["atac"].var_names},
tokenizer_mode_by_mod={"rna": "topk_channels", "atac": "topk_channels"},
layer=-1, top_pairs_per_cell=50, top_n=15)
pd.DataFrame(pairs, columns=["gene", "peak", "summed_attention"]).round(3)
| gene | peak | summed_attention | |
|---|---|---|---|
| 0 | CAMK4 | chr12:52603270-52603954 | 1.226 |
| 1 | CAMK4 | chr6:3623315-3623980 | 1.044 |
| 2 | TXK | chr12:52603270-52603954 | 0.974 |
| 3 | HLA-DQA1 | chr18:55423309-55424046 | 0.947 |
| 4 | MS4A1 | chr18:55423309-55424046 | 0.898 |
| 5 | CD79A | chr18:55423309-55424046 | 0.882 |
| 6 | OSBPL10 | chr18:55423309-55424046 | 0.871 |
| 7 | TRBC2 | chr12:52603270-52603954 | 0.860 |
| 8 | SETBP1 | chr18:55423309-55424046 | 0.856 |
| 9 | FHIT | chr6:3623315-3623980 | 0.843 |
| 10 | BCL2 | chr4:153398580-153399196 | 0.830 |
| 11 | NKG7 | chr1:168509220-168510861 | 0.829 |
| 12 | DENND5B | chr18:55423309-55424046 | 0.809 |
| 13 | CAMK4 | chr18:807940-808918 | 0.781 |
| 14 | TXK | chr11:62637645-62638404 | 0.770 |
A pair-specific ranking uses the centered attention instead: in every cell, the TOP_PAIRS_PER_CELL
pairs with the largest centered attention are counted, and each pair’s count is divided by the number of
cells in which both of its features are tokens (pairs co-present in fewer than MIN_COOCCURRENCE cells
are dropped). A high rate means the model singles out that gene–peak combination whenever it can.
def annotate_pairs(df):
if tss is None or peak_chrom is None:
return df
pidx = test["atac"].var_names.get_indexer(df["peak"])
g = tss.reindex(df["gene"])
same = g["chrom"].to_numpy() == peak_chrom[pidx]
df["same_chromosome"] = same
df["distance_to_tss_kb"] = np.where(same, np.abs(peak_mid[pidx] - g["tss"].to_numpy()) / 1e3, np.nan)
return df
n_c, n_g, n_p = R.shape
k = min(TOP_PAIRS_PER_CELL, n_g * n_p)
top = np.argpartition(-R.reshape(n_c, -1), k - 1, axis=1)[:, :k]
g_feat = np.take_along_axis(GI, top // n_p, axis=1)
p_feat = np.take_along_axis(PI, top % n_p, axis=1)
n_peak_features = test["atac"].n_vars
counts = pd.Series((g_feat.astype(np.int64) * n_peak_features + p_feat).ravel()).value_counts()
cand = counts.index[:2000].to_numpy()
cg, cp = cand // n_peak_features, cand % n_peak_features
gene_present = np.zeros((n_c, test["rna"].n_vars), dtype=bool)
gene_present[np.arange(n_c)[:, None], GI] = True
peak_present = np.zeros((n_c, n_peak_features), dtype=bool)
peak_present[np.arange(n_c)[:, None], PI] = True
co = (gene_present[:, cg] & peak_present[:, cp]).sum(0)
ranked = pd.DataFrame({"gene": test["rna"].var_names[cg], "peak": test["atac"].var_names[cp],
"cells in top pairs": counts.loc[cand].to_numpy(), "cells co-present": co})
ranked = ranked[ranked["cells co-present"] >= MIN_COOCCURRENCE]
if ranked.empty:
print(f"No pair is co-present in at least {MIN_COOCCURRENCE} analyzed cells; lower MIN_COOCCURRENCE or "
"raise MAX_ATTENTION_CELLS.")
ranked["rate"] = ranked["cells in top pairs"] / ranked["cells co-present"]
annotate_pairs(ranked.sort_values(["rate", "cells in top pairs"], ascending=False).head(25).reset_index(drop=True)).round(3)
| gene | peak | cells in top pairs | cells co-present | rate | same_chromosome | distance_to_tss_kb | |
|---|---|---|---|---|---|---|---|
| 0 | TRBC2 | chr12:52603270-52603954 | 11 | 26 | 0.423 | False | NaN |
| 1 | TXK | chr11:62637645-62638404 | 8 | 20 | 0.400 | False | NaN |
| 2 | CAMK4 | chr12:52603270-52603954 | 13 | 33 | 0.394 | False | NaN |
| 3 | NELL2 | chr12:52603270-52603954 | 8 | 22 | 0.364 | True | 7681.764 |
| 4 | CAMK4 | chr11:62637645-62638404 | 7 | 20 | 0.350 | False | NaN |
| 5 | TXK | chr12:52603270-52603954 | 10 | 30 | 0.333 | False | NaN |
| 6 | APBA2 | chr12:52603270-52603954 | 8 | 27 | 0.296 | False | NaN |
| 7 | BCL2 | chr12:52603270-52603954 | 8 | 29 | 0.276 | False | NaN |
| 8 | ACTG1 | chr12:52603270-52603954 | 8 | 31 | 0.258 | False | NaN |
| 9 | INPP4B | chr12:52603270-52603954 | 6 | 28 | 0.214 | False | NaN |
| 10 | TXK | chr20:24065830-24066674 | 5 | 24 | 0.208 | False | NaN |
| 11 | ACTG1 | chr11:62637645-62638404 | 4 | 20 | 0.200 | False | NaN |
| 12 | CAMK4 | chr20:24065830-24066674 | 4 | 22 | 0.182 | False | NaN |
| 13 | CAMK4 | chr4:36208801-36209362 | 4 | 22 | 0.182 | False | NaN |
| 14 | PFN1 | chr14:22534553-22535667 | 3 | 21 | 0.143 | False | NaN |
| 15 | TMSB4X | chr12:52603270-52603954 | 4 | 32 | 0.125 | False | NaN |
| 16 | VIM | chr12:52603270-52603954 | 3 | 26 | 0.115 | False | NaN |
| 17 | PFN1 | chr10:28594055-28594888 | 3 | 27 | 0.111 | False | NaN |
| 18 | IL7R | chr12:52603270-52603954 | 3 | 30 | 0.100 | False | NaN |
| 19 | PFN1 | chr5:102801307-102801847 | 2 | 20 | 0.100 | False | NaN |
| 20 | NKG7 | chr20:24157139-24157672 | 2 | 21 | 0.095 | False | NaN |
| 21 | CAMK4 | chr12:21690585-21691265 | 2 | 21 | 0.095 | False | NaN |
| 22 | LTB | chr12:52603270-52603954 | 3 | 33 | 0.091 | False | NaN |
| 23 | INPP4B | chr12:21690585-21691265 | 2 | 23 | 0.087 | False | NaN |
| 24 | ACTB | chr10:28594055-28594888 | 2 | 23 | 0.087 | False | NaN |
Do gene tokens attend more to their cis peaks?#
A (gene, peak) pair is cis if the peak lies on the gene’s chromosome within CIS_WINDOW_BP of its
TSS. A naive test (“are cis pairs over-represented among the highest attention weights?”) is
confounded by sinks and by genes that send most of their attention to ATAC; on the synthetic stand-in
data used for our automated tests, which has no gene-peak structure at all, the naive version reported a
spurious 2-3x enrichment. The test therefore uses the centered attention: for every gene token with at
least one cis and one non-cis peak token, the AUC with which its centered attention separates cis
from non-cis peaks, averaged over gene tokens (0.5 = no preference).
The same gene–peak pairs recur in many cells, so cells are not independent samples and a bootstrap over
cells would be too optimistic. Instead the statistic is calibrated with a permutation null: peak
coordinates are shuffled across peaks (N_PERMUTATIONS times), which keeps the attention exactly as it
is but breaks its relation to genomic position. Both one-sided p-values are reported; the question asked
in advance is whether cis peaks get more attention.
def mean_cis_auc(R, gi, pi, chrom_of_peak, mid_of_peak):
"""Mean within-gene AUC (cis vs non-cis peak tokens) of centered attention R (cells, genes, peaks)."""
cis = ((g_chrom[gi][:, :, None] == chrom_of_peak[pi][:, None, :]) &
(np.abs(g_tss[gi][:, :, None] - mid_of_peak[pi][:, None, :]) <= CIS_WINDOW_BP) &
gene_ok[gi][:, :, None])
n1 = cis.sum(2)
n0 = cis.shape[2] - n1
ok = (n1 > 0) & (n0 > 0)
ranks = R.argsort(2).argsort(2) + 1.0 # within-row ranks (ties negligible)
u = (ranks * cis).sum(2) - n1 * (n1 + 1) / 2
auc = np.where(ok, u / np.maximum(n1 * n0, 1), np.nan)
return np.nanmean(auc) if ok.any() else np.nan, int(ok.sum())
if tss is None or peak_chrom is None:
print("Gene or peak coordinates are missing; skipping the cis test.")
else:
gene_ok = test["rna"].var_names.isin(tss.index)
g_chrom = np.full(test["rna"].n_vars, "", dtype=object)
g_tss = np.zeros(test["rna"].n_vars, dtype=np.int64)
g_chrom[gene_ok] = tss.loc[test["rna"].var_names[gene_ok], "chrom"].to_numpy()
g_tss[gene_ok] = tss.loc[test["rna"].var_names[gene_ok], "tss"].to_numpy()
observed_auc, n_rows = mean_cis_auc(R, GI, PI, peak_chrom, peak_mid)
rng = np.random.default_rng(0)
null_auc = []
for _ in range(N_PERMUTATIONS):
perm = rng.permutation(len(peak_chrom))
null_auc.append(mean_cis_auc(R, GI, PI, peak_chrom[perm], peak_mid[perm])[0])
null_auc = np.asarray(null_auc)
null_auc = null_auc[np.isfinite(null_auc)]
if not np.isfinite(observed_auc):
print("No gene token had both cis and non-cis peak tokens; widen CIS_WINDOW_BP or raise the token counts.")
else:
p_more = (1 + (null_auc >= observed_auc).sum()) / (1 + len(null_auc))
p_less = (1 + (null_auc <= observed_auc).sum()) / (1 + len(null_auc))
print(f"eligible gene tokens: {n_rows:,} in {len(R):,} cells")
print(f"mean within-gene AUC (cis vs non-cis peaks): {observed_auc:.3f}; permutation null "
f"{null_auc.mean():.3f} ± {null_auc.std():.3f} ({len(null_auc)} permutations)")
print(f"p (cis gets more attention) = {p_more:.3f}; p (cis gets less attention) = {p_less:.3f}")
z_trained = (observed_auc - null_auc.mean()) / null_auc.std()
fig, ax = plt.subplots(figsize=(4, 2.6))
ax.hist(null_auc, bins=25, color="0.7", label="shuffled peak coordinates")
ax.axvline(observed_auc, c="C3", label="observed")
ax.set(xlabel="mean within-gene AUC", ylabel="permutations")
ax.legend(frameon=False, fontsize=7)
plt.show()
eligible gene tokens: 16,216 in 977 cells
mean within-gene AUC (cis vs non-cis peaks): 0.499; permutation null 0.501 ± 0.004 (200 permutations)
p (cis gets more attention) = 0.716; p (cis gets less attention) = 0.289
A reference for how extreme “no knowledge” looks. The permutation test is calibrated, but a single
model can still land a couple of standard deviations from the null by chance. The same statistic is
computed for N_UNTRAINED_CONTROLS untrained copies of the fused model (random weights, so their
attention cannot know anything about genomic position). When we ran 20 such untrained models on
synthetic data, their z-scores ranged from about -2.2 to +2.7, with 1 of 20 one-sided p-values at or
below 0.05, as expected for a calibrated test. Only a trained-model z well outside the controls’ range,
reproduced across training seeds, is worth interpreting.
def cis_z(R_, GI_, PI_, n_perm, seed):
obs, _ = mean_cis_auc(R_, GI_, PI_, peak_chrom, peak_mid)
rng_ = np.random.default_rng(seed)
null = np.array([mean_cis_auc(R_, GI_, PI_, peak_chrom[pp], peak_mid[pp])[0]
for pp in (rng_.permutation(len(peak_chrom)) for _ in range(n_perm))])
null = null[np.isfinite(null)]
return (obs - null.mean()) / null.std() if np.isfinite(obs) and len(null) > 1 else np.nan
if tss is not None and peak_chrom is not None and N_UNTRAINED_CONTROLS > 0 and "z_trained" in globals():
cfg_f, kw_f = specs["fused transformer (v2)"]
control = []
for seed in range(N_UNTRAINED_CONTROLS):
set_seed(1000 + seed)
untrained = UniVIMultiModalVAE(cfg_f, normalize_v1_terms=True, **kw_f).to(device).eval()
Rc, GIc, PIc = [], [], []
for s, tokmap, attn in attention_batches(model=untrained, limit=500):
shift = 1 if tokmap.has_global_cls else 0
a0, a1 = [v + shift for v in tokmap.slices["rna"]]
b0, b1 = [v + shift for v in tokmap.slices["atac"]]
Ac = attn[-1][:, a0:a1, b0:b1].float().cpu().numpy()
Rc.append(Ac - Ac.mean(2, keepdims=True) - Ac.mean(1, keepdims=True) + Ac.mean((1, 2), keepdims=True))
GIc.append(tokmap.meta["rna"]["topk_idx"].cpu().numpy())
PIc.append(tokmap.meta["atac"]["topk_idx"].cpu().numpy())
control.append(cis_z(np.concatenate(Rc), np.concatenate(GIc), np.concatenate(PIc), 50, seed))
print(f"trained fused model: z = {z_trained:+.2f}; untrained controls: "
+ ", ".join(f"{z:+.2f}" for z in control))
trained fused model: z = -0.49; untrained controls: +0.39, -2.06, -0.24, -0.04, -0.52
An observed value inside the permutation distribution and inside the range of the untrained controls
means the last layer shows no preference for cis peaks at this window. A clearly larger value would
suggest the fused encoder links genes to nearby regulatory regions, and a clearly smaller one that it
systematically attends away from them; either is worth checking across seeds and against the earlier
layer (attn[0] above) before interpreting it.
5. Peak–gene links: a benchmark#
Can any of these models point to the peaks that regulate a gene? Every (gene, peak) pair within
CIS_WINDOW_BP of the gene’s TSS is a candidate, and each is scored four ways:
distance to the TSS (closer = higher): the baseline every method has to beat;
correlation across pseudobulks, the standard data-driven approach behind Signac’s
LinkPeaksand ArchR’s peak-to-gene links: training cells are grouped intoN_PSEUDOBULKSclusters, and the gene’s mean expression is correlated with the peak’s mean accessibility across clusters. Each correlation is turned into a z-score against background peaks on other chromosomes with similar accessibility, and pairs with z ≥LINK_Zare the reference links. They are associations in these data, not proven regulatory links;attention: the fused transformer’s centered last-layer attention from the gene token to the peak token, averaged over the cells in which both are tokens (at least
MIN_COPRESENT);attribution: input × gradient of the gene’s predicted expression (from ATAC alone, model
ATTRIBUTION_MODEL) with respect to the peak, averaged overN_ATTRIBUTION_CELLStest cells. Positive means opening the peak raises the model’s prediction for the gene.
Near peaks are more often linked than far ones, so a score can look good simply by preferring near pairs. Each score is therefore also evaluated within distance bins, which removes that advantage.
Two reference sets. Correlation across all pseudobulks is dominated by cell-type differences: for a
lineage marker such as LYZ or CD79A, hundreds of peaks elsewhere in the genome that open in the same cell
type correlate with it as well as its own promoter does, so the background absorbs the signal and markers
get no links (in an earlier run CD79A, CD14, LYZ and SPI1 had none). The second reference therefore
correlates genes and peaks within each lineage in LINEAGES, over pseudobulks built from that lineage’s
cells only, where cell-type co-variation is much weaker; a pair is a within-lineage link if its z-score
reaches LINK_Z in at least one lineage.
from scipy.stats import spearmanr
from sklearn.cluster import MiniBatchKMeans
from sklearn.decomposition import PCA
gene_names, peak_names = test["rna"].var_names, test["atac"].var_names
if tss is None or peak_chrom is None:
raise RuntimeError("Gene and peak coordinates are needed for the peak-gene benchmark.")
g_idx, p_idx, dist_bp = [], [], []
peak_df = pd.DataFrame({"chrom": peak_chrom, "mid": peak_mid})
for c, pk in peak_df.groupby("chrom"):
order = np.argsort(pk["mid"].to_numpy())
mids, cols = pk["mid"].to_numpy()[order], pk.index.to_numpy()[order]
for gi_ in np.flatnonzero(gene_names.isin(tss.index[tss["chrom"] == c])):
t = tss.loc[gene_names[gi_], "tss"]
lo, hi = np.searchsorted(mids, t - CIS_WINDOW_BP), np.searchsorted(mids, t + CIS_WINDOW_BP, side="right")
g_idx += [gi_] * (hi - lo)
p_idx += list(cols[lo:hi])
dist_bp += list(mids[lo:hi] - t)
pairs = pd.DataFrame({"gene": gene_names[g_idx], "peak": peak_names[p_idx], "g": g_idx, "p": p_idx,
"distance_kb": np.asarray(dist_bp, dtype=float) / 1e3})
print(f"{len(pairs):,} candidate pairs: {pairs['gene'].nunique():,} genes, {pairs['peak'].nunique():,} peaks")
# --- correlation across pseudobulks, with an accessibility-matched background on other chromosomes
standardize = lambda X: np.nan_to_num((X - X.mean(0)) / (X.std(0) + 1e-12))
genes_used = np.unique(pairs["g"])
gene_chrom = tss.reindex(gene_names[genes_used])["chrom"].to_numpy()
X_rna_tr = dense(train["rna"].X).astype(np.float32)
X_atac_tr = sp.csr_matrix(train["atac"].X)
pcs = PCA(n_components=30, random_state=0).fit_transform(X_rna_tr)
def pseudobulk_link_z(cells, n_groups, seed=0):
"""Correlation z-scores of all candidate pairs over pseudobulks built from `cells` (indices of training cells)."""
groups = MiniBatchKMeans(n_clusters=n_groups, random_state=seed, n_init=3).fit_predict(pcs[cells])
groups = np.unique(groups, return_inverse=True)[1] # drop empty clusters
M = sp.csr_matrix((np.ones(len(groups)), (groups, np.arange(len(groups)))))
M = sp.diags(1.0 / np.asarray(M.sum(1)).ravel()) @ M # row-normalized
Gz, Az = standardize(np.asarray(M @ X_rna_tr[cells])), standardize(np.asarray((M @ X_atac_tr[cells]).todense()))
n_pb = Gz.shape[0]
r = (Gz[:, pairs["g"]] * Az[:, pairs["p"]]).sum(0) / n_pb
acc = np.asarray((M @ X_atac_tr[cells]).mean(0)).ravel()
acc_bin = pd.qcut(pd.Series(acc).rank(method="first"), 10, labels=False).to_numpy()
rng_bg = np.random.default_rng(seed)
bg_mean, bg_sd = np.zeros((len(gene_names), 10)), np.ones((len(gene_names), 10))
for b in range(10):
cols_b = rng_bg.choice(np.flatnonzero(acc_bin == b), size=min(N_BACKGROUND, int((acc_bin == b).sum())), replace=False)
R_bg = Gz[:, genes_used].T @ Az[:, cols_b] / n_pb # genes x background peaks
R_bg = np.where(gene_chrom[:, None] != peak_chrom[cols_b][None, :], R_bg, np.nan) # other chromosomes only
bg_mean[genes_used, b], bg_sd[genes_used, b] = np.nanmean(R_bg, 1), np.nanstd(R_bg, 1) + 1e-6
z = (r - bg_mean[pairs["g"], acc_bin[pairs["p"]]]) / bg_sd[pairs["g"], acc_bin[pairs["p"]]]
return r, z, n_pb
all_cells = np.arange(len(pcs))
pairs["r"], pairs["corr_z"], n_pb = pseudobulk_link_z(all_cells, min(N_PSEUDOBULKS, len(pcs) // 5))
pairs["link"] = pairs["corr_z"] >= LINK_Z
print(f"{n_pb} pseudobulks; reference links (z >= {LINK_Z:g}): {int(pairs['link'].sum()):,} of {len(pairs):,} "
f"candidate pairs (about {len(pairs) * 0.00135:.0f} would pass by chance if z were standard normal)")
train_types = train["rna"].obs["cell_type"].astype(str).to_numpy()
lineage_z = {}
for lin, members in LINEAGES.items():
cells_l = np.flatnonzero(np.isin(train_types, members))
n_groups = min(N_PSEUDOBULKS_LINEAGE, len(cells_l) // 20)
if n_groups < 15:
print(f" {lin}: {len(cells_l)} training cells, too few for within-lineage pseudobulks; skipped")
continue
_, lineage_z[lin], n_l = pseudobulk_link_z(cells_l, n_groups, seed=1)
print(f" {lin}: {len(cells_l):,} cells, {n_l} pseudobulks, {int((lineage_z[lin] >= LINK_Z).sum()):,} links")
pairs["lineage_z"] = np.max(np.column_stack(list(lineage_z.values())), axis=1) if lineage_z else np.nan
pairs["lineage_link"] = pairs["lineage_z"] >= LINK_Z
print(f"within-lineage links (z >= {LINK_Z:g} in at least one of {len(lineage_z)} lineages): "
f"{int(pairs['lineage_link'].sum()):,} (about {len(pairs) * 0.00135 * max(len(lineage_z), 1):.0f} by chance); "
f"overlap with the all-cell links: {int((pairs['link'] & pairs['lineage_link']).sum()):,}")
# --- attention: centered attention averaged over cells where both tokens are present
pair_id = pairs["g"].to_numpy(np.int64) * len(peak_names) + pairs["p"].to_numpy(np.int64)
order_id = np.argsort(pair_id)
sorted_id = pair_id[order_id]
att_sum, att_n = np.zeros(len(pairs)), np.zeros(len(pairs))
for c0 in range(0, len(R), 200):
ids = (GI[c0:c0 + 200].astype(np.int64)[:, :, None] * len(peak_names) + PI[c0:c0 + 200][:, None, :]).ravel()
vals = R[c0:c0 + 200].ravel()
pos = np.clip(np.searchsorted(sorted_id, ids), 0, len(sorted_id) - 1)
hit = sorted_id[pos] == ids
att_sum += np.bincount(order_id[pos[hit]], weights=vals[hit], minlength=len(pairs))
att_n += np.bincount(order_id[pos[hit]], minlength=len(pairs))
pairs["attention"] = np.where(att_n >= MIN_COPRESENT, att_sum / np.maximum(att_n, 1), np.nan)
pairs["cells_copresent"] = att_n.astype(int)
# --- attribution: input x gradient of predicted gene expression with respect to peak inputs
attr_model = models[ATTRIBUTION_MODEL].eval()
sel_cells = np.random.default_rng(1).choice(test["atac"].n_obs, min(N_ATTRIBUTION_CELLS, test["atac"].n_obs), replace=False)
x_attr = torch.as_tensor(dense(test["atac"].X[sel_cells]), dtype=torch.float32, device=device).requires_grad_(True)
mu_attr, _ = attr_model.encode_modalities({"atac": x_attr})
rna_hat = attr_model.decode_modalities(mu_attr["atac"])["rna"]
if isinstance(rna_hat, dict):
rna_hat = rna_hat.get("mean", rna_hat.get("mu", next(iter(rna_hat.values()))))
attribution = np.full(len(pairs), np.nan)
by_gene = pairs.groupby("g").indices
p_all = pairs["p"].to_numpy()
for k_, (g_, rows_) in enumerate(by_gene.items()):
grad = torch.autograd.grad(rna_hat[:, g_].sum(), x_attr, retain_graph=k_ < len(by_gene) - 1)[0]
ixg = (grad * x_attr).mean(0).detach().cpu().numpy()
attribution[rows_] = ixg[p_all[rows_]]
pairs["attribution"] = attribution
del x_attr, mu_attr, rna_hat
pairs[["distance_kb", "r", "corr_z", "attention", "attribution"]].describe().round(3)
23,839 candidate pairs: 1,859 genes, 20,000 peaks
200 pseudobulks; reference links (z >= 3): 1,214 of 23,839 candidate pairs (about 32 would pass by chance if z were standard normal)
T/NK: 4,480 cells, 60 pseudobulks, 527 links
B: 587 cells, 29 pseudobulks, 253 links
myeloid: 2,618 cells, 60 pseudobulks, 252 links
within-lineage links (z >= 3 in at least one of 3 lineages): 1,016 (about 97 by chance); overlap with the all-cell links: 215
| distance_kb | r | corr_z | attention | attribution | |
|---|---|---|---|---|---|
| count | 23839.000 | 23839.000 | 23839.000 | 1008.000 | 23839.000 |
| mean | 0.360 | 0.157 | 0.660 | -0.000 | 0.000 |
| std | 54.987 | 0.248 | 1.278 | 0.005 | 0.000 |
| min | -99.991 | -0.758 | -2.963 | -0.058 | -0.002 |
| 25% | -44.673 | -0.029 | -0.314 | -0.001 | -0.000 |
| 50% | 0.066 | 0.098 | 0.503 | -0.000 | 0.000 |
| 75% | 45.375 | 0.319 | 1.382 | 0.001 | 0.000 |
| max | 99.997 | 0.951 | 15.821 | 0.064 | 0.012 |
How well does each score recover the reference links?#
AUROC for separating reference links from the other candidate pairs, over all pairs a score is defined for, and averaged within distance bins (0–10, 10–30, 30–60 and 60–100 kb, weighted by pairs). 0.5 is chance. The “within distance” column is the one that says whether a score knows more than distance does.
from sklearn.metrics import roc_auc_score
pairs["distance (closer)"] = -pairs["distance_kb"].abs()
dbins = pd.cut(pairs["distance_kb"].abs(), [-0.1, 10, 30, 60, CIS_WINDOW_BP / 1e3])
scorers = ["distance (closer)", "attention", "attribution"]
def auroc(y, sc_):
ok = np.isfinite(sc_)
return roc_auc_score(y[ok], sc_[ok]) if ok.sum() > 10 and 0 < y[ok].sum() < ok.sum() else np.nan
rows = {}
for ref, name in [(r_, n_) for r_ in ("link", "lineage_link") for n_ in scorers]:
if ref == "lineage_link" and not pairs["lineage_link"].any():
continue
y, sc_ = pairs[ref].to_numpy(), pairs[name].to_numpy(dtype=float)
within, weights = [], []
for b in dbins.cat.categories:
m = (dbins == b).to_numpy()
a = auroc(y[m], sc_[m])
if np.isfinite(a):
within.append(a)
weights.append(int((m & np.isfinite(sc_)).sum()))
ok = np.isfinite(sc_)
rows[("all cells" if ref == "link" else "within lineage", name)] = {"pairs scored": int(ok.sum()), "links among them": int((y & ok).sum()),
"AUROC, all pairs": auroc(y, sc_),
"AUROC, within distance bins": np.average(within, weights=weights) if within else np.nan,
"Spearman with its reference z": spearmanr(sc_[ok], pairs["corr_z" if ref == "link" else "lineage_z"]
.to_numpy()[ok]).statistic if ok.sum() > 10 else np.nan}
link_bench = pd.DataFrame(rows).T
link_bench.index.names = ["reference links", "score"]
display(link_bench.round(3))
fig, axes = plt.subplots(1, 2, figsize=(10, 3.2))
frac = pairs.groupby(pd.cut(pairs["distance_kb"], np.linspace(-CIS_WINDOW_BP / 1e3, CIS_WINDOW_BP / 1e3, 21)),
observed=False)["link"].mean()
axes[0].bar(range(len(frac)), frac.to_numpy(), color="0.5")
axes[0].set_xticks([0, len(frac) // 2, len(frac) - 1], [f"-{CIS_WINDOW_BP // 1000}", "0", f"+{CIS_WINDOW_BP // 1000}"])
axes[0].set(xlabel="peak position relative to TSS (kb)", ylabel="fraction of pairs linked",
title="Reference links by distance from the TSS")
link_bench["AUROC, within distance bins"].unstack(0).plot.barh(ax=axes[1], width=0.75)
axes[1].axvline(0.5, c="k", lw=0.6, ls="--")
axes[1].set(xlabel="AUROC within distance bins", xlim=(0.3, 1.0), title="Which score recovers the links?")
axes[1].legend(frameon=False, fontsize=7, loc="lower right")
plt.tight_layout()
plt.show()
| pairs scored | links among them | AUROC, all pairs | AUROC, within distance bins | Spearman with its reference z | ||
|---|---|---|---|---|---|---|
| reference links | score | |||||
| all cells | distance (closer) | 23839.0 | 1214.0 | 0.611 | 0.547 | 0.191 |
| attention | 1008.0 | 182.0 | 0.488 | 0.493 | -0.038 | |
| attribution | 23839.0 | 1214.0 | 0.673 | 0.653 | 0.238 | |
| within lineage | distance (closer) | 23839.0 | 1016.0 | 0.559 | 0.493 | 0.124 |
| attention | 1008.0 | 74.0 | 0.588 | 0.592 | -0.037 | |
| attribution | 23839.0 | 1016.0 | 0.611 | 0.601 | 0.222 |
Does a score point to the promoter?#
For genes with a promoter peak (within 1 kb of the TSS) among at least three candidates, the fraction whose highest-scoring peak is a promoter peak, next to the fraction expected if the top peak were picked at random. Promoters are the most reliable regulatory element of a gene, so a score that finds them is doing something sensible; distance finds them by construction and is left out.
pairs["promoter"] = pairs["distance_kb"].abs() <= 1.0
prom_rows = {}
for name in ["corr_z", "lineage_z", "attention", "attribution"]:
hits, expect, n_g = 0, 0.0, 0
for g_, d_ in pairs.groupby("gene"):
d_ = d_[np.isfinite(d_[name].to_numpy(dtype=float))]
if len(d_) < 3 or not d_["promoter"].any():
continue
hits += bool(d_.loc[d_[name].astype(float).idxmax(), "promoter"])
expect += d_["promoter"].mean()
n_g += 1
prom_rows[name] = {"genes": n_g, "top peak is a promoter": hits / max(n_g, 1), "expected at random": expect / max(n_g, 1)}
display(pd.DataFrame(prom_rows).T.round(3))
markers_check = [g for g in ["CD79A", "MS4A1", "CD14", "LYZ", "FCGR3A", "NKG7", "GNLY", "IL7R", "CD8A", "SPI1", "LEF1", "TBX21"]
if g in set(pairs["gene"])]
if markers_check:
display(pairs[pairs["gene"].isin(markers_check)].groupby("gene")[["link", "lineage_link"]].sum()
.rename(columns={"link": "links (all cells)", "lineage_link": "links (within lineage)"}).astype(int))
| genes | top peak is a promoter | expected at random | |
|---|---|---|---|
| corr_z | 1002.0 | 0.215 | 0.106 |
| lineage_z | 1002.0 | 0.157 | 0.106 |
| attention | 11.0 | 0.364 | 0.227 |
| attribution | 1002.0 | 0.333 | 0.106 |
| links (all cells) | links (within lineage) | |
|---|---|---|
| gene | ||
| CD14 | 0 | 0 |
| CD79A | 0 | 0 |
| CD8A | 11 | 12 |
| FCGR3A | 1 | 0 |
| GNLY | 4 | 0 |
| IL7R | 2 | 1 |
| LEF1 | 6 | 0 |
| LYZ | 0 | 0 |
| MS4A1 | 3 | 0 |
| NKG7 | 2 | 0 |
| SPI1 | 0 | 1 |
| TBX21 | 9 | 0 |
A locus, three ways#
For the gene with the most reference links, every candidate peak is drawn as an arc from the TSS (orange triangle), once per score: correlation z (top), attention (middle) and attribution (bottom). Arc height follows the size of the score and color its sign (red positive, blue negative), so the three panels can be compared by eye: a model score that agrees with the data draws its tallest red arcs where the correlation panel does.
counts_link = pairs[pairs["link"]].groupby("gene").size()
focus = counts_link.idxmax() if len(counts_link) else pairs.groupby("gene").size().idxmax()
loc = pairs[pairs["gene"] == focus].sort_values("distance_kb")
fig, axes = plt.subplots(3, 1, figsize=(8, 5.4), sharex=True)
t_arc = np.linspace(0, np.pi, 40)
for ax, (name, label) in zip(axes, [("corr_z", "correlation z"), ("attention", "attention"), ("attribution", "attribution")]):
v = loc[name].to_numpy(dtype=float)
ok = np.isfinite(v)
vmax = np.nanmax(np.abs(v)) if ok.any() else 1.0
for x_, val in zip(loc["distance_kb"].to_numpy()[ok], v[ok]):
ax.plot(x_ / 2 + (x_ / 2) * np.cos(t_arc), abs(val) / (vmax + 1e-12) * np.sin(t_arc),
color=plt.cm.RdBu_r(0.5 + 0.5 * val / (vmax + 1e-12)), lw=1.3)
ax.plot(0, 0, marker="^", c="#e67e22", ms=9)
ax.scatter(loc["distance_kb"], np.zeros(len(loc)), marker="|", c="k", s=60)
ax.set_ylim(-0.08, 1.1)
ax.set_yticks([])
ax.set_ylabel(label, fontsize=8)
if not ok.any():
ax.text(0.5, 0.5, "no pairs scored", transform=ax.transAxes, ha="center", fontsize=8)
axes[0].set_title(f"{focus}: candidate peaks within {CIS_WINDOW_BP // 1000} kb "
f"({int(loc['link'].sum())} reference links)", fontsize=9)
axes[-1].set_xlabel("kb from TSS")
plt.tight_layout()
plt.show()
out_dir = Path(OUTPUT_DIR)
out_dir.mkdir(parents=True, exist_ok=True)
pairs.drop(columns=["g", "p"]).to_csv(out_dir / "peak_gene_links.csv.gz", index=False)
print("written:", out_dir / "peak_gene_links.csv.gz")
written: univi_outputs\peak_gene_links.csv.gz
Reading the benchmark. A model score that beats distance within distance bins has learned something about which peaks go with which gene beyond proximity; one that does not is, at best, rediscovering that near peaks matter. UniVI’s objective reconstructs each modality from one pooled embedding and never rewards a gene token for attending to its own regulatory peaks, so pair-specific links are not something it is trained to produce; the benchmark measures how much of them emerges anyway. The reference links come from these same cells and are associations, so agreement with them is necessary but not sufficient. Independent evidence, such as blood eQTLs or published enhancer–gene maps, would be the stronger test.
6. Saving and restoring a transformer model#
save_reference writes the configuration (including the tokenizer options) and the weights;
load_reference rebuilds the model. The check confirms that the restored model reproduces the fused
embedding exactly. (For the position model, call attach_peak_coordinates again after loading: peak
coordinates are data and are not stored with the weights.)
out_dir = Path(OUTPUT_DIR)
out_dir.mkdir(parents=True, exist_ok=True)
save_reference(out_dir / "fused_transformer_reference", fused)
loaded = load_reference(out_dir / "fused_transformer_reference")
restored = (loaded[0] if isinstance(loaded, tuple) else getattr(loaded, "model", loaded)).to(device).eval()
same = np.allclose(encode_fused_adata_pair(restored, {"rna": test["rna"], "atac": test["atac"]}, device=device,
write_to_adatas=False)["mu"], fused_z["fused transformer (v2)"], atol=1e-5)
print("restored model reproduces the fused embedding:", same)
restored model reproduces the fused embedding: True
Notes and limitations#
Objective vs. architecture. Compare “fused transformer (v2)” with “MLP (v2)”, and “transformer (v1)” with “MLP (v1)”; comparing across objectives mixes two changes.
Tokens drop information. Each cell contributes only its top
N_TOKENS_*features to a transformer encoder (128 of 20,000 peaks is 0.6%); the MLP sees every feature. This is the most likely reason a transformer trails the MLP here, and the first setting to vary. Larger token counts cost memory quadratically in the fused model.Peak set.
ATAC_FEATURES = "near_genes"keeps peaks near modeled genes, which the peak–gene analyses need;"most_accessible"keeps the genome-wide most accessible peaks, which may separate cell types better. Compare the Section 3 table between the two.Fused encoders. With UniVI 1.1 and weight decay, a fused model’s per-modality encoders could collapse (Section 2); use UniVI ≥ 1.2.0, and keep the spread check.
Saving.
save_reference/load_referenceround-trip transformer models, including tokenizer options (Section 6); attach peak coordinates again after loading a position model.Genomic coordinates. Compare “transformer + position (v1)” with “transformer (v1)” to see what knowing where a peak lies adds. (In UniVI 1.1 the coordinate embedding ignored the position; fixed in 1.2.0.)