GNN4Colliders / docs /migration.md
ho22joshua's picture
feat: add initial root transformer model
8344837
|
Raw
History Blame Contribute Delete
2.58 kB

Migration and model-family roadmap

Frozen ROOT-GNN baseline

The root-gnn-parity-baseline tag records the completed migration of the active ROOT-GNN behavior into src/gnn4colliders.models.root_gnn. The campaign covered full event preprocessing and graph parity, binary objectives and metrics, deterministic fine-tuning, full-split training, checkpoint reload and resume, reproducibility, and serialized graph-cache checks.

The historical implementation is no longer in the active source tree. Its observable behavior is represented by committed fixtures, tests, and the one-way checkpoint/metadata compatibility adapters. New work must not add imports from historical implementation paths.

Shared contracts

New model families should consume these boundaries:

  • EventSample and named EventMetadata from data;
  • shared collider feature builders from features;
  • a representation-specific sample/batch type from the relevant adapter;
  • task-owned loss, score, prediction, and metric semantics;
  • the shared Trainer, checkpoint, reproducibility, and inference APIs.

The graph path is the current ROOT-GNN representation. A sequence or token model should add a separate representation boundary rather than placing sequence behavior in graph modules or generic data code.

Transformer model-family milestone

The first root_transformer vertical slice is implemented with:

  1. Define a small SequenceSample contract and deterministic fixture.
  2. Implement token construction using shared event/features infrastructure.
  3. Add the transformer model under models/root_transformer/.
  4. Connect it to the existing binary task and trainer on a tiny fixture.
  5. Add checkpoint, prediction, and reproducibility tests.

The first implementation is available as model=root_transformer/transformer. It can train on an existing graph cache by adapting node features into ordered sequence tokens; this is a migration bridge while a native sequence cache and DDP sequence loader are evaluated.

Native sequence-cache storage and distributed sequence loading remain follow-up work; do not generalize shared interfaces until those use cases require it.

Validation requirements

Every new model family must provide unit tests for its representation and model, a small end-to-end integration test, checkpoint reload coverage, and a deterministic repeatability check. Scientific behavior that is intentionally shared with ROOT-GNN should be compared against the frozen reference fixture; architecture-specific behavior should have its own reference outputs.