Distangling Biological Noise in Cellular Images with a focus on Explainability

The process of formation of proteins, or the Central Dogma, is surprisingly not so much of a chemical concept but rather an informational one [2]. The strings of coded information in the DNA provide a template for the structure of protein, but the DNA is enslaved by its geography, i.e, it is only found inside the nucleus while the protein are found in the cytoplasm (outside the nucleus). To carry this information, another molecule called RNA (specifically messenger-RNA) is employed. siRNA are single stranded molecules which bind with mRNA (the intermediary molecule in between DNA and Protein) and stops the formation of the protein [3]. One way to study the function of a particular gene is by silencing it and then studying the resultant cell phenotype. This is also called loss-of-function analyses of the cell [4]. The total silencing renders the action of a particular gene useless, which is termed as gene-knockout and it is a low through put technique as compared to gene-knockdown. Another factor to keep in mind here is the off-target effects of the siRNA used, which might affect not only the mRNA it is designed to inhibit but other mRNAs with partial matches.

The experiment is conducted on single cell type, by choosing one of the four available cell types - HUVEC (Human Umbilical Vein Endothelial Cell), RPE (Retinal Pigment Epithelium), HepG2 (Human Liver Cancer Cell ) and U2OS (Human Bone Osteosarcoma Epithelial Cells). The images are taken from cell cultures.

The siRNA is the object of interest, which is going to create the biological variability, or on a simpler level, changes which we are going to observe. Cell cultures are created in wells of plate which is can hold 384 such well ($16\times 24$). There are in total 51 experiments conducted in the collection of database. One well can be thought as a mini-test tube. There are two imaging sites per well, images from the border wells are discarded as they might be affected by environmental noise like temperature differentials. The image obtained is of the size - $512\times 512\times 6$ i.e, there are six channels per image (for comparison in a standard image there are three channels - Red, Green and Blue).

Controls play a very crucial role in determining the performance of the model as well as calibrating the model

Negative control is a well which is left untreated in the experiments, there is one negative control per plate. A negative control is included in the experiment to distinguish reagent specific effects from non-reagent specific effects in the siRNA-treated cells [4]. When the image from the negative control well is given as an input the model should not make confident predictions, this is a kind of validity check

Positive control is a well which is treated with reagents whose affect on the cell culture is known and well studied [4]. Positive controls are used to measure how efficient the reagent is. In this context when images from the positive control are given as input the model should predict the reagent class with maximum confidence. This again serves as a sanity check.

There are in total $N_{Total}=56416$ images in the data-set. Out of which $N_{Train}=36517$ are in the train set and $N_{Test}=19899$ are in the test set. Since the images comprise of six channels, the total number of gray-scale images in the data-set is $N_{Gray-Scale}=338496$. There in total $N_{Class}=1108$ classes.

The number of samples per class, $n_{Samples/Class}=N_{Train}/N_{Class}$, is a very important number while training the data-set. A data-set might have a high $N_{Total}$ yet a very low $n_{Samples/Class}$ in effect rendering the size of the data-set misleading. The following figure 2 shows the distribution of $n_{Samples/Class}$:

In figure 2 $S_{x}$ is the set of Classes which have same Number of Samples ($=x$). In Table 1 we can see the number of classes which fall in the particular sets.

To get a good bearing on the distribution of pixel values in the data-set, we can look at the distribution of mean pixel value per image according to different channels.

The variance in the mean Pixel Values for Channel 1 is the highest and is the lowest for Channel 4 (Figure 3).

Densenet takes the idea of skip connections (introduced in Resnets [5]) one step further. In Resnet where each layer is connected to the next layer with a shortcut, in Densenet apart from the next layer each layer is connected to every other subsequent layer.

Resnets:

Densenets:

where $x_{l}$ is the feature map after the $l^{th}$ layer. $H_{l}(\cdot)$ is a composite function comprising of activations, non-linearity, pooling etc.

