Custom modalities and likelihoods

Custom modalities and likelihoods#

Open in Colab

UniVI is not limited to RNA + ATAC or RNA + protein. Any number of modalities can be combined, each with its own encoder, decoder and likelihood. This notebook shows how to wire up new assays using small synthetic data (so it runs anywhere in about a minute), covering:

  • three modalities at once

  • a negative binomial likelihood on raw counts

  • a beta-binomial likelihood on methylation-style successes and coverage, which needs extra reconstruction targets

  • a learned gating network that weights modalities per cell

  • an (experimental) transformer encoder for one modality

The numbers produced here only demonstrate the mechanics. For real tri-modal analyses see the TEA-seq and scNMT-seq sections of the paper reproduction pages.

import sys

if "google.colab" in sys.modules:
    %pip install -q "univi[tutorials]>=1.0" "pandas==2.2.3"

import anndata as ad
import numpy as np
import pandas as pd
import scipy.sparse as sp
import torch

from univi import ModalityConfig, TrainingConfig, UniVIConfig, UniVIMultiModalVAE, UniVITrainer
from univi.config import TokenizerConfig, TransformerConfig
from univi.evaluation import (compute_foscttm, cross_modal_predict, encode_adata, encode_fused_adata_pair,
                              encode_moe_gates_from_tensors)
from univi.utils.seed import set_seed
from univi.workflows import make_loader

set_seed(0)
device = "cuda" if torch.cuda.is_available() else "cpu"

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 = 60
N_CELLS = 2000

Synthetic tri-modal data#

Three cell states drive all modalities: RNA counts (300 genes), protein counts (20 antibodies), and a methylation-like assay where each of 150 regions has a number of methylated reads (successes) out of a variable number of covering reads (coverage).

rng = np.random.default_rng(0)
state = rng.integers(0, 3, N_CELLS)
obs = pd.DataFrame({"state": pd.Categorical(state.astype(str))}, index=[f"cell{i}" for i in range(N_CELLS)])

def programs(n_features, scale):
    return rng.normal(0, scale, (3, n_features))[state]

rna_counts = rng.poisson(np.exp(0.5 + programs(300, 1.0))).astype(np.float32)
adt_counts = rng.poisson(np.exp(3.0 + programs(20, 0.8))).astype(np.float32)
coverage = rng.poisson(8, (N_CELLS, 150)).astype(np.float32)
p_meth = 1 / (1 + np.exp(-programs(150, 1.5)))
successes = rng.binomial(coverage.astype(int), p_meth).astype(np.float32)

rna = ad.AnnData(sp.csr_matrix(rna_counts), obs=obs.copy())
adt = ad.AnnData(adt_counts, obs=obs.copy())
clr = np.log1p(adt_counts)
adt.X = (clr - clr.mean(1, keepdims=True)).astype(np.float32)        # CLR for a Gaussian likelihood

meth = ad.AnnData(np.divide(successes, coverage, out=np.full_like(successes, 0.5), where=coverage > 0),
                  obs=obs.copy())                                    # .X: methylated fraction (encoder input)
meth.layers["successes"], meth.layers["coverage"] = successes, coverage

train_idx, val_idx = np.arange(0, int(0.9 * N_CELLS)), np.arange(int(0.9 * N_CELLS), N_CELLS)
subset = lambda d, idx: {k: v[idx].copy() for k, v in d.items()}
data = {"rna": rna, "adt": adt, "meth": meth}

Choosing a likelihood#

The decoder likelihood should match what the modality’s .X (or reconstruction targets) contains:

data in the model input

likelihood

normalized / log / z-scored / CLR / LSI values

"gaussian"

raw counts

"nb", "zinb", "poisson"

binary (e.g. binarized peaks)

"bernoulli"

proportions in (0, 1)

"beta"

successes out of trials (methylation, allele counts)

"binomial", "beta_binomial" + reconstruction targets

integer class codes

"categorical"

The encoder receives the same .X; count models do not log-transform inputs internally. Gaussian likelihoods on normalized inputs gave the best cross-modal alignment in the paper’s integration benchmarks, while count likelihoods are useful when you want decoders that output counts.

Beta-binomial modalities read their successes and trials from layers named in recon_targets_spec:

recon_targets = {"meth": {"successes_layer": "successes", "total_count_layer": "coverage"}}

cfg = UniVIConfig(
    latent_dim=8, beta=1.0, gamma=2.0,
    modalities=[
        ModalityConfig("rna", 300, [128, 64], [64, 128], likelihood="nb"),
        ModalityConfig("adt", 20, [32], [32], likelihood="gaussian", recon_weight=2.0),
        ModalityConfig("meth", 150, [64], [64], likelihood="beta_binomial"),
    ],
)
model = UniVIMultiModalVAE(cfg, loss_mode="v1", v1_recon="avg", normalize_v1_terms=True)
UniVITrainer(model,
             make_loader(subset(data, train_idx), batch_size=128, shuffle=True, drop_last=True,
                         recon_targets_spec=recon_targets),
             make_loader(subset(data, val_idx), batch_size=512, recon_targets_spec=recon_targets),
             TrainingConfig(n_epochs=N_EPOCHS, lr=1e-3, device=device, log_every=20)).fit();

