---
title: MaskFormer-Based Reconstruction
url: https://www.emergentmind.com/topics/maskformer-based-reconstruction
type: topic
---

# MaskFormer-Based Reconstruction

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 [2411.07149]. 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 [2508.20092]. 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 [2411.07149].

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 \(r_a \in \mathbb{R}^4\) for \((\Delta E,\Delta p_T,\Delta \eta,\Delta \phi)\) [2508.20092].

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 \(\phi\) | 6 identical Transformer encoder layers |
| Decoder queries | \(N\) learned \(256\)-D vectors | \(M\) 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 \(256\)-D embedding, processed by a Transformer encoder with sliding-window attention in \(\phi\), 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 [2411.07149]. Per-hit features include Cartesian coordinates \((x,y,z)\), cylindrical coordinates \((r,\phi,\eta)\), conformal mapping \((u,v)\), and cluster properties such as charge fraction, size, and shape. Positional encodings are standard sine/cosine on \(z\) and \(r\), with cyclic positional encoding on \(\phi\).

In GLOW, each input object is embedded by a small per-object MLP, encoded by \(L_e=6\) Transformer encoder layers, and decoded by \(L_d=4\) Transformer decoder layers using masked cross-attention. At decoder layer \(t\), the dynamic mask is
\[
\mathrm{Mask}^t_{a,i}=\sigma\!\left(Q^{t-1}[a]\cdot W^m \cdot X_{\mathrm{enc}}[i]\right),
\]
and the masked cross-attention uses \(\log \mathrm{Mask}^t\) inside the attention logits [2508.20092].

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 \(2100\) for \(p_T>600\) MeV, and queries are initialized as learned \(256\)-D vectors \(Q^{(0)} \in \mathbb{R}^{N\times 256}\). 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 \(Q''_i\) to a mask token \(m_i \in \mathbb{R}^{256}\), and raw mask logits are computed by \(\mathrm{dot}(m_i,H_j)\). After sigmoid activation and thresholding at \(0.5\), the binary mask \(M_{i,j}\in\{0,1\}\) is obtained [2411.07149].

In GLOW, the mask head closely follows MaskFormer, but the output remains a soft assignment:
\[
p_{a,i}=\sigma(m_{a,i}),
\]
where \(m_{a,i}\) is produced by dotting a query-derived prototype vector with the encoded feature of input object \(i\). The predicted soft assignment is trained to approximate the ground-truth incidence matrix
\[
I^\star_{i,a}=E_{i\to a}/E_i,
\]
defined by fractional energy assignments [2508.20092].

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
\[
\hat E_a=\sum_i p_{a,i}\cdot E_i,
\]
with \(\hat p_{T,a},\hat \eta_a,\hat \phi_a\) refined by adding \(r_a\) [2411.07149; 2508.20092].

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 \(C+1\) object classes and \(C=1\) real track plus null, the classification loss is categorical cross-entropy for track versus null,
\[
L_{CE}=-\sum_{i\in \mathrm{matched}} y_i \log p_i.
\]
The mask loss combines Dice and binary focal losses,
\[
L_{\mathrm{Dice}}(M,\hat M)=1-\frac{2|M\cap \hat M|+\epsilon}{|M|+|\hat M|+\epsilon},
\]
\[
L_{\mathrm{Focal}}=-\alpha(1-p)^\gamma \log p-(1-\alpha)p^\gamma \log(1-p),
\]
\[
L_{\mathrm{Mask}}=2L_{\mathrm{Dice}}+50L_{\mathrm{Focal}}.
\]
The regression loss is
\[
L_{\mathrm{Reg}}=\mathrm{SmoothL1}(\vec p-\hat{\vec p})+\mathrm{SmoothL1}(v_z-\hat v_z),
\]
and the total objective is
\[
L=0.1L_{CE}+L_{\mathrm{Mask}}+0.1L_{\mathrm{Reg}}+\sum_{l=1}^8 \lambda_{\mathrm{aux}} L^{(l)},
\]
with auxiliary losses on intermediate decoder outputs and \(\lambda_{\mathrm{aux}}\le 1\) [2411.07149].

