# Dynamic Early Exiting Predictive Coding Neural Networks
††thanks: This work was partly funded by the European Union’s Horizon research
and innovation program under grant agreement No 101070374.

Alaa Zniber, Ouassim Karrakchou, Mounir Ghogho  TICLab, International
University of Rabat, Morocco
<EMAIL_ADDRESS>
###### Abstract

Internet of Things (IoT) sensors are nowadays heavily utilized in various
real-world applications ranging from wearables to smart buildings passing by
agrotechnology and health monitoring. With the huge amounts of data generated
by these tiny devices, Deep Learning (DL) models have been extensively used to
enhance them with intelligent processing. However, with the urge for smaller
and more accurate devices, DL models became too heavy to deploy. It is thus
necessary to incorporate the hardware’s limited resources in the design
process. Therefore, inspired by the human brain known for its efficiency and
low power consumption, we propose a shallow bidirectional network based on
predictive coding theory and dynamic early exiting for halting further
computations when a performance threshold is surpassed. We achieve comparable
accuracy to VGG-16 in image classification on CIFAR-10 with fewer parameters
and less computational complexity.

###### Index Terms:

Dynamic Neural Networks, Predictive Coding, Early Exiting, Edge Devices

##  I Introduction

The Internet of Things (IoT) is nowadays a paradigm of reference for several
applications. IoT is adopted in various domains such as smart farming, drone
imaging, industry 4.0, and safety and is implemented in many challenging
environments like satellites, submarines, and the Large Hadron Collider (LHC).
Hence, huge amounts of data are amassed by sensors and can be overwhelming to
process. For example, the LHC generates approximately one petabyte of
collision data per second [1]. Therefore, IoT sealed an alliance with
Artificial Intelligence (AI) to improve data management pipelines and offer
better analytics to clients. In the context of Artificial Intelligence of
Things (AIoT), Deep Learning (DL) models, mainly fostered in cloud servers,
have proven their efficacy in injecting intelligence into the network of
connected IoT devices [2]. Raw data captured by sensors can thus be used to
produce pertinent insights for the end user’s benefit.

With the rapid development of edge computing, AI services have known a
migration from cloud servers to edge devices (e.g., IoT gateways and fog
nodes). For sensitive applications, such as e-health and smart surveillance,
deploying AI services at the edge has numerous benefits [3]. Data will no
longer need to be transferred through the Internet, increasing data privacy
and reducing security breaches. Latency will also be reduced to account only
for the deployed models’ computational complexity. However, edge devices can
span large computational resource constraints from microcontrollers to
cloudlets passing by single-board computers [4]. Therefore, the observed
heterogeneity in edge device requirements can be troublesome for deploying
state-of-the-art DL models as they usually perform with very deep networks
implicating large numbers of parameters, thus resulting in an inevitable
increase in memory footprint and inference time.

To address this issue, an extensive body of work proposes various techniques
of model compression for lighter memory footprints. For instance, quantization
aims at coding weights on low-precision arithmetic and knowledge distillation
allows the design of a small model (i.e., student) trained on quality features
from a bigger network (i.e., teacher) [5]. Another avenue of research takes
inspiration from the brain, known for its efficiency and plasticity.

The human brain is able to conduct a wide range of tasks dexterously and
process huge amounts of data with around 20 Watts of power [6]. The brain’s
neuronal morphology could explain this frugality in energy consumption. The
brain relies on inhibitory mechanisms that are necessary for quick decision-
making in a survivalist situation [7] as well as bidirectional connections
between higher and lower visual areas that enhance the brain’s abstraction
capacities of the surrounding environment [8]. To understand the interactions
between both visual areas, predictive coding (PC) [9] theory postulates that
the main function of the brain is to minimize a prediction error defined as
the difference between the real and predicted stimuli. PC shows that feed-
forward connections drive prediction errors to higher layers, whereas feedback
connections attempt to predict lower layers’ neural activity. The
bidirectional movement relies on continuously refining the brain’s internal
input representations throughout the visual hierarchy. When implemented in
convolutional neural networks (CNN), PC often yields higher accuracy than its
conventional counterpart [10].

In light of these ideas, our goal is to benefit from PC in designing shallow
CNNs. We aim to show empirically that the PC refinement process allows the
same expressivity and feature diversity that width and depth allow in CNNs
[11]. Furthermore, the refinement of feature representations demands a certain
amount of cyclic processing between lower and higher visual areas until
equilibrium is reached [12]. To alleviate this issue, early exiting techniques
[13] are used to abort further cycling over the feature extractor network once
a performance threshold is reached.

Our contributions can be formulated as follows:

  * •

We apply PC techniques to CNNs to design shallow networks with a considerably
reduced memory footprint that are deployable on edge devices

  * •

We improve PC cyclic processing with an early-exiting mechanism which further
reduces the computational cost and inference time

  * •

We evaluate our proposed model against VGG-16 and achieve for image
classification on CIFAR-10 comparable results with only a 3% difference and
less than 1% of the baseline’s number of parameters

##  II Related Work

This section reviews the main research fields that intersect with our work.
The section covers some solutions to deploy DL models on low-resource edge
devices, previous work done on PC with DL models, and recent applications of
early exiting techniques.

###  II-A DL for Resource-constrained Devices

When computations are fostered in one edge device, compression is a potential
direction for obtaining smaller models. Compression encompasses various
techniques that reduce the model size without a severe drop in performance.
The two most prominent techniques are quantization and pruning. On the one
hand, quantization is the process of representing the model’s weights,
activations, or gradients with a lower number of bits either during or after
training [14]. Rather than optimizing for a target bit-width, mixed
quantization allows the model to be coded on longer bits for early feature
extraction layers and shorter bits for intermediate or concatenation layers
[15]. A recent line of work investigates binary neural networks whose
parameters can be coded only on 2 bits [16]. On the other hand, pruning
entirely removes weights from a network according to a predefined metric of
importance. Thus, numerous metrics exist, such as weight magnitude or impact
on loss [17]. These pruning criteria can either be applied to a pre-trained
network, during training with the possibility of growing back connections, or
even before training at initialization [18].

Instead of working on the model itself with one edge processor, another venue
for deployment in resource-constrained hardware would consider the scenario
where a heterogeneous network is formed via interconnected devices of
different capabilities. Distributed inference [19] is a solution that
leverages the power of internet-connected networks. Rather than deploying a
heavy model on one device, we would cut the model into small executable
segments on the variety of available devices within the network. However, this
solution encounters privacy and latency challenges that Federated Learning has
proposed to raise by avoiding data transfer between the network’s nodes and
encouraging collaborative learning between the devices [20].

Our work complements the aforementioned techniques since our proposed models
can further reduce their memory footprint and number of operations through
quantization and pruning or be deployed in a distributed network.

###  II-B Deep Predictive Coding Networks

PredNet [21] was the first attempt to implement predictive coding in deep
neural networks with the unsupervised task of next-frame video prediction. It
uses multiple layers of recurrent convolutional Long Short-Term Memory cells
that generate a prediction (next frame). This prediction is used to compute an
error representation that is passed to the subsequent layers. However, the
proposed architecture did not benefit from the power of predictive coding as
it did not force the minimization of the bottom-up error [22].

In a supervised learning scheme, Wen et al. enhanced a conventional CNN with
top-down deconvolutional layers using PC updates [10]. It was shown that a
deep PC network yields higher performance than its homologous plain feed-
forward network. Although the network was trained for one particular number of
cycles ($T=6$), they reported a positive correlation between accuracy and
cycling numbers. Finally, PC-based networks were also shown to improve the
robustness of deep models against several types of noises and adversarial
attacks [23]. In this study, pre-trained models were enhanced with feedback
connections, and only the latter were trained.

Our paper is deeply inspired by the advancements in PC-based neural networks.
So far, PC has been utilized to improve model accuracy or robustness by adding
feedback connections to deep feed-forward networks, which causes the number of
parameters to double, and thus the latency to increase. Since our concern is
low-resource edge devices, we envision PC as a technique that might help build
shallow (i.e., few layers and parameters) and expressive (i.e., unrolling the
model through the PC refinement process) networks.

###  II-C Early Exiting Networks

The most straightforward implementation of dynamic neural networks is based on
Early Exiting [13]. It involves mounting tiny decision blocks onto a backbone
model to make quick decisions for easy inputs without resorting to the entire
network. In image classification, the decision block is a tiny multi-layer
perceptron classifier. Therefore, in an early-exit network, a response is
returned if the classifier is sufficiently confident based on a performance
target. Otherwise, the example is passed onto the subsequent layers for finer
processing [24]. In natural language processing, early exiting is used during
the decoding of transformer-based large language models through dynamic per-
token exit decision-making [25].

