Spatial Transformer Networks¶

Spatial Transformer Networks, by Jaderberg, M. et. al explores a novel convolutional neural network (CNN) module designed by the authors at Google Deepmind. Spatial Transformers, aptly named, allow a network to incorporate spatial transformations, typically on the input image, which can have the effect of cropping, reorienting, or undistorting the image in question. The authors show how spatial transformations can take the place of attention networks, and additionally solve previous problems with translational invariance in CNNs. They can be added seamlessly to any neural network. They are differentiable, and therefore can be trained through simple backpropogation.

First, let's consider the motivations behind spatial transformer networks, starting with attention mechanisms in neural architecture.

Attention¶

Attention mechanisms in neural networks give them the ability to focus on subsets of features. In image processing, this takes the shape of spatial localization of the features. In the paper Show,

image from Show, Attend, and Tell

Why is this so effective? Here's the intuition I could come up with: Part of it is the problem space of most testing with CNNs: Single object classification or some single object-based task. Even extending this to several objects, there is often a large amount of irrelevant background information in the images. As such, this is an effective dimensionality reduction technique.

To contrast, it may not be effective for segmentation of a large number of objects, or for semantic segmentation tasks, or for something with extracting texture features. In other words, tasks where you might expect important features to be comprehensively spread through the image.

In CNNs, we can consider attention in other image processing terms, taking the form of a binary or probability mask over an image (or layer), thereby focusing the effective activations within a small region. This naive method of attention works, but with a heavy computational load. Kernels must still be computed over the whole image.

Another proposed option, is to crop the image. However, this method also faces problems, as it so far has been non-differentiable, so some sort of stochastic or baked-in cropping method must be implemented.

As we'll see, image cropping can actually be modeled by a constrained linear transformation.

Translational Invariance¶

$\forall \delta \ A f = A T_\delta f$

Translational invariance is the property that transforming a feature map does not affect the resulting output. Convolutional Neural networks do not have this feature, instead, because of parameter sharing, they have the property that they are translationally equivariant:

$\forall \delta \ T_\delta A f = A T_\delta f$

this would be totally sufficient, except that eventually, we are going to feed the network into a dense layer. Dense layers by definition have no spatial relationships built in, so we need the CNN to be able to feed some proper positional information to the dense layer.

In a sense, max pooling layers create some "translational invariance" by only accepting one input of a local region. Of course, this is only a very slight translation for each layer, so it's not nearly the property we're looking for.

Next, let's quickly go over basic spatial transforms.

The two main concepts involved here are spatial transformations, and sampling:

Spatial Transformations¶

For a much deeper look at these spatial transformations, check out this resource

A "spatial transformation" is a transformation on the coordinate system of an image.

Let a pixel $P$ be defined at indices: $(n, m) \in \mathbb{N}$

then the transformation $T_\theta$ maps the pixel $P$ to a new point $(x', y') \in \mathbb{R}^2$

The most basic forms of spatial transformations are linear or projective transformations, which can be represented with matrix multiplication. Let's look at some examples of these matrices, to have a look at what the parameters of such transformations are.

linear¶

Suppose we consider each pair of integer indices in the image, $U$, to be a set of spatial coordinates. Let's denote the indices as $n$, and $m$. The spatial transformation will map each pair of indices $(n, m)$ to a pair $(x', y')$.

For linear transformations, we can easily represent this transformation as a matrix. Then, depending on how we want to transform the matrix, we can

for a "rigid body" transformation, we have 3 parameters, $t_x, t_y, \theta$:

$$\begin{aligned} \begin{bmatrix} x' \\ y' \\ 1 \end{bmatrix} &= \begin{bmatrix} \cos(\theta) & -\sin(\theta) & -t_x \\ \sin(\theta) & \cos(\theta) & -t_y \\ 0 & 0 & 1 \end{bmatrix} \begin{bmatrix} n \\ m \\ 1 \end{bmatrix} \end{aligned} $$

a "similarity" transformation, with 4 parameters $s, t_x, t_y, \theta$:

$$\begin{aligned} \begin{bmatrix} x' \\ y' \\ 1 \end{bmatrix} &= \begin{bmatrix} s \cos(\theta) & -s \sin(\theta) & -t_x \\ s \sin(\theta) & s \cos(\theta) & -t_y \\ 0 & 0 & 1 \end{bmatrix} \begin{bmatrix} n \\ m \\ 1 \end{bmatrix} \end{aligned} $$

for more general, "affine" transformations, 6 parameters $a, b, c, d, t_x, t_y, \theta$:

