cs.AISep 30, 2026

PTNO: Training Neural Operators with Noisy Monte Carlo Estimates for Particle Transport Problems

Authors: Yubo Cao, Xi Deng, Mengqi Xia, Vignesh Gopakumar, Ander Gray, Anima Anandkumar

Organizations: California Institute of Technology · Yale University · UK Atomic Energy Authority · LIX, CNRS, École polytechnique

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

Particle transport under multiple scattering is central to radiative transfer and plasma physics, yet high-fidelity Monte Carlo (MC) simulations must trace prohibitively many particles. Learning-based surrogates can amortize this cost, but typically train on expensive, well-converged MC solutions. We propose the Particle Transport Neural Operator (PTNO), a neural operator that learns particle transport surrogates directly from noisy, low-cost MC labels. Such labels pose two challenges: (1) high variance, which destabilizes standard supervised learning, and (2) a high dynamic range (HDR) spanning many orders of magnitude. For the first, we learn the solution operator from noisy labels of many configurations, amortizing MC cost and generalizing to unseen configurations. Because MC labels are unbiased, we show that the squared loss on them shares its minimizer with the loss on converged solutions, and our budget-allocation study over training scenes MM, MC samples per render NN, and independent renders per scene KK shows that many noisy scenes beat fewer converged ones. For the second, a nonlinear transform such as the logarithm biases noisy supervision. Instead, PTNO keeps labels in physical space and enforces positivity with a softplus output layer that represents small values effectively. We further train with a pointwise relative L2L_2 loss (PRelL2), the stop-gradient relative loss of HDR denoising and neural rendering, which normalizes each residual by the stop-gradient prediction instead of the noisy label. We demonstrate PTNO on neutron transport in fusion reactors and radiative transfer in participating media. On the two neutronics tasks, PTNO is 10410^4-105×10^5\times faster than converged MC on the same CPU and 10310^3-105×10^5\times cheaper than MC at matched accuracy; on the two radiative-transfer tasks, MC at matched accuracy costs 0.80.8-11×11\times as much as PTNO.

Figures & tables

Figure 1: PTNO learns from extremely high-variance MC labels and, at inference, predicts the well-converged solution. The Particle Transport Neural Operator is trained on MC supervision (left green box, columns 1–2) across varying configurations of the input aa . Bottom row shows the EU-DEMO 1/16 wedge neutron-transport task: columns 1–4 are analog OpenMC runs (no variance reduction) of one held-out configuration (the scene of Figure 2 ) with N=103,104,105,106N{=}10^{3},10^{4},10^{5},10^{6} histories, i.e., simulated neutron trajectories, on one 24-thread Xeon Platinum 8352Y CPU, labeled with their measured transport wall-clock (library initialization, 3535 – 4141 s, excluded). Column 5 is the converged reference of the same configuration, computed with weight windows (a variance reduction that splits or randomly terminates particles to keep their statistical weights within a target band in each region) at a cost of 5,2805{,}280 core-hours. PTNO keeps positivity and a wide dynamic range with a softplus output activation, and our PRelL2\mathrm{PRel}L_{2} loss keeps MC labels in physical space.
Arch.LossN=128N{=}128 ↑\uparrowN=64N{=}64 (vs. N=128N{=}128 ) ↑\uparrowN=4N{=}4 (vs. N=128N{=}128 ) ↑\uparrow
NOL2L_{2}0.660±0.0260.660\pm 0.0260.662±0.0220.662\pm 0.022 ( 1.00×1.00{\times} )0.652±0.0200.652\pm 0.020 ( 0.99×0.99{\times} )
NOlog MSE0.853±0.0030.853\pm 0.0030.812±0.0020.812\pm 0.002 ( 0.95×0.95{\times} )0.470±0.0010.470\pm 0.001 ( 0.55×0.55{\times} )
PTNOPRelL2\mathrm{PRel}L_{2} (ours)0.952±0.002\mathbf{0.952\pm 0.002}0.949±0.002\mathbf{0.949\pm 0.002} ( 1.00×\mathbf{1.00{\times}} )0.917±0.004\mathbf{0.917\pm 0.004} ( 0.96×\mathbf{0.96{\times}} )
Table 1: PTNO + PRelL2\mathrm{PRel}L_{2} loses less accuracy than log MSE as labels get noisier. Far-field radiance with M=106M{=}10^{6} training scenes and NN samples per pixel per label; all other settings shared. Entries are log⁡10\log_{10} SSIM on the 100100 held-out scenes against the converged reference, mean ±\pm std over five seeds (higher is better); parenthesized ratios compare each lower- NN column against N=128N{=}128 .
Figure 4: At matched wall-clock time, MC is still dominated by noise while the NO prediction is close to the converged reference. Each block shows five panels left-to-right: geometry, source distribution, low-budget MC, NO prediction, and converged-MC ground truth (GT), for the two held-out configurations closest to the median per-scene log⁡10\log_{10} SSIM of the seed-42 model of Table 2 . All timings are measured on one node with an NVIDIA H200 GPU and a 24-thread AMD EPYC 9554 CPU; the MC budget is the smallest whose wall-clock comes closest to one NO forward pass. Top—Spherical tokamak neutron flux: NO inference takes 3737 – 3838 ms on the GPU; OpenMC on 24 CPU threads needs 4949 – 5555 ms for N=37N{=}37 – 4343 histories, the smallest budgets it runs near that time. The GT pools 10910^{9} histories from 100100 independent 10710^{7} -history runs ( 6767 – 8181 h of summed wall-clock, each run on 44 CPU cores, that is 270270 – 320320 core-hours). Bottom—Spherical-slab far-field radiance LL : NO inference takes 3.63.6 ms and a Mitsuba render with one sample per pixel 4.84.8 ms on the same GPU; the GT accumulated 10710^{7} SPP in 3131 – 4040 min on one H200.
DatasetMethodlog⁡10\log_{10} rel. L2L_{2} ↓\downarrowlog⁡10\log_{10} SSIM ↑\uparrow
Far-field radianceNO + L2L_{2}0.4782±0.04080.4782\pm 0.04080.6524±0.01960.6524\pm 0.0196
NO + log MSE1.3994±0.00131.3994\pm 0.00130.4699±0.00070.4699\pm 0.0007
PTNO + PRelL2\mathrm{PRel}L_{2} (ours)0.0825±0.0022\mathbf{0.0825\pm 0.0022}0.9165±0.0040\mathbf{0.9165\pm 0.0040}
3D fluenceNO + L2L_{2}0.2824±0.00060.2824\pm 0.00060.7243±0.00040.7243\pm 0.0004
NO + log MSE0.0843±0.00450.0843\pm 0.00450.8977±0.00420.8977\pm 0.0042
PTNO + PRelL2\mathrm{PRel}L_{2} (ours)0.0338±0.0005\mathbf{0.0338\pm 0.0005}0.9518±0.0011\mathbf{0.9518\pm 0.0011}
Table 2: PTNO achieves superior performance across all four datasets. Each block compares loss formulations under matched architectures, data, and optimizer-update budgets; bolded values are the best per dataset. EU-DEMO replaces the linear- L2L_{2} loss with relative L2L_{2} because the flux peak sits near the Adam ϵ\epsilon scale. Mean ±\pm standard deviation over five seeds on held-out configurations against converged references ( 100100 far-field scenes, 100100 3D fluence scenes, 5050 spherical-tokamak configurations, 2424 EU-DEMO configurations). Far-field and EU-DEMO use final checkpoints; the 3D fluence and spherical-tokamak blocks are provisional historical results selected on their evaluation configurations. EU-DEMO is scored only on cells where the converged reference is measured ( 45.4%45.4\% of the mesh; Appendix ).
Appendix figures & tables26 assets

