Improved memory and speed for training geometry aware neural networks
Architecture-Aware Mixed Precision for Replica-Exchange Training of Geometry-Aware Neural Operators with Uncertainty Quantification
Machine Learning
Summary
Training smart computer programs that understand shapes and physical spaces usually takes a lot of computing power and memory. The authors propose a new method that uses different levels of number precision in different parts of the training to save resources without losing accuracy. This approach makes the training faster and uses less memory while still giving reliable results. They tested their method on tasks involving fluid flow and 3D structures and found it works better than standard mixed precision training methods.
What this means in practice
- •For computational engineers: Train large geometry-aware machine learning models for engineering simulations more efficiently with lower memory use and faster runtimes.
- •For cloud infrastructure teams: Optimize GPU resource allocation for training physics-informed models by selectively applying mixed precision to reduce costs and improve throughput.
Authors
Ujin Jeon, Michael Penwarden, Wenqian Chen, Guang Lin
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.