Multi-Modal Track Reconstruction using Graph Neural Networks at Belle II

High backgrounds and detector ageing impact the track finding in the Belle II central drift chamber, reducing both track purity and track efficiency in events. This necessitates the development of new track finding algorithms to mitigate detector performance degradation. Building on our previous success with an end-to-end multi-track reconstruction algorithm for the Belle II experiment at the SuperKEKB collider (arXiv:2411.13596), we have extended the algorithm to incorporate inputs from both the drift chamber and the silicon vertex tracking detector, creating a multi-modal network. We employ graph neural networks to handle the irregular detector structure and object condensation to address the unknown, varying number of particles in each event. This approach simultaneously identifies all tracks in an event and determines their respective parameters. We demonstrate the algorithm’s effectiveness using a realistic full detector simulation, which incorporates beam-induced backgrounds and noise modelled from actual collision data. The simultaneous reconstruction of the information from the two detectors yields a track efficiency improvement from 48.0 % to 74.7 % for uniformly displaced particles up to 100 cm, while increasing the track purity by 5.5 percentage points. We provide a detailed comparison of its track-finding performance against the current Belle II baseline across various event topologies.

💡 Research Summary

The paper presents a unified graph‑neural‑network (GNN) based track finder for the Belle II experiment, called the BA T Finder, which simultaneously processes hits from the central drift chamber (CDC) and the silicon vertex detector (SVD). Traditional Belle II reconstruction proceeds in several stages: CDC‑only seeding, extrapolation to the SVD, a second SVD‑only seeding, and finally a backward extrapolation to attach remaining CDC hits. This staged approach suffers from mismatches between the two detectors, especially under high beam‑induced background and detector ageing, leading to reduced track efficiency and purity.

The BA T Finder replaces the entire staged pipeline with a single inference step. All CDC and SVD hits are represented as nodes in a single graph. Detector‑specific features (CDC drift time, wire geometry, time‑over‑threshold, amplitude; SVD global coordinates, deposited charge, timing, signal‑to‑noise ratios) are concatenated into a fixed‑length feature vector, zero‑padded where a feature does not exist, and a binary flag indicates the detector origin. These vectors are first projected into a common latent space with fully‑connected layers and then processed by a series of GravNet blocks. GravNet dynamically constructs edges based on learned distances, allowing the network to learn local relationships without any explicit ordering of layers or detector types.

Training uses the object‑condensation loss. For each hit the network predicts (i) a point in a learned “cluster space”, (ii) a condensation score β that identifies candidate track seeds, and (iii) initial track parameters (starting position, momentum, charge). The loss combines an attractive term that pulls hits belonging to the same true particle together in cluster space, a repulsive term that pushes hits from different particles apart, and a noise‑suppression term that penalises background hits. The regression of track parameters is weighted by β, so only high‑β (seed) hits are forced to produce accurate parameter estimates. Because SVD hits are far fewer than CDC hits, their contribution to the loss is up‑weighted by a factor of four to avoid bias.

The model is trained on a realistic full‑detector simulation that includes beam‑induced backgrounds and noise modeled from actual Belle II data. Training samples uniformly cover charge, momentum (0.05–6 GeV/c), and vertex displacement up to 100 cm, ensuring exposure to both prompt and highly displaced tracks. Performance is evaluated on events containing two displaced muons per vertex, with vertex radii ranging from a few centimeters to 100 cm.

Results show a substantial improvement over the baseline and the CDC‑only CAT Finder. The baseline algorithm achieves 48.0 % efficiency and 92.2 % purity for displaced muons. The CAT Finder (CDC‑only) raises efficiency to 68.5 % and purity to 93.8 %. The BA T Finder reaches 74.7 % efficiency and 97.6 % purity, a gain of 5.5 percentage points in purity and a 26.7 percentage‑point increase in efficiency relative to the baseline. The improvement is most pronounced at small to intermediate transverse displacements (< 15 cm), where the inclusion of SVD hits directly in the graph allows the network to resolve tracks before they enter the CDC volume. At larger radii, the use of CDC time‑over‑threshold information further suppresses background, maintaining the efficiency advantage across the full displacement range.

Because the network produces a complete ordered list of hits and an initial parameter estimate for each track, the output can be fed directly into the standard Belle II GENFIT2 Kalman filter for final fitting. The single‑step approach eliminates duplicate track creation and mismatched hit assignments that are typical of the staged pipeline, thereby delivering cleaner track collections.

In conclusion, the BA T Finder demonstrates that a multi‑modal GNN can effectively handle heterogeneous detector geometries, learn to cluster hits belonging to the same particle, and predict accurate track seeds in a single forward pass. This approach not only restores and surpasses the performance of the legacy reconstruction under realistic background conditions but also provides a scalable framework for future high‑luminosity operation of Belle II and potentially other experiments with mixed detector technologies. Future work may explore real‑time deployment, hardware acceleration, and adaptation to other experiments such as LHCb or Mu3e.