Supplementary material from the paper’s appendix.

Appendix

SymbolMeaningDefined in
Problem and Monte Carlo labels
aa , Pa\mathcal{P}_{a}input configuration (geometry, materials, source) and its distributionSecs. 3 , 4
U(a)\mathcal{U}(a)converged transport responsesection 3
ξ∈Ω\xi\in\Omega , ν\nu , kkparticle path, path space, path-space measure, number of scattering eventssection 3
ff , pppath contribution function and proposal distributionsection 3
U^N(a,ξ)\hat{\mathcal{U}}_{N}(a,\xi) , ζN(a,ξ)\zeta_{N}(a,\xi)NN -particle MC estimate from the paths ξ=(ξ1,…,ξN)\xi=(\xi_{1},\ldots,\xi_{N}) and its zero-mean noiseSecs. 3 , 4
Appendix
Table 3: Each symbol keeps one meaning throughout the paper. Symbols grouped by where they first appear.
ExperimentReported floor ϵ\epsilonLoss
Far-field radiance (Sec. 5.1 )10−610^{-6}PTNO ++ PRelL2\mathrm{PRel}L_{2}
3D fluence10−1010^{-10}PTNO ++ PRelL2\mathrm{PRel}L_{2}
Spherical-tokamak10−1010^{-10}PTNO ++ PRelL2\mathrm{PRel}L_{2}
EU-DEMO 1/16 wedge10−1710^{-17}PTNO ++ PRelL2\mathrm{PRel}L_{2}
Appendix
Table 4: Each dataset has its own log⁡10\log_{10} floor, and every experiment uses the softplus head with PRelL2\mathrm{PRel}L_{2} . ϵ\epsilon is the lower bound used when reporting log⁡10\log_{10} metrics; EU-DEMO numbers are post-hoc checkpoint evaluations with prediction and reference both floored at 10−1710^{-17} . All experiments use the softplus head Bθ(a)=softplus⁡(Gθ(a))\mathcal{B}_{\theta}(a)=\operatorname{softplus}(\mathcal{G}_{\theta}(a)) with no latent clamp.
HeadDenominatorlog⁡10\log_{10} rel. L2L_{2} ↓\downarrowlog⁡10\log_{10} SSIM ↑\uparrowMedian rel. L2L_{2} ↓\downarrowOutcome
softplus (PTNO)sg prediction0.050\mathbf{0.050}0.928\mathbf{0.928}1.031.03stable
identitysg prediction0.1880.1880.7550.7550.740.748.7%8.7\% negative voxels
softplusper-sample norm0.8010.8010.6140.6141.0001.000collapsed: softplus⁡(−∞)=0\operatorname{softplus}(-\infty)=0
identityper-sample norm0.2310.2310.6670.6670.066\mathbf{0.066}shadow region lost
softpluslive prediction0.1660.1660.8710.871515.8515.8output inflated (flux ∼15×{\sim}15\times )
identitylive prediction0.8010.8010.6140.61461946194collapsed: negative constant
Appendix
Table 5: Only the stop-gradient prediction denominator with the softplus head trains stably on 3D fluence. All runs share data, architecture, and schedule; only the output head and the residual denominator change. One seed per row; 100100 held-out scenes against the converged reference; the linear median is over scenes and includes the emitter voxel. sg denotes stop-gradient (detached).
HeadResiduallog⁡10\log_{10} SSIM ↑\uparrowlog⁡10\log_{10} rel. L2L_{2} ↓\downarrow
softplus (PTNO)pointwise sg ( PRelL2\mathrm{PRel}L_{2} )0.7923±0.0016\mathbf{0.7923\pm 0.0016}0.0839±0.0006\mathbf{0.0839\pm 0.0006}
identitypointwise sg ( PRelL2\mathrm{PRel}L_{2} )0.7078±0.01390.7078\pm 0.01390.0885±0.00470.0885\pm 0.0047
softplusper-sample rel. L2L_{2}0.4135±0.11810.4135\pm 0.11810.2227±0.02320.2227\pm 0.0232
identityper-sample rel. L2L_{2}0.1455±0.00660.1455\pm 0.00660.2310±0.00500.2310\pm 0.0050
Appendix
Table 6: On EU-DEMO, the pointwise stop-gradient residual matters most, and softplus adds a further gain. Five-seed output-head ×\times residual ablation on the EU-DEMO continuous-energy training setup. Entries are mean ±\pm sample standard deviation over five seeds from independent evaluation of final checkpoints on the 2424 held-out configurations at 140×73×146140{\times}73{\times}146 , scored only on cells where the converged reference is measured ( 45.4%45.4\% of the mesh). The softplus– PRelL2\mathrm{PRel}L_{2} and identity–relative- L2L_{2} rows share their runs with Table 2 .
DatasetPhysicsSolverGeometryOutput gridVaried inputs
Far-field radianceradiative transferMC labels; reference by Mitsuba 3 [ 1 ]spherical slab (fixed)40 ⁣× ⁣8040\!\times\!80 far-field radianceσt,c,Q\sigma_{t},c,Q fields, gg
3D fluenceradiative transferMC labels and reference, validated against OpenMC [ 2 ]cube [−1,1]3[-1,1]^{3}64364^{3} Cartesian fluenceσt,c\sigma_{t},c fields, gg , source field
Sph. tokamakneutron transportOpenMC [ 2 ]parametric (Paramak)64364^{3} Cartesian flux20 (geometry and source)
EU-DEMOneutron transportOpenMC [ 2 ]EU-DEMO 1/16 wedge (fixed)140 ⁣× ⁣73 ⁣× ⁣146140\!\times\!73\!\times\!146 flux8 (source only)
Appendix
Table 7: The four main-text datasets span two physical regimes and two output types. Each maps varying scene inputs to a transport solution on a fixed output grid; “Varied inputs” lists what changes between scenes, and held-out scenes are fresh draws from the same design.
Figure 5: PTNO recovers the converged far-field radiance from the scene inputs, while one training label is mostly noise. Held-out scene 2 in an equal-solid-angle layout (azimuth φ\varphi against cos⁡ϑ\cos\vartheta ) on a shared color scale. (a) The seed-42 PTNO + PRelL2\mathrm{PRel}L_{2} checkpoint from Table 2 , using the source, extinction, and albedo tables. (b) A single MC render at the training budget of N=4N{=}4 SPP, i.e. what one training label looks like. (c) The converged MC reference. Panel titles give the rel. L2L_{2} against (c); the cost comparison is in Appendix H .
Figure 6: In the 3D fluence dataset, contrast blocks cast shadows that span orders of magnitude around the point emitters. (a) A representative scene: two σt\mathbf{\sigma}_{t} blocks rendered as translucent boxes inside the cube (axes in normalized units) (purple / orange tint by σt\mathbf{\sigma}_{t} magnitude), with the three point sources as bright spheres. (b) Mid-cut slices ( z ⁣= ⁣32z\!=\!32 , i.e. z ⁣= ⁣0z\!=\!0 in scene units) of the same scene’s σt\mathbf{\sigma}_{t} , albedo, and reference fluence (log scale); only the high- σt\mathbf{\sigma}_{t} /high-albedo block intersects this plane; the dark rectangle in the fluence panel is its cast shadow. (c) Half-cut ( X ⁣≤ ⁣0X\!\leq\!0 visible) view of the same scene’s log-fluence iso-surfaces at {−2,−3,−4}\{-2,-3,-4\} , illustrating the multi-order-of-magnitude attenuation around the two emitters in this half.
SymbolRangeUnitDescription
Paramak radial build (geometry only)
Rcs,inR_{\mathrm{cs,in}}[20, 40][20,\,40]cmcenter-column shield inner radius
Rcs,outR_{\mathrm{cs,out}}[55, 95][55,\,95]cmcenter-column shield outer radius
tblkt_{\mathrm{blk}}[40, 110][40,\,110]cmblanket thickness
wdivw_{\mathrm{div}}[25, 65][25,\,65]cmdivertor width
D-shape descriptors (shared geometry ↔\leftrightarrow source)
Appendix
Table 8: Spherical-tokamak scenes vary 20 geometry and plasma-source parameters. Sampling box of the 20 scalar parameters, drawn from Latin-hypercube designs. The geometry block parameterizes the paramak radial build; D-shape descriptors (R0,ap,κ,δ)(R_{0},a_{\mathrm{p}},\kappa,\delta) are shared between the paramak surface and the plasma source ring; the profile blocks parameterize the H-mode ion-density and ion-temperature radial profiles.
Figure 8: Each spherical-tokamak scene couples a parametric radial build to a D-shaped plasma source, and the flux falls by orders of magnitude from the plasma through the shield and blanket. (a) Paramak radial build for one scene with the D-shaped plasma source outline, major/minor radius, elongation, and triangularity annotated. (b)–(e) Per-scene 64364^{3} mesh views: toroidally averaged source QQ and group-summed flux (log) on the poloidal plane, then Z=0Z{=}0 source and flux slices showing the annular plasma ring and inboard/outboard attenuation. (f) Half-cut 3D view: the toroidal source isosurface inside translucent shield/first-wall/blanket/vessel shells.
SymbolRangeUnitDescription
TT[14.00, 17.00][14.00,\,17.00]keVplasma ion temperature (D–T thermal broadening)
α\alpha[1.20, 1.80][1.20,\,1.80]—radial peaking factor of the source profile
R0R_{0}[715.04, 1,072.56][715.04,\,1{,}072.56]cmmajor radius of the plasma ring
apa_{\mathrm{p}}[230.64, 345.96][230.64,\,345.96]cmminor radius of the plasma ring
κ\kappa[1.32, 1.98][1.32,\,1.98]—plasma elongation
δ\delta[0.2664, 0.3996][0.2664,\,0.3996]—plasma triangularity
Appendix
Table 9: EU-DEMO scenes vary only the plasma source, through 8 shape parameters. Sampling box of these parameters, drawn from a Latin-hypercube space-filling design. Ranges are the upstream EU-DEMO PPS configuration; the geometry and materials are fixed.
Figure 9: The EU-DEMO reference flux falls by more than nine orders of magnitude over one poloidal slice, and the PTNO prediction follows it through the inboard blanket toward the divertor. One held-out validation configuration (the scene of Figure 2 ). (a) Rendering of the DEMO tokamak [ 21 ] ; the wireframe brackets the 1/16 wedge whose poloidal cross-section is plotted in the remaining panels. (b) The z ⁣≃ ⁣0z\!\simeq\!0 slice of the plasma source QQ on the 140×73×146140\times 73\times 146 training mesh. (c) The converged reference flux ϕ\phi on the physical y ⁣= ⁣0y\!=\!0 slice, shown over nine orders of magnitude below its peak. (d) Material classes on the same slice, binned by mass density as in the legend. (e) Iso-contours of the source over the materials; the D-shaped plasma sits inside the vacuum vessel. (f) The PTNO prediction ϕ^\hat{\phi} over three orders of magnitude, with dashed reference iso-contours at 10−110^{-1} , 10−210^{-2} , and 10−310^{-3} of the peak; the predicted field follows the reference through the inboard blanket and toward the divertor.
MetricArch.Loss0.00160.0016 – 0.0130.0130.0130.013 – 0.100.100.100.10 – 0.850.850.850.85 – 6.96.9All
% RMSE ↓\downarrowNOL2L_{2}64.0±6.964.0\pm 6.918.5±1.518.5\pm 1.58.2±0.2\mathbf{8.2\pm 0.2}5.6±0.1\mathbf{5.6\pm 0.1}15.3±1.3\mathbf{15.3\pm 1.3}
NOlog MSE99.9±0.099.9\pm 0.090.6±0.290.6\pm 0.258.2±0.258.2\pm 0.229.7±0.229.7\pm 0.257.6±0.257.6\pm 0.2
PTNOPRelL2\mathrm{PRel}L_{2} (ours)29.1±2.7\mathbf{29.1\pm 2.7}15.7±1.8\mathbf{15.7\pm 1.8}15.4±0.515.4\pm 0.514.4±0.614.4\pm 0.616.4±0.616.4\pm 0.6
log⁡10\log_{10} SSIM ↑\uparrowNOL2L_{2}0.241±0.0790.241\pm 0.0790.490±0.0460.490\pm 0.0460.699±0.0090.699\pm 0.0090.781±0.0090.781\pm 0.0090.652±0.0200.652\pm 0.020
NOlog MSE0.514±0.0010.514\pm 0.0010.509±0.0010.509\pm 0.0010.441±0.0010.441\pm 0.0010.471±0.0020.471\pm 0.0020.470±0.0010.470\pm 0.001
PTNOPRelL2\mathrm{PRel}L_{2} (ours)0.935±0.003\mathbf{0.935\pm 0.003}0.945±0.006\mathbf{0.945\pm 0.006}0.925±0.004\mathbf{0.925\pm 0.004}0.888±0.005\mathbf{0.888\pm 0.005}0.917±0.004\mathbf{0.917\pm 0.004}
Appendix
Table 10: PTNO + PRelL2\mathrm{PRel}L_{2} is the only loss that is accurate in every brightness bin. Far-field radiance, 100100 held-out scenes against the converged reference; columns stratify scenes by their mean reference radiance ( 88 , 1616 , 4242 , and 3232 scenes), and the last column covers all 100100 scenes (two lie outside the bin edges) and reproduces Table 2 . Mean ±\pm standard deviation over the five seeds of Table 2 ; bolded values are best per metric and bin.
Arch.Loss / target% RMSE ↓\downarrowlog⁡10\log_{10} rel. L2L_{2} ↓\downarrowPSNR ↑\uparrowSSIM ↑\uparrowlog⁡10\log_{10} SSIM ↑\uparrow
NOlog MSE56.231.35719.230.5610.474
NOLog-mean MSE ( K=10K{=}10 )25.440.48927.290.8480.787
NO2nd-moment MSE ( K=10K{=}10 )23.210.45628.180.8610.797
NOTaylor-fallback MSE ( 5+55{+}5 )28.390.51625.100.8290.735
NO, 10Gθ10^{\mathcal{G}_{\theta}} headPRelL2\mathrm{PRel}L_{2}17.170.08329.530.9420.914
Appendix
Table 11: Even with explicit Jensen-bias correction, log-target debiasers leave a substantial residual gap. All methods are trained at N=4N{=}4 on the far-field radiance dataset; explicit debiasing methods consume up to K=10K{=}10 renders for target estimation. One run per row, scored on the 100100 held-out scenes against the converged reference; % RMSE, PSNR, and SSIM are linear, the other two columns log⁡10\log_{10} .
ObjectiveBias Correctionrel. L2↓L_{2}\downarrowlog⁡10\log_{10} rel. L2↓L_{2}\downarrowlog⁡10\log_{10} SSIM ↑\uparrow
rel L2L_{2}–1.0000±0.00001.0000\pm 0.00005.0226±0.00005.0226\pm 0.00000.1347±0.00000.1347\pm 0.0000
rel L2L_{2}2nd-order approx. Taylor0.7230±0.00110.7230\pm 0.00110.6571±0.00090.6571\pm 0.00090.4393±0.00520.4393\pm 0.0052
rel L2L_{2}Taylor0.7188±0.00280.7188\pm 0.00280.6486±0.00280.6486\pm 0.00280.4413±0.00250.4413\pm 0.0025
PRelL2\mathrm{PRel}L_{2} (ours)–0.1304±0.0031\mathbf{0.1304\pm 0.0031}0.0624±0.0008\mathbf{0.0624\pm 0.0008}0.9454±0.0019\mathbf{0.9454\pm 0.0019}
Appendix
Table 12: Bias correction improves the label-denominator relative L2L_{2} loss but leaves it at least 5×5\times worse than PTNO. We compare four loss variants on the far-field radiance dataset at (M,N,K)=(106,4,10)(M,N,K)=(10^{6},4,10) and evaluate on the 100100 held-out scenes against the converged reference. The variants are: (1) the baseline relative ( L2L_{2} ) loss, which uses MC labels in the denominator; (2–3) relative ( L2L_{2} ) with bias-correction strategies described in Appendix G ; and (4) PTNO with our prediction-normalized pointwise relative loss. Bias correction improves over the biased baseline but remains at least 5×5\times worse than our method in relative L2L_{2} . Mean ±\pm standard deviation over three seeds.
LabelPixels exactly zeroPixels below η\etaMax pixel weight
One 4-SPP render0.3260.326 ( 0.1400.140 )0.3920.3929.6×10119.6\times 10^{11}
Mean of ten renders (40 SPP)0.1480.148 ( 0.0090.009 )0.2010.2018.0×10118.0\times 10^{11}
Appendix
Table 13: Weak labels have a point mass at zero, so the label-denominator weight 1/(Y+η)21/(Y+\eta)^{2} has no finite mean. 100100 training scenes of the far-field radiance corpus, η=10−6\eta=10^{-6} ; fractions are means over scenes (medians in parentheses), and the maximum pixel weight is each scene’s maximum averaged over scenes.
 MC cost // operator cost 
