TorchJD: Jacobian Descent in PyTorch
summary
TorchJD provides a PyTorch library that computes per-objective Jacobians and aggregates them into update directions, implementing Jacobian descent with functions like autoja.jac, autoja.backward, autoja.jac_to_grad. The demo shows per-sample losses (batch size 16) producing a 16×param Jacobian, then applying an unconflicting projection to guarantee nonnegative inner product with every loss; library ships 15+ aggregation methods and plugs into standard optimizers. Paper: arXiv:2406.16232.
AI-written summary · about
related talks
09mHigh-Level API for Autograd
11mLightning Talk: But How Do You Autograd Through Your Simulator? Making a...
33mStochastic Quasi-Gradient Methods: Variance Reduction via Jacobian Sketching
29mAutomatic differentiation, a tale of two languages
14mJAX: Accelerated Machine Learning Research via Composable Function Transformations in Python
26mTraining Large Deep Learning Models: PyTorch vs JAX
31mHigh performance machine learning with JAX
11mIntro to JAX: Accelerating Machine Learning research