Conservative & Aggressive NaNs Accelerate U-Nets for Neuroimaging

TL;DR

This paper introduces Conservative & Aggressive NaNs, novel methods that identify and skip redundant computations in CNNs for neuroimaging, leading to significant speedups without performance loss.

Contribution

The paper presents two new variants of max pooling and unpooling that leverage numerical uncertainty to accelerate CNN inference in neuroimaging applications.

Findings

Up to 1.67x inference speedup with >50% NaNs in inputs.

Conservative NaNs reduce convolution operations by 30% on average.

Aggressive NaNs skip up to 69.3% of convolutions, with some impact on performance.

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.