White-Box Diffusion Transformer for single-cell RNA-seq generation

TL;DR

This work proposes a hybrid model based on Diffusion model and White-Box transformer that aims to generate synthetic and biologically plausible scRNA-seq data and demonstrates comparable performance of White-Box Diffusion Transformer and Diffusion Transformer in generating scRNA-seq data alongside significant improvements in training efficiency and resource utilization.

Abstract

As a powerful tool for characterizing cellular subpopulations and cellular heterogeneity, single cell RNA sequencing (scRNA-seq) technology offers advantages of high throughput and multidimensional analysis. However, the process of data acquisition is often constrained by high cost and limited sample availability. To overcome these limitations, we propose a hybrid model based on Diffusion model and White-Box transformer that aims to generate synthetic and biologically plausible scRNA-seq data. Diffusion model progressively introduce noise into the data and then recover the original data through a denoising process, a forward and reverse process that is particularly suitable for generating complex data distributions. White-Box transformer is a deep learning architecture that emphasizes mathematical interpretability. By minimizing the encoding rate of the data and maximizing the sparsity of the representation, it not only reduces the computational burden, but also provides clear insight into underlying structure. Our White-Box Diffusion Transformer combines the generative capabilities of Diffusion model with the mathematical interpretability of White-Box transformer. Through experiments using six different single-cell RNA-Seq datasets, we visualize both generated and real data using t-SNE dimensionality reduction technique, as well as quantify similarity between generated and real data using various metrics to demonstrate comparable performance of White-Box Diffusion Transformer and Diffusion Transformer in generating scRNA-seq data alongside significant improvements in training efficiency and resource utilization. Our code is available at https://github.com/lingximamo/White-Box-Diffusion-Transformer

Figures (3)

• Figure 1: Architecture of noise predictor for White-Box Diffusion Transformer. We employ the MSSA layer for data compression, replacing the Attention layer. Additionally, we utilize the ISTA layer to achieve data sparsity and replace the FFN layer.
• Figure 2: Fitting process of White-Box Diffusion Transformer. Scatter plots of the data generated by the White-Box Diffusion Transformer model with varying training epochs and real data are generated after dimensionality reduction using t-SNE.
• Figure 3: a. Scatter plots of generated and real data after t-SNE dimensionality reduction for White-Box DiT and DiT. b. Scatter plots of large batches of generated data and actual data following t-SNE dimensionality reduction.c. Average running time per epoch when training DiT and White-Box DiT.