All Research

Summer Internship · 2023

Application of the Particle Transformer to Quark Flavour Tagging for Future Higgs Factories

Presentation Published Paper

Period July – August 2023

Affiliation Kyushu University, Japan

Supervisor Dr Taikan Suehara

Responsible Professor Prof. Kiyotomo Kawagoe

Python PyTorch Particle Transformer ILD Full Simulation Feature Engineering Transfer Learning ROC Benchmarking

Overview

Future electron–positron Higgs factories — the ILC, FCCee and CEPC — are proposed to produce Higgs bosons in enormous quantities to make precision measurements of their couplings. A critical step in that programme is jet flavour tagging: deciding whether a collimated spray of particles (a jet) was initiated by a bottom (b), charm (c), strange (s) or light (u/d/g) quark. The cleaner that separation, the higher the discovery potential — a factor-of-two gain in background rejection is worth roughly a 40% increase in physics reach, or a doubling of running time.

During this internship at Kyushu University I applied the Particle Transformer (ParT) — a state-of-the-art Transformer architecture for jet tagging — to ILD full-simulation data for the first time, and benchmarked it directly against FCCee fast simulation and the established LCFIPlus algorithm. The headline result: ParT improves background rejection over LCFIPlus by roughly 4–10× on the same data. I then ran a systematic study of how to push performance further — through input-feature selection, logarithmic feature transforms, momentum representation, training sample size, and fine-tuning between datasets. This work fed into the published ICHEP 2024 proceedings on high-level reconstruction with deep learning for the ILD.

Background: Jet Tagging & the Particle Transformer

When quarks are produced in a high-energy collision they cannot exist freely; they immediately radiate and hadronise into a cascade of O(10–100) outgoing particles travelling in roughly the same direction — a jet. The detector measures the position, momentum and trajectory of each of these particles. Jet tagging is the task of inferring the flavour of the original quark from this particle cloud.

Because b and c hadrons travel a measurable distance (~100–400 µm) before decaying, their decay tracks are slightly displaced from the primary collision point. The classical algorithm for Higgs factories, LCFIPlus, reconstructs secondary vertices from these off-axis tracks and feeds vertex and track features into a flavour classifier. It reaches a b-tag efficiency of 80% with 10% charm / 1% light-quark acceptance.

Fig. 1 — Jet tagging pipeline: high-energy collisions produce unstable particles that decay into collimated sprays of outgoing particles; jets are reconstructed from detector hits, then classified by the flavour of the initiating quark.

