Myia: A Differentiable Language for Deep Learning Olivier Breuleux - - PowerPoint PPT Presentation

Institut des algorithmes d’apprentissage de Montréal

Myia: A Differentiable Language for Deep Learning

Olivier Breuleux

Computer Analyst, MILA

Bart van Merriënboer (MILA, Google) Arnaud Bergeron (MILA)

Paradigm

How deep learning and language design intersect

Autodiff

What it is. How it works

Myia

Our proposed solution

Paradigm

How deep learning and language design intersect

Autodiff

What it is. How it works

Myia

Our proposed solution

Deep Learning

(Elgammal & al., 2017)

Deep Learning

Features + Composition + Learning

?

?

?

? ? ?

? ?

? Data Trained structure Gradient descent Initial structure (automatable)

Needs

General purpose: Express complex compositions using control flow. Fast: Leverage parallelism and GPU to process millions of features. Portable: Serializable, support multiple hardware. Goal: a language adapted to the needs of machine learning, past and future Differentiable: Language support for gradient descent.

General purpose

DL algorithms are increasingly complex Feedforward (trivial) Recurrent (loops) Recursive (recursion)

…?

General purpose

DL algorithms are increasingly complex

• More and more language features needed
• Most existing frameworks are limited
• Awkward abstractions
• No recursion
• High level abstraction increases productivity
• Focus on the algorithm over implementation details
• Effortless abstractions encourage their use

Fast

• Scale to millions of parameters
• Lots of parallel operations
• Matrix multiplication, map, reduce
• Can work with low precision (float32, 16, even 8 bits)
• Leverage adapted hardware (GPU, TPU)
• Loop fusion
• Automatic parallelization
• Memory management

Portable

• Serializable models
• Code + parameters (data)
• Run on mobile and embedded systems
• Seamless transfer from research to

production

• Avoid being tied to an ecosystem (e.g.

Python/numpy)

• e.g. ONNX (but more general)

Differentiable

Paradigm

How deep learning and language design intersect

Autodiff

What it is. How it works

Myia

Our proposed solution

Derivative

d dxf(x) = d f dx = f 0(x) = lim

✏!0

f(x + ✏) − f(x) ✏

Derivative

d dxf(x) = d f dx = f 0(x) = lim

✏!0

f(x + ✏) − f(x) ✏

Chain rule

d dxg(f(x)) = dg d f d f dx d dxh(f(x), g(x)) = ∂h ∂f d f dx + ∂h ∂g d f dx

Total derivative

In many dimensions

Jf(x) =   

∂f1 ∂x1

. . .

∂f1 ∂xm

. . . ... . . .

∂fn ∂x1

. . .

∂fn ∂xm

  

Suppose you have a vector of inputs and a vector of outputs.

f : Rm → Rn

The Jacobian is the matrix of partial derivatives.

Chain Rule

y1 = f(x) y2 = g(y1) y3 = h(y2)

f : Rm → Rp g : Rp → Rq h : Rq → Rn

Jhgf(x) | {z }

n⇥m

= Jh(y2) | {z }

n⇥q

Jg(y1) | {z }

q⇥p

Jf(x) | {z }

p⇥m

The derivative of a straight composition of functions is the product of their Jacobians

In what order?

Ordering the Jacobians

( Jh(y2) | {z }

n×q

Jg(y1) | {z }

q×p

| {z }

n×p

) Jf(x) | {z }

p×m

Jh(y2) | {z }

n×q

( Jg(y1) | {z }

q×p

Jf(x) | {z }

p×m

| {z }

q×m

)

Forward Reverse Cost Cost

qpm + nqm

nqp + npm

= m(qp + nq)

= n(qp + pm)

Forward vs Reverse

Forward mode is good when there are few inputs.

• Easy to implement: dual numbers.

x → ✓ y1, dy1 dx ◆ → ✓ y2, dy2 dx ◆ → ✓ y3, dy3 dx ◆

