Research proposal

Research question: How can we improve built-in attention for debugging RL agents that rely on visual input?

Background

Recent years have seen a growing interest in techniques to determine which areas of an input image are most responsible for a neural network's output ("explainable AI", "interpretable AI", "decision attribution", etc.). Such tools are of critical importance for developers trying to identify weaknesses in their trained models. During AISC2, a Tensorboard plugin implementing Greydanus et al. was developed.

The techniques proposed so far are based on a wide range of vastly different ideas, and their performance is typically evaluated by comparison against simple benchmarks (such as gradient-based saliency maps, Simonyan et al.) on standard image classification tasks. Widely used techniques include guided backpropagation (Springenberg et al.), perturbation (partial occlusion or blurring) (Greydanus et al.), noise-adding ensemble methods (Smilkov et al.), layer-wise relevance propagation (Bach et al.), contextual decomposition (Murdoch et al.), and many variations on each of them.

One particular approach involves a set of additional network layers implementing a multiplicative attention mask. This has been studied in an RL context by Yang et al. (CNN, Rainbow, good results) and Sorokin et al. (LSTM, DQN, mixed results).

Goal

For this project, we intend to better characterize the performance of such attention-based saliency maps on RL tasks, and test potential avenues of improving their performance.

One particular concern is that the masks obtained typically appear washed-out and blurry. On one hand, this might be an indication that the RL network truly requires a large, fuzzy area of pixels to make good decisions, instead of focusing on particular entities within the image. On the other hand, if this is false, then much more informative attention masks could be obtained by enforcing a less blurry appearance, without reducing the network performance significantly.

We plan to examine this by imposing various kinds of regularizations on the attention mask, such as:

• Entropy loss: an added loss term representing the entropy of the attention mask. Lower-entropy distributions are more discrete ("lumpy"). It is not clear what is the best way to apply entropy loss to attention maps (applying it per-pixel will reward sharp pixel-sized peaks and likely produce noisy results). Possible ideas are applying entropy to downsampled or blurred images.

• Mixed Gaussian clustering: an added loss term representing how badly the attention mask deviates from a mixture of Gaussians, the peaks of which would correspond to clusters of attention that ideally coincide with meaningful objects in the input image.

Time permitting, we also hope to test different architectures for the attention mechanism, document their impact on the performance, and compare the results to at least some of the more well-known methods listed above.

Metrics

• Earth mover's distance: a simple method to quantify the total difference between maps generated by different approaches
• Remove and Retrain (Hooker et al.): a method to quantify the accuracy of a saliency map (e.g. produced by different variations on the attention mask technique), by zeroing out the supposedly salient parts of the image, feeding it to a second RL network, and comparing the performance of this second agent to that of the first.

Environment

• Atari OpenAI Gym (we would prefer to use one that is less black-box but are not aware of any)

Implementation

stable-baselines (fork of OpenAI Baselines, based on Tensorflow)

Deliverable

A writeup (in some form, either blog post or paper) of how attention built into RL agents behaves after applying extra loss terms to it. Time permitting, also a quantification (with various metrics) of how built-in attention is different from other methods for generating saliency maps.