The early-exiting halting mechanism introduces the concepts of ”easy” and
”hard” samples, which imply that easy samples should exit at earlier
classifiers and harder samples could traverse the full network. If this is
intuitive in image classification, other tasks like object detection need more
precise definitions of sample hardness. Therefore, in [26], a sample for
object detection is labeled ”easy” when the difference in loss (or mean
average precision) between a shallow detector and a deep detector is small.
Early exiting networks are conventionally trained either via scalarization
(i.e., the weighted sum of the internal classifiers losses) or separate
training (i.e., each classifier trained separately). However, both methods
have limitations as scalarization might cause instabilities due to the
accumulation of gradients when the backbone is deep [27], while separate
training hinders the collaboration between classifiers leading to
computational waste [28].

In the proposed models, early exiting is implemented to halt PC cyclic
processing once the user-defined performance target is reached.

##  III Proposed Architecture

As shown in Figure 1, our architecture is formed of a shared backbone
alongside downstream task classifiers. The backbone serves as a feature
extractor. It is built with a bidirectional hierarchy of convolutional and
deconvolutional layers. The backbone is executed for a maximum of $T$ cycles
before outputting the final feature vector for classification. Blue arrows
represent the convolutional forward pass, while red arrows represent the
feedback deconvolutions. Each intermediate representation consists of a
concatenation of a convolutional and a deconvolutional feature maps, referred
to as $C_{l}$ and $D_{l}$ respectively. We perform a variable number of cycles
$t<T$ over the backbone. Once cycling is finished, we feed the last obtained
feature vector $C_{l}(t)$ to the classifier corresponding to the achieved
number of cycles $t$ (green arrow in Figure 1). The first cycle consists of
three consecutive passes: forward, feedback, and forward. Every other cycle
starts with a feedback pass followed by a forward pass.

Figure 1: Proposed model - The backbone is cycling for $t<T$ cycles

###  III-A PC Feature Update Rule

During each cycle, PC adopts a specific feature update rule to merge forward
and feedback feature maps to enrich the internal representation. We adopt the
same formulation proposed in [10]. Suppose the PC architecture is constituted
of $L$ layers, with layer $l=0$ being the input image. We denote the forward
and feedback convolutions by $\mathbf{FF}$ and $\mathbf{FB}$. The
convolutional representations $C_{l}$ are updated on both forward and feedback
passes, while the deconvolutional feature maps $D_{l}$ only change during the
feedback pass. Hence, let $C_{l,f}(t)$ and $C_{l,b}(t)$ be the state of layer
$l$’s convolutional feature map at the end of the forward and feedback passes
of PC cycle $t\in[|1,T|]$. At the beginning of cycle 1, we initialize all
feature maps with a conventional forward pass: $\forall
l\in[|1,L|],C_{l,f}(0)=\mathbf{FF}\left[C_{l-1,f}(0)\right]$. We can thus
define the PC update rules for every cycle $t\in[|1,T|]$ and for every layer
$l\in[|1,L|]$ as follows:

Feedback pass update:

$\displaystyle D_{l-1}(t)=\mathbf{FB}(C_{l,f}(t-1))$ (1) $\displaystyle
C_{l-1,b}(t)=g((1-b_{l-1}).C_{l-1,f}(t-1)+b_{l-1}.D_{l-1}(t))$

Forward pass update:

$\displaystyle
C_{l,f}(t)=g(C_{l,b}(t)+a_{l}.\mathbf{FF}(C_{l-1,b}(t)-D_{l-1}(t)))$ (2)

where $a_{l}$ and $b_{l}$ are trainable non-negative layer-dependent updating
rates, initialized for all layers by 1.0 and 0.5, respectively. However,
$b_{0}=0$ is non-trainable as the input image $C_{0}$ is never updated.
Finally, function $g$ is a Rectified Linear Unit (ReLU) non-linearity.

The PC update rules attempt to minimize the prediction error between forward
and feedback representations. More precisely, after an optimal number of
cycles $T_{opt}$, we will have $C_{l}(t)\simeq D_{l}(t)$. PC update rules thus
tend toward a consistent feature representation: $\forall
l\in[|1,L|],C_{l}(t+1)\simeq C_{l}(t)$. However, this condition is not always
met in practice as it depends on the hardness of the sample being processed.
For instance, hard images might require high cycling numbers to stabilize
representations. Therefore, we do not wait until feature consistency is
achieved to halt the computations as proposed explicitly in [23]. Instead, we
propose to use early exiting to adapt the number of cycles to the sample
hardness.

###  III-B PC Early Exit

We implement early exit on the number of cycles as follows. After the first
cycle, classifier $1$ is applied to the convolutional feature map
$C_{l,f}(1)$. The classification confidence is then compared with a predefined
user threshold. If the confidence is above the threshold, computations are
aborted, and a response is returned. Otherwise, a new cycle is initiated,
followed by another classification and threshold comparison.

In the proposed architecture, the number of classifiers equals the maximum
number of cycles $T$ allowed. The choice of $T$ different classifiers, instead
of one classifier shared by all cycles, is motivated by the fact that feature
vectors are updated from one cycle to another. Hence, a classifier trained on
a 5-cycle-model feature vector will not discern the patterns that a 1-cycle-
model extracts from the same input. Only when feature consistency is achieved
can a shared classifier be sufficient. Nevertheless, since our target hardware
is highly resource-constrained, we usually stop cycling before consistent
feature representations are reached, hence the need for a classifier for each
cycle.

###  III-C PC Training

Our model parameters are the forward, feedback, and the $T$ classifiers’
weights, alongside the $a_{l}$ and $b_{l}$ update rates. These parameters are
learned with classic back-propagation using a single cross-entropy loss
function $\mathcal{L}_{c_{i}}$ per classifier $c_{i}$. Hence, we can train the
network separately for a specific number of cycles corresponding to one of the
classifiers. Moreover, after each backpropagation pass, the weights are frozen
until the chosen number $i$ of PC cycles is achieved before recalculating the
loss function.

To improve the training performance, we can allow classifiers to collaborate
with each other. Therefore, we propose that the backbone and $T$ classifiers
are jointly trained using scalarization adopted from multi-objective
optimization. The total loss is the weighted sum of each classifier’s loss
$\mathcal{L}_{c_{i}}$, as defined in Equation 3:

$\mathcal{L}_{tot}=\sum_{i=1}^{T}\lambda_{i}\mathcal{L}_{c_{i}}$ (3)

where $\lambda_{i}$ is a positive weight for the loss function
$\mathcal{L}_{c_{i}}$.

Our proposed training procedure has the main disadvantage of being time-
consuming. However, not only does the underlying competition launched between
classifiers over the shared weights of the backbone help achieve a Pareto
optimal solution, but it also encourages PC cycles to achieve consistency
early enough as the loss guides the classifiers to collaborate in generating
semantically similar feature vectors. Moreover, the $\lambda_{i}$ coefficients
control the exiting strategy: higher coefficients for the first classifiers
(corresponding to small numbers of cycles) stimulate early exiting, and vice
versa. In our experiments, exits were uniformly weighted: $\forall
i\in[|1,T|],\lambda_{i}=1/T$.

##  IV Experiments

TABLE I: Model configurations A | B | C
---|---|---
Input image 32x32
Conv-16 | Conv-32 | Conv-16
Conv-32 | Conv-64 | Conv-32
Conv-32 | Conv-64 | Conv-32
Conv-64 | Conv-128 | Conv-64
Conv-64 | Conv-128 | Conv-64
- | - | Conv-128
- | - | Conv-128
Global Average Pooling
$T$ Fully Connected layers

###  IV-A Dataset

In order to evaluate our method, we choose the CIFAR-10 dataset [29]. It
includes 60000 32x32 RGB images evenly distributed over 10 classes, with 6000
images per each. CIFAR-10 is adopted by many tiny machine learning benchmarks
[30], and the images simulate well numerous IoT applications using low-
resolution cameras (e.g., surveillance for eyewear protection detection and
smart farming for fruit disease classification). For model learning, the
training set is formed of 50000 images, and the rest is destined for testing.
As a data augmentation procedure, we applied random translation and horizontal
flipping. The training set was clustered in 64-sized batches.