DatasetConverged MCMatch log⁡10\log_{10} SSIMMatch log⁡10\log_{10} rel. L2L_{2}Break-even queries
Far-field radiance1.8×1051.8\times 10^{5}11115.55.51212
3D fluence1.8×1041.8\times 10^{4}2.12.10.80.88787
Sph. tokamak9.9×1049.9\times 10^{4}6.4×104 ∗6.4\times 10^{4}\,{}^{\ast}9.4×1039.4\times 10^{3}197197
EU-DEMO1.2×1041.2\times 10^{4}>1.6×104 ∗>1.6\times 10^{4}\,{}^{\ast}1.6×1031.6\times 10^{3}1212
Appendix
Table 14: At matched accuracy, MC costs 10310^{3} – 105×10^{5}\times the operator on neutronics and 0.80.8 – 11×11\times on radiative transfer. Both methods run on the same device (one GPU for radiative transfer, one 24-thread CPU for neutronics). Matched-accuracy entries are per-scene medians of the MC cost needed to reach the operator’s error, divided by the operator’s cost; EU-DEMO MC uses FW-CADIS weight windows with their generation cost included. ∗ MC does not reach the operator’s log⁡10\log_{10} SSIM within the measured budget on most scenes; the spherical-tokamak entry extrapolates per-scene error curves by at most one order of magnitude, and the EU-DEMO entry is a lower bound.
DatasetNoisy-label setTrainingConverged / sceneConverged / noisyBreak-even
Far-field radiance3.793.79 GPU-h ( 10610^{6} )0.750.75 GPU-h0.3720.372 GPU-h9.8×1049.8\times 10^{4}1212
3D fluence1.971.97 GPU-h ( 3×1053{\times}10^{5} )7.987.98 GPU-h0.1150.115 GPU-h1.75×1041.75\times 10^{4}8787
Sph. tokamak70,30070{,}300 core-h ( 10610^{6} )1.811.81 GPU-h356.8356.8 core-h5.2×1035.2\times 10^{3}197197
EU-DEMO71,60771{,}607 core-h ( 10510^{5} )2.882.88 GPU-h6,0696{,}069 core-h8.5×1038.5\times 10^{3}1212
Appendix
Table 15: Offline cost is repaid after tens to a few hundred queries. Label costs were measured on H200 GPUs (radiative transfer), in 8-core tasks on Xeon Gold 6130 CPUs (spherical tokamak), and in 8-core tasks on Xeon Ice Lake and Skylake nodes (EU-DEMO); training costs are historical five-run means for the same recipes on one H200 GPU, retained for cost accounting rather than the runtimes of the replacement runs. The spherical-tokamak converged cost is stated in the 8-core configuration of its training labels; the 10910^{9} -history reference itself ran as 4-core tasks at 305.5305.5 core-hours per scene. Costs are measured from the scheduler accounting; the EU-DEMO converged cost is the mean over the 2424 reference scenes.
  Match log⁡10\log_{10} rel. L2L_{2}Match log⁡10\log_{10} SSIM
