A generative model for dimensionality reduction with millions of features and few samples

Technical Summary: A Generative Model for Dimensionality Reduction with Millions of Features and Few Samples

Problem Statement
Standard dimensionality reduction methods face significant challenges when applied to datasets with millions of features (e.g., genomic data) and relatively few samples. Linear methods like Principal Component Analysis (PCA) are computationally intensive for high-dimensional data and may fail to capture non-linear structures. Non-linear methods like t-SNE and UMAP are often infeasible for direct application on such high-dimensional data, typically requiring a preliminary PCA step. While autoencoders and Variational Autoencoders (VAEs) offer a flexible, non-linear alternative, there is a prevailing concern that training such models requires a number of samples proportional to the feature dimensionality to avoid overfitting and ensure identifiability. This paper investigates whether a deep generative model can be effectively trained for dimensionality reduction in "high-dimensional, sample-limited" regimes (millions of features, few thousand samples).

Methodology
The authors propose and evaluate a Deep Generative Decoder (DGD), a decoder-only generative model. The core hypothesis is that for networks where the number of parameters scales linearly with the feature dimension, the number of training samples required is essentially independent of the feature dimensionality.

• Theoretical Basis: The authors derive a theoretical estimate for the critical number of samples ($N_c$) needed to train a decoder. For a network with parameters $C \propto d$ (where $d$ is feature dimension), $N_c \approx C/d$, suggesting $N_c$ is constant regardless of $d$.
• Architecture:
  • DGD: A decoder-only network with a latent space ($z$) and a decoder mapping $z$ back to the high-dimensional feature space.
  • Receptive Field Layer: To handle millions of features without infeasible fully connected output layers, the authors employ a "receptive field" layer. This layer connects neurons in the final hidden layer to local windows of the output features with a specific stride and overlap, maintaining sparsity and computational efficiency while preserving local genomic context.
  • Comparative Models: The DGD is compared against the Variational Autoencoder (VAE) (using a symmetric encoder-decoder structure) and PCA.
• Datasets:
  1. Synthetic Data: Non-linear data generated as perturbed points on hyperspherical clusters with dimensions up to 500,000.
  2. 1000 Genomes Project (1KGP): A downsampled dataset of human genetic variants (SNPs) ranging from 100k to ~686k features, with sample sizes from 100 to 1,000.
  3. International Cancer Genome Consortium (ICGC): A curated dataset of somatic mutations across 22 cancer types, containing ~4.4 million features and ~4,000 samples.

Key Contributions

1. Theoretical and Empirical Validation of Sample Independence: The paper demonstrates that the number of samples required to train a decoder-only model is largely independent of the feature dimensionality, provided the network architecture scales parameters linearly with features.
2. Receptive Field Implementation: The introduction of a receptive field layer allows the training of generative models on datasets with millions of features using standard GPU memory (16GB), overcoming the computational bottlenecks of fully connected layers.
3. Performance in High-Dimensional, Low-Sample Regimes: The study provides empirical evidence that DGDs can successfully learn structured latent representations from genomic data with millions of features and only a few thousand samples, outperforming VAEs and PCA in specific downstream tasks.

Results

• Synthetic Data: The DGD consistently achieved lower Mean Squared Error (MSE) than the VAE across all feature dimensions (up to 500k). Crucially, the DGD's test error remained stable and independent of the feature dimension, whereas VAE performance was sensitive to sample size and hidden layer dimensions.
• 1KGP Dataset:
  • Clustering: When applied to population structure clustering (5 continental groups), DGD consistently outperformed VAE in Adjusted Rand Index (ARI) and Normalized Mutual Information (NMI), particularly at lower sample sizes (100–500).
  • PCA Comparison: While PCA showed strong clustering performance (likely due to the linear nature of population structure), DGD achieved comparable results at higher sample sizes (1,000) and demonstrated superior ability to capture non-linear structures in specific ablation experiments (e.g., removing admixed populations).
• ICGC Dataset (4.4M Features):
  • Latent Structure: The DGD produced a latent space with clear clustering by tumor type, whereas the default VAE ($\beta=1$) produced a less organized latent space.
  • Classification: In a tumor type classification task (22 classes) using logistic regression on the latent representations, DGD achieved a balanced accuracy of 67% (95% CI: 64–70%). This significantly outperformed VAE (42%) and PCA (54%).
  • Ablation: Adjusting the VAE's KL weighting ($\beta=0$) improved its performance to 52%, narrowing the gap with PCA but still falling short of DGD. This suggests that while hyperparameter tuning affects VAE performance, the DGD architecture holds a distinct advantage in this regime.
• Computational Efficiency: The DGD was successfully trained on a T4 GPU with 16GB of memory. In contrast, processing 4.4 million features with PCA (even via Incremental PCA) is described as memory-intensive and computationally demanding.

Significance and Claims
The paper claims that training a decoder-only generative model for dimensionality reduction is feasible and advantageous in settings with millions of features and few samples. The primary significance lies in demonstrating that:

1. Feasibility: High-dimensional generative modeling does not strictly require massive sample sizes proportional to feature dimension, challenging the intuition that parameter count dictates sample requirements.
2. Superiority over VAE: The DGD architecture is more robust and stable than VAEs in these specific high-dimensional, low-sample scenarios, producing better-structured latent spaces for downstream tasks like clustering and classification.
3. Efficiency: The receptive field approach enables the processing of massive genomic datasets on consumer-grade hardware (16GB GPU), offering a practical alternative to memory-heavy linear methods like PCA.

The authors maintain a modest tone regarding their claims, noting that they did not perform systematic hyperparameter tuning due to computational costs and limited sample sizes. Consequently, the results are presented as evidence of feasibility rather than an optimized benchmark. They acknowledge that the observed performance gap between DGD and VAE may be partially attributable to specific hyperparameter choices (like $\beta$ in VAE) rather than solely architectural limitations, and that the test-set-based early stopping may lead to optimistic estimates of generalization performance.