But this is all the similarity that exists between the two networks. Dense takes a departure in the way it introduces these connections, instead of summing up the input to the output, it is rather concatenated with it [6]. In equation [2] the vector $[x_{0},x_{1},\ldots,x_{l-1}]$ is the concatenation of the feature maps outputted in layers $0,1,\ldots,l-1$.This can only be done if the dimension of the channels in the input are same as the output, therefore it is hard to maintain these skip connections through the entire network. These dense connectivity layers are packed into a module, called a dense block as shown in Figure 4. This is a miniature feed forward network in itself which has this added advantage of skip connections. These dense blocks are connected using transition layers. It is at these transition layers where the change in the dimensions of the feature map takes place.

Loss functions are a way to get measure of distance in the embedding space. Since the model tries to minimise this function, it encourages certain kind of structures on the embeddings produced by the model. We will trace the evolution of loss function for a particular type of problem - that of Face Recognition, because it is this class where the current task also finds its place. The difficulty which is posed by the problems of this class is that the intravariations within a class can be larger than inter-differences.

Softmax Loss is a go-to function which provides useful supervision signal in object detection problems. But due to the difficulty posed in the previous section, a vanilla softmax loss gives unsatisfactory results.

where $N$ is the batch-size and $K$ is the number of classes.

Based on the observations of Parde et al [7], one can infer about the quality of the input image by the L2-norm of the features learnt by the softmax loss function. This gave rise to L2-softmax [8], which proposes to enforce all the feature to have the same L2-Norm.

There have been other kinds of normalizations proposed for both weights and the features in the softmax loss function and it has become a common strategy in softmax.

While the Softmax Loss did help in solving the difficulty at hand up to some extent, practitioners in the Face Recognition community wanted that the samples should be separated more strictly to avoid misclassifying the difficult samples [9]. Arc Loss provides a way to deal with the misclassification of difficult samples. Arc Loss can be motivated in very straight forward way using the softmax loss function which is written as in equation [3], for simplicity we will assume $b_{j}=0$. Writing $W_{j}^{T}x_{i}=||W_{j}||~{}||x_{i}||\cos(\theta_{j})$, where $\theta_{j}$ is the angle between the weight $W_{j}$ and feature $x_{i}$. Fixing the norm of $W_{j}$ to be $1$ as done in equation [4], following Wang et al [10] we also fix the feature $x_{i}$ by using $l_{2}-normalisation$ and re-scaling it to $s$. Making the changes we get,

further addition of an angular margin penalty $m$ between $x_{i}$ and $W_{j}$ results in intra-class compactness and inter-class discrepancy

Pseudo Labelling is a technique borrowed from the semi-supervised learning domain. The network is trained with labeled as well unlabelled data simultaneously. In place of labels, Pseudo Labels are used which is nothing but picking the class which has the maximum predicted probability [11].

Let $y^{PL}$ be the pseudo label for a particular sample $x$. Let $p$ be the predicted probabilities for the sample $x$, then:

The labelled and unlabelled samples are used simultaneously to calculate the loss value as follows:

where $p_{k}^{j}$ are the predicted probabilities for $M$ unlabelled samples and $\alpha(t)$ is a coefficient balancing the effect of labeled and unlabelled samples on the network.

It is important to achieve a proper scheduling of $\alpha(t)$, a higher value will disrupt the training and a low value will negate any benefits which can be derived from pseudo-labelling [12].

One way to schedule $\alpha(t)$ is as follows:

where $T_{1}$ is the epoch at which the pseudo-labelling starts and increases linearly until it attains its final value at $T_{2}$ where it is equal to $\alpha_{f}$.

The reasons proposed for the success of this technique is Low Density Separation between Classes, according to the cluster assumption the decision boundaries between clusters lie in low density regions, pseudo labelling helps the network output to be insensitive to variations in the directions of low-dimensional manifold [11].

Data augmentation is broad variety of techniques which is used to increase the generalisation capabilities of the network, either by increasing the available train data-set or increasing the feature capturing strength of filters. There are some standard augmentation techniques such as flipping and rotating the input image which try to make the network agnostic of the pose of features in the image. Even after the application of said transforms the ability of the network to learn local features is not enhanced, there are few techniques which try to promote the object localisation capabilities.

One strategy to use is CutMix. A cropped-up part of one image is placed on top of another image and the labels are manipulated accordingly. More area one image has in the final image, the final label is higher for that image too. To achieve this first we need to make a bounding box $B$ which is the region removed from image $x_{A}$. Let the coordinates of the bounding box $B$ = $(r_{x},r_{y},r_{w},r_{h})$. The bounding box is sampled as follows:

