# ReDistill: Residual Encoded Distillation
for Peak Memory Reduction

Fang Chen
University of California, Merced
<EMAIL_ADDRESS>
Gourav Datta
University of Southern California
Mujahid Al Rafi
University of California, Merced
Hyeran Jeon
University of California, Merced
Meng Tang
University of California, Merced
<EMAIL_ADDRESS>
Corresponding author

###### Abstract

The expansion of neural network sizes and the enhancement of image resolution
through modern camera sensors result in heightened memory and power demands
for neural networks. Reducing peak memory, which is the maximum memory
consumed during the execution of a neural network, is critical to deploy
neural networks on edge devices with limited memory budget. 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
2x-3.2x measured peak memory on an edge GPU with negligible degradation in
accuracy for most CNN based architectures. Additionally, our method yields
improved test accuracy for tiny vision transformer (ViT) based models
distilled from large CNN based teacher 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.

##  1 Introduction

Convolutional neural networks (CNN) and vision transformers (ViT) have
demonstrated impressive capabilities across diverge computer vision tasks such
as image recognition [49, 51, 15], object detection [11, 35, 43], semantic
segmentation [37, 44, 14], and image generation [8, 19, 50] etc. 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 it
impractical to execute off-the-shelf deep learning models: ResNet-50 surpasses
the storage limit by $44\times$, while MobileNetV2 exceeds the peak memory
limit by $8\times$. Even the int8 quantized version of MobileNetV2 still
surpasses the memory limit by $2\times$, underscoring a substantial disparity
between desired and available hardware capacity. Hence, it is very important
to reduce the peak memory which is the maximum memory consumed by a neural
network during inference, for edge deployment. Through empirical measurements,
we determined that peak memory usage is predominantly influenced by the
initial layers of convolutional neural networks (CNNs) that are characterized
by large feature maps. For U-shaped CNN architectures, however, the last few
layers also significantly contribute to peak memory consumption.

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 $4{\times}$ lower
theoretical peak memory generates images indistinguishable from the generated
images of a teacher network.

A naive approach to reducing peak memory is aggressive downsampling via
pooling with large kernel size and large stride, which often leads to
unacceptable degradation of network performance due to loss of information in
small feature maps. Given a teacher network with large peak memory, we propose
residual encoded distillation (ReDistill) to train a 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-based models shown in Fig. 1(a), our method reduces the
theoretical peak memory by $4\times{-}5\times$ with a negligible accuracy drop
compared to the baseline teacher network. Additionally, we also show improved
distillation from CNNs to transformers. For diffusion-based image generation
shown in Fig. 1(b), our distilled network generates images extremely similar
to original networks, yet the theoretical peak memory is reduced by $4\times$
on average.

Our ReDistill method outperforms existing response-based or feature-based
distillation methods [18, 1, 58, 41, 29, 40, 2, 54, 5, 30, 52, 25] in terms of
accuracy-memory trade-off. Our method also 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 a large kernel size and stride in the initial pooling layers
and have 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-linear mapping modules termed
residual encoded distillation (RED) blocks between the teacher and student
network during both training and inference. Thirdly, our proposed RED block is
lightweight and effective with additive residual learning and multiplicative
gating mechanism, which adds little computational overhead to student network.
Lastly, we align teacher and student network features asynchronously at layers
with matching feature sizes, while previous approaches align features at
different stages of the networks.

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 a large stride for a
student network, while incurring a negligible 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 a residual encoded distillation (RED)
block to align features between high-peak-memory teacher networks and low-
peak-memory student networks. Our RED block is based on a multiplicative
gating mechanism and additive residual learning and is shown simple and
effective for peak memory reduction with minimum computational overhead.

  * •

For image classification tasks, our distillation method outperforms state-of-
the-art response-based or feature-based distillation methods as shown in
extensive experiments with multiple datasets. Our method yields about
$2\times\sim 3.2\times$ reduction in measured peak memory with a slight
decrease in the classification accuracies for CNN based models. Additionally,
our method improves the accuracy of compact ViT based models, when distilled
from large CNNs.

  * •

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
$4{\times}$ by downsampling the feature maps of the first few encoder layers
and last few decoder layers while maintaining the fidelity and the diversity
of synthesized images.

##  2 Related Work

Memory-constrained deep learning: Limited memory capacity in GPU cards and
edge platforms has been a critical hurdle in CNN training and inference.
Multiple GPUs can be utilized through model and data parallelism [28] to
mitigate the memory bottleneck. Other solutions include optimization methods
such as network quantization [22], compression [13], and pruning [39, 38]
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 [12]. There are also CNNs specifically
designed for resource-constrained applications such as variants of MobileNet
[20] and SqueezeNet [24]. These approaches to memory-constrained deep learning
are orthogonal to our ReDistill framework, which focuses on peak memory
reduction during inference. Nevertheless, a recent work has explored neural
architecture search (NAS) to create networks with minimized peak memory [33].
Reference [32] takes this a step further by leveraging NAS to introduce patch-
based inference and network redistribution [32], consequently shifting the
receptive field to later stages. While NAS significantly exacerbates the
training complexity, patch-based inference necessitates compiler libraries
that may not be compatible with standard GPUs, and incurs additional
computation and latency overhead. Another recent work [4] proposed self-
attention-based 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-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, [18] distills knowledge by matching the prediction probability
distributions of the student architecture and the teacher architecture. DKD
[59] 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 [25] performs logit
distillation through a multi-level alignment based on instance prediction,
input correlation, and category correlation, delivering state-of-the-art
performance. In contrast, feature-based KD methods [1, 58, 41, 29, 54, 2]
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 [40] employs a relation potential function to convey information
from the teacher’s features to the student’s features. ReviewKD [5] aggregates
knowledge of the teacher architecture from different stages into one stage of
the student architecture, the so-called ’knowledge review’, which achieved
impressive performance. KCD [30] iteratively condenses a compact knowledge set
from the teacher to guide the student learning by the Expectation-Maximization
(EM) algorithm, which would empower and be easily applied to other knowledge
distillation algorithms. Thus, existing methods focus on the knowledge
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 a 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 has gained popularity in the realm of diffusion models.
For example, One Step Diffusion [57] defines two score functions, one of the
target distribution and the other of the synthetic distribution being produced
by a one-step generator. By minimizing the KL divergence between these two
score functions, the one-step image generator is enforced to match the
diffusion model at the distribution level and achieves impressive performance.
Adversarial Diffusion [48] utilizes both score distillation loss and
adversarial loss. The score distillation loss occurs between the teacher
diffusion sampler with a 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 [31] 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 on time step reduction and model compression. Our method offers a unique
approach by minimizing the size of intermediate activation maps, which is
orthogonal to and can be easily integrated with existing methods.

