Conservative & Aggressive NaNs Accelerate U-Nets for Neuroimaging

ICLR 2026 Conference Submission19586 Authors

Keywords: Pooling, Unpooling, Convolutions, Deep learning, Optimization, Neuroimaging, Convolutional Neural Networks, Numerical Analysis, Uncertainty Quantification

TL;DR: We introduce Conservative & Aggressive NaNs, techniques that enhance CNN efficiency in neuroimaging by skipping operations on numerically unstable voxels, reducing computational load by up to two-thirds while maintaining model accuracy.

Abstract:

Advancements in deep learning for neuroimaging have resulted in the development of increasingly complex models designed for a wide range of tasks. Despite significant improvements in hardware, enhancing inference and training times for these models remains crucial. Through an analysis of numerical uncertainty in convolutional neural networks (CNNs) inference, we found that a substantial amount of operations in these models are applied to values dominated by numerical noise, with little to no impact on the final output. As a result, up to two-thirds of the floating-point operations executed by some CNNs appear unnecessary.
To address this inefficiency, we introduce Conservative \& Aggressive NaNs ---novel variations of PyTorch's max pooling and unpooling operations. These techniques identify numerically unstable voxels and replace them with NaNs, allowing models to bypass operations on irrelevant data. We evaluated Conservative \& Aggressive NaNs on four models: the FastSurfer and FONDUE CNNs, widely used neuroimaging tools, the Xception CNN, an image classification model, and another CNN designed to classify the MNIST dataset. We observed speedups for data containing at least 50\% NaNs, and most notably, for data with more than two-thirds NaNs (as in many of our use cases), we observed an average speedup of $1.67\times$.
Conservative NaNs reduces the number of convolutions by an average of 30\% across all tested models and datasets, with no measurable degradation in performance.
In some model layers, it can skip up to 64.64\% of convolutions with no performance degradation. The more proactive Aggressive NaNs approach can skip up to 69.30\% convolutions for FastSurfer with no performance degradation, however, it sometimes leads to measurable performance degradation for FONDUE and MNIST. Overall, Conservative \& Aggressive NaNs provide substantial opportunities for runtime acceleration of inference in CNNs, which could potentially reduce the environmental impact of these models.

Primary Area: applications to neuroscience & cognitive science

Submission Number: 19586