1 research outputs found
Automatic Differentiation for Tensor Algebras
Kjolstad et. al. proposed a tensor algebra compiler. It takes expressions
that define a tensor element-wise, such as , and generates the corresponding compute kernel code.
For machine learning, especially deep learning, it is often necessary to
compute the gradient of a loss function with respect
to parameters . If tensor compilers are to be applied in this field,
it is necessary to derive expressions for the derivatives of element-wise
defined tensors, i.e. expressions for .
When the mapping between function indices and argument indices is not 1:1,
special attention is required. For the function , the
derivative of the loss is ; the sum is necessary because index does not appear in the
indices of . Another example is , where is a matrix;
here we have ; the Kronecker delta is
necessary because the derivative is zero for off-diagonal elements. Another
indexing scheme is used by ; here the correct
derivative is , where the range of the
sum must be chosen appropriately.
In this publication we present an algorithm that can handle any case in which
the indices of an argument are an arbitrary linear combination of the indices
of the function, thus all the above examples can be handled. Sums (and their
ranges) and Kronecker deltas are automatically inserted into the derivatives as
necessary. Additionally, the indices are transformed, if required (as in the
last example). The algorithm outputs a symbolic expression that can be
subsequently fed into a tensor algebra compiler.
Source code is provided.Comment: Technical Repor