Architecture-Aware Mixed Precision for Replica-Exchange Training of Geometry-Aware Neural Operators with Uncertainty Quantification

Ujin Jeon; Michael Penwarden; Wenqian Chen; Guang Lin

cs.LGNov 2, 2026

Architecture-Aware Mixed Precision for Replica-Exchange Training of Geometry-Aware Neural Operators with Uncertainty Quantification

Authors: Ujin Jeon, Michael Penwarden, Wenqian Chen, Guang Lin

Organizations: Department of Electrical and Computer Engineering, Purdue University, West Lafayette, 47907, IN, USA · Sandia National Laboratories, 1515 Eubank Blvd SE, Albuquerque, 87123, NM, USA · Pacific Northwest National Laboratory, 3335 Innovation Blvd, Richland, 99354, WA, USA · Department of Mathematics and Mechanical Engineering, Purdue University, West Lafayette, 47907, IN, USA

Abstract

Geometry-aware neural operators face substantial computational and memory demands when learning across complex physical domains. We introduce Architecture-Aware Mixed Precision (AA-MP), a replica-exchange training strategy that allocates numerical precision according to replica roles and architecture-specific operations. By retaining FP32 for refinement while selectively employing BF16 for exploration, AA-MP preserves predictive accuracy while improving training efficiency. Across fluid dynamics and three-dimensional structural mechanics benchmarks, AA-MP achieves comparable accuracy to full-precision training over an ensemble of runs, while reducing hot-replica peak GPU memory by 46.9% and training time by up to 25.8%. It also outperforms naive Automatic Mixed Precision (AMP) in memory efficiency and computational cost, with increasing memory advantages at larger model and geometry scales. Ablation studies show that retaining FP32 in the cold-replica is important for predictive accuracy, while additional FP32 operations in the exploratory replica provide limited benefits relative to their computational cost. Finally, conformal prediction for uncertainty calibration shows that AA-MP achieves near-nominal coverage with narrower prediction intervals than the alternatives. These results advance the scalability of geometry-aware scientific machine learning by reducing computational costs without compromising predictive accuracy or uncertainty quantification.

Figures & tables

Figure 1: Overview of the proposed AA-MP ArGEnT–DeepONet training framework. (a) The DeepONet branch encodes non-geometric physical parameters, while the cross-attention ArGEnT trunk uses query features as queries and the geometry point cloud as keys and values. The exploratory hot replica uses a manually controlled BF16 live model, while the latent reduction, prediction-error and energy reductions, master parameters, and optimizer states are maintained in FP32. The cold replica executes the same network entirely in FP32. (b) Two replicas are evolved using preconditioned reSGLD at $T_{c}<T_{h}$ . The higher-temperature hot replica promotes broader exploration, whereas the FP32 cold replica performs low-temperature refinement and model evaluation. Replica-exchange proposals periodically exchange the parameter and associated optimizer states between the two replica slots.

Operation	Precision	Implementation role
Hot-replica live model parameters	BF16	Explicit reduced-precision model representation without autocast.
Branch/trunk embeddings and MLPs	BF16	Reduced-precision feature computation.
Q/K/V projections, K/V LayerNorm, and RoPE	BF16	Cross-attention and positional feature computation.
Kernel attention, output projection, and feedforward layers	BF16	Computationally intensive Transformer operations.
Branch–trunk elementwise product	BF16	DeepONet feature fusion before latent reduction.
Latent reduction and prediction	FP32	The BF16 branch–trunk product is promoted before summation.

Table 1: Precision structure of the proposed AA-MP configuration. The cold replica remains entirely in FP32. The table summarizes the manually controlled precision policy of the exploratory hot replica.

Parameter	Value
Training iterations	40,000
Case batch size	32
Query-point batch size	3,000
Geometry point-cloud size	3,000
Transformer embedding dimension	128
Initial learning rate	$3\times 10^{-4}$

Table 2: Default configuration for the primary replica-exchange experiments. Unless otherwise stated, the same settings are used for the jet-engine bracket and AirfRANS benchmarks.

Figure 2: Hot-replica computational efficiency for AirfRANS and the jet-engine bracket benchmarks. (a) Hot-replica peak allocated GPU memory with both replicas resident. (b) Hot-replica execution time per training step. Bars show the mean over 10 seeds, with each seed summarized by the median over the final 50% of logged measurement windows; error bars indicate sample standard deviation. Percentage annotations denote reductions relative to FP32/FP32.

