A Simple and Efficient Tensor Calculus
Sören Laue,1,2 Matthias Mitterreiter,1 Joachim Giesen 1
1

Friedrich-Schiller-Universität Jena
Faculty of Mathematics and Computer Science
Ernst-Abbe-Platz 207743 Jena, Germany
Friedrich-Schiller-University Jena
2
Data Assessment Solutions GmbH, Hannover

Abstract
Computing derivatives of tensor expressions, also known as
tensor calculus, is a fundamental task in machine learning. A
key concern is the efficiency of evaluating the expressions andtheir derivatives that hinges on the representation of these expressions. Recently, an algorithm for computing higher orderderivatives of tensor expressions like Jacobians or Hessians
has been introduced that is a few orders of magnitude fasterthan previous state-of-the-art approaches. Unfortunately, theapproach is based on Ricci notation and hence cannot be
incorporated into automatic differentiation frameworks like
TensorFlow, PyTorch, autograd, or JAX that use the simpler
Einstein notation. This leaves two options, to either change
the underlying tensor representation in these frameworks or todevelop a new, provably correct algorithm based on Einsteinnotation. Obviously, the first option is impractical. Hence, we
pursue the second option. Here, we show that using Ricci
notation is not necessary for an efficient tensor calculus anddevelop an equally efficient method for the simpler Einsteinnotation. It turns out that turning to Einstein notation enables
further improvements that lead to even better efficiency.

1

Introduction

Many problems in machine learning are naturally written
in terms of tensor expressions. Any algorithmic method forcomputing derivatives of such expressions is called a tensor calculus. Standard automatic differentiation (deep learning) frameworks like TensorFlow (Abadi and others 2016),
PyTorch (Paszke et al. 2017), autograd (Maclaurin, Duvenaud, and Adams 2015), and JAX (Frostig, Johnson, andLeary 2018) are very efficient when computing derivatives ofscalar-valued functions. However, evaluating the derivativesof non-scalar-valued functions, for instance, Jacobians or
Hessians, in these frameworks is up to three orders of magnitude slower than evaluating the derivatives that are computedby the approach of (Laue, Mitterreiter, and Giesen 2018).
There have been some recent attempts to alleviate this lackof efficiency by accelerating the underlying linear algebrausing automatic batching and optimizing the computationalgraphs of the derivatives (XLA and TensorFlow teams 2017).
These improvements have been incorporated into TensorFlow
Copyright c 2020, Association for the Advancement of Artificial
Intelligence (www.aaai.org). All rights reserved.

and JAX. However, the improvements are rather small andthe efficiency gap of up to three orders of magnitude still
persists.
On the other hand, the approach of (Laue, Mitterreiter, andGiesen 2018) relies crucially on Ricci notation and therefore
cannot be incorporated into standard deep learning frameworks that use the simpler Einstein notation. Here, we remove
this obstacle and provide an efficient tensor calculus in Einstein notation. Already the simple version of our approach is
as efficient as the approach by (Laue, Mitterreiter, and Giesen
2018). We provide further improvements that lead to an evenbetter efficiency.
Ricci notation distinguishes between co- and contravariant
indices, that is, upper and lower indices. This distinction isnecessary in order to compute derivatives in a mathematical correct way. Consider for instance the simple expressionx> Ax. If we want to compute the derivative of this expression with respect to the vector x, then, at some point, we face
the problem of computing the derivative of x> . However, thisderivative in Ricci notation is the delta-tensor δij that cannot
be represented in linear algebra. Note, it is not the identity
matrix which is represented in Ricci notation as δji . Hence,in order to represent the derivative in a mathematical correctway, upper and lower indices are necessary. This problem has
its roots in mathematical tensor analysis, where tensors areused for representing multilinear functions by their values
on a set of basis vectors. These values are stored in a tensor,
that is, a multi-dimensional array. Upper and lower indices
are used to distinguish between vector space and dual vector space components that transform differently under basischanges. In the example expression, x is a vector while x>
is a co-vector from the dual vector space.
In machine learning tensors are typically not used for
representing multi-linear functions, but simply as multidimensional arrays for storing data and parameters. Hence,
there is no need to distinguish different types of components.
Indices can just be used for accessing the different tensor
components. This is basically Einstein notation that is usedin all deep learning frameworks.
The contribution of this paper is an efficient and coherent
method for computing tensor derivatives in Einstein notation together with a correctness proof. In reverse mode auto-matic differentiation, our method is equivalent to the efficientapproach in (Laue, Mitterreiter, and Giesen 2018) for computing higher order derivatives. Additionally, we show thatreverse mode is not optimal. A combination of reverse andforward mode, known as cross-country mode, can be more
efficient. Efficiency can be further improved by compressing
higher order derivatives.
For validating our framework we compute Hessians for several machine learning problems. It turns out that our method,
because of the additional optimizations, outperforms the approach of (Laue, Mitterreiter, and Giesen 2018) which isalready a few orders of magnitude more efficient than TensorFlow, PyTorch, autograd, and JAX.
Related Work. Many details on the fundamentals andmore advanced topics of automatic differentiation can befound in the book by (Griewank and Walther 2008). (Baydin et al. 2018) provide an excellent survey on automaticdifferentiation for machine learning.
Computing derivatives of non-scalar-valued functions is
discussed in (Pearlmutter 1994). In this approach, if the function returns an n-dimensional vector, then its derivative iscomputed by treating each entry as a separate scalar-valued
function. The same idea is employed in almost all implementations for computing derivatives of non-scalar-valued
functions. (Gebremedhin et al. 2009) introduce some optimizations based on graph coloring algorithms.
(Magnus and Neudecker 2007) can compute derivatives
with respect to vectors and matrices. At the core of their approach, matrices are turned into vectors by stacking columnsof a matrix into one long vector. Then, the Kronecker matrix product is used to emulate higher order tensors. Thisapproach works well for computing first order derivativesof scalar-valued functions. However, it is not practicable forcomputing higher order derivatives.
(Giles 2008) collects a number of derivatives for matrix
operators, i.e., pushforward and pullback functions for automatic differentiation. Similarly, (Seeger et al. 2017) provide
methods and code for computing derivatives of Cholesky factorizations, QR decompositions, and symmetric eigenvalue
decompositions. However, they all require that the output
function is scalar-valued, and hence, cannot be generalized
to higher order derivatives.
(Kakade and Lee 2018) consider non-smooth functions andprovide a provably correct algorithm that returns an elementfrom the subdifferential. However, their algorithm is also
restricted to scalar-valued functions.
Another line of research focuses on automatic differentiation from a programming language point of view. The
goal is to incorporate gradient computations into programming languages with the goal of fully general differentiable
programming (LeCun ). Work towards this goal includes the
Tangent package (van Merriënboer, Moldovan, and Wiltschko
2018), the Myia project (van Merriënboer et al. 2018), andthe approach of (Wang et al. 2018). So far this work is again
restricted to scalar-valued functions.
Recently, also second order information has been considered for training deep nets. For instance, this informa-tion can be exploited for escaping one of the many saddle points, but also for turning the final classifier morerobust (Dey et al. 2018). Furthermore, some differencesin the convergence behavior can be explained by looking at the spectrum of the Hessian of the objective function (Ghorbani, Krishnan, and Xiao 2019; Sagun et al. 2018;
Yao et al. 2018). However, so far it has been prohibitive tocompute the full Hessian even for small networks. Here, in
our experiments, we also compute the Hessian of a small
neural net.

