Particle Trajectory Representation Learning with Masked Point Modeling

This paper introduces PoLAr-MAE, a self-supervised masked point modeling framework for Liquid Argon Time Projection Chamber (LArTPC) data that learns physically meaningful trajectory representations from unlabeled events, achieving state-of-the-art segmentation performance with minimal labeled data and demonstrating the potential for a universal foundation model in particle physics.

The Gist

The Big Picture: Teaching a Computer to "See" the Invisible

Imagine you are trying to teach a computer to understand the paths of tiny, invisible particles (like electrons or protons) flying through a giant tank of liquid argon. This tank is called a Liquid Argon Time Projection Chamber (LArTPC).

When these particles fly through the tank, they leave behind a trail of ionized gas, kind of like a jet plane leaving a white contrail in the sky. However, instead of a smooth line, the computer sees a massive, 3D cloud of millions of individual dots (points). Some dots are close together, some are far apart, and they form complex shapes like straight lines, fuzzy clouds, or tiny branching sparks.

The Problem:
Usually, to teach a computer to recognize these shapes, scientists have to spend years creating millions of fake simulations and manually labeling every single dot (e.g., "This dot is a muon," "This dot is a spark"). It's like trying to teach a child to recognize dogs by showing them 100,000 photos of dogs and telling them "This is a dog" for every single one. It's expensive, slow, and if the real world looks slightly different from the fake photos, the computer gets confused.

The Solution:
The authors of this paper invented a new way to teach the computer using Self-Supervised Learning. Instead of needing a teacher to label everything, they let the computer learn by playing a game of "Fill in the Blanks."

---

The Core Idea: The "Blindfolded Artist" Game

Think of the computer as an artist who has been blindfolded.

1. The Setup: The computer is shown a 3D cloud of particle dots.
2. The Game: The computer is told to cover its eyes (mask) over 60% of the dots. It can only see the remaining 40%.
3. The Challenge: The computer must guess what the hidden dots look like based only on the visible ones. It has to predict:
   • Where the missing dots are in 3D space.
   • How much energy those missing dots had.

By playing this game millions of times on raw, unlabeled data, the computer starts to understand the rules of the universe without anyone ever telling it what a "muon" or an "electron" is. It learns that "straight lines usually mean one thing, and fuzzy clouds mean another" just by seeing the patterns repeat.

The Secret Sauce: "C-NMS" (The Smart Grouping)

One of the biggest challenges is that these particle clouds are messy. Some parts are dense (lots of dots), and some are empty. Standard computer vision tools often group these dots poorly, either missing parts of the trail or grouping unrelated dots together.

The authors invented a new tool called C-NMS (Centrality-based Non-Maximum Suppression).

• The Analogy: Imagine you are organizing a crowded party. You want to group people into circles so everyone can talk.
  • Old way: You pick a person, grab everyone within 5 feet, then pick the next person. You might end up with two circles overlapping heavily, or leaving some people out in the cold.
  • C-NMS way: You look for the most "central" person in a cluster, make them the leader of a circle, and then ensure no other circles overlap too much with that leader. It creates perfect, non-overlapping groups that cover the whole party efficiently.

This allows the computer to break the messy particle cloud into neat, manageable "chunks" (patches) that it can study effectively.

The Results: Learning Fast and Smart

The paper shows that this method is incredibly powerful:

1. Data Efficiency: The computer learned so well from the "Fill in the Blanks" game that it only needed 100 labeled examples to become an expert at identifying particle types.
   • Comparison: The old method needed 100,000 labeled examples to reach the same level of skill. That's a 1,000x reduction in the work required!
2. Emergent Intelligence: When the researchers looked at the computer's "brain" (specifically its attention maps), they saw something amazing. The computer had figured out how to separate individual particle tracks on its own. It learned to say, "These dots belong to this specific particle, and those dots belong to that one," even though it was never explicitly taught to do so. It's like a child learning to distinguish between two friends in a crowd just by watching them interact, without being told their names.
3. The "Foundation Model": The authors released a massive new dataset (PILArNet-M) with 1 million events to help other scientists build on this work. They are essentially building a "base model" for particle physics, similar to how large language models (like the one you are talking to now) are built on vast amounts of text.

Why Does This Matter?

In the world of particle physics (like the DUNE experiment looking for neutrinos), understanding these tiny tracks is crucial for discovering new laws of physics.

• Before: Scientists spent years making simulations and labeling data.
• Now: They can use this new "AI artist" that learns the physics of the universe by looking at the data itself. It's faster, cheaper, and more adaptable to real-world experiments.

Summary in One Sentence

The authors taught an AI to understand the complex 3D paths of subatomic particles by playing a "guess the missing pieces" game, allowing it to learn the physics of the universe with 1,000 times less labeled data than ever before.

Technical Summary

1. Problem Statement

Context: Liquid Argon Time Projection Chambers (LArTPCs) are critical detectors in modern neutrino physics (e.g., DUNE, MicroBooNE). They provide high-resolution 3D imaging of charged particle trajectories.
Challenge:

• Data Characteristics: LArTPC data is inherently sparse (>99% empty voxels), complex, and consists of point clouds representing ionization energy depositions.
• Current Limitations: State-of-the-art reconstruction relies on supervised learning trained on massive Monte Carlo simulations (e.g., >100,000 events). This "simulate $\to$ train" paradigm introduces potential biases (domain shift between simulation and real data) and requires significant computational resources for simulation generation.
• Goal: Develop a Self-Supervised Learning (SSL) framework that learns physically meaningful representations directly from unlabeled data, reducing reliance on labeled simulations and enabling "few-shot" adaptation to downstream tasks.

2. Methodology: PoLAr-MAE

The authors propose PoLAr-MAE (Point-based Liquid Argon Masked Autoencoder), a domain-adapted Masked Autoencoder (MAE) for 3D point clouds.

A. Data Preprocessing & Tokenization (C-NMS)

Standard point cloud grouping methods (Farthest Point Sampling + k-NN or Ball Query) fail on LArTPC data due to varying point densities, leading to either missed trajectory segments or excessive patch overlap (which leaks information during masking).

• Solution: The authors introduce Centrality-based Non-Maximum Suppression (C-NMS).
  • Mechanism: It treats group centers as spheres. Using a greedy algorithm, it iteratively selects the most "central" spheres and suppresses overlapping candidates based on a tunable overlap factor ($f$).
  • Benefit: This dynamically determines the number of patches and points per patch, minimizing "missed points" (coverage) and "duplicated points" (overlap), creating a clean, non-overlapping mask structure essential for effective MAE pre-training.

B. Architecture

The model follows a ViT-based Encoder-Decoder structure adapted for point clouds:

1. Patch Embedding: Patches are encoded into latent tokens using a mini-PointNet (MLP + max-pooling), which is permutation-invariant.
2. Transformer Encoder: A heavy encoder processes visible tokens to capture global context.
3. Masking: 60% of patch tokens are randomly masked.
4. Decoder: A lightweight decoder reconstructs the masked tokens using learned mask embeddings and positional encodings.
5. Reconstruction Tasks:
   • Geometric Reconstruction: Predicts the 3D coordinates of masked points using Chamfer Distance loss.
   • Auxiliary Energy Prediction: A novel task where the model predicts the energy deposition ($dE/dx$) of individual points. This is crucial for Particle Identification (PID) as different particles (muons vs. electrons) have distinct energy loss profiles. An Equivariant Mini-PointNet is used here to handle per-point regression while respecting permutation invariance.

3. Key Contributions

1. First SSL on Raw LArTPC Data: Successfully applied masked modeling directly to sparse 3D point clouds from LArTPCs, moving beyond reconstructed physics objects.
2. C-NMS Tokenization: Developed a novel volumetric tokenization strategy specifically for sparse particle trajectories, solving the overlap/coverage trade-off inherent in standard point cloud grouping.
3. Energy Prediction Task: Introduced an auxiliary energy reconstruction head, forcing the model to learn calorimetric features essential for distinguishing particle types.
4. PILArNet-M Dataset: Released a large-scale dataset of 1.2 million simulated LArTPC events (5.2 billion labeled energy depositions) to serve as a benchmark for future research.

4. Results

The model was pre-trained on 1M unlabeled events and evaluated on downstream semantic segmentation (classifying voxels as Track, Shower, Michel electron, or Delta ray).

• Data Efficiency (Few-Shot Learning):
  • PoLAr-MAE fine-tuned on just 100 labeled events achieved performance comparable to a fully supervised Sparse UResNet trained on >100,000 events.
  • Specifically, for Track/Shower classification, PoLAr-MAE reached >99% precision with only 100 events, whereas the supervised baseline dropped to ~33% precision on showers with the same data.
• Linear Probing:
  • Linear SVMs trained on frozen PoLAr-MAE features achieved 99.4% F1 for Tracks and 97.7% F1 for Showers, proving the model learned separable semantic representations without explicit labels.
  • Fine-grained classes (Michel, Delta) were harder (F1 ~44-52%) but still demonstrated learning of physical structures.
• Emergent Instance Segmentation:
  • Attention Maps: Visualization of the Transformer's attention heads revealed emergent instance segmentation. Specific attention heads spontaneously focused on individual particle trajectories (e.g., isolating a short Michel electron track from a background shower) without explicit instance-level supervision.
• Comparison to SOTA:
  • PoLAr-MAE (fine-tuned on 10k events) outperformed the supervised SPINE baseline (trained on 100k events) in Track and Shower classification, while showing a gap in fine-grained Delta/Michel classification.

5. Significance and Future Work

• Foundation Model Potential: The work demonstrates that SSL can build a robust "foundation model" for LArTPC data, capable of serving as a common base for various reconstruction tasks (segmentation, clustering, PID).
• Reducing Simulation Bias: By learning directly from the structure of the data (even if simulated), this approach reduces the dependency on perfect simulation calibration, potentially bridging the "sim-to-real" gap.
• Limitations: The current architecture struggles with fine-grained, sub-token features (like short Delta rays) due to the fixed-resolution tokenization.
• Future Directions: The authors suggest exploring hierarchical architectures (to handle multi-scale features) and point-native transformers to improve the modeling of stochastic, fine-grained structures.

In summary, this paper establishes a new paradigm for LArTPC analysis, proving that self-supervised masked modeling can learn high-fidelity physical representations with drastically less labeled data than traditional supervised methods.