update docs

docs/TECHNICAL.md

@@ -29,7 +29,7 @@ Four task-specific CLS tokens independently attend to the shared backbone output
   - Tau ID: `CrossEntropyLoss` (2-class) with label smoothing
   - Charge: `BCEWithLogitsLoss`
   - Decay mode: `CrossEntropyLoss`
-  - Kinematics: `HuberLoss(δ=1.0)` over 5 regression targets
+  - Kinematics: Combined loss using `HuberLoss(δ=1.0)` for $p_T$, $\eta$, and $m$, plus a chord loss (L2 distance) for the $(\sin\phi, \cos\phi)$ vector. The mass component is weighted by $\lambda_m = 0.2$.
 - **Conditional gating**: Auxiliary losses (charge, decay mode, kinematics) are multiplied by the truth tau label, so only signal jets contribute to those tasks.
 
 ## Input Features
@@ -52,9 +52,9 @@ Four task-specific CLS tokens independently attend to the shared backbone output
 | 12 | `isChargedHadron` | \|PDG\| = 211 (π±) |
 | 13 | `isNeutralHadron` | \|PDG\| = 130 (K⁰L) |
 | 14 | `cand_dz` | Longitudinal impact parameter $d_z$ |
-| 15 | `cand_dz_err` | Error on $d_z$ |
+| 15 | `cand_dz_error` | Error on $d_z$ |
 | 16 | `cand_dxy` | Transverse impact parameter $d_{xy}$ |
-| 17 | `cand_dxy_err` | Error on $d_{xy}$ |
+| 17 | `cand_dxy_error` | Error on $d_{xy}$ |
 
 A maximum of 20 candidates per jet are used (padded/clipped).
 
@@ -67,9 +67,9 @@ A maximum of 20 candidates per jet are used (padded/clipped).
 | 2 | 2, 3, 4 | 1-prong, ≥2 π⁰ |
 | 3 | 5, 10 | 3-prong, 0 π⁰ |
 | 4 | 6–9, 11–14 | 3-prong, ≥1 π⁰ |
-| 5 | 15, 16, -1 | Rare / Other |
+| 5 | 15, -1 | Rare / Other |
 
-Leptonic decay modes (15) and background (-1) are not considered for this classification. Background is tagged in a separate head.
+Leptonic decay modes (15) and background (-1) are mapped to class 5. HPS DM 16 is currently unmapped. Background is tagged in a separate head.
 
 ## Data
 

model_card.md

@@ -49,7 +49,7 @@ The models utilize the **Particle Transformer (ParT)** architecture, which uses
 - **Embedding Dimensions:** `[256, 512, 256]`
 - **Pair Embedding Dimensions:** `[64, 64, 64]`
 - **Attention Heads:** 8
-- **Transformer Layers:** 8 (default) / 2 (specific runs)
+- **Transformer Layers:** 2 (default)
 - **CLS Layers:** 2
 - **Activation:** GELU