Neural Legendre-Fenchel transform with Hessian Preconditioning

Abstract: The Legendre-Fenchel (LF) transform is a fundamental tool in convex analysis and machine learning that maps lower semi-continuous functions to their convex conjugates. In practice, when closed-form formula are not available for expressing convex conjugates of given functions, one must approximate them using various techniques. One recent such versatile numerical method is the deep Legendre transform method which relies on neural networks although it remains challenging particularly for tackling ill-conditioned functions. This work builds on the reformulation of the LF transform as a projective polarity. A notable property of this framework is its affine invariance. We leverage this affine invariance to introduce a Hessian-based preconditioning strategy. Specifically, we apply an affine deformation around a minimizer so that the second-order Taylor approximation of the function coincides with the canonical paraboloid, whose conjugation map is the identity. A residual network initialized near the identity can then learn this simplified mapping, while the original conjugation map is recovered through the inverse deformation. The proposed preconditioning incurs only a modest computational overhead, consisting of a single eigendecomposition during initialization and two matrix-vector multiplications per query. Experiments on a diverse set of convex functions, including high-dimensional benchmarks, demonstrate improved convergence rates and enhanced numerical accuracy of the conjugation, with particularly significant gains for ill-conditioned problems. Finally, we discuss the scope of applicability of our proposed method and highlight several of its limitations.

• The paper introduces a Hessian preconditioning technique that normalizes the local geometry to simplify neural approximations of the Legendre-Fenchel transform.
• It leverages a projective polarity formulation achieving affine invariance, which accelerates convergence for high-dimensional convex conjugation tasks.
• Empirical results demonstrate significant improvements in accuracy and training speed, especially for ill-conditioned functions.

Neural Legendre-Fenchel Transform with Hessian Preconditioning

Overview

The paper "Neural Legendre-Fenchel transform with Hessian Preconditioning" (2606.09077) introduces a theoretically motivated and empirically validated method for improving neural approximations of the Legendre-Fenchel (LF) transform through Hessian-based preconditioning. The LF transform is fundamental in convex analysis and machine learning, central in problems such as optimal transport and exponential family modeling. While neural models have shown promise for amortizing conjugation in high dimensions, the conditioning of the underlying function frequently impedes efficient optimization and approximation quality. The proposed framework leverages the affine invariance of the projective polarity formulation of convex conjugation, yielding a preconditioning scheme that standardizes the second-order geometry of the objective.

Theoretical Motivation and Methodology

The projective polarity perspective frames the LF transform as a polar set operation induced by a matrix-valued quadratic form on extended (homogeneous) coordinates. Critically, this representation exhibits affine invariance: affinely deforming a function is equivalent to affinely transforming its conjugate, and the conjugation learning task remains structurally unchanged. The authors exploit this property by choosing an affine reparameterization that normalizes the function's local Hessian to identity near its minimizer, matching the canonical quadratic paraboloid whose conjugation map is the identity.

The technical steps are as follows:

1. Local Quadratic Matching: Around the minimum $\theta_0$, set $H = \nabla^2 f(\theta_0)$, and define $z = H^{-1/2}(\theta - \theta_0)$. The transformed function $g(z) = f(\theta_0 + H^{-1/2} z) - f(\theta_0)$ approximates $\frac{1}{2} \|z\|^2$ near the origin.
2. Network Parameterization: A residual network parameterizes the conjugation map in the $z$ space, initialized close to identity, thus starting with minimal inductive bias mismatch.
3. Loss and Training: Training minimizes the Fenchel-Young divergence, which in the polarity context corresponds to a bilinear form on primal and dual graph points. The conjugate value reconstruction requires only evaluations of $f$ at the inferred primal variable.
4. Inference: To return to the original coordinates, the output is affinely mapped back, incurring only two matrix-vector multiplications per query. The preconditioning cost is dominated by the one-off eigendecomposition of $H$.

The primary innovation is the principled Hessian preconditioning enabled by the projective polarity formulation, guaranteeing, in the vicinity of the minimizer, that the learning task for conjugation is locally trivialized for the neural network.