Reverse mode is good when there are few outputs.

• Hard to implement: execution is reversed.

x → y1 → y2 → y3 → dy3 dy2 → dy3 dy1 → dy3 dx

Deep Learning

Deep learning involves computing the gradient of a scalar cost with respect to millions of parameters. We need reverse mode.

✓ ← ✓ − ✏@L @✓ where ✓ = (W1, W2, . . . , b1, b2, . . .)

Timeline

Modules SN, Torch, Caffe 1987 2002 2013 Operator overloading Autograd, Chainer, PyTorch 2014 2015 2017 Graphs Theano, TensorFlow, MXNet 2008 2015 2015 Source code transform Tangent, Myia 2017 2018

Timeline

Modules SN, Torch, Caffe 1987 2002 2013 Operator overloading Autograd, Chainer, PyTorch 2014 2015 2017 Graphs Theano, TensorFlow, MXNet 2008 2015 2015 Source code transform Tangent, Myia 2017 2018

Modules

• No automatic differentiation
• Object-oriented approach
• Each operation a stateful object with a forward and a backward method
• Transparent, but tedious

What if we want to connect them in more complicated ways?

Timeline

Modules SN, Torch, Caffe 1987 2002 2013 Operator overloading Autograd, Chainer, PyTorch 2014 2015 2017 Graphs Theano, TensorFlow, MXNet 2008 2015 2015 Source code transform Tangent, Myia 2017 2018

Graphs

• Inspired from computer algebra systems (CAS)
• Program is a directed acyclic graph (DAG)
• A graph transformation automatically creates graph for derivative
• Easy to optimize, high performance, but awkward and not very expressive

f(g(x), h(x))

What about control flow?

Timeline

Modules SN, Torch, Caffe 1987 2002 2013 Operator overloading Autograd, Chainer, PyTorch 2014 2015 2017 Graphs Theano, TensorFlow, MXNet 2008 2015 2015 Source code transform Tangent, Myia 2017 2018

Operator overloading

def f(x): i = 0 while i < 3: i = i + 1 x = tanh(x) x = x * 10 return x i = 0 i = i + 1 x = tanh(x) i = i + 1 x = tanh(x) i = i + 1 x = tanh(x) x = x * 10

Trace Backprop

• Overload every operation to log itself on a tape.
• At the end, we walk the tape backward.
• “Define-by-run”, “Dynamic graph”
• Easy to implement, but lots of overhead
• Discourages composing small & cheap operations

Program Tape

Timeline

Modules SN, Torch, Caffe 1987 2002 2013 Operator overloading Autograd, Chainer, PyTorch 2014 2015 2017 Graphs Theano, TensorFlow, MXNet 2008 2015 2015 Source code transform Tangent, Myia 2017 2018

Source code transformation

• Transform a function that computes a value into a new function that

computes the derivative.

• Operate on source code or intermediate representation
• Standard language optimizations apply.
• Two approaches:
• Tape-based
• Functional AD

SCT: Tape-based

• Static version of operator overloading
• Forward pass computes value and pushes data on tape
• Backward pass pops data from tape and computes gradient
• Push and pop instructions are inserted in the code
• Main approach used by SCT AD frameworks
• Tapenade (C/Fortran), Tangent (Python)
• Issues
• The tape is global state
• Hard to support higher-order derivatives
• Has to deal with push&pop

SCT: Functional AD