$$\begin{aligned} \begin{bmatrix} x' \\ y' \\ 1 \end{bmatrix} &= \begin{bmatrix} a & b & t_x \\ c & d & t_y \\ 0 & 0 & 1 \end{bmatrix} \begin{bmatrix} n \\ m \\ 1 \end{bmatrix} \end{aligned} $$

nonlinear¶

finally, "homographic" or "projective" transformations (our first nonlinear transform), with 8 parameters:

$$\begin{aligned} \begin{bmatrix} x' \\ y' \\ 1 \end{bmatrix} &= \begin{bmatrix} a & b & t_x \\ c & d & t_y \\ e & f & t_z \end{bmatrix} \begin{bmatrix} n \\ m \\ 1 \end{bmatrix} \end{aligned} $$

degrees of freedom¶

These parameters are also called the "degrees of freedom" of the transform, where the degrees of freedom are the minimum number of parameters required to fully specify that type of transformation. In the projective transform, although the number of variables shown is 9, the last row has the constraint of always summing to 1, hence you can remove one parameter.

sampling¶

the caveat with these transforms, is that $(x', y')$ are typically not integers, which conflicts with the digital representation of images. We will need to now find the values of the transformed image at new integer points $(n', m') \in \mathbb{N}$

To handle this, we'll need sampling theory.

Interactive Notebook Kernel: if this notebook is being run on a kernel, we can explore affine transforms and sampling with the interactive widget below. We shall also discuss sampling more later on.

The widget above shows for a particular parameterization (defined on the sliders) of the affine transformation, what the end effect will be on the image. You can choose between the square, an outline of a horse, and a photograph of a coffee cup. I encourage you to test out both the available parameters. You can also toggle different interpolation methods, to see how choice of sampling affects the output image (quadratic does not work due to a bug in the current version of scikit-image).

The Spatial Transformer Module: 3 pieces¶

The Spatial Transformer consists of three pieces, a "localization network", a "grid generator", and a sampler*.

the localization network¶

The localization network is how the spatial transformation is decided for each input in the forward pass. It's simple in concept, it is a standard CNN, where the output units are the free parameters in the transformation model we choose (e.g. affine, similarity, projective, etc.). For example, if we wanted an affine transformation, the localization network should have 6 output units. For a similarity transform, 4 output units, etc. The rest is taken care of by backpropogation! The weight initialization should be small, making the model very close to an identity transform at first.

the grid generator¶

The transformation T is generated from the above network's parameters. Then we need to compute which points on the image we are sampling from after the transformation is applied, then place them into an output layer. The output layer may need not strictly be the same size as the input. In fact, the size of the output layer can be freely chosen.

In the image above we can see the identity transformation contrasted with an affine transformation. The sampling grid generates the points which must be sampled onto. This is illustrated in the example above, where we see the quadrilateral of points on $U$ going directly into the next layer $V$. However, just like in the initial discussion of transformations, the points on the quadrilateral will not, in general, correspond directly to pixel coordinates in $U$. This is why we will need to resample the points.

sampling¶

The remaining issue is to recompute the value of the points we sample from on the image. Because they are in between pixel coordinates, they require sampling. In this context, sampling has nothing to do with statistical sampling, it is interpolation. The widget above does a good job of giving a visual inspection of sampling, let's look at how this occurs in the backpropogation.

backpropagation¶

Take equation (3) from the paper, which gives the tranformed image $V$ from the input image $U$.

$$ V_i^c = \sum_n^H \sum_m^W U_{nm}^c k(x_i^s - m; \Phi_x) k(y_i^s - n; \Phi_y) \forall i \in [1...H'W'] \forall c \in [1...C] $$

$k$: the sampling kernel
$\Phi$: sampling kernel parameters
$H, W, C$: heighth, width, channels of U
$H', W'$: heighth, width of V
$n, m$: indices over heighth and width (pixel coordinates) of U
$i$: flat index over all points in the transformed image, V
$x_i^s, y_i^s$: transformed $x$, $y$ coordinates from $T_\theta(G)$
$c$: index over channels, if relevant

This can seem a little obtuse at first:

the key things to note are:

• in practice, the sampling kernel $k_\Phi$ typically behaves as a delta function, eliminating the contribution of most points in the image. So this is not actually a massive sum over the image, typically only 1, 4, or 16 coordinates for nearest neighbor, bilinear, and bicubic interpolation respectively.

• $x_i^s$, $y_i^s$, correspond to $m$, $n$, except $x_i^s$, $y_i^s$ are continous values where $m$, $n$ are integers. That is the problem which sampling addresses, to generate values for $V$ in the integer domain.