[2026-09-21 01:04:17,892] [UniVITrainer] [INFO] TrainingConfig:
[2026-09-21 01:04:17,893] [UniVITrainer] [INFO]   n_epochs: 60
[2026-09-21 01:04:17,893] [UniVITrainer] [INFO]   batch_size: 256
[2026-09-21 01:04:17,893] [UniVITrainer] [INFO]   lr: 0.001
[2026-09-21 01:04:17,894] [UniVITrainer] [INFO]   weight_decay: 0.0
[2026-09-21 01:04:17,894] [UniVITrainer] [INFO]   device: 'cuda'
[2026-09-21 01:04:17,894] [UniVITrainer] [INFO]   log_every: 20
[2026-09-21 01:04:17,895] [UniVITrainer] [INFO]   grad_clip: None
[2026-09-21 01:04:17,895] [UniVITrainer] [INFO]   num_workers: 0
[2026-09-21 01:04:17,895] [UniVITrainer] [INFO]   seed: 0
[2026-09-21 01:04:17,896] [UniVITrainer] [INFO]   early_stopping: False
[2026-09-21 01:04:17,896] [UniVITrainer] [INFO]   patience: 20
[2026-09-21 01:04:17,897] [UniVITrainer] [INFO]   min_delta: 0.0
[2026-09-21 01:04:17,897] [UniVITrainer] [INFO]   best_epoch_warmup: 0
Training UniVI:   0%|          | 0/60 [00:00<?, ?it/s]
[2026-09-21 01:04:19,884] [UniVITrainer] [INFO] [Epoch 001] Train loss=417.3766 (beta=1.000, gamma=2.000)
[2026-09-21 01:04:19,943] [UniVITrainer] [INFO] [Epoch 001] Val loss=392.5763 (beta=1.000, gamma=2.000)
Training UniVI:   0%|          | 0/60 [00:02<?, ?it/s, beta=1.000, gamma=2.000, train_loss=417.3766, val_loss=392.5763]
[2026-09-21 01:04:19,949] [UniVITrainer] [INFO] [Epoch 001] New best val loss: 392.5763
Training UniVI:   2%|▏         | 1/60 [00:02<02:00,  2.05s/it, beta=1.000, gamma=2.000, train_loss=417.3766, val_loss=392.5763]
Training UniVI:   2%|▏         | 1/60 [00:03<02:00,  2.05s/it, beta=1.000, gamma=2.000, train_loss=363.9108, val_loss=336.5546]
[2026-09-21 01:04:21,257] [UniVITrainer] [INFO] [Epoch 002] New best val loss: 336.5546
Training UniVI:   3%|▎         | 2/60 [00:03<01:33,  1.61s/it, beta=1.000, gamma=2.000, train_loss=363.9108, val_loss=336.5546]
Training UniVI:   3%|▎         | 2/60 [00:04<01:33,  1.61s/it, beta=1.000, gamma=2.000, train_loss=328.5976, val_loss=321.1845]
[2026-09-21 01:04:22,499] [UniVITrainer] [INFO] [Epoch 003] New best val loss: 321.1845
Training UniVI:   5%|▌         | 3/60 [00:04<01:22,  1.44s/it, beta=1.000, gamma=2.000, train_loss=328.5976, val_loss=321.1845]
Training UniVI:   5%|▌         | 3/60 [00:05<01:22,  1.44s/it, beta=1.000, gamma=2.000, train_loss=316.7092, val_loss=311.2253]
[2026-09-21 01:04:23,737] [UniVITrainer] [INFO] [Epoch 004] New best val loss: 311.2253
Training UniVI:   7%|▋         | 4/60 [00:05<01:16,  1.36s/it, beta=1.000, gamma=2.000, train_loss=316.7092, val_loss=311.2253]
Training UniVI:   7%|▋         | 4/60 [00:07<01:16,  1.36s/it, beta=1.000, gamma=2.000, train_loss=308.0895, val_loss=302.8649]
[2026-09-21 01:04:24,947] [UniVITrainer] [INFO] [Epoch 005] New best val loss: 302.8649
Training UniVI:   8%|▊         | 5/60 [00:07<01:11,  1.31s/it, beta=1.000, gamma=2.000, train_loss=308.0895, val_loss=302.8649]
Training UniVI:   8%|▊         | 5/60 [00:08<01:11,  1.31s/it, beta=1.000, gamma=2.000, train_loss=299.6281, val_loss=294.0128]
[2026-09-21 01:04:26,181] [UniVITrainer] [INFO] [Epoch 006] New best val loss: 294.0128
Training UniVI:  10%|█         | 6/60 [00:08<01:09,  1.28s/it, beta=1.000, gamma=2.000, train_loss=299.6281, val_loss=294.0128]
Training UniVI:  10%|█         | 6/60 [00:09<01:09,  1.28s/it, beta=1.000, gamma=2.000, train_loss=291.5870, val_loss=286.8184]
[2026-09-21 01:04:27,407] [UniVITrainer] [INFO] [Epoch 007] New best val loss: 286.8184
Training UniVI:  12%|█▏        | 7/60 [00:09<01:06,  1.26s/it, beta=1.000, gamma=2.000, train_loss=291.5870, val_loss=286.8184]
Training UniVI:  12%|█▏        | 7/60 [00:10<01:06,  1.26s/it, beta=1.000, gamma=2.000, train_loss=284.7932, val_loss=280.9547]
[2026-09-21 01:04:28,655] [UniVITrainer] [INFO] [Epoch 008] New best val loss: 280.9547
Training UniVI:  13%|█▎        | 8/60 [00:10<01:05,  1.26s/it, beta=1.000, gamma=2.000, train_loss=284.7932, val_loss=280.9547]
Training UniVI:  13%|█▎        | 8/60 [00:11<01:05,  1.26s/it, beta=1.000, gamma=2.000, train_loss=280.3788, val_loss=277.5530]
[2026-09-21 01:04:29,892] [UniVITrainer] [INFO] [Epoch 009] New best val loss: 277.5530
Training UniVI:  15%|█▌        | 9/60 [00:11<01:03,  1.25s/it, beta=1.000, gamma=2.000, train_loss=280.3788, val_loss=277.5530]
Training UniVI:  15%|█▌        | 9/60 [00:13<01:03,  1.25s/it, beta=1.000, gamma=2.000, train_loss=277.7192, val_loss=274.8940]
[2026-09-21 01:04:31,152] [UniVITrainer] [INFO] [Epoch 010] New best val loss: 274.8940
Training UniVI:  17%|█▋        | 10/60 [00:13<01:02,  1.25s/it, beta=1.000, gamma=2.000, train_loss=277.7192, val_loss=274.8940]
Training UniVI:  17%|█▋        | 10/60 [00:14<01:02,  1.25s/it, beta=1.000, gamma=2.000, train_loss=275.2480, val_loss=273.3197]
[2026-09-21 01:04:32,386] [UniVITrainer] [INFO] [Epoch 011] New best val loss: 273.3197
Training UniVI:  18%|█▊        | 11/60 [00:14<01:01,  1.25s/it, beta=1.000, gamma=2.000, train_loss=275.2480, val_loss=273.3197]
Training UniVI:  18%|█▊        | 11/60 [00:15<01:01,  1.25s/it, beta=1.000, gamma=2.000, train_loss=273.3574, val_loss=271.8142]
[2026-09-21 01:04:33,636] [UniVITrainer] [INFO] [Epoch 012] New best val loss: 271.8142
Training UniVI:  20%|██        | 12/60 [00:15<00:59,  1.25s/it, beta=1.000, gamma=2.000, train_loss=273.3574, val_loss=271.8142]
Training UniVI:  20%|██        | 12/60 [00:16<00:59,  1.25s/it, beta=1.000, gamma=2.000, train_loss=271.7868, val_loss=270.3079]
[2026-09-21 01:04:34,868] [UniVITrainer] [INFO] [Epoch 013] New best val loss: 270.3079
Training UniVI:  22%|██▏       | 13/60 [00:16<00:58,  1.24s/it, beta=1.000, gamma=2.000, train_loss=271.7868, val_loss=270.3079]
Training UniVI:  22%|██▏       | 13/60 [00:18<00:58,  1.24s/it, beta=1.000, gamma=2.000, train_loss=270.8283, val_loss=269.3782]
[2026-09-21 01:04:36,099] [UniVITrainer] [INFO] [Epoch 014] New best val loss: 269.3782
Training UniVI:  23%|██▎       | 14/60 [00:18<00:57,  1.24s/it, beta=1.000, gamma=2.000, train_loss=270.8283, val_loss=269.3782]
Training UniVI:  23%|██▎       | 14/60 [00:19<00:57,  1.24s/it, beta=1.000, gamma=2.000, train_loss=270.0236, val_loss=268.6918]
[2026-09-21 01:04:37,325] [UniVITrainer] [INFO] [Epoch 015] New best val loss: 268.6918
Training UniVI:  25%|██▌       | 15/60 [00:19<00:55,  1.24s/it, beta=1.000, gamma=2.000, train_loss=270.0236, val_loss=268.6918]
Training UniVI:  25%|██▌       | 15/60 [00:20<00:55,  1.24s/it, beta=1.000, gamma=2.000, train_loss=269.1968, val_loss=267.4535]
[2026-09-21 01:04:38,572] [UniVITrainer] [INFO] [Epoch 016] New best val loss: 267.4535
Training UniVI:  27%|██▋       | 16/60 [00:20<00:54,  1.24s/it, beta=1.000, gamma=2.000, train_loss=269.1968, val_loss=267.4535]
Training UniVI:  27%|██▋       | 16/60 [00:21<00:54,  1.24s/it, beta=1.000, gamma=2.000, train_loss=268.4441, val_loss=267.0032]
[2026-09-21 01:04:39,842] [UniVITrainer] [INFO] [Epoch 017] New best val loss: 267.0032
Training UniVI:  28%|██▊       | 17/60 [00:21<00:53,  1.25s/it, beta=1.000, gamma=2.000, train_loss=268.4441, val_loss=267.0032]
Training UniVI:  28%|██▊       | 17/60 [00:23<00:53,  1.25s/it, beta=1.000, gamma=2.000, train_loss=267.5517, val_loss=266.3224]
[2026-09-21 01:04:41,061] [UniVITrainer] [INFO] [Epoch 018] New best val loss: 266.3224
Training UniVI:  30%|███       | 18/60 [00:23<00:52,  1.24s/it, beta=1.000, gamma=2.000, train_loss=267.5517, val_loss=266.3224]
Training UniVI:  30%|███       | 18/60 [00:24<00:52,  1.24s/it, beta=1.000, gamma=2.000, train_loss=267.1325, val_loss=265.7885]
[2026-09-21 01:04:42,295] [UniVITrainer] [INFO] [Epoch 019] New best val loss: 265.7885
Training UniVI:  32%|███▏      | 19/60 [00:24<00:50,  1.24s/it, beta=1.000, gamma=2.000, train_loss=267.1325, val_loss=265.7885]
[2026-09-21 01:04:43,461] [UniVITrainer] [INFO] [Epoch 020] Train loss=266.3719 (beta=1.000, gamma=2.000)
[2026-09-21 01:04:43,519] [UniVITrainer] [INFO] [Epoch 020] Val loss=265.3527 (beta=1.000, gamma=2.000)
Training UniVI:  32%|███▏      | 19/60 [00:25<00:50,  1.24s/it, beta=1.000, gamma=2.000, train_loss=266.3719, val_loss=265.3527]
[2026-09-21 01:04:43,523] [UniVITrainer] [INFO] [Epoch 020] New best val loss: 265.3527
Training UniVI:  33%|███▎      | 20/60 [00:25<00:49,  1.23s/it, beta=1.000, gamma=2.000, train_loss=266.3719, val_loss=265.3527]
Training UniVI:  33%|███▎      | 20/60 [00:26<00:49,  1.23s/it, beta=1.000, gamma=2.000, train_loss=265.8748, val_loss=264.4891]
[2026-09-21 01:04:44,752] [UniVITrainer] [INFO] [Epoch 021] New best val loss: 264.4891
Training UniVI:  35%|███▌      | 21/60 [00:26<00:48,  1.23s/it, beta=1.000, gamma=2.000, train_loss=265.8748, val_loss=264.4891]
Training UniVI:  35%|███▌      | 21/60 [00:28<00:48,  1.23s/it, beta=1.000, gamma=2.000, train_loss=265.2313, val_loss=264.0228]
[2026-09-21 01:04:45,994] [UniVITrainer] [INFO] [Epoch 022] New best val loss: 264.0228
Training UniVI:  37%|███▋      | 22/60 [00:28<00:46,  1.24s/it, beta=1.000, gamma=2.000, train_loss=265.2313, val_loss=264.0228]
Training UniVI:  37%|███▋      | 22/60 [00:29<00:46,  1.24s/it, beta=1.000, gamma=2.000, train_loss=264.9304, val_loss=263.7187]
[2026-09-21 01:04:47,220] [UniVITrainer] [INFO] [Epoch 023] New best val loss: 263.7187
Training UniVI:  38%|███▊      | 23/60 [00:29<00:45,  1.23s/it, beta=1.000, gamma=2.000, train_loss=264.9304, val_loss=263.7187]
Training UniVI:  38%|███▊      | 23/60 [00:30<00:45,  1.23s/it, beta=1.000, gamma=2.000, train_loss=264.5325, val_loss=263.3748]
[2026-09-21 01:04:48,466] [UniVITrainer] [INFO] [Epoch 024] New best val loss: 263.3748
Training UniVI:  40%|████      | 24/60 [00:30<00:44,  1.24s/it, beta=1.000, gamma=2.000, train_loss=264.5325, val_loss=263.3748]
Training UniVI:  40%|████      | 24/60 [00:31<00:44,  1.24s/it, beta=1.000, gamma=2.000, train_loss=264.2474, val_loss=263.0778]
[2026-09-21 01:04:49,686] [UniVITrainer] [INFO] [Epoch 025] New best val loss: 263.0778
Training UniVI:  42%|████▏     | 25/60 [00:31<00:43,  1.23s/it, beta=1.000, gamma=2.000, train_loss=264.2474, val_loss=263.0778]
Training UniVI:  42%|████▏     | 25/60 [00:33<00:43,  1.23s/it, beta=1.000, gamma=2.000, train_loss=263.8149, val_loss=262.5035]
[2026-09-21 01:04:50,917] [UniVITrainer] [INFO] [Epoch 026] New best val loss: 262.5035
Training UniVI:  43%|████▎     | 26/60 [00:33<00:41,  1.23s/it, beta=1.000, gamma=2.000, train_loss=263.8149, val_loss=262.5035]
Training UniVI:  43%|████▎     | 26/60 [00:34<00:41,  1.23s/it, beta=1.000, gamma=2.000, train_loss=263.2845, val_loss=262.1435]
[2026-09-21 01:04:52,145] [UniVITrainer] [INFO] [Epoch 027] New best val loss: 262.1435
Training UniVI:  45%|████▌     | 27/60 [00:34<00:40,  1.23s/it, beta=1.000, gamma=2.000, train_loss=263.2845, val_loss=262.1435]
Training UniVI:  45%|████▌     | 27/60 [00:35<00:40,  1.23s/it, beta=1.000, gamma=2.000, train_loss=262.8018, val_loss=261.4669]
[2026-09-21 01:04:53,379] [UniVITrainer] [INFO] [Epoch 028] New best val loss: 261.4669
Training UniVI:  47%|████▋     | 28/60 [00:35<00:39,  1.23s/it, beta=1.000, gamma=2.000, train_loss=262.8018, val_loss=261.4669]
Training UniVI:  47%|████▋     | 28/60 [00:36<00:39,  1.23s/it, beta=1.000, gamma=2.000, train_loss=262.6166, val_loss=261.5208]
Training UniVI:  48%|████▊     | 29/60 [00:36<00:38,  1.23s/it, beta=1.000, gamma=2.000, train_loss=262.6166, val_loss=261.5208]
Training UniVI:  48%|████▊     | 29/60 [00:37<00:38,  1.23s/it, beta=1.000, gamma=2.000, train_loss=262.1044, val_loss=260.8174]
[2026-09-21 01:04:55,811] [UniVITrainer] [INFO] [Epoch 030] New best val loss: 260.8174
Training UniVI:  50%|█████     | 30/60 [00:37<00:36,  1.22s/it, beta=1.000, gamma=2.000, train_loss=262.1044, val_loss=260.8174]
Training UniVI:  50%|█████     | 30/60 [00:39<00:36,  1.22s/it, beta=1.000, gamma=2.000, train_loss=262.0019, val_loss=260.7778]
[2026-09-21 01:04:57,054] [UniVITrainer] [INFO] [Epoch 031] New best val loss: 260.7778
Training UniVI:  52%|█████▏    | 31/60 [00:39<00:35,  1.23s/it, beta=1.000, gamma=2.000, train_loss=262.0019, val_loss=260.7778]
Training UniVI:  52%|█████▏    | 31/60 [00:40<00:35,  1.23s/it, beta=1.000, gamma=2.000, train_loss=261.4991, val_loss=260.4848]
[2026-09-21 01:04:58,298] [UniVITrainer] [INFO] [Epoch 032] New best val loss: 260.4848
Training UniVI:  53%|█████▎    | 32/60 [00:40<00:34,  1.23s/it, beta=1.000, gamma=2.000, train_loss=261.4991, val_loss=260.4848]
Training UniVI:  53%|█████▎    | 32/60 [00:41<00:34,  1.23s/it, beta=1.000, gamma=2.000, train_loss=261.2848, val_loss=260.1663]
[2026-09-21 01:04:59,521] [UniVITrainer] [INFO] [Epoch 033] New best val loss: 260.1663
Training UniVI:  55%|█████▌    | 33/60 [00:41<00:33,  1.23s/it, beta=1.000, gamma=2.000, train_loss=261.2848, val_loss=260.1663]
Training UniVI:  55%|█████▌    | 33/60 [00:42<00:33,  1.23s/it, beta=1.000, gamma=2.000, train_loss=260.9317, val_loss=259.6465]
[2026-09-21 01:05:00,745] [UniVITrainer] [INFO] [Epoch 034] New best val loss: 259.6465
Training UniVI:  57%|█████▋    | 34/60 [00:42<00:31,  1.23s/it, beta=1.000, gamma=2.000, train_loss=260.9317, val_loss=259.6465]
Training UniVI:  57%|█████▋    | 34/60 [00:44<00:31,  1.23s/it, beta=1.000, gamma=2.000, train_loss=260.5823, val_loss=259.5702]
[2026-09-21 01:05:01,984] [UniVITrainer] [INFO] [Epoch 035] New best val loss: 259.5702
Training UniVI:  58%|█████▊    | 35/60 [00:44<00:30,  1.23s/it, beta=1.000, gamma=2.000, train_loss=260.5823, val_loss=259.5702]
Training UniVI:  58%|█████▊    | 35/60 [00:45<00:30,  1.23s/it, beta=1.000, gamma=2.000, train_loss=260.2184, val_loss=259.0290]
[2026-09-21 01:05:03,215] [UniVITrainer] [INFO] [Epoch 036] New best val loss: 259.0290
Training UniVI:  60%|██████    | 36/60 [00:45<00:29,  1.23s/it, beta=1.000, gamma=2.000, train_loss=260.2184, val_loss=259.0290]
Training UniVI:  60%|██████    | 36/60 [00:46<00:29,  1.23s/it, beta=1.000, gamma=2.000, train_loss=260.0142, val_loss=259.0664]
Training UniVI:  62%|██████▏   | 37/60 [00:46<00:28,  1.23s/it, beta=1.000, gamma=2.000, train_loss=260.0142, val_loss=259.0664]
Training UniVI:  62%|██████▏   | 37/60 [00:47<00:28,  1.23s/it, beta=1.000, gamma=2.000, train_loss=259.6912, val_loss=258.5705]
[2026-09-21 01:05:05,651] [UniVITrainer] [INFO] [Epoch 038] New best val loss: 258.5705
Training UniVI:  63%|██████▎   | 38/60 [00:47<00:26,  1.22s/it, beta=1.000, gamma=2.000, train_loss=259.6912, val_loss=258.5705]
Training UniVI:  63%|██████▎   | 38/60 [00:48<00:26,  1.22s/it, beta=1.000, gamma=2.000, train_loss=259.4280, val_loss=258.4071]
[2026-09-21 01:05:06,885] [UniVITrainer] [INFO] [Epoch 039] New best val loss: 258.4071
Training UniVI:  65%|██████▌   | 39/60 [00:48<00:25,  1.23s/it, beta=1.000, gamma=2.000, train_loss=259.4280, val_loss=258.4071]
[2026-09-21 01:05:08,046] [UniVITrainer] [INFO] [Epoch 040] Train loss=259.1716 (beta=1.000, gamma=2.000)
[2026-09-21 01:05:08,105] [UniVITrainer] [INFO] [Epoch 040] Val loss=258.1292 (beta=1.000, gamma=2.000)
Training UniVI:  65%|██████▌   | 39/60 [00:50<00:25,  1.23s/it, beta=1.000, gamma=2.000, train_loss=259.1716, val_loss=258.1292]
[2026-09-21 01:05:08,109] [UniVITrainer] [INFO] [Epoch 040] New best val loss: 258.1292
Training UniVI:  67%|██████▋   | 40/60 [00:50<00:24,  1.23s/it, beta=1.000, gamma=2.000, train_loss=259.1716, val_loss=258.1292]
Training UniVI:  67%|██████▋   | 40/60 [00:51<00:24,  1.23s/it, beta=1.000, gamma=2.000, train_loss=258.7762, val_loss=257.6962]
[2026-09-21 01:05:09,320] [UniVITrainer] [INFO] [Epoch 041] New best val loss: 257.6962
Training UniVI:  68%|██████▊   | 41/60 [00:51<00:23,  1.22s/it, beta=1.000, gamma=2.000, train_loss=258.7762, val_loss=257.6962]
Training UniVI:  68%|██████▊   | 41/60 [00:52<00:23,  1.22s/it, beta=1.000, gamma=2.000, train_loss=258.5505, val_loss=257.4986]
[2026-09-21 01:05:10,564] [UniVITrainer] [INFO] [Epoch 042] New best val loss: 257.4986
Training UniVI:  70%|███████   | 42/60 [00:52<00:22,  1.23s/it, beta=1.000, gamma=2.000, train_loss=258.5505, val_loss=257.4986]
Training UniVI:  70%|███████   | 42/60 [00:53<00:22,  1.23s/it, beta=1.000, gamma=2.000, train_loss=258.2787, val_loss=257.0193]
[2026-09-21 01:05:11,794] [UniVITrainer] [INFO] [Epoch 043] New best val loss: 257.0193
Training UniVI:  72%|███████▏  | 43/60 [00:53<00:20,  1.23s/it, beta=1.000, gamma=2.000, train_loss=258.2787, val_loss=257.0193]
Training UniVI:  72%|███████▏  | 43/60 [00:55<00:20,  1.23s/it, beta=1.000, gamma=2.000, train_loss=258.1718, val_loss=256.9307]
[2026-09-21 01:05:13,030] [UniVITrainer] [INFO] [Epoch 044] New best val loss: 256.9307
Training UniVI:  73%|███████▎  | 44/60 [00:55<00:19,  1.23s/it, beta=1.000, gamma=2.000, train_loss=258.1718, val_loss=256.9307]
Training UniVI:  73%|███████▎  | 44/60 [00:56<00:19,  1.23s/it, beta=1.000, gamma=2.000, train_loss=257.6975, val_loss=256.8028]
[2026-09-21 01:05:14,253] [UniVITrainer] [INFO] [Epoch 045] New best val loss: 256.8028
Training UniVI:  75%|███████▌  | 45/60 [00:56<00:18,  1.23s/it, beta=1.000, gamma=2.000, train_loss=257.6975, val_loss=256.8028]
Training UniVI:  75%|███████▌  | 45/60 [00:57<00:18,  1.23s/it, beta=1.000, gamma=2.000, train_loss=257.7495, val_loss=256.3145]
[2026-09-21 01:05:15,589] [UniVITrainer] [INFO] [Epoch 046] New best val loss: 256.3145
Training UniVI:  77%|███████▋  | 46/60 [00:57<00:17,  1.26s/it, beta=1.000, gamma=2.000, train_loss=257.7495, val_loss=256.3145]
Training UniVI:  77%|███████▋  | 46/60 [00:58<00:17,  1.26s/it, beta=1.000, gamma=2.000, train_loss=257.2932, val_loss=256.1866]
[2026-09-21 01:05:16,788] [UniVITrainer] [INFO] [Epoch 047] New best val loss: 256.1866
Training UniVI:  78%|███████▊  | 47/60 [00:58<00:16,  1.24s/it, beta=1.000, gamma=2.000, train_loss=257.2932, val_loss=256.1866]
Training UniVI:  78%|███████▊  | 47/60 [01:00<00:16,  1.24s/it, beta=1.000, gamma=2.000, train_loss=257.1796, val_loss=255.9613]
[2026-09-21 01:05:18,019] [UniVITrainer] [INFO] [Epoch 048] New best val loss: 255.9613
Training UniVI:  80%|████████  | 48/60 [01:00<00:14,  1.24s/it, beta=1.000, gamma=2.000, train_loss=257.1796, val_loss=255.9613]
Training UniVI:  80%|████████  | 48/60 [01:01<00:14,  1.24s/it, beta=1.000, gamma=2.000, train_loss=257.0728, val_loss=255.7202]
[2026-09-21 01:05:19,258] [UniVITrainer] [INFO] [Epoch 049] New best val loss: 255.7202
Training UniVI:  82%|████████▏ | 49/60 [01:01<00:13,  1.24s/it, beta=1.000, gamma=2.000, train_loss=257.0728, val_loss=255.7202]
Training UniVI:  82%|████████▏ | 49/60 [01:02<00:13,  1.24s/it, beta=1.000, gamma=2.000, train_loss=256.9240, val_loss=255.7982]
Training UniVI:  83%|████████▎ | 50/60 [01:02<00:12,  1.25s/it, beta=1.000, gamma=2.000, train_loss=256.9240, val_loss=255.7982]
Training UniVI:  83%|████████▎ | 50/60 [01:03<00:12,  1.25s/it, beta=1.000, gamma=2.000, train_loss=256.5599, val_loss=255.4579]
[2026-09-21 01:05:21,823] [UniVITrainer] [INFO] [Epoch 051] New best val loss: 255.4579
Training UniVI:  85%|████████▌ | 51/60 [01:03<00:11,  1.26s/it, beta=1.000, gamma=2.000, train_loss=256.5599, val_loss=255.4579]
Training UniVI:  85%|████████▌ | 51/60 [01:05<00:11,  1.26s/it, beta=1.000, gamma=2.000, train_loss=256.1473, val_loss=255.2498]
[2026-09-21 01:05:23,034] [UniVITrainer] [INFO] [Epoch 052] New best val loss: 255.2498
Training UniVI:  87%|████████▋ | 52/60 [01:05<00:09,  1.25s/it, beta=1.000, gamma=2.000, train_loss=256.1473, val_loss=255.2498]
Training UniVI:  87%|████████▋ | 52/60 [01:06<00:09,  1.25s/it, beta=1.000, gamma=2.000, train_loss=255.9736, val_loss=254.9805]
[2026-09-21 01:05:24,267] [UniVITrainer] [INFO] [Epoch 053] New best val loss: 254.9805
Training UniVI:  88%|████████▊ | 53/60 [01:06<00:08,  1.24s/it, beta=1.000, gamma=2.000, train_loss=255.9736, val_loss=254.9805]
Training UniVI:  88%|████████▊ | 53/60 [01:07<00:08,  1.24s/it, beta=1.000, gamma=2.000, train_loss=255.7976, val_loss=254.7744]
[2026-09-21 01:05:25,493] [UniVITrainer] [INFO] [Epoch 054] New best val loss: 254.7744
Training UniVI:  90%|█████████ | 54/60 [01:07<00:07,  1.24s/it, beta=1.000, gamma=2.000, train_loss=255.7976, val_loss=254.7744]
Training UniVI:  90%|█████████ | 54/60 [01:08<00:07,  1.24s/it, beta=1.000, gamma=2.000, train_loss=255.4662, val_loss=254.6248]
[2026-09-21 01:05:26,726] [UniVITrainer] [INFO] [Epoch 055] New best val loss: 254.6248
Training UniVI:  92%|█████████▏| 55/60 [01:08<00:06,  1.24s/it, beta=1.000, gamma=2.000, train_loss=255.4662, val_loss=254.6248]
Training UniVI:  92%|█████████▏| 55/60 [01:10<00:06,  1.24s/it, beta=1.000, gamma=2.000, train_loss=255.3694, val_loss=254.3430]
[2026-09-21 01:05:27,970] [UniVITrainer] [INFO] [Epoch 056] New best val loss: 254.3430
Training UniVI:  93%|█████████▎| 56/60 [01:10<00:04,  1.24s/it, beta=1.000, gamma=2.000, train_loss=255.3694, val_loss=254.3430]
Training UniVI:  93%|█████████▎| 56/60 [01:11<00:04,  1.24s/it, beta=1.000, gamma=2.000, train_loss=255.1425, val_loss=253.9804]
[2026-09-21 01:05:29,216] [UniVITrainer] [INFO] [Epoch 057] New best val loss: 253.9804
Training UniVI:  95%|█████████▌| 57/60 [01:11<00:03,  1.24s/it, beta=1.000, gamma=2.000, train_loss=255.1425, val_loss=253.9804]
Training UniVI:  95%|█████████▌| 57/60 [01:12<00:03,  1.24s/it, beta=1.000, gamma=2.000, train_loss=255.2595, val_loss=254.0571]
Training UniVI:  97%|█████████▋| 58/60 [01:12<00:02,  1.24s/it, beta=1.000, gamma=2.000, train_loss=255.2595, val_loss=254.0571]
Training UniVI:  97%|█████████▋| 58/60 [01:13<00:02,  1.24s/it, beta=1.000, gamma=2.000, train_loss=254.7769, val_loss=253.7306]
[2026-09-21 01:05:31,664] [UniVITrainer] [INFO] [Epoch 059] New best val loss: 253.7306
Training UniVI:  98%|█████████▊| 59/60 [01:13<00:01,  1.23s/it, beta=1.000, gamma=2.000, train_loss=254.7769, val_loss=253.7306]
[2026-09-21 01:05:32,825] [UniVITrainer] [INFO] [Epoch 060] Train loss=254.3079 (beta=1.000, gamma=2.000)
[2026-09-21 01:05:32,881] [UniVITrainer] [INFO] [Epoch 060] Val loss=253.5034 (beta=1.000, gamma=2.000)
Training UniVI:  98%|█████████▊| 59/60 [01:14<00:01,  1.23s/it, beta=1.000, gamma=2.000, train_loss=254.3079, val_loss=253.5034]
[2026-09-21 01:05:32,885] [UniVITrainer] [INFO] [Epoch 060] New best val loss: 253.5034
Training UniVI: 100%|██████████| 60/60 [01:14<00:00,  1.23s/it, beta=1.000, gamma=2.000, train_loss=254.3079, val_loss=253.5034]
Training UniVI: 100%|██████████| 60/60 [01:14<00:00,  1.25s/it, beta=1.000, gamma=2.000, train_loss=254.3079, val_loss=253.5034]