2

Einstein Notation

In tensor calculus one distinguishes three types of multiplication, namely inner, outer, and element-wise multiplication.
Indices are important for distinguishing between these types.
For tensors A, B, and C any multiplication of A and B canbe written as
X
C[s3 ] =A[s1 ] · B[s2 ],(s1 ∪s2 )\s3

where C is the result tensor and s1 , s2 , and s3 are the index
sets of the left argument, the right argument, and the result
tensor, respectively. The summation is only relevant for innerproducts that in Ricci calculus are denoted by shared upperand lower indices. If one does not want to distinguish between
upper and lower indices, then the summation must be madeexplicit through the result tensor. The standard way to doso is by excluding the index for summation from the index
set of the result tensor. Hence, the index set of the result
tensor is always a subset of the union of the index sets ofthe multiplication’s arguments, that is, s3 ⊆ (s1 ∪ s2 ). In
the following we denote the generic tensor multiplication
simply as C = A ∗(s1 ,s2 ,s3 ) B, where s3 explicitly represents
the index set of the result tensor. This notation is basically
identical to the tensor multiplication einsum in NumPy,
TensorFlow, and PyTorch, and to the notation used in theTensor Comprehension Package (Vasilache et al. 2018).
The ∗(s1 ,s2 ,s3 ) -notation comes close to standard Einsteinnotation. In Einstein notation the index set s3 of the output isomitted and the convention is to sum over all shared indices in
s1 and s2 . This, however, restricts the types of multiplications
that can be represented. The set of multiplications that can be
represented in standard Einstein notation is a proper subsetof the multiplications that can be represented by our notation.
For instance, standard Einstein notation is not capable ofrepresenting element-wise multiplications directly. Still, in
the following we refer to the ∗(s1 ,s2 ,s3 ) -notation simply as
Einstein notation as it is standard practice in all deep learning
frameworks.
Table 1 shows examples of tensor expressions in standard
linear algebra notation, Ricci calculus, and Einstein notation.
The first group shows an outer product, the second groupshows inner products, and the last group shows examplesof element-wise multiplications. As can be seen in Table 1,
Ricci notation and Einstein notation are syntactically reasonably similar. However, semantically they are quite different.
As pointed out above, Ricci notation differentiates between
co- and contravariant dimensions/indices and Einstein notation does not. While this might seem like a minor difference,it does have substantial implications when computing derivatives. For instance, when using Ricci notation, forward andreverse mode automatic differentiation can be treated in thesame way (Laue, Mitterreiter, and Giesen 2018). This is nolonger the case when using Einstein notation.
vectorized
>

