Structured Model Pruning for Efficient Inference in Computational Pathology

TL;DR

This paper presents a pruning methodology for U-Net models in computational pathology, achieving at least 70% model size reduction with minimal performance loss, thereby enabling more efficient AI inference in healthcare applications.

Contribution

It introduces a novel pruning approach tailored for biomedical U-Net architectures, demonstrating significant compression with negligible accuracy impact.

Findings

Models can be compressed by at least 70%

Pruning maintains high segmentation and classification accuracy

Multiple heuristics evaluated for optimal pruning strategy

Abstract

Recent years have seen significant efforts to adopt Artificial Intelligence (AI) in healthcare for various use cases, from computer-aided diagnosis to ICU triage. However, the size of AI models has been rapidly growing due to scaling laws and the success of foundational models, which poses an increasing challenge to leverage advanced models in practical applications. It is thus imperative to develop efficient models, especially for deploying AI solutions under resource-constrains or with time sensitivity. One potential solution is to perform model compression, a set of techniques that remove less important model components or reduce parameter precision, to reduce model computation demand. In this work, we demonstrate that model pruning, as a model compression technique, can effectively reduce inference cost for computational and digital pathology based analysis with a negligible loss of analysis performance. To this end, we develop a methodology for pruning the widely used U-Net-style architectures in biomedical imaging, with which we evaluate multiple pruning heuristics on nuclei instance segmentation and classification, and empirically demonstrate that pruning can compress models by at least 70% with a negligible drop in performance.