Explainable Object Detection

Purpose: To explore a Machine Learning (ML) model architecture which makes predictions significantly more explainable. Similar to the ProtoPNet, an established architectures in the eXplainable Artificial Intelligence (XAI) literature, this method seeks to reveal explicate patterns that the model is looking for. In essence, I want to let users "open a model," look at the weights, and immediately understand how models made their decisions.  

Disclaimer: This work is unpublished, not peer reviewed, and holds shortcomings which may make this work unsuitable for industrial applications. The intention is to document these shortcomings and meaningful observations for future reference and in case other researchers gain inspiration and decide to invest more time in this pursuit. Many established works on XAI are worth reviewing, and readers are invited to explore XAI surveys for reliable methods. 

Background: Setting the stage, let's say you are a doctor an I am an machine learning engineer. I've developed a great model that can detect cancer in X-Ray images. Moments before meeting with your patient my model reveals terrible news - there's a tumor in the image. Unfortunately, I cannot tell you any specifics. The model location, size, severity... nothing. "But the model has a great track record," I say, "it has been correct in 99.9% of circumstances!" Hopefully my insistence causes some perturbation in your educated mind, as you need both confidence and understanding of the patient's condition before communicating the bad news. Despite the clear value of understanding the model's prediction, explanations in machine learning are often considered secondary to model performance. The following work is inspired by this observation, and seeks to develop a model architecture which provides unique and explicate types of explainability.   

Overview and Dataset: Two primary versions of the model where generated, which corresponds to the sections below in this webpage.  Please note these models were trained on a simple dataset containing randomly oriented shapes. This dataset was self-generated, and was made so that N triangles, rectangles, squares, and circles could be generated in an image with arbitrary rotations, translations, scales, and sheers. An example of the dataset is in Figure 1.  

Fig. 1 Example of images generated in one of several datasets used to train the models. Note that the shapes are dropped from the full combination of shapes (leftmost) to produce individual data samples (all images). As you might imagine, this is a very small dataset so the model is prone to overfitting 

Concept Learning Model: A proof of concept neural network was created to demonstrate that a model's weights can learn/store input patterns of interest. When trained on images, one weight in the model represents the pixel value a model is "looking for" in the input. This approach helps model inspectors understand the exact patterns predicted by the model rather than relying on saliency based methods. To

This model is supervised that includes a reconstruction loss. In short, the model predicts a internal class which determines the set of weights (a saved "concept" or "memory") that it believes is in the image. These weights are multiplied against the probability of the internal class to produce a the image which is compared against the input image to determine a reconstruction loss. Here, the class score is predicted using a standard Convolutional Neural Network (CNN) with several feed forward layers which are connected to a soft-max prediction layer. Importantly, this initial model is trained on static images that contain several shapes. While training samples drop individual shapes, the location of these shapes does not choose. This approach demonstrates feasibility of the model to disentangle certain concepts in the image. If these images contained N shapes, N internal-class predictions were made inside of the model. A more detailed over of the model is represented below in Figure 2. 

Fig. 2  Overview of the concept learning model

And... it worked! Granted after some finagling with activation functions and hyperparameter tuning.  Figure 3 shows the set of shapes present in a larger test, where each other figure represents the model weights corresponding to an internal model class learned by the model. 

Fig. 3 Example of model weights learned by the model - we are directly "opening" the model and looking at what it learned. Note that (a) represents the set of all shapes where sets of shapes are removed per data sample as in Figure, while (b-l) directly visualize the weights learned by the memory section of the model 

This finding is not incredibly novel as we set up the model to make these kinds of predictions - but please consider that point. There are cases where, simply by changing the model structure, we gain special into what the model is doing. Perhaps many other model designs could be derived that provide special insights into model operations.  

Concept Transformation Model: The model above was extended to take the learned concepts, predict transformations of concepts, and apply these transformations before creating the reconstructed image. Of importance, this model is tested on a single concept to resolve issues with predicting transformations. Figure 4 overviews how a single concept is predicted (shape is in image, or not in image), and then both the input image and memorized concept are passed to a second MLP which predicts a transformation that should be applied to the learned concept in order to match the image. Note that the point transformation needs to be implemented in a differentiable manner (represented as the gaussian reconstruction block), which was implemented but is not reported here for brevity.   

Fig. 4 Overview of the concept transformation model

In summary of this work, 

Summary of Findings

Throughout these lab exercises, I refined my ability to work aseptically, handle cells reproducibly, and generate quantitative data for biological analysis. These skills are directly transferable to roles in biomedical research, biotechnology, and clinical laboratory science.