Inside Torch and Autograd: How Deep Learning Libraries Compute Gradients

Wiltschko says the first half of his talk will be practical and the second will dive into Torch Autograd and the concepts shared by all deep learning libraries. He credits the material to colleagues including Soumith Chintala, who he calls the czar of the Torch ecosystem, Hugo Larochelle and Ryan Adams. He describes Torch as an array programming language for Lua, much as NumPy is for Python, with the tensor as its core data type.

He explains why Lua is a reasonable choice. LuaJIT is unreasonably fast, with loops about as fast as C, so it works as convenient glue over high-performance C code. The language is defined by roughly 10,000 lines of C and can be learned in an afternoon, and it was designed to be embedded, which makes calling C libraries easy.

He points to Lua's use in places such as World of Warcraft quest scripting, Adobe Lightroom's interface, and Redis and nginx scripting. He relays the apocryphal story that Torch used Lua because a developer found Python frustrating on embedded chips. Torch also offers first-class interactive GPU computation, so moving data to the GPU means everything after that runs there without hand-written CUDA kernels.

Wiltschko adds that Torch has strong industry support from Twitter, Facebook AI Research and Nvidia but is not owned by any of them. He says research papers often come with high-quality Torch implementations, including image captioning, neural style transfer, sequence-to-sequence models and generative adversarial networks. The community is small compared with Python's, but strong in deep learning, and a Python programmer can become productive in Lua in about a week.

Wiltschko says there is no silver bullet among deep learning libraries. In his rough positioning, TensorFlow is the largest and was engineered for industrial production, Theano for research, and Torch leans toward research, though Twitter has pulled it toward production. He says every piece of media coming into Twitter passes through a Torch model.

Torch's design tenets are immediacy with no compile step, imperative programming, and minimal abstraction, so that it usually takes only one or two jumps from a function's definition to the C code that runs. He then explains the memory model: a tensor is a view onto storage, defined by its size, stride and offset, laid out row-major. Selecting a row or column returns a new tensor with different stride and offset rather than copying data.

He demonstrates the consequence: fill a slice with a number and the original tensor's values are overwritten, because they share memory. The same operations work on GPU tensors after requiring the CUDA package.

Wiltschko sketches the pieces needed to train a network: data loading, training code, the network itself, a cost function and an optimizer. In Torch, the nn package lets you click together layers like Lego blocks inside a sequential container, with convolutions, nonlinearities, pooling and a log-softmax. It shines for linear architectures, has a small API for forward and backward passes, and moves a whole model to the GPU with one call.

For odd architectures, such as the complicated NLP model Richard Socher described the day before, nn becomes awkward. Twitter's Torch Autograd package handles this by letting you glue together pieces as small as addition and multiplication while still getting correct gradients.

The optim package supplies optimizers like stochastic gradient descent, Adagrad and Adadelta, using a functional style in which you pass in a function that evaluates the network and returns gradients. He warns that parameters must live in one linear tensor, so interacting networks need their parameters combined first.

Wiltschko asks the audience to appreciate stable abstractions in scientific computing, such as arrays from Fortran, linear algebra libraries and MATLAB-style collections of routines popularized by NumPy. Machine learning adds another one: automatic differentiation. He quotes a definition from Barak Pearlmutter that it mechanically calculates derivatives of functions expressed as computer programs, at machine precision and with complexity guarantees.

That separates it from finite differences, which are numerically unstable, and from symbolic differentiation, which can produce exploding expressions. He notes it has been rediscovered several times, with early descriptions by Wengert in 1964 and an implementation by Speelpenning in 1980, and that backpropagation is a special case. Fields like weather modeling and computational fluid dynamics have used it for decades with more sophisticated tools than machine learning has.

Machine learning uses reverse mode, he says, because forward mode is a bad idea. Using a small traced program, he shows that forward mode produces one derivative per pass, so a network with millions of parameters would need millions of evaluations. Reverse mode first evaluates the whole program, then walks backwards and obtains all the derivatives at once.

Wiltschko shows that a whole neural network in Autograd is just a plain function for predictions and a loss, plus a call to `grad` on the loss function. Because `grad` is a higher-order function returning a gradient function, training is calling it and updating parameters with stochastic gradient descent. He compares the grad symbol to the nabla in mathematics.

Under the hood, Autograd replaces Torch functions with spy functions that record inputs and outputs while running them normally, building the computational graph just in time. It then walks that list backwards, looking up the partial derivatives that have been written for each Torch function. Since it only records arithmetic operations, it ignores control flow, so if statements, data-dependent for and while loops, and even recursion all differentiate correctly.

He also shows that you can override gradients. Taking the floor of a number has zero derivative almost everywhere, but you can define your own gradient to pass through it, which would matter for things like a differentiable JPEG-style compressor with a quantization step. Autograd provides an easy mechanism for plugging in such custom gradients.

On how libraries compare, Wiltschko first looks at granularity. Some give you a fixed model, as in his "don't be a hero" regime, echoing Andrej Karpathy's advice to reuse others' networks. Middle libraries like nn, Keras and Lasagne let you assemble layers, while Autograd, Theano and TensorFlow let you compose primitive operations because their developers provide partial derivatives at that level.

A second axis is how graphs are built. Static graph libraries such as nn and Caffe fix the structure ahead of time, offering little optimization. Just-in-time libraries such as Autograd and Chainer allow any graph that can change per sample but have little room for compiler optimization, and ahead-of-time ones like TensorFlow and Theano hand a graph to a runtime that can optimize heavily but handles control flow awkwardly. He notes a pure data-flow graph has no room for control flow and mentions Cliff Click's sea-of-nodes representation as an unexplored option.

He shows a small network written three ways in Autograd: with primitive operations, with nn layers turned into functions, and with ready-made standard models. At Twitter, he says, Autograd has accelerated high-risk experiments such as a tree-structured loss that gives partial credit for guessing a related class. Models trained with it run in production, with no speed penalty at test time because the gradient machinery disappears.

Wiltschko lists techniques from automatic differentiation that machine learning has not adopted yet. These include checkpointing, where you recompute cheap values like pointwise nonlinearities instead of storing them, mixing forward and reverse mode for unusual graph shapes, stencils for differentiating image-processing code, source-to-source transformation, and efficient higher-order gradients for Hessian-vector products. He closes by saying Lua and the tools are installable through Anaconda.

In the Q&A, he says Lua's visualization tools lag Python's, with Bokeh the best option. For production, he says calling Torch from C is easy, with no appreciable inference overhead because Autograd's training machinery is absent at test time, and Lua itself is tiny enough for embedded use. Asked why to choose Torch over Keras with TensorFlow, he says people reach for Torch when they want to reason about performance, since heavy compiler infrastructure makes problems harder to diagnose.

On serving models from Java or Python web stacks, he suggests serializing weights and loading them with a native library, or for feed-forward networks hand-writing the matrix multiplies. Twitter instead runs Lua virtual machines inside Java through JNI. He adds that the time to get predictions into a service is not their bottleneck, and he knows of no TensorFlow Serving equivalent for Torch.