\documentclass[runningheads]{llncs}
\usepackage{graphicx}
\usepackage{url}
\usepackage{tikz}
\usepackage{comment}
\usepackage{amsmath,amssymb} %
\usepackage{color}
\begin{document}
\pagestyle{headings}
\mainmatter

\title{Feedback Attention for Cell Image Segmentation} %
\titlerunning{Feedback Attention for Cell Image Segmentation}
\author{Hiroki Tsuda \and
Eisuke Shibuya \and
Kazuhiro Hotta}
\authorrunning{H. Tsuda et al.}
\institute{Meijo University, 1-501 Shiogamaguchi, Tempaku-ku, Nagoya 468-8502, Japan 
\url{http://www1.meijo-u.ac.jp/~kazuhotta/cms_new/} \\
\email{193427019@ccalumni.meijo-u.ac.jp,\\160442066@ccalumni.meijo-u.ac.jp,\\kazuhotta@meijo-u.ac.jp}}
\maketitle

\begin{abstract}
In this paper, we address cell image segmentation task by Feedback Attention mechanism like feedback processing. Unlike conventional neural network models of feedforward processing, we focused on the feedback processing in human brain and assumed that the network learns like a human by connecting feature maps from deep layers to shallow layers. We propose some Feedback Attentions which imitate human brain and feeds back the feature maps of output layer to close layer to the input. U-Net with Feedback Attention showed better result than the conventional methods using only feedforward processing.

\keywords{Cell Image, Semantic Segmentation, Attention Mechanism, Feedback Mechanism}
\end{abstract}

\section{Introduction}
\label{sec:intro}

Deep neural networks has achieved state-of-the-art performance in image classification~\cite{alexnet}, segmentation~\cite{fcn}, detection~\cite{faster-rcnn}, and tracking~\cite{siamesefc}. Since the advent of AlexNet~\cite{alexnet}, several Convolutional Neural Network (CNN)~\cite{lecun1998gradient} has been proposed such as VGG~\cite{vgg}, ResNet~\cite{resnet}, Deeplabv3+~\cite{deeplabv3plus}, Faster R-CNN~\cite{faster-rcnn}, and Siamese FC~\cite{siamesefc}. These networks are feedfoward processing. Neural network is mathematical model of neurons~\cite{widrow1998perceptrons} that imitate the structure of human brain. Human brain performs not only feedfoward processing from shallow layers to deep layers of neurons, but also feedback processing from deep layers to shallow layers. However, conventional neural networks consist of only feedfoward processing from shallow layers to deep layers, and do not use feedback processing to connect from deep layers to shallow layers. Therefore, in this paper, we propose some Feedback Attention methods using position attention mechanism and feedback process.

Semantic segmentation assigns class labels to all pixels in an image. The study of this task can be applied to various fields such as automatic driving \cite{camvid,cordts2016cityscapes}, cartography \cite{ghamisi2014feature,maggiori2016convolutional} and cell biology \cite{sstem,imanishi2018novel,unet}. In particular, cell image segmentation requires better results in order to ensure that cell biologists can perform many experiments at the same time.

In addition, overall time and cost savings are expected to be achieved by automated processing without human involvement to reduce human error. Manual segmentation by human experts is slow to process and burdensome, and there is a significant demand for algorithms that can do the segmentation quickly and accurately without human. However, cell image segmentation is a difficult task because the number of supervised images is smaller and there is not regularity compared to the other datasets such as automatic driving. A large number of supervised images requires expert labeling which takes a lot of effort, cost and time. Therefore, it is necessary to enhance the segmentation ability for pixel-level recognition with small number of training images. 

Most of the semantic segmentation approaches are based on Fully Convolutional Network (FCN)~\cite{fcn}. FCN is composed of some convolutional layers and some pooling layers, which does not require some fully connected layers. Convolutional layer and pooling layer reproduce the workings of neurons in the visual cortex. These are proposed in Neocognitron~\cite{fukushima1982neocognitron} which is the predecessor of CNN. Convolutional layer which is called S-cell extracts local features of the input. Pooling layer which is called C-cell compresses the information to enable downsampling to obtain position invariance. Thus, by repeating the feature extraction by convolutional layer and the local position invariance by pooling layer, robust pattern recognition is possible because it can react only to the difference of shape without much influence of misalignment and size change of the input pattern. Only the difference between CNN and Neocognitron is the optimization method, and the basic elements of both are same structure.

We focused on the relationship between the feature map close to the input and output of the semantic segmentation, and considered that it is possible to extract effective features by using between the same size and number of channels in the feature maps close to the input and output. 
In this paper, we create an attention map based on the relationship between these different feature maps, and a new attention mechanism is used to generate segmentation results. We can put long-range dependent spatial information from the output into the feature map of the input. The attention mechanism is fed back into the feature map of the input to create a model that can be reconsidered in based on the output.

In experiments, we evaluate the proposed method on a cell image datasets ~\cite{sstem}. We confirmed that the proposed method gave higher accuracy than conventional method. We evaluate our method by some ablation studies and show the effectiveness of our method.

This paper is organized as follows. In section~\ref{sec:related}, we describe related works. The details of the proposed method are
explained in section~\ref{sec:proposed}. In section~\ref{sec:experments}, we evaluate our proposed method on segmentation of cell images. Finally, we describe conclusion and future works in section~\ref{sec:conclusions}.

\section{Related works}
\label{sec:related}
\subsection{Semantic Segmentation}
\label{sec:related:seg}

FCNs~\cite{fcn} based methods have achieved significant results for semantic segmentation. The original FCN used stride convolutions and pooling to gradually downsize the feature map, and finally created high-dimensional feature map with low-resolution. This feature map has semantic information but  fine information such as fine objects and correct location are lost. Thus, if the upsampling is used at the final layer, the accuracy is not sufficient. Therefore, encoder-decoder structure is usually used in semantic segmentation to obtain a final feature map with high-resolution. It consists of encoder network that extracts features from input image using convolutional layers, pooling layers, and batch normalization layers, and decoder network that classifies the extracted feature map by upsampling, convolutional layers, and batch normalization layers. Decoder restores the low-resolution semantic feature map extracted by encoder and middle-level features to the original image to compensate for the lost spatial information, and obtains a feature map with high resolution semantic information.

SegNet~\cite{segnet} is a typical network of encoder-decoder structures. Encoder uses 13 layers of VGG16~\cite{vgg}, and decoder receives some indexes selected by max pooling of encoder. In this way, decoder complements the positional information when upsampling and accelerates the calculation by unpooling, which requires no training.

Another famous encoder-decoder structural model is U-net~\cite{unet}. The most important characteristic of U-Net is skip connection between encoder and decoder. The feature map with the spatial information of encoder is connected to the restored feature map of the decoder. This complements the high-resolution information and improves the resolution so that labels can be assigned more accurately to each pixel. In addition, deconvolution is used for up-sampling in decoder.
\subsection{Attention Mechanism}
\label{sec:related:attention}

Attention mechanism is an application of the human attention mechanism to machine learning. It has been used in computer vision and natural language processing. In the field of image recognition, important parts or channels are emphasized.

Residual Attention Network \cite{wang2017residual} introduced a stack network structure composed of multiple attention components, and attention residual learning applied residual learning \cite{resnet} to the attention mechanism. Squeeze-and-Excitation Network (SENet) \cite{senet} introduced an attention mechanism that adaptively emphasizes important channels in feature maps. Accuracy booster blocks \cite{accuracy-booster} and efficient channel attention module \cite{wang2019eca} made further improvements by changing the fully-connected layer in SENet.
Attention Branch Network \cite{fukui2019abn} is Class Activation Mapping (CAM) \cite{cam} based structure to build visual attention maps for image classification.
Transformer \cite{transformer} performed language translation only with the attention mechanism. There are Self-Attention that uses the same tensor, and Source-Target-Attention that uses two different tensors. 
Several networks have been proposed that use Self-Attention to learn the similarity between pixels in feature maps \cite{fu2019dual,huang2019ccnet,stand-alone,wang2018non,sagan}. 
\subsection{Feedback Mechanism using Recurrent Neural Networks}
\label{sec:related:recurrent}

Feedback is a fundamental mechanism of the human perceptual system and is expected to develop in the computer vision in the future. There have been several approaches to feedback using recurrent neural networks (RNNs)~\cite{alom2018recurrent,han2018image,zamir2017feedback}. 

Feedback Network~\cite{zamir2017feedback} uses convLSTM~\cite{xingjian2015convlstm} to acquire hidden states with high-level information and provide feedback with the input image. However, this is intended to solve the image classification task and is not directly applicable to the segmentation task.

RU-Net~\cite{alom2018recurrent} consists of a U-Net~\cite{unet} and a recurrent neural network, where each convolutional layer is replaced by recurrent convolutional layer~\cite{liang2015recurrent}. The accumulation of feature information at each scale by the recurrent convolutional layer gives better results than the standard convolutional layer. However, this is not strictly feedback but the deepening of network.

Feedback U-Net~\cite{shibuya2020feedback} is the segmentation method using convLSTM~\cite{xingjian2015convlstm}. The probability for segmentation at final layer is used as the input image for segmentation at the second round, while the first feature map is used as the hidden state for the second segmentation to provide feedback.

Since RNNs is a neural network that contains loop connections, it can be easily used for feedback mechanisms. However, the problem with RNNs is that the amount of operations increases drastically and a lot of memory is consumed, which makes processing difficult and often results in the phenomenon that information is not transmitted. Thus, we applied RNNs-free feedback mechanism to U-Net,  and excellent performance is shown by the feedback attention mechanism on the segmentation task.

\section{Proposed Method}
\label{sec:proposed}

This section describes the details of the proposed method. Section~\ref{sec:proposed:details} outlines the network of our method. In section~\ref{sec:proposed:feedback}, we describe the details of the proposed attention mechanism.

\begin{figure}[t]
    \centering
    \includegraphics[width=\linewidth]{images/proposed_method.png}
    \caption{Network structure of the proposed method using Feedback Attention}
    \label{fig:proposed}
\end{figure}

\subsection{Network Structure Details}
\label{sec:proposed:details}

The proposed method is based on U-Net~\cite{unet}, which is used as a standard in medical and cell images. Figure~\ref{fig:proposed} shows the detail network structure of our proposed method using U-net. We design to do segmentation twice using U-Net in order to use the feature maps in input and output. Since the proposed method uses the feature maps of input and output, we use the model twice with shared weights. First, we perform segmentation by U-Net to obtain high-resolution important feature maps at the final layer. Then, we connect to Feedback Attention to a feature map that is close to the input with the same size and number of channels as this feature map. In this case, we use the input feature map that was processed two times by convolution. 
The reason is that a feature map convolved twice can extract more advanced features than a feature map convolved once. The details of Feedback Attention is explained in section~\ref{sec:proposed:feedback}. By applying Attention between the feature maps of input and output, we can obtain an input that takes the output into account as feedback control. In training, U-Net is updated by using only the gradients at the second round using feedback attention. In addition, the loss function is trained using Softmax cross-entropy loss.
\begin{figure}[t]
    \centering
    \includegraphics[width=\linewidth]{images/Source-Target.png}\\
    \text{(a) Source-Target-Attention}
    
    \centering
    \includegraphics[width=\linewidth]{images/Self.png}\\
    \text{(b) Self-Attention}
    \caption{Feedback Attention}
    \label{fig:feedback attention}
\end{figure}

\subsection{Feedback Attention}
\label{sec:proposed:feedback}

We propose two kinds of Feedback Attentions to aggregate feature maps with the shape of $C \times H \times W$. Figure~\ref{fig:feedback attention} (a) shows the Source-Target-Attention method that directly aggregates similar features between the feature maps of input and output. Figure~\ref{fig:feedback attention} (b) shows the self-attention method that performs self-attention for output feature map and finally adds it to the feature map of input. Both Feedback Attentions are explained in the following subsections. 

\subsubsection{Feedback Attention using Source-Target-Attention}
\label{sec:proposed:feedback:st}

We use Source-Target-Attention to aggregate the correlation between feature maps based on the relationship between input and output. Since the feature map in the final layer close to the output contains all the information for judging, it can be fed back using attention and effectively extract features again from the shallow input layer. We elaborate the process to aggregate each feature map.

As shown in Figure~\ref{fig:feedback attention} (a), we feed the feature maps of input or output into $1\times1$ convolutions and batch normalization to generate two new feature maps \textbf{Query} and \textbf{Key}, respectively, we are inspired by Self-Attention GAN (SAGAN) \cite{sagan} to reduce the channel number to $C/8$ for memory efficiency. Then, we reshape them to $C/8 \times (H \times W)$. After we perform a matrix multiplication between the transpose of \textbf{Query} and \textbf{Key}, and we use a softmax function to calculate the attention map. Attention map in vector form is as follows.

\begin{equation}
w_{ij}=\frac{1}{Z_i}\exp({Query}_{i}^T ~{Key}_{j}),
\end{equation}
where $w_{ij}$ measures the $i^{th}$ \textbf{Query}'s impact on $j^{th}$ \textbf{Key}. $Z_i$ is the sum of similarity scores as 
\begin{equation}
Z_{i}={\sum_{j=1}^{H \times W} {\exp({Query}_{i}^T ~{Key}_{j})}},
\end{equation}
where $H \times W$ is the total number of pixels in \textbf{Query}. 
By increasing the correlation between two locations, we can create an attention map that takes into account output's feature map. 

On the other hand, we feed the feature map of output into $1\times1$ convolution and batch normalization to generate a new feature map \textbf{Value} and reshape it to $C/2 \times (H \times W)$. Then, we perform a matrix multiplication between attention map and the transpose of \textbf{Value} and reshape the result to $C/2 \times H \times W$.
In addition, we feed the new feature map into $1\times1$ convolution and batch normalization to generate feature map the same size as the feature map of input $C \times H \times W$. Finally, we multiply it by a scale parameter $\alpha$ and perform a element-wise sum operation with the input feature map to obtain the final output as follows.

\begin{equation}
\label{source-target}
A_i=\alpha \sum_{j=1}^{H \times W}{(w_{ij}~ Value_j^T)^T+F_i},
\end{equation}
where $\alpha$ is initialized as 0 and gradually learns to assign more weight \cite{sagan}. $A_i$ indicates the feedbacked output and $F_i$ indicates the feature map of the input. By adding $\alpha \sum_{j=1}^{H \times W}(w_{ij}~ Value_j^T)^T$ to the feature map close to input, we can get the feature map considering feature map of output. The new feature map $A_i$ is fed into the network again, and we obtain the segmentation result. 
From Equation~(\ref{source-target}), it can be inferred that the output $A_i$ is the weighted sum of all positions in output and the feature map of input.  Therefore, the segmentation accuracy is improved by transmitting the information of the output to the input.
\subsubsection{Feedback Attention using Self-Attention}
\label{sec:proposed:feedback:self}

In Source-Target-Attention, the feature map between input and output is aggregated. Thus, the relationship between each feature map can be emphasized.  However, the feature map of the input may not extract enough information and therefore may result in poorly relational coordination. We construct Feedback Attention using Self-Attention that aggregates only the feature map of output.

The structure is shown in Figure~\ref{fig:feedback attention} (b). We feed the feature maps of output into $1\times1$ convolution and batch normalization to generate new feature maps \textbf{Query}, \textbf{Key} and \textbf{Value}. This is similar to Source-Target-Attention. We also reshape \textbf{Query} and \textbf{Key} to $C/8 \times (H \times W)$. Then, we perform a matrix multiplication between the transpose of \textbf{Query} and \textbf{Key}, and use a softmax function to calculate the attention map. Attention map in vector form is as follows.

\begin{equation}
w_{pq}=\frac{\exp({Query}_{p}^T ~{Key}_{q})}{\sum_{q=1}^{H \times W} {\exp({Query}_{p}^T~{Key}_{q})}},
\end{equation}
where $w_{pq}$ measures the $p^{th}$ \textbf{Query}'s impact on $q^{th}$ \textbf{Key}.

We reshape \textbf{Value} to $C/2 \times (H \times W)$. Then, we perform a matrix multiplication between attention map and the transpose of \textbf{Value} and reshape the result to $C \times H \times W$ after $1 \times 1$ convolution. Finally, we multiply it by a scale parameter $\beta$ and perform a element-wise sum operation with the feature maps of input to obtain the final output as follows.

\begin{equation}
\label{self}
A_p=\beta \sum_{q=1}^{H \times W}{(w_{pq}~ Value_q^T)^T+F_p},
\end{equation}
where $\beta$ is initialized as 0 and gradually learns to assign more weight \cite{sagan}. $A_p$ indicates the output, $F_p$ indicates the feature map of input. New feature map $A_p$ is fed into the network again, and we obtain the segmentation result. 

Unlike Equation~(\ref{source-target}), Equation~(\ref{self}) calculates the similarity using only the information of output. In addition, consistency can be improved because information can be selectively passed to the input by the scale parameter.

\begin{figure}[t]
  \centering
    \begin{tabular}{c}
 
      \begin{minipage}[t]{0.19\hsize}
      \centering
          \includegraphics[width=\linewidth]{drosophila/drosophila_input.png}
          \text{Input image}
      \end{minipage}

      \begin{minipage}[t]{0.19\hsize}
        \centering
          \includegraphics[width=\linewidth]{drosophila/drosophila_GT.png}
          \text{Ground truth}
      \end{minipage}

      \begin{minipage}[t]{0.19\hsize}
        \centering
          \includegraphics[width=\linewidth]{drosophila/drosophila_unet.png}
          \text{U-Net\cite{unet}}
      \end{minipage}

      \begin{minipage}[t]{0.19\hsize}
        \centering
          \includegraphics[width=\linewidth]{drosophila/drosophila_feedback_st.png}
          \text{Feedback}\\
          \text{Attention(ST)}
      \end{minipage}
      
      \begin{minipage}[t]{0.19\hsize}
        \centering
          \includegraphics[width=\linewidth]{drosophila/drosophila_feedback_self.png}
          \text{Feedback}\\
          \text{Attention(Self)}
      \end{minipage}
      
    \end{tabular}
    \caption{Examples of segmentation results on ssTEM dataset. ST indicates Source-Target-Attention, Self indicates Self-Attention.}
    \label{fig:sstem}
\end{figure}

\begin{table}[t]
\centering
\caption{Segmentation accuracy (IoU and mIoU) on ssTEM Dataset. ST indicates Source-Target-Attention, Self indicates Self-Attention.}
\label{table:sstem}
\begin{tabular}{l|ccccc}
\hline
Method & Membrane  & Mitochondria  & Synapse  & Cytoplasm & Mean IoU\%  \\ \hline \hline
U-Net\cite{unet} & 74.24 & 71.01 & 43.08 & 92.03 & 70.09 \\
RU-Net\cite{alom2018recurrent} & 75.81 & 74.39 & 43.26 & 92.25 & 71.43 \\
Feedback\\U-Net\cite{shibuya2020feedback} & 76.44 & 75.20 & 42.30 & 92.43 & 71.59 \\
Feedback\\Attention(ST) & {\textbf{76.65}} & {\textbf{78.27}} & {\textbf{43.32}} & {\textbf{92.64}} & {\textbf{72.72}} \\
Feedback\\Attention(Self) & {\color{red} \textbf{76.94}} & {\color{red} \textbf{79.52}} & {\color{red} \textbf{45.29}} & {\color{red} \textbf{92.80}} & {\color{red} \textbf{73.64}} \\\hline
\end{tabular}
\end{table}

\section{Experiments}
\label{sec:experments}

This section shows evaluation results by the proposed method. We explain the datasets used in experiments in section~\ref{sec:experments:dataset}. Experimental results are shown in section~\ref{sec:experments:results}. Finally, section~\ref{sec:experments:ablation studies} describes Ablation studies to demonstrate the effectiveness of the proposed method.

\subsection{Dataset}
\label{sec:experments:dataset}

In experiments, we evaluated all methods 15 times with 5-fold cross-validation using three kinds of initial values on the Drosophila cell image data set \cite{sstem}.  We use Intersection over Union (IoU) as evaluation measure. Average IoU of 15 times evaluations is used as final measure.

This dataset shows neural tissue from a Drosophila larva ventral nerve cord and was acquired using serial section Transmission Electron Microscopy at HHMI Janelia Research Campus \cite{sstem}. This dataset is called ssTEM dataset. There are 20 images of $1024 \times 1024$ pixels and ground truth. In this experiment, semantic segmentation is performed for four classes; membrane, mitochondria, synapses and cytoplasm. We augmented 20 images to 320 images by cropping 16 regions of $256 \times 256$ pixels without overlap from an image. We divided those images into 192 training, 48 validation and 80 test images.
\begin{figure}[t]
  \centering
    \begin{tabular}{c}
 
      \begin{minipage}[t]{0.24\hsize}
      \centering
          \includegraphics[width=\linewidth]{attention/image_in.png}
          \text{Input image}
      \end{minipage}

      \begin{minipage}[t]{0.24\hsize}
      \centering
          \includegraphics[width=\linewidth]{attention/image_gen_st.png}
          \text{Output image(ST)}
      \end{minipage}
      
      \begin{minipage}[t]{0.24\hsize}
      \centering
          \includegraphics[width=\linewidth]{attention/image_st_attention_mem.png}
          \text{Attention map}\\
          \text{Membrane(ST)}
      \end{minipage}
      
      \begin{minipage}[t]{0.24\hsize}
      \centering
          \includegraphics[width=\linewidth]{attention/image_st_attention_cyt.png}
          \text{Attention map}\\
          \text{Cytoplasm(ST)}
      \end{minipage}
      \\
      \begin{minipage}[t]{0.24\hsize}
      \centering
          \includegraphics[width=\linewidth]{attention/mCherry0305.png}
          \text{Ground Truth}
      \end{minipage}
      
      \begin{minipage}[t]{0.24\hsize}
      \centering
          \includegraphics[width=\linewidth]{attention/image_gen_self.png}
          \text{Output image(Self)}
      \end{minipage}
      
      \begin{minipage}[t]{0.24\hsize}
      \centering
          \includegraphics[width=\linewidth]{attention/image_self_attention_mem.png}
          \text{Attention map}\\
          \text{Membrane(Self)}
      \end{minipage}
      
      \begin{minipage}[t]{0.24\hsize}
      \centering
        \includegraphics[width=\linewidth]{attention/image_self_attention_cyt.png}
        \text{Attention map}\\
        \text{Cytoplasm(Self)}
      \end{minipage}
      
      
    \end{tabular}
    \caption{Visualization results of Attention Map on ssTEM dataset. ST indicates Source-Target-Attention, Self indicates Self-Attention.}
    \label{fig:attention maps}
\end{figure}
\subsection{Experimental Results}
\label{sec:experments:results}

Table~\ref{table:sstem} shows the accuracy on ssTEM dataset, and Figure~\ref{fig:sstem} shows the segmentation results. Bold red letters in the Table represent the best IoU and black bold letters represent the second best IoU. Table~\ref{table:sstem} shows that our proposed Feedback Attention  improved the accuracy of all classes compared to conventional U-Net~\cite{unet}. 
We also evaluated two feedback methods using RNNs; RU-Net~\cite{alom2018recurrent} with recurrent convolution applied to U-Net and Feedback U-Net~\cite{shibuya2020feedback} with feedback segmentation applied to U-Net. The result shows that the proposed method gave high accuracy in all classes. In addition, we can see that Self-Attention, which calculates the similarity in the output, is more accurate than Source-Target-Attention which calculates the similarity from the relationship between the input and the output. This indicates that the feature map of the input does not extract enough features and therefore the similarity representation between the input and the output does not work well.

From the yellow frames in Figure~\ref{fig:sstem}, our method using Feedback Attention can identify mitochondria that were detected excessively by conventional methods. In the conventional methods, cell membranes were interrupted, but in our proposed method, we confirm that cell membranes are segmented in such a way that they are cleanly connected. Experimental results show that cell membrane and the mitochondria have been successfully identified even in places where it is difficult to detect by conventional methods.

We visualize some attention maps in Figure~\ref{fig:attention maps} to understand our two kinds of Feedback Attentions. White indicates similarity and black indicates dissimilarity. We find that Self-Attention maps has many similar pixels but Source-Target-Attention maps has fewer pixels. This is because Source-Target-Attention uses the feature maps of input and output, and the feature map near input is different from that of output, so the number of similar pixels are smaller than Self-Attention map. However, the membranes and cytoplasm have different values in the attention map. This means that they are emphasized as different objects. On the other hand, Self-Attention generates attention maps from only the feature map of output. Therefore, as shown in the Figure~\ref{fig:attention maps}, when cell membrane and cytoplasm are selected, they are highlighted as similar pixels.
\begin{table}[t]
\centering
\caption{Comparison of different feedback connections.}
\label{table:connection}
\begin{tabular}{l|ccccc}
\hline
Method & Membrane  & Mitochondria  & Synapse   & Cytoplasm & Mean IoU\%  \\ \hline \hline
Add & 75.56 & 77.36 & 41.84 & 92.46 & 71.81 \\
1$\times$1 Conv & 75.22 & \textbf{78.39} & \textbf{43.46} & 92.49 & 72.39 \\
SE-Net\cite{senet} & 75.89 & 77.31 & 42.92 & 92.49 & 72.15 \\
Light Attention\cite{hiramatsu2020semantic} & 76.20 & 78.27 & 43.18 & 92.57 & 72.56 \\
Feedback\\Attention(ST) & \textbf{76.65} & 78.27 & 43.32 & \textbf{92.64} & \textbf{72.72} \\
Feedback\\Attention(Self) & {\color{red} \textbf{76.94}} & {\color{red} \textbf{79.52}} & {\color{red} \textbf{45.29}} & {\color{red} \textbf{92.80}} & {\color{red} \textbf{73.64}} \\ \hline
\end{tabular}
\end{table}
\subsection{Ablation Studies}
\label{sec:experments:ablation studies}

We performed three ablation studies to show the effectiveness of the proposed method. The first ablation study evaluated the different connection methods. The second ablation study confirmed the effectiveness of connection location from the output to the input. The last ablation study confirmed the effectiveness of before and after Feedback Attention was used.

\subsubsection{Comparison of difference feedback connection}

The effectiveness of the other feedback connection methods from the output to the input was experimentally confirmed. We compare four methods. We compare two methods that do not use the attention mechanism. The first one is that we simply add the feature map in the output to the input. The second one is that we feed the feature map in the output to $1 \times 1$ convolution and then add it to the feature map in the input. Both methods use scale parameter as our propose method.

In addition, we compare two methods using attention mechanism. 
The first one is that we apply SE-Net~\cite{senet}, which suppresses and emphasizes the feature map between channels, to the output feature map, and add it to the input feature map. The second one is that we apply Light Attention~\cite{hiramatsu2020semantic}, which suppresses and emphasizes the important locations and channels in feature map by $3 \times 3$ convolutional processing, to the output feature map and adding it to the input feature map.

From Table~\ref{table:connection}, we can see that the above four methods improve the accuracy from U-Net~\cite{unet} because the feedback mechanism is effective. However, our proposed method is more accurate than those four methods. This shows that our proposed Feedback Attention can use the output's information effectively in the input.

\begin{table}[t]
\centering
\caption{Comparison between different connection locations.}
\label{table:location}
\begin{tabular}{lccccc}
\hline
\multicolumn{1}{l|}{Method} & Membrane  & Mitochondria  & Synapse   & Cytoplasm & Mean IoU\%  \\ \hline \hline
\multicolumn{6}{c}{Feedback Attention using Source-Target-Attention} \\ \hline
\multicolumn{1}{l|}{One conv} & 76.54 & 77.39 & 43.06 & 91.96 & 72.24 \\
\multicolumn{1}{l|}{Two conv(Ours)} & 76.65 & 78.27 & 43.32 & 92.64 & 72.72 \\ \hline
\multicolumn{6}{c}{Feedback Attention using Self-Attention} \\ \hline
\multicolumn{1}{l|}{One conv} & \textbf{76.69} & \textbf{78.73} & \textbf{45.23} & \textbf{92.66} & \textbf{73.33} \\
\multicolumn{1}{l|}{Two conv(Ours)} & {\color{red} \textbf{76.94}} & {\color{red} \textbf{79.52}} & {\color{red} \textbf{45.29}} & {\color{red} \textbf{92.80}} & {\color{red} \textbf{73.64}} \\ \hline
\end{tabular}
\end{table}

\begin{table}[t]
\centering
\caption{Comparison before and after Feedback Attention.}
\label{table:bafore_after}
\begin{tabular}{lccccc}
\hline
\multicolumn{1}{l|}{Method} & Membrane  & Mitochondria  & Synapse   & Cytoplasm & Mean IoU\%  \\ \hline \hline
\multicolumn{6}{c}{Feedback Attention using Source-Target-Attention} \\ \hline
\multicolumn{1}{l|}{First output} & 76.07 & 76.76 & 41.28 & 92.39 & 71.62 \\
\multicolumn{1}{l|}{Second output(Ours)} & \textbf{76.65} & \textbf{78.27} & \textbf{43.32} & \textbf{92.64} & \textbf{72.72} \\ \hline
\multicolumn{6}{c}{Feedback Attention using Self-Attention} \\ \hline
\multicolumn{1}{l|}{First output} & 75.49 & 74.29 & 41.57 & 92.03 & 70.84 \\
\multicolumn{1}{l|}{Second output(Ours)} & {\color{red} \textbf{76.94}} & {\color{red} \textbf{79.52}} & {\color{red} \textbf{45.29}} & {\color{red} \textbf{92.80}} & {\color{red} \textbf{73.64}} \\ \hline
\end{tabular}
\end{table}

\subsubsection{Comparison between different connection locations}

We experimentally evaluated the location of the input feature map which is the destination of feedback. Since the size of feature map should be the same as final layer, the candidates are only two layers close to input. The first one is the feature map closest to the input which is obtained by only one convolution process. The other one is the feature map obtained after convolution is performed two times. We compared the two feature map locations that we use Feedback Attention. 

Table~\ref{table:location} shows that the Feedback Attention to the feature map after two convolution process is better for both Source-Target-Attention and Self-Attention. This indicates that only one convolution process does not extract good features than two convolution processes.
\subsubsection{Comparison before and after Feedback Attention}

When we use Feedback Attention, the output of network is feedback to input as attention. Thus, we get the outputs twice. Although we use the output using Feedback Attention at the second round is used as final result, we compare the results of the outputs at the first and second rounds to show the effectiveness of Feedback Attention. From Table~\ref{table:bafore_after}, the output using Feedback Attention as the second round is better than that at the first round. This demonstrates that the accuracy was improved through the feedback mechanism.

\section{Conclusions}
\label{sec:conclusions}

In this paper, we have proposed two Feedback Attention for cell image segmentation. Feedback Attention allows us to take advantage of the feature map information of the output and improve the accuracy of the segmentation, and segmentation accuracy is improved in comparison with conventional feedforward network, RU-Net~\cite{alom2018recurrent} which uses local feedback at each convolutional layer and Feedback U-Net~\cite{shibuya2020feedback} which uses global feedback between input and output. Ablation studies show that Feedback Attention can obtain accurate segmentation results by choosing the location and attention mechanism that conveys the output information.

In the future, we aim to develop a top-down attention mechanism that directly utilizes ground truth, such as self-distillation~\cite{zhang2019your}. Feedback networks are also categorized as a kind of top-down networks, because the representation of feature extraction will be expanded if the ground truth can be used for direct learning in the middle layer as well. In addition, Reformer~\cite{reformer} using Locality Sensitive Hashing has been proposed in recent years. Since Transformer-based Attention uses a lot of memory, Reformer will work well in our Feedback Attention. These are subjects for future works. 
\bibliographystyle{splncs04}
\bibliography{egbib}
\end{document}


Peer reviews:

content: confidence: 5: The reviewer is absolutely certain that the evaluation is correct and very familiar with the relevant literature
rating: 8: Top 50% of accepted papers, clear accept
review: Summary
-------

This paper presents a feedback mechanism for U-Nets, which re-uses the output
of the U-Net for a second round of processing. For the incorporation of the
output into an earlier feature map of the U-Net, the authors propose two
attention mechanisms (source-target-attention and self-attention). On an
electron microscopy dataset of neural tissue, experiments demonstrate
consistent improvements of the self-attention version of the proposed method
over several baselines, including a vanilla U-Net and a feedback U-Net. In
further ablation studies, the authors investigate the two proposed attention
mechanisms, the choice of the injection point of the output, and the usefulness
of the second round of processing.

Quality and Clarity
-------------------

The technical description of the method is very clear and the provided figures
helpful.

Originality
-----------

Reusing the output of a U-Net for a second round of processing is not an
entirely new idea (as acknowledged by the authors in the Related Work section).
The main contribution here is therefore the addition of an attention mechanism
and the re-use of the same weights for the second round of processing.

Significance
------------

The experimental evaluation is thorough (albeit on only a single
dataset) and addresses key questions about the proposed method. Qualitative
results (especially the attention maps generated by the two proposed
mechanisms) are helpful to understand the contributions of the method.

Pros
----

* clear presentation
* elegant architecture
* convincing results
* thorough analysis of method components in ablation study

Cons
----

* evaluated on only one dataset

Minor Comments
--------------

* line 44: "cell image segmentation is a difficult task because [...] there is not regularity compared to other datasets such as automatic driving" I would personally argue that the opposite is true
* line 420: "menbranes" -> "membranes"
* line 437: "firrst" -> "first"
* line 487: "We" -> "we"
* no "." after "Equation", "Table", or "Figure"
* line 71: "we evaluate the proposed method on two kinds of cell image datasets" results are only presented on one dataset

title: Simple and effective U-Net feedback extension

content: confidence: 3: The reviewer is fairly confident that the evaluation is correct
rating: 8: Top 50% of accepted papers, clear accept
review: ### Quality
This work approaches the addition of a feedback attention mechanism into a U-Net based segmentation method with commendable rigour. While the authors only demonstrate their method on one dataset, they do so in a methodical manner whereby the carefully validate what sort of feedback works best, and also offer explanations as to why this is the case. The authors also validate which layers of the network work most effectively for their chosen network architecture, and that the improvement in segmentation performance is indeed due to the addition of feedback attention rather than just from running data through the U-Net twice.

The figures are generally good quality. The attention maps in Figure 4 are particularly striking and are a good demonstration of the difference between the two proposed feedback attention methods introduced here.

I would have been very interested to see the authors look closer at the attention maps and performance for the mitochondria and synapse classes, as these classes display the most striking segmentation accuracy improvement over the other methods compared. It would also have been interesting for the authors to have discussed cross-applicability of this approach to other microscopy modalities, for example digital pathology data.

### Clarity
It took me a few reads of the paper in order to sufficiently understand the method, and there are quite a lot of grammatical errors and typos throughout.

I personally found that the Query/Key/Value terminology was quite clunky and detracted from my ability to understand section 3.2. For example, the phrase ‘the $i$th Query’s impact’ (line 255) threw me for a while, as I was unclear which dimension (C, H, W) $i$ was indexing. This may well be my own unfamiliarity with the field, but being slightly more explicit in the explanation of the indices would have helped me understand much quicker.

It also felt like there was too much unnecessary repetition between the explanations of the Source-Target-Attention and Self-Attention methods. Again, if I understand correctly, the Source-Target-Attention and Self-Attention methods are identical except for the ‘Query’ being identical to the ‘Key’ in the latter case? If so, this could have been displayed in a more compact mathematical way.

### Originality
This work appears to be the first instance in which attention is integrated into a U-Net via a feedback mechanism. I am not an expert in this field, but a cursory literature search retrieved a [paper discussing the use of attention in medical image segmentation via attention gates](https://arxiv.org/abs/1804.03999) – perhaps this should have been mentioned as ‘Related Work’

### Significance
High-quality semantic segmentation is an undoubtedly significant and impactful challenge for microscopy. The technical merit of this work has been made clear, but I think that the authors could have expanded on the significance of the improved performance in the field of cell imaging. For example, the conclusion does not mention the application at all, and more dwells on further technical adaptations, which feels a little short-sighted. For that reason I also have concerns regarding deployment of this technique and ensuring that biologists are actual able to reap the benefits for their research.

### Pros
* Excellent technical rigour demonstrated in investigating novel method through ablation studies
* Significant increase in segmentation performance achieved with this method

### Cons
* The sections describing the generation of attention maps were quite difficult to follow and understand
* While the performance is very good, the authors do not discuss a route by which this method can actually be used for the benefit of the application demonstrated.
title: Thoroughly-investigated addition of feedback attention to U-Net segmentation