##  3 Proposed Method

###  3.1 Preliminaries

Knowledge Distillation In the task of knowledge distillation, we are given a
dataset $\mathcal{X}$, a high-capacity teacher architecture $\mathcal{T}$ and
a to-learned student architecture $\mathcal{S}$. Specifically, for an input
image $x$ sampled from $\mathcal{X}$, $\pi_{\mathcal{T}}(x)$ and
$\pi_{\mathcal{S}}(x)$ denote the outputs or intermediate features of the
teacher and the student, respectively. Generally, the knowledge distillation
task aims to optimize the student architecture’s parameters $\hat{w}$:

$\displaystyle\hat{w}=\operatorname*{arg\,min}_{w}\sum_{x\in\mathcal{X}}{\mathcal{L}(\pi_{\mathcal{S}}(x;w),\pi_{\mathcal{T}}(x))},$
(1)

where $w$ denotes the trainable weights of $\pi_{\mathcal{S}}$ and
$\mathcal{L}$ denotes the loss function defined by different knowledge
distillation methods. For instance, [18] defines $\pi_{\mathcal{S}}$ and
$\pi_{\mathcal{T}}$ as the logit outputs (without applying the softmax
function) of the student and the teacher, while $\mathcal{L}$ as the Kullback-
Leibler divergence between $\pi_{\mathcal{S}}$ and $\pi_{\mathcal{T}}$ after
applying the softmax function with temperature $t_{p}$:

$\displaystyle\mathcal{L}_{KL}=\quad\sum_{x\in\mathcal{X}}{KL(softmax(\frac{\pi_{\mathcal{S}}(x)}{t_{p}}),softmax(\frac{\pi_{\mathcal{T}}(x))}{t_{p}})}.$
(2)

[58, 40, 5] define different $\pi_{\mathcal{S}}$ and $\pi_{\mathcal{T}}$, such
as the intermediate activation maps from various stages of the student and
teacher, or different $\mathcal{L}$, like the $p$-norm, to achieve various
distillation methods.

Denoising Diffusion Probabilistic Models Diffusion models [19] are latent
variable models of the form $p_{\theta}(x_{0}):=\int
p_{\theta}(x_{0:T})d{x_{1:T}}$, where $x_{1},...,x_{T}$ are latents of the
same dimensionality as the data $x_{0}\sim q(x_{0})$. The joint distribution
$p_{\theta}(x_{0:T})$ is defined as a Markov chain with learned Gaussian
transition starting at $p(x_{T})=\mathcal{N}(x_{T};\textbf{0},\textbf{I})$,
where $T$ is the maximum time step. In the training process, we are given a
noisy input $x_{t}$ which is derived from the data $x_{0}$ and noise
$\epsilon\sim\mathcal{N}(\textbf{0},\textbf{I})$:

$\displaystyle
x_{t}=\sqrt{\bar{\alpha_{t}}}x_{0}+\sqrt{1-\bar{\alpha_{t}}}\epsilon,$ (3)

where $\bar{\alpha}_{t}:=\prod_{s=1}^{t}(1-\beta_{s})$, and $\beta_{s}$ is the
forward process variances fixed as constant in DDPM [19]. The loss of
diffusion model is generally defined as follows:

$\displaystyle\mathcal{L}_{diff}=||\epsilon-\epsilon_{\theta}(x_{t},t)||_{2}^{2}.$
(4)

The noisy input $x_{t}$, accompanied by a time step embedding $t$, is input
into a denoising autoencoder, specifically a U-Net network as in [19], to
estimate the noise component $\epsilon_{\theta}(x_{t},t)$.

###  3.2 Residual Encoded Distillation Block

Figure 2: Residual Encoded Distillation (RED) Block. We use a logit module for
the multiplicative gating mechanism and a residual encoder for additive
residual learning.

The proposed Residual Encoded Distillation (RED) Block, depicted in Fig. 2,
serves as the central module of our framework. It 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. To
accomplish this, the RED Block is 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
a 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:

$\displaystyle f^{\mathcal{D}}$ $\displaystyle=$ $\displaystyle
r(f^{\mathcal{S}})+f^{\mathcal{S}}*l(f^{\mathcal{S}}),$ (5) $\displaystyle
r(\cdot)$ $\displaystyle=$ $\displaystyle ReLU6(BN(Conv_{3\times 3}(\cdot))),$
(6) $\displaystyle l(\cdot)$ $\displaystyle=$ $\displaystyle
Sigmoid(BN(Conv_{1\times 1}(\cdot))),$ (7)

where $f^{\mathcal{S}}$ is the feature map from the student model. $l(\cdot)$
denotes the logit module consisting of a $1\times 1$ convolution layer, a
batch norm layer, and a sigmoid activate function. The logit module generates
element-wise weights $\hat{f}^{\mathcal{S}}$ like a gate to suppress non-
significant components of $f^{\mathcal{S}}$. $r(\cdot)$ denotes the residual
encoder module consisting of a $3{\times}3$ convolution layer, a batch norm
layer, and a relu-6 activation function. The residual encoder module yields
the residual item $f^{\mathcal{R}}$. 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 compensated by the residual
item $f^{\mathcal{R}}$.

###  3.3 Proposed Distillation Framework

Figure 3: 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.

Our proposed framework is illustrated in Fig. 3. The teacher and student have
the same convolutional neural network architecture with different pooling
strategies. In order to reduce the activation peak memory, the initial pooling
layer of the student is assigned with a larger pooling stride. However, we
still keep the same spatial dimensions of the input activations to the final
fully connected layer. Thus, the student has fewer pooling layers but with a
larger pooling stride at the initial pooling layer.

Take an input image $x\in\mathbb{R}^{H\times W\times C}$ as the example. The
teacher and the student network are divided into several stages by pooling
layers as shown in Fig. 3. We assume that all pooling layers of the teacher
have the same pooling stride $\rho$ for simplicity. The initial pooling layer
of the student at stage $1$ is assigned with a stride $k\rho$. Hence, the
student feature map
$f^{\mathcal{S}}_{1}\in\mathbb{R}^{\frac{H}{k\rho}\times{\frac{W}{k\rho}\times
C_{1}}}$ at the output of this pooling layer has the same spatial dimension as
the teacher feature map
$f^{\mathcal{T}}_{k}\in\mathbb{R}^{\frac{H}{k\rho}\times{\frac{W}{k\rho}\times
C_{k}}}$ at the output of the $k$-th pooling layer. Then the matched feature
maps $f^{\mathcal{S}}_{1}$ and $f^{\mathcal{T}}_{k}$ will be fed into the
residual encoded distillation (RED) block detailed in Section 3.2 and
illustrated in Fig. 2, to compute the distillation loss. The output
$f^{D}_{1}$ is to be fed into the following layers of the student network.

