[Paper Review] Second-order optimization for tensors with fixed tensor-train rank

Michael Psenka, Nicolas Boumal|arXiv (Cornell University)|Nov 26, 2020

Tensor decomposition and applications7 references4 citations

TL;DR

This paper presents the first efficient computation of the Riemannian Hessian on the manifold of fixed tensor-train (TT) rank tensors, enabling second-order optimization via the Riemannian trust-region method. The method achieves faster convergence and improved performance in low-rank tensor completion compared to first-order methods and alternating least squares, with a computational complexity of O(d²nr³).

ABSTRACT

There are several different notions of "low rank" for tensors, associated to different formats. Among them, the Tensor Train (TT) format is particularly well suited for tensors of high order, as it circumvents the curse of dimensionality: an appreciable property for certain high-dimensional applications. It is often convenient to model such applications as optimization over the set of tensors with fixed (and low) TT rank. That set is a smooth manifold. Exploiting this fact, others have shown that Riemannian optimization techniques can perform particularly well on tasks such as tensor completion and special large-scale linear systems from PDEs. So far, however, these optimization techniques have been limited to first-order methods, likely because of the technical hurdles in deriving exact expressions for the Riemannian Hessian. In this paper, we derive a formula and efficient algorithm to compute the Riemannian Hessian on this manifold. This allows us to implement second-order optimization algorithms (namely, the Riemannian trust-region method) and to analyze the conditioning of optimization problems over the fixed TT rank manifold. In settings of interest, we show improved optimization performance on tensor completion compared to first-order methods and alternating least squares (ALS). Our work could have applications in training of neural networks with tensor layers. Our code is freely available.

Motivation & Objective

• Address the lack of second-order optimization methods on the manifold of fixed tensor-train (TT) rank tensors.
• Enable superlinear local convergence in optimization by computing the Riemannian Hessian on the TT manifold.
• Develop an efficient algorithm to compute the Riemannian Hessian for tensor completion and related problems.
• Demonstrate improved convergence and performance over first-order methods and alternating least squares (ALS) in low-rank tensor completion tasks.
• Facilitate applications in machine learning, including neural networks with tensor layers, by enabling robust second-order optimization on low-rank tensor manifolds.

Proposed method

• Derive an exact formula for the Riemannian Hessian on the manifold of fixed TT-rank tensors using Riemannian geometry and tensor train structure.
• Leverage the tensor-train format to express the Hessian in terms of core tensors and interface matrices, enabling efficient computation.
• Design an algorithm that computes the Hessian action via forward and backward passes through the TT cores, using recursive matrix products.
• Utilize the structure of the TT format to compute cross-terms and Hessian-vector products in O(d²nr³) operations, avoiding full Hessian assembly.
• Integrate the Hessian computation into the Manopt toolbox for use in Riemannian trust-region methods.
• Apply the method to tensor completion by formulating the optimization problem on the TT manifold and using the Hessian for improved convergence.

Experimental results

Research questions

• RQ1 Can the Riemannian Hessian be efficiently computed on the manifold of fixed tensor-train rank tensors?
• RQ2 Does second-order optimization on the TT manifold lead to faster convergence and better performance than first-order methods in tensor completion?
• RQ3 What is the computational complexity of computing the Riemannian Hessian action for TT-structured tensors?
• RQ4 How does the performance of Riemannian trust-region methods compare to alternating least squares and first-order methods in low-rank tensor completion?
• RQ5 Can second-order optimization on the TT manifold be practically applied to machine learning tasks such as training neural networks with tensor layers?

Key findings

• The paper successfully derives and implements an efficient algorithm to compute the Riemannian Hessian on the fixed TT-rank manifold, enabling second-order optimization.
• The computational complexity of the Hessian-vector product is O(d²nr³), which is scalable and practical for high-order tensors.
• Numerical experiments show that the Riemannian trust-region method achieves faster convergence and better performance than first-order methods and ALS in low-rank tensor completion.
• The method demonstrates superlinear local convergence, confirming the theoretical advantage of second-order optimization on the TT manifold.
• The implementation is publicly available, enabling reproducibility and extension to other applications such as PDEs and neural networks with tensor layers.
• The approach outperforms existing methods in terms of solution accuracy and convergence speed, particularly in ill-conditioned or high-dimensional settings.