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

📄 arXiv: 2611.02427v1 📥 PDF

作者: Ujin Jeon, Michael Penwarden, Wenqian Chen, Guang Lin

分类: cs.LG

发布日期: 2026-11-02

💡 一句话要点

提出架构感知混合精度以提升几何感知神经算子的训练效率

🎯 匹配领域: 支柱八：物理动画 (Physics-based Animation)

关键词: 几何感知神经算子 混合精度训练 不确定性量化 流体动力学 结构力学 计算效率 机器学习

📋 核心要点

  1. 几何感知神经算子在复杂物理领域的学习中面临高计算和内存需求，现有方法难以平衡效率与准确性。
  2. 本文提出架构感知混合精度（AA-MP），通过根据复制角色和架构特定操作分配数值精度，提升训练效率。
  3. AA-MP在流体动力学和结构力学基准测试中，减少了46.9%的内存使用和25.8%的训练时间，同时保持了与全精度训练相当的准确性。

📝 摘要（中文）

几何感知神经算子在复杂物理领域的学习中面临巨大的计算和内存需求。本文提出了架构感知混合精度（AA-MP），这是一种根据复制角色和架构特定操作分配数值精度的复制交换训练策略。AA-MP在保留FP32用于细化的同时，选择性地使用BF16进行探索，从而在提高训练效率的同时保持预测准确性。在流体动力学和三维结构力学基准测试中，AA-MP在多个运行的准确性上与全精度训练相当，同时将热复制峰值GPU内存减少了46.9%，训练时间缩短了25.8%。此外，AA-MP在内存效率和计算成本上优于简单的自动混合精度（AMP），在更大模型和几何规模下内存优势更为明显。最后，AA-MP在不确定性校准的保形预测中实现了接近名义覆盖的结果，且预测区间比替代方案更窄。这些结果推动了几何感知科学机器学习的可扩展性，降低了计算成本而不妥协预测准确性或不确定性量化。

🔬 方法详解

问题定义：本文旨在解决几何感知神经算子在复杂物理领域学习时的高计算和内存需求，现有方法在效率与准确性之间难以取得平衡。

核心思路：提出架构感知混合精度（AA-MP），通过在细化阶段保留FP32精度，同时在探索阶段选择性使用BF16精度，以提高训练效率并保持预测准确性。

技术框架：AA-MP的整体架构包括复制交换训练策略，分为热复制和冷复制，分别负责探索和细化过程。热复制使用BF16进行快速探索，而冷复制则使用FP32进行精确细化。

关键创新：AA-MP的主要创新在于根据复制角色和架构特定操作动态分配数值精度，这与现有的静态混合精度方法（如AMP）有本质区别。

关键设计：在设计中，冷复制保留FP32以确保预测准确性，而在热复制中使用BF16以提高内存效率和计算速度。实验结果表明，冷复制中的FP32保留对准确性至关重要，而在热复制中增加FP32操作的收益相对其计算成本有限。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果显示，AA-MP在流体动力学和三维结构力学基准测试中，减少了46.9%的热复制峰值GPU内存和25.8%的训练时间，同时在多个运行中保持与全精度训练相当的准确性，显著优于传统的自动混合精度方法。

🎯 应用场景

该研究具有广泛的应用潜力，特别是在流体动力学、结构力学等领域，能够有效降低计算资源消耗，同时保持高预测准确性。这将推动科学计算和工程模拟的效率提升，促进更复杂物理现象的建模与分析。

📄 摘要（原文）

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.