Paper deep dive
Lantern: Conflict-Aware Gradient Blending for Physics-Guided Diffusion Models in Calorimeter Simulation
Farzana Yasmin Ahmad, Vanamala Venkataswamy, Geoffrey Fox
Intelligence
Status: succeeded | Model: Gemma-4-26B-A4B | Prompt: intel-v1 | Confidence: 93%
Last extracted: 8/1/2026, 11:04:14 AM
Summary
The paper introduces Lantern, a physics-guided diffusion model for calorimeter simulation that addresses the conflict between statistical denoising objectives and physics-aware auxiliary losses. It proposes the Correlation Frobenius Distance (CFD) metric to evaluate correlation fidelity and two auxiliary losses: a variance-stabilized voxel residual loss and a graph Laplacian loss. The core contribution is GradBlend, a gradient combination rule that anchors step magnitude to the denoising gradient while allowing auxiliary gradients to steer direction, preventing performance degradation seen in standard multi-task methods like PCGrad or GradNorm.
Entities (10)
Relation Signals (8)
LanteRn → uses → GradBlend
confidence 98% · We combine both with denoising through GradBlend... yielding Lantern
LanteRn → evaluateswith → Correlation Frobenius Distance
confidence 96% · We introduce the Correlation Frobenius Distance (CFD)... Lantern improves both FPD and CFD.
LanteRn → trainedon → CaloChallenge Dataset 2
confidence 95% · On CaloChallenge Dataset 2... Lantern improves both FPD and CFD.
LanteRn → uses → Graph Laplacian Loss
confidence 94% · with the Laplacian loss, Lantern improves both FPD and CFD.
LanteRn → uses → Voxel Residual Loss
confidence 94% · We then encode... a variance-stabilized voxel residual loss... yielding Lantern
GradBlend → outperforms → GradNorm
confidence 90% · injecting the physics losses through... GradNorm... inflates FPD... whereas GradBlend admits the same signal
GradBlend → outperforms → PCGrad
confidence 90% · injecting the physics losses through... PCGrad... inflates FPD by 2-100x... whereas GradBlend admits the same signal without regression
Geant4 → isbenchmarkfor → LanteRn
confidence 88% · CFD diagnoses where a surrogate departs from Geant4 after training
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract:Monte Carlo simulation of calorimeter showers is a principal bottleneck for the High-Luminosity LHC, and diffusion models have emerged as fast, high-fidelity surrogates. Their denoising objective is purely statistical, however: a model can minimize it while placing the physics wrong. Existing physics-informed generative methods cannot close this gap, because they assume a closed-form law, a governing PDE residual or a hard per-sample constraint, that a shower does not supply: no per-sample PDE governs a stochastic cascade, and energy conservation fixes only one scalar per shower. Standard metrics ignore the correlation structure across calorimeter layers and voxels, comparing showers only in a physics feature space. We address both gaps. We introduce the Correlation Frobenius Distance (CFD), a single normalized score for correlation fidelity at layer-wise and voxel-wise scales. We then encode the soft per-sample structure available in a shower as two physics-aware auxiliary losses: a variance-stabilized voxel residual loss grounded in counting statistics, and a graph Laplacian loss over the detector geometry. We combine both with denoising through GradBlend, which anchors the step magnitude to the denoising gradient while letting the auxiliary steer its direction, yielding Lantern, a physics-guided diffusion surrogate. On CaloChallenge Dataset 2, injecting the physics losses through task-symmetric rules such as PCGrad, GradNorm, IMTL-G, and ConFIG inflates FPD by 2-100x relative to denoising alone, whereas GradBlend admits the same signal without regression and, with the Laplacian loss, Lantern improves both FPD and CFD. Our ablation on the auxiliary loss scheduler shows that the voxel residual loss, whose gradient conflicts with denoising, requires a terminal denoising-only phase to preserve shower fidelity, whereas the non-conflicting Laplacian loss is insensitive to the schedule.
Tags
Links
- Source: https://arxiv.org/abs/2607.25060v1
- Canonical: https://arxiv.org/abs/2607.25060v1
Trouble viewing inline? Open PDF directly →
Full Text
99,326 characters extracted from source content.
Expand or collapse full text
Lantern: Conflict-Aware Gradient Blending for Physics-Guided Diffusion Models in Calorimeter Simulation Farzana Yasmin Ahmad fa7sa@virginia.edu Computer Science Department, University of VirginiaUSA , Vanamala Venkataswamy v3xu@virginia.edu Visiting Scholar, University of VirginiaUSA and Geoffrey Fox gcfexchange@gmail.com Computer Science Department, Biocomplexity Institute, University of VirginiaUSA Abstract. Monte Carlo simulation of calorimeter showers is a principal computational bottleneck for the High-Luminosity LHC, and diffusion models have emerged as fast, high-fidelity surrogates. Their denoising objective is purely statistical, however: a model can minimize it while placing the physics wrong. Existing physics-informed generative methods cannot close this gap, because they assume a closed-form law, a governing PDE residual or a hard per-sample constraint, that a shower does not supply: no per-sample PDE governs a stochastic cascade, and energy conservation fixes only one scalar per shower. Standard metrics share a blind spot of their own, comparing showers in a physics feature space while ignoring the correlation structure across calorimeter layers and voxels. We address both gaps. We introduce the Correlation Frobenius Distance (CFD), a single normalized score for correlation fidelity at layer-wise and voxel-wise scales. We then encode the soft per-sample structure available in a shower as two physics-aware auxiliary losses. The first is a variance-stabilized voxel residual loss grounded in counting statistics; the second is a graph Laplacian loss over the detector geometry. We combine both with denoising through GradBlend, which anchors the step magnitude to the denoising gradient while letting the auxiliary steer its direction, yielding Lantern, a physics-guided diffusion surrogate. On Dataset 2 of the CaloChallenge, injecting the physics losses through task-symmetric rules such as PCGrad, GradNorm, IMTL-G, and ConFIG inflates FPD by 22–100×100× relative to denoising alone, whereas GradBlend admits the same signal without regression and, with the Laplacian loss, Lantern improves both FPD and CFD. Our ablation study on the auxiliary loss scheduler shows that the voxel residual loss, whose gradient conflicts with denoising, requires a terminal denoising-only phase to preserve shower fidelity, whereas the non-conflicting Laplacian loss is insensitive to the schedule. diffusion models, calorimeter simulation, physics-informed learning, multi-task optimization, high-energy physics †ccs: Computing methodologies Neural networks†ccs: Computing methodologies Simulation types and techniques†ccs: Applied computing Physics 1. Introduction Calorimeter shower simulation models the particle cascade produced when a high-energy particle strikes the detector material, and underpins detector design, calibration, and event reconstruction in high-energy physics. The standard tools are Monte Carlo simulators such as Geant4 (Agostinelli et al., 2003), accurate but computationally expensive; the upcoming High-Luminosity upgrade of the Large Hadron Collider (Apollinari et al., 2017) makes Geant4 alone impractical at the required data volume. The CaloChallenge benchmark (CaloChallenge Collaboration, 2022) spurred deep generative surrogates for full simulation, many producing high-fidelity showers at a fraction of the cost. Evaluating these models is nontrivial: a faithful surrogate must preserve the underlying physics, not only visual plausibility. The standard distributional metrics, the Fréchet and kernel physics distances (FPD and KPD) (Kansal et al., 2023), compare samples in a physics feature space and detect global mismatch, but do not test whether a model reproduces correlations between calorimeter layers and voxels; the Pearson Correlation Coefficient (PCC) (Ahmad et al., 2024) probes these correlations but returns a full matrix rather than a single comparable score. We introduce the Correlation Frobenius Distance (CFD), a single normalized score for layer- and voxel-wise correlation fidelity. CFD diagnoses where a surrogate departs from Geant4 after training; a complementary question is whether physics structure can be built into training itself. The denoising objective of diffusion is purely statistical, and existing physics-informed variants assume a closed-form law that a stochastic shower does not supply (Section 2). Physics-guided training has consequently entered shower generation only through coarse aggregates: the concurrent CaloTrilogy (Jiang et al., 2026) penalizes layer-wise energies with a scalar-weighted sum and offers no mechanism to keep the physics term from degrading the generative objective. We address both gaps: voxel-level losses and a combination rule that protects denoising. Although no conservation law governs a shower voxel by voxel, each sample still carries physical structure: voxel energies whose fluctuations scale with the deposited energy, and local differences between neighboring cells set by the detector geometry. We encode these as two physics-aware auxiliary losses: a variance-stabilized voxel residual loss and a graph Laplacian loss over face-sharing voxels (Section 4). Pairing either loss with denoising turns training into a multi-task problem, and existing combination rules treat the objectives as peers, none protecting denoising when an auxiliary pulls against it. Since denoising keeps the model generative, our second contribution is a gradient blending rule that admits the physics gradient while keeping denoising primary. We study these questions on Dataset 2 of the CaloChallenge (CaloChallenge Collaboration, 2022), a high-granularity calorimeter of 64806480 voxels. Contributions. This paper delivers a metric, two losses, and a combination rule that together make physics guidance safe for a strong generative surrogate. (1) We introduce the Correlation Frobenius Distance (CFD), to our knowledge, the first single-score metric for shower correlation structure that FPD and KPD miss, at layer-wise and voxel-wise scales. (2) We design two physics-aware auxiliary losses, a variance-stabilized voxel residual loss and a graph Laplacian loss, encoding soft per-sample structure that no closed-form constraint expresses. (3) We introduce GradBlend, a denoising-anchored blending rule: the auxiliary steers only the update direction while the magnitude stays anchored to denoising, keeping the generative objective primary by construction. (4) We show this anchoring is necessary: PCGrad, GradNorm, IMTL-G, and ConFIG degrade FPD by 22–100×100×, while GradBlend admits the physics signal and, with the Graph Laplacian loss, improves both FPD and CFD. (5) Because the auxiliary weight cancels from GradBlend’s update, the schedule reduces to a simple on/off gate; we show its terminal denoising-only phase is essential for sample coverage. 2. Related Work Diffusion-based surrogates learn a data distribution through a single denoising objective, equivalent to learning its score (Ho et al., 2020; Song et al., 2021). The generality that lets diffusion fit any distribution from samples alone leaves the physics unconstrained: a surrogate can minimize the loss while misplacing the relationships that make its outputs physically valid. This has motivated writing physics directly into generative training, through strategies that calorimeter showers fit none of. PDE-residual losses (Bastek et al., 2025; Baldan et al., 2025; Zhang and Zou, 2025), inference-time guidance and projection (Chung et al., 2023; Gao et al., 2023; Liang et al., 2025; Trupin et al., 2026; Utkarsh et al., 2026), equivariant architectures (Hassan et al., 2024; Hoogeboom et al., 2022; Tian et al., 2025), and energy-based steering (Li et al., 2026; Vaitl and Klein, 2025) each presuppose a closed-form per-sample object: a governing equation, differentiable constraint, symmetry group, or tractable energy. A calorimeter shower supplies none. It is a single stochastic Geant4 realization whose physical content lives in the ensemble, and its only hard per-sample constraints, total and per-layer energy, are already imposed by construction through conditioning and preprocessing (Section 3), leaving no per-sample target to enforce. Even where a physics loss applies, adding it to a generative objective is fragile: naive weighted sums produce pathological training dynamics when gradient scales and directions disagree (Wang et al., 2021; Liu et al., 2025). How the physics gradient combines with the generative gradient is therefore not an implementation detail but the central obstacle. Expressing ensemble structure as an auxiliary loss alongside denoising turns training into a two-objective gradient problem, for which multi-task learning supplies combination rules. PCGrad (Yu et al., 2020), GradNorm (Chen et al., 2018), IMTL-G (Liu et al., 2021), and ConFIG (Liu et al., 2025) treat the objectives as peers, granting each equal influence over update direction or training rate. Peer treatment suits co-equal tasks, not a setting where one objective must remain primary: splitting the direction evenly forces denoising to cede half of each update to the auxiliary, suppressing the objective that makes the model generative. A separate line treats one task as primary: gradient cosine similarity (GCS) gates the auxiliary loss by the sign or value of its cosine with the main-task gradient, with guaranteed convergence to critical points of the main task (Du et al., 2018), and this primary-anchored gating has since shown gains in physics-informed PDE solving (Yan et al., 2023). GradBlend belongs to this primary-anchored family but differs in mechanism: GCS and its successors modulate the auxiliary’s scalar weight while summing raw gradients, so a large-norm auxiliary still dominates the update direction. GradBlend instead separates direction from magnitude: it blends unit directions so neither gradient wins by norm, anchors the step size to the denoising gradient alone, and gates the auxiliary’s directional share by the conflict angle. Within calorimeter simulation, surrogates almost universally optimize a data-driven objective, encoding physics through architecture rather than the training loss (Amram and Pedro, 2023; Birk et al., 2024; Buhmann et al., 2024; Favaro et al., 2025; Kobylianskii et al., 2024; Leigh et al., 2024a, b; Mikuni and Nachman, 2022, 2024). The one exception, CaloTrilogy (Jiang et al., 2026), adds a per-layer energy constraint to a MeanFlow generator, weighted by a fixed coefficient or an adaptive Lagrange multiplier from the modified differential method of multipliers (Platt and Barr, 1988). Its constraint has a per-sample target (each layer energy has a true value), so once the representation stabilizes the physics and generative gradients agree and scalar magnitude control suffices; still, its reported sensitivity to constraint scheduling shows the combination is delicate in the weakly-conflicting regime. The auxiliaries we study have no per-sample target: they encode correlation structure that per-sample marginals cannot express, whose optimum opposes denoising and places the problem in the conflicting regime. To our knowledge, no prior calorimeter surrogate has treated physics-auxiliary training as a gradient-combination problem; characterizing that conflict and resolving it with a denoising-anchored rule is our contribution. 3. Problem Formulation We represent a calorimeter shower in physical units (MeV) as x∈ℝ≥0Vx _≥ 0^V where V is the number of voxels in a shower and its preprocessed form as x~=(x)∈ℝV x=T(x) ^V, where T is a composition of transformation steps. The network works in the preprocessed space x~ x, and the inverse transform −1T^-1 maps a sample back to physical energy x when a physical quantity is needed. Any quantity with a time index or a hat (xtx_t, x^0 x_0) is in the preprocessed space; the bare symbol x means physical energy. We treat simulation as conditional generation. Geant4 defines the target distribution p(x∣c)p(x c), where x is the shower in physical space and c is the conditioning information. Following prior work (Mikuni and Nachman, 2024; Favaro et al., 2025), we factorize the generation into two stages, each conditioned on its own subset of c, and train a separate diffusion model for each. The energy network, with parameters ϕφ, models pϕ(u∣Einc)p_φ(u E_inc), where EincE_inc is the incident energy of the shower and u∈ℝLu ^L is derived from the per-layer energies of the shower across its L longitudinal layers. The shape network, with parameters θ, models pθ(x∣Einc,u)p_θ(x E_inc,u), the voxelized shower given the incident energy and u. Marginalizing over u recovers the target: (1) p(x∣Einc)=∫pθ(x∣Einc,u)pϕ(u∣Einc)du.p(x E_inc)= p_θ(x E_inc,u)\,p_φ(u E_inc)\,du. The two networks are trained independently and composed at generation time. Our auxiliary losses and gradient method act on the shape network alone; the energy network keeps the standard denoising objective. During shape-network training, u is the ground-truth value, so the network conditions on exact layer ratios, while at generation u is drawn from the energy network. Denoising objectives. Both networks are denoising diffusion probabilistic models (DDPM) (Ho et al., 2020). The forward process perturbs a clean sample z0z_0 under a variance schedule β1,…,βT _1,…, _T; writing αt=1−βt _t=1- _t and α¯t=∏s=1tαs α_t= _s=1^t _s, the noised state at timestep t∈1,…,Tt∈\1,…,T\ is (2) zt=α¯tz0+1−α¯tϵ,ϵ∼(0,I).z_t= α_t\,z_0+ 1- α_t\,ε, ε (0,I). The energy network applies Eq. (2) to the clean ratio vector u and minimizes the standard denoising loss over ϵϕ(ut,t,Einc) _φ(u_t,t,E_inc); it generates the ratios autoregressively following CaloDREAM (Favaro et al., 2025), carries no auxiliary loss, and is detailed in Appendix B.3. All contributions of this paper act on the shape network, which applies Eq. (2) to the preprocessed shower x~0 x_0 and minimizes, with conditioning (Einc,u)(E_inc,u), (3) ℒdenoise(θ)=x~0,ϵ,t[‖ϵ−ϵθ(xt,t,Einc,u)‖22].L_denoise(θ)=E_ x_0,ε,t [ ε- _θ(x_t,t,E_inc,u) _2^2 ]. Physics-aware auxiliary objective. The denoising loss scores the predicted noise, not the shower. To reason about the shower itself, we use the model’s single-step clean estimate (4) x^0(xt,t)=xt−1−α¯tϵθ(xt,t,Einc,u)α¯t. x_0(x_t,t)= x_t- 1- α_t\, _θ(x_t,t,E_inc,u) α_t. We add a physics-aware auxiliary loss on this estimate, (5) ℒaux(θ)=x~0,ϵ,t[ℓphys(x^0(xt,t),x~0)],L_aux(θ)=E_ x_0,\,ε,\,t\! [\, _phys ( x_0(x_t,t),\, x_0 )\, ], where ℓphys _phys is a per-sample penalty that compares the clean estimate with the true shower. We instantiate ℓphys _phys and describe how it is combined with the denoising objective in Section 4. The denoising and auxiliary objectives share one update on the shape network. Writing their gradients as (6) gd=∇θℒdenoise,ga=∇θ(waux(s)ℒaux)=waux(s)∇θℒaux,g_d= _θL_denoise,\,\,\,g_a= _θ (w_aux(s)\,L_aux )=w_aux(s)\, _θL_aux, where waux(s)≥0w_aux(s)≥ 0 weights the auxiliary loss at optimization step s and the denoising loss is never reweighted, we obtain a multi-task gradient problem in which gdg_d and gag_a need not agree. Denoising is the objective that makes the network generative, so it must stay primary; an auxiliary that overrides it degrades the samples it is meant to improve. The problem is therefore twofold. First, construct an update gblend(gd,ga)g_blend(g_d,g_a) that admits the physics signal while keeping denoising primary. Second, decide when during training the auxiliary should act. We address the first with GradBlend (Section 4.2) and the second with the schedule waux(s)w_aux(s) (Section 4.3). 4. Method Our surrogate follows the factorized design of CaloDREAM (Favaro et al., 2025), with the energy and shape networks of Section 3; unlike CaloDREAM’s conditional flow matching, we use DDPM (Ho et al., 2020) for both. The energy network is an autoregressive transformer over the L layers: the ratio uiu_i for layer i is predicted from EincE_inc and the preceding layers through causal masking, matching the forward-propagating structure of a shower. The shape network, which carries all of our contributions, is a conditional Vision Transformer (Dosovitskiy et al., 2020; Peebles and Xie, 2023) trained to predict the injected noise ϵθ(xt,t,Einc,u) _θ(x_t,t,E_inc,u): it patch-embeds the calorimeter volume, processes the patches with transformer blocks conditioned on t, EincE_inc, and u through affine modulation, and projects back to voxel space. The preprocessing transform T (Eq. (18)) is taken unchanged from CaloDREAM; full architectural details and hyperparameters are in Appendix B. 4.1. Physics-Aware Auxiliary Losses 4.1.1. Variance-Stabilized Voxel Residual Loss Here we take ℓphys _phys as a per-voxel residual loss. The denoising loss Eq. (3) weights every voxel equally, but not all calorimeter voxels are equally informative. In a calorimeter shower a few bright cells carry most of the energy, while many dim cells fix the sparsity pattern and the low-energy tail. Shower signals are counting processes, so the fluctuation scale of a voxel grows as the square root of its physical energy (Fabjan and Gianotti, 2003). This motivates an inverse-variance weighting that standardizes each residual by an energy-dependent scale, moving the capacity from the bright core toward the dim tail. (7) ri=x^0,i−x~0,i|x~0,i|+ϵvox,ℓvox=1V∑i=1VHuberδvox(ri),r_i= x_0,i- x_0,i x_0,i + _vox, _vox= 1V _i=1^VHuber_ _vox(r_i), where i∈1,…,Vi∈\1,…,V\ indexes voxels and x^0,i x_0,i and x~0,i x_0,i are the i-th voxel of the model’s clean estimate and of the Geant4 reference, with floor ϵvox=10−6 _vox=10^-6 and Huber transition δvox=1 _vox=1. The true values x~0,i x_0,i are detached from the computational graph, so the denominator acts as a constant per-voxel weight rather than an optimization target. The Huber penalty (8) Huberδ(r)=12r2|r|≤δ,δ(|r|−12δ)|r|>δ,Huber_δ(r)= cases 12r^2& r ≤δ,\\ δ ( r - 12δ )& r >δ, cases stays quadratic on the bulk of residuals and linear beyond δvox _vox, so voxels whose residuals exceed the assumed scale bound their influence on the update rather than dominating it. 4.1.2. Graph-Laplacian Loss The variance-stabilized voxel residual loss of Section 4.1.1 is pointwise: it ignores how each voxel relates to its neighbors, so a model can match every voxel on average yet still get the local correlations wrong. The second auxiliary loss compares the predicted and reference neighborhood structure across physically adjacent voxels. We represent the calorimeter as a graph G with one node per voxel, connected to face-sharing neighbors along the three grid axes: periodic in azimuth φ , open radially and longitudinally. Interior voxels have six neighbors; boundary voxels have fewer, except along φ , which always supplies both. The combinatorial Laplacian L=D−AL=D-A (adjacency A, degree D), applied to a voxel signal y, returns at each node the summed difference to its neighbors, (Ly)i=∑j∼i(yi−yj)(Ly)_i= _j i(y_i-y_j), a measure of local roughness. In graph signal processing the Laplacian acts as a high-pass filter (Shuman et al., 2013): it responds weakly where neighboring voxels agree and strongly where they differ, picking out edges rather than flat regions. A standard regularizer minimizes the Dirichlet energy x^0⊤Lx^0 x_0 L x_0 to drive the signal toward smoothness, but a calorimeter shower is not smooth, so flattening it would erase the structure. We therefore penalize the gap between the predicted and reference Laplacian responses, (9) ℓlap=1λmax(L)21V‖Lx^0−Lx~0‖22, _lap= 1 _max(L)^2\, 1V \|L x_0-L x_0 \|_2^2, which is zero when the model reproduces Geant4’s local differences. The prefactor 1/λmax(L)21/ _ (L)^2 keeps the auxiliary gradient on a scale comparable to the denoising gradient, as the combination scheme of Section 4.2 assumes: the squared Laplacian norm inflates the loss by up to λmax(L)2 _ (L)^2, so dividing by λmax(L)2≈95.5 _ (L)^2≈ 95.5 (λmax(L)≈9.77 _ (L)≈ 9.77, computed once and cached) bounds the operator gain by one. L is a fixed sparse operator applied twice per step, adding only O(V)O(V) cost and no second forward pass. 4.2. GradBlend: A Denoising-Anchored Combination Section 3 posed the combination problem: build a single update from the denoising gradient gdg_d and the auxiliary gradient gag_a without letting the auxiliary displace denoising. A rule that treats the two gradients as peers fails this under conflict. GradBlend separates the update into a direction and a magnitude and assigns the two gradients different roles in each: the auxiliary may steer the direction, but the magnitude is anchored to denoising alone. Let g^d=gd/∥gd∥ g_d=g_d/ g_d and g^a=ga/∥ga∥ g_a=g_a/ g_a be the unit directions of the two gradients, and let (10) θconf=arccos(g^d⋅g^a) _conf= \! ( g_d· g_a ) be the conflict angle between them: below 90∘90 the objectives reinforce, above 90∘90 they oppose. A coefficient ρ, fixed to 11 throughout, sets the auxiliary’s share of a tentative blend direction, (11) g^tent=g^d+ρg^a∥g^d+ρg^a∥,cosφalign=g^d⋅g^tent, g_tent= g_d+ρ\, g_a g_d+ρ\, g_a , _align= g_d· g_tent, where the alignment angle φalign _align measures how far the tentative blend rotates away from the denoising direction. With ρ=1ρ=1 the tentative blend is the bisector of the two unit directions, and the two angles are related in closed form, (12) cosφalign=1+cosθconf2=cos(θconf2), _align= 1+ _conf2= \! ( _conf2 ), so the entire construction is parameterized by the single conflict angle θconf _conf. Magnitude. The step magnitude is anchored to the denoising gradient and contracts with conflict: (13) mref=∥gd∥⋅max(cosφalign, 0.05)=∥gd∥⋅max(cosθconf2, 0.05).m_ref= g_d · \! ( _align,\,0.05 )= g_d · \! ( _conf2,\,0.05 ). The auxiliary enters only through the direction, never the step size: the update equals a pure denoising step when the two gradients agree and shrinks as they conflict, so the auxiliary cannot inflate the effective learning rate, and the update stays conservative under conflict. The floor of 0.050.05 keeps the step from collapsing under near-antiparallel gradients. Direction. The auxiliary contributes to the direction only when the conflict is moderate. Its effective coefficient is gated by θconf _conf: held at full weight below 120∘120 , decayed linearly to zero between 120∘120 and 150∘150 , and suppressed above. The final direction is the unit-normalized ρeff _eff-weighted blend of g^d g_d and g^a g_a; because both inputs are unit vectors, neither gradient can dominate the direction by raw norm. Note that the magnitude rule (13) uses the ungated alignment of Eq. (11), computed at the full coefficient ρ, even in regimes where the gate reduces or removes the auxiliary from the direction. This is deliberate: a large measured conflict signals that the auxiliary disagrees strongly with denoising in the current region of parameter space, and we shrink the step as a stability precaution even when the conflicting direction is discarded. It also decouples the magnitude rule from the gate thresholds, keeping the step size a smooth function of θconf _conf alone. Earlier variants that coupled the magnitude to the auxiliary norm were unstable in training; this rule stabilized them. The GradBlend update. Algorithm 1 gives the full construction.111The implementation adds one numerical guard: it also falls back to g^d g_d if the blended direction is nearly degenerate, ∥g^d+ρeffg^a∥<10−3 g_d+ _eff\, g_a <10^-3. Given the gate, this cannot occur for any θconf _conf (the norm is bounded below by 0.50.5 on θconf≤150∘ _conf≤ 150 ), and it never triggered in our runs. Algorithm 1 GradBlend update 1:denoising gradient gdg_d, auxiliary gradient gag_a 2:blended update gblendg_blend 3:g^d←gd/∥gd∥ g_d← g_d/ g_d ; g^a←ga/∥ga∥ g_a← g_a/ g_a 4:θconf←arccos(g^d⋅g^a) _conf← ( g_d· g_a) 5:mref←∥gd∥⋅max(cos(θconf/2), 0.05)m_ref← g_d · \! ( ( _conf/2),\,0.05 ) ⊳ Eqs. (12)–(13) 6:ρeff←1θconf≤120∘1−θconf−120∘30∘120∘<θconf≤150∘0θconf>150∘ _eff← cases1& _conf≤ 120 \\ 1- _conf-120 30 &120 < _conf≤ 150 \\ 0& _conf>150 cases 7:g^dir←(g^d+ρeffg^a)/∥g^d+ρeffg^a∥ g_dir←( g_d+ _eff\, g_a)\,/\, g_d+ _eff\, g_a 8:gblend←mref⋅g^dirg_blend← m_ref· g_dir 9:return gblendg_blend Relative to a pure denoising step, the squared step-size ratio is R2=max(cos(θconf/2),0.05)2R^2= ( ( _conf/2),0.05)^2, reaching one only at perfect alignment (Appendix D.3). The 150∘150 fallback and the denoising anchor together keep denoising primary. 4.3. Auxiliary Weight Schedule The schedule waux(s)w_aux(s) decides when the auxiliary loss is active. Because GradBlend normalizes each gradient to unit length and anchors the step size to the denoising-gradient norm (Section 4.2), any positive value of waux(s)w_aux(s) cancels from the update. The schedule is therefore a gate: the update is pure denoising when waux(s)=0w_aux(s)=0, and the auxiliary enters at full unit weight otherwise. Lantern consequently exposes no loss-weight hyperparameter; the only design choice is when the gate opens and when it closes, which we study through three gating windows in Section 5.4. We implement the gate with cosine ramps rather than a hard switch: the ramps set the scalar auxiliary loss we report, and they are what the multi-task baselines consume, since their updates scale with waux(s)w_aux(s) directly. Every method that uses the auxiliary loss receives the same schedule, so differences between methods reflect the combination rule and not the schedule. 5. Experiments Dataset All experiments use Dataset 2 of the Fast Calorimeter Simulation Challenge (CaloChallenge) (Faucci Giannelli et al., 2022a), which contains electron showers simulated with Geant4 and incident energies drawn log-uniformly between 1 GeV and 1 TeV. Each shower has V=6,480V=6,480 voxels over 45 layers, with 16 azimuthal (φ ) and 9 radial bins per layer. The first file of 100,000 showers is used for training and the second as the held-out evaluation reference. Dataset 2 is a standard benchmark whose granularity is representative of modern calorimeters, and our method carries over to the finer Dataset 3 (Faucci Giannelli et al., 2022b) from CaloChallenge 2022 geometry without modification. Experimental Setup The surrogate is the two-stage model of Section 4: an energy network that generates the per-layer energy ratios u and a shape network that generates the voxelized shower. The two auxiliary losses and GradBlend attach to the shape network as described there. This subsection fixes the training configuration.222Code is available at https://github.com/Aaheer17/lantern-kdd27. Both networks are denoising diffusion probabilistic models (Ho et al., 2020) with 1,000 diffusion steps. Preprocessing follows the transform defined in Appendix B.1. Training uses a single NVIDIA V100 GPU with 32 GB of memory for both networks. The energy network is trained once for 500 epochs in approximately one hour and is fixed across all configurations. The shape network trains for 800 epochs and dominates the overall training cost, taking approximately 38 hours for most gradient-based methods (ConFIG, GradBlend, IMTL-G, and PCGrad); GradNorm is somewhat more expensive at roughly 50 hours, owing to its additional per-task gradient-norm computation. All remaining hyperparameters are listed in Table 4 of Appendix B. Three shape networks are trained independently with different random seeds. All results are reported as mean and standard deviation over the three seeds. 5.1. Evaluation Metrics We report CFD (Section 1) at two scales: layer-wise, over per-layer energies (Appendix A.7), and voxel-wise, over the longitudinal structure between adjacent layers, defined next. Voxel-wise CFD Voxel-wise CFD measures how well a model reproduces the longitudinal correlation between a voxel and the voxel at the same transverse position in the following layer. For each consecutive layer pair (i,i+1)(i,i+1) with i=1,…,44i=1,…,44, and each transverse position (a,r)(a,r) with angular bin a=1,…,16a=1,…,16 and radial bin r=1,…,9r=1,…,9, we compute the Pearson correlation across showers between the energy at (i,a,r)(i,a,r) and at (i+1,a,r)(i+1,a,r): (14) ρa,r(i)=∑k(Ei,a,r,k−E¯i,a,r)(Ei+1,a,r,k−E¯i+1,a,r)∑k(Ei,a,r,k−E¯i,a,r)2∑k(Ei+1,a,r,k−E¯i+1,a,r)2,ρ^(i)_a,r= _k (E_i,a,r,k- E_i,a,r ) (E_i+1,a,r,k- E_i+1,a,r ) _k (E_i,a,r,k- E_i,a,r )^2\; _k (E_i+1,a,r,k- E_i+1,a,r )^2, where Ei,a,r,kE_i,a,r,k is the energy at position (a,r)(a,r) in layer i for shower k, the index k runs over the K showers in the set, and E¯i,a,r=1K∑kEi,a,r,k E_i,a,r= 1K _kE_i,a,r,k is the mean energy at that position across showers. The superscript (i)(i) identifies the layer pair (i,i+1)(i,i+1), and the subscript (a,r)(a,r) the transverse position. Collecting ρa,r(i)ρ^(i)_a,r over all 44 layer pairs and all 16×916× 9 transverse positions gives a correlation tensor ρ∈ℝ44×16×9ρ ^44× 16× 9 for Geant4 (ρrefρ^ref) and for the generated samples (ρgenρ^gen). We evaluate only the active positions, defined as those where both distributions carry nonzero variance, since the correlation is undefined elsewhere. Voxel-wise CFD is the normalized Frobenius distance between the two correlation tensors, (15) CFDvox=‖ρgen−ρref‖F‖ρref‖F,‖ρ‖F=∑i,a,r(ρa,r(i))2,CFD_vox= \|ρ^gen-ρ^ref \|_F \|ρ^ref \|_F, \|ρ \|_F= _i,a,r (ρ^(i)_a,r )^2\,, with both sums restricted to the active positions. Lower values indicate closer agreement with Geant4’s longitudinal correlation structure. Alongside CFD and FPD/KPD, we report precision, recall, density, and coverage (PRDC) (Naeem et al., 2020), which separate sample quality from diversity, and the binary classifier test (Krause and Shih, 2021), which trains a network to distinguish generated from Geant4 showers on low-level (voxel) and high-level (physics) features; an AUC near 0.5 indicates the two are hard to separate. Protocol details for all metrics in Appendix A. 5.2. Baselines We consider two groups of baselines. Prior surrogates compare the absolute quality of the generated showers; combination baselines isolate the gradient combination rule as the varying component. The prior surrogates are CaloDiffusion (Amram and Pedro, 2023), CaloScore v2 (Mikuni and Nachman, 2024), CaloDREAM (Favaro et al., 2025), CaloDiT (Raikwar et al., 2024) and CaloDiT-2 (Raikwar et al., 2025). CaloDREAM uses conditional flow matching. Standard metrics come from the CaloChallenge report (Krause et al., 2025), which evaluated all submissions with the official pipeline against the same reference file used here. We compute CFD for all models with a single procedure: from the publicly released submission samples for the prior surrogates (Faucci Giannelli et al., 2025), and from our own trained models for the combination baselines. The exception is CaloDiT-2, which postdates the CaloChallenge evaluation: its standard metrics come from its publication (Raikwar et al., 2025) and its CFD is omitted (Appendix A.1). The combination baselines hold the architecture, auxiliary losses, weight schedule, and preprocessing fixed, and vary only the gradient-combination rule. This group contains the denoising-only model, trained without auxiliary losses, and four combination methods: PCGrad (Yu et al., 2020), GradNorm (Chen et al., 2018), IMTL-G (Liu et al., 2021), and ConFIG (Liu et al., 2025). Beyond these gradient-level rules, we also consider two basic loss-weighting schemes, weighted sum and uncertainty weighting (Kendall et al., 2018), which merge both objectives into one scalar loss and backpropagate once. Acting on the summed loss rather than individual gradients, they cannot address the gradient conflict this work targets, so we defer them to Appendix D.2: plain weighted sum and unscheduled uncertainty weighting inflate FPD by orders of magnitude, and only a scheduled uncertainty variant, which reuses our weight schedule (Section 4.3), is competitive (Tables 7 and 8). Base denotes the denoising-only model. Models that add the auxiliary losses under a combination rule are written Base+X, for example Base+PCGrad. The proposed model, Base+GradBlend, is written Lantern. The prior surrogates keep their original names. 5.3. Results and Discussion Table 1. Shower fidelity on CaloChallenge Dataset 2. FPD and KPD are scaled by 10310^3; lower is better for FPD, KPD, and CFDvoxCFD_vox, and for AUC values closer to 0.50.5 are better. Classifier AUC is reported for high-level (HL) and low-level (L) features. CFDvoxCFD_vox is the voxel-wise Correlation Frobenius Distance of Eq. (15), computed in this work for all models including the prior surrogates; the layer-wise variant is defined in Appendix A.7. Remaining standard metrics for the prior surrogates are taken from the CaloChallenge report (Krause et al., 2025). All auxiliary-loss models are trained under schedule S3 (Section 5.4); differences reflect the combination rule alone. Trained models report mean ± std over three seeds, and the best value among the trained models per column is in bold. Model Loss FPD ↓ KPD ↓ CFDvoxCFD_vox ↓ AUC (HL) AUC (L) CaloDiffusion (Amram and Pedro, 2023) N/A 146.9±0.9146.9± 0.9 0.174±0.0420.174± 0.042 0.1550.155 0.591±0.0090.591± 0.009 0.577±0.0040.577± 0.004 CaloScore v2 (Mikuni and Nachman, 2024) N/A 112.5±0.9112.5± 0.9 0.149±0.0570.149± 0.057 0.1760.176 0.666±0.0020.666± 0.002 0.595±0.0030.595± 0.003 CaloDiT (Raikwar et al., 2024) N/A 1691±71691± 7 11.030±0.42711.030± 0.427 0.9790.979 0.912±0.0020.912± 0.002 0.984±0.0010.984± 0.001 CaloDiT-2 (EDM) (Raikwar et al., 2025) N/A 20.06±0.7420.06± 0.74 0.089±0.0900.089± 0.090 – 0.560±0.0050.560± 0.005 0.594±0.0020.594± 0.002 CaloDiT-2 (CD) (Raikwar et al., 2025) N/A 68.83±3.3768.83± 3.37 0.390±0.1500.390± 0.150 – 0.695±0.0040.695± 0.004 0.598±0.0020.598± 0.002 CaloDREAM (Favaro et al., 2025) N/A 24.65±1.0424.65± 1.04 0.023±0.0360.023± 0.036 0.1130.113 0.521±0.0020.521± 0.002 0.531±0.0030.531± 0.003 Base None 30.87±1.2630.87± 1.26 0.002±0.0010.002± 0.001 0.122±0.0040.122± 0.004 0.531±0.0090.531± 0.009 0.545±0.0130.545± 0.013 Base+PCGrad Voxel 68.50±44.5868.50± 44.58 0.137±0.1910.137± 0.191 0.142±0.0080.142± 0.008 0.566±0.0270.566± 0.027 0.560±0.0030.560± 0.003 Base+GradNorm Voxel 60.27±11.1560.27± 11.15 0.102±0.0420.102± 0.042 0.142±0.0140.142± 0.014 0.549±0.0210.549± 0.021 0.563±0.0010.563± 0.001 Base+IMTL-G Voxel 406.9±80.9406.9± 80.9 2.216±0.5382.216± 0.538 0.192±0.0120.192± 0.012 0.807±0.0190.807± 0.019 0.763±0.0110.763± 0.011 Base+ConFIG Voxel 204.7±31.2204.7± 31.2 0.882±0.1100.882± 0.110 0.156±0.0090.156± 0.009 0.636±0.0030.636± 0.003 0.592±0.0020.592± 0.002 Lantern Voxel 32.84±3.2732.84± 3.27 −0.0007±0.0046-0.0007± 0.0046 0.125±0.0050.125± 0.005 0.538±0.0020.538± 0.002 0.558±0.0020.558± 0.002 Base+PCGrad Laplacian 169.0±72.8169.0± 72.8 0.640±0.3920.640± 0.392 0.155±0.0030.155± 0.003 0.621±0.0330.621± 0.033 0.580±0.0210.580± 0.021 Base+GradNorm Laplacian 133.5±51.5133.5± 51.5 0.472±0.2980.472± 0.298 0.147±0.0110.147± 0.011 0.620±0.0360.620± 0.036 0.582±0.0050.582± 0.005 Base+IMTL-G Laplacian 3454±48203454± 4820 23.520±34.50023.520± 34.500 0.510±0.3130.510± 0.313 0.874±0.1060.874± 0.106 0.890±0.0880.890± 0.088 Base+ConFIG Laplacian 262.8±73.9262.8± 73.9 1.209±0.3411.209± 0.341 0.172±0.0070.172± 0.007 0.656±0.0330.656± 0.033 0.623±0.0150.623± 0.015 Lantern Laplacian 25.97±1.4425.97± 1.44 −0.0048±0.0009-0.0048± 0.0009 0.120±0.0050.120± 0.005 0.534±0.0030.534± 0.003 0.554±0.0020.554± 0.002 Table 1 reports shower fidelity on Dataset 2. Among the prior surrogates, CaloDiffusion, CaloScore v2, and the original CaloDiT trail on FPD and CFD, CaloDiT severely so. CaloDiT-2 (EDM) reports the lowest FPD in the table, its single-step CD variant trading fidelity for sampling speed, but both show weaker classifier AUCs than CaloDREAM and their CFD cannot be assessed (Section 5.2). CaloDREAM, the flow-matching model our two-stage architecture adapts, attains the best CFD and classifier AUCs and serves as our reference; Lantern reaches comparable FPD and CFD with a DDPM backbone. Base is a strong baseline in its own right, reaching an FPD of 30.930.9 and a CFD of 0.12150.1215 without any auxiliary loss. This sets a high bar: an auxiliary physics loss is useful only if it can be added without disturbing this baseline. The general-purpose multi-task gradient methods do not meet that bar. Every Base+PCGrad, Base+GradNorm, Base+IMTL-G, and Base+ConFIG row has a higher FPD than Base, in several cases by more than an order of magnitude, with standard deviations large relative to the means. Base+IMTL-G is the worst, diverging on at least one seed under the Laplacian loss. These are unstable outcomes, not small regressions. The gradient diagnostics explain the pattern. Because these methods treat the denoising and auxiliary gradients as co-equal, the combined update drifts from the denoising direction that produces valid samples, degrading fidelity; Section 5.5 reads the per-method allocations from Table 3. The training curves in Figure 1 show the resulting instability, the auxiliary loss rising by orders of magnitude and strongest under the Laplacian loss, whose high-pass gradient is noisier. Lantern is the only combination method that adds the auxiliary loss without this damage. Under the Laplacian loss it reaches the best FPD and CFD among the trained models, slightly ahead of Base; under the voxel loss it stays close to Base, with small standard deviations. Lantern integrates the auxiliary objective and preserves the fidelity of the denoising-only model. (a) Laplacian loss. (b) Voxel residual loss. Figure 1. Unweighted auxiliary loss over training on Dataset 2. The general-purpose multi-task gradient methods produce large fluctuations, rising by several orders of magnitude at points during training. GradBlend (Lantern) keeps the auxiliary loss stable, close to the Base curve that never optimizes it. The dashed line marks the epoch at which the auxiliary weight begins to decay. KPD does not separate the strongest models: Lantern and Base are both near zero, Lantern slightly negative. KPD is an unbiased MMD estimate that can go slightly negative when samples closely match the reference, so this is a close match, not an error, while the failing methods have large positive KPD consistent with their FPD. Classifier AUCs in the Base and Lantern rows are all near 0.50.5. Precision, recall, density, and coverage confirm the failure mode: the multi-task baselines gain precision and density but collapse in recall, the signature of mode collapse, whereas Base and Lantern stay balanced and close to the reference-quality prior surrogates (full breakdown in Appendix D.1, Table 5). 5.4. Ablation Studies We compare three gating windows, holding every other setting fixed. Schedule S1 opens the gate at the first optimization step. S2 opens it after a 50-epoch warmup, reaching full weight by epoch 250 and holding to the end. S3 matches S2 but closes the gate before the end, decaying to zero by epoch 650 for the voxel loss and 550 for the Laplacian loss; the closing epochs were fixed a priori, not tuned. The comparison is factored: S1 versus S2 tests whether denoising must settle before the auxiliary enters, S2 versus S3 whether it must finish alone. Timing decides voxel sample quality. Held to the last epoch (S1, S2), the voxel loss inflates FPD and density far above one while recall collapses (Table 2; full PRDC in Appendix D.1): the generator concentrates mass on part of the real distribution and fails to cover it. Coverage stays near 0.990.99 throughout, so this is a recall collapse, not a coverage failure. S3 reverses it, returning FPD, density, and the AUCs to near their calibrated values; precision drops correspondingly, so the terminal gate trades a little per-sample fidelity for a large gain in coverage and calibration. The Laplacian loss, by contrast, is insensitive to the schedule: FPD spans only 23.523.5 to 26.026.0 and CFDvoxCFD_vox 0.1180.118 to 0.1290.129 across S1 through S3, and its precision, density, recall, and coverage agree across schedules to within run-to-run variation (Table 6). The voxel loss competes with denoising late in training while the Laplacian loss does not, so gating it leaves the optimization trajectory and the samples unchanged. The gate is therefore necessary where the auxiliary conflicts with denoising and inert where it does not. Lantern adopts S3 for both losses, removing the voxel conflict in the final epochs at no cost to the Laplacian loss; this keeps a single schedule across auxiliaries and avoids per-loss tuning. Table 2. Effect of the auxiliary loss schedule on Lantern. Auxiliary Sched. FPD ↓ KPD ↓ CFDvoxCFD_vox ↓ AUCHL AUCLL Voxel S1 173.4173.4 0.9970.997 0.1230.123 0.7770.777 0.7830.783 S2 236.3236.3 1.5281.528 0.1380.138 0.8330.833 0.8380.838 S3 32.8432.84 −0.0007-0.0007 0.1250.125 0.5380.538 0.5580.558 Laplacian S1 25.3425.34 0.0010.001 0.1180.118 0.5340.534 0.5610.561 S2 23.5323.53 0.0010.001 0.1290.129 0.5360.536 0.5600.560 S3 25.9725.97 −0.0048-0.0048 0.1200.120 0.5340.534 0.5540.554 Table 3. Gradient diagnostics over the active window, voxel energy loss on Dataset 2, reported as mean ± std over three seeds. θconf _conf is the conflict angle between the denoising gradient gdg_d and the auxiliary gradient gag_a; R1dR_1^d and R1aR_1^a are first-order progress ratios of the blended step gblendg_blend on the two objectives; Balance is R1d/R1aR_1^d/R_1^a, where a value above one favors the denoising objective; cosd _d and cosa _a are the cosines of gblendg_blend with gdg_d and gag_a; %spike is the fraction of active epochs removed by the auxiliary-norm spike filter. Bold marks the best value among the combination methods per column. Method θconf _conf R1dR_1^d R1aR_1^a Balance cosd _d cosa _a %spike Base+ConFIG 92.0±1.292.0± 1.2 1.42±0.051.42± 0.05 0.78±0.010.78± 0.01 1.81±0.051.81± 0.05 0.69±0.010.69± 0.01 0.69±0.010.69± 0.01 12±012± 0 Base+GradNorm 127.5±3.6127.5± 3.6 0.55±0.090.55± 0.09 0.66±0.070.66± 0.07 0.83±0.050.83± 0.05 0.36±0.020.36± 0.02 0.41±0.030.41± 0.03 9±09± 0 Base+IMTL-G 116.7±3.3116.7± 3.3 0.14±0.010.14± 0.01 0.10±0.010.10± 0.01 1.33±0.111.33± 0.11 0.52±0.020.52± 0.02 0.52±0.020.52± 0.02 13±113± 1 Base+PCGrad 123.9±7.1123.9± 7.1 0.95±0.170.95± 0.17 0.44±0.050.44± 0.05 2.13±0.202.13± 0.20 0.55±0.080.55± 0.08 0.31±0.020.31± 0.02 11±011± 0 Lantern 119.7±0.5119.7± 0.5 0.71±0.020.71± 0.02 0.24±0.070.24± 0.07 3.26±1.163.26± 1.16 0.71±0.020.71± 0.02 0.21±0.040.21± 0.04 ±8± 0 5.5. Training Diagnostics Table 3 reports gradient diagnostics for the voxel energy loss, measured over the active window and averaged over three seeds. Each row summarizes how a combination method allocates its update between the denoising and auxiliary objectives. Definitions of the ratios, the active window, and the spike filter are given in Appendix D.3. The two cosine columns separate the methods. For Base+ConFIG and Base+IMTL-G, cosd _d equals cosa _a in every seed. These methods weight the two gradients equally, so the resulting step makes the same angle with each. Lantern shows the widest gap between the two cosines, 0.710.71 against 0.210.21. Its step stays aligned with the denoising gradient and is held away from the auxiliary gradient. Base+PCGrad shows a partial gap, Base+GradNorm a small one. The Balance column reads which objective each step favors. Lantern keeps the denoising objective ahead in every seed, with Balance above two. Base+GradNorm inverts this order and favors the auxiliary gradient over the denoiser, with Balance below one. A large R1dR_1^d is not a measure of sample quality. Base+ConFIG reaches R1d=1.42R_1^d=1.42, the highest in the table, yet has poor FPD (Table 1). The ratios describe how the step is allocated and should be read together with the evaluation tables. The spike column reports the fraction of active epochs removed by the auxiliary-norm spike filter. Lantern has the lowest rate at eight percent, and Base+IMTL-G the highest. This ordering matches the stability seen in the training curves (Figure 1). Figure 1 shows the unweighted auxiliary loss over training. The general-purpose multi-task methods fluctuate by orders of magnitude as the auxiliary gradient periodically overtakes the update, whereas Lantern keeps it smooth, matching the Base curve that never optimizes it while still carrying the auxiliary signal. 6. Conclusion Physics-informed diffusion typically relies on a closed-form physical law to supervise generation, which calorimeter shower simulation lacks. We presented Lantern, which guides a diffusion surrogate with two per-sample physics losses, a variance-stabilized voxel energy residual and a graph Laplacian smoothness term, and merges their gradients with the denoising gradient through GradBlend, a combination rule that keeps denoising as the trusted reference direction while admitting each physics gradient where it agrees. These losses and GradBlend apply to the second stage, which predicts voxel energies conditioned on the first-stage incident energy and layer ratios. To measure fidelity FPD and KPD miss, we introduced CFD, the Correlation Frobenius Distance, which scores the correlation structure of calorimeter showers at two scales. On CaloChallenge Dataset 2, Lantern integrates both physics losses without the degradation symmetric multi-task rules incur: PCGrad, GradNorm, IMTL-G, and ConFIG raise FPD by factors of two to over one hundred, whereas Lantern stays close to the denoising-only baseline on FPD and KPD and remains competitive with the strongest prior surrogates on CFD and classifier AUC. We treat denoising as the primary objective; the ablations show that an auxiliary loss strengthens the model as long as its gradient stays aligned with it. The voxel gradient turns from cooperative to opposing over training, so holding it at full weight to the final epoch inflates FPD and collapses recall; a terminal phase that trains on the primary loss alone removes this and recovers both. The Laplacian gradient stays aligned throughout training and is insensitive to the schedule, so a single schedule serves both losses without per-loss tuning. We discuss limitations and next steps in Section 7. 7. Limitations and Ethical Considerations Limitations. GradBlend and the two physics losses are demonstrated on a single controlled setting; transfer to other configurations remains future work. The geometries released with LEMURS (Raikwar et al., 2025), the higher-granularity Dataset 3, and hadronic showers are direct targets: adapting the method to them is an application, not a reformulation. The approach carries two sets of empirically set thresholds. The weight schedule’s closing epochs, and GradBlend’s angular thresholds (120∘120 for decay and 150∘150 for the fallback) together with the magnitude floor (0.050.05), were stable across the voxel and Laplacian losses here, but neither is derived, and both the schedule and the direction rule warrant further analysis, ideally tied to the measured conflict angle. We study each auxiliary loss on its own and do not characterize how the voxel-residual and Laplacian objectives interact when trained together. Results are reported over three random seeds of the shape network with the energy network fixed, so the reported variance reflects shape-network initialization alone; three seeds is common but limits power to resolve small metric differences. Physics guidance acts entirely during training and leaves the sampler unchanged, so fidelity gains come at no added generation cost. Reducing the base surrogate’s sampling cost is orthogonal to our contributions and can be combined with them. Ethical considerations. This work develops a surrogate model for calorimeter shower simulation in high-energy physics, a fundamental-physics application trained and evaluated entirely on simulated detector data from Geant4, with no human subjects, personal or sensitive data, or application to individuals. The intended benefit is reduced computational cost for physics simulation, lowering the energy and hardware demands of large-scale Monte Carlo production. Generated samples approximate rather than replace full simulation, as with any surrogate, so downstream physics analyses using them should account for the residual modeling error characterized here, particularly the correlation-structure error CFD is designed to expose. Calorimeter shower generation has no foreseeable harmful dual-use application beyond its intended scientific scope. 8. Generative AI Usage A large language model assisted with manuscript editing and LaTeX formatting, portions of the experiment code, and cross-checking equations and related work. All contributions, including the method, losses, metric, and experimental design, are the authors’ own; the authors have reviewed all AI-assisted material and take full responsibility for the paper’s content, correctness, and originality. References (1) Agostinelli et al. (2003) S. Agostinelli et al. 2003. GEANT4—A Simulation Toolkit. Nuclear Instruments and Methods in Physics Research A 506, 3 (2003), 250–303. Ahmad et al. (2024) Farzana Yasmin Ahmad, Vanamala Venkataswamy, and Geoffrey Fox. 2024. A comprehensive evaluation of generative models in calorimeter shower simulation. arXiv preprint arXiv:2406.12898 (2024). Amram and Pedro (2023) Oz Amram and Kevin Pedro. 2023. Denoising Diffusion Models with Geometry Adaptation for High Fidelity Calorimeter Simulation. Physical Review D 108 (2023), 072014. Apollinari et al. (2017) Giorgio Apollinari, O Brüning, Tatsushi Nakamoto, and Lucio Rossi. 2017. High luminosity large hadron collider HL-LHC. arXiv preprint arXiv:1705.08830 (2017). Baldan et al. (2025) Giacomo Baldan, Qiang Liu, Alberto Guardone, and Nils Thuerey. 2025. Flow matching meets pdes: A unified framework for physics-constrained generation. arXiv preprint arXiv:2506.08604 (2025). Bastek et al. (2025) Jan-Hendrik Bastek, WaiChing Sun, and Dennis M. Kochmann. 2025. Physics-Informed Diffusion Models. In International Conference on Learning Representations (ICLR). arXiv:2403.14404 [cs.LG] Birk et al. (2024) Joschka Birk, Anna Hallin, and Gregor Kasieczka. 2024. OmniJet-α: The first cross-task foundation model for particle physics. Machine Learning: Science and Technology 5, 3 (2024), 035031. doi:10.1088/2632-2153/ad66ad Buhmann et al. (2024) Erik Buhmann, Frank Gaede, Gregor Kasieczka, Anatolii Korol, William Korcari, Katja Krüger, and Peter McKeown. 2024. CaloClouds I: Ultra-Fast Geometry-Independent Highly-Granular Calorimeter Simulation. Journal of Instrumentation 19 (2024), P04020. CaloChallenge Collaboration (2022) CaloChallenge Collaboration. 2022. Fast Calorimeter Simulation Challenge 2022 – Homepage and Evaluation Code. https://github.com/CaloChallenge/homepage. Accessed: 2026-07-13. Chen et al. (2018) Zhao Chen, Vijay Badrinarayanan, Chen-Yu Lee, and Andrew Rabinovich. 2018. GradNorm: Gradient Normalization for Adaptive Loss Balancing in Deep Multitask Networks. In Proceedings of the 35th International Conference on Machine Learning (ICML) (Proceedings of Machine Learning Research, Vol. 80). PMLR, 794–803. Chung et al. (2023) Hyungjin Chung, Jeongsol Kim, Michael T. McCann, Marc L. Klasky, and Jong Chul Ye. 2023. Diffusion Posterior Sampling for General Noisy Inverse Problems. In International Conference on Learning Representations (ICLR). Dosovitskiy et al. (2020) Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. 2020. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.11929 (2020). Du et al. (2018) Yunshu Du, Wojciech M. Czarnecki, Siddhant M. Jayakumar, Mehrdad Farajtabar, Razvan Pascanu, and Balaji Lakshminarayanan. 2018. Adapting Auxiliary Losses Using Gradient Similarity. arXiv:1812.02224. Fabjan and Gianotti (2003) Christian W Fabjan and Fabiola Gianotti. 2003. Calorimetry for particle physics. Reviews of Modern Physics 75, 4 (2003), 1243. Faucci Giannelli et al. (2022a) Michele Faucci Giannelli, Gregor Kasieczka, Claudius Krause, Ben Nachman, Dalila Salamani, David Shih, and Anna Zaborowska. 2022a. Fast Calorimeter Simulation Challenge 2022 - Dataset 2. doi:10.5281/zenodo.6366271 Faucci Giannelli et al. (2022b) Michele Faucci Giannelli, Gregor Kasieczka, Claudius Krause, Ben Nachman, Dalila Salamani, David Shih, and Anna Zaborowska. 2022b. Fast Calorimeter Simulation Challenge 2022 - Dataset 3. doi:10.5281/zenodo.6366324 Faucci Giannelli et al. (2025) Michele Faucci Giannelli, Gregor Kasieczka, Claudius Krause, Benjamin Nachman, Dalila Salamani, David Shih, Anna Zaborowska, Oz Amram, Kerstin Borras, Matthew Buckley, Thorsten Buss, Renato Paulo Da Costa Cardoso, Vijay Ekambaram, Florian Ernst, Luigi Favaro, Frank Gaede, Shih-Chieh Hsu, Kristina Jaruskova, Benno Käch, Jayant Kalagnanam, Dirk Krücker, Qibin Liu, Xiulong Liu, Thandikire Madula, Isabell-Alissandra Melzer-Pellmann, Vinicius Mikuni, Nam Nguyen, Ayodele Ore, Sofia Palacios Schweitzer, Ian Pang, Kevin Pedro, Tilman Plehn, Piyush Raikwar, John Raine, Moritz Alfons Wilhelm Scham, Simon Schnake, Chase Shimmin, Eli Shlizerman, Li Shu, Mudhakar Srivatsa, Sofia Vallecorsa, and Kyongmin Yeo. 2025. Fast Calorimeter Simulation Challenge 2022 - Submissions Dataset 2. doi:10.5281/zenodo.15962050 Favaro et al. (2025) Luigi Favaro, Ayodele Ore, Sofia Palacios Schweitzer, and Tilman Plehn. 2025. CaloDREAM: Detector Response Emulation via Attentive Flow Matching. SciPost Physics 18 (2025), 088. Gao et al. (2023) Zhihan Gao, Xingjian Shi, Boran Han, Hao Wang, Xiaoyong Jin, Danielle C. Maddix, Yi Zhu, Mu Li, and Yuyang Wang. 2023. PreDiff: Precipitation Nowcasting with Latent Diffusion Models. In Advances in Neural Information Processing Systems (NeurIPS). Hassan et al. (2024) Majdi Hassan et al. 2024. ET-Flow: Equivariant Flow-Matching for Molecular Conformer Generation. In Advances in Neural Information Processing Systems (NeurIPS). arXiv:2410.22388 Ho et al. (2020) Jonathan Ho, Ajay Jain, and Pieter Abbeel. 2020. Denoising Diffusion Probabilistic Models. arXiv:2006.11239 [cs.LG] Hoogeboom et al. (2022) Emiel Hoogeboom, Victor Garcia Satorras, Clément Vignac, and Max Welling. 2022. Equivariant Diffusion for Molecule Generation in 3D. In International Conference on Machine Learning (ICML). Oral presentation. Jiang et al. (2026) Cheng Jiang, Sitian Qian, Kevin Pedro, Oz Amram, Huilin Qu, and Maggie Voetberg. 2026. CaloTrilogy: Toward a Breakthrough in One-Step, End-to-End, Physics-Guided Shower Generation for Modern Calorimeters. arXiv preprint arXiv:2606.04165 (2026). Kansal et al. (2023) Raghav Kansal, Anni Li, Javier Duarte, Nadezda Chernyavskaya, Maurizio Pierini, Breno Orzari, and Thiago Tomei. 2023. Evaluating generative models in high energy physics. Physical Review D 107, 7 (2023), 076017. Kendall et al. (2018) Alex Kendall, Yarin Gal, and Roberto Cipolla. 2018. Multi-task learning using uncertainty to weigh losses for scene geometry and semantics. In Proceedings of the IEEE conference on computer vision and pattern recognition. 7482–7491. Kobylianskii et al. (2024) Dmitrii Kobylianskii, Nathalie Soybelman, Etienne Dreyer, and Eilam Gross. 2024. CaloGraph: Graph-Based Diffusion Model for Fast Shower Generation in Calorimeters with Irregular Geometry. Physical Review D (2024). arXiv:2402.11575 Krause et al. (2025) Claudius Krause, Michele Faucci Giannelli, Gregor Kasieczka, Benjamin Nachman, Dalila Salamani, David Shih, Anna Zaborowska, Oz Amram, Kerstin Borras, Matthew R Buckley, et al. 2025. CaloChallenge 2022: a community challenge for fast calorimeter simulation. Reports on Progress in Physics 88, 11 (2025), 116201. Krause and Shih (2021) Claudius Krause and David Shih. 2021. CaloFlow I: Even faster and still accurate generation of calorimeter showers with normalizing flows. arXiv preprint arXiv:2110.11377 (2021). Kynkänniemi et al. (2019) Tuomas Kynkänniemi, Tero Karras, Samuli Laine, Jaakko Lehtinen, and Timo Aila. 2019. Improved precision and recall metric for assessing generative models. Advances in neural information processing systems 32 (2019). Leigh et al. (2024a) Matthew Leigh, Debajyoti Sengupta, Guillaume Quétant, John Andrew Raine, Knut Zoch, and Tobias Golling. 2024a. PC-JeDi: Diffusion for Particle Cloud Generation in High Energy Physics. SciPost Physics 16 (2024), 018. Leigh et al. (2024b) Matthew Leigh, Debajyoti Sengupta, John Andrew Raine, Guillaume Quétant, and Tobias Golling. 2024b. PC-Droid: Faster Diffusion and Improved Quality for Particle Cloud Generation. Physical Review D (2024). arXiv:2307.06836 Li et al. (2026) Yunyang Li, Lin Huang, Luojia Xia, et al. 2026. Elign: Equivariant Diffusion Model Alignment from Foundational Machine Learning Force Fields. arXiv preprint arXiv:2601.21985 (2026). arXiv:2601.21985 [cs.LG] Liang et al. (2025) Jinhao Liang, Yixuan Sun, Anirban Samaddar, Sandeep Madireddy, and Ferdinando Fioretto. 2025. Chance-Constrained Flow Matching. In Advances in Neural Information Processing Systems (NeurIPS). arXiv:2509.25157. Liu et al. (2021) Liyang Liu, Yi Li, Zhanghui Kuang, Jing-Hao Xue, Yimin Chen, Wenming Yang, Qingmin Liao, and Wayne Zhang. 2021. Towards impartial multi-task learning. In International conference on learning representations. Liu et al. (2025) Qiang Liu, Mengyu Chu, and Nils Thuerey. 2025. ConFIG: Towards Conflict-free Training of Physics Informed Neural Networks. In International Conference on Learning Representations (ICLR). arXiv:2408.11104 Mikuni and Nachman (2022) Vinicius Mikuni and Benjamin Nachman. 2022. Score-based generative models for calorimeter shower simulation. Physical Review D 106, 9 (2022), 092009. Mikuni and Nachman (2024) Vinicius Mikuni and Benjamin Nachman. 2024. CaloScore v2: single-shot calorimeter shower simulation with diffusion models. Journal of Instrumentation 19, 02 (2024), P02001. doi:10.1088/1748-0221/19/02/P02001 Naeem et al. (2020) Muhammad Ferjad Naeem, Seong Joon Oh, Youngjung Uh, Yunjey Choi, and Jaejun Yoo. 2020. Reliable fidelity and diversity metrics for generative models. In International conference on machine learning. PMLR, 7176–7185. Peebles and Xie (2023) William Peebles and Saining Xie. 2023. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF international conference on computer vision. 4195–4205. Platt and Barr (1988) John C. Platt and Alan H. Barr. 1988. Constrained Differential Optimization. In Neural Information Processing Systems (NIPS), Dana Z. Anderson (Ed.), Vol. 1. American Institute of Physics, 612–621. Raikwar et al. (2024) Piyush Raikwar, Renato Cardoso, Kristina Jaruskova, Dalila Salamani, Sofia Vallecorsa, Anna Zaborowska, et al. 2024. CaloDiT: Diffusion with transformers for fast shower simulation. Presentation at the 22nd International Workshop on Advanced Computing and Analysis Techniques in Physics Research (ACAT). https://indico.cern.ch/event/1330797/contributions/5796591/. Raikwar et al. (2025) Piyush Raikwar, Anna Zaborowska, Peter McKeown, Renato Cardoso, Mikolaj Piorczynski, and Kyongmin Yeo. 2025. A Generalisable Generative Model for Multi-Detector Calorimeter Simulation. arXiv preprint arXiv:2509.07700 (2025). Shuman et al. (2013) David I Shuman, Sunil K Narang, Pascal Frossard, Antonio Ortega, and Pierre Vandergheynst. 2013. The emerging field of signal processing on graphs: Extending high-dimensional data analysis to networks and other irregular domains. IEEE signal processing magazine 30, 3 (2013), 83–98. Song et al. (2021) Yang Song, Jascha Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. 2021. Score-Based Generative Modeling through Stochastic Differential Equations. arXiv:2011.13456 [cs.LG] Tian et al. (2025) Qingwen Tian, Yuxin Xu, Yixuan Yang, Zhen Wang, Ziqi Liu, Pengju Yan, and Xiaolin Li. 2025. EquiFlow: Equivariant Conditional Flow Matching with Optimal Transport for 3D Molecular Conformation Prediction. In AAAI Conference on Artificial Intelligence. Trupin et al. (2026) Noah Trupin, Rahul Ghosh, and Aadi Jangid. 2026. Learning Flow Distributions via Projection-Constrained Diffusion on Manifolds. arXiv preprint arXiv:2602.17773 (2026). Utkarsh et al. (2026) Utkarsh Utkarsh, Pengfei Cai, Alan Edelman, Rafael Gomez-Bombarelli, and Christopher Rackauckas. 2026. Physics-constrained flow matching: Sampling generative models with hard constraints. Advances in Neural Information Processing Systems 38 (2026), 160217–160252. Vaitl and Klein (2025) Lorenz Vaitl and Leon Klein. 2025. Path Gradients after Flow Matching. In Advances in Neural Information Processing Systems (NeurIPS). arXiv:2505.10139 [cs.LG] Wang et al. (2021) Sifan Wang, Yujun Teng, and Paris Perdikaris. 2021. Understanding and Mitigating Gradient Flow Pathologies in Physics-Informed Neural Networks. SIAM Journal on Scientific Computing 43, 5 (2021), A3055–A3081. Yan et al. (2023) Junjun Yan, Xinhai Chen, Zhichao Wang, Enqiang Zhou, and Jie Liu. 2023. Auxiliary-tasks learning for physics-informed neural network-based partial differential equations solving. arXiv preprint arXiv:2307.06167 (2023). Yu et al. (2020) Tianhe Yu, Saurabh Kumar, Abhishek Gupta, Sergey Levine, Karol Hausman, and Chelsea Finn. 2020. Gradient Surgery for Multi-Task Learning. In Advances in Neural Information Processing Systems (NeurIPS), Vol. 33. 5824–5836. Zhang and Zou (2025) Yi Zhang and Difan Zou. 2025. Physics-informed distillation of diffusion models for pde-constrained generation. arXiv preprint arXiv:2505.22391 (2025). Appendix A Evaluation Protocol and Metric Details A.1. Samples and Reference Data All metrics are computed against the second file of CaloChallenge Dataset 2, which contains 100,000 Geant4 electron showers held out from training (Krause et al., 2025). The same reference file is used for every model, including the prior surrogates we re-evaluate for CFD. For each of our trained models, we generate 100,000 showers with incident energies drawn from the reference distribution, matching the reference sample size. For the prior surrogates, the standard metrics—FPD, KPD, PRDC, and the classifier test—are taken directly from the CaloChallenge report (Krause et al., 2025) (Tables C15, C16, and C19), which evaluated all submissions with the official pipeline against this same reference file. CFD is not reported there, so we compute it ourselves from the public submission samples released on Zenodo (Faucci Giannelli et al., 2025), using the identical procedure applied to our own models (Appendix A.6). CaloDiT-2 postdates the CaloChallenge evaluation; its released artifacts include pre-trained models but no generated sample sets, so its CFD is omitted and we leave regeneration from the Geant4-distributed checkpoints to future work. A.2. Seeds and Aggregation The three-seed protocol and the fixed energy network are described in Section 5; the reported spread therefore reflects shape-network initialization alone. Every metric for a trained model is reported as mean ± std over the three seeds, each evaluated on its own generated sample set. Prior-surrogate rows are single values from the CaloChallenge report, carrying that report’s uncertainties where available. A.3. Standard Metric Pipeline We compute FPD, KPD, and the classifier test with the official CaloChallenge evaluation pipeline (evaluate.py) (CaloChallenge Collaboration, 2022), run unmodified, so that our trained models and the prior surrogates are scored under one protocol. Metric definitions follow the CaloChallenge report (Krause et al., 2025). FPD and KPD are computed on the standard set of physics-motivated shower features and reported scaled by 10310^3, as in the main text. A.4. Classifier Test Protocol The classifier test trains a binary classifier to separate generated showers (label 0) from Geant4 reference showers (label 1) and reports how well it succeeds. We evaluate two feature sets. The low-level classifier uses the voxel energies, each shower normalized by its incident energy. The high-level classifier uses the shower-shape observables of the CaloChallenge protocol (Krause et al., 2025): total deposited energy, per-layer energies, and the energy-weighted centroids and widths in the two angular directions, together with the incident energy. Both feature sets apply the low-energy cut of 0.015 MeV used throughout our evaluation. The classifier is the standard CaloChallenge network (Krause et al., 2025), a multilayer perceptron with two hidden layers of 512 units, trained with the pipeline’s default optimizer, learning rate, and batch size. The combined samples are split into equal-sized classes and then into training, test, and validation partitions. We select the checkpoint with the best test accuracy and report the area under the ROC curve (AUC) on the held-out validation partition; an AUC near 0.5 indicates the two sample sets are indistinguishable. We report AUC averaged over three independent runs, each using samples from a separately seeded shape network, so the spread reflects variation across trained models. Prior-surrogate AUC values are taken from the CaloChallenge report, which averages ten evaluation runs (Krause et al., 2025). A.5. PRDC Settings Precision, recall, density, and coverage (PRDC) (Naeem et al., 2020) are computed with k=5k=5 nearest neighbors, matching the CaloChallenge protocol, on the same feature space as the high-level classifier. Precision and density measure sample quality (generated samples lying on the reference manifold); recall and coverage measure diversity (reference samples reachable from the generated manifold). Prior-surrogate values are taken from Table C19 of the CaloChallenge report (Krause et al., 2025); trained models report mean ± std over three seeds. Results are given in Table 5 in Appendix D.1. A.6. CFD Implementation Details CFD is computed with a single procedure for every model in Tables 1 and 5: from the public Zenodo submission samples for the prior surrogates, and from samples drawn from our own trained models otherwise. Correlations are computed on the voxel energies stored in the sample files (physical units, MeV), which already incorporate the 0.015 MeV low-energy cut; no further thresholding is applied at the correlation stage. For each model the generated and reference sets are truncated to a common size of 100,000 showers. A transverse position enters the voxel-wise computation only if it carries nonzero variance in both the reference and generated samples, since the Pearson correlation is otherwise undefined. For Dataset 2 with 100,000 showers every position satisfies this, so all 44×16×9=6,33644× 16× 9=6,336 layer-pair positions are retained and none are discarded, and both Frobenius sums in Eq. (15) run over this full set. A.7. Layer-wise CFD Layer-wise CFD measures the global correlation structure between per-layer energies. For each pair of layers (i,j)(i,j) with i,j=1,…,Li,j=1,…,L, we compute the Pearson correlation across showers between the layer energies at i and j: (16) ρij=∑k(Ei,k−E¯i)(Ej,k−E¯j)∑k(Ei,k−E¯i)2∑k(Ej,k−E¯j)2, _ij= _k (E_i,k- E_i ) (E_j,k- E_j ) _k (E_i,k- E_i )^2\; _k (E_j,k- E_j )^2, where Ei,kE_i,k is the energy deposited in layer i by shower k, the index k runs over the K showers in the set, and E¯i=1K∑kEi,k E_i= 1K _kE_i,k is the mean layer energy across showers. Collecting ρij _ij over all layer pairs gives an L×L× L correlation matrix ρrefρ^ref for Geant4 and ρgenρ^gen for the generated samples. As in the voxel-wise case, we evaluate only pairs where both layers carry nonzero variance, since the correlation is undefined otherwise. Layer-wise CFD is the normalized Frobenius distance between the two matrices, restricted to the off-diagonal entries, (17) CFDlayer=‖ρgen−ρref‖F,off‖ρref‖F,off,‖ρ‖F,off=∑i≠j(ρij)2,CFD_layer= \|ρ^gen-ρ^ref \|_F,\,off \|ρ^ref \|_F,\,off, \|ρ \|_F,\,off= _i≠ j ( _ij )^2\,, where both sums run over i≠ji≠ j. The unit diagonal is excluded because it is identical in both matrices and would only dilute the normalized distance. Lower values indicate closer agreement with Geant4’s inter-layer correlation structure. The voxel-wise form (Eq. (15), main text) applies the same construction to the per-position correlations between consecutive layers. Appendix B Data Representation and Model Architecture B.1. Data Representation and Preprocessing A shower is stored as physical voxel energies x∈ℝ≥0Vx ^V_≥ 0 in MeV, and the networks act on the preprocessed representation x~=(x) x=T(x) defined in Section 3. We reuse the transform of CaloDREAM (Favaro et al., 2025), built from three invertible maps. Every voxel is first divided by its layer energy, so the shape network sees only the within-layer shares s∈[0,1]s∈[0,1] and the layer energies are handled separately by the ratio vector u. Each share then passes through a regularized logit, (18) s′=logsα1−sα,sα=(1−2α)s+α,α=10−6,s = s_α1-s_α, s_α=(1-2α)\,s+α, α=10^-6, which pushes the many near-zero voxels onto a usable range while staying one-to-one, and each feature is centered and scaled using fixed training-set statistics. Generation runs −1T^-1 and then rescales every layer’s voxels to the energy set by u, fixing the total and per-layer energies of each generated shower exactly. The ratio vector u∈ℝLu ^L expresses the layer energies E0,…,EL−1E_0,…,E_L-1 through EincE_inc, following the parameterization of CaloFlow (Krause and Shih, 2021) adopted by CaloDREAM (Favaro et al., 2025), (19) u0=∑iEifEinc,ui=Ei∑j≥iEj,i=1,…,L−1,u_0= _iE_if\,E_inc, u_i= E_i _j≥ iE_j, i=1,…,L-1, where the constant f=2.85f=2.85 keeps u0∈[0,1]u_0∈[0,1]. For i>0i>0, uiu_i is the share of the energy still undeposited at layer i that layer i receives; the map (Einc,u)↦(E0,…,EL−1)(E_inc,u) (E_0,…,E_L-1) therefore has a closed-form inverse and returns the layer energies without approximation. Table 4. Hyperparameters of the energy and shape networks on Dataset 2. Energy network Parameter DS2 Epochs 500 LR sched. one-cycle Max LR 2×10−32× 10^-3 Batch size 4000 Noise schedule Linear Solver DDPM ancestral sampler Conditioning network (Transformer) Dim embedding 64 Intermediate dim 512 Num heads 4 Encoder layers 4 Decoder layers 4 Normalization LayerNorm Denoising network (dense feed-forward) Intermediate dim 256 Num layers 8 Activation SiLU Training Optimizer AdamW Betas (0.9, 0.999) Timesteps 1000 Shape network Parameter DS2 Epochs 800 LR sched. CosineAnnealing LR / Max LR 5×10−45× 10^-4 / 1×10−31× 10^-3 Batch size 64 Noise schedule sigmoid Solver DDPM ancestral sampler Backbone network (ViT) Input shape [1, 45, 16, 9] Patch shape [3, 16, 1] Hidden dim 240 Depth (layers) 6 Num heads 6 MLP ratio 4.0 Condition dim 46 Attention Standard (non-causal) Prediction type noise Training Optimizer AdamW Betas (0.9, 0.999) Epsilon 1×10−61× 10^-6 Weight decay 1×10−51× 10^-5 Timesteps 1000 B.2. Diffusion Background Both networks are denoising diffusion probabilistic models (Ho et al., 2020). The forward process noises a clean sample z0z_0 into ztz_t as in Eq. (2) of the main text, with αt=1−βt _t=1- _t and α¯t=∏s≤tαs α_t= _s≤ t _s fixed by the variance schedule βtt=1T\ _t\_t=1^T. Sampling inverts this process one step at a time: a noise-prediction network ϵθ(zt,t,c) _θ(z_t,t,c), conditioned on the timestep t and context c (the incident energy, and for the shape network also u), drives the ancestral update (20) zt−1=1αt(zt−1−αt1−α¯tϵθ(zt,t,c))+σtξ,σt2=βt1−α¯t−11−α¯t,z_t-1= 1 _t\! (z_t- 1- _t 1- α_t\, _θ(z_t,t,c) )+ _t\,ξ, _t^2= _t\, 1- α_t-11- α_t, with ξ∼(0,I)ξ (0,I) for t>1t>1 and ξ=0ξ=0 at t=1t=1 for a deterministic final step. Both networks are trained with the standard noise-prediction objective (21) ℒdenoise=z0,ϵ,t[‖ϵ−ϵθ(zt,t,c)‖22].L_denoise=E_z_0,ε,t [\,\|ε- _θ(z_t,t,c)\|_2^2\, ]. B.3. Energy Network: Autoregressive Transformer The energy network predicts the layer-energy ratios u=(u0,…,uL−1)u=(u_0,…,u_L-1) autoregressively, reflecting that the energy reaching a layer is constrained by what has already been deposited upstream. Following the transfusion design of CaloDREAM (Favaro et al., 2025), it couples a transformer encoder for the incident-energy condition with a causally masked decoder over layers and a per-layer denoising head. Conditioning and layer inputs. The incident energy EincE_inc is embedded as a single condition token and so the encoder’s self-attention runs over one element and reduces to the identity. Each layer i is marked by a positional token pip_i that encodes its layer index, zero-padded to the token width; the decoder applies a causal mask, so layer i attends only to the tokens p0,…,pi−1\p_0,…,p_i-1\ that precede it. A cross-attention between the encoded condition and the decoder states produces a per-layer context cic_i, (22) ci=f(Einc)i=0,f(0,…,i−1,Einc)i>0,c_i= casesf(E_inc)&i=0,\\[2.0pt] f(p_0,…,p_i-1,E_inc)&i>0, cases that carries what is needed to predict the ratio at layer i. Per-layer denoising. A single denoising head ϵϕ _φ acts on the noised ratio ui(t)u_i(t) of each layer, conditioned on its context cic_i and the timestep t, and is applied across all L layers, (23) ϵ^(u(t),t,Einc)=(ϵϕ(u0(t),c0,t),…,ϵϕ(uL−1(t),cL−1,t)). ε (u(t),t,E_inc )= ( _φ(u_0(t),c_0,t),\,…,\, _φ(u_L-1(t),c_L-1,t) ). It is trained with the denoising objective of Eq. (21), taking z0=uz_0=u. Because the clean ratios are known for every layer, training evaluates all layers in parallel; sampling instead proceeds layer by layer, each predicted ratio conditioning the next. B.4. Shape Network: Conditional ViT The shape network is a three-dimensional vision transformer (Peebles and Xie, 2023), adapted to calorimeter showers as in CaloDREAM (Favaro et al., 2025), that denoises the preprocessed shower conditioned on (Einc,u)(E_inc,u). The calorimeter volume is split into non-overlapping voxel patches; each patch is linearly embedded, tagged with a learnable positional encoding, and processed by a stack of transformer blocks of multi-head self-attention and feed-forward layers. A final linear layer maps each patch back to voxel space as a noise estimate, and the patches are reassembled into the shower volume. It is trained with the objective of Eq. (21), taking z0=x~z_0= x; the input and patch dimensions are listed in Table 4. The conditioning inputs, the timestep t, the incident energy EincE_inc, and the ratio vector u, are embedded by small dense networks and summed into a single vector that modulates every block. The modulation is affine: the condition sets scale and shift parameters (a,b)(a,b) and a residual gate γ applied to both sublayers, (24) xh=x+γhgh(ahx+bh),xl=xh+γlgl(alxh+bl),x_h=x+ _h\,g_h(a_hx+b_h), x_l=x_h+ _l\,g_l(a_lx_h+b_l), where ghg_h and glg_l are the self-attention and feed-forward operations and (a,b,γ)(a,b,γ) depend on the conditioning vector (Peebles and Xie, 2023). Patch size trades spatial detail against attention cost; the value used for Dataset 2 is given in Table 4. B.5. Training and Hyperparameters As noted earlier, the Autoregressive Layer Network and the Shape Network are trained independently but coupled during generation. All experiments were conducted on a single NVIDIA Tesla V100-SXM2 GPU with 32 GB of memory, using CUDA 12.8 for both training and inference, with all models implemented in PyTorch. Table 4 reports the hyperparameters for both the energy and shape networks. Appendix C Auxiliary-Loss Combination Baselines C.1. Gradient-Combination Baselines The combination baselines isolate the rule that merges the denoising and auxiliary gradients. Architecture, auxiliary losses, warm-up/decay schedule, preprocessing, and the AdamW optimizer are all held fixed at the Base configuration; only the combination rule changes, so any difference between these rows reflects the rule alone. At every optimizer step we backpropagate the denoising loss and each active physics loss separately to full-parameter gradient vectors. We index the objectives by i=1,…,mi=1,…,m, writing ℒiL_i for the i-th loss and gi=∇θℒig_i= _θL_i for its gradient over the shape-network parameters θ; the denoising objective is fixed at i=1i=1, so g1=gdg_1=g_d and the remaining gig_i are the physics-loss gradients (gag_a in the two-objective case). Each rule below maps the set gi\g_i\ to a single update that is written into the parameters before the AdamW step. Unlike the weighted-sum baseline, these methods combine the per-objective gradients directly, so the fixed per-loss weights λi _i play no role (GradNorm instead learns its own weights, described below). If any per-objective gradient is non-finite, we fall back to the pure denoising gradient gdg_d for that step. In our experiments two objectives participate: the denoising loss and one physics loss (voxel or Laplacian). ConFIG. We apply the conflict-free update of (Liu et al., 2025) through the authors’ conflictfree library. Each step backpropagates the denoising loss and the active physics loss separately to full-parameter gradients gdg_d and gag_a, which are passed to ConFIG_update (pseudo-inverse solver, rtol=10−4 rtol=10^-4, default equal-weight and projection-length models). The returned direction has equal, positive projection onto both gradients—the unit bisector of gdg_d and gag_a in our two-objective case—and is written into the parameters before the usual AdamW step. If either gradient is non-finite we fall back to the pure denoising gradient. Architecture, auxiliary loss, schedule, preprocessing, and optimizer are identical to Base, so this configuration isolates the combination rule. PCGrad. We apply gradient surgery (Yu et al., 2020) through the conflictfree library. Given the per-objective gradients, any pair with a negative inner product is deconflicted by projecting one gradient onto the normal plane of the other; the deconflicted gradients are then summed and written into the parameters. We log the number of conflicting pairs resolved per step. As for the other rules, a non-finite gradient triggers the pure-denoising fallback. IMTL-G. We use the impartial multi-task gradient rule (Liu et al., 2021) through the IMTLGOperator of the conflictfree library. It selects the combined direction whose projection onto every objective’s unit gradient u^i=gi/∥gi∥ u_i=g_i/ g_i is equal, so no single objective dominates the shared update. We verify this impartiality online by logging the projection u^i⋅gc u_i· g_c of the combined update gcg_c onto each objective, which the rule drives to a common value. Non-finite gradients trigger the pure-denoising fallback. GradNorm. Unlike the three conflict-based rules, GradNorm (Chen et al., 2018) rebalances gradient magnitudes rather than resolving directional conflict. Each objective receives a weight wiw_i formed as a softmax over learned logits, normalized to sum to the number of objectives and floored at 0.50.5, and the applied update is the weighted sum ∑iwigi _iw_i\,g_i with the weights detached—equivalent to backpropagating ∑iwiℒi _iw_i\,L_i. The logits are trained (Adam, learning rate 10−310^-3) so that each objective’s gradient norm is pulled toward the batch-mean norm rescaled by its relative inverse training rate raised to the power α=1.5α=1.5; the training rate is measured against a per-objective baseline L0,iL_0,i (the mean of ℒiL_i over the first 100100 batches), during which an equal-weight sum is used. The gradient norms driving this weight update are measured over the full set of shape-network parameters, rather than at a single last shared layer as in the original formulation. GradNorm therefore behaves like the weighted-sum baseline but with weights adapted online in place of the fixed λi _i. A non-finite gradient triggers the pure-denoising fallback. C.2. Loss-weighting baselines These methods reduce the objectives to one scalar loss and backpropagate it once. They were our primary investigation before we turned to gradient-level combination: because a scalar weight sets only the magnitude of each loss, they cannot reorient the physics gradient away from a direction that conflicts with denoising—the limitation that motivated the gradient methods above. Weighted sum. The simplest baseline forms a fixed linear combination ℒ=∑iλiℒiL= _i _iL_i and backpropagates it once. We set λ1=1 _1=1 for the denoising loss and a small fixed weight for the physics loss (λlap=10−2 _lap=10^-2 for the Laplacian loss and λvox=10−4 _vox=10^-4 for the voxel loss). These weights are constant throughout training and, unlike the other methods, are not gated by the warm-up/decay schedule. A single scalar per loss fixes both the magnitude and the direction of the auxiliary contribution, so the physics gradient is applied in full even where it opposes denoising. Uncertainty weighting. Following the homoscedastic formulation of Kendall et al. (Kendall et al., 2018), we replace the fixed weights with learned per-loss uncertainties. With a learnable log-variance si=logσi2s_i= _i^2 per objective, each loss enters as 12e−siℒi+12si 12e^-s_iL_i+ 12s_i: the precision e−sie^-s_i scales the loss while the 12si 12s_i term regularizes the variance. The sis_i are initialized to zero and trained jointly with the network by the same AdamW optimizer. We evaluate two variants. The pure form applies the learned weighting throughout training, as in the original method; the scheduled form additionally places the auxiliary term under the same warm-up/decay envelope used by the other methods, so the schedule gates when the physics loss participates while its relative weight remains learned. In both, uncertainty weighting still reduces the two objectives to a single scalar loss, so it rescales the physics gradient but cannot reorient it. Appendix D Additional Results D.1. PRDC Table 5. Precision, density, recall, and coverage (PRDC) on CaloChallenge Dataset 2, computed with k=5k=5 nearest neighbors. Prior-surrogate values are taken from the CaloChallenge report (Krause et al., 2025). Trained models report mean ± std over three seeds. Model Loss Precision ↑ Density↑ Recall ↑ Coverage ↑ CaloDiffusion (Amram and Pedro, 2023) N/A 0.2390.239 1.2361.236 0.2350.235 0.9330.933 CaloScore v2 (Mikuni and Nachman, 2024) N/A 0.2280.228 1.0131.013 0.2280.228 0.9330.933 CaloDiT (Raikwar et al., 2024) N/A 0.5000.500 10.2310.23 0.0600.060 0.9240.924 CaloDREAM (Favaro et al., 2025) N/A 0.2530.253 1.1461.146 0.2200.220 0.9760.976 Base None 0.264±0.0030.264± 0.003 1.330±0.0401.330± 0.040 0.210±0.0020.210± 0.002 0.984±0.0020.984± 0.002 Base+PCGrad Voxel 0.273±0.0130.273± 0.013 1.592±0.3391.592± 0.339 0.199±0.0160.199± 0.016 0.986±0.0040.986± 0.004 Base+GradNorm Voxel 0.280±0.0010.280± 0.001 1.874±0.0731.874± 0.073 0.178±0.0020.178± 0.002 0.994±0.0010.994± 0.001 Base+IMTL-G Voxel 0.326±0.0150.326± 0.015 5.625±0.7415.625± 0.741 0.178±0.0210.178± 0.021 0.891±0.0320.891± 0.032 Base+ConFIG Voxel 0.304±0.0010.304± 0.001 2.674±0.1742.674± 0.174 0.155±0.0100.155± 0.010 0.994±0.0010.994± 0.001 Lantern Voxel 0.256±0.0010.256± 0.001 1.309±0.0521.309± 0.052 0.216±0.0040.216± 0.004 0.981±0.0020.981± 0.002 Base+PCGrad Laplacian 0.293±0.0170.293± 0.017 2.230±0.5482.230± 0.548 0.177±0.0190.177± 0.019 0.991±0.0040.991± 0.004 Base+GradNorm Laplacian 0.293±0.0090.293± 0.009 2.326±0.3452.326± 0.345 0.162±0.0100.162± 0.010 0.996±0.0010.996± 0.001 Base+IMTL-G Laplacian 0.306±0.1770.306± 0.177 167.8±179.4167.8± 179.4 0.047±0.0650.047± 0.065 0.989±0.0160.989± 0.016 Base+ConFIG Laplacian 0.302±0.0050.302± 0.005 2.798±0.0502.798± 0.050 0.154±0.0130.154± 0.013 0.994±0.0030.994± 0.003 Lantern Laplacian 0.259±0.0030.259± 0.003 1.287±0.0361.287± 0.036 0.213±0.0040.213± 0.004 0.983±0.0020.983± 0.002 Table 5 reports the full PRDC breakdown behind the mode-collapse finding of Section 5.3, resolving the aggregate distances FPD, KPD, and CFD into separate fidelity and diversity components. Precision and density measure fidelity; recall and coverage measure diversity (Kynkänniemi et al., 2019; Naeem et al., 2020). Density is normalized so that agreement with the reference yields values near one, and values well above one indicate over-concentration in dense regions of the reference distribution rather than higher fidelity. Base and Lantern keep these balanced, with density near 1.31.3 and recall near 0.210.21, close to CaloDREAM. The general-purpose multi-task methods instead raise density and precision while lowering recall, the signature of mode collapse: samples concentrate in high-density regions and stop covering the full reference distribution. Base+IMTL-G is the most severe. Under the Laplacian loss its density reaches 167.8167.8 while recall falls to 0.0470.047, so its samples pile onto a small region and miss most of the reference. Base+ConFIG and Base+GradNorm show the same trend more mildly, with density above 22 and recall below 0.180.18. The high precision of these methods is therefore a symptom of collapse, not a sign of quality. Lantern shows no such collapse: it keeps the highest recall among the trained models, tied with Base, and a density close to that of the reference-quality prior surrogates. Among the prior surrogates CaloDiT shows the same collapse, with density 10.210.2 and recall 0.0600.060, consistent with its poor FPD. Schedule ablation. Complementing the cross-method comparison above, Table 6 reports PRDC for the three auxiliary-loss schedules of Section 5.4. The voxel loss shows the collapse signature (high density, low recall) under S1 and S2 that the terminal gate of S3 removes, while the Laplacian loss stays near density 1.31.3 and recall 0.210.21 across all three schedules. Table 6. Precision, density, recall, and coverage under the auxiliary loss schedules of Table 2. Auxiliary Sched. Precision Density Recall Coverage Voxel S1 0.5200.520 6.4716.471 0.0810.081 0.9920.992 S2 0.5860.586 10.09010.090 0.0600.060 0.9870.987 S3 0.2560.256 1.3091.309 0.2160.216 0.9810.981 Laplacian S1 0.2660.266 1.3421.342 0.2100.210 0.9830.983 S2 0.2560.256 1.2551.255 0.2170.217 0.9820.982 S3 0.2590.259 1.2871.287 0.2130.213 0.9830.983 D.2. Results for Weighted Sum and Uncertainty Weighting Table 7 separates the three combination rules cleanly. Scheduled uncertainty weighting attains the lowest FPD, KPD, and CFD in both representations and drives the classifier AUCs toward 0.50.5, the point at which generated and reference showers become hard to separate. In voxel space it reaches an FPD of 115.02115.02 and a KPD of 0.3840.384, more than an order of magnitude below weighted-sum (2854.242854.24, 27.12727.127) and close to two orders below pure uncertainty weighting (8372.798372.79, 85.59285.592). The classifier tests agree: the scheduled model sits at AUC 0.6170.617 and 0.5710.571 on the high- and low-level features, whereas the unscheduled variants stay almost perfectly separable above 0.970.97. The Laplacian representation narrows this gap by improving the uncertainty-weighted models rather than weighted-sum. Pure UW falls from an FPD of 8372.798372.79 in voxel space to 237.95237.95, and its high-level AUC drops from 0.9980.998 to 0.6280.628, while weighted-sum remains poor (FPD 1518.221518.22, AUC 0.9760.976). Scheduled uncertainty weighting again gives the strongest distributional scores (FPD 148.77148.77, KPD 0.5410.541, CFD 0.1480.148), and the two uncertainty variants become comparable on the classifier tests. Two observations follow. The scheduled gate is what makes uncertainty weighting usable, since without it pure UW is the worst model in voxel space on every column. CFD follows the same ordering as FPD and KPD but on a compressed scale, and it is the only metric on which the scheduled model improves over the baselines in both representations, which supports its role as a correlation-sensitive score that FPD and KPD do not capture. The low classifier AUCs for the scheduled model should be read alongside its lower PRDC precision in Table 8: it trades per-sample sharpness for a distribution close enough to Geant4 to defeat the classifier. Table 7. Model evaluation metrics. Base+WSum denotes weighted-sum loss balancing, Base+pure UW uncertainty weighting without scheduling, and Base+sched UW scheduled uncertainty weighting. CFD is computed in voxel space. AUC (HL) and AUC (L) are the high- and low-level classifier two-sample tests, where values closer to 0.50.5 indicate samples less distinguishable from data. Model Loss FPD (×103× 10^3) ↓ KPD (×103× 10^3) ↓ CFD ↓ AUC (HL) AUC (L) Base+WSum Voxel 2854.24±14.312854.24± 14.31 27.127±1.53127.127± 1.531 0.2730.273 0.9910.991 0.9790.979 Base+pure UW Voxel 8372.79±21.148372.79± 21.14 85.592±2.83485.592± 2.834 0.3510.351 0.9980.998 0.9910.991 Base+sched UW Voxel 115.02±2.79115.02± 2.79 0.384±0.1670.384± 0.167 0.1450.145 0.6170.617 0.5710.571 Base+WSum Laplacian 1518.22±11.291518.22± 11.29 12.571±0.91312.571± 0.913 0.2190.219 0.9760.976 0.9550.955 Base+pure UW Laplacian 237.95±4.70237.95± 4.70 1.096±0.2911.096± 0.291 0.1810.181 0.6280.628 0.6390.639 Base+sched UW Laplacian 148.77±3.35148.77± 3.35 0.541±0.1780.541± 0.178 0.1480.148 0.6450.645 0.5830.583 Table 8. PRDC Evaluation Metrics Model Loss Precision ↑ Recall ↑ Density ↑ Coverage ↑ Base+WSum Voxel 0.9330.933 0.0110.011 317.0317.0 0.9300.930 Base+pure UW Voxel 0.7320.732 0.2180.218 437.5437.5 0.7450.745 Base+sched UW Voxel 0.2970.297 0.1630.163 2.322.32 0.9970.997 Base+WSum Laplacian 0.8650.865 0.0090.009 123.9123.9 0.9540.954 Base+pure UW Laplacian 0.2690.269 0.1590.159 3.103.10 0.9900.990 Base+sched UW Laplacian 0.2800.280 0.1790.179 2.122.12 0.9950.995 PRDC analysis. The density values far above unity in Table 8, such as 317317 and 437437 for the voxel weighted-sum and pure uncertainty models, do not indicate stronger fidelity; they show that generated samples pile into the densest region of the real manifold, where heavily overlapping neighborhood balls inflate the per-sample count. The accompanying near-zero recall (0.0110.011 for weighted-sum voxel) confirms mode concentration, since the samples are realistic yet enclose almost none of the real distribution. Pure uncertainty weighting raises both precision and density by driving mass into these high-density regions, but its coverage drops to 0.7450.745 because the sparse tail modes receive few generated neighbors. The scheduled variant instead restores a density near 22 and coverage above 0.990.99 at the cost of lower precision, a deliberate exchange of fidelity for broader mode coverage. This behavior is sharpest under the voxel representation: the Laplacian loss keeps the uncertainty models near a density of 22, whereas weighted-sum still collapses. D.3. Gradient Diagnostic Definitions This section defines the quantities in Table 3. All are computed from per-epoch gradient statistics logged during training: the norms of the denoising gradient gdg_d, the auxiliary gradient gag_a, and the applied blended step gblendg_blend, together with the cosines of gblendg_blend with each objective and the conflict angle between gdg_d and gag_a. Progress ratios. For a smooth loss L and a step Δθ θ, the first-order change in L is ⟨∇L,Δθ⟩ ∇ L, θ . We compare the applied step Δθa=−ηgblend _a=-η\,g_blend against a pure step on one objective Δθn=−ηgk _n=-η\,g_k by the ratio of their first-order terms and of their squared sizes, (25) R1k=‖gblend‖gk‖cos(gblend,gk),R2k=(‖gblend‖gk‖)2,R_1^k= \|g_blend\|\|g_k\|\, \! (g_blend,g_k ), R_2^k= ( \|g_blend\|\|g_k\| )^2, with k∈d,ak∈\d,a\ for the denoising and auxiliary objectives. The first-order ratio R1kR_1^k is the share of that objective’s progress the blended step retains. It folds in both the step size and the alignment, so it falls when the step points away from gkg_k. The second-order ratio R2kR_2^k is the squared step-size ratio alone. Balance is R1d/R1aR_1^d/R_1^a; a value above one means the step retains more of the denoising progress than of the auxiliary progress. Cosines and conflict angle. cosd _d and cosa _a are the cosines of the blended step gblendg_blend with gdg_d and gag_a. θconf _conf is the angle between gdg_d and gag_a before blending, in degrees. Active window. The auxiliary weight follows a schedule that ramps from zero, holds, and decays back toward zero. Outside the held region the auxiliary gradient is near zero by design, and the ratios in Eq. 25 diverge from the small denominator. We therefore restrict all statistics to the active window, defined per epoch as those epochs where the auxiliary weight exceeds a fixed fraction of its peak value. We use a floor fraction of 0.010.01, which includes the ramp. Spike filter. Within the active window the auxiliary gradient norm can transiently collapse, which inflates R1aR_1^a and R2aR_2^a through the small denominator. We drop active epochs where the auxiliary gradient norm falls below one quarter of its active-window median, and report ratios on the remaining epochs. The %spike column is the fraction of active epochs removed by this filter. Aggregation. The logged statistics are heavy-tailed, so we summarize each run by the median over its active window. We then report the mean and standard deviation of these per-run medians across three seeds. The energy network is held fixed; the three seeds vary the shape network.