| --- |
| tags: |
| - particle-physics |
| - tau-reconstruction |
| - fcc-ee |
| - key4hep |
| - particle-transformer |
| --- |
| |
| # ML-Tau Model Card |
|
|
| This repository contains models for tau reconstruction and identification at future colliders (FCC), based on the Particle Transformer (ParT) architecture. |
|
|
| ## Dataset |
| - **Name:** `0528_Large_stats` |
| - **Source:** Preprocessed jet-based FCC dataset for hadronic tau reconstruction. |
| - **Physics Processes:** |
| - **Signal:** $Z \to \tau^+\tau^-$ events. |
| - **Background:** $Z \to q\bar{q}$ (light quarks and gluons). |
| - **Generation & Simulation:** |
| - **Generator:** Pythia8 |
| - **Detector Model:** CLD (`CLD_o2_v07`) for FCC-ee. |
| - **Simulation:** Geant4 (via `ddsim`). |
| - **Software Stack:** [Key4hep Project](https://github.com/key4hep) (release 2025-05-29). [Key4hep-sim (v1.2.5)](https://github.com/HEP-KBFI/key4hep-sim/tree/v1.2.5) |
| - **Reconstruction:** Standard CLD reconstruction (`CLDReconstruction.py`). |
| - **Split:** 90% (train+val), 10% (test). |
| - **Input Features:** 17 candidate-level features (kinematics, identification, etc.). |
| - **Jet Composition:** Maximum of 20 candidates per jet. |
|
|
| ## Dataset Statistics |
| - **Total Jets:** 45,219,239 |
| - **Signal (Tau) Jets:** 5,935,398 |
| - **Background (Quark/Gluon) Jets:** 39,283,841 |
| - **Training Set:** 35,355,456 background + 5,341,858 signal jets |
| - **Test Set:** 3,928,385 background + 593,540 signal jets |
|
|
| ## Model Architecture |
| The models utilize the **Particle Transformer (ParT)** architecture, which uses a combination of particle-level and pair-level features to learn jet representations. |
|
|
| ### Variants |
| - **MultiParTau:** A multi-task learning model that simultaneously performs four tasks: |
| - **Tau Identification (`is_tau`):** Binary classification (Signal tau vs. Quark/Gluon jet). |
| - **Charge Classification:** Identification of the tau charge (+1 or -1). |
| - **Decay Mode Classification:** 6-class classification of tau decay modes. |
| - **Kinematics Regression:** Prediction of 5 kinematic corrections: `[log(pt_gen/pt_reco), delta_eta, delta_sin(phi), delta_cos(phi), log(m_gen/m_reco)]`. |
| - **SingleParTau:** Specialized models trained for one of the above tasks individually. |
|
|
| ### Hyperparameters |
| - **Embedding Dimensions:** `[256, 512, 256]` |
| - **Pair Embedding Dimensions:** `[64, 64, 64]` |
| - **Attention Heads:** 8 |
| - **Transformer Layers:** 2 (default) |
| - **CLS Layers:** 2 |
| - **Activation:** GELU |
|
|
| ## Training Scheme |
| - **Optimizer:** AdamW with a weight decay of 1e-2. |
| - **Learning Rate:** 0.001 (base). |
| - **Scheduler:** `OneCycleLR` with cosine annealing. |
| - **Batch Size:** 12288. |
| - **Precision:** 16-mixed (FP16). |
| - **Multi-task Strategy (MultiParTau):** PCGrad (Projected Conflicting Gradients) is employed to handle gradient conflicts between different tasks during training. |
| - **Task Weighting:** |
| - Tau ID: 1.0 |
| - Charge: 1.0 |
| - Decay Mode: 1.0 |
| - Kinematics: 2.0 |
|
|
| ## Trained Models & Git Hashes |
| The models located in the [cld/qq_vs_z_91gev/0612](https://huggingface.co/HEP-KBFI/fcc-tau/tree/main/cld/qq_vs_z_91gev/0612) directories correspond to the following configurations and git hashes: |
|
|
| | Model Name | Task | Git Hash | |
| |------------|------|----------| |
| | `multipartau_full` | Multi-task | `b8483f6` | |
| | `single_charge` | Charge | `b8483f6` | |
| | `single_decaymode` | Decay Mode | `b8483f6` | |
| | `single_kinematics` | Kinematics | `b8483f6` | |
| | `single_tauid` | Tau ID | `b8483f6` | |