For GLOW, supervision is decomposed into segmentation, incidence, and flow terms. The incidence loss is
\[
L_{\mathrm{incidence}}=\frac{1}{NM}\sum_{i=1}^{N}\sum_{a=1}^{M}\left(p_{a,i}-I^\star_{i,a}\right)^2.
\]
After Hungarian matching, the segmentation term for a matched pair \(a\leftrightarrow b\) consists of cross-entropy over object assignments and a soft Dice loss,
\[
L_{CE}=-\sum_{i=1}^{N} I^\star_{i,b}\log p_{a,i},
\]
\[
L_{\mathrm{Dice}}=1-\frac{2\sum_i I^\star_{i,b}p_{a,i}+\epsilon}{\sum_i I^\star_{i,b}+\sum_i p_{a,i}+\epsilon},
\]
combined as
\[
L_{\mathrm{seg}}=L_{CE}+\lambda L_{\mathrm{Dice}}.
\]
The flow term is
\[
L_{\mathrm{type}}=-\log p_a[t_b], \qquad
L_{\mathrm{reg}}=\left\|r_a+\left(\sum_i p_{a,i}x_i^{\mathrm{kin}}\right)-v_b\right\|^2,
\]
with
\[
L_{\mathrm{flow}}=L_{\mathrm{type}}+\gamma L_{\mathrm{reg}},
\qquad
L_{\mathrm{total}}=L_{\mathrm{seg}}+\alpha L_{\mathrm{incidence}}+\beta L_{\mathrm{flow}}.
\]
GLOW identifies permutation invariance through Hungarian matching and query-based decoding as a reason the method removes the need for hand-tuned clustering heuristics [2508.20092].

## 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, \(|\eta|<2.5\), and \(p_T>p_{T\text{-min}}\) of \(600\), \(750\), or \(1000\) MeV are targets. Hits are randomly shuffled in \(\phi\), then sorted by \(\phi\) for sliding-window attention. The encoder uses 12 layers, or 8 for the 1 GeV model, with model dimension \(d=256\), FFN dimension \(2d\), and window size \(w=1024\) for filtering or \(512\) for reconstruction. Training uses batch size 1 event for 30 epochs on an NVIDIA A100, with AdamW at typical settings for vision Transformers, \(lr \simeq 10^{-4}\) and weight decay \(\simeq 10^{-2}\) [2411.07149].

The reported inference times are approximately \(23\) ms/event for hit filtering on about \(60\,\mathrm{k}\) hits, and \(50\) ms for \(p_T>1\) GeV to \(101\) ms for \(p_T>600\) MeV in the MaskFormer reconstruction stage. The combined time is approximately \(100\) ms/event for the \(750\) MeV configuration, with scaling linear in input hit count due to sliding-window attention. The best \(p_T>750\) MeV model reaches plateau efficiency \(\epsilon_{DM}\approx 97.1\%\) with fake rate \(\approx 0.6\%\), \(fakedm_{(p_T>0.9\,\mathrm{GeV})}\approx 0.2\%\), and perfect-match efficiency \(\simeq 94.5\%\) at high \(p_T\). Reported regression residuals include \(p_T\) resolution at few-percent bias/RMS, \(\eta\) and \(\phi\) residuals of order \(10^{-3}\), and \(v_z\) RMS \(\approx 100\,\mu\mathrm{m}\) [2411.07149].

GLOW is trained on 1 M simulated \(e^+e^- \to\) dijet events at \(\sqrt{s}=380\) GeV for the CLIC detector, using PandoraPFA-reconstructed tracks and topoclusters within \(|\eta|<4\). Track features are \((p_T,\eta,\phi,\mathrm{charge})\), cluster features are \((E,\eta,\phi,\mathrm{shape\ moments})\), and augmentations include random \(\phi\) rotations, small Gaussian noise on \(p_T/E\) with \(\sigma\approx 1\%\), and random masking of \(<5\%\) of inputs. Optimization uses AdamW with initial \(lr=1\times 10^{-4}\), weight decay \(1\times 10^{-2}\), linear warm-up for 5 epochs, cosine decay to 0 over 200 epochs, batch size 1024 events across 2 H100 GPUs, dropout \(0.1\) in each FFN, and label smoothing \(\epsilon=0.1\) for type classification [2508.20092].

On evaluation, GLOW reports event-level metrics based on missing \(p_T\) residual and scalar sum \(H_T\), jet-level metrics based on anti-\(k_T\) \(R=0.7\) jets, and clustering purity and efficiency. In all tested metrics, GLOW outperforms HGPflow by approximately \(15\%\) in jet energy resolution and yields narrower \(p_T^{miss}/H_T\) residuals; it also outperforms MLPF by even larger margins, especially at high occupancy. Example median residual magnitudes are event \(p_T^{miss}\): Glow \(0.5\) GeV, HGPflow \(0.8\) GeV, MLPF \(1.2\) GeV, and jet \(\Delta E/E\) width: Glow \(0.04\), HGPflow \(0.047\), MLPF \(0.06\) [2508.20092].

## 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 \(>97\%\) efficiency with \(<1\%\) fake rate at approximately \(100\) 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 [2411.07149].

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 [2508.20092]. 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 \(O(N^2)\) attention cost on large inputs is identified as a bottleneck, with sparse or local attention proposed as a remedy; fixed \(M\) 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 [2508.20092; 2411.07149].

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.

Source: https://www.emergentmind.com/topics/maskformer-based-reconstruction