Natural Riemannian gradient for learning functional tensor networks

Abstract: We consider machine learning tasks with low-rank functional tree tensor networks (TTN) as the learning model. While in the case of least-squares regression, low-rank functional TTNs can be efficiently optimized using alternating optimization, this is not directly possible in other problems, such as multinomial logistic regression. We propose a natural Riemannian gradient descent type approach applicable to arbitrary losses which is based on the natural gradient by Amari. In particular, the search direction obtained by the natural gradient is independent of the choice of basis of the underlying functional tensor product space. Our framework applies to both the factorized and manifold-based approach for representing the functional TTN. For practical application, we propose a hierarchy of efficient approximations to the true natural Riemannian gradient for computing the updates in the parameter space. Numerical experiments confirm our theoretical findings on common classification datasets and show that using natural Riemannian gradient descent for learning considerably improves convergence behavior when compared to standard Riemannian gradient methods.

• The paper presents a natural Riemannian gradient descent framework tailored to optimizing functional low-rank tensor networks in challenging non-quadratic loss settings.
• It leverages block-diagonal and diagonal approximations to balance computational efficiency with Riemannian invariance in high-dimensional tensor models.
• Experimental results demonstrate superior convergence rates and enhanced test accuracy in both least-squares regression and multinomial logistic classification tasks.

Natural Riemannian Gradient for Learning Functional Tensor Networks

Introduction and Context

The paper "Natural Riemannian gradient for learning functional tensor networks" (2604.09263) addresses optimization in machine learning models based on low-rank functional tree tensor networks (FTNs), focusing on the role of the natural gradient in Riemannian manifolds when the learning model is a functional low-rank tensor manifold. Unlike conventional neural networks, FTNs parameterize functions via high-dimensional coefficient tensors, restricted to low-rank tensor formats to maintain tractability and storage efficiency. The main technical challenge arises when generalizing optimization schemes from least-squares to non-quadratic losses (e.g., multinomial logistic regression), where classical alternating optimization breaks down. The authors systematically develop a natural Riemannian gradient descent framework tailored to learning in these tensor models.

Natural Gradient Descent in the Riemannian Setting

The paper begins with a self-contained derivation of natural gradient descent as a reparametrization-invariant steepest descent direction on Riemannian manifolds. The authors consider empirical risk minimization over a hypothesis set $H$ modeled as a finite-dimensional Riemannian manifold, often accessed through a submersive but not injective parametrization $F: M \to H$, where $M$ is another manifold. The natural Riemannian gradient is computed by solving (for the parameter space)

$$DF(\Theta)^* DF(\Theta) [\zeta] = \operatorname{grad}(R \circ F)(\Theta).$$

This yields a search direction in parameter space that corresponds to the steepest descent in hypothesis space, invariant to parametrization and basis choices (first-order invariance). The empirical analog is solved at each iteration, with sample-based approximations of $DF(\Theta)^* DF(\Theta)$ and $\operatorname{grad}(R \circ F)$.

One critical observation is the non-uniqueness of the natural gradient when $F$ has kernel, motivating a quotient manifold treatment. When $F$ is an isometry, the natural gradient coincides with the Riemannian gradient, but this is rarely the case in FTNs due to multilinear, overparameterized representations.

Functional Tensor Networks: Structure and Parameterization

The authors formalize FTNs as mappings $f(x) = \mathsf{A} \{\Phi(x)\}_{1,\dots,d}$, where $\mathsf{A}$ is an order-$F: M \to H$0 coefficient tensor, and $F: M \to H$1 is a tensor-product feature map synthesizing univariate bases. The primary tractability constraint is that $F: M \to H$2 is restricted to a fixed-rank tree tensor network manifold, typically parameterized by a composition of multilinear maps from core-tensor factors.

They discuss the functional Tucker, tensor train (TT), and hierarchical Tucker (HT) formats, as well as the specific structure of balanced binary tree TTNs, and provide explicit parameter manifolds and quotient structures to address gauge freedom and non-uniqueness in tensor decompositions. Notably, the overall hypothesis manifold $F: M \to H$3 is parameterized as $F: M \to H$4, where $F: M \to H$5 applies contraction with the basis feature map and $F: M \to H$6 reconstructs the coefficient tensor from core parameters.

The manuscript contains an in-depth discussion of the invariance of FTNs to basis transformations in the underlying tensor product space, with practical observations on how the conditioning of the natural gradient system depends on basis and sampling measure, especially outside the idealized orthonormal basis/product measure setting.