yx
Ax
y> x
AB
y x
A B
A · diag(x)Ricci

Einsteiniy ∗(i,j,ij) x
A ∗(ij,j,i) xy ∗(i,i,∅) x
A ∗(ij,jk,ik) By ∗(i,i,i) x
A ∗(ij,ij,ij) B
A ∗(ij,i,ij) x

y xj
Aij xjy i xiAij Bkjy i xiAij BjiAij xiFigure 1: Comparison of different linear algebra notations.

Here, the dot product notation D ◦ h is short for the innerproduct D ∗(s1 s2 ,s2 ,s1 ) h, where s1 s2 is the index set of D
and s2 is the index set of h. For instance, if D ∈ Rm1 ×n1 ×n2
and h ∈ Rn1 ×n2 , then s1 = {i, j, k} and s2 = {j, k}.
In the following, we first describe forward and reverse
mode automatic differentiation for expressions in Einsteinnotation, before we discuss extensions like cross-country
mode and compression of higher order derivatives that aremuch easier to realize in Einstein than in Ricci notation. As
can be seen from our experiments in Section 4, the extensions
allow for significant performance gains.

Forward Mode
Any tensor expression has an associated directed acyclic
expression graph (expression DAG). Figure 2 shows the expression DAG for the expression
X > (exp(X · w) + 1)−1exp(X · w)),

(1)
−1

We can show that the generic tensor multiplication operator ∗(s1 ,s2 ,s3 ) is associative, commutative, and satisfies the
distributive property. Our tensor calculus, that we introduce
in the next section, makes use of all three properties. By
s1 s2 we denote the concatenation of the index sets s1 and s2 .
Please see the supplemental material for the easy proofs that
follow directly from our definition of Einstein notation.

where denotes the element-wise multiplication andtheelement-wise multiplicative inverse. The nodes of the DAG∗(ij,i,j)∗(i,i,i)Lemma 1 (Associativity). Let s1 , s2 , s3 , and s4 be index setswith s3 ⊆ s1 ∪ s2 and s4 ∩ (s1 ∪ s2 ) = ∅. Then it holds thatA ∗(s1 ,s2 s4 ,s3 s4 ) B ∗(s3 s4 ,s4 ,s3 ) C =A ∗(s1 ,s2 ,s3 ) B ∗(s2 s4 ,s4 ,s2 ) C .
Unlike standard matrix multiplication tensor multiplicationis commutative.
Lemma 2 (Commutativity). It holds that

y0

∧−1

v4

v3

+ v2

1

x2

∗(ij,j,i)

exp v1v0

A ∗(s1 ,s2 ,s3 ) B = B ∗(s2 ,s1 ,s3 ) A.
Lemma 3 (Distributive property). Let s1 , s2 , and s3 be index
sets with s3 ⊆ s1 ∪ s2 . It holds thatA ∗(s1 ,s2 ,s3 ) B + A ∗(s1 ,s2 ,s3 ) C = A ∗(s1 ,s2 ,s3 ) (B + C) .

3

Tensor Calculus

Now we are prepared to develop our tensor calculus. We startby giving the definition of the derivative of a tensor-valued
expression
with respect to a tensor. For the definition, we use
pP
2
kAk =
s A[s] as the norm of a tensor A.
Definition 4 (Fréchet Derivative). Let f : Rn1 ×n2 ×...×nk →Rm1 ×m2 ×...×ml be a function that takes an order-k tensor
as input and maps it to an order-l tensor as output. Then,
D ∈ Rm1 ×m2 ×...×ml ×n1 ×n2 ×...×nk is called the derivativeof f at x if and only if
lim

h→0

kf (x + h) − f (x) − D ◦ hk
= 0,
khk

where ◦ is an inner tensor product.

X x0

w x1

Figure 2: Expression DAG for Expression (1)that have no incoming edges represent the variables of theexpression and are referred to as input nodes. The nodes of
the DAG that have no outgoing edges represent the functions
that the DAG computes and are referred to as output nodes.
Let the DAG have n input nodes (variables) and m output
nodes (functions). We label the input nodes as x0 , ..., xn−1 ,the output nodes as y0 , ..., ym−1 , and the internal nodes as
v0 , . . . , vk−1 . Every internal and every output node represents either a unary or a binary operator. The arguments ofthese operators are supplied by the incoming edges.
In forward mode, for computing derivatives with respect
to the input variable xj , each node vi will eventually store∂vi
the derivative ∂xwhich is traditionally denoted as v̇i . It isjcomputed from input to out- put nodes as follows: At the
input nodes that represent the variables xi , the derivatives∂xi∂xj are stored. Then, the derivatives that are stored at theremaining nodes, here called f , are iteratively computed bysumming over all their incoming edges as
X ∂f ∂z
X ∂f
∂f
f˙ =
=
·
=
· ż,
∂xj
∂z ∂xj∂zz : (z,f )∈E

z : (z,f )∈E