• channels of the image are all transformed equivalently, so they shouldn't be given too much thought in the process.

Let's look at the bilinear interpolation example, equation (5), to review backpropogation:

$$ V_i^c = \sum_n^H \sum_m^W U_{nm}^c \max(0, 1 - |x_i^s - m|) \max(0, 1 - |y_i^s - n|) $$

Now, taking the gradient of this function with respect to the input layer $U$:

$$ \frac{\partial V_i^c}{\partial U_{mn}^c} = \sum_n^H \sum_m^W \max(0, 1 - |x_i^s - m|) \max(0, 1 - |y_i^s - n|) $$

Then, taking the gradient with respect to $x_i^s$,

$$\begin{aligned} \frac{\partial V_i^c}{\partial x_i^s} &= \frac{\partial}{\partial x_i^s} \sum_n^H \sum_m^W U_{nm}^c \max(0, 1 - |x_i^s - m|) \max(0, 1 - |y_i^s - n|) \\ &= \begin{cases} \sum_n^H \sum_m^W U_{nm}^c \max(0, 1 - |y_i^s - n|) & m - 1 < x_i^s < m \\ -\sum_n^H \sum_m^W U_{nm}^c \max(0, 1 - |y_i^s - n|) & m < x_i^s < m + 1 \\ 0 & \text{else} \end{cases} \end{aligned}$$

We can see here that the kernel heavily constrains the number of points involved in each gradient, so only local affects should play a large role. When $\frac{\partial x}{\partial \theta}$ is considered, it completes the picture for how spatial transformations relate to gradient flow.

By Example¶

Now let's see some Spatial Transformer Networks in action!

We'll start with MNIST, certainly not for showing the effectiveness, but just as a simple proof of concept.

Take a look at your data¶

Let's just take a look at some data using our plotting helper function:

% Work in Progress % Random Preprocessing¶

So this dataset is pretty vanilla. We don't actually expect spatial transformations to do much, if anything, to digits which are already properly oriented.

When I have some extra time I will add some randomization to the mix to test the limitations of the spatial transformer a bit further.

Defining the Model¶

We'll need to set up the pieces of the spatial transformer module, and then plug it into a standard CNN.

First, we'll define the localizer. Remember that this is really the same as a small CNN. In a naive approach, I'm simply putting some parameters at the end, forced into a sigmoid activation, without concern for their meaning within a transformation matrix...

I'll also define a standard cnn for classification. This will help later for comparing STN results to regular networks, as it will serve as both a stand alone CNN, and as the CNN we feed the output of the Spatial Transformer into.

Next, we'll define a special module, using Deepmind Sonnet's tensorflow package. This module, which I'm calling the Spatial Transform Layer, is meant to take care of generating the grid points and the resampling.

I've defined it to accept a tuple of inputs: the input image, and the transform parameters (the output layer of the localization network). The tuple input works more seamlessly with keras.

finally, we have the tools to finalize the model. The format is basic, just to take inputs and transform them.

Testing the Models¶

Now that we've written up the models themselves, let's initialize each a Standard CNN and a Spatial Transformer Module hooked up to a Standard CNN

Since the CNNs have the same basic architecture, we expect to see how STNs may change what goes in, and hopefully, get a glimpse of how they change kernels of the model later.

What is the STN getting wrong?¶

We end up with 83% accuracy in the STN and 98% accuracy in the standard CNN. So what gives?

I should take this time to note, the accuracy of the STN was quite random, indicating a serious problem with initialization, quite frequently it was near chance, 10%. Can you guess why?

let's look at some examples where the STN is having trouble classifying, to see what might be going on.

The Localizer Network is more nuanced than the paper admits¶

The authors don't make any statements about just how sensitive these layers are. Think about it: Out of all the possible transformations, there are a lot of them that actually land outside of the original image entirely, or grab a useless zero region in the MNIST case. To fix this, we redefine the localizer network.

Let's give it conservatively ranged activation functions such as sigmoid and tanh, then add an identity matrix to the final theta layer. The idea here is that the parameters theta should start very small, and slowly learn how to change from an identity transformation (retaining the image $U$ as-is) to finding a convenient transform to $V$ to assist the standard CNN.

Great! The performance is now much better, but let's see what this new model is having trouble with...

below we'll plot some good predictions to see how it performed there.

The performance is similar to the regular network. However, because the localizer trained concurrently with the convolutional network, likely the transformations it learned were somewhat random, and perhaps not very helpful. Then again, this dataset is far too well-behaved. Next we'll have to test it on something more interesting.