MC variantWindow / sceneTransport+ windowTransport+ window
Analog—>1.6×103>1.6\times 10^{3}—no crossing—
FW-CADIS10.210.2 core-h5.4×1025.4\times 10^{2}1.57×1031.57\times 10^{3}unmatched>1.6×104>1.6\times 10^{4}
FW-CADIS →\to MAGIC2,0092{,}009 core-h6.8×1026.8\times 10^{2}2.17×1052.17\times 10^{5}9.1×1059.1\times 10^{5}1.1×1061.1\times 10^{6}
Appendix
Table 16: EU-DEMO matched-accuracy factors depend on the weight window given to MC. Per-scene medians over 2424 scenes of MC cost divided by operator cost, both on one Xeon Platinum 8352Y. “Transport” counts transport only; “+ window” adds the window generation cost. >> marks a lower bound where MC stays short of the operator. All entries are scored only on cells where the reference is measured ( 45.4%45.4\% of the mesh). The MAGIC SSIM factors are extrapolated; the measured end-to-end lower bound is 2.6×1052.6\times 10^{5} .
DatasetVaried per sceneHeld fixed
Far-field radianceEnvironment map at infinity ( 11 – 44 Gaussian lobes and 11 – 44 angular boxes, intensities in [0.1,10][0.1,10] ); piecewise-constant angular σt∈[0.01,10]\mathbf{\sigma}_{t}\in[0.01,10] and c∈[0.01,0.99]c\in[0.01,0.99] ; g∼Unif(−0.99,0.99)g\sim\mathrm{Unif}(-0.99,0.99)Unit-ball geometry
3D fluence22 – 55 axis-aligned boxes with σt∈[0.01,50]\mathbf{\sigma}_{t}\in[0.01,50] and c∈[0.01,0.99]c\in[0.01,0.99] in a near-vacuum background; 11 – 33 point emitters with intensity in [1,20][1,20] , half with a Perlin angular profile; g∼Unif(−0.95,0.95)g\sim\mathrm{Unif}(-0.95,0.95)Cube [−1,1]3[-1,1]^{3} , vacuum boundary
Sph. tokamakAll 2020 axes of Table 8 : radial build, D-shape shared by geometry and plasma source, H-mode density and temperature profiles, Shafranov shift, radial source offsetMaterial compositions, tally mesh
EU-DEMOThe 88 plasma-source axes of Table 9Entire reactor geometry and materials
Appendix
Table 17: Each dataset varies coefficients, sources, or geometry within a fixed family ; held-out sets are fresh draws from the same design.
DatasetScenesSpanϵ\epsilonSpan above ϵ\epsilonExact zeros
Far-field radiance1001003.16±1.273.16\pm 1.2710−610^{-6}3.16±1.273.16\pm 1.270.0%0.0\%
3D fluence10010010.48±1.6910.48\pm 1.6910−1010^{-10}8.85±0.858.85\pm 0.852.3%2.3\%
Sph. tokamak505010.16±0.9910.16\pm 0.9910−1010^{-10}8.21±0.888.21\pm 0.8841.6%41.6\%
EU-DEMO242416.97±0.2616.97\pm 0.2610−1710^{-17}10.90±0.0410.90\pm 0.040.0%0.0\%
Appendix
Table 18: 3D fluence and spherical-tokamak references span about 1010 orders of magnitude, EU-DEMO about 1717 , and far-field radiance about 33 . Per scene, the span is log⁡10\log_{10} of the reference maximum over its smallest positive value; “above ϵ\epsilon ” replaces that smallest value by the dataset’s log⁡10\log_{10} floor ϵ\epsilon when it is lower, which is the span the log⁡10\log_{10} metrics see (their data range, Appendix ). Mean ±\pm standard deviation over the held-out evaluation scenes, on the field as scored: far-field radiance on the 40×8040{\times}80 grid, the spherical tokamak per energy group (over the 350350 scene–group pairs), and EU-DEMO only on cells where the reference is measured.
Mesh and inference pathrel. L2L_{2} ↓\downarrowlog⁡10\log_{10} rel. L2L_{2} ↓\downarrowlog⁡10\log_{10} SSIM ↑\uparrow
Far-field radiance
40×8040{\times}80 (training mesh)0.120\mathbf{0.120}0.056\mathbf{0.056}0.952\mathbf{0.952}
80×16080{\times}160 (native forward)0.1390.1390.0690.0690.9310.931
80×16080{\times}160 ( 40×8040{\times}80 then upsample)0.1650.1650.0940.0940.9010.901
EU-DEMO
70×37×7370{\times}37{\times}73 ( 0.5×0.5\times native)0.1203±0.00880.1203\pm 0.00880.0854±0.00070.0854\pm 0.00070.7895±0.00400.7895\pm 0.0040
Appendix
Table 19: Without retraining, accuracy holds on nearby meshes and degrades with larger resolution steps. Cross-resolution evaluation. The far-field block compares a native forward with block upsampling for one earlier single-seed checkpoint trained on 128128 -SPP labels, not comparable to Table 2 . EU-DEMO entries are mean ±\pm sample standard deviation over the five main-table seeds on the 2424 held-out configurations, scored only on cells where the reference is measured.
Training corpusTrain-grid rel. L2L_{2}Native-direct rel. L2L_{2}Active log MAELowres-up rel. L2L_{2}
Low mesh only0.01730.01731.00001.00005.365.360.16270.1627
Low ++ high meshes0.02070.02070.99980.99982.362.360.16170.1617
Appendix
Table 20: Adding a few high-resolution scenes reduces, but does not remove, the collapse of a direct forward pass at a tenfold grid step. Mixed-resolution diagnostic on an earlier EU-DEMO setup, one run per row. The baseline uses 50,00050{,}000 scenes at 70×55×8070{\times}55{\times}80 ; the mixed run adds 100100 scenes at 350×275×400350{\times}275{\times}400 . Native-direct metrics use three paired scenes, and low-resolution prediction followed by upsampling uses ten.
Figure 27
QuantityDistributionRangeNote
Number of objectscategoricalP(0,1,2,3)=(0.15,0.35,0.30,0.20)P(0,1,2,3)=(0.15,0.35,0.30,0.20)realized (0.152,0.351,0.298,0.199)(0.152,0.351,0.298,0.199)
Shapeuniformsphere 0.50.5 / box 0.50.5box faces are axis-aligned
Sphere radiusuniform[0.12,0.50] H[0.12,0.50]\,H 
Box half-extent (per axis)uniform[0.10,0.50] H[0.10,0.50]\,H 
Center (per axis)uniform[−0.6,0.6] H[-0.6,0.6]\,Hmust fit inside ±0.95H\pm 0.95H
Materialcategoricaldielectric 0.500.50 / mirror 0.250.25 / matte 0.250.25realized 0.499/0.250/0.2510.499/0.250/0.251 of 154,412154{,}412 objects
Appendix
Table 21: Each interface scene adds up to three dielectric, mirror, or matte spheres or boxes. The medium, emitters, and gg follow the 3D fluence design of Table 17 ; H=1H=1 is the half-width of the cube. Realized values are counted over the 10510^{5} training scenes; the shape and center rows give the proposal. nn spans water ( 1.331.33 ) to diamond ( 2.422.42 ).
Modellog⁡10\log_{10} SSIM ↑\uparrowlog⁡10\log_{10} rel. L2↓L_{2}\downarrowSelection-set log⁡10\log_{10} SSIM ↑\uparrowPeak ratio
PTNO, final ( 10510^{5} updates, width 6464 )0.9647\mathbf{0.9647} / 0.9652\mathbf{0.9652}0.0699\mathbf{0.0699} / 0.0682\mathbf{0.0682}0.96130.9613 / 0.96240.96241.001.00 / 1.041.04
Label ceiling (second MC solution)0.99680.99680.01130.0113 1.001.00
PTNO (softplus + PRelL2\mathrm{PRel}L_{2} ), 25,00025{,}000 updates0.95030.9503 / 0.94900.94900.09280.0928 / 0.09420.09420.94790.9479 / 0.94710.94711.471.47 / 1.361.36
NO, log⁡10\log_{10} output + MSE, 25,00025{,}000 updates0.93910.9391 / 0.93640.9364   
Appendix
Table 22: PTNO reaches within 0.0350.035 of the label ceiling on the interface dataset. log⁡10\log_{10} metrics, mean over the 224224 scenes of the stratified comparison set, for seeds 4242 / 4343 with end-of-training checkpoints. The selection-set column scores each model once on a separate 224224 -scene held-out set (the 25,00025{,}000 -update rows on an earlier, identically generated selection set). Peak ratio is the median over scenes of the predicted over the reference maximum. The label ceiling scores one independent Monte Carlo solution of the stratified comparison set against another. The last two rows share one input construction, corpus, and schedule at 25,00025{,}000 updates and differ only in the output head and loss.
 Stratified comparison setTest setLinear, comparison set
