JacNet: Learning Functions with Structured Jacobians

TL;DR

JacNet tackles the challenge of enforcing derivative structure in neural mappings by learning the Jacobian $J^{\boldsymbol{y}}_{\boldsymbol{x}}$ with a neural network and reconstructing the target function via a line integral. By designing activations that enforce properties like non-singularity or $k$-Lipschitzness, it guarantees invertibility or Lipschitz bounds and computes inverses through integration of the inverse Jacobian. The approach leverages Hadamard-type conditions for global invertibility and differentiable neuralODE-based backpropagation to train the model. Initial experiments on simple 1D functions, including $y=\exp(x)$ and $y=|x|$, illustrate accurate inversion and adherence to the specified Jacobian regularities. Overall, JacNet enables explicit incorporation of Jacobian priors into learned mappings, with potential for improved stability and interpretability in derivative-constrained applications.

Abstract

Neural networks are trained to learn an approximate mapping from an input domain to a target domain. Incorporating prior knowledge about true mappings is critical to learning a useful approximation. With current architectures, it is challenging to enforce structure on the derivatives of the input-output mapping. We propose to use a neural network to directly learn the Jacobian of the input-output function, which allows easy control of the derivative. We focus on structuring the derivative to allow invertibility and also demonstrate that other useful priors, such as $k$-Lipschitz, can be enforced. Using this approach, we can learn approximations to simple functions that are guaranteed to be invertible and easily compute the inverse. We also show similar results for 1-Lipschitz functions.

Figures (3)

• Figure 1: A graph of the invertible target function $\exp(\boldsymbol{\textnormal{x}})$, the prediction function at the trained network weights, and the prediction function at the initial network weights. The initial prediction function is inaccurate, while the final one matches the target function closely. Furthermore, the learned inverse of the prediction function closely matches the true inverse function. We include graphs of an unconstrained learned function whose Jacobian can be zero. Note that the unconstrained function is not invertible everywhere.
• Figure 2: A graph of the 1-Lipschitz target function $|\boldsymbol{\textnormal{x}}|$, the prediction function at the trained network weights, and the prediction function at the initial network weights. Note how the initial prediction function is inaccurate while the final prediction function matches the target function closely. Additionally, note how the learned prediction function is 1-Lipschitz at initialization and after training. We include graphs of an unconstrained learned function whose derivative is not bounded by [-1, 1]. This does not generalize well to unseen data.
• Figure 3: A graph of the empirical risk or training loss versus training iteration. As training progresses, we tighten the tolerance of the numerical integrator to continue decreasing the loss at the cost of more computationally expensive iterations. Top: The training dynamics for learning the invertible target function. Bottom: The training dynamics for learning the Lipschitz target function.