Figure 1: Effectiveness of the preconditioning on a randomly initialized ICNN; following Hessian transformation, the function in $z$ closely matches the canonical paraboloid $\frac{1}{2}z_1^2$.

Empirical Evaluation

The experimental campaign targets a suite of convex functions, with varying levels of conditioning and curvature heterogeneity, evaluated across multiple input dimensions ($H = \nabla^2 f(\theta_0)$0). Metrics include RMSE for the learned conjugation map and the recovered conjugate value, both relative to high-precision numerical ground truth obtained via optimization.

Strong claims are supported by the findings:

• Consistent acceleration and accuracy improvements from preconditioning, most pronounced for ill-conditioned objectives (e.g., scaled quadratics with condition number $H = \nabla^2 f(\theta_0)$1), with RMSE ratios exceeding $H = \nabla^2 f(\theta_0)$2 in high dimensions.
• Faster convergence for the preconditioned network, often by an order of magnitude in training steps, with reduced final error floors.

Figure 2: Training losses; preconditioned model achieves lower loss immediately and converges rapidly compared to the baseline.

Figure 3: Reconstructed conjugate $H = \nabla^2 f(\theta_0)$3 vs. ground truth; preconditioned model (green dashed) closely tracks the ground truth (blue) while baseline deviates substantially.

Figure 4: Additional experiments in $H = \nabla^2 f(\theta_0)$4; Hessian preconditioning consistently improves both convergence and final accuracy across diverse function families and properties.

Qualitatively, the method is robust to function type, with exceptions arising for settings where the local quadratic approximation is invalid or preconditioning amplifies numerical instabilities (e.g., nearly singular Hessians, highly nonlocal curvature variation). These edge cases are rigorously documented.

Implications and Limitations

Practical Impact

1. Enabling High-Dimensional Conjugation: The approach efficiently scales to high-dimensional convex conjugation, where existing gradient-based or neural methods often become impractical without careful preconditioning.
2. Low Inference Overhead: The method’s inference step integrates seamlessly into neural architectures, adding only trivial computational cost.
3. Broad Applicability: The affine invariance result generalizes, allowing for other preconditioning strategies if richer structural priors or domain knowledge are available.

Theoretical Consequences

• The affine invariance in the projective polarity regime connects convex analysis with geometric invariant theory, reinforcing the universality of conjugation under linear transformations.
• The method exemplifies how network initialization and problem geometry, if harmonized, fundamentally alter learning trajectories for implicit function approximation.

Limitations

• Requirement of Accessible Minimizer: Preconditioning is contingent on the existence and efficient computation of a global minimizer and its Hessian, which may not be available in all applications.
• Locality of Approximation: For functions exhibiting strong curvature variation away from the minimizer, the quadratic normalization loses global validity, and the benefits diminish accordingly.
• Hessian Degeneracy: Near-singular Hessians yield unstable deformations; the proposed fix (eigenvalue clamping) degrades preconditioning quality.
• Cost of Hessian Calculation: While one-off, exact Hessian computation for large $H = \nabla^2 f(\theta_0)$5 can be nontrivial; efficient approximations or low-rank methods are possible extensions.

Future Directions

Natural directions for extension include:

• Adaptive/Online Preconditioning: Dynamically refining the linearization point ($H = \nabla^2 f(\theta_0)$6) or updating the deformation during training, to cope with nonstationary objectives or drifting data distributions.
• Joint Learning of Deformation: Parametric learning of the best affine transformation as part of network training via bilevel optimization.
• Extension to Stochastic or Set-Valued Problems: Application in entropic OT and stochastic settings, where the primal function is itself an expectation.

Conclusion

This work formalizes and validates a Hessian preconditioning strategy for neural Legendre-Fenchel conjugation (2606.09077). By aligning the local geometry of the function with the canonical paraboloid, the conjugation learning task is simplified, producing demonstrable gains in speed and accuracy, especially for ill-conditioned or high-dimensional problems. The method is theoretically grounded in the affine invariance properties of the Legendre polarity framework and is computationally efficient. The limitations are well-characterized, pointing toward possible refinements that combine geometric, statistical, and algorithmic techniques for improved convex conjugate learning in machine learning and optimization.