SCARF-MCM

Masked contrastive representation learning for high-purity germanium detector waveforms and downstream likelihood-free inference.

Research-code status

This repository is a research project developed for HPGe waveform analysis for the LEGEND experiment. It documents the model and analysis implementation, but it is not currently a standalone reproducible package. Experiment data, metadata, configuration files, trained checkpoints, and parts of the surrounding software environment are not included.

Overview

SCARF-MCM contains a self-supervised learning pipeline for extracting representations from HPGe detector waveform time series. The pretraining approach is inspired by wav2vec 2.0, adapting masked latent modelling and contrastive representation learning from speech waveforms to HPGe detector signals.

The core pretraining method, implemented as Masked Contrastive HPGe-waveform Modeling (MCHM), combines:

• a convolutional waveform encoder
• span masking in the latent sequence
• a Transformer context network
• local or distributed InfoNCE objectives
• multi-GPU training with PyTorch Distributed Data Parallel

MCHM follows the general encoder–masking–context-network structure of wav2vec 2.0, but it is not a direct reproduction of the original method. In particular, MCHM uses the continuous encoder representations directly as contrastive targets rather than introducing the learned quantisation stage used by wav2vec 2.0.

The repository also contains downstream code for:

• binary neural ratio estimation
• frozen and partially fine-tuned representation transfer
• calibration and validation
• inference on experiment-tier LH5 data
• experimental contrastive ratio-estimation variants

Pretraining method

The pretraining pipeline operates on fixed-length waveform windows.

HPGe waveform
      |
      v
1D convolutional encoder
      |
      v
latent sequence z
      |
      +------------------------------+
      |                              |
      v                              v
span masking                   contrastive targets
      |
      v
learned mask token
      |
      v
Transformer context network
      |
      v
context representations c
      |
      v
InfoNCE loss at masked positions

1. Waveform encoder

The encoder is a residual one-dimensional convolutional network that progressively downsamples each waveform into a shorter latent sequence.

The downsampling stack:

• increases the channel dimension to the configured hidden size
• reduces the temporal resolution by approximately a factor of eight
• applies channel-wise layer normalisation
• uses residual pointwise-convolution blocks
• optionally applies activation checkpointing to selected layers

The encoder output is denoted by (z), where every latent token represents a local region of the input waveform.

2. Latent span masking

Random starting positions are sampled independently for each waveform. Contiguous spans of latent tokens are then replaced by a learned mask token.

Mask generation is deterministic with respect to the configured seed, epoch, distributed rank, and world size. This allows training to be resumed while preserving the intended randomisation scheme.

3. Transformer context network

The masked latent sequence is processed by a Transformer context network.

The implementation supports:

• multi-head self-attention
• rotary positional embeddings
• FlashAttention on compatible CUDA systems
• fused Transformer components
• attention and activation checkpointing
• a partial CPU fallback for model inspection and development

The Transformer output is denoted by (c). At masked positions, (c) must identify the corresponding clean encoder representation in (z).

4. Contrastive objective

Training uses an InfoNCE objective rather than waveform reconstruction.

For every masked position:

• the context representation is used as the anchor
• the clean latent at the same position is the positive target
• latent representations from other positions are used as negatives

The implementation provides two negative-sampling modes:

• Local negatives: negatives are selected from masked positions within the same waveform.
• Global negatives: negatives are sampled from the latent pool distributed across all DDP ranks.

The global variant uses custom autograd-aware gather operations so that sampled representations receive gradients across distributed workers.

Cosine similarity or an unnormalised dot product can be selected through the configuration.

Data pipeline

Source data

The data-curation pipeline was written for the LEGEND LH5 data format and its associated metadata hierarchy.

It reads waveform and physical-parameter tiers and can apply:

• file discovery from experiment metadata
• excluded-timestamp filtering
• detector-channel selection
• configurable physical cuts
• calibration-to-physics run mapping
• pole-zero or double-pole-zero correction
• baseline subtraction
• energy normalisation
• waveform alignment
• class or regression-target assignment

The resulting samples are written to NumPy memory-mapped files for subsequent training.

Stored sample format

For the pretraining and fine-tuning pipelines, each stored sample contains:

[waveform samples, alignment index, reconstructed energy, baseline mean, label]

The exact waveform length and selected metadata depend on the experiment configuration.

Runtime preprocessing

During pretraining, the collate function:

• selects a random fixed-length window from each waveform
• uses rank-, worker-, and epoch-dependent random seeds
• applies additional normalisation for large waveform overshoots
• returns waveform and metadata tensors
• prepares pinned host-memory tensors for non-blocking GPU transfer

Distributed training

The pretraining scripts are designed around PyTorch Distributed Data Parallel using the NCCL backend.