[2026-09-21 01:05:32,917] [UniVITrainer] [INFO] Restored best model from epoch 60 (val loss=253.5034)

recon_weight rescales a modality’s reconstruction term, which helps when modalities differ greatly in size (here protein has far fewer features than RNA).

Every encoder places cells in the same space, so any pair of modalities can be compared or translated:

val_data = subset(data, val_idx)
z = {m: encode_adata(model, a, modality=m, device=device, latent="modality_mean") for m, a in val_data.items()}
print({f"FOSCTTM {a}-{b}": round(compute_foscttm(z[a], z[b]), 3) for a, b in [("rna", "adt"), ("rna", "meth"), ("adt", "meth")]})

meth_from_rna = cross_modal_predict(model, val_data["rna"], src_mod="rna", tgt_mod="meth", device=device)
print("predicted methylated fraction, first cell:", np.round(meth_from_rna[0, :5], 3))

{'FOSCTTM rna-adt': 0.165, 'FOSCTTM rna-meth': 0.165, 'FOSCTTM adt-meth': 0.175}
predicted methylated fraction, first cell: [0.753 0.107 0.594 0.517 0.522]

For count and beta-binomial decoders, cross_modal_predict returns the decoder mean (expected counts or expected fractions).

Learned per-cell modality weights#

