K01
ICLR 2026 · Gen² Workshop · Tiny Papers Track

Tensorised Modular Architectures for Multi-Omics Generation.

Arinbjörn Kolbeinsson · Benedikt Kolbeinsson

Single-cell multi-omics generative models typically flatten the feature vector and lose the known module structure of biology. We test whether grouping features into biological modules, plus a Tensor-Train coupling layer, recovers what flat baselines miss on a CITE-seq PBMC dataset.

Abstract

Conditional generative models for single-cell multi-omics typically concatenate features into a flat vector, which discards the known module structure of biology. We compare three autoencoder variants on a CITE-seq PBMC dataset (10,000 cells, 30 cell types, 2,000 RNA highly-variable genes plus 217 antibody-derived tag proteins): a flat dense baseline, a modular encoder with dense coupling, and a modular encoder with a Tensor-Train coupling layer. At comparable parameter budgets around 60–100K, both modular variants substantially beat the flat baseline. The headline number is +48% on conditional RNA–protein correlation for the modular TT model against flat. The Tensor-Train coupling itself contributes a modest, not dominant, gain over dense modular coupling. Confidence bands overlap at several points across the parameter sweep. Results are from a single dataset with three random seeds and should be read as preliminary.

Key findings

  1. Modularity matters more than coupling type. At parameter budgets near 60–100K, both modular variants substantially outperform the flat baseline across conditional RNA–protein correlation, interaction preservation and TSTR macro F1. Modular TT reaches +48% conditional correlation over flat.
  2. Tensor-Train coupling is a modest add over dense. Across a parameter sweep from about 20K to 1M, the TT coupling matches or modestly exceeds dense coupling. The clearest gap sits in the 78K–193K range. At about 195K parameters the conditional correlation is 0.380 for TT against 0.348 for dense. With three seeds, confidence bands overlap at several points.
  3. Cell-type classification is the easy axis. Both modular variants reach TSTR macro F1 above 0.90 throughout the sweep. The harder evaluation is the cross-modal correlation conditioned on cell type, where the gap to flat baselines is widest.
  4. Architecture, not parameter budget, drives the win. The flat baseline does not catch up when given more parameters. Modular variants win at matched budgets, not by being larger.
  5. Whole-genome scaling is hypothesised, not shown. The current experiments use 2,000 RNA features and 217 proteins on one dataset. The authors hypothesise the tensor structure would help more at whole-genome scale. This is not tested.

Why it matters

Multi-omics generative models are used for data augmentation, imputation and in-silico perturbation. Each depends on getting cross-modal relationships right. An RNA expression vector that does not match its associated proteins is no use for downstream biology. Flat dense autoencoders make this hard. They have no way of knowing which features should move together. Building module structure into the encoder lets the model spend capacity on the relationships that matter without scaling parameters everywhere.

The Tensor-Train coupling is a more speculative move. On this dataset the gain over dense modular coupling is small and the confidence bands overlap. If it pays off it will be at scale, where flattening a module-by-module interaction tensor stops being practical. That remains for future work to show.

Scope and limitations

A single dataset (CITE-seq PBMC). One module-construction choice (hierarchical clustering on correlation distance), with no comparison to pathway-based or curated module sets. No random-grouping ablation, so the gain attributed to biological structure could be partly attributable to modularity itself. Three random seeds. Baselines are scVI and totalVI; transformers are not tested. The TT advantage over dense coupling is described as modest by the authors.

Cite

@inproceedings{kolbeinsson2026tensorised,
  title     = {Tensorised Modular Architectures for Multi-Omics Generation},
  author    = {Kolbeinsson, Arinbj{\"o}rn and Kolbeinsson, Benedikt},
  booktitle = {ICLR 2026 Workshop on Generative and Experimental Perspectives for Biomolecular Design (Gen$^2$)},
  year      = {2026},
  url       = {https://openreview.net/forum?id=dgLH2Kuw50}
}