Transformers as Provable Approximators of Sparse PCA

This is the official code repository for STAI-X 2026 strongly accepted paper Transformers as Provable Approximators of Sparse Principal Component Analysis.

Setup

Install the dependencies in REQUIREMENTS.md with the PyTorch/CUDA build that matches your machine.

Check the environment from the repository root:

python check_versions.py

Important

Run experiment scripts from inside their own folder. For example:

cd colon
python data_prepare.py
python main.py

The scripts use relative paths for local data/, results/, and figures/ folders.

Repository Layout

.
|-- model.py              # shared Transformer models
|-- metrics.py            # shared SPCA metrics/utilities
|-- check_versions.py     # environment check
|-- synthetic/            # synthetic SPCA experiments
|-- pitprops/             # Pitprops benchmark
`-- colon/                # Colon gene-expression benchmark

Synthetic

Purpose: generate synthetic SPCA tasks, train/evaluate Transformer SPCA, run classical baselines, and plot synthetic results. Synthetic datasets are already included in the data folder.

Key files:

synthetic/
|-- data_generation.py
|-- main.py
|-- baselines.py
|-- computational_cost.py
|-- plot_barplot.py
|-- plot_ablation_results.py
|-- data/
|-- results/
`-- figures/

Run:

python main.py

Optional: (reproducing Fig. 2, Fig. 3, and Table 1 in paper)

python baselines.py
python plot_barplot.py

python main.py --run ablations --noise 1,2,3,5,10
python plot_ablation_results.py

python computational_cost.py

Pitprops

Purpose: run SPCA experiments on the classical Pitprops correlation-matrix benchmark.

The Pitprops dataset is a classical sparse PCA benchmark. It is provided as a 13x13 correlation matrix computed from 180 observations with 13 explanatory variables. In this project, we load it from the R package elasticnet and save it to a CSV file.

# Install the package if it is not available
if (!requireNamespace("elasticnet", quietly = TRUE)) {
  install.packages("elasticnet")
}
# Load the package and the Pitprops dataset
library(elasticnet)
data(pitprops)
# Save the correlation matrix for later use (replace the path with your own path)
write.csv(pitprops, file = "download/path/to/pitprops_cor.csv", row.names = TRUE)

Key files:

pitprops/
|-- data_prepare.py
|-- main.py
|-- data/
`-- results/

Place the Pitprops correlation CSV at pitprops/data/pitprops_cor.csv。

Run the Pitprops experiment:

python data_prepare.py
python main.py

Colon

Purpose: run SPCA experiments on the Colon Cancer gene-expression benchmark.

The Colon Cancer dataset is a high-dimensional gene expression dataset. In our experiments, we load it from the R package HiDimDA and save it to a CSV file.

# Install the package if it is not available
if (!requireNamespace("HiDimDA", quietly = TRUE)) {
  install.packages("HiDimDA")
}
# Load the package and the Colon Cancer dataset
library(HiDimDA)
data(AlonDS)
# Convert the dataset to a data frame
colon_df <- as.data.frame(AlonDS)
# The downloaded dataset may contain both class labels and gene-expression features.
# In our sparse PCA experiments, we use only the gene-expression feature matrix.

# Remove the first two columns
colon_features <- colon_df[, -c(1, 2)]
# Save the downloaded dataset for later use (replace the path with your own path)
write.csv(colon_features, file = "download/path/to/colon_alon.csv", row.names = TRUE)

Key files:

colon/
|-- data_prepare.py
|-- main.py
|-- plot_tradeoff_loss.py
|-- plot_training_loss_from_logs.py
|-- data/
|-- results/
`-- figures/

Place the raw colon expression CSV at colon/data/colon_alon.csv.

Prepare the variance-screened covariance checkpoint:

python data_prepare.py

Run the colon experiment with variance-screened data:

python main.py

Optional: (reproducing Fig. 4 and Fig. 5 in paper)

python plot_compare_original_top_variance.py
python baselines.py
python plot_tradeoff_loss.py

Outputs

Each experiment writes generated data, checkpoints, logs, CSV results, and figures under its own folder:

<experiment>/data/
<experiment>/results/
<experiment>/figures/