NuGraph2 with Context-Aware Inputs: Physics-Inspired Improvements in Semantic Segmentation

This paper demonstrates that enriching the NuGraph2 architecture with physics-inspired, context-aware input features significantly improves semantic segmentation for underrepresented particles like Michel electrons in Liquid Argon Time Projection Chambers, outperforming auxiliary decoders and energy-based regularization due to the model's current hit-level limitations.

The Gist

Imagine you are trying to solve a massive, three-dimensional jigsaw puzzle inside a giant tank of liquid argon. This isn't a normal puzzle; the pieces are tiny flashes of light and electricity created when invisible particles (like neutrinos) crash into the argon atoms.

The scientists in this paper are using a special type of AI called a Graph Neural Network (GNN). Think of this AI as a very smart detective who looks at all the puzzle pieces (called "hits") and tries to figure out what kind of particle made each one.

The detective is pretty good at spotting big, obvious particles like "MIPs" (which are like straight, clean lines drawn by a ruler). But the detective struggles with Michel electrons. These are tiny, shy particles that appear when a muon stops and decays. They are rare, often get mixed up with other particles, and are very hard to find in the noise.

The researchers asked: "How can we teach our detective to spot these shy Michel electrons better?" They tried three different "training methods," and here is what happened, explained with simple analogies:

1. The "Context Clues" Strategy (Feature Extension)

The Idea: Imagine you are looking at a single dot on a map. If you only see the dot, you don't know if it's a house, a tree, or a car. But if you also know how many roads connect to it and how far the next dot is, you can guess much better.

What they did: They gave the AI extra "context clues" for every single puzzle piece. Instead of just showing the AI the raw data, they added features like:

• How connected is this piece? (Is it a lonely dot or a busy hub?)
• Is it part of a straight line? (Are the neighbors lined up perfectly, suggesting a track?)

The Result: This was the biggest success. By giving the AI these extra clues, it could finally tell the difference between a Michel electron and a regular particle. It was like giving the detective a magnifying glass that showed the shape of the neighborhood, not just the single house. The AI learned that Michel electrons have a specific "neighborhood vibe" that regular particles don't.

2. The "Group Leader" Strategy (Auxiliary Decoders)

The Idea: Imagine you are trying to guess who is in a room. You could try to guess every person individually, or you could ask a "group leader" to first tell you, "Okay, there is definitely a teacher in this room, so now look for students."

What they did: They added a second, smaller AI brain to the main one. This second brain's only job was to count: "How many Michel electrons are in this whole event?" The hope was that if the AI knew a Michel electron had to be there (because it's the child of a muon), it would do a better job finding the specific pieces.

The Result: This didn't work well. It actually confused the main detective. It's like having a manager shouting instructions while the detective is trying to solve the puzzle. The two brains were fighting over the same resources, and the main detective got slightly worse at its job. The "group leader" was trying to do math on the whole picture, but the detective was only looking at individual pieces, so they couldn't agree.

3. The "Energy Budget" Strategy (Physics Regularization)

The Idea: Imagine you are a chef. You know a specific dish (a Michel electron) should weigh exactly 500 grams. If your scale says a dish weighs 2,000 grams, you know something is wrong. So, you decide to punish the chef if they claim a 2,000-gram dish is that specific recipe.

What they did: They added a rule to the AI's training: "If you label a piece as a Michel electron, but the total energy of that piece is too high, you get a penalty." They tried to force the AI to respect the laws of physics regarding how much energy these particles should have.

The Result: This failed and made things worse. The AI became too scared to make a guess. It was like the chef, afraid of the penalty, decided to just say "I don't know" for every dish that looked even slightly heavy. Because the relationship between the "weight" (energy) and the "scale reading" (the data) was a bit fuzzy, the rule was too strict. The AI stopped trying to find the Michel electrons entirely to avoid getting "fired."

The Big Takeaway

The paper teaches us a valuable lesson about teaching AI in science:

It is better to give the AI better eyes (better input data) than to give it a strict rulebook (loss functions) or a nagging manager (auxiliary tasks).

By simply showing the AI the "shape" and "context" of the puzzle pieces, it learned to solve the problem on its own. The other methods tried to force the AI to behave in a certain way, but because the AI was only looking at tiny pieces of the puzzle (not the whole particle), those rules didn't make sense to it.

The scientists conclude that for the next generation of this AI (called NuGraph3), which will be able to look at whole particles instead of just pieces, these "manager" and "rulebook" strategies might work much better. But for now, giving the AI better context clues is the winning strategy.

Technical Summary

1. Problem Statement

The paper addresses the challenge of semantic segmentation in Liquid Argon Time Projection Chambers (LArTPCs), specifically within the NuGraph2 architecture. While Graph Neural Networks (GNNs) have shown promise in reconstructing particle events, they struggle with underrepresented classes, most notably Michel electrons (decay products of muons at rest).

• The Specific Challenge: Michel electrons are critical for energy calibration and muon charge discrimination but are scarce in datasets. In standard NuGraph2, they suffer from low precision and recall because their latent space representations overlap significantly with other classes (particularly Minimum Ionizing Particles, or MIPs).
• The Goal: To improve the classification of Michel electrons and overall semantic segmentation by injecting physics-informed inductive biases into the NuGraph2 model, testing three distinct strategies: input feature augmentation, auxiliary decoders, and loss regularization.