###  IV-B Training & Evaluation

We choose to cycle over the PC backbone for a maximum of $T=5$ cycles. Mini-
batch gradient descent was employed for training via AdamW optimizer
initialized with a learning rate of 0.001 and a weight decay of 0.01.
Regularization was used as a dropout operator of 10%. We run the model for 500
epochs. Moreover, PC models with one cycling number and one classifier were
trained for 200 epochs.

We adopt conventional accuracy as a performance metric for evaluation since
the dataset is balanced. The number of parameters, memory size, and number of
floating-point operations (FLOPs) can be gathered from PyTorch libraries.
Latency was computed for one sample on Google Colaboratory CPU with 500
repetitions in the worst-case scenario (i.e., a sample exiting after all $T=5$
cycles), as it represents the upper bound of our setting.

###  IV-C Model configuration

The model design process was driven by the motivation of exploiting PC
dynamics in order to build shallow networks, in terms of both depth and width,
that could perform as well as the established VGG-16 architecture and that
could be deployed on edge devices of kilobytes (KB) to megabytes (MB) of
memory. We design three models with convolutions sharing the same kernel size
of 3x3 and stride of 1 but different numbers of channels as shown in bold in
Table I. For the sake of readability, we omit PC deconvolutional layers since
they have inverted numbers of channels with respect to the feed-forward
convolutions. Batch normalization is embedded in each convolutional block.
Max-pooling of kernel size 2x2 is activated when the number of channels
changes from one layer to another.

The models presented in Table I are designed to challenge PC feature updates.
Hence, given a shallow Model A, we construct Model B to be wider and Model C
deeper in order to show empirically that recurrence could account for the
improvements that width and depth bring to the neural network. Finally, our
baselines are VGG-16 as well as PC models exclusively trained for $T=5$ cycles
with the aforementioned architectures in Table I, noted PCN-5-{A,B,C}. The
custom baseline models PCN-5 are meant to testify to the validity of the
chosen training scheme.

##  V Discussion

TABLE II: Characteristics for VGG-16 and $T=5$ processing cycles for the proposed models Model | #Params | Size (MB) | FLOPs | Latency (ms)
---|---|---|---|---
VGG-16 | 134M | 512.3 | $8.7\times 10^{8}$ | 55.900
A | 143.416 | 0.6 | $3.3\times 10^{8}$ | 41.438
B | 563.256 | 2.3 | $1.05\times 10^{9}$ | 80.261
C | 590.328 | 2.4 | $4.9\times 10^{8}$ | 54.888

Table II and III aggregate the main results of our experiments from both
efficiency and performance standpoints, respectively.

We observe that the proposed models can fit memory constraints for many edge
processors consuming only a few megabytes of storage. As expected, latency and
the number of operations increase with the model width and depth without
exceeding the reasonable amounts we encounter in edge processing [4]. It is
worth mentioning for Model B that a latency of 80 milliseconds (ms) is rather
acceptable knowing that the network has already done 5 cycles through the
backbone and computed 5 class probabilities, thus, offering better performance
and more accurate classification. Based on Table II and Figure 2, Model A
presents the most desirable properties for an edge application. It is
noticeable that starting from cycle $T=2$, a linear relationship is found
between cycles and FLOPs. It is explained by the fact that one more cycle
implies the same amount of floating point operations. However, it is to note
that the first cycle computes 2 feed-forward passes as the input should first
traverse the network, feedback, and regain again the higher layers near the
classifier. More cycles will thus grow the number of FLOPs linearly. However,
given that our architecture is shallow, attaining reasonably large numbers of
cycles will barely reach the computational cost attained by a conventional
deep network like VGG-16, as seen in the FLOPs column of Table II.

Figure 2: FLOPs and Latency on CPU across exits for Model A

From a performance perspective, we observe from Table III, an increasing
accuracy through cycles as more cycles allow more expressivity to the shallow
network. This is especially the case when classes are hardly separable and for
which the learned patterns must be well distinctive. To appreciate the impact
of PC updates, we note a significant increase of about 20% from $T=1$ to $T=5$
in Model B. Moreover, in order to challenge our training scheme, we report
models trained for the predefined maximum number of cycles $T=5$ and one
classifier. We observe that models using joint training achieve better
accuracy than their homologous PCN-5. This shows that, in a PC setting for
which consistency is sought, joint training encourages collaboration between
classifiers. Furthermore, our proposed models are not very far from VGG-16
performance on CIFAR-10 with only a 3% difference from the 4-cycle and 5-cycle
recurrent processing and a considerable reduction in the numbers of parameters
from $10^{8}$ (i.e., VGG-16) to $10^{5}$ orders of magnitude.

Table III also shows comparable results between the three proposed models.
Wide Model B and deep Model C reach almost the same performance as Model A
which is tinier and faster. Nevertheless, it is worth mentioning that width
helped the model gain high accuracy at later cycles and depth improved the
internal classifiers’ average performance. Overall, we can conclude that PC
dynamics might serve as a complementary key aspect for increasing model
expressivity along with depth and width in feed-forward neural networks.

Figure 3 endorses the claims advanced in Section III-C. We can spot two main
behaviors, other than the collective increase in accuracy which is the goal of
multiobjective optimization. The first is the fluctuations through epochs
which reveal the status of competition between the classifiers. These
fluctuations are severely present in the first epochs since PC dynamics can
still not find common ground between the feature vectors each cycle is
outputting. The second behavior is also related to these fluctuations: we
observe that, near the end of the optimization, the fluctuation amplitude
diminishes and the five curves approach one another. This behavior is
compatible with the underlying idea of PC dynamics in neural networks that
revolves around stability and consistency in the feature representations.
Therefore, with the joint training, a type of implicit knowledge distillation
happens between the cycles. Hence, in numerous cases, a 1-cycle network will
be able to yield a similar feature vector as a 6-cycle network; a property
which is most wanted for reducing computational cost.

TABLE III: Accuracy (%) of baseline models and dynamic PC models at different exits Model | $T=1$ | $T=2$ | $T=3$ | $T=4$ | $T=5$
---|---|---|---|---|---
VGG-16 | 93.25 [31]
PCN-5-A | 85.55
PCN-5-B | 87.53
PCN-5-C | 89.88
A | 79.82 | 85.83 | 88.81 | 89.99 | 89.81
B | 70.21 | 80.9 | 86.72 | 89.85 | 90.82
C | 87.79 | 89.53 | 90.2 | 90.26 | 89.68

Figure 3: Test accuracy per exit for Model A

Finally, Figure 4 emphasizes the motivation behind using early exit. Given
high thresholds, the network can still free with high confidence a large
number of well-classified images at the first exit. Nevertheless, as the user
threshold increases, we notice the importance of more cycles to classify
harder samples correctly.

Figure 4: Ratio of correctly classified images exiting at different cycles
based on a confidence threshold for Model A (ordinates in log scale)

##  VI Conclusion

In this paper, we proposed a shallow network for image classification based on
predictive coding dynamics and early exiting for resource-constrained edge
devices. We found that PC dynamics can play a major role in yielding high
accuracy without resorting to deep models. Since PC is based on minimizing an
objective function that might demand high numbers of cycles, we employed
early-exiting to abort further computation once a user-predefined performance
target is reached. Thus, the paper highlights how PC processing can lead to a
significant reduction in memory footprint while achieving good accuracy. We
intend to continue the present work by putting more emphasis on early exits
through preference-vector-based multiobjective optimization. Furthermore, we
will attempt to define a hardness measure that estimates beforehand the number
of cycles needed in order to avoid calling low-cycle classifiers.