Homogeneously, the feature map $f^{\mathcal{S}}_{2}$ of the student at stage
$2$ also has the same spatial dimension as the feature map
$f^{\mathcal{T}}_{k+1}$ of the teacher at stage $k+1$. They are fed into
another RED block to calculate distillation loss and output $f^{D}_{2}$ to the
following layers. This process is repeated until the last pooling layer
$p_{\mathcal{T}}(\circ;\rho)$ at stage $n$ of the teacher. We assume the size
of the output feature map $f^{\mathcal{T}}_{n}$ of the teacher is identical to
the output feature map $f^{\mathcal{S}}_{i}$ of the student at stage $i$,
where $i-1+k=n$. The following pooling layers of the student at stages
$i+1,i+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.

###  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 and
the RED block’s output. The loss function of a specific RED block is
formulated as follows:

$\displaystyle\mathcal{L}_{RED}(f^{\mathcal{T}},f^{\mathcal{D}})=d_{cos}\\{\frac{\sum_{c\in
C_{\mathcal{T}}}{f^{\mathcal{T}}_{c}}}{|C_{\mathcal{T}}|},\frac{\sum_{c\in
C_{\mathcal{S}}}{f^{\mathcal{D}}_{c}}}{|C_{\mathcal{S}}|}\\},$ (8)

where $f^{\mathcal{D}}$ is calculated by Equation 5, and $d_{cos}$ denotes the
cosine distance function. $C_{\mathcal{T}}$ and $C_{\mathcal{S}}$ denote the
channel dimension of the teacher’s feature map $f^{\mathcal{T}}$ and RED
block’s output $f^{\mathcal{D}}$, respectively. The final loss function of the
proposed method is as follows:

$\displaystyle\mathcal{L}=\alpha\mathcal{L}_{task}+(1-\alpha)\sum_{i=1}^{I}{\beta\mathcal{L}_{RED_{i}}},$
(9)

where $\mathcal{L}_{task}$ is the vanilla loss from the task, such as the
Binary-Cross-Entropy (BCE) loss for image classification. $I$ denotes the
number of RED blocks. $\alpha$ is the hyper-parameter to adjust the
task/distill-oriented loss. $\beta$ is a hyper-parameter to scale RED loss.

###  3.5 Distillation for Diffusion Model

Figure 4: ReDistill for denosing network in DDPM [19]. 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.

In this section, we introduce how to integrate the proposed distillation
framework into a U-Net based denoising network described in DDPM [19], as
shown in Fig. 4. For convenience, let’s assume the teacher model is a 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.e., it just has one down-sample
layer with the stride $4$. We use the output $f_{down}^{\mathcal{S}}$ from the
down-sample layer of the student model and the output $f_{down}^{\mathcal{T}}$
from the second down-sample layer of the teacher model, as they share the same
spatial dimension. $f_{down}^{\mathcal{S}}$ and $f_{down}^{\mathcal{T}}$ are
input into a RED block incorporated into the student model, producing the
output $f_{down}^{\mathcal{D}}$, which is then fed into the subsequent
convolutional layers. Symmetrically, the student model has one up-sample layer
with an expansion ratio of $\times 4$, while the teacher model has two up-
sample layers, each with an expansion ratio of $\times 2$. The input
$f_{up}^{\mathcal{S}}$ of the student’s up-sample layer and the input
$f_{up}^{\mathcal{T}}$ of the teacher’s first up-sample are input into another
RED block, producing the output $f_{up}^{\mathcal{D}}$. This output replaces
the $f_{up}^{\mathcal{S}}$ as the new input of the student’s up-sample layer.
For DDPM distillation, the loss function is defined as Equation 9 while the
$\mathcal{L}_{task}$ is replaced by $\mathcal{L}_{diff}$ defined in Equation
4.

Table 1: Results on ImageNet dataset. The $3^{rd}$ and $4^{th}$ columns show
the results with identical teacher and student architectures, while the
$5^{th}$ column shows the result with different teacher and student
architectures. Top-1 accuracy (%) on testset is used as the evaluation metric.

Method | Teacher | ResNet18 | ResNet50 | ResNet152 | Avg
---|---|---|---|---|---
Acc. (%) | 69.75 | 76.13 | 78.32
Student |  ResNet18$\times 4$ |  ResNet50$\times 4$ |  MbNetV2$\times 2$
Acc. (%) | 61.79 | 69.50 | 62.65
Response |  KD [18] | 63.63 | 70.60 | 62.85 | 65.69
DKD [59] | 63.22 | - | 66.27 | 64.75
MLLD [25] | 64.66 | 70.77 | 68.36 | 67.93
Feature |  FitNet [1] | 62.13 | 71.77 | 60.79 | 64.90
RKD [40] | 61.49 | 66.88 | - | 64.69
ReviewKD [5] | 63.30 | 70.22 | 63.07 | 65.53
CRD [52] | 64.01 | 71.07 | 65.60 | 66.89
RED (ours) | 65.23 | 73.23 | 68.89 | 69.12

‘-’ denotes we do not get reasonable results for the student architecture with
the distillation method under the aggressive pooling setting.

##  4 Experiments

We conduct experiments on various vision tasks to illustrate the effectiveness
of the proposed method with the state-of-the-art distillation methods under
different backbone architectures and datasets. While our evaluation on CNN
based image classification and diffusion model based image generation tasks
are shown below, our evaluation with ViTs are shown in supplementary material
A.5. Datasets used and implementation details are provided in supplementary
material A.1 and A.2 respectively.

###  4.1 Image Classification

Table 1 shows the results on the ImageNet [45] dataset, with the setting that
the teacher model and student model are in identical or different
architectures. ‘$\times n$’ denotes we increase the $1^{st}$ 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 knowledge distillation, as shown in the
$3^{rd}$ and $4^{th}$ columns of Table 1, or different architecture knowledge
distillation, as shown in the $5^{th}$ column of Table 1. We note that some
distillation methods fail to get better metrics even than the student model.
The reason might be that these distillation methods are not designed for the
aggressive pooling setting. For example, some distillation methods would be
sensitive to the resolution of the feature maps, like [1], or require multi-
scale feature maps, like [40]. In our aggressive pooling setting, the student
model has lower activation resolution and fewer pooling layers, i.e., limited-
scale feature maps, which would impede the performance of these traditional
distillation methods. We conduct additional experiments of image
classification on the high resolution STL10 [7] dataset, that are shown in
supplementary material A.3.

###  4.2 Image Generation

