MaskFormer-Based Reconstruction

• MaskFormer-based reconstruction is a query-driven paradigm that employs a Transformer encoder and a MaskFormer-style decoder to generate instance masks and per-object predictions.
• It combines hard segmentation for track building with soft compositional assignments for particle clustering, enabling joint optimization of hit assignments and parameter regressions.
• The approach utilizes permutation-invariant Hungarian matching and scalable attention mechanisms to achieve high efficiency and low false rates in complex detector environments.

MaskFormer-based reconstruction is a query-based detector-reconstruction paradigm in which a Transformer encoder produces contextualized representations of detector objects and a MaskFormer-style decoder maps a fixed set of learned queries to per-object masks together with task-specific predictions. In charged particle tracking, this formulation is used to jointly optimise hit assignments and the estimation of charged-particle properties, with a two-stage Transformer hit filtering network followed by a MaskFormer reconstruction model (Stroud et al., 2024). In particle flow, the same general pattern is combined with incidence matrix supervision, so that each query produces a soft assignment over reconstructed tracks or calorimeter topoclusters, a particle-type prediction, and kinematic corrections (Kobylianskii et al., 27 Aug 2025). Across these settings, MaskFormer-based reconstruction recasts combinatorial reconstruction as set-to-set prediction with permutation-invariant matching.

1. Task formulation and reconstruction targets

In the TrackML tracking system, the MaskFormer reconstruction network takes filtered hits, typically $6\,\mathrm{k}$–$12\,\mathrm{k}$ per event, and outputs $N$ fixed object queries mapped to binary masks over hits together with track properties. The target is simultaneous instance segmentation of detector hits into tracks and regression of per-track parameters $(p_x,p_y,p_z,v_z)$, from which $p_T$, $\eta$, $\phi$, and total $|p|$ are derived (Stroud et al., 2024).

In GLOW, the inputs are reconstructed tracks or calorimeter topoclusters embedded into feature vectors, and the final query embeddings are sent to three parallel heads: a mask head yielding per-input soft mask logits $m_{a,i}$, a classification head producing particle type logits $c_a \in \mathbb{R}^3$ for photon, electron, or hadron, and a regression head producing kinematic corrections $12\,\mathrm{k}$0 for $12\,\mathrm{k}$1 (Kobylianskii et al., 27 Aug 2025).

This suggests that MaskFormer-based reconstruction is not restricted to a single detector domain. The same query-mask abstraction can be instantiated for charged-particle track building or for particle clustering and kinematic refinement, provided that detector objects can be represented as an unordered set and supervised through instance-level assignments.

2. Shared architectural pattern

Both reported systems use an encoder-decoder Transformer with learned queries and mask-guided cross-attention, but they differ in detector inputs, decoder depth, and the semantics of the mask.

Aspect	Tracking system	GLOW
Input objects	3D pixel-detector hits	Reconstructed tracks or calorimeter topoclusters
Encoder	Transformer encoder with sliding-window attention in $12\,\mathrm{k}$2	6 identical Transformer encoder layers
Decoder queries	$12\,\mathrm{k}$3 learned $12\,\mathrm{k}$4-D vectors	$12\,\mathrm{k}$5 learnable particle queries
Decoder depth	8 layers	4 layers
Output masks	Binary masks over hits	Soft assignments over input objects
Additional heads	Objectness classification and regression	Particle type classification and regression

In the tracking model, raw hits are mapped to a $12\,\mathrm{k}$6-D embedding, processed by a Transformer encoder with sliding-window attention in $12\,\mathrm{k}$7, classified as signal or noise, and then passed to a second stage that reuses the same encoder with halved window size before an eight-layer MaskFormer decoder (Stroud et al., 2024). Per-hit features include Cartesian coordinates $12\,\mathrm{k}$8, cylindrical coordinates $12\,\mathrm{k}$9, conformal mapping $N$0, and cluster properties such as charge fraction, size, and shape. Positional encodings are standard sine/cosine on $N$1 and $N$2, with cyclic positional encoding on $N$3.

In GLOW, each input object is embedded by a small per-object MLP, encoded by $N$4 Transformer encoder layers, and decoded by $N$5 Transformer decoder layers using masked cross-attention. At decoder layer $N$6, the dynamic mask is

$N$7

and the masked cross-attention uses $N$8 inside the attention logits (Kobylianskii et al., 27 Aug 2025).

A plausible implication is that MaskFormer-based reconstruction gains its flexibility from separating global contextual encoding from query-specialized decoding. The encoder models long-range dependencies among detector objects, while the decoder progressively specializes queries into object hypotheses whose support is expressed as a mask.