## References

  * [1] Stark, J, ”PLM and the Internet of Things,” Product Lifecycle Management, pp 327–352, vol. 1, Springer Cham., 2022.
  * [2] T. J. Saleem and M. A. Chishti, ”Deep learning for the internet of things: Potential benefits and use-cases,” Digital Communications and Networks, vol. 7, no. 4, pp. 526–542, Nov. 2021.
  * [3] R. Sanchez-Iborra and A. F. Skarmeta, ”TinyML-Enabled Frugal Smart Objects: Challenges and Opportunities,” in IEEE Circuits and Systems Magazine, vol. 20, no. 3, pp. 4-18, 2020.
  * [4] M. G. S. Murshed, C. Murphy, D. Hou, N. Khan, G. Ananthanarayanan et al., ”Machine Learning at the Network Edge: A Survey,” ACM Comput. Surv., vol. 54, no. 8, p. 170:1-170:37, Oct. 2021.
  * [5] J. Lee, L. Mukhanov, A. S. Molahosseini, U. Minhas, Y. Hua et al., ”Resource-Efficient Convolutional Networks: A Survey on Model-, Arithmetic-, and Implementation-Level Techniques,” ACM Comput. Surv., vol. 55, no. 13s, p. 276:1-276:36, Jul. 2023.
  * [6] V. Balasubramanian, ”Brain power,” PNAS, vol. 118, no. 32, p. e2107022118, Aug. 2021.
  * [7] L. Luo, ”Architectures of neuronal circuits,” Science, vol. 373, no. 6559, p. eabg7285, Sep. 2021.
  * [8] C. J. Spoerer, T. C. Kietzmann, J. Mehrer, I. Charest, and N. Kriegeskorte, ”Recurrent neural networks can explain flexible trading of speed and accuracy in biological vision,” PLoS Comput Biol, vol. 16, no. 10, p. e1008215, Oct. 2020.
  * [9] R. P. N. Rao and D. H. Ballard, ”Predictive coding in the visual cortex: a functional interpretation of some extra-classical receptive-field effects,” Nat. Neurosci., vol. 2, no. 1, Art. no. 1, Jan. 1999.
  * [10] H. Wen, K. Han, J. Shi, Y. Zhang, E. Culurciello et al., ”Deep Predictive Coding Network for Object Recognition,” in Proceedings of ICML, Jul. 2018, pp. 5266–5275.
  * [11] R. S. van Bergen and N. Kriegeskorte, ”Going in circles is the way forward: the role of recurrence in visual inference,” Current Opinion in Neurobiology, vol. 65, pp. 176–193, Dec. 2020.
  * [12] Y. Song, T. Lukasiewicz, Z. Xu, and R. Bogacz, ”Can the Brain Do Backpropagation? — Exact Implementation of Backpropagation in Predictive Coding Networks,” in Advances in NeurIPS, 2020, pp. 22566–22579.
  * [13] S. Laskaridis, A. Kouris, and N. D. Lane, ”Adaptive Inference through Early-Exit Networks: Design, Challenges and Directions,” Proceedings of EMDL’21, ACM, Jun. 2021, pp. 1–6.
  * [14] A. Gholami, S. Kim, Z. Dong, Z. Yao, M. W. Mahoney et al., ”A Survey of Quantization Methods for Efficient Neural Network Inference,” in Low-Power Computer Vision, Chapman and Hall/CRC, 2022.
  * [15] Z. Dong, Z. Yao, D. Arfeen, A. Gholami, M. W. Mahoney et al., ”HAWQ-V2: Hessian Aware trace-Weighted Quantization of Neural Networks,” in Advances in NeurIPS, 2020, pp. 18518–18529.
  * [16] H. Qin, R. Gong, X. Liu, X. Bai, J. Song et al., ”Binary neural networks: A survey,” Pattern Recognition, vol. 105, p. 107281, Sep. 2020.
  * [17] U. Evci, T. Gale, J. Menick, P. S. Castro, and E. Elsen, ”Rigging the lottery: making all tickets winners,” in Proceedings of ICML’20. pp. 2943–2952.
  * [18] H. Wang, C. Qin, Y. Bai, Y. Zhang, and Y. Fu, ”Recent Advances on Neural Network Pruning at Initialization,” IJCAI, Jul. 2022, pp. 5638–5645.
  * [19] J. Chen and X. Ran, ”Deep Learning With Edge Computing: A Review,” Proceedings of the IEEE, vol. 107, no. 8, pp. 1655–1674, Aug. 2019.
  * [20] Y. Matsubara, M. Levorato, and F. Restuccia, ”Split Computing and Early Exiting for Deep Learning Applications: Survey and Research Challenges,” ACM Comput. Surv., vol. 55, no. 5, p. 90:1-90:30, Dec. 2022.
  * [21] W. Lotter, G. Kreiman, and D. Cox, ”Deep Predictive Coding Networks for Video Prediction and Unsupervised Learning,” ICLR, Nov. 2016.
  * [22] R. P. Rane, E. Szügyi, V. Saxena, A. Ofner, and S. Stober, ”PredNet and Predictive Coding: A Critical Review,” in Proceedings of ICMR ’20, ACM, Jun. 2020, pp. 233–241.
  * [23] B. Choksi, M. Mozafari, C. Biggs O’ May, B. Ador, A. Alamia et al., ”Predify: Augmenting deep neural networks with brain-inspired predictive coding dynamics,” in Advances in NeurIPS, 2021, pp. 14069–14083.
  * [24] S. Scardapane, M. Scarpiniti, E. Baccarelli, and A. Uncini, ”Why Should We Add Early Exits to Neural Networks?,” Cogn. Comput., vol. 12, no. 5, pp. 954–966, Sep. 2020.
  * [25] T. Schuster, A. Fisch, J. Gupta, M. Dehghani, D. Bahri et al., ”Confident Adaptive Language Modeling,” in Advances in NeurIPS, May 2022.
  * [26] Z. Lin, Y. Wang, J. Zhang, and X. Chu, ”DynamicDet: A Unified Dynamic Architecture for Object Detection,” Proceedings of the IEEE/CVF CVPR, 2023, pp. 6282–6291.
  * [27] H. Li, H. Zhang, X. Qi, R. Yang, and G. Huang, ”Improved Techniques for Training Adaptive Deep Networks,” Proceedings of the IEEE/CVF ICCV, 2019, pp. 1891–1900.
  * [28] M. Wołczyk, B. Wójcik, K. Bałazy, I. Podolak, J. Tabor et al., ”Zero Time Waste: Recycling Predictions in Early Exit Neural Networks,” in Advances in NeurIPS, 2021, pp. 2516–2528.
  * [29] A. Krizhevsky, and G. Hinton, ”Learning Multiple Layers of Features from Tiny Images,” 2009.
  * [30] C. Banbury, V. J. Reddi, P. Torelli, J. Holleman, N. Jeffries et al., ”MLPerf Tiny Benchmark,” Proceedings of the NeurIPS Track on Datasets and Benchmarks, 2021.
  * [31] H. Jónsson, G. Cherubini, and E. Eleftheriou, ”Convergence Behavior of DNNs with Mutual-Information-Based Regularization,” Entropy (Basel), vol. 22, no. 7, p. 727, Jun. 2020.


# SA-Net: Shuffle Attention for Deep Convolutional Neural Networks

###### Abstract

Attention mechanisms, which enable a neural network to accurately focus on all
the relevant elements of the input, have become an essential component to
improve the performance of deep neural networks. There are mainly two
attention mechanisms widely used in computer vision studies, spatial attention
and channel attention, which aim to capture the pixel-level pairwise
relationship and channel dependency, respectively. Although fusing them
together may achieve better performance than their individual implementations,
it will inevitably increase the computational overhead. In this paper, we
propose an efficient Shuffle Attention (SA) module to address this issue,
which adopts Shuffle Units to combine two types of attention mechanisms
effectively. Specifically, SA first groups channel dimensions into multiple
sub-features before processing them in parallel. Then, for each sub-feature,
SA utilizes a Shuffle Unit to depict feature dependencies in both spatial and
channel dimensions. After that, all sub-features are aggregated and a “channel
shuffle” operator is adopted to enable information communication between
different sub-features. The proposed SA module is efficient yet effective,
e.g., the parameters and computations of SA against the backbone ResNet50 are
300 vs. 25.56M and 2.76e-3 GFLOPs vs. 4.12 GFLOPs, respectively, and the
performance boost is more than 1.34% in terms of Top-1 accuracy. Extensive
experimental results on common-used benchmarks, including ImageNet-1k for
classification, MS COCO for object detection, and instance segmentation,
demonstrate that the proposed SA outperforms the current SOTA methods
significantly by achieving higher accuracy while having lower model
complexity. The code and models are available at
https://github.com/wofmanaf/SA-Net.

Index Terms—  spatial attention, channel attention, channel shuffle, grouped
features

##  1 Introduction

Fig. 1: Comparisons of recently SOTA attention models on ImageNet-1k,
including SENet, CBAM, ECA-Net, SGE-Net, and SA-Net, using ResNets as
backbones, in terms of accuracy, network parameters, and GFLOPs. The size of
circles indicates the GFLOPs. Clearly, the proposed SA-Net achieves higher
accuracy while having less model complexity.

