ReDistill: Residual Encoded Distillation for Peak Memory Reduction of CNNs
Authors: Fang Chen
Gourav Datta
Mujahid Al Rafi
Hyeran Jeon
Meng Tang
Venue: Accepted by TMLR ()

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.

Published in Transactions on Machine Learning Research (04/2025)
ReDistill:
Residual Encoded Distillation for Peak Memory Reduction of CNNs
Fang Chen!
fchen20@hcmerred. edu
Gourav Datta²
gourav.datta www.aattttaaae.... edu
Mujahid Al Rafi
wwwff@eucmereed. edu
Hyeran Jeon
hjeon eda
Meng Tang!
mtang4 @wwmerced. edu
"Universily of California Merced
"Case Western Reserve University
Reviewed on OpenReview: https: // openreview. net/ net/ forum? id- akumIm@jNM
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-cosstrinee 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 aggressive down-sampplig 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 teacher-student framework, in which 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 accuracy for most CNN-based architectures. For diffusion-basod image generation, our proposed distillation method yields denoising network with 4x lower theoretical peak memory while maintaining decent diversity and fidelity for image generation. Experiments demonstrate 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 www.iiib..o/mnnggtnnggggggg /mengtang lasttel
1 Introduction
Convolutional neural networks (CNN) have demonstrated impressive capabilities across diverse computer vision tasks such image recognition (Simonyan & Zisserman 2014)), object detection (Redmon & Farhadi
(2018)). ) semantic segmentation (Long et al. (2015)). and image generation (Creswell et al. (2018)). However, the ever-growing network size and image resolution of modern imaging sensors pose significant challenges in deploying neural networks on standard edge devices with limited memory footprint. For example, a standard STM32H5 MCU provides only 640 KB of SRAM and 2 MB of Flash storage. These constraints make impractical to exccute off-tho-shelf deep learning models: ResNet-50 surpasses the storage limit by 44x, while MobileNetV2 exceeds the peak memory limit by 8x. Even the int8 quantized version of MobileNet V2 surpasses the memory limit by 2x, underscoring substantial disparity between desired and

---

Published in Transactions on Machine Learning Research (04/2025)
ImigeNei dataset
Teacher
Studeat (wis Dutiisses
Student (Curs)
76 ResNet50 ni Aggresivee Pool
Students
ResNe Teacher
74-
72
Accirac
Teaeller
Teacher
ReviewKD (C)PR22))
MULD (CVPR223
Aggressive Pool
Studenns
RED (ours)
Peak Semory (MB)
Figure 1: (a) Left: For ImageNet classification, our distillation method significantly reduces the theoretical peak memory of ResNet-based models while achieving accuracy better than existing distillation methods. (b) Right: For diffusion-based image generation, our distilled network with 4x lower theoretical peak memory generates images indistinguishable from the generated images of a teacher network.
available hardware capacity. Hence, it is very important to reduce the peak memory during inference for edge deployment. Note that our primary focus in this work is reducing peak memory usage, as there are existing solutions for addressing other metrics, such as parameter count and the number of operations when deploying CV models at the extreme edge. Similar to (Lin et al.] (2021); Chowdhery et al., (2019)), we estimate the theoretical peak memory by summing the size of the input & output allocation for each operation (e-g., convolution, non-linear activation, pooling). Through empirical measurements, we determined that peak memory usage is prodominantly influenced by the initial layers of convolutional neural networks (CNNs) that are characterized by large fealure maps. For U-shaped CNN architectures, however, the last few layers also significantly contribute to peak memory consumption.
A naive approach to reducing peak memory is aggressive downsampling via pooling with large kernel size and large stride, which often leads unacceptable degradation of network performance due to loss of information in small feature maps. Given teacher network with large peak memory, propose residual encoded distillation (ReDistill) to train student network with significantly lower peak memory. This student network can be considered a variant of the teacher network with aggressive pooling. We demonstrate the effectiveness of our methods for multiple problems including image classification and diffusion- based image synthesis. For image classification with ResNet-basced models shown in Fig. (a), our mothod reduces the thcoretical peak memory by 4-5x with a lower accuracy drop compared to existing distillation methods. For diffusion- image generation shown in Fig. (b), our distilled network generates images similar to original networks, yet the thcoretical peak memory is reduced by 4x on average.
Our ieDistill method outperforms existing response-based or leature-based distillation methods regarding the aceuracy-memory trade-off. Our method differs from these existing counterparts in four regards. Firstly, our distillation method is tailored for peak memory reduction. In contrast, existing distillation techniques focus on transferring knowledge from a high-capability teacher network with a large number of parameters to a student network with fewer parameters. Our student networks apply large kernel size and stride the initial pooling layers with the same number of parameters as the teacher networks while consuming significantly lower peak memory. Secondly, the student network, utilizing aggressive pooling, has fewer pooling layers and consequently fewer stages than the teacher network, resulting in mismatched features at different stages between the two. We add novel non-lincar mapping modules termed residual encoded distillation (RED) blocks between the teacher and sludent network during both training and inference. Thirdly, our proposed RED block lightweight and effective with additive residual learning and multiplicative gating mechanism. We optimize the tradc-off between peak memory and accuracy, while slightly increases the model size due to extra parameters. Lastly, we align teacher and student network features asynchronously at pooling layers with matching feature sizes, while previous approaches align features at different stages of the networks.
2

---