Table 2: Results on U-Net [44] based DDPM [19]. Dataset | Method | IS | FID
---|---|---|---
CIFAR10 | T: U-Net w/ DDPM | 9.06 $\pm$0.10 | 7.23
S: U-Net$\times 2$ w/ DDPM |  7.93 $\pm$ 0.09 | 21.49
S w/ RED (ours) | 8.53 $\pm$ 0.07 | 12.88
Celeb-A | T: U-Net w/ DDPM | 2.97 $\pm$ 0.04 | 19.61
S: U-Net w/ DDPM$\times 2$ |  2.77 $\pm$ 0.02 | 23.03
S w/ RED (ours) | 2.88 $\pm$ 0.03 | 21.43

#### Quantitative Results

Table 2 shows the results on CIFAR10 [27] dataset and Celeb-A [36] dataset.
Our method reduces the fidelity degradation of the student model with a first
pooling stride that is twice that of the teacher model. Specifically, our
method achieves $8.61$ lower FID and $0.6$ 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 dataset, respectively. Due to implementation
specifics, the models are not augmented with EMA [23], leading to results
slightly different from those in the original paper.

#### 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
$\epsilon_{t}$ 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. 5, generally, the images generated
by our method are semantically closer to those generated by the teacher model
although in some cases are closer to those generated by the student model. Due
to the high capability of the teacher model, our method, as an intermediate
model from the teacher to the student, achieves higher fidelity than the
student model.

Figure 5: Generated images on CIFAR10 and Celeb-A. We use the same noise item
$\epsilon_{t}$ for all models in each time step. Generally, our method’s
results are semantically closer to the teacher model’s results, although in
some cases are closer to the student model.

Figure 6: Memory footprint for each layer in teacher and student networks used
for ImageNet classification and Celeb-A generation tasks. Note that both the
teacher and student have identical architecture. The student model is enhanced
with our proposed RED blocks, representing our method.

###  4.3 Memory Footprint and Edge Device Deployment

#### Memory Footprint

To intuitively demonstrate the peak memory reduction enabled by our proposed
method, we trace the memory footprint in layer-wise for all identical
architecture teacher-student pairs, as shown in Fig. 6, and Fig. 7 provided in
supplementary material A.3. Compared to CNN models with vanilla pooling, our
method with aggressive pooling settings achieves significantly lower memory
consumption, particularly in 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. For each
model, 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), similar to [32, 6]. We also measure the actual latency
and model size on an NVIDIA 4090 GPU. Our method reduces the theoretical peak
memory by $3.9\times{\sim}14.7\times$ , as shown in Table 8, 9 for image
classification, and $4\times$, as shown in Table 10 for image generation. All
these three tables are provided in supplementary material A.4.

Table 3: Edge Device Deployment. We measure actual GPU Peak Memory (PkMem), Maximum Power (MaxPower), and Latency for models on an NVIDIA Jetson TX2 device. For each experiment, the batch size is set to the maximum load allowed by the device, which is restricted by the teacher. Model | Input | PkMem (GB) | MaxPower (mW) | Latency (ms)
---|---|---|---|---
T: ResNet18 | $50\times 224^{2}\times 3$ | 2.70 | 3273.0 |  395.42 $\pm$ 4.21
S: ResNet18$\times 4$ | 1.20 | 1297.0 | 151.59 $\pm$ 1.50
S w/ RED (ours) | 1.30 | 1830.0 | 208.43 $\pm$ 1.53
T: ResNet50 | $10\times 224^{2}\times 3$ | 2.00 | 2591.0 |  293.81 $\pm$ 1.09
S: ResNet50$\times 4$ | 1.00 | 1068.0 | 125.49 $\pm$ 1.17
S w/ RED (ours) | 1.10 | 2136.0 | 246.68 $\pm$ 1.73
T: ResNet152 | $8\times 224^{2}\times 3$ | 2.8 | 3880.0 |  612.03 $\pm$ 3.55
S: MobileNetV2$\times 2$ | 0.71 | 305.0 | 23.98 $\pm$ 0.84
S w/ RED (ours) | 0.87 | 611.0 | 72.65 $\pm$ 1.01
T: DDPM | $30\times 32^{2}\times 3$ | 2.90 | 4257.0 |  684.31 $\pm$ 5.39
S: DDPM$\times 2$ | 1.30 | 1983.0 | 249.92 $\pm$ 3.76
S w/ RED (ours) | 1.40 | 2362.0 | 304.97 $\pm$ 1.87

#### Edge Device Deployment

In addition to the empirical peak memory estimated from the activation sizes
in the network, we also measure the actual peak memory consumed on an NVIDIA
Jetson TX2 device. Moreover, we measure the maximum GPU power and latency
incurred by the baseline teacher, student, and our RED models on the same edge
GPU, as shown in Table 3 here, and Table 7 in supplementary material A.3. Our
models yield similar peak memory as the student models (due to similar levels
of aggressive striding) and ${\sim}2{\times}{\sim}{3.2}{\times}$ lower peak
memory for image classification and ${\sim}2\times$ lower peak memory for
image generation tasks compared to the teacher models. Note that our
theoretical and measured peak memory reduction factors are different due to
varying device setup 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.4 Ablation Study

As shown in Table 4, we conduct the ablation study on STL10 [7] dataset with
the ResNext18 backbone by removing logit module, residual encoder, shorcut,
and the whole RED block, respectively. The performance of our proposed model
slightly degrades without the logit module or shortcut but seriously 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 necessity
of integrating RED blocks to the student model. An ablation study on our loss
function is in supplementary material A.6.

Table 4: Module Discussion on STL10 dataset with ResNext18 with $\alpha{=}0.5$, $\beta{=}50$, and cosine distance for the RED loss. LM denotes Logit Module, and RE denotes Residual Encoder. Method | RED block | Top1 Acc. (%)
---|---|---
LM | RE | Shortcut
w/o LM | - | $\checkmark$ | $\checkmark$ | 84.79
w/o RE | $\checkmark$ | - | $\checkmark$ | 82.57
w/o Shortcut | $\checkmark$ | $\checkmark$ | - | 84.54
w/o RED block | - | - | - | 82.31
RED (ours) | $\checkmark$ | $\checkmark$ | $\checkmark$ | 84.80

##  5 Conclusions

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 in edge devices, such as micro-
controllers with tight memory budget, 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 [9, 10], thereby significantly reducing the bandwidth
between the image sensor and the back-end processing unit. Our method is based
on a 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-memory
trade-off. For diffusion-based image generation, our method significantly
reduces the peak memory of the denoising network with slight degradation in
the fidelity and diversity of the generated images.