The Particle Transformer (Qu, Li & Qian, 2022) replaces hand-engineered vertexing with a deep Transformer that ingests all particle-level information directly. Its key innovation is a modified attention mechanism, P-MHA (Particle Multi-Head Attention): instead of the positional encodings used in language models, it adds a learned bias built from pairwise particle interactions (derived from the particles' energy–momentum 4-vectors) directly to the pre-softmax attention weights. On the 100-million-jet JetClass benchmark, ParT surpassed the previous state of the art (ParticleNet) by a large margin.

Fig. 2 — Architecture comparison. ParT stacks Particle Attention Blocks (using the augmented P-MHA, which injects the pairwise interaction matrix U as a bias) followed by Class Attention Blocks that pool information into a global class token for the final flavour scores.

The interaction features (ln Δ, ln k_T, ln z, ln m²) are computed for every pair of particles from their 4-vectors and added as a physics-motivated bias to the attention — letting the network reason about how particles relate, not just their individual properties.

Data Samples

Two simulation chains were used so that a full-simulation detector (ILD) could be compared against a fast-simulation one (FCCee/Delphes):

Source	Simulation	Process	Energy	Statistics
ILD	Full (Geant4)	e⁺e⁻ → qq̄ (q = b,c,uds)	91 GeV	1M jets
ILD	Full (Geant4)	e⁺e⁻ → ννH → ννqq̄	250 GeV	1M jets
FCCee	Fast (Delphes / IDEA)	e⁺e⁻ → ννH → ννqq̄	240 GeV	up to 10M jets

Data were split 80% training / 5% validation / 15% test. The 17 per-particle input features fall into five groups: impact parameters (6), jet distance (2), track errors (15), particle ID (6) and kinematics (4). A notable representation difference is that FCCee provides only scalar momentum p as the interaction feature, whereas ILD provides the full vector (p_x, p_y, p_z) — a difference I exploited later.

Fig. 3 — Relative θ/φ distributions differ between ILD and FCCee: ILD computes them in the detector frame relative to the jet axis, FCCee from the relative track trace. These representation differences propagate into how features interact inside the model.

Training Setup

Models were trained on an NVIDIA TITAN RTX (24 GB). A study of training length showed that both the ROC-AUC and the validation metric plateau by around 20 epochs; while average loss and accuracy kept creeping up, this no longer translated into tagging-performance gains and risked overtraining. 20 epochs (≈3 hours, vs 30 hours for 200) was therefore adopted as the standard.

Fig. 4 — Comparison of 200-epoch and 20-epoch training. The ROC curves are nearly indistinguishable and the validation metric saturates early, justifying the 20-epoch choice to avoid overtraining.

Results

ParT vs LCFIPlus on ILC data

Applied directly to ILD full simulation, ParT immediately and dramatically outperformed LCFIPlus, reducing background acceptance by a factor of 4.05–9.80× at the same signal efficiency.

Method	c-bkg @ b-tag 80%	uds-bkg @ b-tag 80%	b-bkg @ c-tag 50%	uds-bkg @ c-tag 50%
LCFIPlus	10%	1%	10%	2%
ParT	1.29%	0.25%	1.02%	0.43%

Fig. 5 — Mis-identification probability vs tagging efficiency for ParT on ILD data, for both b-tagging and c-tagging against the relevant backgrounds.

Key result: A modern Transformer applied "out of the box" to ILD full simulation beats the long-established LCFIPlus vertexing algorithm by an order of magnitude in background rejection.

ILC vs FCC and the role of input features

With the full feature set, FCCee fast simulation tagged ~3× better than ILD full simulation. To understand why, I ran a feature-ablation scan, switching each feature group on and off. The results were strikingly different between the two datasets:

ILD: the impact parameter dominates; removing all other groups barely changes performance, but removing impact parameters collapses it (c-bkg jumps from 0.64% to 9.69%).
FCCee: particle ID and jet distance both matter significantly, and — counter-intuitively — removing the 15 track-error variables actually improves performance, suggesting they dilute the contribution of the more discriminating features.

Config (features on)	ILC c-bkg @ b80%	FCC c-bkg @ b80%	ILC b-bkg @ c50%	FCC b-bkg @ c50%
All features	0.64%	0.23%	1.09%	0.35%
No particle ID	0.62%	0.47%	1.14%	0.64%
No particle ID / track err.	0.71%	0.24%	1.24%	0.35%
No PID / jet dist.	0.63%	0.75%	1.19%	0.80%
Jet dist. + track err. only	9.69%	2.64%	6.91%	1.58%

Fig. 6 — Feature-ablation ROC scans for ILD (left) and FCCee (right). The two datasets respond differently to feature removal, but when reduced to *impact parameters only* their b-tagging performance becomes very similar — hinting the impact-parameter generation mechanism is comparable across the two simulations.

Feature engineering: log(abs) transforms

Impact-parameter, jet-distance and track-error values cluster tightly around ~10⁻², which is hard for the network to resolve. Taking log(abs) spreads these distributions out. Applying log(abs) to all three groups helped modestly, but the best result came from applying it to the impact parameters only (c-bkg 0.62% → 0.47%); adding it to jet distance and track errors degraded performance, and fully replacing the raw values lost directional information.

Fig. 7 — log₁₀(abs) distributions of three feature groups. The original values gather around 10⁻²; the log transform spreads them into a range the attention mechanism can exploit more effectively.

Best configuration: applying log(abs) to the impact parameters alone gave the strongest tagging performance — adding more transformed features diluted the weight of this most-discriminating group.

Momentum representation & sample size

Replacing scalar momentum p with the full vector (p_x, p_y, p_z) as the interaction feature improved ILD performance (c-bkg 0.62% → 0.49%), though it reduced the marginal benefit of the log(abs) transform. Separately, training-sample size proved to be the single biggest lever of all on FCCee:

Training size	c-bkg @ b-tag 80%	b-bkg @ c-tag 50%
800k	0.23%	0.35%
4M	0.054%	0.20%
8M	0.0076%	0.10%

Fig. 8 — Tagging performance improves sharply and *non-linearly* with training statistics: ~4× from 800k→4M and a further ~5× from 4M→8M jets. Obtaining larger ILD samples is therefore a clear priority.

Fine-tuning between datasets

I tested whether a model pre-trained on a large dataset could be fine-tuned onto a smaller one. Fine-tuning helped only when the simulation setups were similar: pre-training on FCCee (8M, different θ/φ convention) and fine-tuning onto ILD actually hurt, whereas pre-training ILC 250 GeV → fine-tuning ILC 91 GeV gave the best outcome. In all cases fine-tuning accelerated the early epochs of training even when the final performance was worse.

Fig. 9 — Training curves with and without fine-tuning. The acceleration of the initial epochs is most pronounced for the matched-simulation pairing (ILC 250 → ILC 91 GeV).

Published Outcome & Strange Tagging

This study contributed to the ICHEP 2024 proceedings "High Level Reconstruction with Deep Learning using ILD Full Simulation" (Suehara et al., arXiv:2410.08772), where the ParT results were extended to a 6-category tagging including strange (s) jets — a regime where high-momentum kaon identification is decisive. The published 3-category numbers confirm the ~5–10× rejection gain over LCFIPlus, and show that ILD and FCCee are comparable when compatible variables are used:

Configuration	c-bkg @ b80%	uds-bkg @ b80%	b-bkg @ c50%	uds-bkg @ c50%
ILD full sim 1M (LCFIPlus)	6.3%	0.79%	7.4%	1.2%
ILD full sim 1M (ParT)	0.48%	0.14%	0.86%	0.34%
FCCee Delphes 1M (ParT, full)	0.21%	0.054%	0.36%	0.059%
FCCee Delphes 8M (ParT, full)	0.007%	0.006%	0.076%	0.021%

For strange tagging at 80% s-tag efficiency, ILD full simulation with Comprehensive PID (CPID) reached 25.7% gluon / 42.7% down-quark background acceptance — a non-negligible separation, though still behind FCCee (20.3% / 29.6%). Even with truth-level PID, FCCee outperformed ILD, pointing to a residual difference between full and fast simulation that remains to be understood.

Summary

Order-of-magnitude gain: the Particle Transformer improves jet-flavour background rejection over LCFIPlus by ~4–10× on identical ILD full-simulation data.

Features matter, and differently per dataset: impact parameters dominate for ILD; particle ID and jet distance dominate for FCCee. A targeted log(abs) transform of the impact parameters, and the use of vector momentum (p_x,p_y,p_z), each give measurable further gains.

Data is king: tagging performance scales strongly and non-linearly with training statistics; fine-tuning only transfers between datasets with matched simulation setups.

Together these findings established the Particle Transformer as a viable, high-performance flavour tagger for ILD full simulation and mapped out the levers — features, transforms, statistics and transfer learning — for pushing it further. The work fed directly into the ILD Concept Group's deep-learning reconstruction programme for future Higgs factories.

All Research