where ∂f∂z is the partial derivative of node f with respect toz and the multiplication is tensorial. The so called pushforwards ż of the predecessor nodes z of f have been computed
before and are stored at z. Hence, the derivative of eachfunction is stored at the corresponding output node y of theexpression DAG. Obviously, the updates can be done simultaneously for one input variable xj and all output nodes yi .
Computing the derivatives with respect to all input variables
requires n such rounds.
In the following we derive the explicit form of the pushforward for nodes of the expression DAG of a tensor expression.
For such a DAG we can distinguish four types of nodes,
namely multiplication nodes, general unary function nodes,
elementwise unary function nodes, and addition nodes. General unary functions are general tensor-valued functions while
elementwise unary functions are applied to each entry of asingle tensor. The difference can be best explained by thedifference between the matrix exponential function (generalunary function) and the ordinary exponential function applied
to every entry of the matrix (elementwise unary function).
The pushforward for addition nodes is trivially just the sum of
the pushforward of the two summands. Thus, it only remains
to show how to compute the pushforward for multiplication,
general unary functions, and element-wise unary function
nodes. Please refer to the supplemental material for the proofsof the following three theorems, or see the similar proof of
Theorem 8 for the reverse mode in the next section.
Theorem 5. Let x be an input variable with index set s4
and let C = A ∗(s1 ,s2 ,s3 ) B be a multiplication node of theexpression DAG. The pushforward of C isĊ = B ∗(s2 ,s1 s4 ,s3 s4 ) Ȧ + A ∗(s1 ,s2 s4 ,s3 s4 ) Ḃ.
Theorem 6. Let x be an input variable with index set s3 ,let f be a general unary function whose domain has index
set s1 and whose range has index set s2 , let A be a node in
the expression DAG, and let C = f (A). The pushforward of
the node C is Ċ = f 0 (A) ∗(s2 s1 ,s1 s3 ,s2 s3 ) Ȧ, where f 0 is thederivative of f .
In case that the general unary function is simply an elementwise unary function that is applied element-wise to atensor, Theorem 6 simplifies as follows.
Theorem 7. Let x be an input variable with index set s2 , let
f be an elementwise unary function, let A be a node in theexpression DAG with index set s1 , and let C = f (A) where
f is applied element-wise. The pushforward of the node C isĊ = f 0 (A) ∗(s1 ,s1 s2 ,s1 s2 ) Ȧ, where f 0 is the derivative of f .

Reverse Mode
Reverse mode automatic differentiation proceeds similarly tothe forward mode, but from output to input nodes. Each node∂yvi will eventually store the derivative ∂vji which is usuallydenoted as v̄i , where yj is the function to be differentiated.
These derivatives are computed as follows: First, the deriva∂ytives ∂yji are stored at the output nodes of the DAG. Then,the derivatives that are stored at the remaining nodes, here
called z, are iteratively computed by summing over all theiroutgoing edges as follows
XX∂yj∂yj ∂f
∂f
z̄ =
=
·
=
f¯ ·
,
∂z
∂f ∂z
∂z
f : (z,f )∈E

f : (z,f )∈E

where the multiplication is again tensorial. The so-calledpullbacks f¯ have been computed before and are stored at thesuccessor nodes f of z. This means the derivatives of thefunction yj with respect to all the variables xi are stored at
the corresponding input nodes of the expression DAG. Computing the derivatives for all the output functions requires msuch rounds.
In the following we describe the contribution of unary
and binary operater nodes to the pullback of their arguments.
Here we have only two types of binary operators, namely
tensor addition and tensor multiplication. In the addition casethe contribution of C to the pullback of both of its arguments
is simply C̄. In Theorem 8 we derive the explicit form of thecontribution of a multiplication node to the pullback of its
arguments, in Theorem 9 the contribution of a general unaryfunction, and in Theorem 10 we derive the contribution of anelementwise unary function node to its argument.
Theorem 8. Let Y be an output node with index set s4 and let
C = A∗(s1 ,s2 ,s3 ) B be a multiplication node of the expression
DAG. Then the contribution of C to the pullback B̄ of B isC̄ ∗(s4 s3 ,s1 ,s4 s2 ) A and its contribution to the pullback Ā of
A is C̄ ∗(s4 s3 ,s2 ,s4 s1 ) B.
Proof. Here we only derive the contribution of C to the pullback B̄. Its contribution to Ā can be computed analogously.
∂C
The contribution of C to B̄ is C̄ · ∂B
. By Definition 4 we∂Yhave for the derivative C̄ = ∂C of Y with respect to C thatlim

1h̃→0 kh̃k· Y (C + h̃) − Y (C) − C̄ ◦ h̃ = 0.