The training implementation includes:

• synchronous gradient averaging with DDP
• deterministic data sharding across ranks
• separate sharding across DataLoader workers
• staging of training data in job-local shared memory
• bfloat16 autocast for model execution
• FP32 contrastive-loss calculation
• fused mixed-precision LAMB optimisation
• checkpoint saving and continuation
• rank-aware random-number generation
• distributed loss aggregation

The original project was used for multi-node GPU training. Cluster launch configuration and scheduler scripts are environment-specific and are not included in this repository.

Downstream fine-tuning

Binary neural ratio estimation

The finetuning/nre modules reuse representations from the pretrained waveform model for a binary discrimination task.

The downstream model applies attention-based pooling to either:

• encoder representations (z) or
• Transformer context representations (c).

The fine-tuning pipeline supports:

• a completely frozen pretrained model
• fine-tuning selected Transformer blocks
• training the representation model from scratch
• a frozen-head warm-up phase
• independent learning rates for the prediction head and pretrained blocks
• stratified class batches
• reproducible train, validation, and test splits
• shared-memory or SSD-backed loading
• expected-calibration-error and classification diagnostics

The binary classifier logit is used as a neural-ratio-estimation-style decision score.

Experiment-tier inference

The downstream evaluator:

• restores the pretrained and fine-tuned checkpoints
• processes LH5 files in batches
• distributes files across available ranks
• evaluates waveform scores
• writes logits, timestamps, and event identifiers to LH5 output tables

Experimental contrastive ratio-estimation code

The finetuning/nrec directory contains experimental work on contrastive and energy-conditioned ratio estimation.

Some modules in this directory depend on code from related research repositories and are not currently self-contained. They are retained as part of the project history rather than presented as a ready-to-run public interface.

Repository structure

.
├── README.md
├── pyproject.toml
└── src/
    └── scarf_mcm/
        ├── curate_data.py          # Build NumPy memmaps from LH5 data
        ├── trainer.py              # Distributed MCHM pretraining
        ├── lr_sweep.py             # Distributed learning-rate sweep
        ├── evaluate.py             # Pretraining-checkpoint evaluation
        ├── logger.py               # CSV and optional offline W&B logging
        │
        ├── data/
        │   ├── db.py               # Experiment file and metadata discovery
        │   ├── dataloader.py       # LH5 loading and memmap construction
        │   ├── dataset.py          # Rank- and worker-sharded iterable data
        │   ├── collate_fn.py       # Waveform cropping and preprocessing
        │   └── utils.py            # Configuration and path dataclasses
        │
        ├── model/
        │   ├── mchm_model.py       # Complete masked contrastive model
        │   ├── encoder.py          # Convolutional waveform encoder
        │   ├── contextNet.py       # Transformer context network
        │   ├── losses.py           # Local and global InfoNCE objectives
        │   ├── allgather4ddp.py    # Gradient-aware distributed gathering
        │   ├── mha.py              # Modified FlashAttention MHA
        │   └── ...
        │
        └── finetuning/
            ├── nre/                # Binary NRE and transfer learning
            ├── nrec/               # Experimental contrastive NRE variants
            └── evaluate/           # Inference on experiment-tier data

Installation

The package metadata requires Python 3.9 or newer.

git clone https://github.com/nikolay327/scarf_mcm.git
cd scarf_mcm

python -m venv .venv
source .venv/bin/activate

pip install -e .

Runtime dependencies are not currently pinned or declared in pyproject.toml.

The research environment used packages including:

• PyTorch with CUDA and NCCL support
• NumPy
• SciPy
• pandas
• einops
• FlashAttention
• NVIDIA Apex
• lgdo
• dspeed
• optionally Weights & Biases

FlashAttention and Apex must be built against compatible PyTorch and CUDA versions.

Although parts of the attention implementation have CPU fallbacks, the training and evaluation entry points are primarily GPU-oriented.

Required external inputs

The public repository does not include:

• waveform data
• LH5 files
• experiment metadata
• pole-zero parameter databases
• calibration-to-physics mappings
• excluded-timestamp lists
• training or evaluation JSON configurations
• trained model checkpoints

The original working-directory layout was approximately:

WORKING_DIR/
├── meta/
│   ├── file-db.json
│   └── experiment-specific metadata
│
└── trained/
    └── EXPERIMENT/
        └── MODEL_NAME/
            └── VERSION/
                ├── MODEL_NAME_VERSION.json
                ├── checkpoints/
                ├── evaluation/
                └── lr_sweep/

Curated NumPy memory maps are stored separately under MMAP_DIR.

Configuration

Training configurations are JSON files with two kinds of entries:

{
  "model_config": {
    "...": "fields consumed by MCHMConfig"
  },
  "dataset_name": {
    "...": "fields consumed by the LH5 memmap builder"
  }
}

model_config controls, among other things:

• waveform sequence length
• encoder and Transformer dimensions
• number of Transformer layers and attention heads
• masking span and number of masked regions
• InfoNCE temperature and negative sampling
• local versus global contrastive loss
• batch size and optimiser parameters
• random seeds
• data-sampling cycles
• optional activation checkpointing and fused operations

Dataset entries describe LH5 paths, channels, groups, physical cuts, waveform columns, preprocessing parameters, labels, and output memory-map locations.

No generic example configuration is currently included because the existing configurations depend on private experiment metadata.

Running the pipeline

The commands below show the interfaces implemented by the scripts. They are templates rather than a fully reproducible public workflow.

1. Curate the waveform data

python -m scarf_mcm.curate_data \
    <experiment> \
    <model_name> \
    <version> \
    <working_dir> \
    <mmap_dir> \
    <training_config.json> \
    <num_dataloader_workers>

The current curation script contains experiment-specific run selections that should be replaced or moved into configuration before general use.

2. Run a learning-rate sweep

The distributed scripts expect a job-local shared-memory directory:

export JOB_SHMTMPDIR=/path/to/job-local/shared-memory

Then launch the sweep through torchrun or the equivalent scheduler integration:

torchrun \
    --standalone \
    --nproc_per_node=<number_of_gpus> \
    --module scarf_mcm.lr_sweep \
    <experiment> \
    <model_name> \
    <version> \
    <working_dir> \
    <mmap_dir> \
    <training_config.json>

Results are written as CSV files or optional offline W&B runs.

3. Pretrain the model

torchrun \
    --standalone \
    --nproc_per_node=<number_of_gpus> \
    --module scarf_mcm.trainer \
    <starting_epoch> \
    <final_epoch> \
    <experiment> \
    <model_name> \
    <version> \
    <working_dir> \
    <mmap_dir> \
    <training_config.json> \
    <cache_dir>

Use starting_epoch=1 to start from scratch. A later starting epoch restores the preceding checkpoint.

For multi-node execution, replace --standalone with the rank and rendezvous settings supplied by the cluster scheduler.

4. Evaluate pretraining checkpoints

torchrun \
    --standalone \
    --nproc_per_node=<number_of_gpus> \
    --module scarf_mcm.evaluate \
    <first_checkpoint> \
    <last_checkpoint> \
    <experiment> \
    <model_name> \
    <version> \
    <working_dir> \
    <mmap_dir> \
    <training_config.json> \
    <evaluation_config.json> \
    <output_name> \
    <cache_dir>

This evaluates InfoNCE loss on the held-out fraction of the curated waveform data.

5. Fine-tune a binary NRE model

The fine-tuning entry point also expects rank-related environment variables, so it can be launched through torchrun with one process:

torchrun \
    --standalone \
    --nproc_per_node=1 \
    --module scarf_mcm.finetuning.nre.trainer \
    <starting_epoch> \
    <final_epoch> \
    <experiment> \
    <model_name> \
    <version> \
    <working_dir> \
    <mmap_dir> \
    <training_config.json> \
    <keep_data_in_ssd:0_or_1> \
    <memory_mode> \
    <cache_dir>

The fine-tuning configuration specifies the pretrained model version and checkpoint, whether its parameters are frozen, and which Transformer blocks should be updated.

Outputs

Depending on the stage, the pipeline produces:

• PyTorch checkpoints
• optimiser state for training continuation
• InfoNCE training histories
• fine-tuning loss and accuracy histories
• calibration metrics
• CSV evaluation summaries
• optional offline W&B logs
• LH5 tables containing downstream model scores and event identifiers

Large data products, checkpoints, images, notebooks, and LH5 files are intentionally excluded from version control.

Reproducibility

The implementation includes explicit random seeds for:

• data mixing
• dataset shuffling
• waveform cropping
• mask generation
• train, validation, and test splitting
• epoch- and rank-specific sampling

However, complete reproduction of the original analysis additionally requires the original experiment data, metadata, JSON configurations, software versions, and HPC launch environment.

Intended use

This repository is primarily intended to:

• document the implementation used during the research project
• demonstrate masked contrastive modelling for detector waveforms
• preserve the distributed-training and downstream-inference code
• provide a technical reference for related scientific machine-learning work

It is not currently intended to provide a plug-and-play model for arbitrary waveform datasets.

References

1. A. Baevski, H. Zhou, A. Mohamed, and M. Auli,
   wav2vec 2.0: A Framework for Self-Supervised Learning of Speech Representations,
   Advances in Neural Information Processing Systems 33, 2020.