Algorithm finds largest input space mapping through sequential matrix products
The Mode of Null-A: Compositional Computation of a Generalized Inverse
Programming Languages
Summary
Finding all possible inputs that produce outputs within a certain range can be tricky when the transformation involves many steps. The authors developed a method to compute these input sets efficiently when the transformation is broken down into a chain of simpler steps with special matrix shapes. This can help in automatic differentiation, a technique used in machine learning and scientific computing, by giving more information than just a single solution. Their approach works even when the number of variables changes during the process, and they tested it on both small CPU tasks and larger GPU tasks involving common operations like convolution and attention.
What this means in practice
- •For machine learning engineers: Analyze and invert complex transformations in neural networks involving convolution and attention layers to better understand and debug model behavior.
- •For numerical software developers: Implement efficient inverse automatic differentiation for programs where variable counts can change dynamically during computation.
Authors
Barak A. Pearlmutter, Jeffrey Mark Siskind
Abstract
We present a novel algorithm for calculating the preimage of an affine space through a product ${J}={J}_{T-1}\cdots{J}_0$ of matrices ${J}_t$ of special form: finding the largest input space $\mathbf{X}$ such that $\mathbf{x}\in\mathbf{X}$ implies ${J}\mathbf{x}\in \mathbf{Y}$, where $\mathbf{Y}$ is a given output affine space. These special matrices arise in AD, where the Jacobians $J$ describing the linearized computation have precisely this structure: the product of a series of linearized primitive numeric operations. This allows us to use the new algorithm to formulate Null-A mode preimage AD, which finds the affine preimage through the Jacobian or Jacobian transpose of a numeric computation. This is a generalization of the inverse AD problem of solving ${J}\acute{\mathbf{x}}^{\ast}=\acute{\mathbf{y}}^{\ast}$ or ${J}^{T}\grave{\mathbf{y}}^{\ast}=\grave{\mathbf{x}}^{\ast}$. The key is to represent affine spaces in a fashion which lends itself to efficient preimage calculation, in a compositional and quasi-local fashion, through a succession of matrices ${J}_t$. Unlike previous methods, Null-A preimage mode AD allows the ${J}_t$ matrices to be non-square, corresponding to a computer program whose number of active variables swells and shrinks during the computation. When ${J}$ is square and the initial affine space is a single point, this finds the conventional inverse. But in the more general case, having the entire affine space provides freedom which can be leveraged in a problem-specific manner. We apply the method to small problems on-CPU where the $J_t$ are linearized scalar unary or binary numeric functions; and to larger problems on-GPU where the $J_t$ are linearized aggregate array operations like convolution and attention.