Computation of the Natural Riemannian Gradient

Least-Squares Regression

For the quadratic loss, the empirical risk minimization in FTNs simplifies structurally, and the natural gradient corresponds to a Gauss-Newton direction in tensor parameter space. Through repeated use of the chain rule and inner products on $F: M \to H$7 function spaces, the authors derive explicit formulas for the matrix $F: M \to H$8, showing it decomposes (under product measures and orthonormal bases) into a tensor product of Gram matrices of univariate bases, which considerably reduces computational complexity.

Multinomial Logistic Regression

Classification introduces more complexity. Here, the hypothesis space is further composed with a softmax mapping; the output is a simplex-valued function, and the loss is the negative log-likelihood (cross-entropy). The canonical Riemannian structure is now the Fisher-Rao geometry, where tangent spaces correspond to mean-zero perturbations and the natural metric involves the Fisher information. The computation of the natural gradient for the composed softmax-FTN involves propagation of derivatives through the softmax and weighting by Fisher information, resulting in empirical systems with sample-dependent, structured preconditioners.

Hierarchy of Practical Approximations

Direct computation of the natural Riemannian gradient is infeasible for large FTNs due to the scale of $F: M \to H$9. The authors propose a hierarchy of structured approximations balancing computational tractability and fidelity:

1. Block-Diagonal Approximation (BD-ngrad): Only intra-core blocks are kept, yielding $M$0 smaller independent systems, akin to blockwise natural gradient/K-FAC in neural networks.
2. Rank-One Surrogates (BDO-ngrad): Sample-based compression of Fisher-blocks to a single outer product, producing a further per-class/per-core reduction in complexity.
3. Fully Diagonal (D-ngrad): Each block is replaced by a scaled identity ($M$1), with $M$2 estimated robustly (e.g., via power iteration). This yields a preconditioned SGD/ADAM-like scheme with Riemannian momentum, including appropriate vector transport.

All algorithms are formulated in pseudo-code, emphasizing matrix-free iterative solution strategies essential for scaling to moderately high-dimensional tensor models.

Numerical Experiments

Experiments are conducted on least-squares recovery (where FTN ground-truth is known), as well as deterministic and stochastic multinomial regression (digits and MNIST). The results demonstrate:

• The natural Riemannian gradient approach exhibits superior convergence rates in iteration count, with less sensitivity to basis conditioning compared to standard Riemannian gradient descent.
• While each natural gradient iteration is more expensive, convergence in wall-clock time is competitive and often better, especially in block-diagonal or diagonalized variants.
• In stochastic regimes (large-scale minibatch), the diagonal (D-ngrad) method achieves both faster convergence and higher test accuracy on MNIST relative to standard Riemannian gradient with momentum.

Theoretical and Practical Implications

The work establishes a principled pipeline for Riemannian natural gradient optimization in the context of functional low-rank tensor models. The proposed methods close the gap between information geometry and practical large-scale optimization, specifically in non-Euclidean parameterizations relevant to quantum-inspired models, high-dimensional regressions, and classification tasks where tensor train-type networks are attractive for their memory and computational complexity.

On the theoretical side, the approach clarifies the connection between statistical natural gradients and Riemannian geometry for both least-squares (Gauss-Newton) and categorical (Fisher-Rao) losses, providing a generalized language that accommodates arbitrarily complex chains of manifold parameterizations. The block-diagonal and diagonal approximations are justified as cost-sensitive relaxations that still capture most of the invariance and conditioning benefits of the natural gradient.

Practically, the diagonal method (D-ngrad) enables scalable minibatch optimization with momentum and vector transport, analogous to ADAM/Fisher-ADAM, but built from genuine Riemannian constructions.

Future Directions

The paper identifies multiple open research directions: (1) design of rank-adaptive FTN architectures guided by natural gradient-based optimization, (2) extension to compositional FTN models (deep tensor networks), incorporating the proposed approximations in multilayer architectures, and (3) further analysis of convergence rates and generalization error under different empirical Fisher approximations in the Riemannian setting.

Conclusion

This work provides a rigorous and practically relevant framework for natural Riemannian gradient descent in the context of learning functional tensor networks. By establishing both the foundational geometric underpinning and a suite of efficient approximations for high-dimensional settings, the paper advances the state of the art for optimization in low-rank tensor models, with clear implications for scalable machine learning in structured, non-Euclidean parameter spaces.