## References

  * [1] 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.
  * [2] 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 pattern recognition, pages 9163–9171, 2019.
  * [3] 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/CVF Conference on Computer Vision and Pattern Recognition, pages 22669–22679, 2023.
  * [4] 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 on Applications of Computer Vision, pages 3974–3983, 2023.
  * [5] Pengguang Chen, Shu Liu, Hengshuang Zhao, and Jiaya Jia.  Distilling knowledge via knowledge review.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 5008–5017, 2021.
  * [6] Aakanksha Chowdhery, Pete Warden, Jonathon Shlens, Andrew Howard, and Rocky Rhodes.  Visual wake words dataset.  arXiv preprint arXiv:1906.05721, 2019.
  * [7] Adam Coates, Andrew Ng, and Honglak Lee.  An analysis of single-layer networks in unsupervised feature learning.  In Proceedings of the fourteenth international conference on artificial intelligence and statistics, pages 215–223. JMLR Workshop and Conference Proceedings, 2011.
  * [8] Antonia Creswell, Tom White, Vincent Dumoulin, Kai Arulkumaran, Biswa Sengupta, and Anil A Bharath.  Generative adversarial networks: An overview.  IEEE signal processing magazine, 35(1):53–65, 2018.
  * [9] Gourav Datta, Souvik Kundu, Zihan Yin, Ravi Teja Lakkireddy, Joe Mathai, Ajey P Jacob, Peter A Beerel, and Akhilesh R Jaiswal.  A processing-in-pixel-in-memory paradigm for resource-constrained tinyml applications.  Scientific Reports, 12(1):14396, 2022.
  * [10] 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), pages 1–5, 2023.
  * [11] Ross Girshick.  Fast r-cnn.  In Proceedings of the IEEE international conference on computer vision, pages 1440–1448, 2015.
  * [12] Song Han, Xingyu Liu, Huizi Mao, Jing Pu, Ardavan Pedram, Mark A Horowitz, and William J Dally.  Eie: Efficient inference engine on compressed deep neural network.  ACM SIGARCH Computer Architecture News, 44(3):243–254, 2016.
  * [13] Song Han, Huizi Mao, and William J Dally.  Deep compression: Compressing deep neural networks with pruning, trained quantization and huffman coding.  International Conference on Learning Representations, 2016.
  * [14] Kaiming He, Georgia Gkioxari, Piotr Dollár, and Ross Girshick.  Mask r-cnn.  In Proceedings of the IEEE international conference on computer vision, pages 2961–2969, 2017.
  * [15] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Deep residual learning for image recognition.  In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770–778, 2016.
  * [16] 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, pages 3779–3787, 2019.
  * [17] Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter.  Gans trained by a two time-scale update rule converge to a local nash equilibrium.  Advances in neural information processing systems, 30, 2017.
  * [18] Geoffrey Hinton, Oriol Vinyals, and Jeff Dean.  Distilling the knowledge in a neural network.  arXiv preprint arXiv:1503.02531, 2015.
  * [19] Jonathan Ho, Ajay Jain, and Pieter Abbeel.  Denoising diffusion probabilistic models.  Advances in neural information processing systems, 33:6840–6851, 2020.
  * [20] Andrew Howard, Mark Sandler, Grace Chu, Liang-Chieh 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, pages 1314–1324, 2019.
  * [21] Zehao Huang and Naiyan Wang.  Like what you like: Knowledge distill via neuron selectivity transfer.  arXiv preprint arXiv:1707.01219, 2017.
  * [22] Itay Hubara, Matthieu Courbariaux, Daniel Soudry, Ran El-Yaniv, and Yoshua Bengio.  Binarized neural networks.  Advances in neural information processing systems, 29, 2016.
  * [23] J Stuart Hunter.  The exponentially weighted moving average.  Journal of quality technology, 18(4):203–210, 1986.
  * [24] Forrest N Iandola, 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.
  * [25] Ying Jin, Jiaqi Wang, and Dahua Lin.  Multi-level logit distillation.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 24276–24285, 2023.
  * [26] Jangho Kim, SeongUk Park, and Nojun Kwak.  Paraphrasing complex network: Network compression via factor transfer.  Advances in neural information processing systems, 31, 2018.
  * [27] Alex Krizhevsky, Geoffrey Hinton, et al.  Learning multiple layers of features from tiny images.  Technical Report, University of Toronto, 2009.
  * [28] Matthias Langer, Zhen He, Wenny Rahayu, and Yanbo Xue.  Distributed training of deep learning models: A taxonomic perspective.  IEEE Transactions on Parallel and Distributed Systems, 31(12):2802–2818, 2020.
  * [29] Seung 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), pages 335–350, 2018.
  * [30] 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, pages 19–35. Springer, 2022.
  * [31] Lijiang Li, Huixia Li, Xiawu Zheng, Jie Wu, Xuefeng Xiao, Rui Wang, Min Zheng, Xin Pan, Fei Chao, and Rongrong Ji.  Autodiffusion: Training-free optimization of time steps and architectures for automated diffusion model acceleration.  In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 7105–7114, 2023.
  * [32] Ji Lin, Wei-Ming Chen, Han Cai, Chuang Gan, and Song Han.  Mcunetv2: Memory-efficient patch-based inference for tiny deep learning.  arXiv preprint arXiv:2110.15352, 2021.
  * [33] 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.
  * [34] 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, 35:22941–22954, 2022.
  * [35] Wei Liu, Dragomir Anguelov, Dumitru Erhan, Christian Szegedy, Scott Reed, Cheng-Yang Fu, and Alexander C Berg.  Ssd: Single shot multibox detector.  In European conference on computer vision, pages 21–37. Springer, 2016.
  * [36] 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.
  * [37] Jonathan Long, Evan Shelhamer, and Trevor Darrell.  Fully convolutional networks for semantic segmentation.  In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 3431–3440, 2015.
  * [38] Yiheng Lu, Ziyu Guan, Yaming Yang, Wei Zhao, Maoguo Gong, and Cai Xu.  Entropy induced pruning framework for convolutional neural networks.  In Proceedings of the AAAI Conference on Artificial Intelligence, volume 38, pages 3918–3926, 2024.
  * [39] 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.
  * [40] Wonpyo Park, Dongju Kim, Yan Lu, and Minsu Cho.  Relational knowledge distillation.  In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 3967–3976, 2019.
  * [41] Nikolaos Passalis and Anastasios Tefas.  Learning deep representations with probabilistic knowledge transfer.  In Proceedings of the European Conference on Computer Vision (ECCV), pages 268–284, 2018.
  * [42] Ilija Radosavovic, Raj Prateek Kosaraju, Ross Girshick, Kaiming He, and Piotr Dollár.  Designing network design spaces.  In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 10428–10436, 2020.
  * [43] Joseph Redmon and Ali Farhadi.  Yolov3: An incremental improvement.  arXiv preprint arXiv:1804.02767, 2018.
  * [44] Olaf Ronneberger, Philipp Fischer, and Thomas Brox.  U-net: Convolutional networks for biomedical image segmentation.  In Medical image computing and computer-assisted intervention–MICCAI 2015: 18th international conference, Munich, Germany, October 5-9, 2015, proceedings, part III 18, pages 234–241. Springer, 2015.
  * [45] 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.
  * [46] Tim Salimans, Ian Goodfellow, Wojciech Zaremba, Vicki Cheung, Alec Radford, and Xi Chen.  Improved techniques for training gans.  Advances in neural information processing systems, 29, 2016.
  * [47] Mark Sandler, Andrew Howard, Menglong Zhu, Andrey Zhmoginov, and Liang-Chieh Chen.  Mobilenetv2: Inverted residuals and linear bottlenecks.  In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 4510–4520, 2018.
  * [48] Axel Sauer, Dominik Lorenz, Andreas Blattmann, and Robin Rombach.  Adversarial diffusion distillation.  arXiv preprint arXiv:2311.17042, 2023.
  * [49] Karen Simonyan and Andrew Zisserman.  Very deep convolutional networks for large-scale image recognition.  arXiv preprint arXiv:1409.1556, 2014.
  * [50] Jiaming Song, Chenlin Meng, and Stefano Ermon.  Denoising diffusion implicit models.  arXiv preprint arXiv:2010.02502, 2020.
  * [51] Christian Szegedy, Wei Liu, Yangqing Jia, Pierre Sermanet, Scott Reed, Dragomir Anguelov, Dumitru Erhan, Vincent Vanhoucke, and Andrew Rabinovich.  Going deeper with convolutions.  In 2015 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), volume 1, pages 1–9, 2015.
  * [52] Yonglong Tian, Dilip Krishnan, and Phillip Isola.  Contrastive representation distillation.  In International Conference on Learning Representations, 2020.
  * [53] Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Herve Jegou.  Training data-efficient image transformers & distillation through attention.  In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 10347–10357. PMLR, 18–24 Jul 2021.
  * [54] Frederick Tung and Greg Mori.  Similarity-preserving knowledge distillation.  In Proceedings of the IEEE/CVF international conference on computer vision, pages 1365–1374, 2019.
  * [55] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Ł ukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
  * [56] 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 on computer vision and pattern recognition, pages 1492–1500, 2017.
  * [57] Tianwei Yin, Michaël Gharbi, Richard Zhang, Eli Shechtman, Fredo Durand, William T Freeman, and Taesung Park.  One-step diffusion with distribution matching distillation.  arXiv preprint arXiv:2311.18828, 2023.
  * [58] 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.
  * [59] Borui Zhao, Quan Cui, Renjie Song, Yiyu Qiu, and Jiajun Liang.  Decoupled knowledge distillation.  In Proceedings of the IEEE/CVF Conference on computer vision and pattern recognition, pages 11953–11962, 2022.