Attention mechanisms have been attracting increasing attention in research
communities since it helps to improve the representation of interests, i.e.,
focusing on essential features while suppressing unnecessary ones [1, 2, 3,
4]. Recent studies show that correctly incorporating attention mechanisms into
convolution blocks can significantly improve the performance of a broad range
of computer vision tasks, e.g., image classification, object detection, and
instance segmentation.

There are mainly two types of attention mechanisms most commonly used in
computer vision: channel attention and spatial attention, both of which
strengthen the original features by aggregating the same feature from all the
positions with different aggregation strategies, transformations, and
strengthening functions [5, 6, 7, 8, 9]. Based on these observations, some
studies, including GCNet [1] and CBAM [10] integrated both spatial attention
and channel attention into one module and achieving significant improvement
[10, 4]. However, they generally suffered from either converging difficulty or
heavy computation burdens. Other researches managed to simplify the structure
of channel or spatial attention [1, 11]. For example, ECA-Net [11] simplifies
the process of computing channel weights in SE block by using a 1-D
convolution. SGE [12] groups the dimension of channels into multiple sub-
features to represent different semantics, and applies a spatial mechanism to
each feature group by scaling the feature vectors over all locations with an
attention mask. However, they did not take full advantage of the correlation
between spatial and channel attention, making them less efficient. “Can one
fuse different attention modules in a lighter but more efficient way?”

To answer this question, we first revisit the unit of ShuffleNet v2 [13],
which can efficiently construct a multi-branch structure and process different
branches in parallel. Specifically, at the beginning of each unit, the input
of $c$ feature channels are split into two branches with $c-c^{\prime}$ and
$c^{\prime}$ channels. Afterwards, several convolution layers are adopted to
capture a higher-level representation of the input. After these convolutions,
the two branches are concatenated to make the number of channels as same as
the number of input. At last, the “channel shuffle” operator (defined in [14])
is adopted to enable information communication between the two branches. In
addition, to increase calculation speed, SGE [12] introduces a grouping
strategy, which divides the input feature map into groups along the channel
dimension. Then all sub-features can be enhanced parallelly.

Based on these above observations, this paper proposes a lighter but more
efficient Shuffle Attention (SA) module for deep Convolutional Neural
Networks(CNNs), which groups the dimensions of channel into sub-features. For
each sub-feature, SA adopts the Shuffle Unit to construct channel attention
and spatial attention simultaneously. For each attention module, this paper
designs an attention mask over all the positions to suppress the possible
noises and highlight the correct semantic feature regions as well.
Experimental results on ImageNet-1k (which are shown in Figure 1) have shown
that the proposed simple but effective module containing fewer parameters has
achieved higher accuracy than the current state-of-the-art methods.

The key contributions in this paper are summarized as follows: 1) we introduce
a lightweight yet effective attention module, SA, for deep CNNs, which groups
channel dimensions into multiple sub-features, and then utilizes a Shuffle
Unit to integrate the complementary channel and spatial attention module for
each sub-feature. 2) extensive experimental results on ImageNet-1k and MS COCO
demonstrate that the proposed SA has lower model complexity than the state-of-
the-art attention approaches while achieving outstanding performance.

##  2 Related Work

Multi-branch architectures. Multi-branch architectures of CNNs have evolved
for years and are becoming more accurate and faster. The principle behind
multi-branch architectures is “split-transform-merge”, which eases the
difficulty of training networks with hundreds of layers. The InceptionNet
series [15, 16] are successful multi-branch architectures of which each branch
is carefully configured with customized kernel filters, in order to aggregate
more informative and multifarious features. ResNets [17] can also be viewed as
two-branch networks, in which one branch is the identity mapping. SKNets [2]
and ShuffleNet families [13] both followed the idea of InceptionNets with
various filters for multiple branches while differing in at least two
important aspects. SKNets utilized an adaptive selection mechanism to realize
adaptive receptive field size of neurons. ShuffleNets further merged “channel
split” and “channel shuffle” operators into a single element-wise operation to
make a trade-off between speed and accuracy.

Grouped Features. Learning features into groups dates back to AlexNet [18],
whose motivation is distributing the model over more GPU resources. Deep Roots
examined AlexNet and pointed out that convolution groups can learn better
feature representations. The MobileNets [19, 20] and ShuffleNets [13] treated
each channel as a group, and modeled the spatial relationships within these
groups. CapsuleNets [21, 22] modeled each grouped neuron as a capsule, in
which the neuron activity in the active capsule represented various attributes
of a particular entity in the image. SGE [12] developed CapsuleNets and
divided the dimensions of channel into multiple sub-features to learn
different semantics.

Attention mechanisms. The significance of attention has been studied
extensively in the previous literature. It biases the allocation of the most
informative feature expressions while suppressing the less useful ones. The
self-attention method calculates the context in one position as a weighted sum
of all the positions in an image. SE [23] modeled channel-wise relationships
using two FC layers. ECA-Net [11] adopted a 1-D convolution filter to generate
channel weights and significantly reduced the model complexity of SE. Wang et
al. [24] proposed the non-local(NL) module to generate a considerable
attention map by calculating the correlation matrix between each spatial point
in the feature map. CBAM [10], GCNet [1], and SGE [12] combined the spatial
attention and channel attention serially, while DANet [3] adaptively
integrated local features with their global dependencies by summing the two
attention modules from different branches.

##  3 Shuffle Attention

In this section, we firstly introduce the process of constructing the SA
module, which divides the input feature map into groups, and uses Shuffle Unit
to integrate the channel attention and spatial attention into one block for
each group. After that, all sub-features are aggregated and a “channel
shuffle” operator is utilized to enable information communication between
different sub-features. Then, we show how to adopt SA for deep CNNs. Finally,
we visualize the effect and validate the reliability of the proposed SA. The
overall architecture of SA module is illustrated in Figure 2.

Fig. 2: An overview of the proposed SA module. It adopts “channel split” to
process the sub-features of each group in parallel. For channel attention
branch, using GAP to generate channel-wise statistics, then use a pair of
parameters to scale and shift the channel vector. For spatial attention
branch, adopting group norm to generate spatial-wise statistics, then a
compact feature is created similar to the channel branch. The two branches are
then concatenated. After that, all sub-features are aggregated and finally we
utilize a “channel shuffle” operator to enable information communication
between different sub-features.

Feature Grouping. For a given feature map $X\in\mathbb{R}^{C\times H\times
W}$, where $C$, $H$, $W$ indicate the channel number, spatial height, and
width, respectively, SA first divides $X$ into $G$ groups along the channel
dimension, i.e., $X=[X_{1},\cdots,X_{G}],X_{k}\in\mathbb{R}^{C/G\times H\times
W}$, in which each sub-feature $X_{k}$ gradually captures a specific semantic
response in the training process. Then, we generate the corresponding
importance coefficient for each sub-feature through an attention module.
Specifically, at the beginning of each attention unit, the input of $X_{k}$ is
split into two branches along the channels dimension, i.e., $X_{k1}$,
$X_{k2}\in\mathbb{R}^{C/2G\times H\times W}$. As illustrated in Figure 2, one
branch is adopted to produce a channel attention map by exploiting the inter-
relationship of channels, while the other branch is used to generate a spatial
attention map by utilizing the inter-spatial relationship of features, so that
the model can focus on “what” and “where” is meaningful.

Channel Attention. An option to fully capture channel-wise dependencies is
utilizing the SE block proposed in [23]. However, it will bring too many
parameters, which is not good for designing a more lightweight attention
module in terms of a trade-off between speed and accuracy. Also, it is not
suitable to generate channel weights by performing a faster 1-D convolution of
size $k$ like ECA [11] because $k$ tends to be larger. To improve, we provide
an alternative, which firstly embeds the global information by simply using
global averaging pooling (GAP) to generate channel-wise statistics as
$s\in\mathbb{R}^{C/2G\times 1\times 1}$, which can be calculated by shrinking
$X_{k1}$ through spatial dimension $H\times W$:

$s=\mathcal{F}_{gp}(X_{k1})=\frac{1}{H\times
W}\sum_{i=1}^{H}\sum_{j=1}^{W}X_{k1}(i,j)$ (1)

Furthermore, a compact feature is created to enable guidance for precise and
adaptive selection. This is achieved by a simple gating mechanism with sigmoid
activation. Then, the final output of channel attention can be obtained by

$X^{\prime}_{k1}=\sigma(\mathcal{F}_{c}(s))\cdot
X_{k1}=\sigma(W_{1}s+b_{1})\cdot X_{k1}$ (2)

