abstract: A capsule is a group of neurons whose activity vector represents the instantiation parameters of a specific type of entity such as an object or an object part. We use the length of the activity vector to represent the probability that the entity exists and its orientation to represent the instantiation parameters. Active capsules at one level make predictions, via transformation matrices, for the instantiation parameters of higher-level capsules. When multiple predictions agree, a higher level capsule becomes active. We show that a discrimininatively trained, multi-layer capsule system achieves state-of-the-art performance on MNIST and is considerably better than a convolutional net at recognizing highly overlapping digits. To achieve these results we use an iterative routing-by-agreement mechanism: A lower-level capsule prefers to send its output to higher level capsules whose activity vectors have a big scalar product with the prediction coming from the lower-level capsule.
author: Sara Sabour
Nicholas Frosst
Geoffrey E. Hinton
Google Brain
Toronto
`{sasabour, frosst, email@example.com`
bibliography: nips.bib
date: May 2017
title: Dynamic Routing Between Capsules

# Introduction

Human vision ignores irrelevant details by using a carefully determined sequence of fixation points to ensure that only a tiny fraction of the optic array is ever processed at the highest resolution. Introspection is a poor guide to understanding how much of our knowledge of a scene comes from the sequence of fixations and how much we glean from a single fixation, but in this paper we will assume that a single fixation gives us much more than just a single identified object and its properties. We assume that our multi-layer visual system creates a parse tree-like structure on each fixation, and we ignore the issue of how these single-fixation parse trees are coordinated over multiple fixations.

Parse trees are generally constructed on the fly by dynamically allocating memory. Following , however, we shall assume that, for a single fixation, a parse tree is carved out of a fixed multilayer neural network like a sculpture is carved from a rock. Each layer will be divided into many small groups of neurons called "capsules" () and each node in the parse tree will correspond to an active capsule. Using an iterative routing process, each active capsule will choose a capsule in the layer above to be its parent in the tree. For the higher levels of a visual system, this iterative process will be solving the problem of assigning parts to wholes.

The activities of the neurons within an active capsule represent the various properties of a particular entity that is present in the image. These properties can include many different types of instantiation parameter such as pose (position, size, orientation), deformation, velocity, albedo, hue, texture, etc. One very special property is the existence of the instantiated entity in the image. An obvious way to represent existence is by using a separate logistic unit whose output is the probability that the entity exists. In this paper we explore an interesting alternative which is to use the overall length of the vector of instantiation parameters to represent the existence of the entity and to force the orientation of the vector to represent the properties of the entity[^1]. We ensure that the length of the vector output of a capsule cannot exceed $1$ by applying a non-linearity that leaves the orientation of the vector unchanged but scales down its magnitude.

The fact that the output of a capsule is a vector makes it possible to use a powerful dynamic routing mechanism to ensure that the output of the capsule gets sent to an appropriate parent in the layer above. Initially, the output is routed to all possible parents but is scaled down by coupling coefficients that sum to $1$. For each possible parent, the capsule computes a "prediction vector" by multiplying its own output by a weight matrix. If this prediction vector has a large scalar product with the output of a possible parent, there is top-down feedback which increases the coupling coefficient for that parent and decreasing it for other parents. This increases the contribution that the capsule makes to that parent thus further increasing the scalar product of the capsule's prediction with the parent's output. This type of "routing-by-agreement" should be far more effective than the very primitive form of routing implemented by max-pooling, which allows neurons in one layer to ignore all but the most active feature detector in a local pool in the layer below. We demonstrate that our dynamic routing mechanism is an effective way to implement the "explaining away" that is needed for segmenting highly overlapping objects.

Convolutional neural networks (CNNs) use translated replicas of learned feature detectors. This allows them to translate knowledge about good weight values acquired at one position in an image to other positions. This has proven extremely helpful in image interpretation. Even though we are replacing the scalar-output feature detectors of CNNs with vector-output capsules and max-pooling with routing-by-agreement, we would still like to replicate learned knowledge across space. To achieve this, we make all but the last layer of capsules be convolutional. As with CNNs, we make higher-level capsules cover larger regions of the image. Unlike max-pooling however, we do not throw away information about the precise position of the entity within the region. For low level capsules, location information is "place-coded" by which capsule is active. As we ascend the hierarchy, more and more of the positional information is "rate-coded" in the real-valued components of the output vector of a capsule. This shift from place-coding to rate-coding combined with the fact that higher-level capsules represent more complex entities with more degrees of freedom suggests that the dimensionality of capsules should increase as we ascend the hierarchy.

# How the vector inputs and outputs of a capsule are computed

There are many possible ways to implement the general idea of capsules. The aim of this paper is not to explore this whole space but simply to show that one fairly straightforward implementation works well and that dynamic routing helps.

We want the length of the output vector of a capsule to represent the probability that the entity represented by the capsule is present in the current input. We therefore use a non-linear **"squashing"** function to ensure that short vectors get shrunk to almost zero length and long vectors get shrunk to a length slightly below $1$. We leave it to discriminative learning to make good use of this non-linearity. $${\bf v}_j = \frac{||{\bf s}_j||^2}{1+||{\bf s}_j||^2} \frac{{\bf s}_j}{||{\bf s}_j||}
\label{squash}$$ where ${\bf v}_j$ is the vector output of capsule $j$ and ${\mathbf{s}}_j$ is its total input.

For all but the first layer of capsules, the total input to a capsule ${\bf s}_j$ is a weighted sum over all "prediction vectors" ${\bf \hat{u}}_{j|i}$ from the capsules in the layer below and is produced by multiplying the output ${\bf u}_i$ of a capsule in the layer below by a weight matrix ${\bf W}_{ij}$ $${\bf s}_j = \sum_i c_{ij} {\bf \hat{u}}_{j|i} \ , \ \ \ \ \ \ \
{\bf \hat{u}}_{j|i} = {\bf W}_{ij}{\bf u}_i$$where the $c_{ij}$ are coupling coefficients that are determined by the iterative dynamic routing process.

The coupling coefficients between capsule $i$ and all the capsules in the layer above sum to $1$ and are determined by a "routing softmax" whose initial logits $b_{ij}$ are the log prior probabilities that capsule $i$ should be coupled to capsule $j$. $$c_{ij} = \frac{\exp(b_{ij})}{\sum_k \exp(b_{ik})}
\label{softmax}$$ The log priors can be learned discriminatively at the same time as all the other weights. They depend on the location and type of the two capsules but not on the current input image[^2]. The initial coupling coefficients are then iteratively refined by measuring the agreement between the current output ${\bf v}_j$ of each capsule, $j$, in the layer above and the prediction ${\bf \hat{u}}_{j|i}$ made by capsule $i$.

The agreement is simply the scalar product $a_{ij} = {\bf v}_j . {\bf \hat{ u}}_{j|i}$. This agreement is treated as if it was a log likelihood and is added to the initial logit, $b_{ij}$ before computing the new values for all the coupling coefficients linking capsule $i$ to higher level capsules.

In convolutional capsule layers, each capsule outputs a local grid of vectors to each type of capsule in the layer above using different transformation matrices for each member of the grid as well as for each type of capsule.

# Margin loss for digit existence

We are using the length of the instantiation vector to represent the probability that a capsule's entity exists. We would like the top-level capsule for digit class $k$ to have a long instantiation vector if and only if that digit is present in the image. To allow for multiple digits, we use a separate margin loss, $L_k$ for each digit capsule, $k$: $$L_k = T_k \ \max(0, m^{+} - ||{\bf v}_k||)^2 + \lambda \ (1-T_k) \ \max(0, ||{\bf v}_k|| - m^{-})^2
\label{digit-loss}$$ where $T_k=1$ iff a digit of class $k$ is present[^3] and $m^{+} = 0.9$ and $m^{-} = 0.1$. The $\lambda$ down-weighting of the loss for absent digit classes stops the initial learning from shrinking the lengths of the activity vectors of all the digit capsules. We use $\lambda=0.5$. The total loss is simply the sum of the losses of all digit capsules.

# CapsNet architecture

A simple CapsNet architecture is shown in Fig. . The architecture is shallow with only two convolutional layers and one fully connected layer. Conv$1$ has $256$, $9 \times 9$ convolution kernels with a stride of 1 and ReLU activation. This layer converts pixel intensities to the activities of local feature detectors that are then used as inputs to the *primary* capsules.

The primary capsules are the lowest level of multi-dimensional entities and, from an inverse graphics perspective, activating the primary capsules corresponds to inverting the rendering process. This is a very different type of computation than piecing instantiated parts together to make familiar wholes, which is what capsules are designed to be good at.

The second layer (PrimaryCapsules) is a convolutional capsule layer with $32$ channels of convolutional $8$D capsules (*i.e.* each primary capsule contains 8 convolutional units with a $9 \times 9$ kernel and a stride of 2). Each primary capsule output sees the outputs of all $256 \times 81$ Conv$1$ units whose receptive fields overlap with the location of the center of the capsule. In total PrimaryCapsules has $[32 \times 6 \times 6]$ capsule outputs (each output is an $8$D vector) and each capsule in the $[6 \times 6]$ grid is sharing their weights with each other. One can see PrimaryCapsules as a Convolution layer with Eq.  as its block non-linearity. The final Layer (DigitCaps) has one $16$D capsule per digit class and each of these capsules receives input from all the capsules in the layer below.