Dataset	Configuration	Peak alloc (GB)	Hot sec/step
AirfRANS	FP32/FP32	$5.727\pm 0.000$	$0.0804\pm 0.0024$
	FP32/AMP	$3.356\pm 0.000$	$0.0638\pm 0.0016$
	FP32/AA-MP	$3.042\pm 0.000$	$0.0603\pm 0.0006$
Bracket	FP32/FP32	$5.709\pm 0.000$	$0.0861\pm 0.0018$
	FP32/AMP	$3.347\pm 0.000$	$0.0660\pm 0.0027$
	FP32/AA-MP	$3.034\pm 0.000$	$0.0639\pm 0.0007$

Table 3: Hot-replica resource measurements corresponding to Fig. 2 . Values are mean

\pm

sample standard deviation across 10 seeds after per-seed aggregation.

Figure 3: AirfRANS reSGLD per-field test relative

\ell_{2}

error. Curves show the mean over 10 independent seeds, while shaded regions indicate empirically calibrated pointwise bands based on across-seed variability.

Figure 4: Jet-engine bracket reSGLD cold-replica per-field test relative

\ell_{2}

error. Curves show the mean over 10 independent seeds, while shaded regions indicate empirically calibrated pointwise bands based on across-seed variability.

Dataset	Field	FP32/FP32	FP32/AA-MP	FP32/AMP	FP32/Naive-BF16
AirfRANS (reSGLD)	$u$	$0.1055\pm 0.0155$	$\mathbf{0.1022}\pm 0.0210$	$0.1145\pm 0.0158$	$0.1250\pm 0.0179$
	$v$	$0.0765\pm 0.0069$	$\mathbf{0.0736}\pm 0.0078$	$0.0795\pm 0.0080$	$0.0865\pm 0.0104$
	$p$	$\mathbf{0.0501}\pm 0.0011$	$0.0511\pm 0.0008$	$0.0518\pm 0.0014$	$0.0548\pm 0.0035$
	$\nu_{t}$	$\mathbf{0.1226}\pm 0.0054$	$0.1229\pm 0.0032$	$0.1244\pm 0.0043$	$0.1363\pm 0.0102$
Bracket (reSGLD)	$u_{x}$	$\mathbf{0.1543}\pm 0.0073$	$0.1546\pm 0.0069$	$0.1553\pm 0.0064$	$0.1548\pm 0.0075$
	$u_{y}$	$0.1954\pm 0.0055$	$\mathbf{0.1945}\pm 0.0043$	$0.1947\pm 0.0041$	$0.1959\pm 0.0060$

Table 4: Per-field test relative

\ell_{2}

errors across the evaluated precision strategies. For the AirfRANS and bracket reSGLD blocks, the cold replica remains in FP32 and the column header denotes the hot-replica precision policy. The bracket SGD block reports single-chain precision-policy controls in the corresponding column order; the cold/hot replica notation in the column headings applies only to the reSGLD results. Values are mean

\pm

sample standard deviation over 10 seeds. Naive-BF16 is included as a reduced-precision ablation. Bold indicates the lowest mean error for each field.

Figure 5: Illustrative AirfRANS solution fields obtained with AA-MP. For each predicted quantity, the columns show the prediction, reference solution, and pointwise percentage error normalized by the maximum absolute reference value for that field. Prediction and reference fields use identical color limits within each field.

Figure 6: Illustrative jet-engine bracket solution fields obtained with AA-MP. For each predicted quantity, the columns show the prediction, reference solution, and pointwise percentage error normalized by the maximum absolute reference value for that field. Prediction and reference fields use identical color limits within each field.

Precision policy	$u$	$v$	$p$	$\nu_{t}$
AA-MP	$0.0593\pm 0.0048$	$0.0684\pm 0.0046$	$0.0524\pm 0.0016$	$0.0988\pm 0.0032$
AA-MP + FP32 RoPE	$0.0582\pm 0.0062$	$\mathbf{0.0642}\pm 0.0056$	$\mathbf{0.0500}\pm 0.0018$	$\mathbf{0.0936}\pm 0.0026$
AA-MP + FP32 LayerNorm	$0.0596\pm 0.0072$	$0.0674\pm 0.0046$	$0.0512\pm 0.0012$	$0.0953\pm 0.0032$
AA-MP + FP32 input path	$0.0598\pm 0.0085$	$0.0668\pm 0.0064$	$0.0510\pm 0.0025$	$0.0941\pm 0.0032$
AA-MP + FP32 projections	$\mathbf{0.0581}\pm 0.0056$	$0.0669\pm 0.0055$	$0.0508\pm 0.0020$	$0.0959\pm 0.0030$
FP32/FP32 control	$0.0609\pm 0.0094$	$0.0681\pm 0.0089$	$0.0516\pm 0.0017$	$0.0952\pm 0.0029$