Output head + losslog⁡10\log_{10} SSIM ↑\uparrowlog⁡10\log_{10} rel. L2↓L_{2}\downarrowlog⁡10\log_{10} SSIM ↑\uparrowlog⁡10\log_{10} rel. L2↓L_{2}\downarrowCore rel. L2↓L_{2}\downarrowPeak ratio
PTNO (softplus + PRelL2\mathrm{PRel}L_{2} )0.9647\mathbf{0.9647} / 0.9652\mathbf{0.9652}0.0699\mathbf{0.0699} / 0.0682\mathbf{0.0682}0.9598\mathbf{0.9598} / 0.9611\mathbf{0.9611}0.0735\mathbf{0.0735} / 0.0720\mathbf{0.0720}1.1551.155 / 0.0860.0861.7991.799 / 1.6271.627
NO, log⁡10\log_{10} output + MSE0.94670.9467 / 0.94620.94620.12890.1289 / 0.12890.12890.94300.9430 / 0.94300.94300.12860.1286 / 0.12870.12870.0780.078 / 0.0840.0840.9850.985 / 1.0161.016
NO, identity head + linear L2L_{2}0.22510.2251 / 0.23540.23540.89290.8929 / 0.89570.89570.21010.2101 / 0.21980.21980.94580.9458 / 0.95020.95020.9640.964 / 0.9620.9620.0010.001 / 0.0010.001
Identity head + PRelL2\mathrm{PRel}L_{2}0.29190.2919 / 0.27260.27260.70560.7056 / 0.79520.79520.26530.2653 / 0.24540.24540.72030.7203 / 0.82710.82710.9500.950 / 0.9500.9500.0010.001 / 0.0010.001
Softplus + per-sample rel. L2L_{2}0.16200.1620 / 0.16200.16203.65523.6552 / 3.65523.65520.16580.1658 / 0.16580.16583.79433.7943 / 3.79433.79431.0001.000 / 1.0001.0000.0000.000 / 0.0000.000
Appendix
Table 23: At equal budget, the softplus head with PRelL2\mathrm{PRel}L_{2} outperforms plain losses on the interface dataset. Every row uses the final model’s construction, corpus, schedule, and 10510^{5} updates and changes only the output head and loss; seeds 4242 / 4343 , end-of-training checkpoints, means over 224224 scenes. The PTNO row rescores the final model of Table 22 with the evaluator used for the other rows (differences below 10−410^{-4} ). Core rel. L2L_{2} is the linear relative L2L_{2} with each emitter voxel and its 2626 neighbors removed; peak ratio is the mean over scenes of the predicted over the reference maximum.