By default the fused posterior weights modalities by their posterior precision. A learned gating network can re-weight them per cell. It is trained only if the fused posterior enters the loss, so pair use_moe_gating=True with v1_recon="moe" (or loss_mode="v2").

gated_cfg = UniVIConfig(latent_dim=8, beta=1.0, gamma=2.0, modalities=cfg.modalities,
                        use_moe_gating=True, moe_gating_hidden=[32])
gated = UniVIMultiModalVAE(gated_cfg, loss_mode="v1", v1_recon="moe", normalize_v1_terms=True)
UniVITrainer(gated,
             make_loader(subset(data, train_idx), batch_size=128, shuffle=True, drop_last=True,
                         recon_targets_spec=recon_targets),
             None, TrainingConfig(n_epochs=N_EPOCHS, lr=1e-3, device=device, log_every=20)).fit();

x = {m: a.X for m, a in val_data.items()}
router = encode_moe_gates_from_tensors(gated, x, device=device, kind="router_x_precision")
precision = encode_moe_gates_from_tensors(gated, x, device=device, kind="effective_precision")
pd.DataFrame({"router x precision": router["per_modality_mean"],
              "precision only": precision["per_modality_mean"]}).round(3)

[2026-09-21 01:05:33,005] [UniVITrainer] [INFO] TrainingConfig:
[2026-09-21 01:05:33,006] [UniVITrainer] [INFO]   n_epochs: 60
[2026-09-21 01:05:33,006] [UniVITrainer] [INFO]   batch_size: 256
[2026-09-21 01:05:33,006] [UniVITrainer] [INFO]   lr: 0.001
[2026-09-21 01:05:33,007] [UniVITrainer] [INFO]   weight_decay: 0.0
[2026-09-21 01:05:33,007] [UniVITrainer] [INFO]   device: 'cuda'
[2026-09-21 01:05:33,007] [UniVITrainer] [INFO]   log_every: 20
[2026-09-21 01:05:33,008] [UniVITrainer] [INFO]   grad_clip: None
[2026-09-21 01:05:33,008] [UniVITrainer] [INFO]   num_workers: 0
[2026-09-21 01:05:33,009] [UniVITrainer] [INFO]   seed: 0
[2026-09-21 01:05:33,009] [UniVITrainer] [INFO]   early_stopping: False
[2026-09-21 01:05:33,009] [UniVITrainer] [INFO]   patience: 20
[2026-09-21 01:05:33,009] [UniVITrainer] [INFO]   min_delta: 0.0
[2026-09-21 01:05:33,010] [UniVITrainer] [INFO]   best_epoch_warmup: 0
Training UniVI:   0%|          | 0/60 [00:00<?, ?it/s]
[2026-09-21 01:05:33,929] [UniVITrainer] [INFO] [Epoch 001] Train loss=417.2287 (beta=1.000, gamma=2.000)
Training UniVI:   0%|          | 0/60 [00:00<?, ?it/s, beta=1.000, gamma=2.000, train_loss=417.2287]
Training UniVI:   2%|▏         | 1/60 [00:00<00:54,  1.09it/s, beta=1.000, gamma=2.000, train_loss=417.2287]
Training UniVI:   2%|▏         | 1/60 [00:01<00:54,  1.09it/s, beta=1.000, gamma=2.000, train_loss=365.1456]
Training UniVI:   3%|▎         | 2/60 [00:01<00:51,  1.12it/s, beta=1.000, gamma=2.000, train_loss=365.1456]
Training UniVI:   3%|▎         | 2/60 [00:02<00:51,  1.12it/s, beta=1.000, gamma=2.000, train_loss=327.4522]
Training UniVI:   5%|▌         | 3/60 [00:02<00:49,  1.14it/s, beta=1.000, gamma=2.000, train_loss=327.4522]
Training UniVI:   5%|▌         | 3/60 [00:03<00:49,  1.14it/s, beta=1.000, gamma=2.000, train_loss=313.8234]
Training UniVI:   7%|▋         | 4/60 [00:03<00:48,  1.15it/s, beta=1.000, gamma=2.000, train_loss=313.8234]
Training UniVI:   7%|▋         | 4/60 [00:04<00:48,  1.15it/s, beta=1.000, gamma=2.000, train_loss=302.1901]
Training UniVI:   8%|▊         | 5/60 [00:04<00:47,  1.15it/s, beta=1.000, gamma=2.000, train_loss=302.1901]
Training UniVI:   8%|▊         | 5/60 [00:05<00:47,  1.15it/s, beta=1.000, gamma=2.000, train_loss=292.9939]
Training UniVI:  10%|█         | 6/60 [00:05<00:46,  1.16it/s, beta=1.000, gamma=2.000, train_loss=292.9939]
Training UniVI:  10%|█         | 6/60 [00:06<00:46,  1.16it/s, beta=1.000, gamma=2.000, train_loss=285.9724]
Training UniVI:  12%|█▏        | 7/60 [00:06<00:45,  1.16it/s, beta=1.000, gamma=2.000, train_loss=285.9724]
Training UniVI:  12%|█▏        | 7/60 [00:06<00:45,  1.16it/s, beta=1.000, gamma=2.000, train_loss=281.2038]
Training UniVI:  13%|█▎        | 8/60 [00:06<00:45,  1.15it/s, beta=1.000, gamma=2.000, train_loss=281.2038]
Training UniVI:  13%|█▎        | 8/60 [00:07<00:45,  1.15it/s, beta=1.000, gamma=2.000, train_loss=277.7510]
Training UniVI:  15%|█▌        | 9/60 [00:07<00:44,  1.16it/s, beta=1.000, gamma=2.000, train_loss=277.7510]
Training UniVI:  15%|█▌        | 9/60 [00:08<00:44,  1.16it/s, beta=1.000, gamma=2.000, train_loss=275.9164]
Training UniVI:  17%|█▋        | 10/60 [00:08<00:43,  1.16it/s, beta=1.000, gamma=2.000, train_loss=275.9164]
Training UniVI:  17%|█▋        | 10/60 [00:09<00:43,  1.16it/s, beta=1.000, gamma=2.000, train_loss=273.5066]
Training UniVI:  18%|█▊        | 11/60 [00:09<00:42,  1.16it/s, beta=1.000, gamma=2.000, train_loss=273.5066]
Training UniVI:  18%|█▊        | 11/60 [00:10<00:42,  1.16it/s, beta=1.000, gamma=2.000, train_loss=272.2236]
Training UniVI:  20%|██        | 12/60 [00:10<00:41,  1.16it/s, beta=1.000, gamma=2.000, train_loss=272.2236]
Training UniVI:  20%|██        | 12/60 [00:11<00:41,  1.16it/s, beta=1.000, gamma=2.000, train_loss=270.6353]
Training UniVI:  22%|██▏       | 13/60 [00:11<00:40,  1.16it/s, beta=1.000, gamma=2.000, train_loss=270.6353]
Training UniVI:  22%|██▏       | 13/60 [00:12<00:40,  1.16it/s, beta=1.000, gamma=2.000, train_loss=269.7929]
Training UniVI:  23%|██▎       | 14/60 [00:12<00:39,  1.16it/s, beta=1.000, gamma=2.000, train_loss=269.7929]
Training UniVI:  23%|██▎       | 14/60 [00:12<00:39,  1.16it/s, beta=1.000, gamma=2.000, train_loss=269.0286]
Training UniVI:  25%|██▌       | 15/60 [00:12<00:38,  1.16it/s, beta=1.000, gamma=2.000, train_loss=269.0286]
Training UniVI:  25%|██▌       | 15/60 [00:13<00:38,  1.16it/s, beta=1.000, gamma=2.000, train_loss=268.1733]
Training UniVI:  27%|██▋       | 16/60 [00:13<00:37,  1.16it/s, beta=1.000, gamma=2.000, train_loss=268.1733]
Training UniVI:  27%|██▋       | 16/60 [00:14<00:37,  1.16it/s, beta=1.000, gamma=2.000, train_loss=267.3939]
Training UniVI:  28%|██▊       | 17/60 [00:14<00:37,  1.16it/s, beta=1.000, gamma=2.000, train_loss=267.3939]
Training UniVI:  28%|██▊       | 17/60 [00:15<00:37,  1.16it/s, beta=1.000, gamma=2.000, train_loss=266.8654]
Training UniVI:  30%|███       | 18/60 [00:15<00:36,  1.16it/s, beta=1.000, gamma=2.000, train_loss=266.8654]
Training UniVI:  30%|███       | 18/60 [00:16<00:36,  1.16it/s, beta=1.000, gamma=2.000, train_loss=266.3379]
Training UniVI:  32%|███▏      | 19/60 [00:16<00:35,  1.16it/s, beta=1.000, gamma=2.000, train_loss=266.3379]
[2026-09-21 01:05:50,310] [UniVITrainer] [INFO] [Epoch 020] Train loss=265.7747 (beta=1.000, gamma=2.000)
Training UniVI:  32%|███▏      | 19/60 [00:17<00:35,  1.16it/s, beta=1.000, gamma=2.000, train_loss=265.7747]
Training UniVI:  33%|███▎      | 20/60 [00:17<00:34,  1.16it/s, beta=1.000, gamma=2.000, train_loss=265.7747]
Training UniVI:  33%|███▎      | 20/60 [00:18<00:34,  1.16it/s, beta=1.000, gamma=2.000, train_loss=265.2659]
Training UniVI:  35%|███▌      | 21/60 [00:18<00:33,  1.16it/s, beta=1.000, gamma=2.000, train_loss=265.2659]
Training UniVI:  35%|███▌      | 21/60 [00:19<00:33,  1.16it/s, beta=1.000, gamma=2.000, train_loss=264.5494]
Training UniVI:  37%|███▋      | 22/60 [00:19<00:32,  1.16it/s, beta=1.000, gamma=2.000, train_loss=264.5494]
Training UniVI:  37%|███▋      | 22/60 [00:19<00:32,  1.16it/s, beta=1.000, gamma=2.000, train_loss=264.3029]
Training UniVI:  38%|███▊      | 23/60 [00:19<00:32,  1.15it/s, beta=1.000, gamma=2.000, train_loss=264.3029]
Training UniVI:  38%|███▊      | 23/60 [00:20<00:32,  1.15it/s, beta=1.000, gamma=2.000, train_loss=263.9684]
Training UniVI:  40%|████      | 24/60 [00:20<00:31,  1.15it/s, beta=1.000, gamma=2.000, train_loss=263.9684]
Training UniVI:  40%|████      | 24/60 [00:21<00:31,  1.15it/s, beta=1.000, gamma=2.000, train_loss=263.5113]
Training UniVI:  42%|████▏     | 25/60 [00:21<00:30,  1.16it/s, beta=1.000, gamma=2.000, train_loss=263.5113]
Training UniVI:  42%|████▏     | 25/60 [00:22<00:30,  1.16it/s, beta=1.000, gamma=2.000, train_loss=263.0846]
Training UniVI:  43%|████▎     | 26/60 [00:22<00:29,  1.15it/s, beta=1.000, gamma=2.000, train_loss=263.0846]
Training UniVI:  43%|████▎     | 26/60 [00:23<00:29,  1.15it/s, beta=1.000, gamma=2.000, train_loss=262.5995]
Training UniVI:  45%|████▌     | 27/60 [00:23<00:28,  1.16it/s, beta=1.000, gamma=2.000, train_loss=262.5995]
Training UniVI:  45%|████▌     | 27/60 [00:24<00:28,  1.16it/s, beta=1.000, gamma=2.000, train_loss=262.1023]
Training UniVI:  47%|████▋     | 28/60 [00:24<00:27,  1.15it/s, beta=1.000, gamma=2.000, train_loss=262.1023]
Training UniVI:  47%|████▋     | 28/60 [00:25<00:27,  1.15it/s, beta=1.000, gamma=2.000, train_loss=262.0164]
Training UniVI:  48%|████▊     | 29/60 [00:25<00:27,  1.14it/s, beta=1.000, gamma=2.000, train_loss=262.0164]
Training UniVI:  48%|████▊     | 29/60 [00:26<00:27,  1.14it/s, beta=1.000, gamma=2.000, train_loss=261.7366]
Training UniVI:  50%|█████     | 30/60 [00:26<00:26,  1.13it/s, beta=1.000, gamma=2.000, train_loss=261.7366]
Training UniVI:  50%|█████     | 30/60 [00:26<00:26,  1.13it/s, beta=1.000, gamma=2.000, train_loss=261.3628]
Training UniVI:  52%|█████▏    | 31/60 [00:26<00:25,  1.13it/s, beta=1.000, gamma=2.000, train_loss=261.3628]
Training UniVI:  52%|█████▏    | 31/60 [00:27<00:25,  1.13it/s, beta=1.000, gamma=2.000, train_loss=260.9697]
Training UniVI:  53%|█████▎    | 32/60 [00:27<00:24,  1.14it/s, beta=1.000, gamma=2.000, train_loss=260.9697]
Training UniVI:  53%|█████▎    | 32/60 [00:28<00:24,  1.14it/s, beta=1.000, gamma=2.000, train_loss=260.6790]
Training UniVI:  55%|█████▌    | 33/60 [00:28<00:23,  1.14it/s, beta=1.000, gamma=2.000, train_loss=260.6790]
Training UniVI:  55%|█████▌    | 33/60 [00:29<00:23,  1.14it/s, beta=1.000, gamma=2.000, train_loss=260.1455]
Training UniVI:  57%|█████▋    | 34/60 [00:29<00:22,  1.15it/s, beta=1.000, gamma=2.000, train_loss=260.1455]
Training UniVI:  57%|█████▋    | 34/60 [00:30<00:22,  1.15it/s, beta=1.000, gamma=2.000, train_loss=260.0006]
Training UniVI:  58%|█████▊    | 35/60 [00:30<00:21,  1.15it/s, beta=1.000, gamma=2.000, train_loss=260.0006]
Training UniVI:  58%|█████▊    | 35/60 [00:31<00:21,  1.15it/s, beta=1.000, gamma=2.000, train_loss=259.7604]
Training UniVI:  60%|██████    | 36/60 [00:31<00:20,  1.15it/s, beta=1.000, gamma=2.000, train_loss=259.7604]
Training UniVI:  60%|██████    | 36/60 [00:32<00:20,  1.15it/s, beta=1.000, gamma=2.000, train_loss=259.4353]
Training UniVI:  62%|██████▏   | 37/60 [00:32<00:19,  1.15it/s, beta=1.000, gamma=2.000, train_loss=259.4353]
Training UniVI:  62%|██████▏   | 37/60 [00:32<00:19,  1.15it/s, beta=1.000, gamma=2.000, train_loss=259.3033]
Training UniVI:  63%|██████▎   | 38/60 [00:32<00:19,  1.15it/s, beta=1.000, gamma=2.000, train_loss=259.3033]
Training UniVI:  63%|██████▎   | 38/60 [00:33<00:19,  1.15it/s, beta=1.000, gamma=2.000, train_loss=258.8386]
Training UniVI:  65%|██████▌   | 39/60 [00:33<00:18,  1.15it/s, beta=1.000, gamma=2.000, train_loss=258.8386]
[2026-09-21 01:06:07,714] [UniVITrainer] [INFO] [Epoch 040] Train loss=258.6115 (beta=1.000, gamma=2.000)
Training UniVI:  65%|██████▌   | 39/60 [00:34<00:18,  1.15it/s, beta=1.000, gamma=2.000, train_loss=258.6115]
Training UniVI:  67%|██████▋   | 40/60 [00:34<00:17,  1.15it/s, beta=1.000, gamma=2.000, train_loss=258.6115]
Training UniVI:  67%|██████▋   | 40/60 [00:35<00:17,  1.15it/s, beta=1.000, gamma=2.000, train_loss=258.2700]
Training UniVI:  68%|██████▊   | 41/60 [00:35<00:16,  1.15it/s, beta=1.000, gamma=2.000, train_loss=258.2700]
Training UniVI:  68%|██████▊   | 41/60 [00:36<00:16,  1.15it/s, beta=1.000, gamma=2.000, train_loss=258.0853]
Training UniVI:  70%|███████   | 42/60 [00:36<00:15,  1.15it/s, beta=1.000, gamma=2.000, train_loss=258.0853]
Training UniVI:  70%|███████   | 42/60 [00:37<00:15,  1.15it/s, beta=1.000, gamma=2.000, train_loss=257.7725]
Training UniVI:  72%|███████▏  | 43/60 [00:37<00:14,  1.15it/s, beta=1.000, gamma=2.000, train_loss=257.7725]
Training UniVI:  72%|███████▏  | 43/60 [00:38<00:14,  1.15it/s, beta=1.000, gamma=2.000, train_loss=257.5811]
Training UniVI:  73%|███████▎  | 44/60 [00:38<00:13,  1.15it/s, beta=1.000, gamma=2.000, train_loss=257.5811]
Training UniVI:  73%|███████▎  | 44/60 [00:39<00:13,  1.15it/s, beta=1.000, gamma=2.000, train_loss=257.2817]
Training UniVI:  75%|███████▌  | 45/60 [00:39<00:12,  1.16it/s, beta=1.000, gamma=2.000, train_loss=257.2817]
Training UniVI:  75%|███████▌  | 45/60 [00:39<00:12,  1.16it/s, beta=1.000, gamma=2.000, train_loss=257.1293]
Training UniVI:  77%|███████▋  | 46/60 [00:39<00:12,  1.16it/s, beta=1.000, gamma=2.000, train_loss=257.1293]
Training UniVI:  77%|███████▋  | 46/60 [00:40<00:12,  1.16it/s, beta=1.000, gamma=2.000, train_loss=256.8682]
Training UniVI:  78%|███████▊  | 47/60 [00:40<00:11,  1.16it/s, beta=1.000, gamma=2.000, train_loss=256.8682]
Training UniVI:  78%|███████▊  | 47/60 [00:41<00:11,  1.16it/s, beta=1.000, gamma=2.000, train_loss=256.7763]
Training UniVI:  80%|████████  | 48/60 [00:41<00:10,  1.16it/s, beta=1.000, gamma=2.000, train_loss=256.7763]
Training UniVI:  80%|████████  | 48/60 [00:42<00:10,  1.16it/s, beta=1.000, gamma=2.000, train_loss=256.3408]
Training UniVI:  82%|████████▏ | 49/60 [00:42<00:09,  1.16it/s, beta=1.000, gamma=2.000, train_loss=256.3408]
Training UniVI:  82%|████████▏ | 49/60 [00:43<00:09,  1.16it/s, beta=1.000, gamma=2.000, train_loss=256.2361]
Training UniVI:  83%|████████▎ | 50/60 [00:43<00:08,  1.15it/s, beta=1.000, gamma=2.000, train_loss=256.2361]
Training UniVI:  83%|████████▎ | 50/60 [00:44<00:08,  1.15it/s, beta=1.000, gamma=2.000, train_loss=256.0570]
Training UniVI:  85%|████████▌ | 51/60 [00:44<00:07,  1.15it/s, beta=1.000, gamma=2.000, train_loss=256.0570]
Training UniVI:  85%|████████▌ | 51/60 [00:45<00:07,  1.15it/s, beta=1.000, gamma=2.000, train_loss=255.7032]
Training UniVI:  87%|████████▋ | 52/60 [00:45<00:06,  1.16it/s, beta=1.000, gamma=2.000, train_loss=255.7032]
Training UniVI:  87%|████████▋ | 52/60 [00:45<00:06,  1.16it/s, beta=1.000, gamma=2.000, train_loss=255.5835]
Training UniVI:  88%|████████▊ | 53/60 [00:45<00:06,  1.16it/s, beta=1.000, gamma=2.000, train_loss=255.5835]
Training UniVI:  88%|████████▊ | 53/60 [00:46<00:06,  1.16it/s, beta=1.000, gamma=2.000, train_loss=255.3054]
Training UniVI:  90%|█████████ | 54/60 [00:46<00:05,  1.16it/s, beta=1.000, gamma=2.000, train_loss=255.3054]
Training UniVI:  90%|█████████ | 54/60 [00:47<00:05,  1.16it/s, beta=1.000, gamma=2.000, train_loss=255.1259]
Training UniVI:  92%|█████████▏| 55/60 [00:47<00:04,  1.15it/s, beta=1.000, gamma=2.000, train_loss=255.1259]
Training UniVI:  92%|█████████▏| 55/60 [00:48<00:04,  1.15it/s, beta=1.000, gamma=2.000, train_loss=254.9013]
Training UniVI:  93%|█████████▎| 56/60 [00:48<00:03,  1.14it/s, beta=1.000, gamma=2.000, train_loss=254.9013]
Training UniVI:  93%|█████████▎| 56/60 [00:49<00:03,  1.14it/s, beta=1.000, gamma=2.000, train_loss=254.5619]
Training UniVI:  95%|█████████▌| 57/60 [00:49<00:02,  1.15it/s, beta=1.000, gamma=2.000, train_loss=254.5619]
Training UniVI:  95%|█████████▌| 57/60 [00:50<00:02,  1.15it/s, beta=1.000, gamma=2.000, train_loss=254.6233]
Training UniVI:  97%|█████████▋| 58/60 [00:50<00:01,  1.15it/s, beta=1.000, gamma=2.000, train_loss=254.6233]
Training UniVI:  97%|█████████▋| 58/60 [00:51<00:01,  1.15it/s, beta=1.000, gamma=2.000, train_loss=254.2427]
Training UniVI:  98%|█████████▊| 59/60 [00:51<00:00,  1.15it/s, beta=1.000, gamma=2.000, train_loss=254.2427]
[2026-09-21 01:06:25,051] [UniVITrainer] [INFO] [Epoch 060] Train loss=253.8308 (beta=1.000, gamma=2.000)
Training UniVI:  98%|█████████▊| 59/60 [00:52<00:00,  1.15it/s, beta=1.000, gamma=2.000, train_loss=253.8308]
Training UniVI: 100%|██████████| 60/60 [00:52<00:00,  1.16it/s, beta=1.000, gamma=2.000, train_loss=253.8308]
Training UniVI: 100%|██████████| 60/60 [00:52<00:00,  1.15it/s, beta=1.000, gamma=2.000, train_loss=253.8308]
router x precision precision only
rna 0.665 0.372
adt 0.000 0.296
meth 0.335 0.332