3. Mask semantics, query specialization, and assignment structure

In the tracking model, the number of queries is set to the maximum reconstructable tracks in the training events, for example $N$9 for $(p_x,p_y,p_z,v_z)$0 MeV, and queries are initialized as learned $(p_x,p_y,p_z,v_z)$1-D vectors $(p_x,p_y,p_z,v_z)$2. Each query is intended to specialise to one track. After self-attention and mask-guided cross-attention in each decoder layer, a small MLP maps the updated query $(p_x,p_y,p_z,v_z)$3 to a mask token $(p_x,p_y,p_z,v_z)$4, and raw mask logits are computed by $(p_x,p_y,p_z,v_z)$5. After sigmoid activation and thresholding at $(p_x,p_y,p_z,v_z)$6, the binary mask $(p_x,p_y,p_z,v_z)$7 is obtained (Stroud et al., 2024).

In GLOW, the mask head closely follows MaskFormer, but the output remains a soft assignment: $(p_x,p_y,p_z,v_z)$8 where $(p_x,p_y,p_z,v_z)$9 is produced by dotting a query-derived prototype vector with the encoded feature of input object $p_T$0. The predicted soft assignment is trained to approximate the ground-truth incidence matrix

$p_T$1

defined by fractional energy assignments (Kobylianskii et al., 27 Aug 2025).

A common misconception is that query-mask reconstruction enforces exclusive assignments. In the TrackML system, a hit may be assigned to multiple queries, explicitly allowing shared hits. In GLOW, the mask is inherently fractional and is used to build kinematics through incidence-weighted sums such as

$p_T$2

with $p_T$3 refined by adding $p_T$4 (Stroud et al., 2024, Kobylianskii et al., 27 Aug 2025).

Taken together, these formulations indicate that MaskFormer-based reconstruction can represent both hard-ish segmentation and soft compositional assignment. The former is useful when track candidates are naturally expressed as subsets of hits, while the latter is useful when a detector object can contribute to multiple overlapping particles.

4. Matching and multi-task optimization

Both systems use Hungarian matching to align unordered query predictions with unordered ground-truth objects, thereby making the training objective permutation-invariant.

For tracking, the total loss is computed after optimal bipartite assignment between predicted queries and ground-truth tracks. With $p_T$5 object classes and $p_T$6 real track plus null, the classification loss is categorical cross-entropy for track versus null,

$p_T$7

The mask loss combines Dice and binary focal losses,

$p_T$8

$p_T$9

$\eta$0

The regression loss is

$\eta$1

and the total objective is

$\eta$2

with auxiliary losses on intermediate decoder outputs and $\eta$3 (Stroud et al., 2024).

For GLOW, supervision is decomposed into segmentation, incidence, and flow terms. The incidence loss is

$\eta$4

After Hungarian matching, the segmentation term for a matched pair $\eta$5 consists of cross-entropy over object assignments and a soft Dice loss,

$\eta$6

$\eta$7

combined as

$\eta$8

The flow term is

$\eta$9

with

$\phi$0

GLOW identifies permutation invariance through Hungarian matching and query-based decoding as a reason the method removes the need for hand-tuned clustering heuristics (Kobylianskii et al., 27 Aug 2025).

5. Training regimes and reported empirical performance

The tracking system is trained on TrackML inner-pixel only data, with 4 barrel and 7 disks per side and on average 4–6 hits per track. Only particles with at least 3 pixel hits, $\phi$1, and $\phi$2 of $\phi$3, $\phi$4, or $\phi$5 MeV are targets. Hits are randomly shuffled in $\phi$6, then sorted by $\phi$7 for sliding-window attention. The encoder uses 12 layers, or 8 for the 1 GeV model, with model dimension $\phi$8, FFN dimension $\phi$9, and window size $|p|$0 for filtering or $|p|$1 for reconstruction. Training uses batch size 1 event for 30 epochs on an NVIDIA A100, with AdamW at typical settings for vision Transformers, $|p|$2 and weight decay $|p|$3 (Stroud et al., 2024).

The reported inference times are approximately $|p|$4 ms/event for hit filtering on about $|p|$5 hits, and $|p|$6 ms for $|p|$7 GeV to $|p|$8 ms for $|p|$9 MeV in the MaskFormer reconstruction stage. The combined time is approximately $m_{a,i}$0 ms/event for the $m_{a,i}$1 MeV configuration, with scaling linear in input hit count due to sliding-window attention. The best $m_{a,i}$2 MeV model reaches plateau efficiency $m_{a,i}$3 with fake rate $m_{a,i}$4, $m_{a,i}$5, and perfect-match efficiency $m_{a,i}$6 at high $m_{a,i}$7. Reported regression residuals include $m_{a,i}$8 resolution at few-percent bias/RMS, $m_{a,i}$9 and $c_a \in \mathbb{R}^3$0 residuals of order $c_a \in \mathbb{R}^3$1, and $c_a \in \mathbb{R}^3$2 RMS $c_a \in \mathbb{R}^3$3 (Stroud et al., 2024).

