Alice's Adventures in a differentiable Wonderland
Presenter: Simone Scardapane
Differentiability: the key ingredient of AI?
01
“For, you see, so many out-of-the-way things had happened lately, that Alice had begun to think that very few things indeed were really impossible.”
—Chapter 1, Down the Rabbit-Hole
“This position paper proposes an architecture and training paradigms with which to construct autonomous intelligent agents.”
“A system architecture for autonomous intelligence. All modules in this model are assumed to be differentiable.”
What is an “artificial neural network”?
No biology, please.
« computing systems vaguely inspired by the biological neural networks that constitute animal brains »
— Wikipedia
Convolutional neural networks
Transformers blocks
Implicit layers
Neural computers
A deep network is a differentiable function ...
def my_network(x: tensor) -> tensor:
...
...
...
return y
Input type
Output type
Sequence of differentiable primitives
… that can be optimized from data.
Optimizer
Automatic Differentiation
Loss goes in
A “better” function comes out
Corollary: deep networks are composable
Automatic Differentiation
def f(x):
y = my_network(x)
y = another_network(y)
y = yet_another_network(y)
return y
The how: automatic differentiation
02
“Would you tell me, please, which way I ought to go from here?”
“That depends a good deal on where you want to get to,” said the Cat.
“I don’t much care where” said Alice.
“Then it doesn’t matter which way you go,” said the Cat.
—Chapter 6, Pig and Pepper
Preliminaries: Derivative(s)
For a scalar function, the rate of change for an infinitesimally small displacement.
Preliminaries: Gradient(s)
For a function with an n-dimensional vector in input, we can use partial derivatives:
i-th basis vector
The vector of all partial derivatives is called the gradient:
Preliminaries: Jacobian(s)
Consider now a function with an n-dimensional input and a m-dimensional output, its Jacobian is defined as:
m rows (one for each output)
n columns (one for each input)
The Jacobian wrt what?
Let us go back to a classical fully-connected layer:
The Jacobian wrt x is (m,n), but the Jacobian w.r.t. W is rank-3 (m,m,n). In general, full Jacobians in real layers can be quite cumbersome.
We will return to this point later on.
Chain rule of Jacobians
Like classical derivatives, Jacobians also have a chain rule:
The gradient of function composition is the multiplication of the corresponding Jacobians.
Neural network primitives
Neural networks are composed of simple differentiable primitives:
Input
(Trainable) parameters
For each primitive, we know how to compute the input Jacobian and the weight Jacobian.
Neural networks
Simple NNs are a sequence of primitive operations:
In the more general case, we can have a DAG (not a sequence), and also parameter sharing between layers.
Computational graph
The goal of autodiff
Note that our last output, almost always, is scalar (e.g., sum of the per-element losses).
What we need is a way to efficiently, simultaneously compute all weight Jacobians (up to numerical precision):
Symbolic: build a symbolic formula for the gradients from the original symbolic program.
Numeric: numerically evaluate the derivatives using the definition.
One worked-out example
Consider a very simple example:
For example, this could be a one-layer neural network, cross-entropy loss, and final sum.
Considering each instruction in isolation, we have 4 input Jacobians and 3 weight Jacobians:
Then, we can use the chain rule to “stitch” them together.
Jacobians of f1
Jacobians of f2
Jacobians of f3
Performing multiplications left-to-right: forward-mode autodiff
Performing multiplications right-to-left: reverse-mode autodiff
Repeated!
Forward-mode autodiff
Forward-model autodiff can be implemented easily: all operations can be performed in parallel to the original program (i.e., we can devise a new program returning the original outputs and the gradients).
However, all operations will scale linearly w.r.t. number of parameters, which is impractical for today’s neural networks. On the good side, it requires little memory because previous operations can be discarded.
Forward-mode autodiff
Reverse-mode autodiff
Reverse-mode autodiff collects all intermediate operations of the program (tracing), and then “unrolls” all the gradient operations from right-to-left.
It is significantly more efficient (only vector-matrix operations above), but it requires to store all intermediate outputs, making it highly memory consuming.
Autodiff in ML is almost always in reverse-mode (backpropagation).
Reverse-mode autodiff
Quick-note: autodiff or backprop?
A few commonly accepted milestones:
Vector-Jacobian products
Importantly, we do not need to know how to compute the Jacobians of the primitives, but only their vector-Jacobian products (VJP):
(To see this, transpose all the equations before!)
Forward-mode autodiff can equivalently be written with Jacobian-vector products (JVPs), by computing the final gradients one value at a time.
VJPs vs. Jacobians
The previous result is important, because sometimes VJPs are easier than the full Jacobians:
The input Jacobian is a rank-3 tensor, while:
The nitty-gritty details
03
“The best way to explain it is to do it.”
—Chapter 3, A Caucus-Race and a Long Tale
Deep learning frameworks
Tensor primitives + their VJPs (matrix-multiplication, etc.).
Layer 0
Autodiff module
One primitive may be implemented with multiple kernels depending on the supported hardware (CPU, GPU, TPU, IPU, …).
High-level constructs (layers, optimizers, losses, metrics, …)
Layer 1
Layer 2
Ecosystem
(hubs, libraries, extensions, …)
Dispatchers
Some history and terminology
The revival of AD in neural networks started with Theano (2008), to which followed a Cambrian explosion of frameworks (TensorFlow 1.0, PyTorch, Caffe, JAX, …).
Theano and TF 1.0 focused heavily on performance. The user implemented the computational graph with a small domain specific language (DSL), and execution was decoupled from the definition (define-then-run).
Most frameworks today are more dynamic (define-by-run), leaving compilation as a separate, optional step (JIT in JAX / PyTorch, tf.function in TF).
Flavours of implementations
While most frameworks implement similar things, the way they implement them can make some use cases considerably faster or easier.
Anatomy of a PyTorch tensor
dtype = …
requires_grad = True
data
grad
grad_fn
Identifies explicitly tensors that require gradients
The actual storage of the tensor
Gradient is accumulated here
A pointer to the parent operation to be able to traverse the DAG
Note: PyTorch also has a functional variant, very similar to the JAX implementation.
High-level APIs
Most frameworks (TensorFlow, PyTorch) implement an object-oriented API:
class MyModule(nn.Module):
def __init__(self):
self.params = Parameter( torch.tensor ( … ) )
def forward(self, x):
return self.params @ x
Parameters are properties of the object
Simple polymorphism allows for module compositionality
Forward logic is a method
Materials to learn more
PyTorch internals : ezyang’s blog (slightly outdated)
Automatic differentiation (Matthieu Blondel, 2020)
GitHub - mattjj/autodidact: A pedagogical implementation of Autograd
GitHub - geohot/tinygrad: You like pytorch? You like micrograd? You love tinygrad! ❤️
GitHub - MINI-PYTORCH/MINI-TORCH: Mini-pytorch implemented from scratch using Python
The spelled-out intro to neural networks and backpropagation: building micrograd
Our discussion will be highly simplified, leaving out many important topics (e.g., tracing and distributed implementations).
Notebook time!
Notebook 1: simple experiments with PyTorch and JAX, side-by-side.
Notebook 2: building a toy autodiff framework, PyTorch-style.
https://colab.research.google.com/drive/1AbNRRjL0DMoj4VPnul7QIYJC5nLU3Fn2?usp=sharing
Limits of OOP
By default, JAX takes a fully functional paradigm: everything (layers, losses) is a function.
�Sometimes, this makes it easier to explicitly manipulate parameters and to compose different transformations.
Most high-level frameworks in JAX define a layer by a pair of init/apply functions (with some exceptions, see Equinox).
Notebook time!
Notebook 3: moving from PyTorch to JAX and .vice versa
https://colab.research.google.com/drive/1AbNRRjL0DMoj4VPnul7QIYJC5nLU3Fn2?usp=sharing
Advanced topics
04
““Curiouser and curiouser!” cried Alice (she was so much surprised, that for the moment she quite forgot how to speak good English).
—Chapter 2, The Pool of Tears
Layer sharing
1 layer
2 layers
2 layers (replicated)
Deep equilibrium layers
What happens if we replicate the layer infinite times?
The output of the layer is now implicitly defined by a fixed-point equation:
…
Input
Output (can be initialized to zero)
Solving fixed-point equations
Writing a fixed-point layer is easy (of course, there are faster alternatives):
class FixedPointLayer(nn.Module):
def __init__(self):
self.w = ...
def forward(self, x):
z = torch.zeros_like(x)
while self.check_convergence():
z = f(self.w, z, x)
return z
By default, however, AD requires to store all intermediate steps of the while loop and backpropagate through them, which is expensive.
Implicit function theorem
By the implicit function theorem, there exists a continuous function z* such that:
Differentiating everything:
We can differentiate the layer by computing two gradients at the optimum z (no need to store any intermediate layers).
VJPs of implicit functions
In order to implement the layer as a primitive in a framework, we need its VJP with some vector u. Fascinatingly, this can be expressed as another fixed-point equation:
Read more
Implicit differentiation is the key for several topics:
Check out this tutorial for more: Deep Implicit Layers
Gradient checkpointing
To save memory, gradient checkpointing is now popular. Outputs of red nodes are stored (checkpoints). When back-propagating through a non-checkpointed node, its output is recomputed starting from the previous checkpoint in memory.
Bilevel optimization
Outer problem
Inner problem
Examples:
Physics-informed NNs
Code 2.0
?
Thanks for listening!
“Tut, tut, child!” said the Duchess. “Everything’s got a moral, if only you can find it.”
—Chapter 9, The Mock Turtle’s Story