ParticleTransformer: Tau Lepton Identification

• Tau lepton identification algorithms are advanced ML solutions designed to distinguish hadronically decaying tau leptons from QCD jets using integrated binary classification, kinematic regression, and decay-mode classification.
• The ParticleTransformer network employs a multi-task transformer architecture with dedicated MLP heads, achieving sub-percent bias in momentum resolution and decay-mode accuracy up to 95%.
• The approach demonstrates strong domain-shift resilience and optimized loss strategies, paving the way for improved tau reconstruction in collider experiments under varying conditions.

Tau lepton identification algorithms are designed to discriminate hadronically decaying tau leptons (τₕ) from the overwhelming background of quark- and gluon-initiated jets, as well as to reconstruct their visible momentum and decay modes. Modern approaches employ unified machine learning frameworks that simultaneously address identification, kinematic regression, and decay-mode classification within a single architecture. The ParticleTransformer network, applied to simulated $e^+e^-$ collisions with realistic detector response via PandoraPFA, exemplifies the state of the art in unified tau identification, substantially surpassing heuristic, rule-based approaches in resolution, classification accuracy, and domain stability (Tani et al., 9 Jul 2024).

1. Task Decomposition in Tau Reconstruction

The reconstruction of hadronically decaying tau leptons is approached as a composite problem comprising three tightly related sub-tasks:

1. Tau Identification ("isTau") A binary classification targeted at separating genuine τₕ objects from generic QCD jets. This forms the basis for downstream tau-specific analyses.
2. Kinematic Reconstruction Regression of the visible transverse momentum, $p_T^{\text{vis}}$, of the tau candidate. This excludes the momentum carried away by neutrinos and focuses on charged and neutral hadrons produced in the decay.
3. Decay-Mode Classification Multi-class classification into discrete decay modes such as $h^\pm$, $h^\pm\pi^0$, $3h^\pm$, and so on, reflecting the number and nature of charged and neutral prongs. The full scheme partitions into 16 classes, including rare higher-prong combinations, as determined by prong multiplicities.

These tasks are implemented within a single multi-task network backbone, using separate output heads for each objective. This allows for joint or sequential training regimes, with the possibility of weighted multi-task loss optimization.

2. ParticleTransformer Network Architecture

The ParticleTransformer architecture processes up to $N_{\text{max}}=16$ particle-flow (PF) candidates per jet, selected by descending $p_T$. Each PF candidate is encoded by a composite feature vector:

$$x_i = \{ p_x, p_y, p_z, E; q; \text{one-hot(pid)}; \log p_T, \log E; \Delta\eta, \Delta\phi; \log(p_T/p_{T,\text{jet}}), \log(E/E_{\text{jet}}) \} \in \mathbb{R}^{D_{\text{in}}}$$

A learnable linear embedding maps $x_i$ into a $d_{\text{model}}$-dimensional latent space:

$$e_i = W_{\text{emb}} x_i + b_{\text{emb}}, \quad e_i \in \mathbb{R}^{d_{\text{model}}}$$

The transformer encoder then applies $L$ layers (e.g., $L=6$) of multi-head self-attention (with $H=8$ heads per layer):

$$\text{head}_h = \text{softmax}\left( Q_h K_h^T / \sqrt{d_k}\right) V_h , \quad Q_h, K_h, V_h \in \mathbb{R}^{N \times d_k}$$

Each attention layer is followed by a position-wise feed-forward network of width $d_{\text{ff}} = 4 \times d_{\text{model}}$, and all sub-layers are wrapped in residual connections and LayerNorm.

After $L$ layers, global pooling combines the per-particle embeddings by both mean and max to generate a jet-level embedding vector $e_{\text{jet}} \in \mathbb{R}^{2 d_{\text{model}}}$. This is routed through three distinct multi-layer perceptron heads, each task-specific:

• isTau head: $e_{\text{jet}} \to$ [128] $\to$ [64] $\to$ 1 (sigmoid)
• Momentum regression head: $e_{\text{jet}} \to$ [128] $\to$ [64] $\to$ 1 (linear)
• Decay-mode classification head: $e_{\text{jet}} \to$ [128] $\to$ [64] $\to K$ (softmax, $K=16$)

3. Loss Functions and Optimization Strategies

Each task corresponds to a distinct loss function:

• Tau ID (Binary Cross-Entropy):

$$L_{\text{ID}} = -\frac{1}{N} \sum_{j=1}^N \left[ y^{(j)} \log \hat y^{(j)} + (1 - y^{(j)}) \log (1 - \hat y^{(j)}) \right]$$

• Decay-Mode Classification (Categorical Cross-Entropy):

$$L_{\text{mode}} = -\frac{1}{N} \sum_{j=1}^N \sum_{c=1}^K t_c^{(j)} \log p_c^{(j)}$$

$t_c$ is a one-hot ground truth, $p_c$ is softmax output.

• Kinematic Regression (Huber Loss on $\log$-ratio):

$$\Delta_j = \log \frac{p_{T,\text{pred}}^{(j)}}{p_{T,\text{true}}^{(j)}}$$

