# Tree-Preconditioned Differentiable Optimization and Axioms as Layers

arXiv: https://arxiv.org/abs/2601.06036 · HTML: https://mlsanity.com/paper/2601.06036/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 4 of its 20 references that are in the map.
- First posted on arXiv 3 Dec 2025.

## Builds on (most cited first, 4 of 4)

- [OptNet: Differentiable Optimization as a Layer in Neural Networks](https://mlsanity.com/paper/1703.00443/) (2017)
- [Differentiable Convex Optimization Layers](https://mlsanity.com/paper/1910.12430/) (2019)
- [Convex Relaxation of Optimal Power Flow, Part I: Formulations and Equivalence](https://mlsanity.com/paper/1405.0766/) (2014)
- [Convex Relaxation of Optimal Power Flow, Part II: Exactness](https://mlsanity.com/paper/1405.0814/) (2014)

## Abstract (arXiv)

This paper introduces a differentiable framework that embeds the axiomatic structure of Random Utility Models (RUM) directly into deep neural networks. Although projecting empirical choice data onto the RUM polytope is NP-hard in general, we uncover an isomorphism between RUM consistency and flow conservation on the Boolean lattice. Leveraging this combinatorial structure, we derive a novel Tree-Preconditioned Conjugate Gradient solver. By exploiting the spanning tree of the constraint graph, our preconditioner effectively "whitens" the ill-conditioned Hessian spectrum induced by the Interior Point Method barrier, achieving superlinear convergence and scaling to problem sizes previously deemed unsolvable. We further formulate the projection as a differentiable layer via the Implicit Function Theorem, where the exact Jacobian propagates geometric constraints during backpropagation. Empirical results demonstrate that this "Axioms-as-Layers" paradigm eliminates the structural overfitting inherent in penalty-based methods, enabling models that are jointly trainable, provably rational, and capable of generalizing from sparse data regimes where standard approximations fail.