🤖 machine learning

Lantern: Conflict-Aware Gradient Blending for Physics-Guided Diffusion Models in Calorimeter Simulation

The paper introduces Lantern, a physics-guided diffusion model for calorimeter simulation that employs GradBlend to resolve gradient conflicts between statistical denoising and physics-aware auxiliary losses, thereby significantly improving both correlation fidelity and physics accuracy compared to existing methods.

Original authors: Farzana Yasmin Ahmad, Vanamala Venkataswamy, Geoffrey Fox

Published 2026-07-29

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Farzana Yasmin Ahmad, Vanamala Venkataswamy, Geoffrey Fox

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine trying to predict the chaotic splash of a single raindrop hitting a puddle, but the "puddle" is a massive, high-tech detector buried deep underground, and the "raindrop" is a particle of light traveling at nearly the speed of light. This is the world of high-energy physics, specifically the hunt for new particles at the Large Hadron Collider (LHC). To understand what happens when these particles crash into the detector, scientists rely on a digital simulation called a "Monte Carlo" simulation. Think of it as a super-accurate, but incredibly slow, video game engine that calculates every tiny collision. The problem? The upcoming upgrade to the collider will produce so much data that running these simulations would take longer than the age of the universe. Scientists need a shortcut: a "surrogate" model that can guess the outcome almost instantly without losing the physics.

Enter Diffusion Models, a type of AI that has become a superstar in generating images. You might know them from tools that turn a sketch of a cat into a photorealistic masterpiece. These models work by learning to reverse a process of "noise," slowly turning static into a clear picture. In the world of particle physics, researchers have tried using these AI models to simulate particle showers. However, there's a catch: the AI is a master of statistics but a novice at physics. It can learn to make pictures that look like particle showers, but it might get the underlying rules wrong—like making a car that looks real but has no engine. The challenge is to teach the AI the rules of the game without breaking its ability to create new, realistic scenarios.

This is where the paper introduces Lantern, a new method designed to guide these AI models with the laws of physics without dimming their creative spark.

The Problem: The AI vs. The Physics Teacher

Imagine you are teaching a student (the AI) to draw a complex city skyline. The student is great at copying the general vibe and colors (the "denoising" part), but they keep getting the perspective of the buildings wrong. You want to help them by giving them a ruler and a protractor (the "physics" part).

The problem is that if you just tell the student, "Draw the buildings straight," they might get so focused on the ruler that they forget how to draw the buildings at all. In the world of AI training, this is called a "gradient conflict." The AI's natural instinct to learn the data distribution (the skyline) fights against the new rule you're trying to enforce (the ruler). Previous attempts to fix this were like having two teachers shouting different instructions at the same time; the student got confused, and the drawings got worse. Some methods tried to force the physics rules in, but it was like trying to steer a car by yanking the steering wheel while someone else is driving; the car would swerve wildly and crash.

The Solution: Lantern's "Conflict-Aware" Steering

The authors of this paper, working with data from the "CaloChallenge" (a competition to simulate particle detectors), realized they needed a smarter way to blend these two instructions. They introduced a new system called Lantern, which uses a technique called GradBlend.

