Hard Attention

Task

The images of handwritten digits, MNIST dataset, might be treated as a set of maps or two-dimensional environments which are traveled by an agent. The agent can see only 2 pixels in each direction and must learn a traveling policy to classify environments - digits. The policy trajectories are limited by the allowed number of walking steps.

Solution

We try to solve this task via Active Inference. Active Inference assumes that a policy is computed by maximizing the likelihood of desired observations. In this setting the rewards or state values, that are used in reinforcement learning and optimal control, are replaced by prior beliefs on expected observations or system states. And policy learning becomes equivalent to learning of expected observations or state transitions. On policy, at each step a selected action maximizes likelihood of expected observation. The expected observations can be set by demonstration.

This task is similar to learning an attention policy as described in “Recurrent Models of Visual Attention” http://arxiv.org/abs/1406.6247 (RAM). So we would learn attention policy in the same way however the model would be different.

The differences:

1. No glimpse network, glimpse is a single patch (4x4), instead a collection of multiple resolution patches as in RAM

2. Actions, saccades, are discrete: move up, up-right, right, down-right, down, down-left, left from the current location. RAM actions are locations of patches.

3. Additional network is used to identify initial location from where glimpses start. The location network receives downsized image as input.

Action and Initial Location networks are trained by a simple policy gradient algorithm, such as REINFORCE.

In baseline setup, the reward is class probability value at the last step of an episode.

In demo setup, we walk the agent through an environment to set observation expectations from the demonstration. The demonstration is a sequence of observations (patches), which agent may learn or memorize. The reward is likelihood of observation at each step of an episode, not just at last step as in baseline setup.

Baseline setup results

20 glimpses of 4x4 size with action step size 2, 10 epochs train – 97.48%, test – 95.56%

As you can see below, in many cases learned trajectories reproduce the form of digits. This result confirms optimality of using the same state for generating (actions) and recognition (perception) of trajectories, which is main postulate of active inference.

	Images
original image
trajectories (gray)

The model parameters

RNN

LSTM, one layer of 512 cells

Action Network

1 hidden layer with 512 neurons + soft max + stochastic multinomial module

Initial Location Network

one layer, 4 x downsized input, soft max + stochastic multinomial module that outputs location coordinates

Regularization

All networks use dropout 0.7. The dropout is essential for network stability and allows finding better action policies.

Demo setup results

In main setup we use a single demonstration trajectory. The reward is cross entropy between the patches generated by policy and demo. This reward alone appeared insufficient for map classification because demo trajectory patches are very similar to each other as demo trajectory follows digit curves. Thus agent may just go back and forth at the same location. An additional reward of classification success is still required. In our tests, the composite reward of class probability plus likelihood of an observation gives the same performance as in baseline setup described above. So to verify if observation expectation is plausible as a reward we need other toy task that produces distinguishable patches at each step of episode.