Explore similar work

CardsList
Jun 16, 2026cs.LG

Operator Boosting Produces Pareto-Efficient PDE Surrogates

Neural operators are widely used as surrogate solution maps for partial differential equations (PDEs), but full-size models can be costly to store, deploy, and evaluate in many-query scientific workflows. This work introduces Operator Boosting, a stagewise residual-learning framework for constructing compact neural-operator surrogates directly, rather than training a large model and compressing it afterward. Starting from the empirical mean predictor in normalized output coordinates, the method trains a sequence of tiny same-family neural operators on residual fields and incorporates each correction through validation-selected shrinkage. We instantiate the framework with Fourier neural operators (FNOs), DeepONets, and convolutional neural operators (CNOs), and compare boosted tiny stacks against full-size monolithic baselines across one-, two-, and three-dimensional PDE benchmarks from PDEBench, APEBench, and The Well. Across 30 dataset-architecture pairs, 21 show positive mean accuracy gains and 17 have positive confidence intervals, while all boosted stacks reduce trainable parameter count by approximately 72-95%. Best-model comparisons show empirical Pareto improvements on 7 of 10 completed PDE benchmarks, including two-dimensional Navier-Stokes, shallow-water dynamics, Darcy flow, one-dimensional transport and reaction systems, and three-dimensional compressible Navier-Stokes. These results show that Operator Boosting often improves the empirical accuracy-parameter Pareto frontier of neural PDE surrogates, while also exposing PDE- and architecture-dependent regimes where residual boosting fails to offset compression.

