Branch-Extended Automatic Differentiation
Abstract
We present branch-extended automatic differentiation (BEAD), a software package that provides information about inactive branches of piecewise-differentiable functions, intended for use in nonsmooth optimization algorithms. BEAD is implemented in Python, building on the popular JAX framework. For supported nonsmooth operations such as abs and max, BEAD identifies alternative branches within tolerance from the given input, and computes gradients of the resulting branches extended to the given input. Such information has been previously used for robust optimization methods, but has required error-prone and tedious manual differentiation, unlike BEAD, which is fully automated for functions written in NumPy-like JAX code. Our experiments show that our approach is as accurate as manually written derivative code, and offers more flexibility in modifying the function and choosing between branch-discovery rules. While we were originally motivated by the manifold sampling algorithm for nonsmooth composite objectives, we also integrate BEAD branch-extension gradients into a nonsmooth BFGS-SQP solver. On two small nonsmooth test problems, the resulting path-aware search resolves failures or substantially reduces the number of evaluations required by the BFGS-SQP solver.
Related
- Derivative Computation in PINNs: Automatic Differentiation, Finite Differences and Beyond
- Coordinate-Free Automatic Differentiation
- Automatic Differentiation with Constraints in Forward and Inverse Problems
- asdex: Automatic Sparse Differentiation in JAX
- Shape calculus and automatic differentiation for multi-phase level-set topology optimisation with unfitted finite elements
BibTeX
@misc{huckelheim2026branchextended,
title = {{Branch-Extended Automatic Differentiation}},
author = {Jan Hückelheim and Jeffrey Larson and Matt Menickelly},
howpublished = {arXiv},
year = {2026},
month = nov,
doi = {10.48550/arxiv.2611.08637},
eprint = {2611.08637},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2611.08637},
note = {CC BY 4.0},
}
RIS
TY - UNPB
TI - Branch-Extended Automatic Differentiation
AU - Jan Hückelheim
AU - Jeffrey Larson
AU - Matt Menickelly
PY - 2026
DA - 2026/11/05
PB - arXiv
DO - 10.48550/arxiv.2611.08637
UR - https://arxiv.org/abs/2611.08637
AB - We present branch-extended automatic differentiation (BEAD), a software package that provides information about inactive branches of piecewise-differentiable functions, intended for use in nonsmooth optimization algorithms. BEAD is implemented in Python, building on the popular JAX framework. For supported nonsmooth operations such as abs and max, BEAD identifies alternative branches within tolerance from the given input, and computes gradients of the resulting branches extended to the given input. Such information has been previously used for robust optimization methods, but has required error-prone and tedious manual differentiation, unlike BEAD, which is fully automated for functions written in NumPy-like JAX code. Our experiments show that our approach is as accurate as manually written derivative code, and offers more flexibility in modifying the function and choosing between branch-discovery rules. While we were originally motivated by the manifold sampling algorithm for nonsmooth composite objectives, we also integrate BEAD branch-extension gradients into a nonsmooth BFGS-SQP solver. On two small nonsmooth test problems, the resulting path-aware search resolves failures or substantially reduces the number of evaluations required by the BFGS-SQP solver.
ER -