Published in Transactions on Machine Learning Research (04/2025)
Our key contributions are summarized below.
We propose ReDistill, a distillation framework tailored for reducing the peak memory of convolutional neural networks. Our method allows aggressive downsampling of feature maps via pooling layers with large stride for student network while incurring less accuracy drop, To the best of our knowledge, ReDistill is the first distillation method focused on peak memory reduction for efficient deep learning.
The core of our ReDistill framework is residual encoded distillation (RED) block to align features between bigh-peak-memory teacher networks and low-peak-memory student networks Our RED block is based on multiplicative gating mechanism and additive residual learning and is shown to be simple and effective for peak memory reduction with minimum computational overhead.
For image classification tasks, our distillation method outperforms state-of-the-an response-based feature-based distillation methods when applied to the same student network assigned with a large pooling stride, as shown in extensive experiments with multiple datasets. Our method yields 4x ~ 5x reduction in theoretical peak memory with a slight. decrease in classification accuracies for CNN based models.
We also show the versatility of our distillation method for denoising diffusion probabilistic models for image generation. For a U-Net based denoising network, our method reduces the theoretical peak memory by 4x by downsampling the feature maps of the first few encoder layers and last few decoder layers while maintaining the fidelity and diversity of synthesized images.
2 Related Work
Memory-constraineed deep learning Limited memory capacity in GPU cards and edge platforms has been critical hurdle in CNN training and inference. Multiple GPUs can be utilized through model and data parallelism (Langer et al. (2020)) to mitigate the memory bottleneck. Other solutions include optimization methods such as network quantization (Hubara et al. (2016)), compression (Han et al. (2016b)), ) and pruning [Molchanvv et al. (2017) Lu et al. (2024)), which focus on maintaining essential bits of weights or parameters while minimizing accuracy loss. To produce correct outputs with compressed data, these solutions are typically designed with specialized accelerators to accommodate meta-data processing (Han ee al. (2) There are also CNNs specifically designed for resource-consllaineee applications, such variants of MobileNet (Howard et al. (2019t) and SqueezeNet. (Tandola et al. (2017)). These approaches to cemmory-constrineee deep learning orthogonal to our ReDistill framework, which focuses peak memory reduction during inference. Nevertheless, recent work has explored neural architecture search (NAS) to create networks with minimized peak memory (Lin et al. (2020)). ). Reference (T.nn et al. (2021)) takes this a step further by leveraging NAS to introduce patch-based inference and network redistribution (Lin et al. (2021)) consequently shifting the receptive field to later stages. While NAS significantly exacerbates the training complexity, patch-bace based inference necessitates compiler libraries that may not be compatible with standard GPUs and incurs additional computation and latency overhead. Another recent work (Chen et al. (2023) proposed self-attention--aase pooling to aggressively compress the activation maps in the first few layers to reduce the peak memory at the cost of increased compute complexity.
Knowledge distillation for image classification can be roughly categorized into two groups: response-based KD and feature-based KD. Response-basse based KD methods derive the distillation loss by leveraging the logit outputs from the fully connected layers of the student model and the teacher model. For example, KD (Hinton et al. (2015) ) distills knowledge by matching the prediction probability distributions of the student architecture and the teacher architecture. DKD (Zhao et al.
2022 ) decouples the classical KD loss into two parts, target class knowledge distillation (TCKD) and non- target class knowledge distillation (NCKD) enhancing training efficiency and flexibility. MLLD (Jin et al. 2023) performs logit distillation through a multi-level aligument based on instance prediction, input correlation, and category correlation, delivering state-of-the--ad art performance. In contrast, bassed KD methods (Adriana et al. (2015): Zagoruyko Komodakis (2016): Passalis & Tefas (2018); Loe al. (2018); Tung & Mori (2019); Ahn et al. (2019)) reduce the disparity between features in the teacher and student models, compelling the student model to replicate the teacher model at the feature level. RKD (Park et al. (2019) employs
3

---

Published in Transactions on Machine Learning Research (04/2025)
relation potential function to convey information from the teacher's features to the student's features. ReviewKD (Chen et al. (202T)) aggregates knowledge of the teacher from different stages into one stage of the student, the so-called knowledge review', which achieved impressive performance. KCD (Li et al. (2022) iteratively condenses a compact knowledge set from the teacher to guide the student learning by the Expeetaaion--Mxximiaaioooo (EM) algorithm, which would empower and be easily applied to other knowledge distillation algorithms. Existing methods focus on the distillation from a high-capacity teacher model with a large amount of parameters to an efficient student model with limited parameters. In this work, the student model, employing large kernel size and stride in the initial pooling layer possesses the same number of parameters as the teacher architecture but incurs significantly lower peak memory.
Knowledge distillation for diffusion models has gained popularity. For example, One Step Diffu- sion (Yin et al. (2023)) defines two score functions, one the target distribution and the other of the synthetic distribution produced by a one-step generator. By minimizing the KL divergence between these two score functions, the one-step generator is enforced to match the diffusion model at the distribution level and achieves impressive performance. Adversarial Diffusion (Sauer et al. (2023) utilizes both score distillation loss and adversarial loss. The score distillation loss occurs between the teacher diffusion sampler with large number of T steps and the student diffusion sampler with one or two steps. Meanwhile, the adversarial loss originates from a discriminator trained to differentiate between generated samples and real images. Auto Diffusion (Li et al. (2023) searches for the optimal time steps and compressed models in a unified framework to achieve effective image generation for diffusion models. In summary, existing KD methods for diffusion models mainly focus time-step reduction and model compression. We offer unique and orthogonal approach by minimizing peak memory that can be casily integrated with existing methods.
3 Proposed Method
3.1
Preliminaries
Knowledge Distillation We are given a dataset X, a high-capacity teacher architecture T and a to-learned student architecture S. For an input image sampled from X, TT(x) and TT(x) denote the outputs or intermediate features of the teacher and student, respectively. The knowledge distillation task aims to optimize the student's parameters wn
① arg IiiL
Σ (() DEX
(1)
where w denotes the trainable weights of as and L denotes the loss function defined by different knowledge distillation methods. For instance, Hinton et al. (2015) defines Ts and गт as the logit outputs (without applying the softmax function) of the student and the teacher, while L as the Kullback-Leibler divergence between as and Ny after applying the softmax function with temperature tp:
LKL
Σ KT(softmant
TT(x), ), softmar TT(x)) tp
Some other methods (Zagoruyko & Komodakis (2016); Park et al. (2019) Chen et al. (2021)) define different Ts and ny, such as the intermediate activation maps from various stages of the student and teacher, or different L, like the p-norm, to achieve various distillation methods.
KD methods can be categorized as response-based methods with uT defined as the logit outputs, or feature- baser methods with Ts, गт defined as the intermediate activation maps. Response- based KD methods generally keep the same peak memory with the student model, since they don't alter the student architecture. For feature-based KD methods, some might increase the student's peak memory due to extra trainable modules. However, none of these KD methods can lead to peak memory lower than the student, which is lower bound. Our method reaches such a lower bound and achieves the highest accuracy compared to existing KD methods, as shown in comprchensive experiments.
Denoising Diffusion Probabilistic Models Diffusion models Ho et al. (020 are latent variable models of the form (o == p (Po where T1,..., PT are latents of the same dimensionality the data
4

---

Published in Transactions on Machine Learning Research (04/2025)
Teucher Archilecture
f
feix
1i CCR77
clusifier
X
stage
stage
pe((sp)
Distillation Blocks
H 1f e R "
f
RED Black
RED Block
ps(F:p)
Ps(::
stage
stage
staget I
stager ++1
R
Student Architecture
activation
cotratttion layer
p(o.p)
pooling kayer
dimilition
Our proposed residual encoded distillation framework (ReDistill). RED blocks are incorporated into the student model following the pooling layers to minimize the discrepancy between the down-sampled features of the student and teacher models.
Figure 2:
#g ~ (xr)). The joint distribution (= is defined as Markov chain with learned Gaussian transition starting at p(=T) = N(xr; 0, I), where T is the maximum time step. In the training process, are given a noisy input Ht, which is derived from the data 26 and noise ~ (O,
Vorzo vI oie,
(3)
where ae == [T Bs), and Bg is the forward process variances fixed constant in DDPM (Ho et al. (2200) The loss of diffusion model is generally defined as follows:
Lass ||e (0(
The noisy input ,t, accompanied by time step embedding t, is input into a denoising autoencoder, specifically a U-Net network as in Ho et al. (2020). to estimate the noise component co(x,, t).
3.2 Proposed Distillation Framework
Our proposed framework is illustrated in Fig. 2 To reduce the activation peak memory, the initial pooling layer of the student is assigned larger pooling stride. However, still keep the same spatial dimensions of the activations to the final fully connected layer for the teacher and student. Thus, the student has fewer pooling layers but with a larger pooling stride at the initial pooling layer compared to the teacher.
Take an input image x E Hx XC as the example. The teacher and the student network are divided into several stages by pooling layers as shown in Fig, 2 We assume that all pooling layers of the teacher have the same pooling stride p for simplicity. The initial pooling layer of the student at stage 1 is assigned with stride p²_ Hence, the student feature map ff R xC, at the output of this pooling layer has the same spatial dimension the teacher fcature map J * XC2, at the output of the k-th pooling layer. Then the matched feature maps I and f will be fed into the residual encoded distillation (RED) block illustrated in Fig. 3 to compute the distillation loss. The output /P is to be fed into the following layers of the student network.
Similarly, the feature map f of the student stage also has the same spatial dimension as the feature map f241 of the teacher at stage + 1. They are fed into another RED block to calculate distillation loss
5

---

Published in Transactions on Machine Learning Research (04/2025)
RED Block
f
shortcut
logit
Rosichel
Ehooder
fx
LEEO
Figure 3: Residual Encoded Distillation (RED) Block. We use logit module for the multiplicative gating mechanism and a residual encoder for additive residual learning.
and output f2 to the following layers. This process is repeated until the last pooling layer pr(o; p) at stage of the teacher. We assume the size of the output feature map S of the teacher is identical to the output feature map 1 of the student at stage i, where i - 1 + k n. The following pooling layers of the student at stages i + 1, + 2.... n are all assigned with stride 1. Hence, these pooling layers do not change the spatial dimension and are similar to standard convolution. As a result, the final aggregated features of the student and the teacher have the same spatial dimensions. In Fig. 5 we show an example of ResNet18 with aggressive
pooling and how integrate our proposed RED blocks.
3.3 Residual Encoded Distillation Block
The proposed Residual Encoded Distillation (RED) Block, depicted in Fig. serves as the central module of our framework. It is designed ensure that the output of pooling layer closely resembles the distribution of the input while preserving essential features at reduced spatial dimension. To accomplish this, the RED Block designed to be lightweight, enabling it to modify the feature space distribution of the student's pooling layer effectively. This allows the student model to learn the down-sampled features of the teacher model. Meanwhile, this block introduces non-linearity to the pooling layer, enabling the student's pooling layer to aggregate features like standard pooling layer and adjust the feature distribution similar to a convolutional layer with an activation function. Specifically, we use a logit module for the multiplicative gating mechanism and a residual encoder for additive residual learning. The residual encoded distillation block could be formulated as follows:
fd Residual ReLU BN (Contssaa))))) logit Sigmoid BN (Convisa())
j + SS
(5)
(6)
(7)
where fs is the feature map from the student model. The logit module consists of a x convolution layer, a batch norm layer, and a sigmoid activation function, generating element- weights fs like a gate to suppress non-significant components of fs. The residual encoder module consists of a 3x3 convolution layer, batch norm layer, and relu-6 activation function, yielding the residual item fR, We use relu-6 to bound activations to prevent exploding gradients. Besides, relu-6 is widely used in efficient neural network design (Sandler et al. (2018) Howard et al. (2019); Lin et al. (2020)) since it is particularly useful for fixed-point or low-precision inference in quantization. For kernel size in RED module, we choose x 1 kernel for the logit module (LM), because LM is a gating mechanism to select critical components of the activations. We choose BatchNorm layer because the proposed framework primarily used for the distillation of CNNs, for which BatchNorm is commonly used. We hypothesize that the output of the student pooling layers might lack some crucial information compared to the teacher's down-sampled features, which could be compensatod by the residual item R

---

Published in Transactions on Machine Learning Research (04/2025)
Teacher
Denoising U-Net
Xt
X2
63
ddown
Pup
Student
Denoising U-Net
Aggressive Pool
X
Flock
Block
X2
shortes
doanssanple Jayer
costrolation fiper
up-sapplo layer
Figure ReDistill for denosing network DDPM (Ho et al. 2020 We integrate RED blocks into the student model after the down-sampling layer in the encoder and before the up-sample layer in the decoder.
3.4 Loss Function
The RED loss first calculates the mean value alongside the channel dimension and then minimizes the cosine distance between the teacher's feature map f and the RED block's output fD, which is formulated
CRED(F .88) = df Execr Sf. Eeecs Se [Csl
(8)
|Crl
where pP is calculated by Equation and d denotes the distance measurement. OT and Cs denote the channel dimension of teacher's feature map f and the RED block's output f", respectively, The final loss function of the proposed method is as follows:
Ciusk Σ aLRED,
where Lrask the vanilla loss from the task, such as the Binary-Cross- (BCF) loss for image classification. denotes the number of RED blocks. a is an experimental syper-parameter to scale RED loss.
3.5 Distillation for Diffusion Model
In this section, we introduce how to integrate the proposed distillation framework into U-Net based denoising network described in DDPM (Ho et al. (0200), as shown in Fig. 4 For convenience, let's 's assume the teacher model is U-Net with two down-sample layers, each having a stride of 2, while the student is a U-Net with the aggressive pooling setting, i.c., it just has one down-sample layer with the stride 4. We use the output is SSon from the down-sample layer of the student model and the output yoo from the second down-sample layer of the teacher model, as they share the same spatial dimension. fåown and fTow are input into a RED block incorporated into the student model, producing the output fP which is then fed into the subsequent convolutional layers. Symmetrically, the student model has one up-sample layer with expansion ratio of x4, while the teacher model has two up-sample layers, each with an expansion ratio of x2. The input f of the student's up-sample layer and the input. Súp of the teacher's first up-sample are input inlo another RED block, producing the output fo This output replaces the f as the new input of the student's up-sample layer. For DDPM distillation, the loss function is defined as Equation while the Crask is replaced by Lasfs defined in Equation

---

Published in Transactions on Machine Learning Research (04/2025)
4 Experiments
In Sections 4.1 and 4.2 datascts and implementation details for image classification and image generation
lasks are introduced separately. In Sections 4.3 and 4.4 we conduct experiments on various vision lasks to illustrate the effectiveness of the proposed method with state-of-the-art distillation methods under different. backbone architectures and datasets. We also compare the memory footprint of our method with the teacher and the student architectures, and deploy our method on the edge device. Results are reported in Section 4.5 At last, some ablation study related with module discussion, loss function, and distillation strategy are reported in Section 4.6
4.1 Datasets
Datasets for Image Classification 1) STL10 (Coates et al. 2011 contains 5K training images with 10 classes and 8K testing images of resolution 96 x 96 pixels. Specifically, we resize the image resolution to 128 x 128 pixels for aggressive pooling. 2) ImageNet (Russakovsky et al. (2015)) is a widely-used dataset of classification, which provides 1.2 million images for training and 50K images for validation over 1,000 classes. We keep the same resolution of 224 224 pixels as the origin for aggressive pooling.
Datasets for Image Generation 1) CIFAR-10 (Krizhessky et al. (2009) comprises 60,000 color images of 32x32 resolution across 10 classes, with each class containing 6,000 images. The dataset is divided into 50,000 training images and 10,000 test images. We keep the original resolution of 32 x 32 in our experiments. 2) Celeb-A (Liu et al. (2015)) is a large-scale face attributes dataset containing over 200,000 celebrity images, each annotated with 40 attributes. We use the resized resolution of 64 64, which is widely used in diffusion-based methods (Song et al. (2000); Bao et al. (023) in our experiments.
4.2 Implementation Details and Baselines
Details for Image Classification Unlike traditional distillation tasks, we only modily the pooling layer strides instead of the depth and width of the network to get the student model, which is called the aggressive pooling setting. The advantage of the aggressive pooling setting is to reduce the peak memory and also reduce the computational complexity and inference time of the network. Specifically, we increase the first pooling layer stride x2 ~ x8 times and adjust the last several pooling layer strides to ensure the final output of the student model has the same information density the teacher model. All experiments are implemented in Pytorch and evaluated on 4 NVIDIA A100 GPUs.
On STL10 dataset, we experiment with three representative and widely-used network architectures, including MobileNetV2 [Sandler et al.] (2018)), ) MobileNetV3 (Howard et al. (2019)), and ResNext (Xic et al. (201))) p). The proposed method is compared with several distillation methods (Hinton el al. (2015); Adriana el al. (2015): ; Zagoruyko & Komodakis (2016); Tung & Mori (2019); ; Ahn et al. (2019); ; Park et al. (2019); Heo et al. (019)) Kim et al. (2018); Huang & Wang (2017)) Specifically, the student architecture is trained from scratch being distilled from pre-trained teacher architecture by different methods for 300 epochs. The batch size is set to 8 and the dropout rate is set. to 0.2. The SGD with momentum equal to 0.9 is used as the optimizer. The initial learning rate is set to 0.01, which is reduced by a factor 0.2 at the 180th 240th and 2h epoch, respectively. The a in Equation g is set to 50.
On ImageNet dataset, we experiment on the MobileNet (Sandler et al. (2018)) and ResNet (He al. (2016)), ) which are widely used in distillation benchmarks (Jin al. (2023) : Chen al. (2021); Tian al.
(2020)
The proposed method is compared with based methods like KD (Hinton et al. (0155)) )
DKD Zhao et al. (2022)) ), MLLD (Jin et al. (2023)) ) LSKD Sun et al. (2024) ) and feature-based hoth
like Fit (Adriana et al. (2015)). RKD (Park et al. (2099) ReviewKD (Chen et al. (2021)) CRD Tian
2020 ), and OFAKD (Hao el al. (002)). ). All these methods are widely used in knowledge distillation and, to the best of our knowledge, yield SOTA performance. We use the same experiment settings as (Jin
et al.
ct al. (2023) ) but keep training for 300 epochs and decay the learning rate at the 180 240th and 27
epoch with factor 0.1, since the student architectures with the aggressive pooling setting generally require more epochs to converge. The in Equation g is set to be 1.
8

