ReDistill: Residual Encoded Distillation for Peak Memory Reduction of CNNs

TL;DR

ReDistill introduces a residual encoded distillation technique that significantly reduces peak memory in CNNs and diffusion models while maintaining performance, enabling deployment on resource-constrained devices.

Contribution

The paper proposes a novel residual encoded distillation method within a teacher-student framework for peak memory reduction in CNNs and diffusion models.

Findings

Achieves 4x-5x peak memory reduction in CNNs with minimal accuracy loss.

Reduces peak memory by 4x in diffusion-based image generation models.

Outperforms existing distillation methods in memory efficiency and performance.

Abstract

The expansion of neural network sizes and the enhanced resolution of modern image sensors result in heightened memory and power demands to process modern computer vision models. In order to deploy these models in extremely resource-constrained edge devices, it is crucial to reduce their peak memory, which is the maximum memory consumed during the execution of a model. A naive approach to reducing peak memory is aggressive down-sampling of feature maps via pooling with large stride, which often results in unacceptable degradation in network performance. To mitigate this problem, we propose residual encoded distillation (ReDistill) for peak memory reduction in a teacher-student framework, in which a student network with less memory is derived from the teacher network using aggressive pooling. We apply our distillation method to multiple problems in computer vision, including image classification and diffusion-based image generation. For image classification, our method yields 4x-5x theoretical peak memory reduction with less degradation in accuracy for most CNN-based architectures. For diffusion-based image generation, our proposed distillation method yields a denoising network with 4x lower theoretical peak memory while maintaining decent diversity and fidelity for image generation. Experiments demonstrate our method's superior performance compared to other feature-based and response-based distillation methods when applied to the same student network. The code is available at https://github.com/mengtang-lab/ReDistill.