• Works in pure functional languages
• Loops = tail recursion (no stack growth)
• Transform a function by making it return (output, backpropagator)
• Bprop takes gradient wrt output, returns gradients wrt inputs
• Bprop has access to original variables as free variables
• Closures are chained together in a sort of “tape”
• Higher-order derivatives
• Stalingrad (Scheme), DiffSharp (F#)

See: Reverse-Mode AD in a Functional Framework: Lambda the Ultimate Backpropagator (Pearlmutter & Siskind, 2008)

Paradigm

How deep learning and language design intersect

Autodiff

What it is. How it works

Myia

Our proposed solution

Reminder: Needs

General purpose: Express complex compositions with control flow. Fast: Leverage parallelism and GPU to process millions of features. Portable: Serializable, support multiple hardware. Goal: a language adapted to the needs of machine learning, past and future Differentiable: Language support for gradient descent.

Comparison

General Fast Portable Differentiable TensorFlow (graph) No ✔ ✔ Partially PyTorch (overloading) ✔ Partially Partially (Tracing) ✔ Tangent (SCT) ✔ Partially Python-specific ✔ Myia (SCT) ✔ ✔ ✔ ✔

Pipeline

Representation

• Functional
• Function abstraction,

recursion, closures

• No side-effects
• Graph representation
• Represent data flow directly
• Flexible scheduling (unlike SSA or CPS)
• Independent operations can run in parallel
• See: Sea of Nodes (Click 1995)
• Direct pointers to free variables
• See: Thorin (Leißa, 2015)

def fact(x): if x <> 1: return 1 else: return x * fact(x - 1)

Output Operation Input Constant

Parsing Python

def pow(x, n): r = 1 while n > 0: r = r * x n = n - 1 return r

• Parser based on SSA converter, adapted for functional representation
• Conditional branches become thunks (functions with no arguments)
• Loops converted to tail recursions

Type System

• Inference: unification-based
• Python interface: infer from types of arguments on Python side
• Recompile when types change
• Fundamental types:
• Scalars: Int/UInt/Float<8/16/32/64>, Bool, Char
• Tuple<T1, T2, ../>
• Struct<field1=T1, field2=T2, ../>
• List<T>
• Linked list, efficient append
• NDArray<T, Shape<D1, D2, ../>?
• D1, D2, ... need not be known statically
• Function<Args<TIn1, TIn2, ../>, TOut> (Function types)

Primitives

• Scalar primitives
• Usual arithmetic: add, mul, div, ../
• Commonly used: log, exp, ../
• Numeric stability: log1p, expm1, ../
• Array primitives
• broadcast, gemm, ../
• “Higher-order” primitives
• Apply a scalar-typed function on arrays following a pattern
• map, reduce
• Implicitly parallel: can be run efficiently on GPU

Automatic differentiation

def f(x, y): a = x *+ 3 b = y *+ 4 c = a * b return c

IR grad(f) grad(grad(f))

Optimization

• AD transform creates a lot of “junk”
• Tuple packing and unpacking
• Unnecessary gradients
• Aggressive inlining of primitive adjoints and pattern-matching
• ptimizations (peephole) can go a long way
• Stability optimizations
• Stabilize simple patterns
• Can be applied prior to AD

def f(x, y): a = x *+ 3 b = y *+ 4 c = a * b return c

Backend(s)

• Use existing backends to generate GPU kernels
• NNVM, XLA, Tensor Compr., TensorRT, NGraph, …
• Offload many optimizations to them (e.g. loop fusion)
• Leverage man-centuries of efforts
• But: not fully general (no recursion)
• Custom VM for general control flow
• Higher-level
• Python-based VM (for debugging)
• LLVM backend (for production)

Conclusion

• Roadmap to beta
• Support gradient operator in type system
• Improve optimization of the gradient
• Finalize support for array primitives and broadcasting
• Finalize support for backend

Myia aims to be a language adapted to the needs of machine learning, past and future:

• General purpose through support for recursion
• Fast and portable through NNVM backend and general purpose VM
• Differentiable using functional source code transformation

Thanks!

https://github.com/mila-udem/myia

Follow our progress!

is hiring!

Want to work at the confluence of academia and industry? We seek:

*

Professors

*

Software engineers

*

Director of software

*

R&D & technology transfer

*

Linux sysadmins

https://tinyurl.com/mila-jobs

Plus: free French classes!