Think of GradBlend as a very smart co-pilot. The main driver (the AI's natural learning) is always in control of how fast the car goes (the "magnitude" of the step). The co-pilot (the physics rules) can only suggest which way to turn (the "direction"). If the co-pilot suggests a turn that is too sharp or goes against the driver's intuition, the system gently ignores the turn or softens it, ensuring the car never swerves off the road. This allows the AI to learn the physics without forgetting how to generate realistic images.

To make this work, the team invented two new "physics lessons" for the AI:

The Variance-Stabilized Voxel Loss: This teaches the AI that not all parts of the particle shower are equally important. Just like a camera sensor is more sensitive to bright light than dim shadows, this rule tells the AI to pay extra attention to the "dim" parts of the shower that are often ignored but crucial for accuracy.
The Graph Laplacian Loss: This teaches the AI about the neighborhood. In a particle shower, energy doesn't just appear randomly; it flows from one cell to its neighbors. This rule ensures that if one part of the shower has energy, the cells right next to it should have a logical, smooth relationship with it, like ripples in a pond.

The New Ruler: Measuring Correlation

The paper also points out that the old ways of grading these AI models were flawed. They used metrics that checked if the total energy was right, but they missed how the energy was connected across the detector. It's like grading a painting by only checking if the total amount of paint used was correct, ignoring whether the colors blended nicely.

The authors introduced a new metric called the Correlation Frobenius Distance (CFD). Imagine this as a "relationship score." It checks if the AI understands how different parts of the particle shower talk to each other. If the AI gets the total energy right but the connections wrong, the CFD score will catch it.

The Results: A Smoother Ride

When the team tested Lantern on a dataset of 6,480 tiny detector cells (voxels), the results were striking.

The Old Way: When they tried to add physics rules using standard methods (like PCGrad or GradNorm), the AI's performance tanked. The "physics" rules overwhelmed the AI, causing the simulation to become 2 to 100 times worse at reproducing the true particle behavior. It was like the co-pilot grabbed the wheel and crashed the car.
The Lantern Way: By using GradBlend, the AI was able to absorb the physics lessons without losing its grip. In fact, with the "Graph Laplacian" rule, Lantern actually improved the simulation, making it more accurate than the previous best models.

The team also discovered a crucial timing trick. They found that for the "Voxel" rule, it was best to turn it off during the very last part of training. It's like letting the student practice with the ruler for most of the class, but then letting them draw freehand for the final exam to ensure they haven't become too rigid. This "terminal denoising-only phase" was essential to keep the simulation diverse and realistic.

The Bottom Line

This paper doesn't claim to have solved all of physics, nor does it say the AI is perfect. Instead, it shows that by treating the AI's natural learning as the primary driver and using a "conflict-aware" system to gently guide it with physics rules, we can create faster, more accurate simulations. Lantern proves that you don't have to choose between a fast AI and a physics-accurate one; you just need a better way to blend their instructions. This could help scientists at the Large Hadron Collider process the massive amounts of data expected in the coming years, allowing them to discover new secrets of the universe without waiting for the computers to catch up.

Technical Summary: Lantern: Conflict-Aware Gradient Blending for Physics-Guided Diffusion Models in Calorimeter Simulation

1. Problem Statement

Monte Carlo (MC) simulation of calorimeter showers (e.g., using Geant4) is a computational bottleneck for the High-Luminosity Large Hadron Collider (HL-LHC). While diffusion models have emerged as fast, high-fidelity surrogates, their standard denoising objective is purely statistical. Consequently, a model can minimize the loss while failing to reproduce the underlying physical correlations, particularly across calorimeter layers and voxels.

Existing physics-informed generative methods are ill-suited for this domain because they typically assume:

A closed-form governing law or PDE residual (which stochastic showers do not possess).
Hard per-sample constraints (shower physics is defined by ensemble statistics, not per-sample conservation laws beyond total energy).
Standard multi-task learning (MTL) combination rules (e.g., PCGrad, GradNorm) that treat the generative objective and physics auxiliary objectives as peers. This "peer treatment" often degrades the generative fidelity because the auxiliary gradients can override the denoising gradient, leading to mode collapse or unstable training.

Furthermore, standard evaluation metrics like the Fréchet Physics Distance (FPD) and Kernel Physics Distance (KPD) compare samples in a feature space but fail to capture the specific correlation structure between layers and voxels.

2. Methodology

The proposed framework, Lantern, addresses these gaps through a novel metric, two physics-aware auxiliary losses, and a conflict-aware gradient blending rule. The architecture follows a two-stage design (Energy Network and Shape Network) based on CaloDREAM, with all innovations applied to the Shape Network.

2.1 Correlation Frobenius Distance (CFD)

To address the blind spot in standard metrics, the authors introduce the Correlation Frobenius Distance (CFD).

Definition: A single normalized score measuring the fidelity of correlation structures at both layer-wise and voxel-wise scales.
Mechanism: It computes the Pearson correlation coefficients between adjacent layers and transverse positions for both the reference (Geant4) and generated samples, forming a correlation tensor. The CFD is the normalized Frobenius distance between the generated and reference correlation tensors.
Utility: It diagnoses where a surrogate departs from the true correlation structure, which FPD and KPD may miss.

2.2 Physics-Aware Auxiliary Losses

Two soft, per-sample physics losses are encoded to capture structure that no closed-form constraint expresses:

Variance-Stabilized Voxel Residual Loss:
- Based on counting statistics, where fluctuation scales with the square root of energy.
- It applies an inverse-variance weighting to the residual between the model's clean estimate ( $\hat{x}_0$ ) and the ground truth ( $\tilde{x}_0$ ).
- Uses a Huber penalty to prevent outliers from dominating the update.
- Conflict: This loss often conflicts with the denoising objective, particularly in later training stages.
Graph Laplacian Loss:
- Encodes local smoothness constraints based on the detector geometry (face-sharing neighbors).
- Instead of minimizing the signal's roughness (which would erase shower structure), it minimizes the difference between the Laplacian of the predicted signal and the Laplacian of the reference signal.
- Conflict: This loss generally aligns well with the denoising objective.

2.3 GradBlend: Conflict-Aware Gradient Blending

To combine the denoising gradient ( $g_d$ ) and the auxiliary gradient ( $g_a$ ) without degrading the generative objective, the authors propose GradBlend. Unlike peer-based MTL methods, GradBlend anchors the update magnitude to the denoising gradient while allowing the auxiliary to steer the direction only when beneficial.

Mechanism:
1. Direction: Computes unit vectors for both gradients. A tentative blend direction is formed. If the conflict angle ( $\theta_{conf}$ ) is high (gradients oppose), the auxiliary's contribution to the direction is gated (linearly decayed to zero between $120^\circ$ and $150^\circ$ ).
2. Magnitude: The step size is anchored strictly to the norm of the denoising gradient ( $\|g_d\|$ ), scaled by a factor dependent on the conflict angle ( $\max(\cos(\theta_{conf}/2), 0.05)$ ).
3. Result: The auxiliary can steer the update direction but cannot inflate the step size or override the denoising objective. If gradients conflict strongly, the step shrinks, preserving stability.

2.4 Training Schedule

Because GradBlend normalizes gradients, the auxiliary weight acts as a simple on/off gate. The authors employ a schedule (S3) where:

The auxiliary is active after a warmup.
Crucially, the auxiliary is gated off (turned off) in the final epochs of training.
This "terminal denoising-only phase" is essential for the voxel loss (which conflicts late in training) to ensure full sample coverage and prevent recall collapse. The Laplacian loss is insensitive to this schedule.

3. Key Contributions

CFD Metric: The first single-score metric for shower correlation fidelity at layer and voxel scales.
Physics-Aware Losses: Two auxiliary losses (variance-stabilized voxel residual and graph Laplacian) that encode soft per-sample structure without requiring closed-form constraints.
GradBlend: A gradient blending rule that keeps denoising primary by anchoring step magnitude to the denoising gradient and gating the auxiliary's directional influence based on conflict.
Empirical Validation: Demonstration that standard MTL rules (PCGrad, GradNorm, IMTL-G, ConFIG) degrade FPD by 2–100 $\times$ when applied to this problem, whereas GradBlend admits the physics signal without regression.

4. Results

Experiments were conducted on Dataset 2 of the CaloChallenge (6,480 voxels, 45 layers).

Fidelity vs. Baselines:
- Standard MTL methods (Base+PCGrad, Base+GradNorm, etc.) significantly degraded the Fréchet Physics Distance (FPD) compared to the denoising-only baseline (Base). For example, Base+IMTL-G with the Laplacian loss resulted in an FPD of 3454 (vs. 30.87 for Base).
- Lantern (Base+GradBlend) maintained FPD comparable to the denoising-only baseline (25.97 with Laplacian, 32.84 with Voxel) while improving or matching the CFD.
- Lantern achieved the best CFD among trained models (0.120 with Laplacian), outperforming the denoising-only baseline (0.122).
Stability: Training curves showed that standard MTL methods suffered from massive fluctuations in auxiliary loss (orders of magnitude), whereas Lantern kept the auxiliary loss stable and close to the baseline curve.
Ablation: Removing the terminal denoising-only phase (S1/S2 schedules) caused the voxel loss to inflate FPD and collapse recall (mode collapse). The Laplacian loss was robust to the schedule, but the voxel loss required the terminal phase to preserve sample coverage.
Classifier AUC: Lantern and the denoising baseline achieved AUCs near 0.5 (indicating indistinguishability from Geant4), while failing MTL baselines showed high AUCs (e.g., 0.874 for Base+IMTL-G Laplacian), indicating easy separation from the true distribution.

5. Significance and Claims

The paper claims that Lantern successfully bridges the gap between statistical generative modeling and physics-guided training for stochastic cascades where no per-sample governing equation exists.

Primary Claim: Physics guidance is safe for strong generative surrogates only if the combination rule protects the primary generative objective. Treating physics and generation as peers leads to catastrophic degradation of sample fidelity.
Mechanism: By decoupling direction from magnitude and anchoring the step size to the denoising gradient, GradBlend allows the model to incorporate soft physics constraints (correlation structure) without sacrificing the statistical validity of the generated showers.
Metric Contribution: The introduction of CFD provides a necessary diagnostic tool for evaluating correlation structures that standard distributional metrics miss.
Limitations: The authors note that the method is currently demonstrated on a single controlled setting (Dataset 2). The thresholds for GradBlend (angular gates, magnitude floor) and the schedule closing epochs are empirically set and require further analysis. The work does not claim to replace full simulation but to provide a high-fidelity surrogate for specific high-energy physics applications.