Table 5: Per-field predictive accuracy for the architecture-level FP32-protection ablation on AirfRANS. All configurations use the same 20,000-iteration training protocol and 10 paired seeds. Values are mean per-case test relative

\ell_{2}

error, reported as mean

\pm

sample standard deviation across seeds. The FP32-protection rows are cumulative in their displayed order: RoPE, K/V LayerNorm, the input-feature pathway, and Q/K/V projections. Each successive row retains the protections of the preceding row. Bold indicates the lowest mean error for each field.

Precision policy	Peak alloc. (GB)	Mem. red.	Hot sec/step	Time red.
AA-MP	3.051	46.8%	0.0580	25.5%
AA-MP + FP32 RoPE	3.785	34.0%	0.0655	15.8%
AA-MP + FP32 LayerNorm	4.153	27.6%	0.0685	12.0%
AA-MP + FP32 input path	4.151	27.6%	0.0690	11.4%
AA-MP + FP32 projections	4.357	24.0%	0.0728	6.5%
FP32/FP32 control	5.734	–	0.0778	–

Table 6: Hot-replica computational cost for the AirfRANS FP32-protection ablation. For each seed, resource measurements are summarized by the median over the final 50% of logged measurement windows, using the same aggregation protocol as Table 3 . Values are then averaged over 10 seeds. The FP32 protections are cumulative in the same order as Table 5 . Memory and step-time reductions are relative to the FP32/FP32 control.

Field	Method	Mean cov.	Mean full width	Interval score	5th pct.
$u$	FP32/FP32	96.24	14.570	16.880	91.73
	FP32/AMP	95.94	23.300	26.400	91.59
	FP32/AA-MP	94.92	3.328	4.593	82.36
	FP32/Naive-BF16	95.95	28.564	31.537	93.45
$v$	FP32/FP32	96.24	14.370	16.690	92.80
	FP32/AMP	96.16	28.600	30.030	89.25

Table 7: Post-training AirfRANS conformal-UQ results for four precision configurations at nominal 95% case-averaged node coverage (98 evaluation cases). Coverage and 5th-percentile per-case coverage are percentages; mean full prediction-interval width and interval score use de-normalized dataset output units. Smaller interval widths indicate sharper predictions, while lower interval scores indicate a more favorable combined sharpness–coverage trade-off. Bold indicates coverage closest to the nominal target, the smallest interval width and score, or the highest 5th-percentile coverage, as appropriate.

Figure 7: Empirical coverage of conformal prediction intervals across AirfRANS flow fields. Markers indicate the mean fraction of covered spatial nodes across 98 held-out CFD cases for each precision configuration. Horizontal error bars represent 95% bootstrap confidence intervals obtained by resampling evaluation cases. The dashed vertical line denotes the nominal 95% coverage target. All four configurations achieve mean coverage near the target, with FP32/AA-MP exhibiting coverage closest to 95% across the predicted fields.

Figure 8: Spatial conformal uncertainty quantification for FP32/AA-MP on AirfRANS evaluation case 911. Rows show x-velocity $u$ , y-velocity $v$ , reduced pressure $p$ , and turbulent viscosity $\nu_{t}$ . Columns show the reference solution, predictive mean, full conformal-calibrated prediction-interval width $W_{k}=2H_{k}$ , and spatial interval-coverage diagnostic. Gray indicates locations covered by the interval, red indicates reference values above its upper bound, and blue indicates values below its lower bound.

Training configuration	Purpose
FP32/FP32	Full-precision replica-exchange baseline used to quantify the effect of low precision.
AA-MP/AA-MP	Evaluates the effect of applying reduced precision to both the cold and hot chain.
AA-MP/FP32	Evaluates the effect of applying low precision to the cold chain while retaining full precision in the hot chain.
FP32/AA-MP	Proposed replica-dependent precision configuration, retaining full precision for cold-chain refinement while using low precision for hot-chain exploration.

Table 8: Ablation configurations used to evaluate replica-dependent precision allocation. All preconditioned reSGLD configurations use identical hyperparameters.

Figure 9: AirfRANS cold-replica test relative

\ell_{2}

error for four cold/hot precision allocations: FP32/FP32, AA-MP/FP32, AA-MP/AA-MP, and FP32/AA-MP. Curves show the mean over 10 independent seeds, while shaded regions indicate empirically calibrated pointwise bands based on across-seed variability. Configurations are denoted as cold-replica precision / hot-replica precision.

Figure 10: Jet-engine bracket cold-replica test relative