where $x_{A},~{}x_{B}\in\mathbb{(}R)^{W\times H\times C}$ and $\lambda$ is the CutMix coefficient which is sampled from a beta distribution $Beta(\alpha,\alpha)$ with $\alpha$ usually set to $1$. Also note $\frac{r_{w}r_{h}}{WH}=1-\lambda$ which is the cropped area ratio. Using this we get our mask $M\in\{0,1\}^{W\times H}$ which is equal to $0$ for all the points inside the bounding box $B$. Therefore the final image $\hat{x}$ and label $\hat{y}$ thus become,

There have been various techniques [13], [14] proposed in recent times to understand the workings of deep-CNNs. One major break through in this has been the technique called Class Activation Maps (CAM) [15], which proposes to produce these visualizations using global average pooling. A better and more sophisticated version has been proposed based on this work called Grad-CAM.

To go forward, we have to answer a pertinent question - What makes a good visual explanation? This can be answered in two parts. First, the technique should be able to clearly show the difference in localization properties of model for different classes, further it should show the discriminative behaviour in these localizations. Second, the visual explanations should a good resolution, therefore capturing the fine grain details about models attention. These two points are the focus of Grad-CAM.

Grad-CAM is gradient weighted global average pooling technique [16], which works by first computing the importance weights $\alpha_{k}^{c}$,

An $\alpha_{k}^{c}$ is computed for a particular class $c$ and Activation Map $A_{ij}^{k}$ where the partial differentiation term $\frac{\partial y^{c}}{\partial A_{ij}^{k}}$ inside the global average pooling function is the gradient of the class activation, $y^{c}$(before the softmax) with respect to activation map $A_{ij}^{k}$

After this summation of each activation map, weighted by the importance in taken, to generate the Grad-CAM visualization $L_{Grad-CAM}^{c}$

As compared to Resnet, a pre-trained densenet (trained on Imagenet dataset) performed better in general and was used for further training. The version used is the memory-efficient implementation of Densenet-161. The input is a six channel image of the dimensions $6\times 512\times 512$, the dimensions of the transforming image along with the operation acting on it are tracked in the following table:

The way a siRNA reagent interacts with the cell to produce the end results, depends upto some extent on the cell-type as well. This information is incorporated in the model after the Densenet operations. There are four unique cell types, this is represented as a one-hot vector $c\in\{0,1\}^{4}$. The output of Densenet (a tensor of size $2208\times 16\times 16$) is flattened using adaptive average pooling to bring it down to size ($2208\times 1\times 1$) which is then further reduced in dimensions to produce a vector $output_{densenet}\in\mathbb{R}^{2208}$. Concatenating $c$ and $output_{densenet}$ we end up with a vector, $combined\in\mathbb{R}^{2212}$.

The loss function acts on a embedding, $emb\in\mathbb{R}^{1024}$ which is the output the final layer of the model. The final loss function is composite function made up of Softmax [3.2] and Arc-Margin [3.2] loss functions. Using the functions defined in [3] and [6] we get the composite loss,

The hyperparameter $s=30$ and $m=0.5$ in Equation 6, the value of $c$ is $0.2$.

A composite loss function is used to produce a stable training regime. A loss function composed entirely of $L_{arc-margin}$ ended up in an oscillatory behavior. The inclusion $L_{softmax}$ helped stabilize the training process.

As can be seen in the Figure 6 the generalising capabilities of the network were improved almost instantaneously with the introduction of Pseudo-Labelling [3.3].

As defined in section 3.4, CutMix helps the model generalise better. The coefficient $\lambda$ in equation [9] was sampled from the distribution $Beta(\alpha,\alpha)$ with the value of $\alpha=1$.

Pseudo-Labelling proved better at helping the model generalise on the validation set. Though after the introduction of pseudo-labelling the training loss and training accuracy [6] both took a hit. We will get a better insight into the benefits of pseudo-labelling in the next section.

