Efficient Automatic Differentiation of Matrix Functions

Abstract

Forward and reverse mode automatic differentiation methods for functions that take a vector argument make derivative computation efficient. However, the determinant and inverse of a matrix are not readily expressed in the language of vectors. The derivative of a function f(X) for a d ×d matrix X is itself a d ×d matrix. The second derivative, or Hessian, is a d ² ×d ² matrix, and so computing and storing the Hessian can be very costly. In this paper, we present a new calculus for matrix differentiation, and introduce a new matrix operation, the box product, to accomplish this. The box product can be used to elegantly and efficiently compute both the first and second order matrix derivatives of any function that can be expressed solely in terms of arithmetic, transposition, trace and log determinant operations. The Hessian of such a function can be implicitly represented as a sum of Kronecker, outer, and box products, which allows us to compute the Newton step efficiently. Whereas the direct computation requires $\mathcal{O}({d}^{4})$ storage and $\mathcal{O}({d}^{6})$ operations, the indirect representation of the Hessian allows the storage to be reduced to $\mathcal{O}(k{d}^{2})$, where k is the number of times the variable X occurs in the expression for the derivative. Likewise, the cost of computing the Newton direction is reduced to $\mathcal{O}(k{d}^{5})$ in general, and $\mathcal{O}({d}^{3})$ for k = 1 and k = 2.