Tree-Preconditioned Differentiable Optimization and Axioms as Layers

TL;DR

The paper addresses enforcing economic rationality via differentiable optimization by embedding Random Utility Model (RUM) axioms as a layer inside neural networks. It develops a hyperplane-based reformulation that converts the NP-hard RUM projection into a large but tractable quadratic program, solvable with a novel Tree-Preconditioned Interior Point Method that yields strong spectral whitening. Axioms-as-Layers are realized by differentiating through the solver using the Implicit Function Theorem, with the gradient computed through a shared SPD system, enabling end-to-end training that is provably rational. Empirical results demonstrate efficient convergence, numerical stability under extreme conditioning, and improved generalization in low-data regimes, illustrating a scalable path toward axiom-informed neuro-symbolic AI.

Abstract

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.

Figures (9)

• Figure 1: Boolean lattice $\mathcal{G}$ and one possible spanning tree/co-tree partition.
• Figure 2: Network Construction
• Figure 3: Verification of Theorem \ref{['thm:convergence']}. Under the Frozen Barrier protocol ($n=10$, ambient dim $\approx 5120$), the PCG iteration count scales linearly with the effective rank of the observed data. The slope is significantly less than 1, indicating strong spectral clustering of the preconditioned operator.
• Figure 4: Convergence Rate Comparison under Extreme Ill-Conditioning We compare the residual norm convergence on a static problem where barrier weights vary from $10^{-2}$ to $10^6$. The Baseline and Jacobi methods stagnate, unable to resolve the structural coupling of the constraints.
• Figure 5: Structural Overfitting in Penalty Methods ($n=8$, $N_{\text{dim}}=1024$).

Theorems & Definitions (24)

• Definition 1: The Boolean Lattice Graph $\mathcal{G}$
• Remark 1
• Definition 2: Flows on $\mathcal{G}$
• Theorem 1
• Theorem 2
• proof
• Remark 2: Geometric Duality and Complexity Separation
• Theorem 3: Adjoint Gradient Computation
• proof
• Remark 3: Computational Equivalence