We have routing only between two consecutive capsule layers (e.g. PrimaryCapsules and DigitCaps). Since Conv$1$ output is $1$D, there is no orientation in its space to agree on. Therefore, no routing is used between Conv$1$ and PrimaryCapsules. All the routing logits ($b_{ij}$) are initialized to zero. Therefore, initially a capsule output (${\bf u}_i$) is sent to all parent capsules (${\bf v}_0...{\bf v}_{9}$) with equal probability ($c_{ij}$).
Our implementation is in TensorFlow () and we use the Adam optimizer () with its TensorFlow default parameters, including the exponentially decaying learning rate, to minimize the sum of the margin losses in Eq. .

## Reconstruction as a regularization method

```latex
\begin{tabular}{>{\centering\arraybackslash} m{1cm} | >{\centering\arraybackslash} m{1.5cm} | >{\centering\arraybackslash} m{1.5cm} | >{\centering\arraybackslash}  m{1.5cm} | >{\centering\arraybackslash}m{1.5cm} ? >{\centering\arraybackslash} m{1.5cm} | >{\centering\arraybackslash} m{1.5cm}}
$(l,p,r)$ & $(2,2,2)$ & $(5,5,5)$ & $(8,8,8)$ & $(9,9,9)$ & $(5,3,5)$ & $(5,3,3)$\\ \hline
\pbox{1cm}{Input \\ \\ \\ \\ Output}&
%input &
\includegraphics[height=3cm]{recons/2_258} &
\includegraphics[height=3cm]{recons/5_153} &
\includegraphics[height=3cm]{recons/8_226} &
\includegraphics[height=3cm]{recons/9_125} &
\includegraphics[height=3cm]{recons/5_2035} &
\includegraphics[height=3cm]{recons/5_2035p}
\end{tabular}
```

We use an additional reconstruction loss to encourage the digit capsules to encode the instantiation parameters of the input digit. During training, we mask out all but the activity vector of the correct digit capsule. Then we use this activity vector to reconstruct the input image. The output of the digit capsule is fed into a decoder consisting of $3$ fully connected layers that model the pixel intensities as described in Fig. . We minimize the sum of squared differences between the outputs of the logistic units and the pixel intensities. We scale down this reconstruction loss by $0.0005$ so that it does not dominate the margin loss during training. As illustrated in Fig.  the reconstructions from the $16$D output of the CapsNet are robust while keeping only important details.

# Capsules on MNIST

Training is performed on $28 \times 28$ MNIST () images that have been shifted by up to $2$ pixels in each direction with zero padding. No other data augmentation/deformation is used. The dataset has $60$K and $10$K images for training and testing respectively.

We test using a single model without any model averaging. achieves $0.21$% test error with ensembling and augmenting the data with rotation and scaling. They achieve $0.39$% without them. We get a low test error ($\bm{0.25}$%) on a $3$ layer network previously only achieved by deeper networks. Tab.  reports the test error rate on MNIST for different CapsNet setups and shows the importance of routing and reconstruction regularizer. Adding the reconstruction regularizer boosts the routing performance by enforcing the pose encoding in the capsule vector.

The baseline is a standard CNN with three convolutional layers of $256, 256, 128$ channels. Each has 5x5 kernels and stride of 1. The last convolutional layers are followed by two fully connected layers of size $328, 192$. The last fully connected layer is connected with dropout to a $10$ class softmax layer with cross entropy loss. The baseline is also trained on 2-pixel shifted MNIST with Adam optimizer. The baseline is designed to achieve the best performance on MNIST while keeping the computation cost as close as to CapsNet. In terms of number of parameters the baseline has 35.4M while CapsNet has 8.2M parameters and 6.8M parameters without the reconstruction subnetwork.

|  Method  | Routing | Reconstruction |        MNIST (%)        | MultiMNIST (%) |     |
|:--------:|:-------:|:--------------:|:-----------------------:|:--------------:|:---:|
| Baseline |   \-    |       \-       |         $0.39$          |     $8.1$      |     |
| CapsNet  |    1    |       no       |   $0.34_{\pm 0.032}$    |       \-       |     |
| CapsNet  |    1    |      yes       |   $0.29_{\pm 0.011}$    |     $7.5$      |     |
| CapsNet  |    3    |       no       |   $0.35_{\pm 0.036}$    |       \-       |     |
| CapsNet  |    3    |      yes       | $\bm{0.25}_{\pm 0.005}$ |   $\bm{5.2}$   |     |

CapsNet classification test accuracy. The MNIST average and standard deviation results are reported from $3$ trials.

## What the individual dimensions of a capsule represent

| Scale and thickness |  |
|:---|:---|
| Localized part |  |
| Stroke thickness |  |
| Localized skew |  |
| Width and translation |  |
| Localized part |  |

Since we are passing the encoding of only one digit and zeroing out other digits, the dimensions of a digit capsule should learn to span the space of variations in the way digits of that class are instantiated. These variations include stroke thickness, skew and width. They also include digit-specific variations such as the length of the tail of a 2. We can see what the individual dimensions represent by making use of the decoder network. After computing the activity vector for the correct digit capsule, we can feed a perturbed version of this activity vector to the decoder network and see how the perturbation affects the reconstruction. Examples of these perturbations are shown in Fig. . We found that one dimension (out of $16$) of the capsule almost always represents the width of the digit. While some dimensions represent combinations of global variations, there are other dimensions that represent variation in a localized part of the digit. For example, different dimensions are used for the length of the ascender of a 6 and the size of the loop.

## Robustness to Affine Transformations

Experiments show that each DigitCaps capsule learns a more robust representation for each class than a traditional convolutional network. Because there is natural variance in skew, rotation, style, etc in hand written digits, the trained CapsNet is moderately robust to small affine transformations of the training data.

To test the robustness of CapsNet to affine transformations, we trained a CapsNet and a traditional convolutional network (with MaxPooling and DropOut) on a padded and translated MNIST training set, in which each example is an MNIST digit placed randomly on a black background of $40\times40$ pixels. We then tested this network on the affNIST[^4] data set, in which each example is an MNIST digit with a random small affine transformation. Our models were never trained with affine transformations other than translation and any natural transformation seen in the standard MNIST. An under-trained CapsNet with early stopping which achieved 99.23% accuracy on the expanded MNIST test set achieved 79% accuracy on the affnist test set. A traditional convolutional model with a similar number of parameters which achieved similar accuracy (99.22%) on the expanded mnist test set only achieved 66% on the affnist test set.

# Segmenting highly overlapping digits

