\section{Introduction}

Diffusion MRI provides valuable insights into the $3$D geometric structure of brain neural fiber tracts, allowing tract-based analysis (TBA) of brain white matter microstructure in vivo \cite{ciccarelli2008diffusion}. TBA usually focuses on specific fiber bundles that are extracted from the whole-brain tractograms. Therefore, effective methods for segmenting whole-brain tractograms into fiber bundles of interest are desirable.

Existing fiber parcellation approaches can be divided into two categories: ROI-based and streamline-based.
ROI-based approaches~\cite{yendiki2011automated, wassermann2016white} first parcellate the brain surface based on an atlas and then use the parcellation ROIs to extract different fiber bundles. Streamline-based approaches~\cite{garyfallidis2015robust,zhang2018anatomically} employ streamline registration methods to directly transfer the fiber parcellation information from an atlas to tractograms of an individual. These methods are directly affected by registration accuracy, which is, however, hard to guarantee in practice due to factors such as inter-subject variability, noise, and artifacts. Moreover, the registration procedure is usually time-consuming, making it unsuitable for large-scale studies and real-time applications. Recently, Wasserthal et al. \cite{wasserthal2018tractseg} proposed to predict the binary mask of a fiber bundle from the whole brain fiber peaks using convolutional neural networks. This method avoids the registration procedure but is limited to generating a binary tract segmentation mask rather than the parcellation labels of the fiber tracts. In addition, this method utilizes the fiber orientation information at each voxel rather than the actual geometries of the fiber tracts.

In this work, we propose a deep-learning approach, called DeepBundle, to parcellate whole-brain tractograms without the time-consuming registration. Specifically, we view the coordinates of the points on each fiber tract as functions defined on a graph. The point coordinates are then fed to a graph convolutional neural networks (GCNN) to extract the latent geometric features of fiber tracts for bundle recognition. Our network is trained end-to-end with the point coordinates as inputs and parcellation labels as outputs, thus avoiding having to manually craft features for the purpose of bundle parcellation. During the testing stage, the point coordinates extracted from fiber tracts of the whole brain are directly fed to the trained networks, thus avoiding the use of atlases and any registration method. Extensive experiments performed using data from the Human Connectome Project (HCP) indicate that DeepBundle yields fiber parcellation labels with remarkably improved accuracy, confirming that the geometric features extracted by GCNNs are effective for fiber bundle parcellation.

\section{Methods}

In this section, we will first show how fiber streamlines can be represented using graphs, and then, introduce the theory of spectral graph convolution and graph pooling. Finally, we will describe our network in detail.

\subsection{Graph Representation of Fiber Tracts}

Considering an undirected graph as $\mathcal{G}=(\mathcal{V},\mathcal{E},\mathcal{W})$, $\mathcal{V}=\left\{v_1, \cdots, v_n \right\}$ is a set of $n$ vertices, $\mathcal{E}\subseteq \mathcal{V} \times \mathcal{V} $ is the edge set, and $\mathcal{W}=\left(w_{\left(i,j\right)}\right)$ is the $n\times n$ adjacency matrix, which is symmetric, i.e., $w_{\left(i,j\right)}=w_{\left(j,i\right)}$.
We uniformly sample $n$ points on each fiber tract, and thus this discrete fiber tract can be represented by a line-type graph, where the edge weights are given by
\begin{equation}
w_{i,j}=\begin{cases}
1, & \text{if}~i,j~\text{are connected},\\
0, & \text{otherwise}.
\end{cases}
\end{equation}
The graph $\mathcal{G}$ now encodes the geometric relationships between sampling points on the fiber tract~\cite{bronstein2017geometric}. We then view the point coordinates extracted from a fiber tract as graph-structured data.

\subsection{Spectral Graph Convolution}
To extract the underlying geometry-invariance features of each fiber tract, we employ spectral graph convolution~\cite{bruna2014spectral}, which generalizes conventional convolution in Euclidean space using graph Fourier transformation. The transformation relies on the eigendecomposition of the graph Laplacian $\Delta$, which is defined as
\begin{equation}
    \Delta = \varPhi \Lambda \varPhi^{\text{T}} = I_n - \mathcal{D}^{-1/2}\mathcal{W}\mathcal{D}^{-1/2},
\label{e1}
\end{equation}
where $\varPhi = \left( \varphi_1, \cdots, \varphi_n \right) $ is a matrix of orthonormal eigenvectors ($\varPhi^{\text{T}}\varPhi=I_n$),  $\varLambda=\text{diag}(\lambda_1,\cdots,\lambda_n)$ is a diagonal matrix of corresponding eigenvalues, $I_n$ is an identity matrix, and $\mathcal{D}=\sum_{j\ne i} w_{(i,j)}$ is the degree matrix. Such eigenvectors $(\varphi_1, \cdots, \varphi_n )$ can be interpreted as the Fourier bases, and $\varPhi^T$ is used to transform the features from the spatial domain to spectral domain. The graph Fourier transformation is formulated by
\begin{equation}
\text{f}*\text{g} = \varPhi \left(\varPhi^{\text{T}}\text{g}\,\right)\odot\left(\varPhi^{\text{T}}\text{f}\,\right)=\varPhi \text{diag}\left(\hat{g}_1,\cdots, \hat{g}_n\right)\hat{\text{f}},
\label{e2}
\end{equation}
where $\text{f}=(f_1, \cdots, f_n)^{\text{T}}$ is the input signal, which denotes the geometric features of one streamline on the vertices of graph $\mathcal{G}$. Its Fourier transformation is given by $\hat{\text{f}}=\left(\varPhi^{\text{T}}\text{f}\,\right)$;
$\text{g}$ is a convolutional filter in the spatial domain, and the spectral convolution can be defined as element-wise product $\left(\varPhi^{\text{T}}\text{g}\right)\odot\left(\varPhi^{\text{T}}\text{f}\,\right)$, which can also be written by $\text{diag}\left(\hat{g}_1, \cdots, \hat{g}_n\right)\hat{\text{f}}$. Hereby, $\text{diag}\left(\hat{g}_1, \cdots, \hat{g}_n\right)$ is the corresponding convolutional filter in the spectral domain.

Utilizing the spectral definition of the convolution, GCNNs generalize CNNs to graphs. The $\ell$-th spectral convolution layer can be written by
\begin{equation}
\text{f}_{k}^{(\ell+1)} = \xi\left(\varPhi\sum_{k'=1}^{p}\text{diag}\left(\hat{g}_{(k,k',1)}^{(\ell)}, \cdots, \hat{g}_{(k,k',n)}^{(\ell)}\right)\varPhi^{T}\text{f}_{k'}^{(\ell)}\right),
\label{e3}
\end{equation}
where $\text{F}^{(\ell)} =\left(\text{f}_{1}^{(\ell)},\cdots,\text{f}_{p}^{(\ell)} \right)$ and $\text{F}^{(\ell+1)} =\left(\text{f}_{1}^{(\ell+1)},\cdots,\text{f}_{q}^{(\ell+1)} \right)$ are the $n \times p$ and $n \times q$ matrices, representing the input and output features of $\ell$-th spectral convolution layer; $\text{f}_{k'}^{(\ell)}$ and $\text{f}_{k}^{(\ell+1)}$ denote the $k'$-th column of the input matrix and the $k$-th column of the output matrix. $\varPhi=\left(\varphi_{1},\cdots, \varphi_{n}\right)$ is an $n\times n$ matrix, and $\text{diag}\left(\hat{g}_{(k,k',1)}^{(\ell)},\cdots,\hat{g}_{(k,k',n)}^{(\ell)}\right)$ is a $n\times n$ diagonal matrix which denotes learnable filters in the spectral domain; and $\xi$ denotes the nonlinearity, e.g.,\,ReLU.

\subsection{Fast Graph Pooling}
We utilize the Graclus~\cite{dhillon2007weighted,defferrard2016convolutional} to coarsen the input graph into multi-scales, which is similar to the pooling operation in conventional CNNs. In practice, the multi-scale graph is organized as a binary tree, which becomes coarse from the leaf layer to the root layer. When constructing this tree, Graclus introduces fake nodes, rearranges the nodes of the streamline, and thus, we can perform graph pooling using fast $1$D spatial pooling. Cousin nodes in one layer are aggregated into one parent node in the upper layer of the binary tree.
The pooling process is illustrated in Fig.~\ref{fig:net}.

\begin{figure}[!t]
	\centering
	\includegraphics[width=0.98\textwidth]{Network_architectures}
	\caption{The architecture of the employed spectral GCNNs. Red arrows denote the spectral convolution layers, blue arrows denote the fast max-pooling layers, a green arrow denotes one fully connected layer, and the yellow arrow denotes the softmax layer.}
	\label{fig:net}
\end{figure}

\subsection{Network Architecture of DeepBundle}
Fig.~\ref{fig:net} illustrates the network architecture of DeepBundle. We employ a spectral GCNNs with the architecture of GC$32$-P$2$-GC$64$-P$2$-FC$512$, where GC$c$ denotes a graph convolution layer which has $c$ filters, P$2$ denotes a graph pooling layer with a factor $2$, and FC$512$ denotes a fully connected layer with $512$ hidden nodes. Each graph convolutional layer is activated by the ReLU function. The last layer is the softmax regression layer, and we employ the cross-entropy loss with an $l_2$ regularization term. We separately trained networks for segmenting different fiber bundles of interest from the whole-brain tracts.

In Fig.~\ref{fig:net}, the length of the input is increased from $n$ to $k$ due to the introduction of fake nodes (black dots) by Graclus. We employed $32$ and $64$ filters in the first convolution and the second convolution layer respectively. In the pooling layer, two neighbor nodes are aggregated into one.
The input with a length $k$ is reduced to $k/4$ after going through two max-pooling layers.

\section{Results}
\subsection{Implementation Details}
For evaluation, we used the publicly available HCP fiber tract dataset~\cite{wasserthal2018tract}.
This dataset contains $105$ subjects. Each subject has $72$ fiber bundles, containing streamlines of  different lengths. We uniformly resampled them to have the same number of points, i.e.,\,$100$ points. We then represented all streamlines using a common graph. In our experiments, we randomly selected $25$ subjects from the database for training, $2$ for validation, and $11$ for testing. In preparing the training data, a positive label is assigned to all streamlines from a bundle of interest. The negative samples are collected by 1) selecting streamlines from spatially neighboring bundles, and 2) randomly selecting an equal number of streamlines from all other bundles. For testing, all streamlines of each testing subject were fed separately into the trained network.