Lennon J. Shikhman

Georgia Institute of Technology College of Computing Atlanta, Georgia, USA

Neural OperatorsNeural Partial Differential Equation SolversSurrogate ModelsBoosting
AskPDF
Aug 10, 2026cs.LG

MoNo: Multiscale Optimal Transport Neural Operator for Solving PDEs on General Geometries

Transformer-based neural operators have achieved substantial progress in solving Partial Differential Equations (PDEs) by projecting spatial observations into compact latent tokens and learning physical interactions in latent spaces. However, we reveal that existing learnable projection mechanisms cannot ensure stable and balanced assignments from observation points to latent tokens, causing some latent tokens to be over-assigned while others remain underutilized. This limitation further restricts the design of hierarchical architectures, as assignment imbalance is continuously inherited and amplified across latent spaces, eventually causing severe token collapse in deeper spaces. To address these issues, we propose MoNo (Multiscale Optimal Transport Neural Operator), a progressive multiscale neural operator that efficiently solves PDEs on general geometries through stable latent-space construction. At its core is CoTAP (Cross-scale Optimal Transport Assignment and Projection), a novel latent-space construction method that formulates cross-space assignment between adjacent spaces as an entropy-regularized optimal transport problem, thereby constructing balanced bidirectional projections and stable latent spaces. CoTAP also ensures stable information transfer across multiple latent spaces, further enabling multiscale architectures on general geometries, which in turn support more efficient learning of long-range physical interactions. Extensive experiments demonstrate that MoNo outperforms existing state-of-the-art neural operators in both prediction performance and computational efficiency. Code is available at https://github.com/ZijiangY1116/MoNo.