By specializing h̃ = A ∗(s1 ,s2 ,s3 ) h we get
Y (C + h̃) − Y (C) − C̄ ◦ h̃
= Y (A ∗(s1 ,s2 ,s3 ) B + A ∗(s1 ,s2 ,s3 ) h)
− Y (A ∗(s1 ,s2 ,s3 ) B) − C̄ ◦ (A ∗(s1 ,s2 ,s3 ) h)
= Y (A ∗(s1 ,s2 ,s3 ) (B + h)) − Y (A ∗(s1 ,s2 ,s3 ) B)
− C̄ ◦ (A ∗(s1 ,s2 ,s3 ) h)
= Y (A ∗(s1 ,s2 ,s3 ) (B + h)) − Y (A ∗(s1 ,s2 ,s3 ) B)
− C̄ ∗(s4 s3 ,s3 ,s4 ) (A ∗(s1 ,s2 ,s3 ) h)
= Y (A ∗(s1 ,s2 ,s3 ) (B + h)) − Y (A ∗(s1 ,s2 ,s3 ) B)
− (C̄ ∗(s4 s3 ,s1 ,s4 s2 ) A) ∗(s4 s2 ,s2 ,s4 ) h)
= Y (A ∗(s1 ,s2 ,s3 ) (B + h)) − Y (A ∗(s1 ,s2 ,s3 ) B)
− (C̄ ∗(s4 s3 ,s1 ,s4 s2 ) A) ◦ h),where the first equality follows from the definitions of C andh̃, the second from Lemma 3, the third from the definitionof ◦, the fourth from Lemma 1, the fifth from Lemma 2, andthe last again from the definition of ◦. Hence, we have for∂C∂Y∂C · ∂B that1
· Y (C + h̃) − Y (C) − C̄ ◦ h̃h̃→0 h̃1= lim· Y (A ∗(s1 ,s2 ,s3 ) (B + h))
h→0 khk

0 = lim− Y (A ∗(s1 ,s2 ,s3 ) B) − (C̄ ∗(s4 s3 ,s1 ,s4 s2 ) A) ◦ h)Thus, the contribution of C to the pullback B̄ is
∂Y ∂C
∂C
·
= C̄ ·= C̄ ∗(s4 s3 ,s1 ,s4 s2 ) A.
∂C ∂B
∂B
If the output function Y in Theorem 8 is scalar-valued,
then we have s4 = ∅ and the pullback function coincides
with the function implemented in all modern deep learning
frameworks including TensorFlow and PyTorch. Hence, ourapproach can be seen as a direct generalization of the scalar
case.
Theorem 9. Let Y be an output function with index set s3 ,let f be a general unary function whose domain has index
set s1 and whose range has index set s2 , let A be a node in
the expression DAG, and let C = f (A). The contribution of
the node C to the pullback Ā isf¯ ∗(s3 s2 ,s2 s1 ,s3 s1 ) f 0 (A),where f 0 is the derivative of f .
In case that the general unary function is simply an elementwise unary function that is applied element-wise to atensor, Theorem 9 simplifies as follows.
Theorem 10. Let Y be an output function with index set s2 ,let f be an elementwise unary function, let A be a node in theexpression DAG with index set s1 , and let C = f (A) where
f where f is applied element-wise. The contribution of thenode C to the pullback Ā isf¯ ∗(s2 s1 ,s1 ,s2 s1 ) f 0 (A),where f 0 is the derivative of f .
The proofs of Theorem 9 and 10 are similar to the proofsof Theorems 6 and 8. They can be found in the supplemental
material.

Beyond Forward and Reverse Mode
Since the derivative of a function y with respect to an input
variable x is the sum over all partial derivatives along allpaths from x to y (Griewank and Walther 2008), we can
combine forward and reverse mode. Using that v̄ = ∂y∂v and∂vv̇ = ∂x, we getX∂y=v̄ ∗(s1 sv ,sv s2 ,s1 s2 ) v̇,∂xv∈S

where sv is the index set of node v, s1 is the index set of theoutput function y, s2 is the index set of the input node x, andS is the set of nodes in a cut of the expression DAG. General combinations of forward and reverse mode lead to theso-called cross-country mode. We will show that the differentiation of tensor expressions becomes even more efficient
by a special instantiation of the cross-country mode and bycompressing higher order derivatives.
Cross-Country Mode. In both forward and reverse mode,
derivatives are computed as sums of products of partial derivatives. In general, the time for evaluating the derivatives depends on the order by which the partial derivatives are multiplied. The two modes multiply the partial derivatives in
opposite order. Derivatives are multiplied from input to output nodes in forward mode and the other way around inreverse mode.
If the output function is scalar-valued, then reverse mode isefficient for computing the derivative with respect to all input
variables. It is guaranteed that evaluating the derivative takes
at most six times the time for evaluating the function itself. In
practice, usually a factor of two is observed (Griewank andWalther 2008). However, this is no longer true for non-scalarvalued functions. In the latter case, the order of multiplying
the partial derivatives has a strong impact on the evaluation
time, even for simple functions (Naumann 2004). Reordering
the multiplication order of the partial derivatives is known
as cross-country mode in the automatic differentiation literature (Bischof, Hovland, and Norris 2002). Finding an optimalordering is NP-hard (Naumann 2008) in general.
However, it turns out that significant performance gains
for derivatives of tensor expressions can be obtained by there-ordering strategy that multiplies tensors in order of theirtensor-order, that is, multiplying vectors first, then matrices,and so on. We illustrate this strategy on the following example
f (x) = B · g(h(Ax)),(2)