---

Published in Transactions on Machine Learning Research (04/2025)
Details for Image Generation For the teacher model, we keep the same experiment settings as DDPM Ho et al. (2020)) ) with applying T 1000, 31 = 10-1, 3rr = 0.02, and the U-Net backbone with 4 different feature map resolutions (32 x 32 to 4 x 4 for CIFAR-10, while 64 x 64 to 8 x 8 for Celeb-A). For the student model, we increase the first pooling layer stride of the U-Net backbone x2 times while adjusting the last pooling layer stride to keep the same latent feature resolution. The same stride modification is symmetrically applied to the up-sample layers of the U-Net, and thus with different fcature map resolutions (16 x 16 to 4 x for CIFAR-10, while 32 32 to x 8 for Celeb-A). For our method, the RED blocks are inserted after not only the down-sample layers but also the up-sample layers. For CIFAR-10 dataset, train all models for 1000K iterations and sample 50K images for FID (Heusel et al. (2017)) & IS (Salimans et al. (2016) evaluation. For Celeb-A dataset, we train all models with 250K iterations and sample 50K images for FID (Heusel et al. (2017)) & IS (Salimans el al. (2016)) evaluation. All experiments are implemented in Pytorch and evaluated on an NVIDIA 4090 GPU.
Details for Theoretical Peak Memory For the theoretical peak memory analysis, we follow previous work, MCUNet (Lin et al. (2020) 2021 2022)), for calculating peak memory. Generally, it traces the memory consumption for each layer. For the standard convolutional layer, the maxpool layer, or the average pool layer, the memory consumption is the summation of the input activation memory size and output activation memory size. For a group convolutional layer, the memory consumption is the summation of the input activation memory size, the output activation memory size, and a buffer with size equal to one channel convolutional kernel in the group. For the residual connection in the network, the residual activations memory size will be added into memory tracing until the residual item is added with the output activations.
Details for Aggressive Pooling Sctting. To explain how exactly the image classífication models are configured, we list the network conlig details for the 'T: ResNet18' and 'S: ResNet18x4' distillation pair that used in the third column of Table [ As shown in Fig. 5 we aggressively increase the stride of the first downsampling layer 'Conv 1 from 2 to 8 following the aggressive pooling setting, while setting the stride to 1 of maxpool layer and the last downsampling layer, i.c., the first conv layer in 'Stage 4', for keeping the same activation size for average pool and fe layer. Then the student network possesses only three downsampling layers less than five downsapling layers of the teacher network. We integrate the proposed RED block after each downsampling layer of the student in our ReDistill framework to improve the student's performance while maintaining low peak memory.
4.3 CNN-based Image Classification
Table shows the results on the ImageNet (Russakoskky et al. (2015)) dataset, with the setting that the teacher model and student model are in identical architecture familics or different architectures. 'xn' donotes that we increase the 1st pooling layer stride of this architecture with n times, and the best results are highlighted in boldface. Our method achieves the best performance compared with different response-based and feature- based distillation methods, no matter for identical architecture family knowledge distillation, as shown in the 3rd and 4thh columns of Table 1 or different architecture knowledge distillation, as shown in the 5th column of Table 1
We find some distillation methods perform even worse than the student model itself without any distillation since these methods are not specially designed for the aggressive pooling setting and would be sensitive to the resolution of the feature maps, like Fet, or require multi-scale feature maps, like RKD. Generally, conventional foaturo-based methods match the activations of the tcacher model and student model stage by stage. In each stage, these activations possess the same resolution. However, in our aggressive pooling setting, the resolution of student activations and teacher activations are mismatched in each stage, and thus, traditional distillation methods are not guaranteed to be positively effective in knowledge distillation with aggressive pooling. This observation also illustrates the necessity of the proposed ReDistill framework.
Table shows the classification results on STL10 dataset with the setting that the teacher model and student model are in identical architecture family. Same as ImageNet, 'xn' denotes we increase the 1st pooling layer stride of this architecture with n times, and the best results are highlighted in boldface. Some KD methods in Table I require customized dataloader with contrastive version of the input data, like CRD Tian et al.
9

---

Published in Transactions on Machine Learning Research (04/2025)
ResNet18
S: ResNet18x4
RED (ours)
Network Config
layer config
activation size
layer config
activation size
layer config
activation size
Input
224X226
224X
224x224
Conv
conv 7x7: c_64,
112x112
comp 7x7: c_64,5_8
28x
7x7: 6,64, RED block
28x28
maxpool 3x3: s_2
maxpool 3X3: x_1
maxpoul 3x3: x_1
3x3:c_64,5,
3x3.c_64,5__
3x3:
Max Pool Stuge
Looay 3x3: 64,
56x56
Loone 3x3:: 64,8
28x28
Leonv 3x3: 64,
28x
connp 3x3:c_64,,__1 1,1
[conp 3x3:c_64,5.1
icome 3x3: c,64,,1 1oomv3x3:c_64...1
conv 3x3:c_64,
conp 3x3:c_64,5_1
conr 3x3:c_128,
coov 3x3c__2_2_
foonur 3x3:c. 128,
cuner 3x3 c_128.
conv 3x3:c_128,__1
+ +RED btock
Stage 2
[conv 3x3:c_12__,
conr cx28c_128,__1
conp33x3:c_128..1
ccone 3x3;c_128,5_1
conv 3x3:c_128.5_1
[conv 3x3:c_100_1
conr 3x3:c_256.
coou coor 3x3:c.256,,22
[ooRu 3x3:c
cuner c,
conv 3x3:c_256,;_11
RED black
Stage 3
conn 3x3:c 256
ccans 3x.3: 256
monnv3x3:: 256
com.3333._256.
Lconv 3x3:c_256,511
[conv 2x3_,xx:_2566,51
coiu 3x3:c_512,5.2 conv 3x3,5222_.5
[ooaP [omm 3x3:c_51222._]
[canv 3x3:c_512,5_1
convv 3x3:c_512.5_1
[coRv 3x3:c_5122,551
Stage 4
[oonu 3x3:c.6122.. cune c
icanu c.12 (conv 512.
[EVRF 3x3:c_512.a.11
Avg Pool + FC
average pool, 1000-d
1x1
average pool. 1000-d
1x1
average pool, 1000-d éc.
1x1
Figure 5: Example of our proposed Aggressive Pooling Setting with ResNet18. We highlight all downsampling layers (either conv layer maxpool layer with the stride larger than 1) in red color. For example, in the "Conv 1' cell of 'T: ResNet18', 'conv 7 x 7 64, 2' denotes this convolution layer is with kernel size 7 x 7, number of channels 64, and stride 2. In the following row, 'maxpool 3 3 s__" denotes the max pool layer with kernel size 3 x 3 and stride 2. In the student network 'S: ResNet18x44; we increase the stride of the first downsampling layer 'Conv 1' to 4x from s s_2 to s s_8, while setting the stride to 1 of the maxpool layer and the last downsampling layer (i.e., the first conv layer in Stage 4), in order get the same activation size before average pool and fe layer.
Table 1: Top-1 accuracy (%) on ImageNet. The 3rd and 4th columns show the results with identical teacher and student architecture families, while the 5th column shows the results with different teacher and student. architectures.
Teacher
ResNet18
ResNet50
ResNet152
Top1 Acc. (%)
69.75
76.13
78.32
Method
Student
ResNet18x4
ResNet50x
MbNetV2x2
Topl Acc. (%)
61.79
69.50
62.65
KD [Hinton et al.
2015
63.63
70.60
62.85
DKD
(Zha et al
2022
63.22
66.27
Response
MLLD
2023
64.66
70.77
68.36
LSKD Sun et al.
2024
63.81
72.28
66.21
FitNet (Adriana et (015
62.13
71.77
60.79
RKD Park et al. 2019)
61.49
66.88
ReviewKD
CChnn et al
(2021]
63.30
70.22
63.07
Feature
CRD Tian et al. (2020)
64.01
71.07
65.60
OFAKD Hao et al.
2023
64.83
68.49
RED (ours)
65.23
73.23
68.89
< denotes that get reasonable resulta for the student architecture with the distillation method under the aggressive pooling setting.
10

---

Published in Transactions on Machine Learning Research (04/2025)
2023 However, they don't provide the corresponding dataloader for STL10 dataset in their original code im implementations. Thus, compare our method on STL.10 dataset with some feasible KD methods in Table I like KD, FitNet and RKD, and also compare with some other commonly-used KD methods, like AT, VID, NST, etc. Under the identical architecture family setting, our method performs the best among all state-of-the-art, distillation methods. We also measure the average distillation time for different distillation methods, as shown in the last column of Table 2 Our method achieves similar time consumption compared to most distillation methods. For different architecture knowledge distillation, shown in Table our method achieves better performance compared with other distillation methods as well.
2020 MLLD Jin et al.
4.4 Image Generation
Quantitative Results Table 4 shows the results on CIFAR10 (Krizhessky et al. 2009)) dataset and Celeb-A (Liu et al. (2015)) datasct. Our method reduces the fidelity degradation of the student model with first pooling stride that is twice that of the teacher model. Specifically, our method achieves 3.43 lower FID and 0.51 higher IS score than the student model on CIFAR10 dataset, while 1.6 lower FID and 0.11 higher IS score than the student model on Celeb-A, respectively. We also compare to a simple distillation method, which termed MSE-Distill-DD SiaE--Ditill-DDDDM.. DDPM, by directly matching the final output activations of the teacher model and student model using MSE loss. On CIFAR10 dataset, the student with and 9.62 IS score, while our method achieves a better 10.85 FID and 9.63 IS score. Due to implementation specifics, the models are not augmented with EMA (Hunter (986)), leading to results slightly different from those the original paper.
Table 4: Results on U-Net (Ronneberger et. al. (0015) based DDPM (Ho et al. 20200).
Dataset
Method
ISt
FIDI
T: U-Net w/ DDPM
9.70 00133
3.75
U-Netx2 w/ DDPM
9.12 1 0.14
14.28
CIFAR1O
S + MSE-Distill
9.62 + 0.11
12.04
S w/ RED (ours)
9.63 0.10
10.85
T: U-Net w/ DDPM
2.97 + 0.04
19.61
U-Netx2 w/ DDPM
2.77 + 0.02
23.03
Celeb-A
S + MSE-Distill
2.71 1 0.02
22.34
S w/ RED (ours)
2.88 1 0.03
21.16
achieves 12.04 FID
Visualization To further illustrate how the proposed method improves the fidelity of image generation, we visualize some samples generated by the teacher model, the student model, and our method. Specifically, we utilize the same noise item in each time step for all these three models. In this way, the generated images are expected to be visually similar if the two models have comparable capabilities. As shown in Fig. 6 generally, the images generated by our method are semantically closer to those generated by the
Table 2: Top-1 accuracy (%) on STL10 with identical teacher and student architocture family.
T: MbNetV2
T: M.NNeV33-Small
T: ResNext18
85.34
83.74
85.12
Distill Time
Method
S: MbNet V2X4
S: Mb.ett33SSx4
S: ResNext18x4
(s/epoch)
76.18
71.27
79.07
KD (Hinton et al. (2015)
79.26
71.56
81.27
3.23
FitNet
[Adriana et all (2015)
78.55
71.93
80.21
3.28
AT
Zagoruyko & Komodakis (0166)
81.35
75.88
82.90
3.33
SP
Tung & Mori
(2019
72.42
77.59
3.31
VID
Ahn al.
2019]
77.53
71.83
76.99
3.40
RKD
Park et al.
2019))
78.00
67.27
75.80
3.33
AB
Heo et
80.79
72.94
81.42
3.69
FT
Kim et al.
2018
74.74
72.11
80.64
3.35
NST
Huang &
Wang (2177)
82.66
76.11
81.89
9.33
RED
(ours)
83.97
77.31
84.80
3.88
1 denotes that do not reasonable results for the student architecture with the distillation method under the aggressive pooling setting.
11

---

Published in Transactions on Machine Learning Research (04/2025)
Table 3: Top-1 accuracy (%) on STL10 with different teacher and student architecture.
T: 18sNext18
T: ResNext18
T: MbNetV2
85.12
85.12
85.34
Method
S: MbNetV2xx
S: Vo-Net33-Smallle
S:
76. 18
71.27
71.27
KD
(Hinton (2015)
78.
69.98
70.95
FitNet [Adriana et al.
(2155)
80.46
73.23
72.88
AT
Zagoruyko Komodakis (2016))
80.63
74.58
73.39
SP
Tung & Mori
20191)
67.56
63.91
Alun et al.
74.40
69.19
71.69
RKD
Park et
2019
70.63
72.10
68.78
AB
Heo al.
2019
81.46
73.98
75.20
FT
Kim et al.
20188
77.33
69.79
66.85
NST
(uang
Wang 2017
79.09
65.04
73.26
RED
(ours)
83.23
77.15
77.19
denotes that we don't get reasonable results for the student architecture with the distillation method under the aggressive pooling setting.
Shudent Teachrr
CIFARI0
Ours
Teacher
Celeb-A Student
Our
Figure 6: Generated images CIFAR10 and Celeb-A with the same noise item for all models. Generally, our results are semantically closer to the teacher's.
teacher model, although in cases, they closer to those generated by the student model. Due to the high capability of the teacher model, method, intermediate model from the teacher to the student, achieves higher fidelity than the student model.
4.5 Memory Footprint and Edge Device Deployment
Memory Footprint To intuitively demonstrate the peak memory reduction enabled by proposed method, we trace the memory footprint in layer-wisc for all identical architecture family teacher-student pairs, as shown in Fig. 7 Compared to CNN models with vanilla pooling, our method with aggressive pooling settings achieves significantly lower memory consumption, particularly the initial layers where peak memory usage occurs. For DDPM, which utilizes a U-Net architecture with both down-sampling and up-sampling layers, our method reduces memory consumption in both the initial layers and the final layers where peak memory usage occurs.
12

---

Published in Transactions on Machine Learning Research (04/2025)
Memory Footprint- MobileNetV2
Memory Footprint- MobileNetv3 Small
Memory Footprint- R8sNeXt18
Baseline Ours
Baselino
Baseline
1.25
Ours
Ours
Memory (MB)
Memory (MB)
Memony (MB)
0.50
2-
0.25
0.00
2
10
12
14
16
2
10
2
7
Memocy Footprint- R 8
Memory Fooiprint ResNet50
Memory Fooiprint DDPM
Bareline
Baseline Ouss
Daseline
Menoor (MB)
32
Ouss
Our
Memory
Memory
x
4 Layer
2
6
10
12
14
16
20
24
Layer
Layer
Figure 7: Memory footprint for each layer in teacher and student networks (top-all: STL10 classification task; bottom-left, bottom-mid: ImageNet classification task; bottom-right: Celeb-A generation task). The student model is the teacher model's modified version of being assigned with the aggressive pooling strategy and enhanced with our proposed RED blocks, representing our method.
Edge Device Deployment In Table 5 we mcasure theorctical peak memory consumption, actual peak memory consumption, and other efficiency related metrics like model size, maximum power, and latency, for all the models we use STL10 & ImageNet classification and image generation tasks. Specifically, the 1st_gth rows are for models we use in STL10 classification, and the 10th-18th row are for models we use in ImageNet classification, while the last rows are for models we use in image generation task. We estimate the ical peak memory by summing the size of the input & output allocation for each operation and assume a batch size of 1 for theoretical estimation to emulate most inference use-cases, similar to (Lin et al. (2021); Chowdhery et al. (2019)). Our method achieves 3.9 x 14.77x reduction in theoretical peak memory for image classification and 4x reduction in theoretical peak memory for image generation, as shown in Table 5 In addition to the theoretical peak memory, we also measure the actual peak memory consumed on an NVIDIA Jetson TX2 device. Moreover, we measure the maximum GPU power (to ensure we satisfy the power budget of edge devices), model size, and latency incurred by the baseline teacher, student, and our RED models the same edge GPU, Our models yield similar peak memory the student models (due to similar levels of aggressive striding) and ~2x~332x lower peak memory for image classification and ~2x lower peak memory for image generation lasks compared to the teacher models. Note that our theoretical and measured peak memory reduction factors are different due to varying device setups and buffer allocations. However, our models incur worse latency and power compared to the student models due to the additional RED blocks and improved latency compared to the teacher models.
4.6 Ablation Study
Module Discussion As shown in Table [ we conduct the ablation study on STL10 dataset with the MobileNetV3-Smmall backbone by removing the logit module, residual encoder, shortcut, and the whole RED block, respectively. The performance of our proposed model degrades slightly without the logit module but scriously degrades without the residual encoder, which illustrates that the residual encoder plays a more significant role in the RED block. Without RED blocks, the student model has poor performance when only applying RED loss on the student activation maps and teacher activation maps, which illustrates the nccessity to integrate RED blocks into the student model. Naively stacking the logit module and residual encoder, i.e., without Shorteut, performs even worse than without RED blocks, illustrating the effectiveness of our designed shortcuts in the logit. module and residual encoder. We also conduct the ablation study on
13

---

Published in Transactions on Machine Learning Research (04/2025)
Table 5: Edge Device Deployment. We analyze the theoretical Peak Memory (T-PkMem), and measure actual GPU Peak Memory (A-PkMem), Model Size (MS), Maximum Power (MxP), and Latency for models on an NVIDIA Jetson TX2 device. For theoretical Peak Memory analysis, the batch size assumed to be 1. For each actual measurement, the batch size (BS) reported in the table and set to the maximum load allowed by the device, which is restricted by the teacher. The 1 rows are for models classification, and the 10th row are for models we use in ImageNet classification, while the last 3 rows are for models we use in image generation task. The overheads of RED modules on top of student networks is minimal particular for cases with more aggressive pooling, c.g., distillation of MbNetV2 with x8 pooling.
use in STL10
Theoretical Analysis
Actual Measurement
Model
T-PkMem (MB)
BS
A- PkMem (GB)
MS (MB)
MxP (mW)
Latency (ms)
T: MbNetV2
7.50
2.30
13.50
1830
221.53
2.34
S: MbNett2xx
0.51
20
0.78
13.50
458
49.51 + 1.98
w/ RED (ours)
0.51
0.78
14.25
55.43 = 1.80
T: MbNetV3-Snnn
1.41
2.30
9.75
1526
195.63 1.09
S: V V3-Smallx4
0.31
50
1.10
9.75
611
69.76 = 1.55
w/ RED (ours)
0.31
1.10
10.18
687
83.74 = 1.40
T: ResNext18
5.00
2.40
21.47
3276
448.38 + 6,97
ResNext188x
0.75
50
21.47
1450
173.23 1 1.92
8 w/ RED (ours)
0.75
1.10
25.56
2058
275.02 1 3.25
T: ResNet18
3.83
2.70
44.63
3273
395.42 + 4.21
S: 1
0.77
50
1.20
44.63
1297
151.59 + 1.50
w/ RED (ours)
0.77
1.20
48.72
1830
208.43 = 1.53
T: ResNet50
9.19
2.00
97.70
2591
293.81 1.09
S: ResNet5004
2.30
10
1.00
97.70
1068
125.49 = 1.17
w/ RED (ours)
2.30
1.00
160.44
2136
246.68 = 1.73
T: ResNet152
9.19
2.80
230.20
3880
612.03 3.55
S: MobileNetV2x2
1.15
8
0.71
16.23
305
23.98 = 0.84
w/ RED (ours)
1.15
0.71
32.73
611
72.65 1.01
T: UNet w/ DDPM
3.50
2.90
133.09
4257
684.31 5.39
S: UNetx2 w/ DDPM
0.88
30
1.30
133.09
1983
249.92 = 3.76
w/ RED (ours)
0.88
1.30
149.922
2362
304.97 + 1.87
Table 6: Module Discussion STL10 dataset with 3 with a=50,, and cosine distance for the RED loss. LM denotes Logit Modulc, and RE denotes Residual Encoder. ks denotes the kernel size of the convolution layer in RE.
RED block
Method
Top1 Acc. (%)
LM
RE
Shortcut
ks
w/o LM
3
76.66
w/o RE
4
72.30
w/o Shortcut
3
52.42
w/o RED block
69.31
RF w/ ks=1
73.86
RE w/ ks=5
5
76.26
RED (ours)
3
77.31
the kernel size of the residual module of our RED block. It shows that the proposed method achieves the best performance with kernel size 3.
Distillation Strategy We conduct the ablation study to compare different strategies during distillation, including the feature alignment strategy and pooling stride of the initial layer. Specifically, we train the ResNext18 network on the STL10 dataset with three different strategies while enlarging the stride of the first pooling layer to x2, x4, and x8. The results shown in Table 7 REDD w/o Distillation' is the baseline of
14

---

Published in Transactions on Machine Learning Research (04/2025)
Table 7: Ablation study distillation strategies on STL10 dataset with ResNext 18.
Backbone: ResNext18
Topl
Acc.
(%)
1st Pooling Stride
x2
x4
x8
RED w/o Distillation
81.43
79.07
73.16
RED w/ Stage-Align
80.59
81.99
74.00
RED w/ Pooling-Align (ours)
85.51
84.80
77.48
training aggressive pooled ResNext18 without distillation. 'RED w/ Stage-Align' applies the traditional stage-based feature alignment between the student and the teacher activations for distillation. 'RED w/ Pooling-Align' is the proposed pooling-based feature aligmment strategy. Compared to traditional stage-based alignment, the proposed pooling- bosseng-based alignment performs much better for the aggressive pooling strategy.
Loss Function We also conduct the abla- tion study for the proposed RED loss func- tion, an equally important portion of distil- lation. The results are based on ResNext18 backbone and STL10 dataset, as shown in Table Without the RED loss, i.e. setting a 0 and then the loss function the student model is just integrated with RED blocks but without distillation. In this circunstance, the poor performance of the student model illustrates the need to apply the loss of RED. The proposed method per- forms worse when applying the Euclidean distance instead of the cosine distance to cal- culate the RED loss. Averaging along the channel dimension with different Cf and Cs in Equation g might mitigate the absolute difference in Euclidean space, while cosine distance measures the angle between two vec- tors, preserving their relative difference. By adjusting from 200 to 1, the performance degradation shows that the proposed method is sensitive to the hyper-parameter in Equation g Besides, evaluate the proposed method's performance when incorporated with other knowledge distillation methods, such KD (Hinton et al. (2015)). Specifically, we first evaluate the student model integrated with the proposed RED blocks by applying only KD loss, which is denoted as FREDw/oLRES KDD Then we apply both RED loss and KD loss into the same student model, denoted as 'RED (DD'))+KD.. Performance improvement illustrates the effectiveness of the proposed RED loss.
Table & Loss Function Discussion STL10 dataset with NesNex118x4. We conduct ablation study on a and distance measurement for CRED. We also compare our method with the student integrated with the proposed RED block but distilled only by LKD and distilled by both our CRED and LKD.
Lrask,
CRED
Method
LKD
Topl Acc. (%)
Distance
82.26
50
Euclidean
82.66
Cosine
82.04
Cosine
82.86
RED (ours)
10
Cosine
83.41
50
Cosine
84.80
100
Cosine
85.10
200
Cosine
84.95
RED w/fCRED +KD
81.94
RED (ours)+KD
50
Cosine
84.82
Distillation w/o Aggressive Pooling In addition to aggressive pooling, we conduct experiments to compare several distillation mothods from the teacher network ResNext50 to the student network MobileNet V3-Small without aggressive pooling on STL10 dataset. As shown in Table g our method still achieves the best performance, illustrating the generalization ability of the proposed framework.
Table 9: Distillation w/o Aggressive Pooling Setting.
Method
Top1 Acc. (%)
T: ResNext50
83.54
S: MbNet
72.62
KD
77.06
Comparison to Non-distillation Method We compare our (QAT) method by using
FitNet
76.35
methods to crantizatiion-wwar--rreee
AT
76.76
the pytorch native architecture optimization toolkit TorchAO ftor
VID
74.83
chao maintainers & contributors (2024)) and bits QAT method, INQ (Aojun Zhou| (2017)). We also deploy Intel Neural Compres-
RKD
73.96
AB
77.50
sor (INC) [Cheng et al. 2023) to conduct int8 fully precision and
NST
73.78
RED (ours)
79.71
15

---

Published in Transactions on Machine Learning Research (04/2025)
int8/fp32 mixed precision wewttrrninngguuantzzttnn (PTQ) methods. As shown in Table 10 the ResNext18 quantized by INC to int8 achieves 20.12% accuracy and 1.25MB peak memory. The ResNext18 quantized by INC to mixed precision achieves 84.18% accuracy and 4.19 MB peak memory. The ResNext18 quantized by INQ to 6 bits achieves 84.07% accuracy and 5.00 MB peak memory. Our distillation method achieves the optimal trade-off between accuracy and peak memory, with 84.80% accuracy and 0.75 MB peak memory. Due the extra RED blocks, our method possesses the higher model size 25.557 MB. However, our method orthogonal to these quantization methods and thus can be applied with the quantization method to further optimize the student model's overhead. As shown in the last row of Table I0 the student model ResNext 8 distilled by our method and quantized by INC to mixed precision achieves 84.78% accuracy, 0.44 MB peak memory, and 5.621 MB model size.
Table 10: Comparison w/ Non-distillation Method.
Method
Top1 Acc.(%)
T- I-PkMem (MB)
Model Size (MB)
T: Nextex fp32
85.12
5.00
21.47
w/
TorchAO dynamic, QAT
83.37
5.00
19.513
w/
INC int8, PTQ
20.12
1.25
4.960
INC mixed prec.; PTQ
84.18
4.19
5.015
INQ bits, QAT
84.07
5.00
4.026
St ResNext18x4 fp32
79.07
0.75
21.47
w/ INC mixed prec., PTQ
78.50
0.44
5.015
8 + RED (ours) fp32
84.80
0.75
25.557
w/ INC mixed prec., PTQ
84.78
0.44
5.621
5 Conclusions and Future Work
We propose ReDistill, a novel residual encoded distillation method to reduce the peak memory of convolutional neural networks during inference. Our method enables the deployment of these networks edge devicos, such as micro-comte controllnsssssss with tight memory budgets, while accommodating high-resolution images necessary for intricate vision tasks. The reduced peak memory can also enable these networks to be implemented with recently proposed in-sensor computing systems (Datta al. (2022) 2023)), thereby significantly reducing the bandwidth between the image sensor and the back-end processing unit. Our method is based teacher-student distillation framework, where the student network using aggressive pooling with reduced peak memory is distilled from the teacher network. For image classification, our method outperforms existing response-based and feature-based distillation methods in terms of accuracy-memor trade-off. For diffusion- based image generation, our method significantly reduces the peak memory of the denoising network with slight dogradation in the fidclity and diversity of the generated images. While we focus on distillation of CNNs in this work, we leave as future work peak memory reduction of vision transformer via distillation for image classification (Wu al. 2022 and image generation (Peebles & Xie (2023)
This work is supported by Department of Defense under funding award W911NF-244 0295, National Science Foundation Career award #2341039, and Google Cloud research credits program.
Acknowledgement
References
Romero Adriana, Ballas Nicolas, K Samira Ebrahimi, Chassang Antoine, Gatta Carlo, and Bengio Yoshua. Fitnets: Hints for thin deep nets. In Proceedings of the International Conference on Learning Representations, 2(3):1, 2015.
Sungsoo Ahn, Shell Xu Hu, Andreas Damianou, Neil D Lawrence, and Zhenwen Dai. Variational information distillation for knowledge transfer. In Proceedings of the IEEE/CVF conference on computer vision and pollern recognition, pp. 9163 9171, 2019.
16

---

Published in Transactions on Machine Learning Research (04/2025)
Yiwen Guo Lin Xu Yurong Chen Aojun Zhou, Anbang Yao. Incremental network quantization: Towards lossless cnns with low-precision weights. In International Conference on Learning Representations, TCT720017 2017.
Fan Bao, Shen Nie, Kaiwen Xue, Yue Cao, Chongxuan Li, Hang Su, and Jun Zhu. All are worth words: A vit backbone for diffusion models. In Proceedings of the IEEE/CVE Conference on Computer Vision and Poltern Recognition, pp. 22669 22679, 2023.
Fang Chen, Gourav Datta, Souvik Kundu, and Peter A Beerel. Self-attentive pooling for efficient deep learning- In Proceedings of the IEEE/CVF Winter Conference Applications of Compuler Vision, pp. 1174-3983, 2023.
Pengguang Chen, Shu Liu, Hengshuang Zhao, and Jiaya Jia. Distilling knowledge via knowledge review. In Proceedings of the TEEE/CV Conference on Computer Vision and Pattern Recognation, pp. 5008 5017, 2021.
Wenhua Cheng, Yiyang Cai, Kaokao Lv, and Haihao Shen. Teq: Trainable equivalent transformation for quantization of Ilms, 2023. URL httpss:://xx....rr////////// abgg/bbs//3/0 10944
Aakanksha. Chowdhery, Pete Warden, Jonathon Shlens, Andrew Howard, and Rocky Rhodes. Visual wake words dataset. Xive preprint SSSS:19966007211, 2019.
Adam Coates, Andrew Ng, and Honglak Lee. An analysis of single-layer networks in unsupervised feature learning. In Proceedings of the fourleenth international conference on artificial intelligence and stalistics, pp. 215-223, JMLR Workshop and Conference Proceedings, 2011.
Antonia Creswell, Tom White, Vincent Dumoulin, Kai Arulkumaran, Biswa Sengupta, and Anil A Bharath. Generative adversarial networks: An overview. IEFE signal processing magazine, 35(1):53 65, 2018.
Gourav Datta, Souvik Kundu, Zihan Yin, Ravi Teja Lakkireddy, Joe Mathai, Ajey Jacob, Peter A Beerel, and Akhilesh R Jaiswal. processing-in in- iicptlg------i-ee-iee bexel-in-memool paradigm for resouree-cnnstranned tinyml applications. Scientific Reports, 12(1):14396, 2022.
Gourav Datta, Zeyu Liu, Md Abdullah-Al Kaiser, Souvik Kundu, Joe Mathai, Zihan Yin, Ajey P. Jacob, Akhilesh R. Jaiswal, and Peter A. Beerel. In-sensor & neuromorphic computing are all you need for energy efficient computer vision. In ICASSP 2023 2023 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 1 1-5, 2023. doi: 10.L109/ICASSSP49933 10094900...
Song Han, Xingyu Liu, Huizi Mao, Jing Pu, Ardavan Pedram, Mark Horowitz, and William Dally. Eie: Ffficient inference engine compressed deep neural network. ACM SIGARCH Computer Architecture News, 44(3):243 254, 2016a.
Song Han, Huizi Mao, and William J Dally. Deep compression: Compressing deep neural networks with pruning, trained quantization and huffman coding. International Conference Learning Representations, 2016b.
Zhiwei Hao, Jianyuan Guo, Kai Han, Ychui Tang, Han Hu, Yunhe Wang, and Chang Xu. One-for-all: Bridge the gap between heterogeneous architectures in knowledge distiilation. In Advances in Neural Information Processing Systems, 2023.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEFF conference on computer vision and pattern recognition, pp. 770-778, 2016.
Byeongho Heo, Minsik Lee, Sangdoo Yun, and Jin Young Choi. Knowledge transfer via distillation of activation boundaries formed by hidden neurons. In Proceedings of the AAAI conference on artificial intelligence, volume 33, pp. 3779 3787, 2019.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by two time-scale update rule converge to local nash equilibrium. Advances in neural information processing systems, 30, 2017.
17

---

Published in Transactions on Machine Learning Research (04/2025)
Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in neural network. arXiv preprint 1rXiu:1503.0255 2015.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advonces in newral information processing systems, 33:6840-685) 2020.
Andrew Howard, Mark Sandler, Grace Chu, Liang- Chig-hhihh Chen, Bo Chen, Mingxing Tan, Weijun Wang, Yukun Zhu, Ruoming Pang, Vijay Vasudevan, et al. Searching for mobilenetv3. In Proceedings of the IEEE/CVF international conference on computer vision, pp. 1314-1324, 2019.
Zehao Huang and Naiyan Wang. Like what you like: Knowledge distill via neuron selectivity transfer. arXiv preprint trXio:1777.01/) 012.01 2017.
Itay Hubara, Matthieu Courbariaux, Daniel Soudry, Ran El- Yaniv, and Yoshua Bengio. Binarized neural networks. Advances in neural information processing systems, 29, 2016.
Stuart Hunter. The exponentially weighted moving average. Journal of quality technology, 18(4):203-21) 1986.
Forrest N landola, Song Han, Matthew W Moskewicz, Khalid Ashraf, William J Dally, and Kurt Keutzer. Squeezenet: Alexnet-level accuracy with 50x fewer parameters and< 0.5 mb model size. International Conference on Learning Representations, 2017.
Ying Jin, Jiaqi Wang, and Dahua Lin. Multi-level logit distillation. In Proceedings of the IEEE/CVE Conference on Computer Vision and Pottern Recognition, pp. 12276-24885, 2023.
Jangho Kim, ScongUk Park, and Nojun Kwak. Paraphrasing complex network: Network compression via factor transfer. Advances in neural information processing systems, 31, 2018.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. Technical Report, University of Toronto, 2009.
Matthias Langer, Zhen He, Wenny Rahayu, and Yanbo Xue. Distributed training of deep learning models: A (axonomic perspective. IEEE Transactions on Parallel and Distribuled Systems, (1(1))22802 2818, 2020.
Soung Hyun Lee, Dae Ha Kim, and Byung Cheol Song. Self-supervised knowledge distillation using singular value decomposition. In Proceedings of the European conference on computer vision (ECCV), pp. 335-350, 2018.
Chenxin Li, Mingbao Lin, Zhiyuan Ding, Nie Lin, Yihong Zhuang, Yue Huang, Xinghao Ding, and Liujuan Cao. Knowledge condensation distillation. In European Conference on Computer Vision, pp. 19-35. Springer, 2022.
Lijiang Li, Huixia Li, Xiawu Zheng, Jie Wu, Xueleng Xiao, Rui Wang, Min Zheng, Xin Pan, Fei Chao, and Rongrong Ji. Autodiffusion: Training-free free optimization of time steps and architectures for automated diffusion model acceleration. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 7105-7114, 2023.
Ji Lin, Wei-Ming Chen, Yujun Lin, Chuang Gan, Song Han, et al. Mcunet: Tiny deep learning on iot devices. Advances in Neural Information Processing Systems, 33:11711-11722, 2020.
Ji Lin, Wei-Ming Chen, Han Cai, Chuang Gan, and Song Han. Mcunetv2: Memory-effiiien patch-based inference for tiny deep learning. arXiv preprint arXiv:2110. 15352, 2021.
Ji Lin, Ligeng Zhu, Wei-Ming Chen, Wei-Chen Wang, Chuang Gan, and Song Han. On-device training under 256kb memory. Advances in Neural Information Processing Systems, 53:22941-22054 2022.
Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In Proceedings of International Conference on Computer Vision (ICCV), December 2015.
18

---

Published in Transactions on Machine Learning Research (04/2025)
Jonathan Long, Evan Shelhamer, and Trevor Darrell. Fully convolutional networks for semantic segmentation. In Proceedings of the IEEF conference on computer vision and poltern vecognition, pp. 3431-3440,, 2015.
Yiheng Lu, Ziyu Guan, Yaming Yang. Wei Zhao, Maoguo Gong, and Cai Xu. Entropy induced pruming framework for convolutional neural networks. In Proceedings of the AAAT Conference on Artificial Intelligence, volume 38, pp. 3918-399199 2024.
Pavlo Molchanov, Stephen Tyree, Tero Karras, Timo Aila, and Jan Kautz. Pruning convolutional neural networks for resource efficient inference. International Conference on Learning Representations, 2017.
Wonpyo Park, Dongju Kim, Yan Lu, and Minsu Cho. Relational knowledge distillation. In Proceedings of the IEEE/CVF conference on computer vision and paltern recognilion, pp. 3967- 3976, 2019.
Nikolaos Passalis and Anastasios Tefas. Learning deep representations with probabilistic knowledge transfer. In Proceedings of the European Conference on Compuler Vision (ECCV), pp. 284, 2018.
William Peebles and Saining Xie. Scalable diffusion models with transformers. Proceedings of the IEEE/CVF international conference on computer vision, pp. 4195 4205, 2023.
Joseph Redmon and Ali Farhadi. Yolov3: An incremental improvement. arXiv preprint arXiv:1804.0276; 2018.
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In Medical image computing and computer-assister intervention-MICCA. 2015: 18th international conference, Munich, Germany, October 5-9, 2015, proceedings, part III 18, pp. 234-241. Springer, 2015.
Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, Alexander C. Berg, and Li Fei-Fei. ImageNet Large Scale Visual Recognition Challenge. International Journal of Computer Vision (IJCV), 115(3):211-252, 2015. doi: 10. /s1 1263-01--0816-y
Tim Salimans, Ian Goodfellow, Wojeiech Zaremba, Vicki Cheung, Alec Radford, and Xi Chen. Improved techniques for training gans. Advances in neural information processing systems, 29, 2016.
Mark Sandler, Andrew Howard, Menglong Zhu, Andrey Zhmoginov, and Liang-Chiel Chen. Mobilenet
Inverted residuals and linear bottlenecks. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 4510-4520, 2018.
Axel Sauer, Dominík Lorenz, Andreas Blattmann, and Robin Rombach. Adversarial diffuusion distillation.
Xrxiv preprinl 1xX2:2231..10042, 2023.
Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recognition. arXiv preprint XrXiv:1409, 1556, 2014.
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. ar preprinl 200002200,, 2020.
Shangquan Sun, Wenqi Ren, Jingzhi Li, Rui Wang, and Xiaochun Cao. Logit standardization in knowledge distillation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp- 15731 15740, 2024.
Yonglong Tian, Dilip Krishnan, and Phillip Isola. Contrastive representation distillation. In International Conference on Learning Representations, 2020.
torchao maintainers and contributors. torchao: Pytorch native quantization and sparsity for training and inference, October 2024. URL
Frederick Tung and Greg Mori. international conference on computer vision pp. 1365-1374 2019.
knowledge distillation. In Proceedings of the IEEE/CVF
19

---

Published in Transactions on Machine Learning Research (04/2025)
Kan Wu, Jinnian Zhang, Houwen Peng, Mengchen Liu, Bin Xiao, Jianlong Fu, and Lu Yuan. Tinyvit: Fast pretraining distillation for small vision transformers. In European conference on computer vision, pp. 68 85. Springer, 2022.
Saining Xie, Ross Girshick, Piotr Dollár, Zhuowen Tu, and Kaiming He. Aggregated residual transformations for deep neural networks. In Proceedings of the IEEE conference computer vision and pottern recognition, pp. 1492 1500, 2017.
Tianwei Yin, Michaël Gharbi, Richard Zhang, Eli Shechtman, Fredo Durand, William Freeman, and Taesung Park. One-step diffusion with distribution matching distillation. arviv preprind 23x:2311..8822. 2023.
Sergey Zagoruyko and Nikos Komodakis. Paying more attention to attention: Improving the performance of convolutional neural networks via attention transfer. arXiv preprint arXiv:1612 03928, 2016.
Borui Zhao, Quan Cui, Renjie Song, Yiyu Qiu, and Jiajun Liang. Decoupled knowledge distillation. In Proceedings of the IEEE/CVE Conference on compuler vision and pattern recognition, pp. 11953 11962, 2022.
20


Reviews:
[{"review_id": "PAAhevWeex", "reviewer": "Reviewer_TGXm", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n\nThe RED student models demonstrate top performance across various computer vision tasks.\n\nThe manuscript is well written and easy to follow.\n\n\nWeaknesses\n\nTable 6 indicates that in some cases, the actual peak memory of RED-based models exceeds that of other student techniques.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes a novel model distillation method for computer vision tasks that emphasizes minimizing peak memory consumption. The approach employs aggressive pooling combined with newly proposed RED blocks and is optimized using a dedicated loss function. The experimental results suggest that RED-based distilled models achieve competitive accuracy while maintaining bounded peak memory usage.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "It is recommended that the authors provide more detailed memory consumption metrics—both theoretical estimates and empirical measurements—for all compared techniques (e.g., NST, DeiT). This clarification could be incorporated into Table 6 or presented in an additional section to better contextualize the comparative performance.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "s6oi070W4y", "reviewer": "Reviewer_fHB7", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n1. The authors introduce a new component, the RED block, that significantly improves the performance of a student with lower memory requirements (4x in theory, 2x on a NVIDIA Jetson).\n\n2. The presented method shows improvements across both image classification and generation.\n\n3. In the case of classification (but not generation), the authors compare their method with many other established distillation methods.\n\nWeaknesses:\n\n1. Section 4.4 \"ViT-based Image Classification\": since the RED module \"is designed to ensure that the output of a pooling layer closely resembles the distribution of the input while preserving essential features at a reduced spatial dimension\", it is not clear why it is even applied on the ViT architecture. In the case of the ViT, there is no such reduced spatial dimension, from what I can tell, and also is there no reduction in theoretical peak memory requirement. It also seems contra-intuitive to train a base-sized transformer on STL10 (only 5k training examples!), but to use the tiny-sized transformer for ImageNet (with 1.2M training examples). Note that [(Touvron, 2020)] reports 74.5 for ViT-Ti/16 distilled, and 72.2 for ViT-Ti/16 from scratch in table 5 (vs. 72.5 reported here for both distillation with and without RED).\n\n2. I do appreciate the comparison with many other distillation methods, but I would really like to see the results anchored in some previously published baselines. From what I can tell, tables 1-3 don't report numbers from other papers, and table 5 reports a FID=7.23 on CIFAR10 for the teacher network, but when looking at [(Ho, 2020)] Table 1 I find FID=3.17 (note that the text in Section 4.5 says \"results slightly different from those in the original paper\", but this difference is unacceptably large, assuming that I'm indeed comparing to the correct value).\n\n3. The authors mention related work for \"memory-constrained deep learning\", such as quantization, but then the authors only compare existing distillation methods in their setup (with more aggressive pooling). I could only find a short mention of QAT in Section 4.7, and Table 11 mentions model size (which is reduced by only 10%, so I wonder what kind of quantization was used), but not peak memory requirement.\n\n4. Even though the paper is well written overall, I found some of parts hard to parse, and some important details missing (see requested changes below).\n\n[(Touvron, 2020)]: https://arxiv.org/pdf/2012.12877\n[(Ho, 2020)]: https://arxiv.org/abs/2006.11239", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors introduce a new method called \"residual encoded distillation\" (ReDistill) that allows distilling a CNN teacher network into a CNN student network with 4-5x reduced theoretical peak memory requirements, with the goal to make the models more suitable for deployment on memory-constrained devices.\n\nTheir method is based on a more aggressive initial pooling, and the addition of a \"residual encoded distillation (RED)\" block for feature distillation.\n\nThe authors claim improvements for image classification on ImageNet and STL10, as well as for image generation for CIFAR10 and Celeb-A (using DDPM).", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Critical changes:\n\n1. Better ablations: I would be interested to see how different RE kernel sizes perform, and similarly for alpha values different from 0, 1, and 50 (is 50 optimal? does the same optimal value work equally well for every architecture and pooling reduction factor?).\n\n2. Explain how exactly the image classification models were configured. I could find no details on the \"aggressive pooling setting\", other than the \"first pooling layer stride ×2 ∼ ×8 times and adjust the last several pooling layer strides to ensure the final output of the student model with the same information density as the teacher model\" (Section 4.2). How exactly were the networks configured? I thought (after reading Section 3.2 and Figure 2) that the pooling layers were all the same, other than the initial pooling layer that incorporates multiple pooling layers from the teacher architecture. Also, it's not clear what the layers correspond to in Figure 6 (e.g. there are only 9 layers shown for ResNet18, and 17 layers for ResNet50; also, every layer has lower memory footprint, even though the previous description only talks about memory footprint reduction in the initial layers). Why the decision to look at \"x4\" variants for the classification of the experiments (\"x2\" is only mentioned once in Tables 1-3) and \"x2\" for the generation experiments? In this context it would also be helpful to explain how the peak memory reduction is computed?\n\n3. Fix the discrepancy of the CIFAR10 FID in Table 5 (7.23 vs 3.17 mentioned in [(Ho, 2020)]).\n\n4. Compare ReDistill with other distillation methods in image generation (Section 4.5), or explain why this is not possible.\n\n5. I find the values in Table 3 surprising: Why is ReDistill compared to a set of distillation methods different from Tables 1 and 2? How come that many of the reported values are *worse* than the student architecture without distillation? Please explain this better in the main text, or remove the table.\n\nMedium importance changes:\n\n1. I think the presented method in its current form only makes sense for CNNs, so I think it would help the reader and add clarity to mention \"convolution\" already in the title.\n\n2. I don't find that the section about ViTs adds anything to the paper. If the authors want to keep the section, then I would like to see my points above (under weaknesses) addressed.\n\n3. I find it confusing that the RED block (in Figure 2) has $f^T$ as an *input*, since it is only used to compute the loss. I would recommend having the RED block with simply $f^S$ and $f^D$ as input and output (ideally using some notation that makes it clear that they have the same dimensions), and then have an arrow going from $f^D$ to box $L_{RED}$ (with another arrow coming from $f^T$).\n\n4. I find Figure 2 hard to parse: Why does it list stages 1,k,k+1,n for the teacher and stages 1,2,i,i+1 for the student? From the text I believe to understand that only stages 1..k-1 are skipped and from stage k on, every stage is found in both teacher and student (so no need to list stage n). Also, the text says that teacher and student stages have the same spatial dimensionality in the output, but might be different wrt number of channels – but the figure is exactly the other way around (i.e. illustration of teacher and student feature maps have the same number of channels, but different size).\n\n5. Explain why applying QAT (Table 11) only saves 10% in model size and compare QAT with the presented method in terms of peak memory reduction – theoretical, and ideally also practical. I think that comparing ReDistill with other methods mentioned in related work \"memory-constrained deep learning\" would make the submission more practically relevant.\n\n[(Ho, 2020)]: https://arxiv.org/abs/2006.11239\n\n\nMinor comments:\n\n1. In the introduction, the authors motivate their approach with a STM32H5 MCU (640 kb of RAM), but then in the practical experiments, they run the networks on a NVIDIA Jetson TX2 (8GB of RAM) – it seems that the presented method is not at all applicable to the kind of edge device mentioned in the introduction.\n2. \"However, we still keep the same spatial dimensions of the input activations to the final fully connected layer.\" (section 3.2) – what exactly is kept the same here? Is it the spatial dimensions of the inputs (i.e. the input image), or the spatial dimensions of the teacher (which is the input to the student during distillation). Please clarify.\n3. \"Homogeneously\" (section 3.2) – consider writing \"similarly\", or \"correspondingly\" instead\n4. Equations (5,6,7) – why not reuse symbols from Figure 3 and text? i.e. $f^R=r(·)$ and $\\hat{f}^S=l(·)$.\n5. Table 6 mentions MbNetV2x8 but table Table 2 only has results for MbNetV2x4 – maybe a typo?", "broader_impact_concerns": "None.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "E7gOhxTtfK", "reviewer": "Reviewer_NDk8", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- Reducing peak memory is a good research topic considering the memory-bound scenario of current large-scale transformer-based model deployment.\n- The presentation quality is satisfying. The target and architecture of the proposed ReDistill framework are introduced clearly.\n- Comprehensive experiments are conducted on various models and tasks. Especially the image generation with diffusion model is not included in most previous work but is highlighted in this paper.\n- Edge device deployment results align with the theoretical peak memory analysis, improving the credibility of results and robustness of this work.\n\nWeakness:\n- KD baselines in the experiments are somewhat outdated, please consider including more recent KD baselines in the experiments such as [1] and [2].\n\n[1] One-for-All: Bridge the Gap Between Heterogeneous Architectures in Knowledge Distillation, NeurIPS 2023\n\n[2] Logit Standardization in Knowledge Distillation, CVPR 2024", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposed residual encoded distillation (ReDistill) for peak memory reduction and efficient deployment of large-scale models on resource-constrained devices. While most previous knowledge distillation work focuses on model size compression or improving the performance of a given model with a larger one, ReDistill is the first targeting peak memory reduction. ReDistill derives the lightweight student model from the teacher network using aggressive pooling and a residual encoded distillation (RED) block to align features between teacher and student. The proposed ReDistill is extensively verified on various architectures (CNNs, ViTs, Diffusion Models) and tasks (classification, generation), demonstrating effectiveness and efficiency.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- Some typos and grammar errors could be fixed: \n  - Sec.1 ...a less accuracy drop $\\rightarrow$ ...a lower accuracy drop\n  - Sec.4.2 all models 250K iterations $\\rightarrow$ all models with 250K iterations\n  - Tab.6 Actural $\\rightarrow$ Actual, 2.8 $\\rightarrow$ 2.80\n  - ...", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "QvW4ad0Gun", "reviewer": "Reviewer_zkR4", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "## Strengths\n\n1. The study of peak memory reduction is meaningful in practice. The proposed approach, based on aggressive downsampling, is novel and intuitive.\n\n2. The proposed approach achieves significant peak memory reduction.\n\n3. Both image classification and generation tasks are considered and evaluated.\n\n## Weaknesses\n\n1. Overhead of RED modules: As shown in Table 6, students with RED may have up to 2x model size, 2x power consumption, and 3x latency, compared with the models without RED. Although the main focus of this work is the peak memory, overheads in other aspects should also be considered.\n\n2. Low-resolution image datasets: Many of the datasets tested in this work are a bit old and contain low-resolution images, such as STL, CIFAR-10, and Celeb-A. They may no longer be representative of modern mobile applications. It would be better to train and evaluate models on more modern datasets such as iNaturalist (for classification) and CelebA-HQ (for generation). Also, high-resolution images may more significantly show the advantage of RED in reducing the peak memory.\n\n3. Unintuitive module design: Although the reviewer can understand the RED architecture (Section 3.3) and find its effectiveness in the ablation study (Table 7), the model design still seems to be arbitrary. Why is the residual activation function ReLU6 (not regular ReLU)? Why does one convolution layer have a 3x3 kernel while the other has 1x1? Why is BatchNorm employed in RED rather than other normalization layers (e.g., LayerNorm seems more conventional when there is a vision transformer). A few sentences that explain the authors’ intuition may be helpful in the paper.\n\n4. Effectiveness in ViT-based image classification: Table 4 only shows the accuracy comparison. Without details about computation resource consumptions, this section seems incomplete. The reviewer agrees that it might be good to investigate peak memory reduction for transformers in the future.\n\n## Minor suggestions\n1. In Section 3.2, the dimension of the teacher feature map at the $k$-th pooling layer should be $\\frac{H}{\\rho^k}\\times\\frac{W}{\\rho^k}\\times C_k$. The student’s initial stride and layer indices may also need to be revised.\n\n2. Citing format: When the referred work appears as a part of the sentence, it might be better to drop the parenthese. For example in Section 3.1:\n    - “Reference (Hinton et al. (2015)) defines…” - > “Hinton et al. (2015)) define”\n    “specifically a U-Net network as in (Ho et al. (2020))” -> “specifically a U-Net network as in Ho et al. (2020)”", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper proposes a knowledge distillation method Residual Encoded Distillation (ReDistill) to reduce the peak memory of visual models during inference, specifically for resource-constrained edge devices. This work uses aggressive downsampling in the student network’s first pooling layer and introducing Residual Encoded Distillation (RED) blocks that align feature maps between the teacher and student networks. The proposed method is evaluated on tasks such as image classification and diffusion-based image generation, achieving significant reductions in peak memory usage with minimal performance degradation.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Please see the Weaknesses and Minor suggestions above. In particular, it is encouraged to include some more explanation about the RED computational overhead, and results on newer datasets.", "broader_impact_concerns": "The reviewer does not have any broader impact concerns.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: All reviewers were positive about the work in the post-rebuttal phase, and they especially liked the proposed module and distillation method. Besides, the authors have provided extensive new experiments in response to the reviewers' suggestions. Hence, the decision is to recommend the paper for acceptance with a minor revision for incorporating the new experiments.
Decision: Accept (poster)