Zijiang Yang, Xiaomeng Wu, Dongmei Fu

School of Automation and Electrical Engineering, University of Science and Technology Beijing · 2Beijing Engineering Research Center of Industrial Spectrum Imaging

Neural OperatorsNeural Partial Differential Equation SolversPartial Differential EquationsSpatio-Temporal Transformers
AskPDF
May 15, 2026cs.LG

Martingale Neural Operators: Learning Stochastic Marginals via Doob-Meyer Factorization

Neural operators excel as deterministic surrogates, but inevitably collapse to the conditional mean when applied to stochastic PDEs, discarding the variance and tail structure upon which uncertainty quantification depends. Recovering this structure typically requires Monte Carlo rollouts or grafted generative models, both of which surrender the one-shot efficiency and resolution invariance that define the operator paradigm. To resolve this, we draw on the Doob-Meyer theorem, which establishes that any semimartingale fundamentally decomposes into a predictable drift and an unpredictable, zero-mean martingale. Translating this theorem into an architectural prior, we introduce the Martingale Neural Operator (MNO). MNO maps an initial condition directly to the conditional mean and covariance of the terminal law, parameterized by a drift-like mean and a low-rank factor BφB_φ with Bφ⊤BφB_φ^\top B_φ positive semi-definite by construction. For our experiments, we use a Gaussian residual instantiation. Across 1D SPDEs, rough volatility, and 2D operator tasks, MNO reduces Wasserstein distance by up to 120×120\times on φ4φ^4 field theory and 68×68\times on stochastic Burgers, evaluating ∼3×\sim 3\times faster than a conditional diffusion baseline at matched wall-clock training budgets. On 2D tasks, MNO is comparable to FNO on zero-shot resolution transfer and turbulent flow, while quasi-deterministic systems such as Gray-Scott remain a failure mode.

Kai Hidajat

Department of Applied Mathematics University of Washington Seattle, WA, USA

Neural OperatorsStochastic Differential EquationsDiffusion Models
AskPDF