A transformer encoder (experimental)#

Any modality can use a transformer encoder instead of an MLP. The tokenizer turns each cell into a set of tokens (here the 64 highest-valued features, each carrying its value, rank and a dropout indicator).

tf_cfg = UniVIConfig(
    latent_dim=8, beta=1.0, gamma=2.0,
    modalities=[
        ModalityConfig("rna", 300, [128, 64], [64, 128], likelihood="nb", encoder_type="transformer",
                       tokenizer=TokenizerConfig(mode="topk_channels", n_tokens=64,
                                                 channels=("value", "rank", "dropout")),
                       transformer=TransformerConfig(d_model=64, num_heads=4, num_layers=2, dim_feedforward=128)),
        ModalityConfig("adt", 20, [32], [32], likelihood="gaussian"),
    ],
)
tf_model = UniVIMultiModalVAE(tf_cfg, loss_mode="v1", v1_recon="avg", normalize_v1_terms=True)
two = {"rna": rna, "adt": adt}
UniVITrainer(tf_model, make_loader(subset(two, train_idx), batch_size=128, shuffle=True, drop_last=True),
             make_loader(subset(two, val_idx), batch_size=512),
             TrainingConfig(n_epochs=max(1, N_EPOCHS // 3), lr=1e-3, device=device, log_every=10)).fit();
encode_adata(tf_model, val_data["rna"], modality="rna", device=device).shape

[2026-09-21 01:06:25,126] [UniVITrainer] [INFO] TrainingConfig:
[2026-09-21 01:06:25,126] [UniVITrainer] [INFO]   n_epochs: 20
[2026-09-21 01:06:25,127] [UniVITrainer] [INFO]   batch_size: 256
[2026-09-21 01:06:25,127] [UniVITrainer] [INFO]   lr: 0.001
[2026-09-21 01:06:25,127] [UniVITrainer] [INFO]   weight_decay: 0.0
[2026-09-21 01:06:25,127] [UniVITrainer] [INFO]   device: 'cuda'
[2026-09-21 01:06:25,128] [UniVITrainer] [INFO]   log_every: 10
[2026-09-21 01:06:25,128] [UniVITrainer] [INFO]   grad_clip: None
[2026-09-21 01:06:25,128] [UniVITrainer] [INFO]   num_workers: 0
[2026-09-21 01:06:25,129] [UniVITrainer] [INFO]   seed: 0
[2026-09-21 01:06:25,129] [UniVITrainer] [INFO]   early_stopping: False
[2026-09-21 01:06:25,129] [UniVITrainer] [INFO]   patience: 20
[2026-09-21 01:06:25,130] [UniVITrainer] [INFO]   min_delta: 0.0
[2026-09-21 01:06:25,130] [UniVITrainer] [INFO]   best_epoch_warmup: 0
Training UniVI:   0%|          | 0/20 [00:00<?, ?it/s]
[2026-09-21 01:06:25,825] [UniVITrainer] [INFO] [Epoch 001] Train loss=438.6072 (beta=1.000, gamma=2.000)
[2026-09-21 01:06:25,860] [UniVITrainer] [INFO] [Epoch 001] Val loss=400.6721 (beta=1.000, gamma=2.000)
Training UniVI:   0%|          | 0/20 [00:00<?, ?it/s, beta=1.000, gamma=2.000, train_loss=438.6072, val_loss=400.6721]
[2026-09-21 01:06:25,863] [UniVITrainer] [INFO] [Epoch 001] New best val loss: 400.6721
Training UniVI:   5%|▌         | 1/20 [00:00<00:13,  1.37it/s, beta=1.000, gamma=2.000, train_loss=438.6072, val_loss=400.6721]
Training UniVI:   5%|▌         | 1/20 [00:01<00:13,  1.37it/s, beta=1.000, gamma=2.000, train_loss=363.5851, val_loss=332.1391]
[2026-09-21 01:06:26,543] [UniVITrainer] [INFO] [Epoch 002] New best val loss: 332.1391
Training UniVI:  10%|█         | 2/20 [00:01<00:12,  1.42it/s, beta=1.000, gamma=2.000, train_loss=363.5851, val_loss=332.1391]
Training UniVI:  10%|█         | 2/20 [00:02<00:12,  1.42it/s, beta=1.000, gamma=2.000, train_loss=324.7584, val_loss=320.2292]
[2026-09-21 01:06:27,224] [UniVITrainer] [INFO] [Epoch 003] New best val loss: 320.2292
Training UniVI:  15%|█▌        | 3/20 [00:02<00:11,  1.45it/s, beta=1.000, gamma=2.000, train_loss=324.7584, val_loss=320.2292]
Training UniVI:  15%|█▌        | 3/20 [00:02<00:11,  1.45it/s, beta=1.000, gamma=2.000, train_loss=318.4940, val_loss=315.8456]
[2026-09-21 01:06:27,915] [UniVITrainer] [INFO] [Epoch 004] New best val loss: 315.8456
Training UniVI:  20%|██        | 4/20 [00:02<00:11,  1.45it/s, beta=1.000, gamma=2.000, train_loss=318.4940, val_loss=315.8456]
Training UniVI:  20%|██        | 4/20 [00:03<00:11,  1.45it/s, beta=1.000, gamma=2.000, train_loss=315.4977, val_loss=313.8099]
[2026-09-21 01:06:28,615] [UniVITrainer] [INFO] [Epoch 005] New best val loss: 313.8099
Training UniVI:  25%|██▌       | 5/20 [00:03<00:10,  1.44it/s, beta=1.000, gamma=2.000, train_loss=315.4977, val_loss=313.8099]
Training UniVI:  25%|██▌       | 5/20 [00:04<00:10,  1.44it/s, beta=1.000, gamma=2.000, train_loss=313.7196, val_loss=312.2919]
[2026-09-21 01:06:29,292] [UniVITrainer] [INFO] [Epoch 006] New best val loss: 312.2919
Training UniVI:  30%|███       | 6/20 [00:04<00:09,  1.45it/s, beta=1.000, gamma=2.000, train_loss=313.7196, val_loss=312.2919]
Training UniVI:  30%|███       | 6/20 [00:04<00:09,  1.45it/s, beta=1.000, gamma=2.000, train_loss=312.4509, val_loss=311.3230]
[2026-09-21 01:06:29,979] [UniVITrainer] [INFO] [Epoch 007] New best val loss: 311.3230
Training UniVI:  35%|███▌      | 7/20 [00:04<00:08,  1.45it/s, beta=1.000, gamma=2.000, train_loss=312.4509, val_loss=311.3230]
Training UniVI:  35%|███▌      | 7/20 [00:05<00:08,  1.45it/s, beta=1.000, gamma=2.000, train_loss=311.3750, val_loss=310.2835]
[2026-09-21 01:06:30,673] [UniVITrainer] [INFO] [Epoch 008] New best val loss: 310.2835
Training UniVI:  40%|████      | 8/20 [00:05<00:08,  1.45it/s, beta=1.000, gamma=2.000, train_loss=311.3750, val_loss=310.2835]
Training UniVI:  40%|████      | 8/20 [00:06<00:08,  1.45it/s, beta=1.000, gamma=2.000, train_loss=309.7627, val_loss=308.3967]
[2026-09-21 01:06:31,363] [UniVITrainer] [INFO] [Epoch 009] New best val loss: 308.3967
Training UniVI:  45%|████▌     | 9/20 [00:06<00:07,  1.45it/s, beta=1.000, gamma=2.000, train_loss=309.7627, val_loss=308.3967]
[2026-09-21 01:06:31,996] [UniVITrainer] [INFO] [Epoch 010] Train loss=307.8888 (beta=1.000, gamma=2.000)
[2026-09-21 01:06:32,032] [UniVITrainer] [INFO] [Epoch 010] Val loss=306.0486 (beta=1.000, gamma=2.000)
Training UniVI:  45%|████▌     | 9/20 [00:06<00:07,  1.45it/s, beta=1.000, gamma=2.000, train_loss=307.8888, val_loss=306.0486]
[2026-09-21 01:06:32,036] [UniVITrainer] [INFO] [Epoch 010] New best val loss: 306.0486
Training UniVI:  50%|█████     | 10/20 [00:06<00:06,  1.46it/s, beta=1.000, gamma=2.000, train_loss=307.8888, val_loss=306.0486]
Training UniVI:  50%|█████     | 10/20 [00:07<00:06,  1.46it/s, beta=1.000, gamma=2.000, train_loss=305.0382, val_loss=302.7785]
[2026-09-21 01:06:32,725] [UniVITrainer] [INFO] [Epoch 011] New best val loss: 302.7785
Training UniVI:  55%|█████▌    | 11/20 [00:07<00:06,  1.46it/s, beta=1.000, gamma=2.000, train_loss=305.0382, val_loss=302.7785]
Training UniVI:  55%|█████▌    | 11/20 [00:08<00:06,  1.46it/s, beta=1.000, gamma=2.000, train_loss=303.1881, val_loss=306.5779]
Training UniVI:  60%|██████    | 12/20 [00:08<00:05,  1.47it/s, beta=1.000, gamma=2.000, train_loss=303.1881, val_loss=306.5779]
Training UniVI:  60%|██████    | 12/20 [00:08<00:05,  1.47it/s, beta=1.000, gamma=2.000, train_loss=302.0521, val_loss=304.3819]
Training UniVI:  65%|██████▌   | 13/20 [00:08<00:04,  1.46it/s, beta=1.000, gamma=2.000, train_loss=302.0521, val_loss=304.3819]
Training UniVI:  65%|██████▌   | 13/20 [00:09<00:04,  1.46it/s, beta=1.000, gamma=2.000, train_loss=300.9462, val_loss=298.2224]
[2026-09-21 01:06:34,762] [UniVITrainer] [INFO] [Epoch 014] New best val loss: 298.2224
Training UniVI:  70%|███████   | 14/20 [00:09<00:04,  1.47it/s, beta=1.000, gamma=2.000, train_loss=300.9462, val_loss=298.2224]
Training UniVI:  70%|███████   | 14/20 [00:10<00:04,  1.47it/s, beta=1.000, gamma=2.000, train_loss=300.4201, val_loss=301.9197]
Training UniVI:  75%|███████▌  | 15/20 [00:10<00:03,  1.46it/s, beta=1.000, gamma=2.000, train_loss=300.4201, val_loss=301.9197]
Training UniVI:  75%|███████▌  | 15/20 [00:10<00:03,  1.46it/s, beta=1.000, gamma=2.000, train_loss=299.2633, val_loss=297.9133]
[2026-09-21 01:06:36,121] [UniVITrainer] [INFO] [Epoch 016] New best val loss: 297.9133
Training UniVI:  80%|████████  | 16/20 [00:10<00:02,  1.47it/s, beta=1.000, gamma=2.000, train_loss=299.2633, val_loss=297.9133]
Training UniVI:  80%|████████  | 16/20 [00:11<00:02,  1.47it/s, beta=1.000, gamma=2.000, train_loss=297.9183, val_loss=299.5851]
Training UniVI:  85%|████████▌ | 17/20 [00:11<00:02,  1.47it/s, beta=1.000, gamma=2.000, train_loss=297.9183, val_loss=299.5851]
Training UniVI:  85%|████████▌ | 17/20 [00:12<00:02,  1.47it/s, beta=1.000, gamma=2.000, train_loss=297.6540, val_loss=296.4494]
[2026-09-21 01:06:37,470] [UniVITrainer] [INFO] [Epoch 018] New best val loss: 296.4494
Training UniVI:  90%|█████████ | 18/20 [00:12<00:01,  1.48it/s, beta=1.000, gamma=2.000, train_loss=297.6540, val_loss=296.4494]
Training UniVI:  90%|█████████ | 18/20 [00:13<00:01,  1.48it/s, beta=1.000, gamma=2.000, train_loss=297.3117, val_loss=295.4856]
[2026-09-21 01:06:38,155] [UniVITrainer] [INFO] [Epoch 019] New best val loss: 295.4856
Training UniVI:  95%|█████████▌| 19/20 [00:13<00:00,  1.47it/s, beta=1.000, gamma=2.000, train_loss=297.3117, val_loss=295.4856]
[2026-09-21 01:06:38,865] [UniVITrainer] [INFO] [Epoch 020] Train loss=296.0893 (beta=1.000, gamma=2.000)
[2026-09-21 01:06:38,901] [UniVITrainer] [INFO] [Epoch 020] Val loss=294.4854 (beta=1.000, gamma=2.000)
Training UniVI:  95%|█████████▌| 19/20 [00:13<00:00,  1.47it/s, beta=1.000, gamma=2.000, train_loss=296.0893, val_loss=294.4854]
[2026-09-21 01:06:38,905] [UniVITrainer] [INFO] [Epoch 020] New best val loss: 294.4854
Training UniVI: 100%|██████████| 20/20 [00:13<00:00,  1.43it/s, beta=1.000, gamma=2.000, train_loss=296.0893, val_loss=294.4854]
Training UniVI: 100%|██████████| 20/20 [00:13<00:00,  1.45it/s, beta=1.000, gamma=2.000, train_loss=296.0893, val_loss=294.4854]

[2026-09-21 01:06:38,938] [UniVITrainer] [INFO] Restored best model from epoch 20 (val loss=294.4854)

(200, 8)

Checklist for a new assay#

  1. One AnnData per modality; identical, identically ordered obs_names across paired modalities.

  2. Put the model input in .X (or an .obsm key passed as X_key) and keep raw data in a layer.

  3. Fit any learned transform (feature selection, scaling, SVD) on training cells only.

  4. Pick the likelihood from the table above; add recon_targets_spec for binomial-type data.

  5. Size encoders to the input: wide inputs (thousands of features) get wider first layers.

  6. Balance modalities with recon_weight if one dominates the loss.

Contents