Learning on ancestral recombination graphs with transformers

ARGformer is a transformer encoder based on ModernBERT for Ancestral Recombination Graph (ARG) data. It uses the FlexBERT architecture with YAML-based configuration.

The codebase builds upon MosaicBERT and the fork with Flash Attention 2 under Apache 2.0 license.

For ModernBERT details, see the release blog post.

Setup

Install pixi

If you do not have pixi installed or are not familiar with it, pixi is a fast alternative to tools like conda and pip for managing environments and Python packages. You can install it with:

curl -fsSL https://pixi.sh/install.sh | sh

Once pixi is installed, create and enter the environment with:

pixi install
pixi shell

Verify that the solve produced a CUDA-enabled PyTorch build before installing FlashAttention:

python -c "import torch; print(torch.__version__); print(torch.version.cuda); print(torch.cuda._is_compiled())"

The base pixi environment does not include flash-attn. Install it separately inside the pixi shell if you want FlashAttention 2 support:

pip install "flash_attn==2.6.3" --no-build-isolation

For H100 GPUs, optionally install Flash Attention 3:

git clone https://github.com/Dao-AILab/flash-attention.git
cd flash-attention/hopper
python setup.py install

Overview

ARGformer supports:

• Pretraining: Masked language modeling on ARG sequences
• Contrastive Learning: Fine-tuning for retrieval and similarity tasks
• Embeddings: Extracting embeddings for downstream analysis
• Retrieval: Finding similar sequences in large corpora

Data Format

ARG data structure:

/path/to/arg/data/
├── train/
│   ├── tokenized_train_sequences_and_vocab.pkl
│   └── labels.pkl  # Optional: for contrastive learning
└── val/
    ├── tokenized_val_sequences_and_vocab.pkl
    └── labels.pkl  # Optional: for contrastive learning

The ARGDataset class supports pretokenized sequences with vocabulary mappings for node IDs and special tokens ([PAD], [CLS], [SEP]).

Workflow

1. Prepare Pretraining Data

Extract sequences from tree files using src/data/prepare_data_pretrain.py:

python src/data/prepare_data_pretrain.py

Edit the script to configure input paths, output directory, and train/val split.

2. Pretraining

Configure yamls/mlm.yaml with dataset paths and model parameters, then run:

composer main.py yamls/mlm.yaml

3. Contrastive Fine-tuning

Configure yamls/contrastive.yaml with pretrained checkpoint path and run:

python sequence_contrastive.py yamls/contrastive.yaml

4. Extract Embeddings

Use the checkpoint's training YAML, with paths updated to your local files. Keep the same model architecture, node vocabulary/IDs, sequence length, and skip_extant_tokens setting used during training. Node tokens are specific to the ARG used to build that vocabulary.

For a self-supervised masked language modeling (MLM) checkpoint:

python embeddings.py \
  --config yamls/mlm.yaml \
  --checkpoint /path/to/checkpoint.pt \
  --include_splits train,val \
  --pooling first_node --batch_size 256 --max_samples 0 \
  --output_path outputs/paths.pt

This encodes unmasked paths in evaluation mode and extracts final encoder hidden states, before the MLM prediction head. Labels are not required. --max_samples 0 includes all paths; a positive value caps the total number of rows (the default is 10,000). Extraction collects the result in memory.

--pooling first_node selects the first node after [CLS]. With skip_extant_tokens: true, the input is [CLS], first ancestor, …, [PATH-SEP], so this returns the first-ancestor embedding. Other options are cls (position 0), mean (all non-padding tokens, including special tokens), and all (all token vectors). For a classification/contrastive checkpoint, use its training YAML; head and head_norm use the trained pooling head, with the latter applying L2 normalization. These two options do not apply to MLM checkpoints.

The saved PyTorch dictionary contains embeddings ([paths, hidden_size], or [paths, sequence_length, hidden_size] for all), aligned extant_node_ids (raw leaf IDs), input_ids, attention_masks, and extraction metadata.

Aggregate before visualization

Self-supervised embeddings from real ARGs can be noisy at the local path level. Averaging embeddings from multiple unique marginal paths of the same haplotype or individual often reveals clearer structure in visualizations.

The preparation script already deduplicates full leaf-to-root paths. If using another preparation workflow, deduplicate those paths before extraction, including across the selected splits. Do not deduplicate by embedding values or truncated model inputs: distinct full paths can yield identical inputs. The following helper gives each unique path equal weight within a haplotype; it does not weight by genomic span or perform deduplication itself.

python aggregate_embeddings.py --input outputs/paths.pt --output outputs/haplotypes.pt

For individuals, provide a CSV with extant_node_id,individual_id columns, one row per haplotype (both haplotypes for diploid individuals), then run:

python aggregate_embeddings.py --input outputs/paths.pt \
  --individual-mapping /path/to/node_to_individual.csv --output outputs/individuals.pt

Individual vectors average the haplotype means, giving each haplotype equal weight even when path counts differ. The helper requires embeddings for all mapped haplotypes of each represented individual. Outputs contain aligned IDs, path_counts, and (for individuals) haplotype_counts. Use paths from a single ARG/node namespace and a single checkpoint in each aggregation.

For example, project the aggregated vectors onto two principal components using PyTorch and plot them with Matplotlib:

import matplotlib.pyplot as plt
import torch

data = torch.load("outputs/haplotypes.pt", map_location="cpu", weights_only=True)
x = data["embeddings"].float()  # Also usable as downstream model features.
centered = x - x.mean(dim=0)
_, _, axes = torch.linalg.svd(centered, full_matrices=False)
xy = (centered @ axes[:2].T).numpy()  # Requires at least two rows and features.
plt.scatter(xy[:, 0], xy[:, 1], s=12)
plt.xlabel("PC1")
plt.ylabel("PC2")
plt.show()

5. Retrieval

python retrieve.py [arguments]

See the script for usage examples.

Configuration

Training uses composer with YAML configuration files in yamls/:

• mlm.yaml: Pretraining configuration
• contrastive.yaml: Contrastive learning configuration

Key configuration sections:

• model: Model architecture and checkpoint paths
• train_loader / eval_loader: Dataset paths and data loading settings
• optimizer / scheduler: Training hyperparameters
• loggers: WandB logging configuration

Citation

If you use ARGformer in your research, please cite our paper:

@article{bonet2026argformer,
  author={Bonet, David and Shanks, Cole and Cara, Marçal Comajoan and Abante, Jordi and Ioannidis, Alexander G.},
  title={{ARGformer}: learning on ancestral recombination graphs with transformers},
  journal={Bioinformatics},
  year={2026},
  volume={42},
  pages={btag438},
  doi={10.1093/bioinformatics/btag438},
  url={https://doi.org/10.1093/bioinformatics/btag438}
}