2. Methodology

The authors investigated three complementary approaches to inject domain knowledge into the NuGraph2 pipeline, which processes detector hits as nodes in a graph connected via Delaunay triangulation.

A. Feature Extension (Context-Aware Inputs)

The authors augmented the original four input features (hit time, wire, width, integral) with four new physics-inspired features derived from detector geometry and track continuity:

1. Node Degree: The number of edges connected to a node, distinguishing "hub" nodes from isolated ones.
2. Shortest Edge Length: Euclidean distance to the nearest neighbor in the time-vs-wire plane to identify isolated hits vs. continuous tracks.
3. $\Delta_{wire}$: A double difference of wire coordinates between a node and its two nearest neighbors.
4. $\Delta_{time}$: A double difference of time coordinates between a node and its two nearest neighbors.

• Rationale: These features encode structural context (e.g., linear tracks vs. dispersed noise) that standard message-passing GNNs struggle to learn from raw data alone. The double-difference formulation ($\Delta x = 2x_i - x_{j1} - x_{j2}$) was specifically chosen to highlight local directionality while minimizing global domain shifts between different neutrino beam directions.

B. Auxiliary Decoders (Multi-Task Learning)

The standard NuGraph2 has a semantic decoder (5 classes) and a filter decoder (signal vs. background). The authors added a third decoder to predict class-level correlations:

• Motivation: Michel electrons are conditional on the presence of MIPs (muons).
• Implementation: They tested three variations: a binary node-level Michel classifier, a graph-level Michel counter, and a graph-level class distribution predictor (outputting the percentage of each class per event).
• Mechanism: A global attention mechanism was used to aggregate node embeddings into a fixed-size vector for the Multi-Layer Perceptron (MLP) decoder.

C. Energy-Based Regularization

A physics-based penalty term was added to the loss function to constrain Michel electron predictions based on energy conservation laws.

• Mechanism: The model was penalized if the sum of "deposited energies" (estimated via the waveform Integral feature) for predicted Michel hits exceeded a threshold ($E_{cut}$) or deviated from the expected Langauss distribution (or empirical simulation distribution).
• Constraint: The regularization weight was kept low to avoid suppressing the semantic loss, as the relationship between the Integral feature and true energy has high variance.

3. Key Contributions & Results

The study was evaluated on MicroBooNE public datasets. The results, summarized in Tables 1 and 2 of the paper, reveal a clear hierarchy of effectiveness:

Strategy	Impact on Michel Electrons	Impact on Other Classes	Overall Conclusion
Baseline	Precision: 0.60 Recall: 0.75	High performance	Standard NuGraph2 performance.
Feature Extension	Precision: 0.64 (+4 pts) Recall: 0.78 (+3 pts)	Improved across all classes	Most Effective. Successfully disentangled overlapping latent spaces.
Auxiliary Decoders	Precision: 0.48 Recall: 0.68	Degraded (e.g., MIP Precision dropped to 0.98)	Ineffective. Introduced competing gradients that hurt the primary task.
Energy Regularization	Precision: 0.58 Recall: 0.74	Neutral/Slight degradation	Ineffective. Made the model overly conservative; unreliable due to high variance in energy estimation.

• Feature Extension Success: The inclusion of context-aware features provided the largest gains. By explicitly encoding structural roles (degree) and track continuity ($\Delta$ features), the model could better distinguish Michel electrons from MIPs, which previously shared similar latent representations.
• Failure of Auxiliary Tasks: The auxiliary decoders failed to improve performance. The authors attribute this to the hit-level nature of NuGraph2, which lacks explicit particle- or event-level representations. The optimization trajectory was disrupted by competing loss terms, and the coarse granularity (event-level counts) did not align well with node-level predictions.
• Failure of Regularization: The energy regularization caused the model to become "overly cautious," reducing recall. This was due to the high variance in the correlation between the waveform Integral and true deposited energy, and the lack of a one-to-one mapping between graph nodes and physical particles in the current architecture.

4. Significance and Future Directions

• Input vs. Loss: The primary finding is that embedding physics context directly into node-level inputs is significantly more effective than imposing task-specific auxiliary losses or regularization terms for this specific architecture.
• Architectural Implications: The limitations of the auxiliary and regularization approaches highlight the need for hierarchical architectures. The authors suggest that NuGraph3 (currently in development), which will include explicit particle- and event-level reasoning, will provide a more natural setting for advanced decoders and physics-based regularization.
• Domain Knowledge Integration: The work demonstrates that hand-crafted, physics-inspired features can successfully guide GNNs to learn representations that respect known physical structures (like track continuity), effectively resolving class ambiguities in sparse data without requiring massive increases in model complexity.

In conclusion, while multi-task learning and physics-informed losses are theoretically sound, their success depends heavily on the architectural granularity. For NuGraph2, enriching the input representation with structural context is the optimal strategy for improving the segmentation of rare particle classes like Michel electrons.