##  Appendix A Supplementary Materials

###  A.1 Datasets

#### Datasets for image classification

1) STL10 [7] contains 5K training images with 10 classes and 8K testing images
of resolution $96\times 96$ pixels. Specifically, we resize the image
resolution to $128\times 128$ pixels for aggressive pooling. 2) ImageNet [45]
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\times 224$ pixels as the origin for aggressive
pooling.

#### Datasets for image generation

1) CIFAR-10 [27] 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\times 32$ in our experiments. 2) Celeb-A [36] 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\times 64$ which is
widely-used in diffusion-based methods [50, 3] in our experiments.

###  A.2 Implementation Details

#### Details for Image Classification

Different from traditional distillation tasks, we only modify the pooling
layer strides instead of the depth and width of the network to get the student
model, which is called aggressive pooling setting. The advantage of the
aggressive pooling setting is to reduce the peak memory and also reduces the
computational complexity and inference time of the network, which is discussed
in Section 4.3. Specifically, we increase the first pooling layer stride
$\times 2\sim\times 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. 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 [47], MobileNetV3 [20], and
ResNext [56]. The proposed method is compared with several distillation
methods [18, 1, 58, 54, 2, 40, 16, 26, 21]. Specifically, the student
architecture is trained from scratch as being distilled from pre-trained
teacher architecture by different methods for 300 epochs. The batch size is
set to 8 and drop out 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 factor $0.2$ at the $180^{th}$, $240^{th}$ and $270^{th}$ epoch,
respectively. The $\alpha,\beta$ in Equation 9 is set to be $0.5,50$,
respectively.

On ImageNet dataset, we experiment on the MobileNetV2 [47], and ResNet [15],
which are widely used in distillation benchmarks [25, 5, 52]. The proposed
method is compared with response-based methods like KD [18], DKD [59], MLLD
[25] and feature-based methods like FitNet [1], RKD [40], ReviewKD [5] and CRD
[52]. 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 [25] but keep training for 300 epochs and decay the learning rate
at the $180^{th}$, $240^{th}$ and $270^{th}$ epoch with factor $0.1$, since
the student architectures with the aggressive pooling setting generally
require more epochs to converge. The $\alpha,\beta$ in Equation 9 is set to be
$0.5,1$, respectively.

#### Details for Image Generation