where A and B are two matrices, x is a vector and g(.) andh(.) are vector-valued functions that also take a vector as input. The derivative in this case is B diag(u) diag(v)A, whereu = g 0 (h(Ax)), v = h0 (Ax), and diag(u) is the diagonal
matrix with u on its diagonal. Reverse mode multiplies thesematrices from left to right while forward mode multiplies
them from right to left. However, it is more efficient to firstmultiply the two vectors u and v element-wise and then to
multiply the result with the matrices A and B.
Actually, the structure of Example 2 is not contrived, but
fairly common in second
P order derivatives. For instance, consider the expression g(h(Ax)), where g and h are as above
and the sum is over the vector components of the vectorvalued expression g(h(Ax)). Many machine learning problems feature such an expression as subexpression, where A
is a data matrix and the optimization variable x is a parameter vector. The gradient of this expression has the form of
Example 2 with B = A> . As can be seen in the experiments
in Section 4, reordering the multiplications by our strategy
reduces the time for evaluating the Hessian by about 30%.

Compressing Derivatives. Our compression schemebuilds on the re-ordering scheme (cross-country mode) fromabove and on the simple observation that in forward as wellas in reverse mode the first partial derivative is always a unit
tensor. It is either, in reverse mode, the derivative of the output nodes with respect to themselves or, in forward mode,the derivative of the input nodes with respect to themselves.
This unit tensor can always be moved to the end of the multiplications, if the order of multiplication is chosen exactly
as in our cross-country mode strategy that orders the tensors
in increasing tensor-order. Then, the multiplication with theunit tensor at the end is either trivial, i.e., amounts to a multiplication with a unit matrix that has no effect and thus can beremoved, or leads to a compactification of the derivative.
For an example, consider the loss functionf (U ) = T − U V >2of the non-regularized matrix factorization problem. Here,T ∈ Rn×n , U, V ∈ Rn×k and n is usually large while k issmall. The Hessian of f is the fourth order tensor
H = 2(V ∗(ij,ik,jk) V ) ∗(jl,ik,ijkl) I ∈ Rn×k×n×k ,where I is the identity matrix. Newton-type algorithms forthis problem solve the Newton system which takes timein O (nk)3 . However, the Hessian can be compressed to2(V ∗(ij,ik,jk) V ) which is a small matrix of size k × k.
This matrix can be inverted in O(k 3 ) time. The performancegain realized by compression can be significant. For instance,
solving the compressed Newton system needs only about
10 µsec whereas solving the original system needs about
1 sec for a problem of size n = 1000 and k = 10. For moreexperimental results refer to Section 4.
As another example, consider a simple neural net witha fixed number of fully connected layers, ReLU activation
functions, and a softmax cross-entropy output layer. The Hessian of each layer is a fourth order tensor that can be written
as A ∗(ijl,ik,ijkl) I for a suitable third order tensor A. In thiscase, the Hessian can be compressed from a fourth ordertensor to a third order tensor. We provide expression trees
for both derivatives, compressed and uncompressed, in the
supplemental material. Computing with the compact representation of the Hessian is of course more efficient which we
confirm experimentally in the next section.

4

Experiments

We have implemented both modes of the tensor calculus fromthe previous section together with the improvements that canbe achieved by cross-country mode and the compactificationof higher order derivatives. State-of-the-art frameworks like
TensorFlow and PyTorch only support reverse mode since
it allows to compute derivatives with respect to all input
variables at the same time. Similarly to all other frameworks,
our implementation performs some expression simplification
like constant folding and removal of zero and identity tensors.
An anonymized version of our implementation can be foundin the supplemental material.

Experimental Setup. We followed the experimental setupof (Laue, Mitterreiter, and Giesen 2018). However, we added
one more experiment, a small neural net. We have implemented our algorithms in Python. To evaluate expressions
we used NumPy 1.16 and CuPy 5.1. We compared ourframework with the state-of-the-art automatic differentiation
frameworks TensorFlow 1.14, PyTorch 1.0, autograd 1.2, and
JAX 0.1.27 used with Python 3.6 that were all linked against
the Intel MKL. All these frameworks support reverse mode
automatic differentiation for computing first order derivatives.
For scalar-valued functions the reverse mode of each of theseframeworks coincides with the reverse mode of our approach.
For non-scalar-valued functions all the frameworks computethe derivative for each entry of the output function separately.
The expression DAGs that are generated by our reverse modefor general tensor expressions coincide with the derivatives
computed by the approach of (Laue, Mitterreiter, and Giesen
2018). The experiments were run in a pure CPU setting (Intel
Xeon E5-2643, 8 cores) as well as in a pure GPU setting(NVIDIA Tesla V100), except for autograd that does notprovide GPU support.
We computed function values, gradients, and Hessians foreach set of experiments. We computed Hessians on the CPU
as well as on the GPU. To avoid the issue of sparsity patterns
we generated dense, random data for each experiment. In thissetting the running time does not depend on whether the data
are synthetic or real world.
Logistic Regression. Let X ∈ Rm×n be a data matrix and y ∈ {−1, +1}m the corresponding binary label
logistic regression
 function is fn(w) =P vector. The (i)(i), where w ∈ R is the