GLOW is trained on 1 M simulated $c_a \in \mathbb{R}^3$4 dijet events at $c_a \in \mathbb{R}^3$5 GeV for the CLIC detector, using PandoraPFA-reconstructed tracks and topoclusters within $c_a \in \mathbb{R}^3$6. Track features are $c_a \in \mathbb{R}^3$7, cluster features are $c_a \in \mathbb{R}^3$8, and augmentations include random $c_a \in \mathbb{R}^3$9 rotations, small Gaussian noise on $12\,\mathrm{k}$00 with $12\,\mathrm{k}$01, and random masking of $12\,\mathrm{k}$02 of inputs. Optimization uses AdamW with initial $12\,\mathrm{k}$03, weight decay $12\,\mathrm{k}$04, linear warm-up for 5 epochs, cosine decay to 0 over 200 epochs, batch size 1024 events across 2 H100 GPUs, dropout $12\,\mathrm{k}$05 in each FFN, and label smoothing $12\,\mathrm{k}$06 for type classification (Kobylianskii et al., 27 Aug 2025).

On evaluation, GLOW reports event-level metrics based on missing $12\,\mathrm{k}$07 residual and scalar sum $12\,\mathrm{k}$08, jet-level metrics based on anti-$12\,\mathrm{k}$09 $12\,\mathrm{k}$10 jets, and clustering purity and efficiency. In all tested metrics, GLOW outperforms HGPflow by approximately $12\,\mathrm{k}$11 in jet energy resolution and yields narrower $12\,\mathrm{k}$12 residuals; it also outperforms MLPF by even larger margins, especially at high occupancy. Example median residual magnitudes are event $12\,\mathrm{k}$13: Glow $12\,\mathrm{k}$14 GeV, HGPflow $12\,\mathrm{k}$15 GeV, MLPF $12\,\mathrm{k}$16 GeV, and jet $12\,\mathrm{k}$17 width: Glow $12\,\mathrm{k}$18, HGPflow $12\,\mathrm{k}$19, MLPF $12\,\mathrm{k}$20 (Kobylianskii et al., 27 Aug 2025).

6. Suitability, limits, and emerging directions

For large-scale collider environments, the TrackML study concludes that the two-stage Transformer+MaskFormer system achieves state-of-the-art tracking on the full TrackML benchmark, combining $12\,\mathrm{k}$21 efficiency with $12\,\mathrm{k}$22 fake rate at approximately $12\,\mathrm{k}$23 ms/event. It further states that linear scaling in hit multiplicity and GPU-friendly FlashAttention kernels make the method competitive for HL-LHC trigger-level deployment. The same study highlights that the unified instance-segmentation plus regression framework naturally handles hit sharing, joint parameter estimation, and offers a single end-to-end learned model that is potentially extensible to calorimeter clustering, vertexing, and beyond (Stroud et al., 2024).

GLOW frames MaskFormer as a natural interface for instance-level particle clustering integrated with incidence-matrix supervision and per-particle regression, and argues that a single unified transformer architecture can effectively address diverse reconstruction tasks in particle physics (Kobylianskii et al., 27 Aug 2025). This suggests that MaskFormer-based reconstruction is less a single algorithm than a family of detector-reconstruction formulations organized around query-conditioned masking.

The current limitations are explicit. In GLOW, the $12\,\mathrm{k}$24 attention cost on large inputs is identified as a bottleneck, with sparse or local attention proposed as a remedy; fixed $12\,\mathrm{k}$25 queries must be chosen high enough to avoid missed particles; and reliance on fully supervised incidence labels may motivate semi-supervised or contrastive pretraining. In the tracking study, future extensions include adding strip layers, expanding query capacity, quantisation or pruning for faster inference, and joint training of hit filtering and reconstruction (Kobylianskii et al., 27 Aug 2025, Stroud et al., 2024).

An objective reading of these results is that MaskFormer-based reconstruction is most compelling when reconstruction can be expressed as instance segmentation over detector objects plus per-instance prediction. Its principal strengths, as documented in the cited works, are compatibility with global attention, permutation-invariant matching, shared or fractional assignments where needed, and direct coupling of grouping with downstream regression and classification.