For the teacher model, we keep the same experiment settings as DDPM [19] with
applying $T=1000$, $\beta_{1}=10^{-4}$, $\beta_{T}=0.02$, and the U-Net
backbone with $4$ different feature map resolutions ($32\times 32$ to $4\times
4$ for CIFAR-10, while $64\times 64$ to $8\times 8$ for Celeb-A). For the
student model, we increase the first pooling layer stride of the U-Net
backbone $\times 2$ 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 $3$
different feature map resolutions ($16\times 16$ to $4\times 4$ for CIFAR-10,
while $32\times 32$ to $8\times 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 by obeying the same rule as described in Section 3.5. For
CIFAR-10 dataset, we train all models 800K iterations and sample 50K images
for FID [17] & IS [46] evaluation. For Celeb-A dataset, we train all models
250K iterations and sample 50K images for FID [17] & IS [46] evaluation. All
experiments are implemented in Pytorch and evaluated on a NVIDIA 4090 GPU.

###  A.3 Additional Experiments on STL10

Table 5: Top-1 accuracy (%) on STL10 with identical teacher and student
architecture.

Method | T: MbNetV2 | T: MbNetV3-Small | T: ResNext18 | Avg.
---|---|---|---|---
85.34 | 83.74 | 85.12
S: MbNetV2$\times 4$ |  S: MbNetV3-Small$\times 4$ |  S: ResNext18$\times 4$
76.18 | 71.27 | 79.07
KD [18] | 79.26 | 71.56 | 81.27 | 77.36
FitNet [1] | 78.55 | 71.93 | 80.21 | 76.90
AT [58] | 81.35 | 75.88 | 82.90 | 80.04
SP [54] | - | 72.42 | 77.59 | 75.01
VID [2] | 77.53 | 71.83 | 76.99 | 75.45
RKD [40] | 78.00 | 67.27 | 75.80 | 73.69
AB [16] | 80.79 | 72.94 | 81.42 | 78.38
FT [26] | 74.74 | 72.11 | 80.64 | 75.83
NST [21] | 82.66 | 76.11 | 81.89 | 80.22
RED (ours) | 83.97 | 77.31 | 84.80 | 82.04

‘-’ denotes we do not get reasonable results for the student architecture with
the distillation method under the aggressive pooling setting.

Table 6: Top-1 accuracy (%) on STL10 with different teacher and student
architecture.

Method | T: ResNext18 | T: ResNext18 | T: MbNetV2 | Avg
---|---|---|---|---
85.12 | 85.12 | 85.34
S: MbNetV2$\times 4$ |  S: MbNetV3-Small$\times 4$ |  S: MbNetV3-Small$\times 4$
76.18 | 71.27 | 71.27
KD [18] | 78.14 | 69.98 | 70.95 | 73.02
FitNet [1] | 80.46 | 73.23 | 72.88 | 75.52
AT [58] | 80.63 | 74.58 | 73.39 | 76.2
SP [54] | 67.56 | 63.91 | - | 65.74
VID [2] | 74.40 | 69.19 | 71.69 | 71.76
RKD [40] | 70.63 | 72.10 | 68.78 | 70.50
AB [16] | 81.46 | 73.98 | 75.20 | 76.88
FT [26] | 77.33 | 69.79 | 66.85 | 71.32
NST [21] | 79.09 | 65.04 | 73.26 | 72.46
RED (ours) | 83.23 | 77.15 | 77.19 | 79.19

‘-’ denotes we don’t get reasonable results for the student architecture with
the distillation method under the aggressive pooling setting.

Figure 7: Memory footprint for models used in STL10 classification. We measure the memory footprint of the baseline model, which is the teacher model, and the student model enhanced with RED blocks, representing our method.  Table 7: Edge Device Deployment of STL10. We measure actual GPU Peak Memory (PkMem), Maximum Power (MaxPower), and Latency for models we use in STL10 classification on an NVIDIA Jetson TX2 device. For each experiment, the batch size is set to the maximum load allowed by the device, which is restricted by the teacher. Model | Input | PkMem (GB) | MaxPower (mW) | Latency (ms)
---|---|---|---|---
T: MbNetV2 | $20\times 128^{2}\times 3$ | 2.30 | 1830.0 | 221.53 $\pm$ 2.34
S: MbNetV2$\times 8$ | 0.78 | 458.0 | 49.51 $\pm$ 1.98
S w/ RED (ours) | 0.81 | 534.0 | 55.43 $\pm$ 1.80
T: MbNetV3-Small | $50\times 128^{2}\times 3$ | 2.30 | 1526.0 | 195.63 $\pm$ 1.09
S: MbNetV3-Small$\times 4$ | 1.10 | 611.0 | 69.76 $\pm$ 1.55
S w/ RED (ours) | 1.10 | 687.0 | 83.74 $\pm$ 1.40
T: ResNext18 | $50\times 128^{2}\times 3$ | 2.40 | 3276.0 | 448.38 $\pm$ 6.97
S: ResNext18$\times 4$ | 1.10 | 1450.0 | 173.23 $\pm$ 1.92
S w/ RED (ours) | 1.20 | 2058.0 | 275.02 $\pm$ 3.25

#### Image Classification Results

Table 5 shows the classification results on STL10 dataset with the setting
that the teacher model and student model are in identical architecture. Same
as ImageNet, ‘$\times n$’ denotes we increase the $1^{st}$ pooling layer
stride of this architecture with $n$ times, and the best results are
highlighted in boldface. Under the identical architecture setting, our method
performs the best among all state-of-the-art distillation methods. For
different architecture knowledge distillation, as shown in Table 6, our method
achieves better performance compared with other distillation methods as well.

#### Memory Footprint

Fig. 7 shows the traced memory footprint for identical architecture teacher-
student pairs we use in STL10 classification. Similar to Fig. 6, our method
achieves much lower memory consumption in layer-wise compared to the teacher
baseline model.

#### Edge Device Deployment

Table 7 shows the actual peak memory, maximum GPU power, and latency measured
on an NVIDIA Jetson TX2 device for models we use in STL10 classification. Our
models achieve similar peak memory as the student models, while incurring
worse latency and power compared to the student models, and improved latency
compared to the teacher models.

###  A.4 Efficiency Analysis

Table 8: Efficiency Analysis on STL10 dataset. We measure Latency for all models with batch size $100$ under the resolution of $128^{2}\times 3$ on a NVIDIA 4090 GPU. Model | PkMem (MB) | Latency (ms) | Model Size (MB) | Top1 Acc. (%)
---|---|---|---|---
T: MbNetV2 | 7.50 | 20.47 $\pm$ 0.02 | 13.50 | 85.34
S: MbNetV2$\times 8$ | 0.51 | 6.90 $\pm$ 0.03 | 13.50 | 69.41
S w/ RED (ours) | 0.51 | 10.93 $\pm$ 0.24 | 14.25 | 73.65
T: MbNetV3-Small | 1.41 | 8.60 $\pm$ 0.03 | 9.75 | 83.74
S: MbNetV3-Small$\times 4$ | 0.31 | 2.59 $\pm$ 0.17 | 9.75 | 71.27
S w/ RED (ours) | 0.31 | 3.49 $\pm$ 0.22 | 10.18 | 77.31
T: ResNext18 | 5.00 | 12.80 $\pm$ 0.01 | 21.47 | 85.12
S: ResNext18$\times 4$ | 0.75 | 3.35 $\pm$ 0.01 | 21.47 | 79.07
S w/ RED (ours) | 0.75 | 6.32 $\pm$ 0.02 | 25.56 | 84.80

Table 8, Table 9, and Table 10 show the results of efficiency analysis on
image classification and image generation tasks for the teacher model, the
student model, and the distilled student model integrated with RED blocks
denoted by ‘S w/ RED’, respectively. The best two results are emphasized in
boldface. Specifically, we analyze the theoretical peak memory [33, 32, 34]
(denoted as ‘PkMem’), and measure the inference latency, and model size of
each model under the corresponding dataset setting. We assume a batch size of
1 for theoretical peak memory estimation to emulate most inference use-cases,
same as [33]. For convenience, the top1 accuracy / FID of each model is
reported as well, and the best two metrics are emphasized in boldface. As
shown in Table 8 and Table 9, student models generally have much lower peak
memory and inference latency but worse top1 accuracy than teacher models. It’s
necessary to note that the student model has the same model size as the
teacher model because we only change the pooling layers’ stride and keep the
network’s depth, width, and architecture totally the same, i.e., identical
architecture setting, leading to the same number of parameters but lower peak
memory and latency. For the last three rows in Table 9, the student model and
teacher model have different architectures, and thus they have different model
sizes. In general, our method would have larger model size and higher latency
compared to the unmodified student model because RED blocks include extra
parameters and computational complexity. However, we seek for the trade-off of
peak memory, latency, and model’s performance. The proposed method achieves
much lower peak memory and latency compared to the teacher model, and much
higher top1 accuracy compared to the student model, which demonstrates the
effectiveness of our method. In the distillation of diffusion models, as
indicated in Table 10, our method attains a lower FID with the same peak
memory as the student model, while also exhibiting reduced peak memory and
latency compared to the teacher.

Table 9: Efficiency Analysis on ImageNet dataset. We measure Latency for all models with batch size $100$ under the resolution of $224^{2}\times 3$ on a NVIDIA 4090 GPU. Model | PkMem (MB) | Latency (ms) | Model Size (MB) | Top1 Acc. (%)
---|---|---|---|---
T: ResNet18 | 3.83 | 13.25 $\pm$ 0.01 | 44.63 | 69.75
S: ResNet18$\times 4$ | 0.77 | 4.13 $\pm$ 0.03 | 44.63 | 61.79
S w/ RED (ours) | 0.77 | 6.78 $\pm$ 0.04 | 48.72 | 65.23
T: ResNet50 | 9.19 | 44.61 $\pm$ 0.04 | 97.70 | 76.13
S: ResNet50$\times 4$ | 2.30 | 12.67 $\pm$ 0.03 | 97.70 | 69.50
S w/ RED (ours) | 2.30 | 21.70 $\pm$ 0.04 | 160.44 | 73.23
T: ResNet152 | 9.19 | 94.99 $\pm$ 0.05 | 230.20 | 78.32
S: MbNetV2$\times 2$ | 1.15 | 4.01 $\pm$ 0.02 | 16.23 | 62.65
S w/ RED (ours) | 2.30 | 9.69 $\pm$ 0.01 | 32.73 | 68.89
Table 10: Efficiency Analysis on CIFAR10 dataset for DDPM. We measure Latency for all models with batch size $200$ under the resolution of $32^{2}\times 3$ on a NVIDIA 4090 GPU. Model | PkMem (MB) | Latency (ms) | Model Size (MB) | FID
---|---|---|---|---
T: DDPM | 3.50 | 71.30 $\pm$ 0.04 | 133.09 | 7.23
S: DDPM$\times 2$ | 0.88 | 24.33 $\pm$ 0.19 | 132.34 | 21.49
S w/ RED (ours) | 0.88 | 41.40 $\pm$ 0.15 | 149.92 | 12.88

###  A.5 Distillation for Vision Transformer

Table 11: Evaluated of our proposed ReDistill method on Vision Transformers [55]. Dataset | Method | Top1 Acc. (%)
---|---|---
STL10 | T: MbNetV2 [47] | 85.34
S: ViT-B-16 [55] | 65.65
DeiT [53] | 66.60
RED (ours) + DeiT | 71.63
ImageNet | T: RegNetY-160 [42] | 82.97
S: ViT-T-16 [53] | 71.34
DeiT [53] | 72.47
RED (ours) + DeiT | 72.52

It is interesting to explore the possibility of combining the proposed method
with existing distillation methods in vision transformer pre-training. In ViT
[55], the patch embedding layer can be seen as a special pooling layer since
it re-arranges the image pixels by patch-wise and the resolution decreases by
the square of patch size in the spatial dimension, even though it increases in
the channel dimension. Specifically, we insert a RED block after the patch
embedding layer to apply the proposed distillation framework. Then the
modified ViT model is integrated with DeiT [53], a widely-used distillation
method for ViT pre-training. DeiT utilizes the logit output of CNN-based
teacher model to supervise a distill token that interacts with other tokens,
which is orthogonal to and easily combined with our method. As shown in Table
11, our method brings improvements to DeiT on STL10 dataset and ImageNet
dataset. Due to implementation specifics, all models are not applied with EMA
[23], leading to results slightly different from those reported in the
original paper. ViT doesn’t have a good performance on STL10 dataset because
transformer-based model usually requires a large amount of data for pre-
training [53]. Even DeiT fails to perform well on STL10 dataset due to the
limited training set and it might be insufficient to utilize merely the logit
output of the CNN-based teacher model. Our method provides a new perspective
to utilize the features of CNN-based models for ViT pre-training and achieves
obvious improvements when combined with DeiT.

We leave thorough investigation of peak memory reduction for transformers as
future work. Unlike CNNs with peak memory at the first few layers with large
feature maps, the peak memory of transformers can be dominated by attention
maps, which are of size $N\times N$ for $N$ tokens. Our distillation method is
intended to reduce the peak memory of convolutional neural networks.

###  A.6 Ablation Study on Loss Function

We also conduct the ablation study for the loss function, an equally important
portion of distillation. The results are shown in Table 12. Without the RED
loss, i.e. setting $\alpha$ to $1$ and then the loss function
$\mathcal{L}=\mathcal{L}_{task}$, the student model is just integrated with
RED blocks but without distillation. Under this circumstance, the poor
performance of the student model illustrates the necessity of applying RED
loss. The proposed method performs worse when applying the Euclidean distance
instead of the cosine distance to calculate the RED loss, and it is sensitive
to the hyper-parameter $\beta$ in Equation 9.

Table 12: Loss Function Discussion on STL10 dataset with ResNext18. Method | RED loss | RED block | Top1 Acc. (%)
---|---|---|---
$\alpha$ | $\beta$ | Distance | Logit | Residual Encoder | Shortcut
w/o $\mathcal{L}_{RED}$ | 1 | 0 | Cosine | $\checkmark$ | $\checkmark$ | $\checkmark$ | 82.26
Euclidean Distance | 0.5 | 50 | Euclidean | $\checkmark$ | $\checkmark$ | $\checkmark$ | 82.66
$\beta=1$ | 0.5 | 1 | Cosine | $\checkmark$ | $\checkmark$ | $\checkmark$ | 82.04
RED (ours) | 0.5 | 50 | Cosine | $\checkmark$ | $\checkmark$ | $\checkmark$ | 84.80

###  A.7 Broader Impact

Our work significantly reduces the peak memory of convolutional neural
networks for image classification and image generation, which enables the
deployment of these networks on edge devices and saves power consumption.
However, image generation models can be misused for various purposes, e.g., by
creating fake face images. We urge users to adhere to usage guidelines for
image generators.