log
exp
−y
(X
w)
+
1
i
weight vector, X (i) is the i-th data point (i-th row of X), andy (i) its label. We set m = 2n in the experiments.
Matrix Factorization. Let T = Rm×m be some target
matrix, Ω ∈ {0, 1}m×m be an indicator matrix that defines
which elements of T are known, and let U, V ∈ Rm×k be
low-rank factor matrices. Matrix factorization is the problem
of solving minU,V kT − U V > k2Ω . For the experiments, we
set k = 5 and compute the gradient and Hessian with respect
to U .
Neural Net. We have created a small neural net with ten
fully connected layers, ReLU activation functions, and asoftmax cross-entropy output layer. The weight matrices forthe different layers all had the same size, n × n. Here, wereport running times for computing the Hessian of the firstlayer.
Evaluation. In the case of scalar-valued functions all theframeworks basically work in the same way. Thus, it is not
surprising that their running times for computing function
values and gradients are almost the same, see Figure 1 in the
supplemental material.
The situation is different for Hessians. First, it can be seenthat the reverse mode in our approach, whose results agree

Figure 3: Running times on a CPU (top row) and on a GPU (bottom row) for computing the Hessian of the logistic regressionfunction (left), matrix factorization (middle), and a small neural net (right).
with the results in (Laue, Mitterreiter, and Giesen 2018), is a
few orders of magnitude faster than current state-of-the-art
frameworks like TensorFlow, PyTorch, autograd, and JAX.
This holds true for all experiments on the CPU and on theGPU, see Figure 3. For the logistic regression problem ourcross-country mode is about 30% faster on the CPU, while its
effect on the GPU is negligible because of the GPU overhead.
The performance gain of our compression scheme can beseen on the matrix factorization problem and on the neural net(middle and right column of Figure 3). Computing Hessians
for small neural nets has now become feasible.
In our experiments, the effect of recent efforts in speeding
up deep learning frameworks turned out to be rather small.
Actually, enabling XLA for TensorFlow on the CPU slowed
the computation down by a factor of two. Hence, we omit
these running times in Figure 3. Enabling XLA on the GPU
provided only marginal improvements. JAX which relies on
XLA did not finish computations but raised memory errors
indicating that it went out of main memory. These errors
seem to be caused by the automatic batching function vmap
that is used by JAX for auto-vectorization when computing
Hessians. This was surprising to us since JAX is meant tobe more memory efficient than TensorFlow. There is one
exception, on the GPU, JAX finished the computation forthe logistic regression problem. However, as can be seen in
Figure 3, even in this case it is not significantly more efficientthan the other deep learning frameworks. This was to be
expected since JAX relies on XLA and the authors of XLA
report a speed-up of only 15% in the GPU setting (XLA2019).

5

Conclusion

We have developed a simple, efficient and provably correctframework for computing derivatives of general tensor expressions that is much simpler than previous approaches.
Furthermore, it can be easily integrated into state-of-theart frameworks like TensorFlow and PyTorch that use thesame tensor representation, but are a few orders of magnitude slower than our approach when computing higher
order derivatives. We have also demonstrated that reverse
mode automatic differentiation is not optimal for computing higher order derivatives. Significant speed ups can beachieved by a special instantiation of the cross-country modeand by compressing higher order derivatives. The algorithms
presented here also form the basis of an online tool for computing derivatives of matrix and tensor expressions which isaccessed by more than 30,000 users per year.

Acknowledgments
Sören Laue has been funded by Deutsche Forschungsgemeinschaft (DFG) under grant LA 2971/1-1.