where $W_{1}\in\mathbb{R}^{C/2G\times 1\times 1}$ and
$b_{1}\in\mathbb{R}^{C/2G\times 1\times 1}$ are parameters used to scale and
shift $s$.

Fig. 3: PyTorch code of the proposed SA module Fig. 4: Validation on the
effectiveness of SA.

Spatial Attention. Different from channel attention, spatial attention focuses
on “where” is an informative part, which is complementary to channel
attention. First, we use Group Norm (GN)[25] over $X_{k2}$ to obtain spatial-
wise statistics. Then, $\mathcal{F}_{c}(\cdot)$ is adopted to enhance the
representation of $\hat{X}_{k2}$. The final output of spatial attention is
obtained by

$X^{\prime}_{k2}=\sigma(W_{2}\cdot GN(X_{k2})+b_{2})\cdot X_{k2}$ (3)

where $W_{2}$ and $b_{2}$ are parameters with shape $\mathbb{R}^{C/2G\times
1\times 1}$.

Then the two branches are concatenated to make the number of channels as the
same as the number of input, i.e.
$X^{\prime}_{k}=[X^{\prime}_{k1},X^{\prime}_{k2}]\in\mathbb{R}^{C/G\times
H\times W}$.

Aggregation. After that, all the sub-features are aggregated. And finally,
similar to ShuffleNet v2 [13], we adopt a ”channel shuffle” operator to enable
cross-group information flow along the channel dimension. The final output of
SA module is the same size of $X$, making SA quite easy to be integrated with
modern architectures.

Note that $W_{1},b_{1},W_{2},b_{2}$ and Group Norm hyper-parameters are only
parameters introduced in the proposed SA. In a single SA module, the number of
channels in each branch is $C/2G$. Therefore, the total parameters are $3C/G$
(typically $G$ is 32 or 64), which is trivial compared with the millions of
parameters of the entire network, making SA quite lightweight. The overall
architecture of SA module is illustrated in Figure 2.

Implementation. SA can be easily implemented by a few lines of code in PyTorch
and TensorFlow where automatic differentiation is supported. Figure 3 shows
the code based on PyTorch.

SA-Net for Deep CNNs. For adopting SA into deep CNNs, we exploit exactly the
same configuration with SENet [23] and just replace SE block with SA module.
The generated networks are named as SA-Net.

Fig. 5: Sample visualization on ImageNet-1k val split generated by GradCAM.
All target layer selected is “layer4.2”.

###  3.1 Visualization and Interpretation

In order to verify whether SA can improve semantic feature representation by
utilizing feature grouping and channel shuffle, we first train SA-Net50B
(without “channel shuffle”) and SA-Net50 (with “channel shuffle”) on the
ImageNet-1k training set. Assume $I$ is the original input, we calculate Top-1
accuracy of $I\times X_{k}$ in each group at SA_5_3 (i.e., the last bottleneck
according to the following scheme: SA_stageID_blockID.) before and after using
of the SA module. We use the accuracy scores as indicator and plot the
distribution for different classes (“pug”,“goldfish” and “plane”) in the
ImageNet-1k validation set in Figure 4(a, b). For comparison, we also plot the
distribution of classification scores across all the 1000 classes.

As shown in Figure 4 (a, b), the Top-1 accuracy statistically increases after
SA, which means feature grouping can significantly enhance the semantic
representations of feature maps. In addition, the average score in each group
gains ($\approx 0.4\%$) with “channel shuffle”, which demonstrates the
effectiveness of “channel shuffle”.

To fully validate the effectiveness of SA, we plot the distribution of average
activations (the mean value of the channel-wise feature maps in each group,
similar to SE) across three classes (“pug”, “goldfish”, and “plane”) at
different depths in SA-Net50 (with shuffle). The results are shown in Figure
4(c). We make some observations about the role of SA module: (1) the
distribution across different classes is very similar to each other at the
earlier layers(e.g., SA_2_3 and SA_3_4), which suggests that the importance of
feature groups is likely to be shared by different classes in the early
stages; (2) at greater depths, the activation of each group becomes much more
class-specific as different classes exhibit different performance to the
discriminative value of features(e.g., SA_4_6 and SA_5_3); (3) SA_5_2 exhibits
a similar pattern over different classes, which means SA_5_2 is less important
than other blocks in providing recalibration to the network.

In order to validate the effectiveness of SA more intuitively, we sample 9
images from ImageNet-1k val split. We use GradCAM [26] to visualize their
heatmaps at SA_5_3 on SA-Net50. For comparison, we also draw their heatmaps of
ResNet50 at “layer4.2”. As shown in Figure 5, our proposed SA module allows
the classification model to focus on more relevant regions with more object
details, which means the SA module can effectively improve the classification
accuracy.

Therefore, the proposed SA module is validated to indeed enhance the
representation power of networks.

##  4 Experiments

Experiment Setup. All experiments are conducted with exactly the same data
augmentation and hyper-parameters settings in [17] and [23]. Specifically, the
input images are randomly cropped to $224\times 224$ with random horizontal
flipping. The numbers of groups $G$ in SA-Net and SGE-Net are both set 64. The
initialization of parameters in $\mathcal{F}_{c}(\cdot)$ are set to 0 for
weights ($W_{1}$ and $W_{2}$) and 1 for biases ($b_{1}$ and $b_{2}$) to obtain
better results. All the architectures are trained from scratch by SGD with
weight decay 1e-4, momentum 0.9, and mini-batch size 256 (using 8 GPUs with 32
images per GPU) for 100 epochs, starting from the initial learning rate 0.1
(with a linear warm-up [27] of 5 epochs) and decreasing it by a factor of 10
every 30 epochs. For the testing on the validation set, the shorter side of an
input image is first resized to 256, and a center crop of $224\times 224$ is
used for evaluation.

Table 1: Comparisons of different attention methods on ImageNet-1k in terms of network parameters (Param.), giga floating point operations per second (GFLOPs), and Top-1/Top-5 accuracy (in %). The best records and the improvements are marked in bold and $\uparrow$, respectively. Attention Methods | Backbones | Param. | GFLOPs | Top-1 Acc (%) | Top-5 Acc (%)
---|---|---|---|---|---
ResNet [17] | ResNet-50 | 25.557M | 4.122 | 76.384 | 92.908
SENet [23] | 28.088M | 4.130 | 77.462 | 93.696
CBAM [10] | 28.090M | 4.139 | 77.626 | 93.660
SGE-Net [12] | 25.559M | 4.127 | 77.584 | 93.664
ECA-Net [wang2020eca] | 25.557M | 4.127 | 77.480 | 93.680
SA-Net (Ours) | 25.557M | 4.125 | 77.724 ($\uparrow 1.34$) | 93.798 ($\uparrow 0.89$)
ResNet [17] | ResNet-101 | 44.549M | 7.849 | 78.200 | 93.906
SENet [23] | 49.327M | 7.863 | 78.468 | 94.102
CBAM [10] | 49.330M | 7.879 | 78.354 | 94.064
SGE-Net [12] | 44.553M | 7.858 | 78.798 | 94.368
ECA-Net [wang2020eca] | 44.549M | 7.858 | 78.650 | 94.340
SA-Net (Ours) | 44.551M | 7.854 | 78.960 ($\uparrow 0.76$) | 94.492 ($\uparrow 0.59$)

Classification on ImageNet-1k. We compare SA-Net with the current SOTA
attention methods. Evaluation metrics include both efficiency(i.e., network
parameters and GFLOPs) and effectiveness(i.e., Top-1/Top-5 accuracy). As shown
in Table 1, SA-Net shares almost the same model complexity (i.e., network
parameters and FLOPs) with the original ResNet [17], but achieves 1.34% gains
in terms of Top-1 accuracy and 0.89% advantages over Top-5 accuracy (on
ResNet-50). When using ResNet-101 as the backbone, the performance gains are
0.76% and 0.59%, respectively. Comparing with SOTA counterparts, SA obtains
higher accuracy while benefiting lower or similar model complexity.
Specifically, when using ResNet-101 as backbone, the recently SOTA SE [23]
module increases 4.778M parameters, 14.34FLOPs, and 0.268 top-1 accuracy, but
our SA only increase 0.002M parameters, 5.12FLOPs, and top-1 0.76 accuracy,
which can demonstrate that SA is lighter and more efficient.

