Nephrobase Cell+: Multimodal Single-Cell Foundation Model for Decoding Kidney Biology (2509.26223v1)

Abstract: Background: Large foundation models have revolutionized single-cell analysis, yet no kidney-specific model currently exists, and it remains unclear whether organ-focused models can outperform generalized models. The kidney's complex cellular architecture further complicate integration of large-scale omics data, where current frameworks trained on limited datasets struggle to correct batch effects, capture cross-modality variation, and generalize across species. Methods: We developed Nephrobase Cell+, the first kidney-focused large foundation model, pretrained on ~100 billion tokens from ~39.5 million single-cell and single-nucleus profiles across 4,319 samples. Nephrobase Cell+ uses a transformer-based encoder-decoder architecture with gene-token cross-attention and a mixture-of-experts module for scalable representation learning. Results: Nephrobase Cell+ sets a new benchmark for kidney single-cell analysis. It produces tightly clustered, biologically coherent embeddings in human and mouse kidneys, far surpassing previous foundation models such as Geneformer, scGPT, and UCE, as well as traditional methods such as PCA and autoencoders. It achieves the highest cluster concordance and batch-mixing scores, effectively removing donor/assay batch effects while preserving cell-type structure. Cross-species evaluation shows superior alignment of homologous cell types and >90% zero-shot annotation accuracy for major kidney lineages in both human and mouse. Even its 1B-parameter and 500M variants consistently outperform all existing models. Conclusions: Nephrobase Cell+ delivers a unified, high-fidelity representation of kidney biology that is robust, cross-species transferable, and unmatched by current single-cell foundation models, offering a powerful resource for kidney genomics and disease research.

• The paper introduces a transformer-based, multimodal foundation model tailored for kidney single-cell and spatial omics data that enhances annotation and cross-species generalization.
• It leverages advanced architecture with gene-token cross-attention, a Mixture-of-Experts module, and specialized loss functions to robustly integrate diverse datasets.
• Benchmarking shows superior clustering metrics and zero-shot cell-type annotation accuracy over 90%, matching expert manual curation across species.

Nephrobase Cell+: A Multimodal Single-Cell Foundation Model for Kidney Biology

Introduction

Nephrobase Cell+ introduces a transformer-based, multimodal foundation model specifically tailored for kidney single-cell and spatial omics data. The model is trained from scratch on an extensive dataset comprising approximately 39.5 million single-cell and single-nucleus profiles from 4,319 samples, spanning four mammalian species (human, mouse, rat, pig) and multiple assay modalities (scRNA-seq, snRNA-seq, snATAC-seq, spatial transcriptomics). This organ-centric approach addresses the limitations of general-purpose single-cell models, which often fail to capture the hierarchical and microenvironmental complexity of kidney tissue, and struggle with batch effects, cross-modality variation, and species generalization.

Model Architecture and Training

Nephrobase Cell+ employs a transformer encoder-decoder architecture augmented with gene-token cross-attention and a Mixture-of-Experts (MoE) module. The input is a cell-by-gene matrix, where each gene is represented by its identity, normalized expression value, and optional metadata. Gene indices are embedded via a trainable layer, and expression counts are sum-log normalized and embedded through a non-linear MLP. The encoder applies self-attention to the tokenized input, while the decoder uses cross-attention to reconstruct gene expression profiles and predict cell types.

Key architectural components include:

• Zero-Inflated Negative Binomial (ZINB) Reconstruction Head: Models count overdispersion and excess zeros typical of single-cell data.
• Focal Loss Classification Head: Addresses class imbalance and emphasizes difficult examples.
• Mixture-of-Experts (MoE) Layer: Implements sparse top-k routing to specialized experts, with a load-balancing auxiliary loss to ensure even expert utilization.
• Elastic Cell Similarity (ECS) Regularizer: Enforces a target level of dissimilarity between cell embeddings, preventing representational collapse.
• Supervised Contrastive Loss: Pulls together embeddings of cells with the same annotation and pushes apart those with different labels.
• Adversarial Discriminators with Gradient Reversal: Remove assay and batch signals from learned features, producing assay-invariant representations.

Training is performed using Fully Sharded Data Parallelism (FSDP) across 4 H100 GPUs, with Adam/AdamW optimization, learning rate scheduling (ReduceLROnPlateau or CosineAnnealingLR), warmup, gradient clipping, and mixed precision. The model is initialized with Xavier uniform for linear layers.

Data Integration and Preprocessing

The training dataset integrates public and in-house data, including scRNA-seq, snRNA-seq, snATAC-seq, and high-plex spatial transcriptomics (CosMx, Xenium). Gene orthology mapping is performed to harmonize feature spaces across species, resulting in a unified set of 32,768 highly variable genes. Rigorous quality control is applied, including ambient RNA correction (CellBender), median absolute deviation filtering, and total count normalization.