References[Abadi and others 2016] Abadi, M., et al. 2016. Tensorflow:
A system for large-scale machine learning. In USENIX Conference on Operating Systems Design and Implementation(OSDI), 265–283. USENIX Association.
[Baydin et al. 2018] Baydin, A. G.; Pearlmutter, B. A.; Radul,
A. A.; and Siskind, J. M. 2018. Automatic differentiationin machine learning: a survey. Journal of Machine Learning
Research 18(153):1–43.
[Bischof, Hovland, and Norris 2002] Bischof, C. H.; Hovland, P. D.; and Norris, B. 2002. Implementation of automaticdifferentiation tools. In ACM SIGPLAN Notices, volume 37,
98–107. ACM.
[Dey et al. 2018] Dey, P.; Nag, K.; Pal, T.; and Pal, N. R. 2018.
Regularizing multilayer perceptron for robustness. IEEE
Trans. Systems, Man, and Cybernetics: Systems 48(8):1255–
1266.
[Frostig, Johnson, and Leary 2018] Frostig, R.; Johnson,
M. J.; and Leary, C. 2018. Compiling machine learning
programs via high-level tracing. In Conference on Systems
and Machine Learning (SysML).
[Gebremedhin et al. 2009] Gebremedhin, A. H.; Tarafdar, A.;
Pothen, A.; and Walther, A. 2009. Efficient computation ofsparse hessians using coloring and automatic differentiation.
INFORMS Journal on Computing 21(2):209–223.
[Ghorbani, Krishnan, and Xiao 2019] Ghorbani, B.; Krishnan, S.; and Xiao, Y. 2019. An investigation into neuralnet optimization via hessian eigenvalue density. In International Conference on Machine Learning (ICML).
[Giles 2008] Giles, M. B. 2008. Collected matrix derivative
results for forward and reverse mode algorithmic differentiation. In Advances in Automatic Differentiation, 35–44.
Springer Berlin Heidelberg.
[Griewank and Walther 2008] Griewank, A., and Walther, A.
2008. Evaluating derivatives - principles and techniques ofalgorithmic differentiation (2. ed.). SIAM.
[Kakade and Lee 2018] Kakade, S. M., and Lee, J. D. 2018.
Provably correct automatic sub-differentiation for qualified
programs. In Advances in Neural Information Processing
Systems (NeurIPS). 7125–7135.
[Laue, Mitterreiter, and Giesen 2018] Laue, S.; Mitterreiter,M.; and Giesen, J. 2018. Computing higher order derivativesof matrix and tensor expressions. In Advances in Neural
Information Processing Systems (NeurIPS).
[Laue, Mitterreiter, and Giesen 2019] Laue, S.; Mitterreiter,M.; and Giesen, J. 2019. GENO – GENeric Optimization
for Classical Machine Learning. In Advances in Neural
Information Processing Systems (NeurIPS).
[LeCun ] LeCun, Y. Deep Learning est mort. Vive Differentiable Programming! https://www.facebook.com/yann.lecun/
posts/10155003011462143.

[Maclaurin, Duvenaud, and Adams 2015] Maclaurin, D.; Duvenaud, D.; and Adams, R. P. 2015. Autograd: Effortless
gradients in numpy. In ICML AutoML workshop.
[Magnus and Neudecker 2007] Magnus,
J.
R.,
and
Neudecker, H.
2007.
Matrix Differential Calculus
with Applications in Statistics and Econometrics. John Wileyand Sons, third edition.
[Naumann 2004] Naumann, U. 2004. Optimal accumulationof jacobian matrices by elimination methods on the dual
computational graph. Math. Program. 99(3):399–421.
[Naumann 2008] Naumann, U. 2008. Optimal jacobian accumulation is np-complete. Math. Program. 112(2):427–441.
[Paszke et al. 2017] Paszke, A.; Gross, S.; Chintala, S.;
Chanan, G.; Yang, E.; DeVito, Z.; Lin, Z.; Desmaison, A.;
Antiga, L.; and Lerer, A. 2017. Automatic differentiation inpytorch. In NIPS Autodiff workshop.
[Pearlmutter 1994] Pearlmutter, B. A. 1994. Fast exact multiplication by the hessian. Neural Computation 6(1):147–160.
[Sagun et al. 2018] Sagun, L.; Evci, U.; Güney, V. U.;
Dauphin, Y.; and Bottou, L. 2018. Empirical analysis of
the hessian of over-parametrized neural networks. In ICLR(Workshop).
[Seeger et al. 2017] Seeger, M. W.; Hetzel, A.; Dai, Z.; andLawrence, N. D. 2017. Auto-differentiating linear algebra.
In NIPS Autodiff workshop.
[van Merriënboer et al. 2018] van Merriënboer, B.; Breuleux,O.; Bergeron, A.; and Lamblin, P. 2018. Automatic differentiation in ml: Where we are and where we should begoing. In Advances in Neural Information Processing Systems (NeurIPS).
[van Merriënboer, Moldovan, and Wiltschko 2018] van Merriënboer, B.; Moldovan, D.; and Wiltschko, A. 2018. Tangent:
Automatic differentiation using source-code transformation
for dynamically typed array programming. In Advances in
Neural Information Processing Systems (NeurIPS).
[Vasilache et al. 2018] Vasilache, N.; Zinenko, O.; Theodoridis, T.; Goyal, P.; DeVito, Z.; Moses, W. S.; Verdoolaege, S.;
Adams, A.; and Cohen, A. 2018. Tensor comprehensions:
Framework-agnostic high-performance machine learning abstractions. arXiv preprint arXiv:1802.04730.
[Wang et al. 2018] Wang, F.; Decker, J.; Wu, X.; Essertel, G.;
and Rompf, T. 2018. Backpropagation with callbacks: Foundations for efficient and expressive differentiable programming. In Advances in Neural Information Processing Systems(NeurIPS).
[XLA and TensorFlow teams 2017] XLA
andTensorFlow teams. 2017. XLA — TensorFlow, compiled.
https://developers.googleblog.com/2017/03/xla-tensorflowcompiled.html.
[XLA 2019] 2019. XLA: Optimizing Compiler for TensorFlow. https://www.tensorflow.org/xla.
[Yao et al. 2018] Yao, Z.; Gholami, A.; Keutzer, K.; and Mahoney, M. W. 2018. Hessian-based analysis of large batch
training and robustness to adversaries. In Advances in Neural
Information Processing Systems (NeurIPS).