Table 2: Performance comparisons of SA-Net (using ResNet-50 as backbones) with four options (i.e., eliminating Group Norm, eliminating Channel Shuffle, eliminating $\mathcal{F}_{c}(\cdot)$ and utilizing $1\times 1$ Conv to replace $\mathcal{F}_{c}(\cdot)$) on ImageNet-1k in terms of GFLOPs and Top-1/Top-5 accuracy (in %). The best records are marked in bold.  Methods | GFLOPs | Top-1 Acc (%) | Top-5 Acc (%)
---|---|---|---
origin | 4.125 | 77.724 | 93.798
w/o_gn | 4.125 | 77.372 | 93.804
w/o_shuffle | 4.125 | 77.598 | 93.758
w/o_$\mathcal{F}_{c}(\cdot)$ | 4.125 | 77.608 | 93.886
$1\times 1$ Conv | 4.140 | 77.684 | 93.840

Ablation Study. We report the ablation studies of SA-Net50 on ImageNet-1k, to
thoroughly investigate the components of the SA. As shown in Table 2, the
performance drops significantly when eliminating Group Norm, which indicates
the distribution of features generated by different samples from the same
semantic group is inconsistent. It is difficult to learn robust significant
coefficients without normalization. When eliminating channel shuffle, the
performance drops a little, which demonstrates that information communication
among different groups can enhance the features representation. There is no
doubt if $\mathcal{F}_{c}(\cdot)$ is eliminated, the performance drops since
$\mathcal{F}_{c}(\cdot)$ is adopted to enhance the representation of features.
However, the performance does not improve if we replace
$\mathcal{F}_{c}(\cdot)$ with $1\times 1$ Conv. This may because the number of
channels in each sub-feature is too few, so that it is unnecessary to exchange
information among different channels.

Object Detection on MS COCO. We further train the current SOTA detectors on
the COCO train2017 split, and evaluate bounding box Average Precision (AP) for
object detection, and mask AP for instance segmentation. We implement all
detectors using the MMDetection toolkit with the default settings and trained
them within 12 epochs(namely, ‘1× schedule’). For a fair comparison, we only
replace the pre-trained backbones on ImageNet-1k and transfer them to MS COCO
by fine-tuning, keeping the other components in the entire detector intact. As
shown in Table 3, integration of either SE block or the proposed SA module can
improve the performance of object detection using either one-stage or two-
stage detector by a clear margin. Meanwhile, our SA can outperform SE block
with lower model complexity. Specifically, adopting Faster R-CNN [28] as the
basic detector, SA outperforms SE by 1.0% and 1.4% in terms of AP by using
ResNet-50 and ResNet-101, respectively. If we use RetinaNet as the base
detector, the gains both increase by 1.5%.

Table 3: Object detection results of different attention methods on COCO val2017. The best records and the improvements are marked in bold and $\uparrow$, respectively.  Methods | Detectors | Param. | GFLOPs | AP50:95 | AP50 | AP75 | APS | APM | APL
---|---|---|---|---|---|---|---|---|---
ResNet-50 | Faster R-CNN | 41.53M | 207.07 | 36.4 | 58.4 | 39.1 | 21.5 | 40.0 | 46.6
\+ SE | 44.02M | 207.18 | 37.7 | 60.1 | 40.9 | 22.9 | 41.9 | 48.2
\+ SA (Ours) | 41.53M | 207.35 | 38.7 ($\uparrow 2.3$) | 61.2 | 41.4 | 22.3 | 42.5 | 49.8
ResNet-101 | 60.52M | 283.14 | 38.5 | 60.3 | 41.6 | 22.3 | 43.0 | 49.8
\+ SE | 65.24M | 283.33 | 39.6 | 62.0 | 43.1 | 23.7 | 44.0 | 51.4
\+ SA (Ours) | 60.53M | 283.60 | 41.0($\uparrow 2.5$) | 62.7 | 44.8 | 24.4 | 45.1 | 52.5
ResNet-50 | Mask R-CNN | 44.18M | 275.58 | 37.3 | 59.0 | 40.2 | 21.9 | 40.9 | 48.1
\+ SE | 46.67M | 275.69 | 38.7 | 60.9 | 42.1 | 23.4 | 42.7 | 50.0
\+ SA (Ours) | 44.18M | 275.86 | 39.4($\uparrow 2.1$) | 61.5 | 42.6 | 23.4 | 42.8 | 51.1
ResNet-101 | 63.17M | 351.65 | 39.4 | 60.9 | 43.3 | 23.0 | 43.7 | 51.4
\+ SE | 67.89M | 351.84 | 40.7 | 62.5 | 44.3 | 23.9 | 45.2 | 52.8
\+ SA (Ours) | 63.17M | 352.10 | 41.6($\uparrow 2.2$) | 63.0 | 45.5 | 24.9 | 45.5 | 54.2
ResNet-50 | RetinaNet | 37.74M | 239.32 | 35.6 | 55.5 | 38.3 | 20.0 | 39.6 | 46.8
\+ SE | 40.25M | 239.43 | 36.0 | 56.7 | 38.3 | 20.5 | 39.7 | 47.7
\+ SA (Ours) | 37.74M | 239.60 | 37.5($\uparrow 1.9$) | 58.5 | 39.7 | 21.3 | 41.2 | 45.9
ResNet-101 | 56.74M | 315.39 | 37.7 | 57.5 | 40.4 | 21.1 | 42.2 | 49.5
\+ SE | 61.49M | 315.58 | 38.8 | 59.3 | 41.7 | 22.1 | 43.2 | 51.5
\+ SA (Ours) | 56.64M | 315.85 | 40.3($\uparrow 2.6$) | 61.2 | 43.2 | 23.2 | 44.4 | 53.5
Table 4: Instance segmentation results of various state-of-the-arts attention modules using Mask R-CNN on COCO val2017.  Methods | AP50:95 | AP50 | AP75 | APS | APM | APL
---|---|---|---|---|---|---
ResNet-50 | 34.2 | 55.9 | 36.2 | 18.2 | 37.5 | 46.3
\+ SE | 35.4 | 57.4 | 37.8 | 17.1 | 38.6 | 51.8
\+ ECA | 35.6 | 58.1 | 37.7 | 17.6 | 39.0 | 51.8
\+ SGE | 34.9 | 56.9 | 37.0 | 19.1 | 38.4 | 47.3
\+ SA (Ours) | 36.1($\uparrow 1.9$) | 58.7 | 38.2 | 19.4 | 39.4 | 49.0
ResNet-101 | 35.9 | 57.7 | 38.4 | 19.2 | 39.7 | 49.7
\+ SE | 36.8 | 59.3 | 39.2 | 17.2 | 40.3 | 53.6
\+ ECA | 37.4 | 59.9 | 39.8 | 18.1 | 41.1 | 54.1
\+ SGE | 36.9 | 59.3 | 39.4 | 20.0 | 40.8 | 50.1
\+ SA (Ours) | 38.0($\uparrow 2.1$) | 60.0 | 40.3 | 20.8 | 41.2 | 51.7

Instance Segmentation on MS COCO. Instance segmentation using Mask R-CNN [29]
on MS COCO are shown in Table 4. As shown in Table 4, SA module achieves
clearly improvement over the original ResNet and performs better than other
state-of-the-arts attention modules (i.e., SE block, ECA module, and SGE
unit), with less model complexity. Particularly, SA module achieves more gains
for small objects, which are usually more difficult to be correctly detected
and segmented. These results verify our SA module has good generalization for
various computer vision tasks.

##  5 Conclusion

In this paper, we propose a novel efficient attention module SA to enhance the
representation power of CNN networks. SA first groups the dimensions of
channel into multiple sub-features before processing them in parallel. Then,
for each sub-feature, SA utilizes a Shuffle Unit to capture feature
dependencies in both spatial and channel dimensions. Afterward, all sub-
features are aggregated, and finally, we adopt a “channel shuffle” operator to
make information communication between different sub-features. Experimental
results demonstrate that our SA is an extremely lightweight plug-and-play
block, which can significantly improve the performance of various deep CNN
architectures.

In the future, we will further explore the spatial and channel attention
modules of SA and adopt them into more CNN architectures, including the
ShuffleNet family, the SKNet [2], and MobileNetV3 [20].

