Implement Multi-Track Reconstruction Using Auto-Regressive Single-Track Finder

[dhshin04]

Summary

The goal of this task is to extend the existing single track reconstruction model into a fully functional multi-track reconstruction framework. The multi-track finder will follow an auto-regressive approach, repeatedly invoking the single track finder on a modified event representation until a learned confidence or quality metric indicates that no valid tracks remain.

This issue covers both model-level changes and algorithm-level logic needed to make the multi-track finder robust, accurate, and stable.

---

Motivation

The current single track finder is capable of identifying one track per event, but realistic events can contain zero, one, or many tracks. A multi-track finder enables full event reconstruction by iteratively extracting tracks while suppressing already identified signals.

A key challenge is defining when the auto-regressive process should stop. This requires a principled and learnable stopping criterion rather than a fixed iteration count or heuristic threshold.

---

High-Level Approach

1. Introduce a confidence head in the single track finder model that predicts whether a valid track exists in the current event state.
2. Use this confidence score as the stopping criterion for the auto-regressive multi-track finder.
3. After each successful track extraction, update the event representation by subtracting the contribution of the identified track in a soft and detector-aware manner.
4. Repeat until the confidence score indicates no valid tracks remain.

---

Current Status

Components Already Implemented

• Single track finder model
• High-level multi-track finder base class with run and evaluation hooks

---

Tasks to Be Completed

1. Add Confidence Head to Single Track Finder

Proposal A: Stop-or-Go Confidence Head (Event Level)

Location:

• QTracker_training/models/TrackFinder.py in dev branch

Description:
Add an additional output head to the existing single track finder model that predicts a confidence score indicating whether the current track proposal is valid.

Target Definition:

• Binary confidence indicator:
  • 1: There is at least one valid track present in the event
  • 0: There are no remaining tracks in the event

Training Strategy:

• Update the dataset to support a variable number of ground truth tracks per event (0, 1, 2, ...).
• Ground truth confidence label:
  • 1 if number of GT tracks >= 1
  • 0 if number of GT tracks = 0
• Use Binary Cross Entropy (BCE) loss for the confidence head.
• Carefully weight the confidence loss relative to existing losses to preserve single track reconstruction performance.

Evaluation Requirements:

• Validate that single track performance (residuals, momentum, q-metric) is not degraded.
• Evaluate confidence head quality using:
  • Accuracy
  • Precision and recall
  • Correlation with reconstruction quality and residuals

Proposal B: Track Correctness Confidence Head (Prediction Level)

Purpose
Estimate whether the produced track is a valid, well reconstructed physical track.
This corresponds to the intuitive notion of a confidence score.

Output

• Single scalar confidence logit (sigmoid)

Target Definition
Identify the hit overlap between the predicted track and the best matching ground truth track, using either a binary threshold or the overlap score itself as a soft label.

Let:

• P be the set of predicted hits
• Gk be the set of hits for GT track k

Define hit overlap:

• Precision = |P ∩ Gk| / |P|
• Recall = |P ∩ Gk| / |Gk|
• F1 = 2PR / (P + R)

Match prediction to best GT track:

• k* = argmax_k F1(P, Gk)

Label options:

• Soft label (recommended): y = F1(P, Gk*)

Loss

• BCE (hard or soft target) or MSE (soft target) between soft label calculated above and model's confidence score prediction
• Weighted conservatively relative to reconstruction losses

Evaluation

• Accuracy and precision recall (hard labels)
• Correlation with F1, residuals, and momentum error
• Reconstruction quality vs confidence threshold

Notes
Ground truth overlap is used only to generate training targets.
At inference, confidence is predicted solely from the model’s internal track representation.

---

2. Implement Hit Subtraction Logic in Multi-Track Finder

Location:

• QTracker_main/src/models/multi_track_finder.py in dev branch

Problem:
The current multi-track finder does not remove or suppress hits belonging to already identified tracks. Without this, repeated runs can produce identical or confused track proposals.

Naively zeroing out hits is insufficient due to:

• Error propagation
• Overlapping tracks
• Shared detector elements

Proposed Solution:
Implement soft subtraction based on detector-wise softmax scores.

Algorithm:
For each detector in an event:

1. Identify the detector element ID predicted to be a signal hit for the current track.
2. Retrieve the highest softmax probability for that detector.
3. Subtract this probability from the corresponding hit value.

Example:

• Detector: I
• Event: J
• Element ID: K
• Predicted signal hit softmax: 0.9
• Original hit value: 1.0
• Updated hit value: 1.0 - 0.9 = 0.1

Safety Constraint:

• Enforce a lower bound of zero to avoid negative values:
  new_value = max(0.0, original_value - softmax_value)

This handles cases where overlapping mu+ and mu- tracks share the same detector element.

---

Codebase Overview

Branch:

• We will use the dev branch in QTracker_basic repository as our main version control for this feature.
• Try not to commit directly to dev. Instead, create a new branch, commit your changes, then submit a PR. Ideally, code should only be merged to dev after one review.

Folder Structure:

1. QTracker_basic/ / QTracker_prod/

   • Imported directly from main
   • Contains legacy learning scripts
   • Not used for new development

2. QTracker_training/

   • Contains training scripts and latest single track finder model
   • Confidence head implementation goes here: models/TrackFinder.py
   • Includes momentum and q-metric training pipelines

3. QTracker_main/

   • Contains skeleton implementation of the multi-track finder
   • Required setup:
     • Create checkpoints/ subfolder and add model checkpoints
     • Install dependencies:
       pip install uv
uv sync
conda install -c conda-forge root
   • Hit subtraction logic implemented in:
     • src/models/multi_track_finder.py

---

Acceptance Criteria

• Single track finder includes a trained and validated confidence head.
• Confidence score reliably indicates presence or absence of valid tracks.
• Multi-track finder correctly suppresses previously identified hits using soft subtraction.
• Auto-regressive process terminates based on confidence score without manual iteration limits.
• No regression in single track reconstruction performance.
• Multi-track reconstruction produces distinct, non-duplicated tracks.

---

Notes

• Loss weighting and stability of training are critical.
• Special care is required to handle overlapping tracks and shared detector hits.
• This implementation forms the foundation for scalable, event-level reconstruction.