Dynamic routing can be viewed as a parallel attention mechanism that allows each capsule at one level to attend to some active capsules at the level below and to ignore others. This should allow the model to recognize multiple objects in the image even if objects overlap. Hinton et al. propose the task of segmenting and recognizing highly overlapping digits ( and others have tested their networks in a similar domain (, , ). The routing-by-agreement should make it possible to use a prior about the shape of objects to help segmentation and it should obviate the need to make higher-level segmentation decisions in the domain of pixels.

## MultiMNIST dataset

```latex
\begin{tabular}{c|c|c|c?c|c?c|c}
R:$(2,7)$ &  R:$(6,0)$ & R:$(6,8)$ & R:$(7,1)$ & *R:$(5,7)$ & *R:$(2,3)$ & R:$(2,8)$ & R:P:$(2,7)$\\
L:$(2,7)$ &  L:$(6,0)$ & L:$(6,8)$ & L:$(7,1)$ & L:$(5,0)$ & L:$(4,3)$ &L:$(2,8)$ & L:$(2,8)$\\ \hline
\includegraphics[height=3cm]{recons/27} &
\includegraphics[height=3cm]{recons/60} &
\includegraphics[height=3cm]{recons/68} &
\includegraphics[height=3cm]{recons/71} &
\includegraphics[height=3cm]{recons/0_5_5_0_332} &
\includegraphics[height=3cm]{recons/4_3_3_4_397} &
\includegraphics[height=3cm]{recons/2_8_2_7_152} &
\includegraphics[height=3cm]{recons/2_8_2_7_153p} \\ \hline
R:$(8,7)$ & R:$(9,4)$ & R:$(9,5)$ & R:$(8,4)$ & *R:$(0,8)$ & *R:$(1,6)$ & R:$(4,9)$ & R:P:$(4, 0)$\\
L:$(8,7)$ & L:$(9,4)$ & L:$(9,5)$ & L:$(8,4)$ &  L:$(1,8)$ & L:$(7,6)$ & L:$(4, 9)$ &L:$(4, 9)$\\ \hline
\includegraphics[height=3cm]{recons/87} &
\includegraphics[height=3cm]{recons/94} &
\includegraphics[height=3cm]{recons/95} &
\includegraphics[height=3cm]{recons/84} &
\includegraphics[height=3cm]{recons/1_8_8_1_264} &
\includegraphics[height=3cm]{recons/7_6_6_7_4} &
\includegraphics[height=3cm]{recons/4_9_4_0_453} &
\includegraphics[height=3cm]{recons/4_9_4_0_454p}  \\ \hline
\end{tabular}
```

We generate the MultiMNIST training and test dataset by overlaying a digit on top of another digit from the same set (training or test) but different class. Each digit is shifted up to $4$ pixels in each direction resulting in a $36\times36$ image. Considering a digit in a $28\times28$ image is bounded in a $20\times20$ box, two digits bounding boxes on average have $80$% overlap. For each digit in the MNIST dataset we generate $1$K MultiMNIST examples. So the training set size is $60$M and the test set size is $10$M.

## MultiMNIST results

Our $3$ layer CapsNet model trained from scratch on MultiMNIST training data achieves higher test classification accuracy than our baseline convolutional model. We are achieving the same classification error rate of $5.0$% on highly overlapping digit pairs as the sequential attention model of achieves on a much easier task that has far less overlap ($80$% overlap of the boxes around the two digits in our case vs $<4$% for ). On test images, which are composed of pairs of images from the test set, we treat the two most active digit capsules as the classification produced by the capsules network. During reconstruction we pick one digit at a time and use the activity vector of the chosen digit capsule to reconstruct the image of the chosen digit (we know this image because we used it to generate the composite image). The only difference with our MNIST model is that we increased the period of the decay step for the learning rate to be $10\times$ larger because the training dataset is larger.

The reconstructions illustrated in Fig.  show that CapsNet is able to segment the image into the two original digits. Since this segmentation is not at pixel level we observe that the model is able to deal correctly with the overlaps (a pixel is on in both digits) while accounting for all the pixels. The position and the style of each digit is encoded in DigitCaps. The decoder has learned to reconstruct a digit given the encoding. The fact that it is able to reconstruct digits regardless of the overlap shows that each digit capsule can pick up the style and position from the votes it is receiving from PrimaryCapsules layer.

Tab.  emphasizes the importance of capsules with routing on this task. As a baseline for the classification of CapsNet accuracy we trained a convolution network with two convolution layers and two fully connected layers on top of them. The first layer has $512$ convolution kernels of size $9\times9$ and stride $1$. The second layer has $256$ kernels of size $5\times5$ and stride $1$. After each convolution layer the model has a pooling layer of size $2\times2$ and stride $2$. The third layer is a $1024$D fully connected layer. All three layers have ReLU non-linearities. The final layer of $10$ units is fully connected. We use the TensorFlow default Adam optimizer () to train a sigmoid cross entropy loss on the output of final layer. This model has $24.56$M parameters which is $2$ times more parameters than CapsNet with $11.36$M parameters. We started with a smaller CNN ($32$ and $64$ convolutional kernels of $5\times5$ and stride of $1$ and a $512$D fully connected layer) and incrementally increased the width of the network until we reached the best test accuracy on a $10$K subset of the MultiMNIST data. We also searched for the right decay step on the $10$K validation set.

We decode the two most active DigitCaps capsules one at a time and get two images. Then by assigning any pixel with non-zero intensity to each digit we get the segmentation results for each digit.

# Other datasets

We tested our capsule model on CIFAR10 and achieved 10.6% error with an ensemble of 7 models each of which is trained with $3$ routing iterations on $24 \times 24$ patches of the image. Each model has the same architecture as the simple model we used for MNIST except that there are three color channels and we used $64$ different types of primary capsule. We also found that it helped to introduce a "none-of-the-above" category for the routing softmaxes, since we do not expect the final layer of ten capsules to explain everything in the image. 10.6% test error is about what standard convolutional nets achieved when they were first applied to CIFAR10 ().

One drawback of Capsules which it shares with generative models is that it likes to account for everything in the image so it does better when it can model the clutter than when it just uses an additional "orphan" category in the dynamic routing. In CIFAR-10, the backgrounds are much too varied to model in a reasonable sized net which helps to account for the poorer performance.

We also tested the exact same architecture as we used for MNIST on smallNORB () and achieved $2.7\%$ test error rate, which is on-par with the state-of-the-art (). The smallNORB dataset consists of 96x96 stereo grey-scale images. We resized the images to 48x48 and during training processed random 32x32 crops of them. We passed the central 32x32 patch during test.

We also trained a smaller network on the small training set of SVHN () with only 73257 images. We reduced the number of first convolutional layer channels to 64, the primary capsule layer to 16 $6D$-capsules with $8D$ final capsule layer at the end and achieved $4.3\%$ on the test set.

# Discussion and previous work

For thirty years, the state-of-the-art in speech recognition used hidden Markov models with Gaussian mixtures as output distributions. These models were easy to learn on small computers, but they had a representational limitation that was ultimately fatal: The one-of-n representations they use are exponentially inefficient compared with, say, a recurrent neural network that uses distributed representations. To double the amount of information that an HMM can remember about the string it has generated so far, we need to square the number of hidden nodes. For a recurrent net we only need to double the number of hidden neurons.

Now that convolutional neural networks have become the dominant approach to object recognition, it makes sense to ask whether there are any exponential inefficiencies that may lead to their demise. A good candidate is the difficulty that convolutional nets have in generalizing to novel viewpoints. The ability to deal with translation is built in, but for the other dimensions of an affine transformation we have to chose between replicating feature detectors on a grid that grows exponentially with the number of dimensions, or increasing the size of the labelled training set in a similarly exponential way. Capsules () avoid these exponential inefficiencies by converting pixel intensities into vectors of instantiation parameters of recognized fragments and then applying transformation matrices to the fragments to predict the instantiation parameters of larger fragments. Transformation matrices that learn to encode the intrinsic spatial relationship between a part and a whole constitute viewpoint invariant knowledge that automatically generalizes to novel viewpoints. proposed transforming autoencoders to generate the instantiation parameters of the PrimaryCapsule layer and their system required transformation matrices to be supplied externally. We propose a complete system that also answers "how larger and more complex visual entities can be recognized by using agreements of the poses predicted by active, lower-level capsules".

Capsules make a very strong representational assumption: At each location in the image, there is at most one instance of the type of entity that a capsule represents. This assumption, which was motivated by the perceptual phenomenon called "crowding" (), eliminates the binding problem () and allows a capsule to use a distributed representation (its activity vector) to encode the instantiation parameters of *the* entity of that type at a given location. This distributed representation is exponentially more efficient than encoding the instantiation parameters by activating a point on a high-dimensional grid and with the right distributed representation, capsules can then take full advantage of the fact that spatial relationships can be modelled by matrix multiplies.

Capsules use neural activities that vary as viewpoint varies rather than trying to eliminate viewpoint variation from the activities. This gives them an advantage over "normalization" methods like spatial transformer networks (): They can deal with multiple different affine transformations of different objects or object parts at the same time.

Capsules are also very good for dealing with segmentation, which is another of the toughest problems in vision, because the vector of instantiation parameters allows them to use routing-by-agreement, as we have demonstrated in this paper. The importance of dynamic routing procedure is also backed by biologically plausible models of invarient pattern recognition in the visual cortex. proposes dynamic connections and canonical object based frames of reference to generate shape descriptions that can be used for object recognition. improves upon dynamic connections and presents a biologically plausible, position and scale invariant model of object representations.

Research on capsules is now at a similar stage to research on recurrent neural networks for speech recognition at the beginning of this century. There are fundamental representational reasons for believing that it is a better approach but it probably requires a lot more small insights before it can out-perform a highly developed technology. The fact that a simple capsules system already gives unparalleled performance at segmenting overlapping digits is an early indication that capsules are a direction worth exploring.

**Acknowledgement.** Of the many who provided us with constructive comments, we are specially grateful to Robert Gens, Eric Langlois, Vincent Vanhoucke, Chris Williams, and the reviewers for their fruitful comments and corrections.

# How many routing iterations to use?

In order to experimentally verify the convergence of the routing algorithm we plot the average change in the routing logits at each routing iteration. Fig.  shows the average $b_{ij}$ change after each routing iteration. Experimentally we observe that there is negligible change in the routing by $5$ iteration from the start of training. Average change in the $2^{nd}$ pass of the routing settles down after 500 epochs of training to 0.007 while at routing iteration 5 the logits only change by $1e-5$ on average.

We observed that in general more routing iterations increases the network capacity and tends to overfit to the training dataset. Fig.  shows a comparison of Capsule training loss on Cifar10 when trained with 1 iteration of routing vs $3$ iteration of routing. Motivated by Fig.  and Fig.  we suggest 3 iteration of routing for all experiments.

[^1]: This makes biological sense as it does not use large activities to get accurate representations of things that probably don't exist.

[^2]: For MNIST we found that it was sufficient to set all of these priors to be equal.

[^3]: We do not allow an image to contain two instances of the same digit class. We address this weakness of capsules in the discussion section.

[^4]: Available at <http://www.cs.toronto.edu/~tijmen/affNIST/>.


abstract: Theano is a Python library that allows to define, optimize, and evaluate mathematical expressions involving multi-dimensional arrays efficiently. Since its introduction in  it has been one of the most used CPU and GPU mathematical compilers – especially in the machine learning community  – and has shown steady performance improvements . Theano is being actively and continuously developed since 2008, multiple frameworks have been built on top of it and it has been used to produce many state-of-the-art machine learning models.
  .
  The present article is structured as follows. Section  provides an overview of the Theano software and its community. Section  presents the principal features of Theano and how to use them, and compares them with other similar projects. Section  focuses on recently-introduced functionalities and improvements. Section  compares the performance of Theano against Torch7  and TensorFlow  on several machine learning models. Section  discusses current limitations of Theano and potential ways of improving it.
bibliography: theano2016.bib
title: Theano: A Python framework for fast computation of mathematical expressions

[^1]

# Overview

## Vision

Theano allows a user to symbolically define mathematical expressions and have them compiled in a highly optimized fashion either on CPUs or GPUs (the latter using CUDA)[^2], just by modifying a configuration flag. Furthermore, Theano can automatically compute symbolic differentiation of complex expressions, ignore the variables that are not required to compute the final output, reuse partial results to avoid redundant computations, apply mathematical simplifications, compute operations in place when possible to minimize the memory usage, and apply numerical stability optimization to overcome or minimize the error due to hardware approximations. To achieve this, the mathematical expressions defined by the user are stored as a graph of variables and operations, that is pruned and optimized at compilation time.

The interface to Theano is Python, a powerful and flexible language that allows for rapid prototyping and provides a fast and easy way to interact with the data. The downside of Python is its interpreter, that is in many cases a poor engine for executing mathematical calculations both in terms of memory usage and speed. Theano overcomes this limitation, by exploiting the compactness and ductility of the Python language and combining them with a fast and optimized computation engine.

Theano's API mimics NumPy , a widely adopted Python library that provides an n-dimensional array data type and many functions for indexing, reshaping, and performing elementary computations (exp, log, sin, etc.) on entire arrays at once. This allows Python users to rapidly switch to Theano using a familiar syntax and set of instructions – extended with advanced features, such as automatic gradient computation, numerical stability improvements and optimization – and generate a high-performance code for CPU as well as for GPU, without requiring changes to the user code. Theano has also been designed for easy and fast extensibility through the definition of custom graph expressions written in Python, C++, or CUDA.

## Community

Theano is a free, open-source software, licensed under the New (3-clause) BSD license. It relies on a wide and very active community of developers and users worldwide.

The main communication channels with the developers are the project's GitHub page[^3] for bug reports, feature requests, and pull requests, and the theano-dev mailing list,[^4] which has 675 subscribers. Support for users is provided by the community at theano-users[^5] (more than 3000 members) and on StackOverflow[^6] (more than 1000 questions asked). PyPI[^7] counted 38k downloads of Theano packages during the last month.

Since the project development migrated to GitHub in 2011, Theano has been forked 1280 times. Around 250 developers have actively contributed to the code base, and numerous others have played a role in the community, asking, answering or curating questions, helping discussing the development needs, and writing documentation, tutorials,[^8] or even full-fledged software projects based on Theano.

## Software based on Theano

Several software packages have been developed to build on the strengths of Theano, with a higher-level user interface, more suitable for certain goals. For instance, machine learning and deep learning packages, such as Pylearn2 , Blocks , Lasagne , and Keras , have been developed with the goal of making it easier to express the architecture of deep learning models, and training algorithms, as mathematical expressions to be evaluated by Theano.

Another example is PyMC3 , a probabilistic programming framework that uses Theano to derive expressions for gradients automatically, and to generate C code for fast execution.

# Main features

Theano defines a *language* to represent mathematical expressions and manipulate them (Section ), a *compiler* to create functions that can compute values for these expressions (Section ), and a *library* which will execute these functions when evaluated on numeric values (Section ). We also explain how Theano can be extended (Section ). Finally, we provide some comparison points with related software (Section ).

## Mathematical expressions

### Graph structure

Theano represents symbolic mathematical expressions as directed, acyclic graphs. These graphs are also bipartite, containing two kinds of nodes:

-   **Variable** nodes (or variables), which represent *data*, usually tensors;

-   **Apply** nodes, which represent the application of *mathematical operations*.

In practice, variables are used for graph inputs and outputs, as well as for intermediate values. During the execution phase, values will be provided for input variables, and computed for intermediate and output ones. An Apply node has inputs and outputs, which are Variable nodes; it represents the application of a mathematical operation (or Op) on its input variables. A Variable node can be the input to several Apply nodes, but can be the output of at most one (graph inputs are not the result of any computation). This corresponds to the single static assignment (SSA) form in compiler design, in that a variable is the result of only one assignation.

This structure is similar to dataflow graphs , where Apply nodes would correspond to operations nodes (the only kind of nodes), and Variable nodes would correspond to arcs in the dataflow graph. The main difference is that a single intermediate Variable node can be an input to several Apply nodes, whereas a dataflow graph would require different arcs, one for each of the next operations.

Variables are strongly typed, they enforce some conditions on the values that can be associated with them. These types are known since the construction of the graph. The main categories of types are:

-   `TensorType`, which represents n-dimensional arrays in the main memory, the values associated with variables of that type are NumPy `ndarray` objects;

-   `CudaNdarrayType`, which represents n-dimensional arrays in GPU memory, associated with `CudaNdarray` objects, used in the legacy GPU back-end;

-   `GpuArrayType`, associated with `GpuArray` objects, its equivalent in the new GPU back-end;

-   `Sparse`, for main-memory sparse matrices, represented by SciPy CSC or CSR matrices.

The number of dimensions and the data type (float32, int64, etc.) are part of the type, as well as what we call the *broadcastable pattern*, which indicates which dimensions are guaranteed to have a shape of 1. Otherwise, the shape is not part of the type, and neither is the memory layout (strides).

### Building a graph

A computation graph is usually constructed by creating free symbolic variables first, corresponding to the inputs of the graph. Since variables are strongly typed in Theano, the type of these variables has to be specified at creation time. By calling Python functions on variables, the user can then interact with them in a direct and natural way. This is reflected under the hood by the creation of Apply nodes and new Variable nodes that extend the graph. The `tensor` module exposes many of the functions provided by NumPy for tensor operations, to present a familiar interface to users. Some of these add a single Apply node and its output to the graph, returning the output Variable node, while other build more complex graphs with Apply nodes corresponding to different Ops, combined in such a way that the returned variable represents the expected result.

It is also possible to clone an existing graph, or a part of it. In that case, what was an intermediate variable in the original graph could become a free input, or an output, of the cloned graph. It is also possible to clone with replacements, which make it possible to plug together different disconnected graphs, making inputs into intermediate Variable nodes.

### Symbolic differentiation

A useful way of deriving gradients is by applying the chain rule backwards through the graph, from a scalar cost towards the inputs (or parameters). This procedure is known as gradient back-propagation, or as the backward or reverse mode of differentiation. For instance, if we have three functions $f: \mathbb R^M \rightarrow \mathbb R$, $g: \mathbb R^N \rightarrow \mathbb R^M$, and $C: \mathbb R^N \rightarrow \mathbb R$ so that $C(x) = f(g(x))$, then: $$\left.\frac{\partial C}{\partial x}\right|_x =
              \left.\frac{\partial f}{\partial g}\right|_{g(x)}
              \cdot \left.\frac{\partial g}{\partial x}\right|_x$$ Instead of computing (and storing in memory) explicitly the whole $M \times N$ Jacobian matrix, $\left.\frac{\partial g}{\partial x}\right|_x$, all we need is a function $\nabla g_x: \mathbb R^M \rightarrow \mathbb R^N, v \mapsto v \cdot \left.\frac{\partial g}{\partial x}\right|_x$ that computes the vector-Jacobian dot product for any vector $v$. This can be generalized easily to functions with several inputs, which can be multi-dimensional arrays.

Most of Theano Ops implement a `grad` method that, given symbolic variables for $x$ and $v$, will return a symbolic expression of $\nabla g_x(v)$, where $g$ is the function represented by that Op. `theano.grad` traverses the graph following the usual back-propagation algorithm, calling the `grad` method on each Apply node's Op, passing that node's input as $x$ and the gradient coming from the subsequent operations as $v$. This builds a symbolic expression for the gradient of the cost with respect to variables. These gradients are symbolic variables that are part of the graph as well, so it is possible to use them as parts of other symbolic expressions (to express a learning rule, for instance), and even to traverse the graph again to obtain higher-order derivatives.

Many Theano Ops also implement an `R_op` method, computing a symbolic expression for the the Jacobian-vector dot product, $R g_x: {\mathbb R^N \rightarrow \mathbb R^M}, {v \mapsto \left.\frac{\partial g}{\partial x}\right|_x \cdot v}$. This is the R-operator introduced by , and corresponds to the forward mode of differentiation. `theano.Rop` traverses the graph from inputs to outputs, calling the `R_op` method on each Apply node's Op.

### Scan: Symbolic loops

Since the computation graph is acyclic, and its structure is fixed and independent from the actual data, it can be a challenge to express loops symbolically. One option, when the number of steps in the loop is fixed, is to explicitly unroll the loop, adding to the computation graph the computation of each of the iterations multiple times. Unfortunately, this makes it impossible to iterate over a sequence of unknown length, or to iterate a variable number of times depending on the value of the data.

To sidestep these issues, Theano implements a special Op called *Scan*, which abstracts the entire loop in a single Apply node in the graph. That single node contains a full computation graph, isolated from the main one, that represents the computation done during each iteration of the loop. The scan node handles the communication between the *external* or *outer* computation graph it belongs to, and the *internal* or *inner* graph. It is also responsible to manage the bookkeeping between the different iterations.

The gradient of a Scan operation is implemented as another Scan operation, which iterates over reversed sequences, computing the same gradient as if the loop had been unrolled, implementing what is known as *back-propagation through time*. Similarly, the R operator is also a Scan operation that goes through the loop in the same order as the original Scan.

## The compilation phase

The compilation phase produces a Theano function (a Python callable object) able to compute values for specified *output* symbolic variables, given values for *input* variables. The set of input and output variables have to be provided when compiling the function, but the inputs do not have to be inputs to the full computation graph, and outputs do not have to be ultimate outputs either. It is possible to compile a function going from some intermediate variables of the graph to other intermediate variables, as long as the set of inputs contains all the information to compute the set of outputs. Several Theano functions can be compiled, computing different parts of the same computation graph.

During the compilation of a Theano function, first the relevant portion of the computation graph is cloned, then it gets rewritten by the application of *graph optimizations*, next some optimized C++ or CUDA code gets generated and compiled if necessary, and finally a callable object is built and returned to the user.

### Graph optimizations

The computation graph structure makes it possible to replace parts of the graph. For instance, a Variable node which is the output of one particular Apply node could be replaced by the output of a different Apply node, as long as they have the same type. Optimizations specify how to perform replacements of variables by other variables representing an equivalent computation. Some of them are *local*, which means they only look at one Apply node and can replace its outputs, some of them are *global*, and can examine the whole computation graph and perform arbitrary substitutions. Optimizations are mostly organized into the stages described below, even if there is some overlap.

-   *Canonicalize:* Put the graph in a canonical form, to ease the task of subsequent optimizations (for instance, $x*x \Rightarrow x^2$). It performs some simplifications as well, like removing duplicate computations, removing some unnecessary computations ($xy/y \Rightarrow x$), and computing the value of expressions if all their inputs are known (constant-folding, $2 + 2 \Rightarrow 4$).

-   *Stabilize:* Increase numerical stability, for instance $\log{(1+x)} \Rightarrow \mathrm{log1p}(x)$, where log1p is a stable implementation for small $x$.

-   *Specialize:* Insert faster implementations of operations. For instance, successive element-wise operations are fused together to avoid having to loop over a tensor several times.

-   *GPU:* Replace the default version of Ops and variables by GPU-specific versions, using either the old or new back-end, if a GPU is requested. Transfer Ops (CPU-to-GPU or GPU-to-CPU) are inserted so that the type of inputs and outputs is preserved, and around CPU-only operations.

-   *Inplace:* Replace the default version of Ops by a version that can work in-place, as a view or destructive operation over its inputs. The array types used by Theano, like `ndarray`, support arbitrarily-strided arrays, so all transposition operations, as well as basic slicing, can happen in place, in constant time. Some operations, like most element-wise ones, can overwrite their input and return it, to avoid allocating memory. Since destructive operations introduce additional dependencies between Apply nodes (a value can only be overwritten by the *last* operation to read it), dependency cycles have to be detected and prevented.

-   *Scan:* Optimize performance and memory use of Scan nodes. For instance, only keep the value for the last step of an output in memory if the whole sequence is not needed, merge different Scan nodes to perform computations only once, and move invariants out of the loop.

While individual optimizations or groups of optimizations can be individually enabled or disabled, some optimizers (sets of optimizations) are predefined: `'None'` does not include any optimization, `'fast_compile'` includes only canonicalization and transfer to the GPU, and `'fast_run'` (the default) includes most optimizations except for experimental and "unsafe" ones (removing assertions).

### Shared variables

Shared variables are symbolic variables that are associated with persistent values, that are shared between Theano functions. They can only be input variables (not intermediate ones), since their value is not the result of the computation of an Apply node. Shared variables are implicit inputs to all the Theano functions using them.

When compiling a Theano function, it is possible to specify *update expressions* for shared variables. These expressions are symbolic variables that represent the new value to assign the the shared variables at the end of each function execution. They are implicit outputs of the function, and will be computed along with the other outputs, before the value gets updated. Such update rules make it possible to update the array in-place in some cases, rather than returning a different array.

It is also possible to explicitly assign a new value to an existing shared variable, outside of a Theano function, as long as it is compatible with its type. Since the shape is not part of the type, it is possible for the shape of a shared variable to change. If a GPU is enabled, shared variables will be created on the GPU by default, to avoid transfers (this only works for `float32` arrays in the old back-end).

### C code compilation and caching

The code to compute output values given input values for each Op can be implemented either in Python or in C++ (or CUDA for GPU Ops), using the C API from Python and NumPy (and from CudaNdarray or GpuArray for GPU).

After the function graph is optimized, each Op generates the C++ or CUDA code for a Python module implementing that computation (including reading and writing from the right storage map), which is then compiled, and imported.

A persistent cache on disk makes it possible to avoid generating code twice for the same Op, and to avoid compiling again when different Ops generate the same code (this can happen for the same operation applied on different data types, or different numbers of dimensions, for instance).

## Function execution

Theano includes a runtime engine that, upon a Theano function call, determines the computation to be executed on which data and in what order, and orchestrate their evaluation. This was originally done by forward-traversing graphs from input to output, requiring all branches to be evaluated before outputs could be returned. The default runtime now uses a virtual machine (VM) system. By running small code units (each corresponding to an Apply node for one Op) and ignoring branches not necessary for correct computations, lazy evaluation is now possible.

The runtime uses a data structure containing pointers to storage for each variable (inputs and outputs of each Apply node), ordering constraints, pointers to the functions performing the computations, and information on what has been computed and needs to be computed in the current call. If the speed of execution is more important than memory usage, it is possible to keep references to ndarrays containing intermediate results, to prevent Python's garbage collection from freeing them, and to re-use it for the next run of the function, through the configuration flag `allow_gc=False`. The default is to allow the garbage collector to free the storage of intermediate values.

The C implementation of that VM (CVM) is the default runtime. Not only does this increase performance by running the runtime loop in C, if a C implementation of an Op is available, the CVM can directly execute it. This eliminates the overhead from a Python function call, which is especially advantageous when performing many operations on small operands.

A Python implementation is also available. It is more flexible and easier to instrument, which is useful to collect more profiling information (for instance, memory usage) and add callbacks for debugging.

## Extending Theano

If the existing Theano library does not include the operations required for a particular model, the framework was designed for easy extensibility. New Ops can be written by specifying the type of their input and output variables, and providing Python code to perform the evaluation. That Python code can use bindings to external high-performance libraries, or Cython, for instance. Methods can also be added to specify expressions for gradients and the R-operator (see Section ), and shape inference. Theano's self-testing functions can be used to validate outputs and check symbolic gradients against numeric evaluations among others.

As mentioned above, operators can also be implemented directly in C++ or CUDA. The raw code can be supplied as a string that the Python code uses to produce the code used by the graph compiler. For added convenience, Theano can now load code from an external C-like file with the `COp` class. The file is divided into sections that map to the different pieces of code that Theano requires. Keeping the Python and C code separate allows more readable code with better indentation. It also enables a clearer view of the C code itself since you can use your favorite C editor to modify that file with syntax highlighting.

A user can then write a new *optimization* to automatically insert that optimized operation in the computation graph, instead of the more naïve or slow version. This is especially useful when implementing an operation on GPU.

## Related software

Although Theano is developed and mainly used for research in machine learning and deep learning, it is not a deep learning framework in itself (see Section  for some machine learning frameworks based on Theano). However, it makes sense to compare the core features of such systems with Theano, as they all support the definition of a mathematical model in a symbolic way, and implement some automatic gradient computation.

TensorFlow  has a core in C++ and includes most of the features from Theano, in particular the graph-compiling approach, and symbolic differentiation (on full layers as well as on elementary operations), all directly accessible from Python through the API. In addition, it has a focus on distributed, multi-node computation. Even though a graph-rewriting engine is present (and used to distribute computation across devices, for instance) it does not seem to be used for mathematical expressions simplification or kernel fusion at the moment.

Torch7  has a different approach: it implements efficient CPU and GPU computation kernels in C and makes them available in Lua, but does not provide gradient expressions for elementary operations. Instead, packages like 'nn' and 'cunn' feature higher-level *layers* that can store parameters and provide methods to compute values for forward propagation, gradient back-propagation, and parameter updates. Many packages extend Torch's features, in particular Autograd[^9] provides automatic differentiation of code written in Torch, by building a graph that records the evaluation of expressions (even through loops and conditionals), and playing those records back to build an expression graph for gradients. That graph is symbolic as well, making it possible to express higher-order gradients. Moreover, an optimizer can rewrite the graph to make it more efficient to evaluate.

MXNet  and Caffe , both written in C++, feature the same kind of higher-level layers as Torch. MXNet can also express the gradients through those layers as symbolic layers themselves, giving more flexibility for the dispatching of the computation to different devices, and for memory reuse. It also allows distributed computation over multiple nodes. Caffe2[^10] is an experimental rewrite of Caffe that features explicit symbolic gradients in the computation graph, rather than a "backward" method of the layers.

Neon[^11] and Chainer  are two other machine learning frameworks written in Python, with GPU kernels, that feature symbolic computation graphs and symbolic differentiation. Neon's most prominent feature is its collection of highly-optimized GPU kernels, in particular for operations used in neural networks. Chainer instead builds its computation graph dynamically at the same time as its first evaluation, making it easier to express loops and conditionals.

# New features

Over the last couple of years, multiple improvements have been made in Theano, in particular for faster execution, including support for more operations on the GPU and multiple-GPU support (Section ), faster graph optimization, especially for larger graphs (Section ), and ease of use, with better error messages and tools for introspection, visualization, and debugging (Section ).

## Increased performance

### Abstract Ops and 2D convolutions

Convolution operations are at the core of Convolutional Neural Networks (CNNs) that have lead to spectacular advances in machine learning problem involving visual data . A more detailed description of the convolution operations can be found in .

The multiplication of available implementations for convolution (CPU-GEMM, GPU-cuDNN, GPU-GEMM, FFT, …) available in Theano has increased the need of a flexible convolution interface that easily allows to switch between those implementations, each implementation having different speed and memory trade-off, as well as different software dependencies. To suit this need, Theano 0.8 introduces abstract Ops that allows to disentangle the interface of an Op to their actual implementation. An abstract Op introduces is a place-holder Apply node in the graph, corresponding to a given operation, that does not provide an actual implementation. For each optimized implementation of that operation, there is an optimization that will insert an Apply node for that optimized Op instead of the abstract Apply node during the compilation phase.

In particular, Theano proposes three abstract Ops for convolution: `AbstractConv2d`, `AbstractConv2d_gradInputs`, and `AbstractConv2d_gradWeights`, that correspond respectively to the forward convolution, the convolution gradient w.r.t. inputs and the convolution gradient w.r.t. weights. Each abstract Op can be replaced by one of the different implementations. By default, if a GPU is enabled and cuDNN is available, Theano will use it (see Section ), otherwise it will fall back to using the GEMM version. A slow, Python-only implementation is part of the abstract Ops for debugging purposes. The optimizations can be included or excluded using the configuration flags, which makes it possible to manually select a specific convolution implementation.

### Using cuDNN

Efficient CUDA primitives for neural networks are implemented in the cuDNN library , in particular convolutions, pooling, and their gradients. Several implementation of convolutions (and gradients) are provided, with the same interface, with performance and memory usage that depends on the actual shape of the data and filters. Since the best implementation can be different for different convolutions in the same model (depending on their size) and on different hardware (depending on the available memory), cuDNN also provides a heuristic to guess the best algorithm given shapes, and to actually time the different implementations (that are feasible given the available free memory) and select the fastest one.

Theano wraps cuDNN 2D and 3D convolutions and their gradients, and provide options to select the algorithm to use, either explicitly or using one of the following special values: `'guess_once'`, `'guess_on_shape_change'`, `'time_once'`, or `'time_on_shape_change'`. This selection can be done individually for each Apply node in the graph, and configuration flags select the global default for the forward convolution, the gradient w.r.t. the data, and the gradient w.r.t. the weights. Theano also wraps pooling operations, as well as softmax and log-softmax operations. More operations will be added in the future.

### CNMeM integration

Another improvement to the GPU performance comes integrating the CNMeM library,[^12] and using the allocator and deallocator it provides. The main issue was that calling `cudaFree` is synchronous, so it forces the synchronization of all the streams on the device, waiting for them to finish, which seriously limited the potential for parallel execution of different kernels. A previous option was to keep memory allocated for intermediate values between calls, as mentioned in Section , but the amount of memory typically available on GPU devices is limited.

CNMeM works by allocating large memory pools using `cudaMalloc`, returning chunks of it when its allocator is called, and keeping track of which ones are released by its deallocator. Theano makes it possible to reserve part of the GPU memory from the start, using `lib.cnmem=0.9` to reserve 90% of the memory for CNMeM. The new GPU back-end does not use CNMeM, but implements a similar strategy, with asynchronous allocator and deallocator and a memory pool.

### Improvements in Scan

Important speed improvements have been made to Scan, in addition to making it more stable, and supporting more cases. The time to optimize and compile graphs containing Scan Apply nodes has been reduced a lot, and the execution time of the resulting function has improved as well.

The optimizations related to Scan (pushing computation out of the loop, removing useless computation) have been improved so they can be applied faster. Additional optimizations have been added, so that more computation can be moved out of the loop, for increased execution speed.

The execution back-end of Scan has been made more efficient as well, by removing some of the bookkeeping overhead, and making the internal function write directly into the right output buffer at each execution step, rather than having to copy the intermediate results each time.

The `grad` method of Scan has been rewritten to scale better in the case of large numbers of input and output variables, and to generate a cleaner graph. That cleaner graph can lead to a faster optimization time, since less rewriting is needed and the inner graph is smaller, and faster execution as well. In the case of nested symbolic loops, the observed speed up in compilation time was sometimes huge, going from hours to minutes.

Finally, an additional keyword, `strict`, has been added to the `scan` function. It prevents shared variables from being implicitly added as non-sequence inputs to the inner function. This forces the user to explicitly provide all non-sequences needed in the inner function, which may not be the shared variables themselves, but rather outputs of some computation done of them. In that case, doing so prevents pulling that computation inside the loop, which can speed up the optimization as well as the execution.

### New gpuarray-based back-end

Theano now features a new GPU backend based on libgpuarray . This new back-end brings in several improvements over the previous one. The most visible improvement is that it supports all the usual data types, instead of being limited to float32 data. In particular, it supports half-precision floating point values (float16). As did the previous back-end, this one supports views and strides to avoid copies and reuse memory whenever possible.

libgpuarray[^13] is a separate project with the aim of providing a ndarray-like object on the GPU. It has a C interface so that it can be reused in other projects that don't use Python. It also supports 64-bit indexing, so that arrays with more than $2^{32}$ elements are supported.

Another noticeable improvement is that we have basic support for OpenCL, however a sizable portion of the GPU Ops in Theano do not currently support it. This could be fixed with some porting effort.

The new back-end also allows using multiple GPUs in the same function to do model parallelism. One example of such a model is the two-stack variant of AlexNet . This however may be hampered by the Python Global Interpreter Lock (GIL) in some cases, meaning that one will get correct results, but may lose parallelism.

Several new features that help performance are present, but not obvious. One of these is that all computations are transparently asynchronous, which allows the CPU part of the Ops to execute in parallel with the GPU part. There is a mechanism keeping track of the dependencies between operations to ensure that the right data is always used. Data transfers are automatically done on a separate stream, so they can overlap with the computation.

The new back-end is now fully functional, and well tested for correctness. It supports almost all the operations of the old back-end on CUDA-capable devices, including wrapping cuDNN for efficient convolutions, but we are still in the process of tuning some of its kernels for a better performance. In particular, int64-based indexing can be significantly slower than int32, so some adjustments have to be made.

### Data parallelism with Platoon

To take advantage of multiple computing devices, there are two main approaches: model parallelism and data parallelism. Model parallelism consists in splitting the model itself into multiple parts and have those parts computed by different devices. It requires a careful balancing of the size of the parts and of the communication costs to ensure optimal performance. Data parallelism on the other hand is about splitting your input data in multiple parts, and running multiple copies of the model. It requires attention to model synchronization so that the copies don't drift apart too much during training, and to the way of aggregating the results produced.

Usually, data parallelism on a single machine is done using multiple threads, but this approach is unworkable in Python because of the Python GIL. Because of this, we have to turn to multiple processes and this presents a new set of challenges. Platoon[^14] is a package that has been developed to to address those challenges and help train Theano models faster by using data parallelism.

Platoon features a central controller process, that communicates with different worker processes, each using Theano to train a copy of the model on a CPU or GPU. It uses shared memory to share model parameters between workers, in order to avoid inter-process communication overhead. The communications with the central controller are sent asynchronously, so that the worker does not have to wait for a reply. There is also a script to launch all the workers and monitor them while running that provides a central "job" to wait for on clusters.

Two ways of performing the updates on the central parameters are currently implemented: Asynchronous SGD (ASGD), similar to Downpour SGD , and Elastic Averaging SGD (EASGD) . Other algorithms can be added by implementing additional parameter synchronization rules.

## Faster compilation of graphs

### Faster, simpler optimizer

As mentioned in Section , some sets of optimizations are pre-defined and can be easily specified. One of these optimizers, `'fast_compile'`, has recently been upgraded to include the optimizations that transfer computation to a GPU, as well as the optimizations necessary to make those optimizations apply. This drastically shortens the graph optimization time, at the cost of a slightly slower execution time and increased memory usage. That option can speed up the development or prototyping phase of a model, allowing the developer to iterate faster.

### Swapping updates without recompiling

It is now possible to copy functions using the `function.copy()` method. This can be useful when creating functions that are similar but use different shared variables or update parameters, for instance when creating test and validation functions. Most importantly, the optimized graph of the original function is copied, meaning compilation only occurs once.

The interface for `copy` lets users specify which shared variables to swap, and whether or not updates are carried over. It is also possible to have copied functions share intermediate storage in memory (storage that is not input or output). When this is combined with disabled garbage collection, this can increase execution speed and save memory.

### Save and reload optimized graphs

Optimized computation graphs, such as the ones in Theano functions, can now be serialized using the `pickle` module, and get de-serialized without being optimized again. It is possible to force the re-optimization, for instance if the set of optional dependencies available has changed between saving and reloading, in which case the function may not run (if a dependency has been removed) or be sub-optimal (if one has been added). This is especially useful when check-pointing and restoring running experiments. Note that the C++ or CUDA code may still need to be recompiled.

## Visualization, debugging, and diagnostic tools

Since the definition of Theano functions is separate from their execution, some specific tools have been developed to help users visualize parts or the whole of the computation graph, pinpoint the origin of errors, and understand what is happening at execution time.

### Interactive visualization with d3viz

Interactive visualization of computation graphs is now possible with the `d3viz` module, which extends Theano's printing module. Instead of outputting a text representation (like `debugprint`) or creating a static picture (like `pydotprint`), it creates an HTML file, which can be opened with current web browsers. An example is shown in Figure .

Several features are supported. Users can zoom different regions, move graphs via drag and drop, and position nodes both manually and automatically. The visualisation can retrieve additional information about nodes and edges such as their data type or definition in the source code, edit node labels and visualize profiling information. Nested graphs such as `OpFromGraph` nodes can also be explored by expanding or shrinking the nodes as needed.

Internally, `d3viz` represents a compute graph in the Graphviz DOT language, using the pydot package, and defines a front-end based on the d3.js library to visualize it. However, any other Graphviz front-end can be used, which allows to export graphs to different formats such as PNG and PDF.

### Test values

Detecting errors in the way a mathematical expression is implemented in Theano can be a challenge, since it is not possible to directly map an intermediate Variable node to the value that will be associated to it at execution time. To mitigate this problem, it is possible to associate a *test value* to input variables, and to compute automatically values associated to intermediate variables as soon as they are defined. This makes it much easier to detect shape mismatches, for instance, or unexpected values.

Note that these values are computed only once, when the graph is built. That means that stability optimizations will not be applied to these values, so NaN (not-a-number) values could be produced during that phase, even if they would not be present when evaluating the optimized graph.

### NanGuardMode

A frequent symptom of issues when optimizing a model is the appearance of NaN (not-a-number), infinity, or very large values. They can indicate a wide range of issues, e.g., use of un-initialized memory, lack of numerical stability in the computation, divergence of the algorithm itself.

To help diagnosing the appearance of such values, NanGuardMode is an instrumented version of the runtime environment that can check the values of inputs and outputs of each Apply node during execution, and raise an error when some problematic values are detected.

### The PdbBreakPoint Op

`PdbBreakPoint` is an Op designed to check the value of a *condition*, which is a symbolic expression, during the execution of a Theano function. If the condition is met, then the program will drop into the Python debugger (`pdb`), and make available the values associated to a list of pre-defined *monitored* variables. This is especially useful when something goes wrong during the training of a model, but only after a number of iterations, so it is not practical to log all values all the time.

### Keeping the creation stack trace

When a variable is created, part of the stack trace is recorded, in particular the line of the call that created it. For instance, if variable `z` is created by calling `z = a + b`, then the line where that expression is called is associated to `z`. If evaluating that expression fails, for instance because `a` and `b` have incompatible shapes, then the error message will mention that file, line, and line number.

A challenge of that mechanism is that, when optimizations are applied, the replacement variables are not created at the same place as the ones they replace (or that "correspond" to them in a more general sense). In fact, they are created inside the optimization, so no stack trace is associated to them. For instance, if the expression above is optimized to move `a` and `b` to a GPU, and `z` gets replaced by `host_from_gpu(gpu_z)` where `gpu_z = gpu_add(gpu_a, gpu_b)`, then the replacement for `z` can easily retain the original stack trace, but `gpu_z` would not.

To improve this feature, we are currently in the process of going through all optimizations, so that they assign the creation stack trace of the original variable (or variables) to the "corresponding" or equivalent one when they create replacements or new intermediate variables.

# Benchmarks

This section aims at giving a sense of the performance might expect from Theano against some of its largest competitors among machine learning research software, on different kinds of models. We used publicly-available software to compare against, when possible. We have made some of the benchmarking code public as well already, and will try to provide the remaining code as well in the future.

The goal of having more extensive benchmarks, on a wider variety of models and frameworks, is more easily attained by online projects, that can provide a picture more up-to-date. Among these projects, we can cite convnet-benchmarks,[^15] rnn-benchmarks,[^16] and hopefully DeepMark[^17] in the future.

We benchmarked Theano against Torch and TensorFlow (Section ), on three kinds of popular machine learning models: convolutional networks (Section ), recurrent neural networks (Section ), and recurrent neural networks for sequence-to-sequence mapping (Section ). Finally, we show how the computation speed scales when using multiple GPUs with Platoon (Section ).

## Setup

All the benchmarks were run on a NVIDIA Digits DevBox, with 4 Titan X GPUs, and a Core i7-5930K CPU. All the benchmarks except for data-parallelism were run on only one GPU, which was not the one used for running the X server (using `CUDA_VISIBLE_DEVICES`). We used Cuda 7.5.17, with cuDNN v4 (version 4007), and data type float32, for all frameworks and all experiments.

The compared software were installed as follow:

-   Theano was installed from the development version, at commit `1bd371c`. The following configuration flags were used: `floatX=float32`, `lib.cnmem=0.45`, `device=gpu0`, `optimizer_including=unsafe`, `dnn.conv.algo_fwd=time_once`, `dnn.conv.algo_bwd_filter=time_once`, `dnn.conv.algo_bwd_data=time_once`. For fast_compile experiments, the additional option `optimizer=fast_compile` was provided.

-   TensorFlow 0.8 was installed from the binary package.

-   Torch7 was installed from <https://github.com/torch/distro> at commit `ffffc39`.

## Convolutional networks

We measure the performance of four different convolutional models, that have been successfully used on the Imagenet dataset:

-   AlexNet, the one-column variant from , with a batch size of 128;

-   OverFeat, the *fast* variant from , with a batch size of 128;

-   VGG, also known as OxfordNet, model A , with a batch size of 64;

-   GoogLeNet V1 , with a batch size of 128.

We used the code from <https://github.com/soumith/convnet-benchmarks> at commit `84b5bb1` for Theano, Torch, and TensorFlow. We report the processing time per minibatch, for the forward and the backward pass.

The results, presented in Figure , show that Theano is slightly slower than Torch and TensorFlow, but the performance is comparable, both for the forward and the backward passes. Furthermore, using the `fast_compile` optimizer shows a slow-down between 10% and 25% only, which is a reasonable trade-off when developing or exploring a new model.

## Recurrent neural networks: LSTM on Penn Treebank

To showcase recurrent network models, we benchmarked variants of the LSTM model applied to the Penn Treebank dataset described in . We compared:

-   the Torch implementation available at <https://github.com/wojzaremba/lstm>;

-   the TensorFlow implementation showcased at <https://www.tensorflow.org/versions/r0.8/tutorials/recurrent/>;[^18] and

-   the Theano implementation available at <https://github.com/caglar/rnn_benchmarks>.

We measured words per second during training, and report results on the following models:

-   Small: Single Layer, 200 hidden units, sequence length: 20;

-   Medium: Single Layer, 600 hidden units, sequence length: 40;

-   Large: Two Layers, 650 hidden units each, sequence length: 50.

All three models used dropout on non-recurrent connections during training, following . The batch size was set to 20.

Figure  shows that Theano comes second behind TensorFlow for the small model, but is slightly faster on the medium and large model. Torch was slower than Theano on all three models, and perhaps more surprisingly, slower than the fast_compile version of Theano on the two larger models.

## Sequence-to-sequence: Caption generation from video

In this section, we use the sequence-to-sequence mapping model from . The input is a series of video frames and the output is a one-sentence English description of the input. Each input video frame is preprocessed by a GoogLeNet that was pre-trained for classification on ImageNet. The representation of the frame is thus a 1024 vector. The entire input is therefore represented by (M, F, 1024) where M is the minibatch size, and F is the number of frames. The output size is (M, L), where M is the minibatch size and L the sentence length (padding is used within a minibatch to ensure the same length, but different minibatches could have different L). Specifically, the model is written as $P(S|V)$, an LSTM on the sentence $S$, conditioned on the video $V$. $V$ is a weighted sum of frames representations.

The original code for  is available at <https://github.com/yaoli/arctic-capgen-vid>. We used simplified versions, in Theano and TensorFlow, instrumented for profiling, which will be made public in the future. There was no publicly available implementation in Torch. Theano with fast_compile could not run because it was requiring too much memory. We report the processing time per minibatch, for the forward and backward passes, using three different batch sizes.

Figure  shows a small advantage to Theano for the forward pass, but a disadvantage for the backward pass. The total time was comparable overall, with Theano being slightly faster on smaller batches, and TensorFlow being faster on larger ones. As expected, the time per minibatch grows slower than the minibatch size, because the potential for parallel computation is greater with larger batches.

## Data parallelism for LSTM

We re-use the models from Section , this time using Platoon to train on multiple GPUs on the same machine, using ASGD. We report results for 2 GPUs (using devices `gpu1` and `gpu2`) and 4 GPUs, compared against the results on 1 GPU obtained without Platoon and reported in Section . We measured the overall processing speed (words per second) during training when synchronizing the models after every minibatch, and when synchronizing only every 100 batches. The benchmarking code using Platoon will be made public soon.

Figure  shows a consistent increase in processing speed when adding more GPUs. As can be seen on the left, communication and synchronization overhead make that scaling sub-linear when synchronizing after every single batch, we found a speed-up between 1.6 and 1.7 for 2 GPUs and around 3.2 for 4 GPUs across all three models. Synchronizing only every 100 batches, on the right, brings the computation speed-up close to the theoretical optimum, at 2 for 2 GPUs and between 3.9 and 4 for 4 GPUs.

# Limitations and challenges

Despite the progress made in recent years and our best efforts, there remain some limitations or shortcomings in Theano. Some of these issues have been addressed by competing frameworks mentioned in Section , and by other projects like CGT (Computation Graph Toolkit).[^19]

## Limitations from Python

Since Theano uses Python as its core language, and uses NumPy arrays and other Python objects to store values, it is affected by Python's limitations. The main one is the Python GIL, that limits concurrent execution of threads. We have seen that it is possible to make single-threaded execution fast by compiling binary modules that are then loaded in Python (Sections  and ), and it would also be possible to release the GIL during the execution of these functions. However, the GIL has to be acquired again each time references to Python objects are added or removed, when using the C API of Python and NumPy. Since the execution of such functions is usually quite short, most threads would spend their time waiting for the lock instead of performing actual computation.

Since Python has a concept of threads and expects to be in charge of threading, it is also not possible to launch different, independent Python interpreters in different threads of the same process, as is possible with Lua for instance.

To avoid that issue, we could use a different n-dimensional array structure, that is accessible directly from C++ without actually being a Python object, like the one libgpuarray provides on the GPU. It would require Theano to explicitly manage memory allocation and deallocation, in a thread-safe way. It would also require to rewrite all the C++ and CUDA code for existing Ops, so that they use a different interface for reading their input data and writing their output data. Finally, it could make it harder to create new Ops by integrating existing Python code.

## Graph optimization time

The execution time of the graph optimization phase is not scaling well with graph size. Currently, it is scaling supra-linearly relative to the number of nodes. One issue is that some groups of local optimizations try to apply over and over, until none of them can be applied any more, and the graph stops changing. In practice, it can force a number of passes through the whole graph that becomes bigger for bigger graphs (the chances of some local optimization applying somewhere are higher).

An option would be to completely reorganize the existing optimizations so that they are more lightweight, and can be applied in a fixed number of passes through the graph. It could be possible, for instance, to use a one-pass or two-pass optimization phase, like CGT does. Doing that without any regressions in the stability optimizations could be a large-scale project.

## Code compilation time

Currently, the same Theano Op can generate a large quantity of different C++ or CUDA modules, depending on its properties at compile time, such as the data type of inputs and outputs, whether it will run in place, and other flags determining its behaviour. Compiling and loading those modules can take time and add a load on the file system.

To alleviate those issues, it would be possible in most cases to pass that information dynamically at runtime, instead of hard-coding it in the generated code. This approach is already being used in the new back-end to specify which GPU should be used for the execution of a particular Apply node, but it could be generalized.

## Loops and control-flow structures

Using Scan for loops, and the `ifelse` lazy Op for conditionals, has proven a useful way of expressing control-flow operations. However, with an increasing need for more flexibility (attention mechanisms, nested loops, recursive loops, changes in shape between iterations of the same loop), we may need a more principled way of expressing these structures.

One appealing way would be to use *switch* and *merge* Apply nodes in the computation graph, like in a dataflow graph . This is the approach taken by TensorFlow  for symbolic loops. This would require adding support for cycles in the computation graph in these circumstances, extending the runtime to be able to recompute values inside the loop, and rewriting all the graph optimizations currently existing for Scan, including the ones limiting memory consumption.

## Multi-node parallelism

Scaling model execution and training to multiple machines is outside of the scope of Theano's core, but additional packages could be developed to interface with Theano, in the same way Platoon does for multiple GPUs in a single node. In fact, tools like parameter servers and coordinators do not have to be specific to Theano, and could be common to different frameworks.

## Improving memory usage

Given the limited availability of on-board GPU memory, memory consumption is often a bottleneck for training machine learning algorithms. This can limit the size and modelling power of trainable models, and make the processing power of GPUs under-used, for instance when batch sizes have to be reduced. In addition to storing intermediate values in a lower-precision format (for instance, storing data as float16 is supported in Theano's new GPU back-end), different options could be explored and combined:

-   Change the order of execution of computations, so the peak memory usage is reduced. This can be done statically before the function is executed, or dynamically, for instance by detecting that memory is insufficient and waiting for some other computation to finish and free intermediate values.

-   Move intermediate values to the main (CPU) memory, or to another GPU's memory, if it is not needed for a while, and transfer it back before it is used again. This method has been successfully implemented by .

-   Free intermediate values, and recompute them when they are needed again. This approach has been used in , and can be especially useful for fast operations that have large outputs.

## The future of gradient-based computation frameworks

Tools like Theano and TensorFlow are compilers for mathematical expressions, in that they require the code (or computation graph) to be defined first, and then executed. On the other hand, Torch works more like an interpreter: the computation is done as soon as the expression is called. It could be interesting to explore how to apply JIT (just-in-time) compiler ideas to the computation graph, to combine the immediate response and flexibility of an interpreter (including using control flow statements like `if`, `for`, `while`, from the language directly), and the performance gains of a compiler when an expression has to be evaluated multiple times.

Most machine-learning frameworks can now share efficient implementations of GPU kernels, such as the ones published by NVIDIA (cuDNN) and Nervana. Graph optimizations could be another component shared between projects, maybe through a common language to define computation graphs and such optimizations. It could be common to machine learning frameworks and computer algebra systems (CAS) such as SymPy  and SympyCore.[^20]

# Conclusion

Theano pioneered ideas for efficient gradient-based computation that are now part of most mainstream machine-learning research libraries, for instance combining a high-level scripting language with highly-optimized computation kernels, especially using GPUs, symbolic computation graph, and symbolic differentiation. Some other features of Theano, like graph rewriting and optimizations, and automatic generation and compilation of kernels, are starting to become more widely used as well.

Continuous improvements have been made to Theano's functionality, usability, and performance, for instance wrapping libraries like cuDNN, and integrating ideas that have been successfully explored and implemented by other frameworks, like data parallelism and model parallelism for distributed computation. Computation performance is on par with other major research software, like Torch and TensorFlow.

There are ways to improve Theano (and other frameworks as well) by taking inspiration from other machine learning software (sometimes more experimental). Longer-term improvements could be the result of collaborations with other fields, for instance CAS, and language and compiler design, in order to build a next generation of mathematical computation software.

[^1]: code available at <https://github.com/Theano>

[^2]: Some OpenCL support is available in the new GPU back-end, but it is still limited and experimental.

[^3]: <https://github.com/Theano/Theano/>

[^4]: <https://groups.google.com/group/theano-dev/>

[^5]: <https://groups.google.com/group/theano-users/>

[^6]: <http://stackoverflow.com/questions/tagged/theano>

[^7]: <https://pypi.python.org/pypi>

[^8]: For instance, the deep learning tutorials at <http://deeplearning.net/tutorial/>

[^9]: <https://github.com/twitter/torch-autograd/>

[^10]: <https://github.com/Yangqing/caffe2>

[^11]: <http://neon.nervanasys.com/>

[^12]: The original code is available at <https://github.com/NVIDIA/cnmem>, Theano includes a copy of it.

[^13]: <http://deeplearning.net/software/libgpuarray/>, code available at <https://github.com/Theano/libgpuarray>

[^14]: <https://github.com/mila-udem/platoon>

[^15]: <https://github.com/soumith/convnet-benchmarks/>

[^16]: <https://github.com/glample/rnn-benchmarks>

[^17]: <https://github.com/DeepMark/deepmark>

[^18]: Code at <https://github.com/tensorflow/tensorflow/tree/master/tensorflow/models/rnn/ptb>

[^19]: <http://rll.berkeley.edu/cgt/>

[^20]: <https://github.com/pearu/sympycore>