\subsection{Experimental Results}

\input{bundle.tex}

Fig.~\ref{fig:orig} shows the fiber parcellation results of one randomly selected subject ($\#623844$).
$12$ bundles of interest are parcellated from more than one million whole-brain fiber tracts. Although fiber bundles in the left and right hemispheres may share similar shapes, DeepBundle is able to distinguish between them.
For some small-size bundles, e.g.,\,FX, which have only hundreds of streamlines, DeepBundle consistently gives promising parcellation results. In addition, we show the quantitative parcellation results of the $12$ fiber bundles in Table~\ref{t1}, further confirming the conclusions that can be drawn from Fig.~\ref{fig:orig}.

\begin{table}[bt]
	\centering
	\caption{\label{t1} Classification results for $12$ bundles of subject $\#623844$.}
	\begin{tabular}{m{1.3cm}m{1.55cm}m{1.55cm}m{1.55cm}m{1.55cm}m{1.55cm}m{1.55cm}@{}ccccccc@{}}
		\toprule
		\small{Counts}               & CST left   & CST right  & CC 1       & CC 2       & CC 6       & CA           \\
		\cmidrule(r){2-7}
		TP                & $4755$     & $7726$     & $4419$     & $51967$    & $30082$    & $1330$       \\
		FP                & $244$      & $395$      & $143$      & $1100$     & $602$      & $0$          \\
		FN                & $54$       & $382$      & $29$       & $153$      & $320$      & $0$          \\
		\midrule
		\small{Counts}                  & UF left    & UF right    & AF left    & AF right   & FX left    & FX right   \\
		\cmidrule(r){2-7}
		TP                & $3535$     & $4916$      & $54448$    & $42635$    & $94$       & $123$      \\
		FP                & $0$        & $0$         & $199$      & $495$      & $6$        & $7$        \\
		FN                & $1$        & $3$         & $190$      & $220$      & $1$        & $1$        \\
		\bottomrule
	\end{tabular}
\end{table}

\begin{table}[bt]
	\centering
	\caption{\label{t2} Classification performance (mean $\pm$ standard deviation in $\%$) on $12$ fiber bundles from $11$ subjects.}
	\begin{tabular}{m{1.6cm}m{1.6cm}m{1.6cm}m{1.6cm}m{1.6cm}m{1.6cm}m{1.6cm}@{}ccccccc@{}}
		\toprule
		\            & CST left      & CST right     & CC 1          & CC 2          & CC 6          & CA             \\
		\cmidrule(r){2-7}
		Precision    & $90.5\pm 7.1$ & $91.1\pm 7.9$ & $96.5\pm 2.1$ & $97.6\pm 2.2$ & $95.6\pm 2.3$ & $99.9\pm 0.1$  \\
		Recall       & $88.4\pm 11.6$& $86.6\pm 11.6$& $95.5\pm 4.3$ & $98.4\pm 1.3$ & $98.1\pm 1.8$ & $100 \pm 0.0$  \\
		\midrule
		\            & UF left      &   UF right     & AF left       & AF right      & FX left       & FX right \\
		\cmidrule(r){2-7}
		Precision    & $99.7\pm 0.6$ & $99.8\pm 0.2$ & $99.2\pm 1.0$ & $99.4\pm 0.6$ & $87.3\pm 12.7$& $90.7\pm 14.1$  \\
		Recall       & $99.2\pm 2.1$ & $99.6\pm 0.4$ & $99.4\pm 0.8$ & $96.5\pm 2.8$ & $97.5\pm 3.9$ & $96.6\pm 3.8$  \\

		\bottomrule
	\end{tabular}
\end{table}

Table~\ref{t2} shows the quantitative results computed across $10$ testing subjects. A large \textit{precision}/\textit{recall} value indicates more accurate fiber parcellation. It can be observed that DeepBundle achieves high \textit{precision} and \textit{recall} values for all bundles of interest, indicating promising fiber parcellation accuracy. It is interesting to note that the recall numbers of left CST and right CST are smaller than the other fiber bundles. This is due to the fact their neighboring bundles, such as POPT and FPT, are very similar to CST in shape.
Such streamlines induce ambiguities in fiber parcellation, even when performed manually.

We also compared DeepBundle with a popular method called RecoBundles~\cite{garyfallidis2018recognition}. We mapped all parcellated fiber streamlines to created volumetric visitation maps and computed their Dice scores with respect to the gold standard. The results, shown in Table~\ref{t3}, indicate that DeepBundle significantly improves the Dice score and outperforms RecoBundles for all bundles of interest. RecoBundles gives significantly lower Dice scores CA, left FX, and right FX. DeepBundle yields relatively smaller Dice scores for left CST and right CST, but still significantly outperforms RecoBundles.

\begin{table}[bt]
	\centering
	\caption{\label{t3} Dice scores (mean $\pm$ standard deviation in $\%$) of $12$ fiber bundles from $11$ testing subjects.}
	\begin{tabular}{m{2.1cm}m{1.55cm}m{1.55cm}m{1.55cm}m{1.55cm}m{1.55cm}m{1.55cm}@{}ccccccc@{}}
		\toprule
		\                 & CST left         & CST right        & CC 1             & CC 2            & CC 6             & CA           \\
		\cmidrule(r){2-7}
		RecoBundles       & $45.3\pm9.7$     & $40.6\pm6.1$     & $51.3\pm13.2$    & $63.6\pm8.6$   & $60.7\pm6.31$     & $25.2\pm16.1$     \\
		DeepBundle        & $80.7\pm10.1$    & $86.9\pm5.1$     & $93.0\pm3.3$     & $96.8\pm1.2$    & $97.1\pm1.4$     & $99.1\pm1.6$       \\
		\midrule
		\                 & UF left          & UF right         & AF left          & AF right        & FX left          & FX right    \\
		\cmidrule(r){2-7}
		RecoBundles       & $46.9\pm6.3$     & $52.1\pm7.7$     & $61.3\pm9.1$     & $63.1\pm12.7$   & $8.1\pm2.8$      & $9.5\pm3.6$      \\
		DeepBundle        & $96.4\pm5.1$     & $98.1\pm1.3$     & $95.8\pm3.8$     & $96.2\pm4.0$    & $82.4\pm9.3$     & $86.5\pm8.8$      \\
		\bottomrule
	\end{tabular}
\end{table}

\section{Conclusion}
In this paper, we have proposed a framework for fiber bundle parcellation using a GCNN. Our method directly predicts a tract parcellation label from the point coordinates extracted from a fiber tract. GCNNs are capable of extracting robust geometric features for tract parcellation in an end-to-end manner. Experiments on HCP data demonstrate that our method, DeepBundle, yields promising tract bundle parcellation results with high precision and recall rates. DeepBundle also achieves much higher Dice scores compared with RecoBundles. Our results also suggest that DeepBundle is even capable of effectively parcellating small tract bundles.

\section*{Acknowledgment}
This work was supported in part by NIH grants NS093842 and the Xi'an Science and Technology Project funded by the Xi'an Science and Technology Bureau under Grant 201805060ZD11CG44.

\bibliographystyle{llncs_splncs}


\section{Introduction}
In the last years, an increasing amount of effort is invested into executing DNN inference in low-precision arithmetic.
While very effective in cutting down on memory, compute and power usage, this comes at a price of degraded network performance caused by weight- and activation-rounding errors. Quantization, the conversion of a net to its low-precision version, is performed according to a "scheme", a conceptual model of the hardware execution environment.  In this paper we follow the 8-bit integer quantization scheme used in \cite{Jacob2017} which is widely employed across both cloud- and edge-devices due to its easy and efficient implementation on hardware. For most network architectures, 8-bit quantization carries a minimal degradation penalty \cite{KrishnaGoogle} leading recent research to double-down on more aggressive schemes, using 4 bits or less for activations, weights or both \cite{IBMFAQ2018discover4bit,Soudry4bit, choi2018pact_ReluAlphaTrained, Jung2018}. Nonetheless, some architectures, such as Mobilenet \cite{2017arXiv170404861H,2018arXiv180104381S}, Inception-V1 \cite{Szegedy2014} and Densenet \cite{Huang2016} still exhibit significant degradation when quantized to 8 bits. Mobilenets are of particular interest as they were designed specifically for embedded applications and as such are often deployed in their quantized form. In order to successfully deploy these networks a quantization-aware training phase \cite{QuallcomMobilenet, KrishnaGoogle, Jacob2017} is usually introduced to regain the network accuracy.  This phase is costly in time and requires access to the original dataset on which the network was trained. Moreover, while effective, the training offers little insight into why there was significant degradation in the first place. Recently, there has been more focus on methods that do not rely on an additional training phase \cite{SameSame, Soudry4bit, zhao2019ChannelSplitting, Migacz2017} but an exhaustive 8-bit quantization scheme that excludes training remains an open challenge.
\newline\indent In this paper, we begin by showing that some networks exhibit a significant shift in the mean activation value following quantization, which we henceforth denote as MAS (mean activation shift). We note that in order for such a shift to have an impact on the performance of the network it must be of significant magnitude with respect to the mean activation itself. The shift is the result of quantization rounding errors being highly unbalanced - a 'small numbers effect' that is statistically implausible for layers with many parameters, but can become the main error source for layers with a small amount of parameters (e.g. depthwise convolution with 9 parameters per channel). Note that the frugal use of parameters is especially prevalent in network architectures aimed at embedded applications \cite{2017arXiv170404861H,2018arXiv180104381S,Zhang2017} making them more susceptible to quantization induced shifts.
We then show that the shift introduced by the quantization can be canceled by using the bias parameter of the channels. Since the bias term is additive, any constant added to it will shift the activation distribution of the channel. Once the task is defined as fixing the shifts, the problem then becomes how to estimate them. To this end, we develop two procedures - one based on direct estimation over a set of test images and the other performing a fine-tuning phase involving only the bias terms of the network.
The main contributions of this paper are:
\begin{itemize}
	\setlength\itemsep{1em}
	\item \textbf{Mean activation shift (MAS)}: We explore both the statistical origins and the impact of MAS on layer-level quantization errors. We show that MAS can arise and be responsible for a large component of the error when quantizing network architectures relying on layers with a very small number of parameters (e.g. MobileNet's depthwise layers).
	\item \textbf{Shift compensation using bias terms}: We show MAS can be compensated by using the bias terms of the network layers, and this alone can drastically reduce degradation, further establishing the previous claim. We propose two algorithms to that effect - Iterative Bias Correction (IBC) and Bias Fine Tune (BFT) - and we analyze their performance on a variety of Imagenet trained classifiers.
\end{itemize}
While our experiments use convolutional neural networks (CNNs) for classification, we expect both the techniques and the analysis to extend into other tasks (e.g. Object-Detection) building on these classifiers as their feature extractors.
\section{Previous Work}

\subsection{Quantization Procedures}
A slew of ideas on DNN quantization were published in recent years, with widely and subtly diverging assumptions about the capabilities of the underlying deployment hardware . For instance, support for non-uniform quantization \cite{xilinxSmallFloat,baskin2018uniq, SongHanDeepCompression}, fine-grained mixed precision \cite{park2018value_mixedprec} or per-channel ("\textit{channelwise}") quantization \cite{KrishnaGoogle, Soudry4bit, choukrounHuawei2019low, goncharenko2018fast}. Leaving the batch-normalization layer unfolded \cite{IBMFAQ2018discover4bit} is another nuance since, similarly to \textit{channelwise quantization}, it enables per-channel scaling at the cost of additional computations. In this work, we restrict ourselves to the simplest setting \textit{folded-batch normalization, layerwise, UINT8-activations, INT8-weights} that was introduced by \cite{Jacob2017}. The attractiveness of this simple setting is that it is supported \textit{efficiently} by most existing hardware on which deep-learning applications are deployed.
The literature can also be dissected by the cost and complexity of the quantization procedures, falling into two main categories:

\paragraph{Post-training quantization:} these methods work directly on a pre-trained network without an additional training phase to compensate for the effects of quantization. Generally lean, these methods have minimal requirements in terms of data, expertise and effort (of both man and machine). Most works are focused on how to best estimate the dynamic range of a layer given its activation distribution. These range from simple heuristics for outlier removal  \cite{KrishnaGoogle, IBMFAQ2018discover4bit}, to more advanced methods employing statistical analysis \cite{Migacz2017, PolinoDistillation,choi2018pact_ReluAlphaTrained, Soudry4bit, choukrounHuawei2019low, goncharenko2018fast}. In this work we use naive min/max statistics (without clipping) to determine dynamic range in order to keep the comparative experiments on IBC/BFT as clean as possible, however, it can be combined with more advanced methods. Other approaches try to lower the dynamic range of the layer by reducing the intra-layer channel range. In \cite{zhao2019ChannelSplitting} it was proposed to improve layerwise quantization by splitting channels with outliers, improving quantization at the cost of additional computations. Finally, in \cite{SameSame} it was proposed to employ inversely-proportional factorization to equalize the dynamic ranges of channels within a layer.

\paragraph{Quantization-aware training} involves optimizing the quantized weights using a large training set, with net-specific hyperparameters \cite{IBMFAQ2018discover4bit} usually taking the pre-trained net as a starting point. The optimization target may use the ground-truth labels \cite{Jacob2017} as in the normal training, and/or an adaptation of knowledge distillation \cite{hinton2015distilling} loss, namely using the full-precision net as the target ("Teacher"), with the quantized ("Student") net punished for distance of its output to the Teacher's \cite{PolinoDistillation, Apprentice}. Rendering the training "quantization-aware" is non-trivial since the rounding and clipping operations are non-differentiable. The simplest and standard approach is that of a Straight-Through-Estimator (STE) \cite{Bengio_STE, PolinoDistillation} which essentially consists of using the quantized version of the net on forward-pass while using the full-precision network on the backward-pass. One of the methods developed in this work (BFT) employs knowledge-distillation and a STE within a very short training procedure (micro-training) restricted to the network biases only. Concurrently to our work, in \cite{choukrounHuawei2019low} a micro-training of scaling factors (multiplicative rather than additive as in ours) was proposed. Like us, they categorize this as post-training quantization since very little data and training time is used.

\subsection{MobileNets Quantization}
Mobilenets \cite{2017arXiv170404861H, 2018arXiv180104381S} are CNNs based on "separable" convolutions - intermittent depthwise and pointwise convolution layers. Their compact and low-redundancy design makes quantization challenging as they have less resilience to noise. The basic \textit{layerwise} 8-bit quantization of Mobilenet-v1 leads to complete loss of accuracy \cite{QuallcomMobilenet, KrishnaGoogle}. A major source of the degradation is due to issues caused by the folding of batch-normalization to degenerate channels (i.e. channel with constant zero output) \cite{QuallcomMobilenet}. Zeroing out these channels reduces degradation to ~$~10\%$ and we employ this technique as well. Quantization-aware training methods \cite{Jacob2017, KrishnaGoogle} are able to reduce degradation to $~1\%$ for both Mobilenet-v1 and v2. Current \textit{post-training quantization} methods can achieve similarly good results (i.e. $<~1\%$) only by using channelwise quantization \cite{KrishnaGoogle} which is not supported on all hardware. Authors of \cite{QuallcomMobilenet} achieve $~2.5\%$ degradation in a layerwise post-training quantization setting but resort to remodeling the network architecture into a "quantization-friendly" version by changing the activation function and modifying the separable-convolution building block. In \cite{SameSame} a $~3\%$ degradation is achieved without either remodelling or retraining by using inversely-proportional factorization to equalize the dynamic range of channels; this further reduces to $~0.6\%$ when the BFT method described in this paper is applied on top, setting the new state-of-the-art for Mobilenet-v1 and v2 "basic-gemmlowp" quantization while being easy and fast to run.

\section{Problem Statement and Analysis}\label{sec:Problem Statement}
A general sentiment underlying quantization of neural networks is that \cite{PeteWarden17}: \begin{quote} "...as long as the errors generally cancel each other out, they'll just appear as the kind of random noise that the network is trained to cope with and so not destroy the overall accuracy by introducing a bias..."\end{quote}
Contrary to the above assumption, in this work we show that when comparing the activations of a feature map in a full-precision and quantized net (QNN), a significant "DC" shift between the distributions can be observed (Fig. \ref{fig:DistributionShift}).
This conceptual gap could arise from thinking on the level of a single tensor, whose rounding error distribution is assumed to be uniform and symmetric; then assuming this property to propagate through the net. However, the law of large numbers doesn't always hold, and small asymmetries might get amplified and accumulate across layers.

In the rest of this section, we first (\ref{subsection3.1}) define the quantization error and it's "DC" component (denoted as MAS). Then in \ref{subsection3.2} we quantify in precise terms the MAS contribution to the error, empirically demonstrating its significance. Finally, in \ref{subsection3.3} we explore further the statistical mechanism of MAS generation and its dependence on layer structure.

\subsection{Quantization Error Measures}\label{subsection3.1}
For a given layer $l$ and feature $ch$ we define the quantization error as the element-wise difference between activations in the original DNN ($x$) and those in the quantized DNN (QNN) ($x^{(q)}$) :
\begin{equation}\label{define_error}
e_{l, ch} = (x^{(q)}_{l, ch}-x_{l, ch}).
\end{equation}
Where $e_{l, ch}, x^{(q)}_{l, ch},x_{l, ch}$ are vectors (for brevity we sometimes drop the subscripts where they are understood from context). Some works make a distinction between rounding "noise" and clipping "distortion"; in this work both are treated indivisibly and we use the terms "noise" and "error" interchangeably. For every channel we define the signal energy, $E(x^2)$, as:
\begin{equation}\label{define_signal_energy}
E(x^2) = \frac{1}{N} \sum{x^2_{l,ch}}
\end{equation}
The expectations $E()$ here and below are to be understood as taken across all pixels of the testing set.
We use the mean rather than sum (as in standard $L_2$ norm definition) so that our analysis will be invariant to the number of elements within a layer. Similarly, we define $E(e^2)$ as the energy of the quantization error. Next, we define the \textbf{Mean Activation Shift} (MAS) as:
\begin{equation}\label{define_shift}
\Delta_{l, ch} = E(x^{(q)}_{l, ch}) - E(x_{l, ch}) = E(x^{(q)}_{l, ch} - x_{l, ch}) = E(e_{l, ch})
\end{equation}

For convenience of exposition we also define we also define the following two quantites: inverse root quantization to noise ratio (rQNSR) and Mean activation Shift to Signal Ratio (MSSR) as:
\begin{equation}\label{define_QNSR}
rQNSR_{l, ch} = \sqrt{\frac{E(e^2)} {E(x^2)}}
\end{equation}
\begin{equation}\label{define_MSSR}
MSSR_{l, ch} = \frac{E(e_{l, ch})} {\sqrt{E(x_{l,ch}^2)}}
\end{equation}
MSSR measures how significant the MAS is compared to the average activation. The rQSNR is a measure of the overall noise level in a channel. Note that MSSR and rQNSR scale similarly, allowing for easy comparison.

All definitions above being random variables' expectations, in practice we use estimates generated by sampling the activations of the network across inferences on a batch of images.

\begin{figure}[ht]
		\vskip 0.2in
		\begin{center}
			\centerline{\includegraphics[width=\columnwidth]{fig1final1.png}}
			\caption{Mean activation shift in two different layers of Mobilenet 1 1.0 224: \textbf{Left:} weight rounding errors (vertical axis) vs. full precision (FP) weights (horizontal axis). An ideal transformation would place all the dots on the null horizontal line. Black indicates a weight whose quantized value is greater than the FP, and magenta indicates weights whose quantized value is smaller or equal to the FP value. \textbf{Right:} 32-image activation histograms for the channel corresponding to the weights on the left. \textbf{Top} row is for the \textit{conv1/channel1} kernel slice, with 27 weight elements (\textbf{1.a}) and a small (0.029) relative mean kernel shift, producing a small MAS (\textbf{1.b}). \textbf{Bottom} row is for the are for \textit{depthwise1} layer (\textit{channel12}), with 9 weight elements (\textbf{1.c}). The relative mean kernel error is large, (0.239), resulting in a large MAS. (\textbf{1.d}). The vertical blue and red lines mark the mean activation values for FP and quantized networks, with a clear shift in \textbf{1.d} and an indiscernible shift in \textbf{1.b} }
		\label{fig:DistributionShift}
		\end{center}
		\vskip -0.2in
\end{figure}

\subsection{The Effect of Mean Activation Shifts}\label{subsection3.2}
One view of MAS as defined in eq. \eqref{define_shift} is as the activations' distribution shift - between the original and the quantized net (see fig. \ref{fig:DistributionShift}). Another analysis avenue relates it to the MSE decomposition:
\begin{equation}\label{mean_var_decompose}
MSE = E(error^2) = {Mean^2}(error) + Var(error)
\end{equation}
In this framework, the question of the MAS's  significance can be formulated as gauging the \textbf{relative contribution} of the error mean to error magnitude, which in our notation can be written as: $\frac{E^2(e_{l, ch})}{E(e^2_{l, ch})}$. To answer this question we estimate these quantities on data samples obtained by running the full precision and the quantized nets over batches of few tens of images each to obtain signal and error vectors; see figs. \ref{fig:MeanShiftContribution}, \ref{fig:MeanShiftContributionLayers} (curves for different choices of batch are very similar; we drop all but one from plot for visual brevity). We can see that for many layers the contribution is significant, sometimes dominating the error energy. Fortunately, the degradation-complicit $\Delta_{l,ch}$ lends itself to a convenient correction, one that doesn't entail changing the NN's computational graph but only adjusting its parameters - namely, subtracting MAS from the (per-channel) biases $B_{l,ch}$ to be added pre-activation
\begin{footnote} {
 Referring to a generic convolutional layer with kernel of size $K$:

 $ X_{l,ch}=F_{act} \left(\sum_{ch',dx,dy}^{K} {X_{l-1,ch'}W_{l,ch}^{ch',dx,dy} + B_{l,ch}} \right) $
}
\end{footnote}
. For a single linear ($F_{act}=I$) layer, this bias-fixing operation can indeed be proven as the optimal error reduction by changing biases alone - effectively removing the Mean contribution in eq. \eqref{mean_var_decompose}. Note also that in contrast to the weights (and their potential adjustments), the bias is quantized at high precision in most implementations (the cost for this being rather low), so the above correction can be done relatively precisely.
 The natural mitigation taken together with the significance of the MAS contribution strikes a great cost-benefit balance, compared to e.g. quantization-aware training methods (and actually performing on-par with those in some cases as we shall see in section \ref{sec:Results}).

Note that for a deep networks the situation is more complex since MAS for different channels in the first layers are mixed in deeper layers, causing input errors that may be comparable with (or larger than) the MAS of the deeper layers.

We will see how these issues are tackled by our optimization methods in section \ref{sec:Optimization Methods}.

\begin{figure}[ht]
		\vskip 0.2in
		\begin{center}
			\centerline{\includegraphics[width=\columnwidth]{Fig2A.png}}
			\caption{32-image estimates of \textbf{(a}) $E(e_{l, ch})$ (the MAS) and \textbf{(b),(c)} $\sqrt{E(x^2_{l, ch})}$, $\sqrt{E(e^2_{l, ch})}$ (L2 norms of activation signal, and its quantization error), per channel, for typical layers of two nets strongly differing in quantization-friendliness. When quantization error is large, the mean-shift tends to be larger even in relative terms, becoming the major contribution to q-error.
			}
		\label{fig:MeanShiftContribution}
		\end{center}
		\vskip -0.2in
\end{figure}

\begin{figure}[ht]
		\vskip 0.2in
		\begin{center}
			\centerline{\includegraphics[width=\columnwidth]{Fig2new.png}}
			\caption{32-image estimates of the MSSR, QNSR, and their ratio ((a)/(b), (c)/(b), (a)/(c) ratios in terms of fig. (\ref{fig:MeanShiftContribution})) L2-aggregated (RMS) over channels, for all layers of same 2 nets. Comparing \textbf{top} and \textbf{center} we can see again, now at the whole net level, that for mobilenet the overall error is greater and it's driven in large part by the MAS. Some typical behaviors can be spotted - e.g. average-pooling layer reduces QNSR but the mean-shift contribution increases. Comparing \textbf{center} and \textbf{bottom} plots, the great error reduction following the BFT procedure is evident (IBC gives a similar result); being driven by reducing MAS, it's not surprising that the relative contribution of MAS (black line) goes down as well
			}
		\label{fig:MeanShiftContributionLayers}
		\end{center}
		\vskip -0.2in
\end{figure}

\subsection{MAS Generation Mechanism in MobileNets}\label{subsection3.3}
We now wish to connect the number of weight elements in a layer and the typical magnitude of MSSR in that layer and show that small kernels tend to increase MAS. Let us consider layer $l$ in which the calculation of an output channel activation involves $k$ kernel elements and $k$ input activations. We define a \textbf{weight rounding error} as \textbf{$\delta_{W_i}$} where $i$ denotes the weight elements in the set of size $k$.
The weight rounding errors are i.i.d and follow the distribution:
\begin{equation}\label{weight_dist}
f(\delta_{W_i}) = U[-\frac{max(|W|)}{2^{N-1}}, \frac{max(|W|)}{2^{N-1}}]
\end{equation}
The mean and standard-deviation of the sum of rounding errors is given by:

\begin{equation}\label{rounderrormean}
E(\sum^k(\delta_{W_i})) = 0
\end{equation}
\begin{equation}\label{rounderrorstd}
\sigma(\sum^k\delta_{W_i}) = C \sqrt{(k/12)}
\end{equation}
with $C$ depending only on $max(|W|)$ and $N$ where $max(|W|)$ is the maximum element magnitude in the weight kernel and $N$ is the number of bits used for quantized representation.
It is worth noting that the assumption of a given $max(|W|)$ means that quantization grid is constant. The rounding error is the results of decimating the kernel values to this grid and so the expectation in the above equations is taken with respect to possible kernel values.

The activation shift for an input image, defined in \ref{define_error} can be expressed as:
\begin{equation}\label{mean_shift_full}
e_{l,ch} = \sum_{i=1}^k (x_{in}^i + \delta_{x_{in}}^i) \times (w_i + \delta_{w_i}) - \sum_{i=1}^k x_{in}^i \times w_i
\end{equation}
or,
\begin{equation}\label{mean_shift_full2}
e_{l,ch} = \sum_{i=1}^k \delta_{x_{in}}^i \times w_i + \sum_{i=1}^k x_{in}^i \times \delta_{w_i}
\end{equation}

In the first layer of the network, the first term in the right-hand side is negligible since when the input rounding errors are small such as in the first layer. We consider only the second term in the following analysis.
Averaging over input data set and the spatial dimensions of layer $l$ to achieve the numerator in \ref{define_MSSR}, we get:
\begin{equation}\label{Emean_shift_full2}
E(e_{l,ch}) = E(\sum_{i=1}^k x_{in}^i \times \delta_{w_i}) = E(x_{in}) \times \sum_{i=1}^k \delta_{w_i}
\end{equation}
where we made the assumption that the input data $x_i$ are i.i.d. The separation of terms in the mean is possible since for a given channel the $\delta_{w_i}$ are constant and independent of the data.

Let us rewrite $MSSR_{l,ch}$ as:
\begin{equation}\label{MSSR2}
MSSR_{l, ch} = \frac{E(x_{in})} {\sqrt{E(x_{out}^2)}} \times \sum_{i=1}^k \delta_{w_i}
\end{equation}
We regard the first term of the product as the 'data term'. If we treat the convolution as a matched filter it can be approximated as:
\begin{equation}\label{rms_prod}
\frac{E(x_{in})} {\sqrt{E(x_{out}^2)}} \approx 1/k
\end{equation}
Using the above approximation we express the mean and variance of $MSSR_{l,ch}$ with regards to possible kernel values in layer $l$ as:
\begin{equation}\label{MSSR2}
E(MSSR_{l, ch}) = C_1\times 1/k \times E(\sum_{i=1}^k\delta_{w_i}) = 0
\end{equation}
\begin{align}\label{MSSR2}
\begin{split}
\sigma(MSSR_{l, ch}) = C_21\times 1/k  \times \sigma(\sum_{i=1}^k\delta_{w_i}) \\
\approx \sqrt{k}/k=1/\sqrt{k}
\end{split}
\end{align}

Equation \ref{MSSR2} tells us that layers with small kernels will tend to have larger MAS energy to activation energy ratio.
This conclusion is in agreement with our observation that Mobilenet architecture exhibits strong mean activation shifts, since it incorporates many depthwise (dw) convolutions, involving small (9 elements) kernel slices. In contrast, conventional 2D convolutions typically involve hundreds of elements for the calculation of each output channel.

This is only one possible source of MAS and the occurrence of significant mean activation shift in Inception v1 despite the absence of depthwise layers suggests other possible sources for MAS. One such source can be the neglected term in \ref{mean_shift_full2} which becomes significant once the input of the layer exhibits large noise.

\section{Optimization Methods}\label{sec:Optimization Methods}
We now present two methods that compensate for the QNN MAS (sec. \ref{sec:Problem Statement}) by updating the bias terms. While the first (IBC) is very fast and requires a very small dataset, the second (BFT) is more general and utilizes a stronger optimizer.

\subsection{Iterative Bias Correction (IBC)}
Estimating the MAS of all the layers and channels and updating the bias terms of the network in one pass cannot be done, since correcting the bias terms of layer $l$ changes the input to layer $l + 1$, which, in turn, gives a different MAS than the one calculated before. Therefore, the correction process has to be performed \textbf{iteratively} according to algorithm \ref{IBC algorithm}, where $act^{orig}_{l, ch}$ and $act^{quant}_{l, ch}$ are the activation values of layer $l$, channel $ch$ in the original and quantized networks, respectively. The inner 'for' loop in the algorithm can be vectorized and done on all of the channels of a layer $l$ simultaneously. The set of input images used for the Iterative Bias Correction (short-named \textit{IBC-batch}), has a typical size of 8-64 images, and can be drawn from the same set used for the QNN calibration. We find that the Iterative Bias Correction effectiveness can vary with IBC-batch size, and larger batches do not necessarily produce better results.
The strength of IBC lies in its simplicity and high speed. With as little as 8 unlabeled images (that can be just the quantization batch) and just a few minutes computation we were able to achieve the results in table \ref{results table}.

\begin{algorithm}
 \caption{Iterative Bias Correction (IBC). The inputs to the IBC algorithm are the original and quantized networks, a topologically sorted list of layers, where 0 is the input layer and the last layer is the output layer. In each iteration the output of the current layer for both the original and quantized networks is evaluated and averaged over $n$ images the spatial dimensions to give a single value per output channel $ch$. The MAS is computed and added to the layer bias.}
 \SetAlgoLined
 \KwResult{Corrected QNN checkpoint}
 evaluate $act^{orig}_{l, ch}$ for all $l$, $ch$ in original net\;
 $l \gets 0$\;
 \While{not end of original net}{
 evaluate $act^{quant}_{l, ch}$\;
 \For{$ch \in l$}{
 $\Delta_{l, ch} \gets <act^{orig}_{l, ch}> - <act^{quant}_{l, ch}> $\;
 $bias^{quant}_{l,ch} \gets bias^{quant}_{l,ch} + \Delta_{l, ch}$\;
 }
 $l \gets l+1$\;
 }
 \label{IBC algorithm}
\end{algorithm}

\subsection{Bias Fine Tuning (BFT)} \label{BFTdesc}
The observations discussed in section 3.2 and in fig. \ref{fig:MeanShiftContribution} suggest that modifications of the QNN's biases, reducing or otherwise modulating the MAS, have the potential to reduce noise energy, impact downstream QSNR and eventually the final layer's output accuracy. This line of thought suggests attempting a generic optimization of the biases alone, with the eventual objective of reducing the overall QNN's loss and accuracy degradation. To this end, we use a quantization-aware training procedure, which starts from the pre-trained weights and biases, but restricts the trainable variables to be the biases only. This restriction enables fast training on a minuscule part of the training set, as small as \textbf{1K} images, without significant overfit, since the number of trainable variables is several orders of magnitude less than with a full training. Our procedure uses global optmization in contrast to the local estimation procedure used in IBC. Drawing inspiration from \cite{Apprentice, PolinoDistillation}, we use pure teacher-student distillation loss, namely the cross-entropy between the logits of the full-precision and the quantized net. This enables a label-free training. To make the training quantization-aware, we use quantized weights for the trainable net, and add $FakeQuant$ operations of the \textit{TensorFlow} framework on the activations' path, implementing the standard Straight-Through Estimator (STE) approach \cite{Bengio_STE, understanding_STE}. Since the weights are fixed for the procedure, our method be also viewed as a miniature variant of Incremental Training \cite{zhou2017IntelINQincremental, zhuang2018GradualBitReduction}, in the sense that biases are optimized to compensate for the error generated by quantization of the weights. Finalizing our procedure, the fine-tuned full precision biases are re-quantized to 16 bits which is precise enough for this final step to have no accuracy impact.

\section{Experiments}\label{sec:Results}
For our tests we used the following setup and procedure (all experiments were done with the \textit{TensorFlow} framework):
\begin{enumerate}
  \item Publicly available pre-trained full-precision models from \href{https://github.com/tensorflow/models/tree/master/research/slim}{tf-slim collection} is ingested; any BatchNorm ops are folded onto preceding layer's weights.
  \item For mobilenets: dead channels (channels with constant output / zero variance) are dropped.
  \item Weights are quantized onto a uniform, symmetric 8-bit (INT8), per-layer grid. Biases are quantized to 16-bit representation.
  \item Activations are quantized (at run-time) onto a fixed, uniform, asymmetric 8-bit (UINT8), per-layer grid over a range defined by simple $min,max$ statistics on a 64-image calibration batch.

  \item QNN accuracy is evaluated on the standard 50K ImageNet validation set.
\end{enumerate}
Table 1 summarizes the results of simple-case tests over several neural nets. The second column gives the Top-1 accuracy of the original, floating-point representation net. The third column gives the top-1 accuracy degradation using the basic quantization described in \cite{Jacob2017}. The last 2 columns give the top-1 accuracy degradation of the quantized nets after applying the IBC and BFT optimization tools.

\begin{table}[]
 	\caption{Top-1 accuracy of original nets and accuracy degradation with basic 8-bit quantization, Iterative Bias Correction (IBC), and Bias Fine Tune (BFT). All numbers are in units of absolute percentage (\%). Best IBC results for were achieved with IBC batch size of 8.}

 	\begin{tabular}{|l||l||l|l|l|}\hline\label{results table}
 		\textbf{\begin{tabular}[c]{@{}l@{}}Network \\ Name\end{tabular}}&
 		\textbf{\begin{tabular}[c]{@{}l@{}} Original  \\ Top-1 \\ accuracy \end{tabular}}&
 		\textbf{\begin{tabular}[c]{@{}l@{}}Basic \\ Quant.\end{tabular}}& \textbf{\begin{tabular}[c]{@{}l@{}}IBC \end{tabular}}&
 		 \textbf{\begin{tabular}[c]{@{}l@{}}BFT\end{tabular}} \\
 		 \hline
 		Mobilenet-v1-1.0  &$71.02$& $7.90$   & $0.92$ & $1.03$ \\ \hline
 		Mobilenet-v2-1.0  &$71.8$& $16.44$   & $1.42$ & $1.2$ \\ \hline
 		Mobilenet-v2-1.4  &$74.95$& $6.42$   & $1.1$ & $0.85$ \\ \hline
 		Inception-v1      &$69.76$& $2.26$   & $0.44$ & $0.47$ \\ \hline
 		Resnet-50-v1      &$75.2$& $0.27$   & $>10$ & $0.30$ \\ \hline
  	\end{tabular}
\end{table}

\subsection{Experiments With IBC}
Testing IBC with different sizes of IBC batches between 8 to 64 images shows a variance of a few 0.1\% in results. Out of the cases we tested, the best results were achieved using an IBC batch of 8 images from the calibration batch.
While MobileNet and Inception architectures benefit from employing IBC, Resnets show large degradation by approximately 10\% when utilizing IBC. We were unable to understand the source of this degradation.

The IBC algorithm discussed throughout this paper compares \textbf{post-activation} means in the original and quantized nets. In Mobilenet v1 the activation function is $Relu6$, clipping the output at 0 and 6, which results in a loss of information regarding the convolution output distribution. It would be reasonable to assume that updating the bias terms according to \textbf{pre-activation} distribution would utilize the full information available and produce even better results than those shown in table \ref{results table}. Indeed, testing a variant of algorithm \ref{IBC algorithm} in which $\Delta_{l, ch}$ is calculated \textbf{pre-activation} on Mobilenet v1, yields a degradation of 0.76\% -  $0.16\%$ better than the post-activation degradation given in table \ref{results table}).

\subsection{Experiments With BFT}
 We trained using the standard Adam optimizer, with a learning rate schedule scanning a wide range as recommended in \cite{IBMFAQ2018discover4bit}: in our default schedule we use $10^{-3},10^{-4},10^{-5},10^{-6}$ rates for 16 mini-epochs each, for a total of 64 mini-epochs using the same 1K images. This schedule is used across all nets presented here (see table \ref{results table}), proving a high degree of robustness. That stands in contrast to the normally high appetite of training-based methods for expertise and data. Given the above, and especially the relaxed input requirement of only 1K label-less images, (similar in size to calibration sets typically used \cite{Migacz2017, KrishnaGoogle}), we claim that the procedure should be seen as a post-training quantization method despite sharing the "backprop+STE+SGD" approach with of quantization-aware training methods.

The results (table \ref{results table}) are quite similar to IBC (suggesting that both utilize well a common "resource" discussed above) - with the exception of ResNet which isn't degraded with BFT. On both Mobilenet-v1 and v2, we perform on par with the state-of-the-art of 1\% degradation reached with full quantization-aware training \cite{KrishnaGoogle}. The restriction to biases apparently enables harvesting the low-hanging fruits of quantization-aware training - which happen to include the lion's share of the degradation to be reduced - using a fraction of the cost.
When combining \textit{ChannelEqualization} \cite{SameSame} with BFT on the Mobilenet-v2 nets, we achieve state-of-the-art quantized net accuracy of 71.1\% (v2-1.0) and 74.3\% (v2-1.4).

\section{Discussion}
Mean activation shift (MAS) occurs in several neural net architectures after quantization. One progenitor for this shift was demonstrated to be a non-symmetrical distribution of the rounding error of weights, more pronounced when the weight kernels involved are small, such as in depthwise convolutions. The phenomenon was observed in Mobilenet v1 and v2 which incorporate depthwise convolutions. It was also observed in Inception v1 which incorporates 2D convolutions and concatenation layers, and no depthwise layers, leading to the conclusion that there are more sources for MAS.

We presented two methods that compensate for mean activation shift of nets by modifying biases - \textbf{Iterative Bias Correction (IBC)}, and \textbf{ Bias Fine Tuning (BFT)}.
Employing either one significantly improves the accuracy of Mobilenet v1, v2 and Inception v1, bringing their degradation down by an order of magnitudes relative to the baseline quantization. For Mobilenet-v1/v2, post-\textbf{BFT} degradation is on par with the ~1\% state-of-the-art \cite{KrishnaGoogle} achieved by resource-intensive quantization-aware training, tested with the same 8-bit quantization scheme. Both tools require very little in terms of tuning effort, run-time and data. BFT performs a robust micro-training procedure, incorporating a strong optimizer (gradient descent) and requiring $\mbox{\tt\raisebox{-.5ex}{\symbol{126}}}{1000}$ unlabeled images; typically taking $\mbox{\tt\raisebox{-.5ex}{\symbol{126}}}{20min}$ on a single GPU. IBC uses a simpler direct-estimation, runs even faster (2-3min is typical) and requires as little as 8 unlabeled images.
We expect these methods to readily extend to mobilenet-based nets used as a base for other tasks (detection, segmentation, etc.) and at other quantization schemes (e.g. 4-bit weights). Both may permit further improvement by using more data and tuning their parameters in a scheme- and net- specific way; however, we refrain from it here, emphasizing instead the "without-bells-and-whistles" effectiveness of the methods, which highlights our basic insight (sec. \ref{subsection3.2}) about the underlying issue. The default setting seems to be enough for sub-$1\%$ degradation on most presented nets (table \ref{results table}). On the other hand, in yet more challenging cases (e.g. quantization to $n<8$ bits) the prospect for the methods described is to be used as a part of a wider post-training quantization toolset. We leave that to future work.

{\small
\bibliographystyle{ieee}


\section{Introduction}
Segmentation annotation maps are crucial for supervised training a deep learning based segmentation model. For segmentation annotation maps, besides the class label dimension, there are spatial dimensions that contain rich information of object size, shape, and between-object/between-class relations.

Previous work has proposed some methods for modifying annotation maps for training better deep learning based segmentation models. Directional map \cite{uhrig2016pixel} was proposed to generate additional training loss based on the relative positions of the pixels to the centers of their corresponding objects. Deep watershed transform \cite{bai2017deep} provided a similar approach that converts an annotation map to a watershed energy map to guide the training of a segmentation model. These efforts demonstrated that changing segmentation annotation maps to include additional information (e.g., relative position, stronger instance-level information) can help train better deep learning models for segmentation tasks.

In medical image segmentation, different classes of objects often have strong mutual locations and spatial correlations. Due to these correlations, learning representation and feature transform for one object class can often indicate the existence of some other object classes (possibly nearby). A conventional way of using multi-class annotation maps is to treat a full annotation map as a whole subject and use spatial cross-entropy loss function to compare it with the model's outputs in back propagation \cite{ronneberger2015u,chen2016deep,zheng2018new}. Due to spatial correlations among different object classes, directly using annotations of all object classes to train a deep network may cause a deep network not to be able to fully exploring its representation learning ability for every object class, especially for those classes with small sizes and unclear/confusing appearance. Furthermore, there may be multiple distinct structures/clusters under one class of objects, and each sub-class structure may better utilize a unique set of feature representations. In principle, we believe that modeling individual classes and sub-classes of structures or objects can encourage deep learning models to learn richer and more comprehensive feature transforms and data representations for segmentation problems.

 \begin{figure}[t]
  \centering
  \includegraphics[width=4.5in]{figure_0.pdf}
  \caption{Segmentation annotation map decomposition based on object classes.}
\label{fig_decomposition_based_on_class}
\end{figure}

 In this paper, we propose to systematically decompose the original
 annotation maps to encourage deep networks to learn richer and possibly more disentangled feature transforms and representations. Our new scheme consists of two main stages: {\it decompose} and {\it integrate}.  {\it Decompose}: by annotation map decomposition, the original segmentation problem is decomposed into multiple segmentation sub-problems (e.g., see Fig.~\ref{fig_decomposition_based_on_class}); these new segmentation sub-problems are modeled by training multiple deep learning modules, each with its own set of feature transforms. {\it Integrate}: a procedure summarizes the solutions of the modules in the previous stage; a final solution is then formed for the original segmentation problem.  This decompose-and-integrate scheme allows to explicitly enforce a deep learning model to learn representations for every object class. Besides, it can also be applied to learn feature transforms and representations for (human expert defined) meaningful sub-class data clusters and structures (see Fig.~\ref{fig_decomposition_based_on_shape}).

In Section \ref{method_section}, we present different ways to decompose annotation maps for different scenarios, and develop a new $K$-to-1 deep network model for implementing our new
 learning scheme. In Section \ref{sec:exp}, we evaluate our decompose-and-integrate learning scheme utilizing multiple state-of-the-art fully convolutional networks (FCNs) on three medical image segmentation datasets, and examine several proposed annotation decomposition (AD) methods.

\begin{figure}[t]
  \centering
  \includegraphics[width=4.2in]{figure_1.pdf}
  \caption{Segmentation annotation map decomposition based on object shape property.}
\label{fig_decomposition_based_on_shape}
\end{figure}

\section{Decompose-and-Integrate Learning}\label{method_section}

Consider a $K$-class segmentation training dataset $\{(x_i,y_i)$, $i = 1,2, \ldots, h$\}, $x_i \in R^{m \times n}$ is a raw image, and $y_i \in \{1,2, \ldots, K\}^{m\times n}$ is a segmentation annotation map containing all the annotations of the $K$ classes of interest. Supervised learning for $K$-class segmentation tasks aims to learn a function $f \in F$ that transforms $x$ to $y$. Note that each $y_i$ can be denoted as $\{y_i^1,y_i^2, \ldots, y_i^K\}$, where $y_i^k\in \{0,1\}^{m\times n}$ is an annotation map for object class $k$, for $k = 1,2, \dots, K$.

For a segmentation problem with two foreground object classes $A$ and $B$, suppose modeling $p(y^B|x)$ for class $B$ is more difficult than modeling $p(y^A|x)$ for class $A$. This means that learning a robust latent representation $R^B(x)$ for class $B$ takes more computational effort (e.g., more training iterations/gradient descent effort) than learning a robust latent representation $R^A(x)$ for class $A$.
Note that $R^B(x)$ and $R^A(x)$ are {\bf not} necessarily disjoint.
When $p(y^A)$ and $p(y^B)$ have moderate or high spatial correlations, using joint annotations of these two classes for training a deep learning model can lead to:
(1) $p(y^A|x)$ is quite
likely to be modeled using $R^A(x)$; (2) $p(y^B|x)$ would be modeled with help from $R^A(x)$, and not mainly by using $R^B(x)$; (3) $R^B(x)$ is not fully explored during model training, due to the ``help'' of the annotations from class $A$.
For a better representation and feature learning performance, such ``help'' is undesired. Besides multi-class segmentation scenarios, when an object class has distinct meaningful underlining sub-class structures/clusters, having a separate modeling for each individual structure/cluster enforces a deep network to learn more meaningful and useful data representations and feature transforms for such a class.

\subsection{Segmentation annotation map decomposition}\label{sec_AD}

\subsubsection{Based on object classes.}
For a $K$-class segmentation problem, we can decompose $y_i$ into $K$ binary annotation maps $y_i^k$, $k={1,2, \ldots, K}$. Algorithm \ref{alg:AD_class} gives the exact procedure. Fig.~\ref{fig_decomposition_based_on_class} shows an image illustration of the effect of this annotation decomposition (AD). In medical image segmentation problems, the number of object classes is usually small, and is much smaller than in natural scene images.
A general guideline is that the decomposed segmentation maps and their associated extra computational costs should be under a manageable level.
Table \ref{tabHVSMR1} shows that object-class based AD can effectively improve segmentation performance for segmentation problems with multiple foreground classes.

\vspace{-0.2in}
\begin{algorithm}
  \caption{Object-class based annotation decomposition}
  \label{alg:AD_class}
\footnotesize
  \begin{algorithmic}[1]
    \Function{AnnotationDecomposition1} {$y_i \in \{1,2, \ldots, K\}^{m\times n}$}
        \For{$k \gets 1 \textrm{ to } K$}
            \State{$y_i^k$ = a new array of size $m\times n$ with all 0;}
            \State {$y_i^k$ [where($y_i==k$)] $\gets 1$;}
        \EndFor
        \State \Return{$y_i^1, y_i^2, \dots ,y_i^K$}
    \EndFunction
  \end{algorithmic}
\end{algorithm}
\vspace{-0.4in}

\subsubsection{Based on object shapes.}
Annotation maps can also be decomposed based on different shape structures in the annotation maps. This type of decomposition can be applied to 2-class segmentation or even $K$-class segmentation for $K>2$.

Shape information contains valuable cues for segmentation tasks. Decomposing annotation maps based on different object shapes can encourage a deep learning model to learn feature transforms that encode the raw images into different shape-guided representations. In histology image analysis, morphological features such as shape convexity play an important role in object detection, segmentation, and diagnosis. Thus, we propose to decompose segmentation annotation maps based on shape convexity of objects in the annotation maps. Specifically, two sub-segmentation maps are generated from an original segmentation annotation map, one containing convex-like shape objects and the other containing concave-like shape objects. {\color{black}This decomposition provides additional information that directly helps a learning model to perceive object information at a higher (object shape) level.} The detailed procedure and image illustration are provided in Algorithm \ref{alg:AD_shape} and Fig.~\ref{fig_decomposition_based_on_shape}. In practice, we set $T_{shape}$ as 0.9.
Table \ref{gland_results_1} demonstrates the usefulness of shape based AD when segmentation problems contain objects with several shape types.

\vspace{-0.1in}
\begin{algorithm}
  \caption{Object-shape based annotation decomposition}
  \label{alg:AD_shape}
\footnotesize
  \begin{algorithmic}[1]
    \Function{AnnotationDecomposition2}{$y_i \in \{1,2, \ldots, K\}^{m\times n}$, $T_{shape}$}
            \State{$y_i^{convex}$ $\gets$ a new array of size $m\times n$ with all 0;}
            \State{$y_i^{concave}$ $\gets$ a new array of size $m\times n$ with all 0;}
            \For{every object $p$ in $y_i$}
                \State Compute the convex hull $p^{convex}$ of $p$
                \State $ratio$ = size($p$)/size($p^{convex}$)
                \If{$ratio > {T}_{shape}$ ($p$ is of a convex-like shape)}
                    \State{Add object $p$ to $y_i^{convex}$}
                \Else
                    \State{Add object $p$ to $y_i^{concave}$}
                \EndIf
            \EndFor
        \State \Return{$y_i^{convex}$ and $y_i^{concave}$}
    \EndFunction
  \end{algorithmic}
\end{algorithm}
\vspace{-0.3in}

\subsubsection{Based on image-level information.}
Image-level information, statistics, and cues can be utilized for annotation map decomposition. For example, if images contain one or multiple foreground objects, we can decompose the segmentation maps based on the number of objects appeared in an image. As the number of objects could only be revealed at a global level or deeper layer in a deep learning model, this decomposition method pushes a learning model to be more aware of global and higher-level information when generating segmentation results. An exact annotation decomposition procedure based on the image-level number of objects is given in Algorithm \ref{alg:AD_imagelevel}.
In Table \ref{lymph_node_results}, we show the effectiveness of image-level information based AD for lymph node segmentation in ultrasound images.

\vspace{-0.1in}
\begin{algorithm}
  \caption{Image-level information based annotation decomposition}
  \label{alg:AD_imagelevel}
\footnotesize
  \begin{algorithmic}[1]
    \Function{AnnotationDecomposition3}{$y_i \in \{1,2, \ldots, K\}^{m\times n}$}
            \State{$y_i^{single\_obj}$ $\gets$ a new array of size $m\times n$ with all 0;}
            \State{$y_i^{multiple\_obj}$ $\gets$ a new array of size $m\times n$ with all 0;}
            \If{$y_i$ contains only one object}
                \State{$y_i^{single\_obj} \gets y_i$}
            \EndIf
            \If{$y_i$ contains multiple objects}
                \State{$y_i^{multiple\_obj} \gets y_i$}
            \EndIf
        \State \Return{$y_i^{single\_obj}$ and $y_i^{multiple\_obj}$}
    \EndFunction
  \end{algorithmic}
\end{algorithm}
\vspace{-0.4in}

\begin{figure}[t]
  \centering
  \includegraphics[width=4.0in]{overview-2.pdf}
  \caption{The $K$-to-1 deep network framework for our new decompose-and-integrate learning scheme. $y^{k}$, $k=1,2, \ldots, K$, are decomposed segmentation annotation maps, and $y$ is the original segmentation annotation map.}
\label{fig_K-1}
\end{figure}

\subsection{The $K$-to-1 deep network for decompose-and-integrate learning}\label{sec_K1}
 Suppose every original annotation map $y_i$, $i = 1,2, \ldots, h$, is decomposed into $K$ annotation maps $y^k_i$, $i = 1,2, \ldots, h$ and $k = 1,2, \ldots, K$. We aim to model each sub-segmentation problem using a deep learning segmentation module with its own set of parameters. Then another modeling procedure is applied on top of these $K$ modules to form the final solution of the original segmentation problem.

 Thus, we propose a new $K$-to-1 deep network framework for implementing our above decompose-and-integrate learning scheme. Fig.~\ref{fig_K-1} shows an overview of our $K$-to-1 deep network. The modules (e.g., Seg-Module 1.1, Seg-Module 2) used in this
 network can be changed according to the type of images (e.g., 2D or 3D images) of the specific segmentation problem. The full model can be trained in end-to-end manner. Let the function of the overall $K$-to-1 network be denoted as $f_{complete}$, and the function of Seg-Module $1.k$ be denoted as $f_{1.k}$. The overall loss for the decompose-and-integrate learning scheme is defined as:

\begin{equation}\label{eq5}
\frac{1}{h}\sum_{i=1}^{h}(\mathcal{L}(f_{complete}(x_i),y_i) + \lambda \sum_{k=1}^{K}\mathcal{L}(f_{1.k}(x_i),y^k_i))
\end{equation}
\noindent

where $\mathcal{L}$ is the spatial cross entropy loss, and $\lambda$ is set as simple as a normalization term $\frac{1}{K}$. We aim to minimize the above function with respect to the parameters of $f_{complete}$ and $f_{1.k}$ for $k = 1,2, \dots, K$.

\section{Experiments and Results}\label{sec:exp}
We conduct experiments on three datasets. The 3D cardiovascular segmentation dataset \cite{pace2015interactive} contains two classes of foreground objects (myocardium and great vessels), which have close spatial relations. Thus, we apply object-class based annotation decomposition (AD) to this dataset. The gland segmentation dataset \cite{sirinukunwattana2017gland} contains glands that have quite different shapes (from concave shape to convex shape); hence shape convexity based annotation decomposition (AD) is applied to this dataset. Our in-house lymph node dataset
contains the lymph node areas of 237 patients in ultrasound images (one image may contain one or more lymph nodes). Thus, image-level information based AD is applied to this dataset.

\begin{table}[t]
\centering
\caption{Comparison of segmentation results on the HVSMR dataset.}
\label{tabHVSMR1}
\scriptsize
\resizebox{\columnwidth}{!}{
\begin{tabular}{|c ccc ccc c|}
    \hline

    \multirow{2}{*}{ Method } &   \multicolumn{3}{c}{Myocardium} & \multicolumn{3}{c}{Blood pool}  &  \multirow{2}{*}{\shortstack{Overall \\ score}} \\  \cline{2-7}
    & Dice   &  ADB     &  Hausdorff   & Dice   &  ADB    &  Hausdorff  &  \\
    \hline
    3D U-Net \cite{cciccek20163d}
    &  $0.694$  &  $1.461$  & $10.221$  &   $ 0.926$ & $ 0.940 $  &  $ 8.628$   &  $ -0.419 $ \\
    \hline
    VoxResNet \cite{chen2017voxresnet}  & $ 0.774 $   &  $ 1.026  $  &  $ 6.572  $  &  $ 0.929  $ & $ 0.981 $  & $ 9.966 $     & $ -0.202 $  \\
    \hline
    DenseVoxNet  \cite{yu2017automatic} &  $ 0.821  $  &  $ 0.964  $  &  $ 7.294  $  & $ 0.931  $  &  $ 0.938  $  &  $ 9.533  $    &  $ -0.161 $ \\
    \hline
     Ensemble Meta-learner \cite{zheng2018new}& 0.823  &   0.685    &  3.224   &  0.935  &   0.763   &  5.804  &  0.215  \\ %
     \hline
   Class-AD + $K$-to-1 DenseVoxNet ({\bf ours})&  $ 	0.839  $  &  $ 0.744  $  &  $ 3.500  $  & $ 0.941  $  &  $ 0.658  $  &  $  5.973  $    &  $ {\bf 0.223} $ \\
    \hline
{Ablation Study}: &  &  &  &  &  & &  \\
    \hline

     Large DenseVoxNet &  $ 0.804 $  &  $ 0.847  $  &  $ 3.980  $  & $ 	0.935  $  &  $ 	0.756  $  &  $  7.706  $    &  $ {0.079} $ \\
      \hline
          2-stacked DenseVoxNet&  $ 0.837  $  &  $ 0.797  $  &  $ 3.405 $  & $ 0.939  $  &  $ 0.629  $  &  $  7.529  $    &  $ {0.167} $ \\
    \hline
    $K$-to-1 DenseVoxNet w/o AD &  $ 0.824  $  &  $ 0.776  $  &  $ 3.619  $  & $ 0.940  $  &  $ 0.677  $  &  $  6.632  $    &  $ {0.177} $ \\
    \hline
\end{tabular}
    }
\end{table}

{\bf Implementation details.}
The input window size of the deep learning segmentation models we use is set as $64\times64\times64$ for 3D experiments and $192\times192$ for 2D experiments. During training, random cropping, rotation, and flipping are applied. Since the images in each dataset are larger than the
model window size, there are virtually many more
samples for model training than the number of images in each dataset. The Adam optimizer is used for model training. The mini-batch size is set as 8. The maximum number of training iteration is set to 60000. We find that usually 60000 iterations using Adam are sufficient for an FCN-type model to converge for a moderate sized training set. The learning rate is set as 0.0005 initially, and decreased to 0.00005 after 30000 iterations.

{\bf 3D cardiovascular segmentation in MR images.}
The HVSMR dataset \cite{pace2015interactive} seeks to segment myocardium and great vessels (blood pool) in 3D cardiovascular MR images. The ground truth of the test data is not available to the public; the evaluations are done by submitting segmentation results to the organizers' server. We experiment with the object class based AD for this dataset. Table \ref{tabHVSMR1} shows that our AD combined with $K$-to-1 network (utilizing DenseVoxNets) achieves state-of-the-art performance on this dataset. In the ablation study part of Table \ref{tabHVSMR1}, we compare our full model with K-to-1 network without AD (where $y_i^k=y_i, i = 1,2, \ldots, h$, and $k = 1,2, \ldots, K$), a 2-stacked DenseVoxNets, and a large-size DenseVoxNet that uses a similar amount of parameters as the $K$-to-1 DenseVoxNets. The ablation study results confirm the effectiveness of our decompose-and-integrate learning scheme.

{\bf Gland segmentation in H\&E stained images.} This dataset \cite{sirinukunwattana2017gland} contains 85 training images (37 benign (BN), 48 malignant (MT)), 60 testing images (33 BN, 27 MT) in part A, and 20 testing images (4 BN, 16 MT) in part B. We modify the original CUMedNet \cite{chen2016deep} to make it deeper with two more encoding and decoding blocks (denoted as CUMedNet$^+$). We run all the experiments for the $K$-to-1 network and ablation study 5 times. Table \ref{gland_results_1} shows the mean performance and standard derivations. Compared with the state-of-the-art models, our AD + $K$-to-1 network (utilizing CUMedNet$^+$) yields considerably better segmentation results. In ablation study (the bottom part of Table \ref{gland_results_1}), we compare AD + $K$-to-1 network with $K$-to-1 network without AD, a 2-stacked CUMedNet$^+$, and a large-size CUMedNet$^+$.

\begin{table*}[t]
\small
\centering
	\caption{Comparison of segmentation results on the gland segmentation dataset.}
	\label{gland_results_1}
	\scriptsize
\resizebox{\columnwidth}{!}{
\begin{tabular}{|ccccccc|}

  \hline
  \multirow{2}{*}{Method}&\multicolumn{2}{c }{$F_1$ Score}&\multicolumn{2}{c }{ObjectDice}&\multicolumn{2}{c| }{ObjectHausdorff}\\\cline{2-7}
  &\multicolumn{1}{c }{part A}&\multicolumn{1}{c }{part B}&\multicolumn{1}{c }{part A}&\multicolumn{1}{c }{part B}&\multicolumn{1}{c }{part A}&\multicolumn{1}{c|}{part B}\\
    \hline
    CUMedVision \cite{chen2016dcan}& 0.912 & 0.716 & 0.897  &0.718  &45.418  &160.347 \\
  \hline

  Multichannel2 \cite{xu2016gland}& 0.893  & 0.843  & 0.908  &0.833&{44.129}&116.821 \\
   \hline
   MILD-Net \cite{graham2018mild} &0.914 & {0.844} & \textbf{0.913} &0.836&{41.54}&105.89  \\
  \hline
CUMedNet \cite{chen2016deep}$^+$ & 0.907$\pm$0.007  & 0.835$\pm$0.009   & 0.893$\pm$0.007  &0.832$\pm$0.008 & 49.97$\pm$2.12 &113.40$\pm$7.22 \\

    \hline
  Shape-AD + $K$-to-1 {\bf(ours)} & \textbf{0.923}$\pm$0.002  & \textbf{0.861}$\pm$0.004   & 0.910$\pm$0.004  &\textbf{0.846}$\pm$0.001 & \textbf{40.79}$\pm$1.72 & \textbf{101.42}$\pm$1.49 \\
    \hline
{Ablation Study}: &  &  &  &  &  &  \\
    \hline
    Large CUMedNet$^+$ &0.918$\pm$0.005  & 0.817$\pm$0.021   & 0.903$\pm$0.002  &0.827$\pm$0.012 & 43.81$\pm$1.39 &{109.43$\pm$5.39} \\

  \hline

  2-stacked CUMedNet$^+$ & 0.914$\pm$0.002  & 0.830$\pm$0.009   & 0.908$\pm$0.001  &0.844$\pm$0.002 & 45.32$\pm$1.05 &{101.43$\pm$2.25} \\
    \hline
    $K$-to-1 w/o AD & 0.915$\pm$0.007  & 0.829$\pm$0.008   & 0.898$\pm$0.007  &0.831$\pm$0.004 & 45.23$\pm$3.71 &{108.92$\pm$4.74} \\

  \hline

\end{tabular}
}
\end{table*}

{\bf Lymph node segmentation in ultrasound images.} We collected
patients' lymph node ultrasound images. We use 137 images for model training, and
100 images for model testing. The image size is $1080 \times 768$. There is no identity overlap between the training data and testing data. The AD procedure follows Algorithm \ref{alg:AD_imagelevel}. Table \ref{lymph_node_results} demonstrates that AD + $K$-to-1 network can effectively improve lymph node segmentation performance in ultrasound images.

\begin{table*}[t]
\centering
	\caption{Comparison of segmentation results on the lymph node segmentation dataset.}
	\label{lymph_node_results}
\scriptsize

\begin{tabular}{|ccccc|}
  \hline
  {Method}&{IoU}&{Precision}&{Recall}&{$F_1$ Score}\\
    \hline
U-Net \cite{ronneberger2015u}& {0.661} & {0.834} & {0.7607}& {0.7957}\\\hline
Deeper U-Net  & {0.7369} & {0.8555} & {0.8416} & {0.8485}\\\hline
CUMedNet \cite{chen2016deep}$^+$& {0.7595} & {0.8472} & {0.8801
} & {0.8633} \\
    \hline
Image-level-AD + $K$-to-1 ({\bf ours})& \bf{0.8102} & \bf{0.9012} & \bf{0.8893} &\bf{0.8952}\\

    \hline
Ablation Study: &  &  &  &\\
    \hline
        Large CUMedNet$^+$& {0.7795} & {0.8808} & { 0.8714} & {0.8761} \\
        \hline
    2-stacked CUMedNet$^+$& {0.7759} & {0.876} & { 0.8716} & {0.8738} \\
        \hline
 $K$-to-1 w/o AD & {0.7842} & {0.8798} & {0.8783} & {0.8790} \\
    \hline

\end{tabular}
\end{table*}

\section{Conclusions}
In this paper, we developed a new decompose-and-integrate learning scheme for medical image segmentation. Our new learning scheme is well motivated, sound, and quite flexible. Comprehensive experiments on multiple datasets show that our new learning scheme is effective in improving segmentation performance.

\bibliographystyle{splncs03}