## References

  * [1] Yue Cao, Jiarui Xu, Stephen Lin, Fangyun Wei, and Han Hu,  “Gcnet: Non-local networks meet squeeze-excitation networks and beyond,”  CoRR, vol. abs/1904.11492, 2019.
  * [2] Xiang Li, Wenhai Wang, Xiaolin Hu, and Jian Yang,  “Selective kernel networks,”  in IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2019, Long Beach, CA, USA, June 16-20, 2019, 2019, pp. 510–519.
  * [3] Jun Fu, Jing Liu, Haijie Tian, Yong Li, Yongjun Bao, Zhiwei Fang, and Hanqing Lu,  “Dual attention network for scene segmentation,”  in IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2019, Long Beach, CA, USA, June 16-20, 2019, 2019, pp. 3146–3154.
  * [4] Yuhui Yuan and Jingdong Wang,  “Ocnet: Object context network for scene parsing,”  CoRR, vol. abs/1809.00916, 2018.
  * [5] HyunJae Lee, Hyo-Eun Kim, and Hyeonseob Nam,  “SRM : A style-based recalibration module for convolutional neural networks,”  CoRR, vol. abs/1903.10829, 2019.
  * [6] Hengshuang Zhao, Yi Zhang, Shu Liu, Jianping Shi, Chen Change Loy, Dahua Lin, and Jiaya Jia,  “Psanet: Point-wise spatial attention network for scene parsing,”  in Computer Vision - ECCV 2018 - 15th European Conference, Munich, Germany, September 8-14, 2018, Proceedings, Part IX, 2018, pp. 270–286.
  * [7] Xia Li, Zhisheng Zhong, Jianlong Wu, Yibo Yang, Zhouchen Lin, and Hong Liu,  “Expectation-maximization attention networks for semantic segmentation,”  CoRR, vol. abs/1907.13426, 2019.
  * [8] Zhen Zhu, Mengde Xu, Song Bai, Tengteng Huang, and Xiang Bai,  “Asymmetric non-local neural networks for semantic segmentation,”  CoRR, vol. abs/1908.07678, 2019.
  * [9] Zilong Huang, Xinggang Wang, Lichao Huang, Chang Huang, Yunchao Wei, and Wenyu Liu,  “Ccnet: Criss-cross attention for semantic segmentation,”  CoRR, vol. abs/1811.11721, 2018.
  * [10] Sanghyun Woo, Jongchan Park, Joon-Young Lee, and In So Kweon,  “CBAM: convolutional block attention module,”  in Computer Vision - ECCV 2018 - 15th European Conference, Munich, Germany, September 8-14, 2018, Proceedings, Part VII, 2018, pp. 3–19.
  * [11] Qilong Wang, Banggu Wu, Pengfei Zhu, Peihua Li, Wangmeng Zuo, and Qinghua Hu,  “Eca-net: Efficient channel attention for deep convolutional neural networks,”  in 2020 IEEE/CVF Conference on Computer Vision and Pattern Recognition, CVPR 2020, Seattle, WA, USA, June 13-19, 2020. 2020, pp. 11531–11539, IEEE.
  * [12] Xiang Li, Xiaolin Hu, and Jian Yang,  “Spatial group-wise enhance: Improving semantic feature learning in convolutional networks,”  CoRR, vol. abs/1905.09646, 2019.
  * [13] Ningning Ma, Xiangyu Zhang, Hai-Tao Zheng, and Jian Sun,  “Shufflenet V2: practical guidelines for efficient CNN architecture design,”  in Computer Vision - ECCV 2018 - 15th European Conference, Munich, Germany, September 8-14, 2018, Proceedings, Part XIV, 2018, pp. 122–138.
  * [14] Xiangyu Zhang, Xinyu Zhou, Mengxiao Lin, and Jian Sun,  “Shufflenet: An extremely efficient convolutional neural network for mobile devices,”  in 2018 IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2018, Salt Lake City, UT, USA, June 18-22, 2018. 2018, pp. 6848–6856, IEEE Computer Society.
  * [15] Christian Szegedy, Vincent Vanhoucke, Sergey Ioffe, Jonathon Shlens, and Zbigniew Wojna,  “Rethinking the inception architecture for computer vision,”  in 2016 IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2016, Las Vegas, NV, USA, June 27-30, 2016, 2016, pp. 2818–2826.
  * [16] Christian Szegedy, Sergey Ioffe, Vincent Vanhoucke, and Alexander A. Alemi,  “Inception-v4, inception-resnet and the impact of residual connections on learning,”  in Proceedings of the Thirty-First AAAI Conference on Artificial Intelligence, February 4-9, 2017, San Francisco, California, USA, 2017, pp. 4278–4284.
  * [17] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun,  “Deep residual learning for image recognition,”  in 2016 IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2016, Las Vegas, NV, USA, June 27-30, 2016, 2016, pp. 770–778.
  * [18] Alex Krizhevsky, Ilya Sutskever, and Geoffrey E. Hinton,  “Imagenet classification with deep convolutional neural networks,”  in Advances in Neural Information Processing Systems 25: 26th Annual Conference on Neural Information Processing Systems 2012. Proceedings of a meeting held December 3-6, 2012, Lake Tahoe, Nevada, United States, 2012, pp. 1106–1114.
  * [19] Mark Sandler, Andrew G. Howard, Menglong Zhu, Andrey Zhmoginov, and Liang-Chieh Chen,  “Mobilenetv2: Inverted residuals and linear bottlenecks,”  in 2018 IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2018, Salt Lake City, UT, USA, June 18-22, 2018, 2018, pp. 4510–4520.
  * [20] Andrew Howard, Ruoming Pang, Hartwig Adam, Quoc V. Le, Mark Sandler, Bo Chen, Weijun Wang, Liang-Chieh Chen, Mingxing Tan, Grace Chu, Vijay Vasudevan, and Yukun Zhu,  “Searching for mobilenetv3,”  in 2019 IEEE/CVF International Conference on Computer Vision, ICCV 2019, Seoul, Korea (South), October 27 - November 2, 2019. 2019, pp. 1314–1324, IEEE.
  * [21] Sara Sabour, Nicholas Frosst, and Geoffrey E. Hinton,  “Dynamic routing between capsules,”  in Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, 4-9 December 2017, Long Beach, CA, USA, 2017, pp. 3856–3866.
  * [22] Geoffrey E. Hinton, Sara Sabour, and Nicholas Frosst,  “Matrix capsules with EM routing,”  in 6th International Conference on Learning Representations, ICLR 2018, Vancouver, BC, Canada, April 30 - May 3, 2018, Conference Track Proceedings. 2018, OpenReview.net.
  * [23] Jie Hu, Li Shen, and Gang Sun,  “Squeeze-and-excitation networks,”  in 2018 IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2018, Salt Lake City, UT, USA, June 18-22, 2018, 2018, pp. 7132–7141.
  * [24] Xiaolong Wang, Ross B. Girshick, Abhinav Gupta, and Kaiming He,  “Non-local neural networks,”  in 2018 IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2018, Salt Lake City, UT, USA, June 18-22, 2018, 2018, pp. 7794–7803.
  * [25] Yuxin Wu and Kaiming He,  “Group normalization,”  in Computer Vision - ECCV 2018 - 15th European Conference, Munich, Germany, September 8-14, 2018, Proceedings, Part XIII, 2018, pp. 3–19.
  * [26] Ramprasaath R. Selvaraju, Michael Cogswell, Abhishek Das, Ramakrishna Vedantam, Devi Parikh, and Dhruv Batra,  “Grad-cam: Visual explanations from deep networks via gradient-based localization,”  in IEEE International Conference on Computer Vision, ICCV 2017, Venice, Italy, October 22-29, 2017. 2017, pp. 618–626, IEEE Computer Society.
  * [27] Priya Goyal, Piotr Dollár, Ross B. Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, and Kaiming He,  “Accurate, large minibatch SGD: training imagenet in 1 hour,”  CoRR, vol. abs/1706.02677, 2017.
  * [28] Shaoqing Ren, Kaiming He, Ross B. Girshick, and Jian Sun,  “Faster R-CNN: towards real-time object detection with region proposal networks,”  in NIPS 2015, December 7-12, 2015, Montreal, Quebec, Canada, Corinna Cortes, Neil D. Lawrence, Daniel D. Lee, Masashi Sugiyama, and Roman Garnett, Eds., 2015, pp. 91–99.
  * [29] Kaiming He, Georgia Gkioxari, Piotr Dollár, and Ross B. Girshick,  “Mask R-CNN,”  in IEEE International Conference on Computer Vision, ICCV 2017, Venice, Italy, October 22-29, 2017, 2017, pp. 2980–2988.