  Applications such as unbalanced and fully shuffled regression can be
approached by optimizing regularized optimal transport (OT) distances, such as
the entropic OT and Sinkhorn distances. A common approach for this optimization
is to use a first-order optimizer, which requires the gradient of the OT
distance. For faster convergence, one might also resort to a second-order
optimizer, which additionally requires the Hessian. The computations of these
derivatives are crucial for efficient and accurate optimization. However, they
present significant challenges in terms of memory consumption and numerical
instability, especially for large datasets and small regularization strengths.
We circumvent these issues by analytically computing the gradients for OT
distances and the Hessian for the entropic OT distance, which was not
previously used due to intricate tensor-wise calculations and the complex
dependency on parameters within the bi-level loss function. Through analytical
derivation and spectral analysis, we identify and resolve the numerical
instability caused by the singularity and ill-posedness of a key linear system.
Consequently, we achieve scalable and stable computation of the Hessian,
enabling the implementation of the stochastic gradient descent (SGD)-Newton
methods. Tests on shuffled regression examples demonstrate that the second
stage of the SGD-Newton method converges orders of magnitude faster than the
gradient descent-only method while achieving significantly more accurate
parameter estimations.