\ell_{2}

Figure 11: Geometry point-cloud-size resource scaling for FP32/FP32, FP32/AMP, FP32/Naive-BF16, and FP32/AA-MP. The geometry point-cloud size is varied from 5,000 to 25,000.

Figure 12: Embedding-dimension resource scaling for FP32/FP32, FP32/AMP, FP32/Naive-BF16, and FP32/AA-MP. The Transformer embedding dimension is varied from 128 to 512.

Appendix figures & tables3 assets

Supplementary material from the paper’s appendix.

Appendix

Figure A.1 : FP32/FP32: spatially resolved post-training conformal UQ for AirfRANS evaluation case 804. For $u$ , $v$ , $p$ , and $\nu_{t}$ (rows), the columns display truth, predictive mean, calibrated 95% half-width, and interval coverage. Blue denotes truth below the interval, gray denotes truth inside, and red denotes truth above.

Figure A.2 : FP32/AMP: spatially resolved post-training conformal UQ for the same AirfRANS evaluation case 804. Panel definitions and categorical coverage colors are identical to those of Fig. A.1 .

Figure A.3 : FP32/Naive-BF16: spatially resolved post-training conformal UQ for AirfRANS evaluation case 804. The panels show truth, predictive mean, calibrated 95% half-width, and categorical interval coverage. Aggregate calibration and interval-score results are reported in Table 7 .

Mixed-Precision Training

See all 8 papers

Bit-Accurate Modeling of GPU Matrix Multiply-Accumulate Units: Demystifying Numerical Discrepancy and Accuracy
Nov 14, 2025Peichen Xie, Shuotao Xu, Yang Wang +2GPU Acceleration
Geometry-Aware Post-Hoc Uncertainty Quantification in Operator Learning
Jun 16, 2026Oriol Vendrell-Gallart, Nima Negarandeh, Ramin BostanabadUncertainty Quantification PDE Surrogate Modeling
Department of Mechanical and Aerospace Engineering University of California, Irvine Irvine, CA 92617
PGD-NO: A Neural Operator with Precomputed Geometry Decomposition for 3D Million-scale Physics Simulations
Jul 9, 2026Weiheng Zhong, Jing Bi, Victor Oancea +1PDE Surrogate Modeling PDE Operator Learning
Department of Civil and Environmental Engineering, Grainger College of Engineering, University of Illinois Urbana-Champaign, Urbana, IL, US · Dassault Systemes Americas, Johnston, RI, US · Siebel School of Computing and Data Science, Grainger College of Engineering, University of Illinois Urbana-Champaign, Urbana, IL, US

Nov 14, 2025J/K move · Enter open · S save

Bit-Accurate Modeling of GPU Matrix Multiply-Accumulate Units: Demystifying Numerical Discrepancy and Accuracy

Peichen Xie, Shuotao Xu, Yang Wang, Fan Yang+1

GPU Acceleration

Open paper AskPDF arXiv

Modern AI accelerators rely on matrix multiply-accumulate units (MMAUs), such as NVIDIA Tensor Cores and AMD Matrix Cores, to accelerate deep neural network workloads. MMAUs expose only instruction-level or API-level interfaces of matrix multiply-accumulate (MMA) operations, while leaving internal floating-point arithmetic behavior undocumented. Consequently, MMAUs across vendors and architectural generations often produce numerical discrepancies for identical inputs, and sometimes exhibit reduced numerical accuracy that can cause training instability. Diagnosing and understanding the root causes of these effects is challenging without white-box models of their arithmetic behavior. This paper proposes closed-loop feature probing (CLFP), a generic and systematic framework for constructing bit-accurate arithmetic behavior models of MMA operations. Based on this framework, we analyze all MMA instructions on ten GPU architectures spanning NVIDIA Volta through RTX Blackwell and AMD CDNA1 through CDNA3, and derive the first bit-accurate arithmetic models for these MMAUs. Our models explain previously observed cross-platform numerical discrepancies and accuracy issues, enable white-box numerical error analysis, reveal four types of precision bottlenecks and one type of numerical asymmetry, and inform software workarounds as well as design suggestions for future MMAUs. This work is open-source at https://github.com/microsoft/MMA-Sim

Abstract

Mixed-Precision Training

Explore similar work

Bit-Accurate Modeling of GPU Matrix Multiply-Accumulate Units: Demystifying Numerical Discrepancy and Accuracy

Geometry-Aware Post-Hoc Uncertainty Quantification in Operator Learning

PGD-NO: A Neural Operator with Precomputed Geometry Decomposition for 3D Million-scale Physics Simulations