$$L_{\text{kin}} = \frac{1}{N}\sum_{j=1}^N \begin{cases} \frac{1}{2} \Delta_j^2 &\vert\Delta_j\vert < \delta \ \delta (|\Delta_j| - \frac{1}{2}\delta) & \text{otherwise}, \end{cases}$$

with $\delta=1.0$.

The joint multi-task loss applies weights $\{\alpha, \beta, \gamma\}$ to each term:

$$L_{\text{total}} = \alpha L_{\text{ID}} + \beta L_{\text{kin}} + \gamma L_{\text{mode}}$$

Possible optimization schedules include staged training (e.g., pre-training the identification head, then freezing) or simultaneous joint minimization with tuned loss weights.

4. Dataset, Data Representation, and Training Procedure

The publicly available FuTauTure dataset underpins both the algorithm development and benchmarking:

• Events: $e^+e^-$ collisions at $\sqrt{s}=380$ GeV ($ZH$, $WW$, $H\nu\nu$) with $\sim2$ million events per channel.
• Detector: Full Geant4 simulation using the CLICdet geometry, reconstructed with PandoraPFA.
• PF objects: Each candidate records $(p_x, p_y, p_z, E)$, electric charge $\in\{-1, 0, +1\}$, and PF-type $\{\text{e}, \mu, \gamma, h^\pm, h^0\}$.
• Jet construction: $ee\_$genkt algorithm ($p=-1$, $R=0.4$, $p_T>5$ GeV).
• Ground-truth assignment: jets matched to generator-level $\tau_h$ via $\Delta R < 0.3$, with stored isTau label, decay-mode label (0…15), and true visible $p_T$.

The training set is composed of $WW$ and $H\nu\nu$ events, while $ZH$ forms the test split to quantify domain-shift resilience. Training uses AdamW with $10^{-2}$ weight decay, initial learning rate $10^{-3}$ (cosine annealed over 100 epochs), batch size 1024, dropout 0.1, label smoothing 0.01, and early stopping on validation loss. Each configuration is repeated three times with independent random seeds for statistical robustness.

5. Performance Metrics and Comparative Analysis

The performance of the ParticleTransformer and alternative architectures is quantified via momentum resolution, decay-mode precision, and ROC area under curve (AUC) for τ-ID:

Model	Momentum Resolution (IQR $\log(p_{T,\text{pred}}/p_{T,\text{true}})$)	Decay-Mode Precision	τ-ID AUC
ParticleTransformer	2.1–3%	80–95%	$\approx 0.995$
LorentzNet	2.3–3.5%	78–93%	–
DeepSet	3–4.5%	70–88%	$\approx 0.98$
HPS baseline	3.5–10%	$\sim$60–90%	$\approx 0.96$

ParticleTransformer demonstrates sub-percent bias, $2$–$3\%$ momentum resolution, and $80$–$95\%$ per-class decay-mode accuracy—exceeding the HPS heuristic baseline, especially for modes with high $\pi^0$ multiplicities. At $p_T \in [20,200]$ GeV, the model robustly generalizes to the $ZH$ test domain without re-training, with $\mathcal{O}(1\%)$ degradation.

6. Domain-Shift Resilience and Future Directions

ParticleTransformer maintains high accuracy under mild kinematic domain shifts between training ($WW$, $H\nu\nu$) and held-out testing ($ZH$), indicating strong generalization. Several future directions are outlined:

• Overlay of realistic beam-induced backgrounds (e.g., $\gamma\gamma \rightarrow$ hadrons) to evaluate model robustness in the presence of pileup and underlying event contamination.
• Incorporation of full impact parameter information ($d_z$, $d_{xy}$) for enhanced lifetime discrimination in high-occupancy scenarios.
• Pre-training the transformer backbone on generic jet tagging or substructure tasks to enable fine-tuning on highly specialized $\tau$ reconstruction, facilitating transfer learning and efficiency in data-constrained environments.
• Benchmarking lightweight transformer or quantized network variants for ultra-fast FPGA-based deployment and exploration of physics-informed/self-supervised architectures.

The FuTauTure dataset and associated pipelines are intended as a community standard for further advancement of ML-based tau identification methodologies.

7. Implications, Significance, and Limitations

The unified machine learning formulation adopted here obviates the need for expert-crafted sequence-of-cuts or decay-mode hypotheses, treating $\tau_h$ as a special case of a highly collimated, low-multiplicity jet. This enables direct optimization on identification, regression, and classification objectives—a paradigm that integrates seamlessly with the broader context of jet tagging. The approach attains significant gains over previous—particularly heuristic or BDT-based—algorithms in momentum resolution and decay-mode fidelity, with minimal sensitivity to modest domain shifts. However, the presented studies are based on clean $e^+e^-$ environments with full simulation. Realistic $pp$ environments with pileup, beam backgrounds, and detector non-idealities may necessitate further domain-specific adaptation and validation.

In summary, the ParticleTransformer-based tau lepton identification algorithm represents a mature, end-to-end solution that sets a new performance benchmark for hadronic tau reconstruction across identification, kinematic, and decay-mode axes in collider experiments (Tani et al., 9 Jul 2024).