Benchmarking and Performance

Embedding Visualization and Clustering

Nephrobase Cell+ produces tightly clustered, biologically coherent embeddings in both human and mouse kidney datasets, outperforming Geneformer, scGPT, UCE, PCA, and autoencoders. Quantitative metrics demonstrate superior cell-type isolation (KMeans ARI: 0.82 for 1B model vs. 0.40 for scGPT, 0.22 for Geneformer), higher normalized mutual information (NMI: 0.78 vs. 0.48 for scGPT), and better silhouette scores. Batch-mixing metrics (cLISI: 1.00, iLISI: ~0.17) indicate effective removal of batch effects while preserving cell-type structure.

Cross-Species Generalization

Nephrobase Cell+ achieves near-perfect alignment of homologous cell types across human and mouse, with zero-shot annotation accuracy exceeding 90% for major kidney lineages. The model's embeddings are species-invariant (cLISI: 1.00, iLISI: 0.01), and clustering metrics (NMI: 0.75, ARI: 0.72 for 500M model) surpass those of Geneformer and UCE. Graph connectivity and PCR values are also highest for Nephrobase Cell+, reflecting well-connected and unbiased latent spaces.

Zero-Shot Cell-Type Annotation

Confusion matrices and Sankey diagrams reveal that Nephrobase Cell+ matches expert-level manual curation for major cell groups in both human and mouse kidneys. Over 90% of proximal tubule cells are correctly annotated, with minor misclassifications restricted to closely related subtypes. The model generalizes to new data without retraining, enabling robust annotation consistent with manual standards.

In Silico Perturbation

Gene perturbation experiments demonstrate that Nephrobase Cell+ can simulate biologically coherent transcriptional responses. For example, CCL2 upregulation induces proinflammatory and chemotactic signatures, VCAM1 perturbation enriches adhesion and immune activation pathways, GDF15 modulates growth control and MAPK signaling, and SOX4 affects developmental and bioenergetic processes. These results validate the model's capacity to capture gene regulatory networks and predict cellular responses to perturbation.

Implementation Considerations

• Computational Requirements: Training Nephrobase Cell+ requires substantial resources (multi-GPU, FSDP, mixed precision). The 1B-parameter and 500M-parameter variants both outperform existing models, but the larger model offers marginal gains at increased computational cost.
• Data Diversity: The model's robustness is contingent on the diversity and scale of the training dataset. Rare cell types and extreme pathological states may be underrepresented, limiting generalization in those contexts.
• Modality and Species Coverage: While the model integrates RNA and ATAC modalities across four mammalian species, it does not yet support proteomics, metabolomics, or non-mammalian species.
• Deployment: Pretrained Nephrobase Cell+ models can be fine-tuned for downstream tasks (annotation, trajectory inference, perturbation simulation) with minimal additional data. The learned embeddings are general-purpose and transferable.

Theoretical and Practical Implications

Nephrobase Cell+ demonstrates that organ-specific foundation models can outperform generalized models in accuracy, interpretability, and cross-species transferability. The transformer-based architecture, multimodal integration, and adversarial training components collectively enable unified, high-fidelity representations of kidney biology. This facilitates large-scale atlas construction, disease mechanism discovery, and simulation of cellular responses, addressing key challenges in nephrology such as batch effects, sparse data, and inter-sample variability.

The model's attention weights and latent spaces offer opportunities for mining novel regulatory circuits and prioritizing candidate biomarkers. Iterative updating with new data, including patient-derived biopsies and clinical outcomes, could bridge the gap between molecular signatures and patient prognosis.

Limitations and Future Directions

• Coverage: Expansion to additional species (e.g., non-human primates) and modalities (spatial multi-omics, proteogenomics) is warranted.
• Rare Cell Types: Targeted data collection and fine-tuning on underrepresented cell types or disease states could improve performance.
• Causality: While the model captures correlations and regulatory structure, experimental validation remains essential for causal inference.
• Scalability: Further optimization of training strategies and model architectures may reduce resource requirements and facilitate broader adoption.

Conclusion

Nephrobase Cell+ establishes a new benchmark for kidney single-cell analysis, delivering robust, multimodal, and cross-species transferable representations. Its specialized transformer architecture and organ-scale pretraining enable superior clustering, batch correction, and annotation performance compared to existing foundation models. The model provides a versatile resource for kidney genomics and disease research, with significant implications for atlas construction, biomarker discovery, and therapeutic development. Future work should focus on expanding modality and species coverage, refining model architectures, and integrating clinical outcomes to further advance precision nephrology.