One way to look at the deep-learning model is through the lens of template and feature vectors. The first part of the model takes the image, $I$ and does the convolution and pooling operations, giving a feature vector $x\in\mathbb{R}^{F}$. After this a dot product is taken between $x$ and $w_{i}$, which is the template for the $i^{th}$ class. The resulting operation $w_{i}^{\top}x$ can be summarised for every class using,

where $W=[w_{1},w_{2},\ldots,w_{C}]$ and $W\in\mathbb{R}^{C\times F}$, where $C$ is the number of classes. After taking softmax of $p^{\prime}$ we get $p$. The value of $p_{i}$ is the probability that the image $I$ belongs to class $i$. The key take away here is that, $p_{i}\propto w_{i}^{\top}x$

Therefore a model will generalise well for which $w_{i}$’s are spread apart. Further all the feature vectors of the same class should form a cluster and feature vectors of different classes should be far apart.

In this model we have separate template vectors for the two types of losses. The Figure 7 shows the difference between the template vectors obtained obtained with and without pseudo-labelling. As compared to the vectors in left Figure 7 the vectors are more spread apart in the right Figure 7.

Figure 8 are for softmax template vectors. Apart from the spread of points, another thing to notice here are the blue points. The points shown with blue color are templates for Junk Classes. These are the classes for which there are no samples present in the data-set. The templates generated from the model trained with pseudo-labels has made a completely separate cluster for such classes. The distinction between the Junk Classes and Actual Classes is visible all the four Figure, 7, 8, but the distinction is best visible in the right Figure 8.

Another interesting plot to look for is that of the feature vectors. The plot without pseudo-labelling (Figure 9) has vague clusters and almost in-existent decision boundary. On the other hand the clusters in Figure 9 are much more well defined and tightly packed with a clear decision boundary separating the clusters. The significance further increases because of the fact that these three classes are the most prevalent in the data-set.

In the scatter Plot of class template vectors and feature vectors initially $w_{i},x\in\mathbb{R}^{1024}$ which is reduced to two dimensions using t-SNE [17]

Convolution layers act as filter which can be tuned in the training process to pick-up certain features and dampen others. The filter is usually square shaped and extends through all the channels of the image on which it is acting. Lets consider a square input image for simplicity (also the input image is square shaped in case). Let the input image be $I_{in}\in\mathbb{R}^{W\times W\times C_{in}}$ and the convolution filter be $f\in\mathbb{R}^{n\times n\times C_{in}}$ where $C$ is the number of input channels. The action of one filter produces a new Image with a single channel. When we club a lot of these kind of filters, we get that many images as output which can be stacked to form channels of the output image $I_{out}\in\mathbb{R}^{W^{\prime}\times W^{\prime}\times C_{out}}$.

We are interested in filters of shape $1\times 1\times C_{in}$. The Figures in 10 shows two such filter layers. The image on the left [10] shows the $1\times 1$ part of the Denselayer (6) of Denseblock (1) [2].

Another way to think of a single $1\times 1$ convolution operation is as taking the weighted average of the channels of the input image $I_{in}$.

where, $[w_{1,k},w_{2,k},\ldots,w_{C_{in},k}]$ is the $k^{th}$ $1\times 1$ filter of the convolution layer, $I_{out,k}$ is the $k^{th}$ channel of the output image and $I_{in,i}$ is the $i^{th}$ channel of the input image.

What is important to note in both the figures 10, is the vertical bands of one-color, which are repeated horizontally, forming contiguous blocks of same color.

A band of same color signifies that instead of weighing each input channel ($I_{in,i}$) differently they multiplied by the same number, hence an average of sorts, lets call such and output $I_{avg,k}$ if this happens for the $k^{th}$ filter.

A block of same color signifies that these average images $I_{avg,k}$ are repeated multiple times in the output image. If this happens for $z$ such filters, a lot of variability in the feature map is lost, which is not a good sign.

This phenomenon is observed predominantly in two of the layers shown in Figure 10.

Visualization techniques such as Grad-CAM [3.5] provide us tools to better understand which part of the section of the image impacts the prediction in a positive way. This information is crucial to understand the feature selection capabilities of the network.

In figure 11, it is quite evident that the model seem to paying attention very specific features in the image. Further it would not be wrong to say that these features are of significance, given that it is the Endoplasmic Reticulum in the left, Mitochondria in the middle and Nucleolus in the right.