Autism_FM

A research codebase for modeling Autism Spectrum Disorder (ASD) from genetic variant data using transformer-based aggregation over per-variant embeddings. The project combines large-scale SPARK genomic data, variant windows, and pretrained nucleotide embeddings to build sample-level representations for downstream classification and analysis.

Project Goal

The repository explores whether a transformer-based model can learn a biologically meaningful representation of each individual from the set of variants in their genome. The core idea is:

• Represent each variant with a fixed-length embedding (for example, a 1280-dimensional representation from a nucleotide model)
• Aggregate all variant embeddings for a participant into a sample-level representation
• Use the resulting representation for ASD vs control classification and exploratory genomic analysis

This is compared against simpler baselines such as mean pooling and logistic regression models.

Key Research Themes

• ASD risk modeling from variant-level genomic signals
• Transformer aggregators over variant sets
• Comparison of sample-level embeddings against mean-pooling baselines
• PCA / UMAP / clustering analysis for biological interpretation
• Candidate gene and SFARI-related analyses
• Investigation of genomic flank lengths and confounding factors

Repository Layout

Autism_FM/
├── .gitignore
├── requirements.txt
├── SFARI_genes/
│   └── ... SFARI gene tables and BED coordinates
├── SPARKDataRelease_2025-12-15/
│   └── supplemental data and metadata
├── analysis/
│   └── exploratory scripts for PCA, UMAP, confounding checks, etc.
├── category1_genes/
│   ├── analysis/
│   ├── experiments/
│   ├── preprocessing/
│   └── transformer_aggregation/
├── category1_test/
│   ├── experiments/
│   └── production/
├── classifier/
│   └── classification models and baselines
├── hierarchical_transformer/
│   └── hierarchical transformer experiments
├── mean_aggregation/
│   └── mean-pooling aggregation experiments
├── miscellaneous/
│   └── utility notes / helper scripts
├── pilot_study/
│   └── earlier pilot experiments and preprocessing
├── pilot_v1/
│   └── earlier variant modeling experiments
├── preprocessing/
│   └── variant-window generation and data preprocessing
├── project_memory/
│   └── summary notes, historical context, and project logs
├── transformer_aggregation/
│   └── main transformer pipeline
└── README.md

Main Pipeline

The primary active workflow is centered around the transformer_aggregation/ directory.

Typical pipeline:

1. Prepare variant windows and aligned embeddings
2. Create a processed dataset of per-sample variant embeddings
3. Train a transformer aggregator with a CLS token
4. Extract CLS representations for each sample
5. Evaluate the embeddings with downstream classifiers and analysis scripts

A high-level flow is:

NT embeddings / variant vectors
        ↓
create_transformer_dataset.py
        ↓
processed_embeddings_*.pt
        ↓
train.py
        ↓
extract.py
        ↓
representations + downstream analysis

Data Requirements

This repository is built around genomic datasets and embeddings that are not fully included in GitHub. Expected data include:

• SPARK participant metadata and labels
• Variant-level embedding arrays in HDF5 / PyTorch formats
• Aligned mappings between samples and variant records
• SFARI / category-based gene annotations

The project memory notes reference a dataset layout such as:

• SPARKDataRelease_2025-12-15/individuals_registration-2025-12-15.csv
• embeddings_flank1024/embeddings_flank1024_subset5k_merged.h5
• processed .pt files used by the transformer pipeline

You will need to supply or generate the local data files expected by the scripts before running the model end-to-end.

Environment Setup

Use the repository's pinned dependencies:

python -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt

The project depends on scientific Python tooling and deep learning libraries, including:

• PyTorch
• Transformers
• scikit-learn
• NumPy / pandas
• PyArrow / HDF5 tools
• Scanpy / seaborn / matplotlib
• Biopython and genomic utilities

Training the Transformer

From the main transformer directory:

cd transformer_aggregation
python train.py --no_ddp

For multi-GPU runs:

torchrun --nproc_per_node=4 train.py

The training script:

• builds a sample-level transformer over variant embeddings
• uses stratified cross-validation or train/validation splits
• trains with class-weighted cross-entropy and warmup/cosine scheduling
• saves the best validation checkpoint

Extracting Representations

After training, embeddings can be extracted for downstream tasks:

cd transformer_aggregation
python extract.py --checkpoint transformer_runs/best_fold1.pt --flank 1024

This produces sample representations that can be used for:

• logistic regression evaluation
• UMAP visualization
• clustering or biological interpretation
• comparison with mean-pooled baselines

Analysis Modules

The analysis/ directory contains scripts for:

• PCA and UMAP analysis
• variance and confounding investigation
• diagonal cluster diagnostics
• lineages / variant-level checks
• summary generation tools

Representative scripts include:

• analysis/pca_analysis.py
• analysis/umap_analysis.py
• analysis/pc1_confounding_audit.py
• analysis/check_spark.py
• analysis/lineage_evaluation.py

Baselines and Related Experiments

The repo includes multiple modeling variants beyond the main transformer pipeline:

• mean_aggregation/ — simple mean-pooling baseline models
• classifier/ — direct classification experiments
• hierarchical_transformer/ — hierarchical attention architecture variants
• pilot_study/ and pilot_v1/ — earlier exploratory runs
• category1_genes/ and category1_test/ — gene-set-specific experiments

Notes on Scope

This is a research repository, not a packaged library. The scripts are intended to be run in a local research environment with access to the required data and compute resources. Some experiments are highly data- and GPU-intensive and may require substantial RAM and/or multiple GPUs.

Recommended Workflow

1. Prepare or connect the SPARK data and embedding files locally
2. Install the Python environment from requirements.txt
3. Run preprocessing and dataset-building steps
4. Train the transformer on transformer_aggregation/train.py
5. Extract embeddings and evaluate with logistic regression or analysis scripts
6. Compare against baseline mean-pooled representations

Summary

Autism_FM is a genomics and machine learning research repository focused on learning ASD-related sample representations from per-variant embeddings. The main technical contribution is a transformer-based aggregation model that converts a set of variant embeddings into a single person-level representation for downstream phenotype classification and biomedical interpretation.

If you are onboarding to the project, start in transformer_aggregation/, then move into analysis/ for downstream interpretation.