The Mode of Null-A: Compositional Computation of a Generalized Inverse

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.