Conservative & Aggressive NaNs Accelerate U-Nets for Neuroimaging

TL;DR

This work investigates computational inefficiency in CNN-based neuroimaging models due to numerical noise, revealing that a large fraction of embeddings behave as irrelevant numerical noise. It introduces Conservative NaNs and Aggressive NaNs to identify numerically unstable voxels via modified pooling/unpooling and a NaN-Convolution mechanism that skips operations when NaN density is high, all without altering model architectures. Across FastSurfer, FONDUE, MNIST, and Xception, the methods yield substantial speedups—up to around 2× in many MRI-focused models—with Conservative NaNs offering reliable accuracy preservation and Aggressive NaNs providing larger gains in homogeneous MRI data but with potential trade-offs in others. The results demonstrate that exploiting numerical uncertainty can meaningfully reduce redundant computation in CNNs for neuroimaging, and point toward future integration with sparse representations and hardware-aware execution for even greater efficiency and uncertainty quantification.

Abstract

Deep learning models for neuroimaging increasingly rely on large architectures, making efficiency a persistent concern despite advances in hardware. Through an analysis of numerical uncertainty of convolutional neural networks (CNNs), we observe that many operations are applied to values dominated by numerical noise and have negligible influence on model outputs. In some models, up to two-thirds of convolution operations appear redundant. We introduce Conservative & Aggressive NaNs, two novel variants of max pooling and unpooling that identify numerically unstable voxels and replace them with NaNs, allowing subsequent layers to skip computations on irrelevant data. Both methods are implemented within PyTorch and require no architectural changes. We evaluate these approaches on four CNN models spanning neuroimaging and image classification tasks. For inputs containing at least 50% NaNs, we observe consistent runtime improvements; for data with more than two-thirds NaNs )common in several neuroimaging settings) we achieve an average inference speedup of 1.67x. Conservative NaNs reduces convolution operations by an average of 30% across models and datasets, with no measurable performance degradation, and can skip up to 64.64% of convolutions in specific layers. Aggressive NaNs can skip up to 69.30% of convolutions but may occasionally affect performance. Overall, these methods demonstrate that numerical uncertainty can be exploited to reduce redundant computation and improve inference efficiency in CNNs.

Figures (16)

• Figure 1: Significant Digit Maps for FastSurfer Model Embeddings in Selected Model Layers for Numerically Unstable Data.
• Figure 2: Comparison of Max Pooling vs Conservative & Aggressive NaNs in the presence of numerical uncertainty. Green indicates numerically stable values, red unstable values.
• Figure 3: Average Speed Up Between Standard and NaN Convolutions (left plot). Ratio of skipped convolutions (right plot).
• Figure 4: Comparison of Standard and NaN Convolutions in bfloat16 and float32 for GPU
• Figure 5: Ratio of Skipped Convolutions Across FONDUE (top) and FastSurfer (bottom) for Architecture (left) and Axial Brain Slices (right). For Thresholds 1, 0.8, 0.5 and 0.4, 6.1 %, 10.95%, 26.59% and 31.95% of total convolutions were skipped respectively for FastSurfer. For FONDUE, 0 %, 26.85%, 33.97% and 33.77 % of total convolutions were skipped respectively.