arXiv CSopen access

Natural Riemannian gradient for learning functional tensor networks

Unknown · 2026 · arxiv_cs

arXiv CS · Papers · License: Open Access · 2026

machine learning, deep learning, neural networks

arXiv:2604.09263v1 [math.OC] 10 Apr 2026

Natural Riemannian gradient for learning functional tensor networks Nikolas Klug∗ Michael Ulbrich† André Uschmajew‡ Marius Willner∗

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.

1 Introduction Many machine learning methods are based on (empirical) risk minimization. Let X ⊆ Rdx , Y ⊆ Rdy , µ be a joint probability measure on X × Y and H be a set of hypotheses, also called the learning model. The goal is to find h ∈ H which minimizes the risk R, defined as Z R : H → R , R(h) := E(x,y)∼µ [ℓ(h, x, y)] = ℓ(h, x, y)µ(dx, dy) . (1.1) X ×Y

Here, ℓ : H × X × Y → R is called loss function and must be sufficiently regular, that is, ℓ(h, ·, ·) ∈ L1 (X × Y, µ) must hold for all h ∈ H. Usually, ℓ takes only nonnegative values. In this work we focus on the case where H is a (finite-dimensional) real differentiable Riemannian manifold, that is, we consider the problem Find

h∗ ∈ arg min R(h) .

(1.2)

h∈H

In practice, the manifold H is usually accessed through a parametrization F : M → H, where M is another, more tractable (finite-dimensional) Riemannian manifold. The parametrization F ∗

Institute of Mathematics, University of Augsburg, 86159 Augsburg, Germany Department of Mathematics, Technical University of Munich, 85748 Garching b. München, Germany ‡ Institute of Mathematics & Centre for Advanced Analytics and Predictive Sciences, University of Augsburg, 86159 Augsburg, Germany †

is usually not unique; there can be many different parametrizations and a “bad” choice can severely impact the behavior of optimization algorithms, in particular first-order methods based on gradient descent. A popular approach to address the influence of the parametrization F is the concept of the natural gradient [4]. Several common classes of learning models H are used in machine learning. Neural networks in various flavors belong to the most prominent examples and, depending on the problem, have proven to be quite successful. In this work, we consider a different class of learning models based on low-rank functional tree tensor networks (TTNs). These models represent functions f : Ω ⊆ Rd → Rn0 in the form f (x) = ⟨A, Φ(x)⟩1,...,d , (1.3) where A ∈ Rn0 ×···×nd is an order-(d + 1) tensor, Φ : Ω ⊆ Rd → Rn1 ×···×nd is a feature map corresponding to point evaluations in a tensor product basis (see Section 3.1), and ⟨·, ·⟩1,...,d denotes tensor contractions along the indices 1, . . . , d. Because of high-dimensionality one can usually not allow arbitrary coefficient tensors A in practice. In low-rank models one hence restricts A to tensors in certain low-rank tensor decompositions which can be efficiently stored in memory and, moreover, allow an efficient computation of tensor contractions. In this work, we consider tensors contained in fixed-rank tree tensor network manifolds M, treated in this work mostly as the quotient manifold w.r.t. a multilinear parametrization. Note that, the learning model given by (1.3) is fairly simple in the sense that it is linear in the parameter tensor A, but nonetheless possesses high expressivity because it parameterizes functions in high-dimensional tensor product spaces, depending on the choice of Φ. By restricting the coefficient tensor A to a low-rank manifold, the learning model becomes nonlinear. While the effect of this restriction to the expressivity is in general difficult to assess rigorously, functional tensor models can offer a more systematic approach to machine learning problems because the overall mathematical theory for tensor methods is already well-developed. Restricting the tensor A in (1.3) to a fixed-rank manifold M naturally yields a parametrization F : M → H of the functional tensor network (FTN) learning model. In this work we show how to efficiently compute natural gradients for such parameterized low-rank FTN models. We develop the theory along two typical machine learning problems: least-squares regression and classification via multinomial logistic regression. In (least-squares) regression, on usually assumes a functional relationship between X and Y, which means that for each x ∈ X , there is a unique y = y(x). The hypotheses are functions h : X → Y and the goal is to find h i h∗ ∈ arg min Ex∼µX ∥h(x) − y(x)∥2Y , h∈H

where µX is the marginal probability measure of µ on X . For such problems, a natural gradient descent approach for compositional functional tensor trains was recently proposed in [12]. In classification, the vectors in X are to be classified into one of n0 classes. For simplicity we assume unique true labels, that is, for every x ∈ X there is again a unique y = y(x) (although our approach is also applicable to the more general case). In multinomial logistic regression (also known as softmax regression), the loss function is the negative log-likelihood of a categorical distribution, which results in the problem to find   n0 X h∗ ∈ arg min Ex∼µX − y(x)j log(h(x)j ) , h∈H

j=1

where y(x) ∈ {0, 1}n0 with y(x)j = 1 iff x belongs to class j.

Note that whereas in least-squares regression, standard low-rank tensor algorithms based on alternating least-squares optimization are possible and are likely to outperform gradient based methods, this is no longer the case for multinomial logistic regression: Here, the subproblems for the individual cores are not linear least-squares problems anymore, necessitating different algorithms such as (Riemannian) gradient descent. Contributions In this work, we investigate how natural gradient descent can be efficiently applied to machine learning tasks with low-rank FTNs as the learning model. In order to rigorously account for the fact, that the parametrization of the learning model H can itself be defined on a manifold M, we first present a self-contained derivation of the natural gradient descent algorithm in a Riemannian framework. The idea behind this is to leave some flexibility regarding the actual optimization methods used the parameter manifold M, explicitly enabling the tools from Riemannian optimization [1, 9]. In particular, while we consider fixed-rank tree tensor networks through the quotient manifold M in the space of tensor factors, one could in principle also consider the embedded manifolds in tensor space(see Section 3.1). We therefore adopt the terminology of natural Riemannian gradient throughout this work. Within our framework, we then derive formulas for computing and approximating the natural Riemannian gradient both for least-squares regression and the multinomial logistic regression setting for classification. In the context of (low-rank) FTNs computing the natural Riemannian gradient poses a central challenge, since one has to solve a linear system (see Equation (2.2)) which can be extremely large even for moderately sized models. To address this problem, we propose a hierarchy of heuristic approximations, which balance approximation quality and computational cost. At their core, our heuristics are driven by a block-diagonalization of the linear system, leveraging the multilinear parametrization of the underlying FTN. The resulting algorithms are applicable to the deterministic learning setting, where all samples can be treated in one batch. In addition, we propose an algorithm for stochastic natural Riemannian gradient descent, where further consideration is required to obtain stable estimates for the natural gradient. While the stochastic version of our algorithm is currently based on an ad hoc approach and not systematically developed, we consider it an additional contribution of our work. In the numerical experiments we compare the proposed algorithms with standard Riemannian gradient descent in practice. For least-squares regression, we conduct tests for a recovery problem on a toy dataset and validate that the choice of the tensor product basis significantly influences the convergence rate for standard Riemannian gradient descent. In comparison, our natural Riemannian gradient approach reduces this effect considerably and requires a fewer number iterations, although in its deterministic version these are more expensive. For experiments in a more realistic setting, we test our deterministic and stochastic algorithms on two standard classification datasets, digits and MNIST. Again, we observe that the algorithms based on natural Riemannian gradient lead to considerably faster convergence w.r.t. the number of iterations. However, even with approximations, natural Riemannian gradient descent still comes at a higher computational cost per iterations than standard Riemannian gradient descent. Despite this, we can achieve faster convergence w.r.t. absolute time. Moreover, in the stochastic regime, our proposed algorithm not only converges faster but also results in improved test accuracy. Related work The concept of a natural gradient as a descent direction based on the (Riemannian) geometry of statistical learning models was introduced by Amari [4] and is nowadays well known. We also

refer to the more recent survey by Martens [27] on the natural gradient, and the considerations about K-FAC [28], which inspired parts of the algorithmic design in this work. Further concerning algorithmics, one of our fully-diagonal approximation schemes is related but not equivalent to the Fisher ADAM algorithm used in [19]. Recently, a formula for evaluating natural Riemannian gradients was proposed by [18], who justify their method using tools from information geometry in a bottom-up approach. This is different from our top-down treatment, where the natural gradient arises naturally from functional considerations. Functional low-rank tensor formats have been initially proposed for high-dimensional applications in quantum chemistry and PDEs; see the survey articles [6, 7, 14] and monographs [16, 20]. In particular, tree tensor networks come in various flavors, the notable examples being the Tucker format, the tensor train (TT) format [31] and the hierarchical Tucker (HT) format [15]. Beginning with seminal works such as [30, 38], low-rank functional tensor networks have received increasing attention for machine learning and are still an active field of research. The work [38], similar as several subsequent ones, e.g. [10, 13, 22], followed an alternating optimization approach similar to DMRG type algorithms for quantum systems. Notably, for the special case of the least-squares loss, subproblems can be solved optimally. While this is not possible for, e.g. classification via logistic regression, they still can be treated via alternating optimization by applying nonlinear solvers to subproblems [10]. Recently, another approach has been taken in [42] by via an expectation-maximization alternating least squares algorithm. In contrast, [30] directly employed Riemannian gradient descent on the fixed-rank TT manifold. Such algorithms are based on the manifold properties of fixed-rank tree tensor networks as established in [39] and had been initially successfully applied for tensor completion problems [11, 23, 36] and also eigenvalue problems [33]. The more recent work by Willner et al. [41] provides a framework for Riemannian gradient descent based on the quotient manifold formalism for functional tensor networks, on which the algorithmic of this work is based. Among the mentioned references, let us particularly highlight the work by Da Silva and Herrmann [11], where for the problem of tensor completion in the hierarchical Tucker format a Gauss-Newton-based algorithm is suggested based on a quotient manifold formalism. This shares similarity with our approach in the sense that it aims at undoing the reparametrization effect. However, their work does not consider functional tensor networks and therefore does not capture the natural gradient in a statistical sense. The treatment of FTNs is more involved because of the extra layer of complexity added by the feature map Φ in Equation (1.3); see also the discussion at the end of Section 3.3. In addition, the empirical approximation with samples requires further considerations which we develop in this work. There is also recent interest in compositional FTNs [34]. These models compose several functional tensor networks (e.g. functional tensor trains) as layers, similar to neural networks, to form one larger model. For compositional functional tensor networks, natural gradient descent was considered in the recent work [12], where the authors compute the natural gradient on the embedded manifold of low-rank tensor coefficients for least-squares problems. Compared to [12], we do not consider composition of tensor networks, but solely focus on the core problem of minimizing a function on a functional low-rank tensor manifold from a Riemannian perspective. Moreover, we compute the natural gradient directly through the multilinear parametrization of the tensor network, which can usually be implemented more conveniently. Outline The work is structured as follows. In Section 2 we present a self-contained derivation of the natural Riemannian gradient descent algorithm for empirical risk minimization and its empirical counter part. The main formulas are (2.2) and (2.9), respectively, describing the Gauss–Newton

type linear systems that need to be solved for computing the natural Riemannian gradient. In Section 3 we then apply the concepts to optimization on functional low-rank tensor manifolds. These manifolds are introduced in Section 3.1 and Section 3.2. As guiding examples the TT format and the balanced binary tree format are discussed. Notably, our representation of these manifolds is a quotient structure in the space of core tensors (see Definition 3.4). An interesting aspect of functional tensor models is the choice of basis in the tensor product space, which has an interesting parallel interpretation as feature map into rank-one tensors. Due to an invariance of low-rank tensors under change of tensor product, the particular choice does not affect the learning model but may still influences practical computations as discussed in Section 3.2. In Section 3.3 and Section 3.4 the computation of natural Riemannian gradient is discussed separately for least-squares regression and multinomial regression, respectively. Section 3.5 then presents our main practical contributions including block diagonal approximation of the Gauß-Newton system, a heuristic stochastic version natural Riemannian gradient descent employing stabilizing momentum transport and further approximation of block diagonals by scaled identities. The numerical experiments are presented in Section 4. They include a more conceptual study for least-squares recovery of an exact FTN model via in Section 4.1, in which we also inspect the influence of the basis choice. The more challenging classification problems via multinomial logistic regression are treated in Section 4.3 where we also employ the full hierarchy of proposed approximations.

2 Natural Riemannian gradient descent The concept of a natural gradient for parametric optimization has been introduced by Amari [4] and is well understood, in particular in the context of statistical models. Mathematically, it shares strong similarities with the idea of Gauß-Newton methods. Nevertheless, in order to clearly work out some often omitted details, we present here a rather self-contained derivation from a Riemannian perspective where we view the natural gradient purely as a parametrizationindependent descent direction for the empirical risk optimization problem (1.2). For this, we assume that H is a finite-dimensional real differentiable manifold with a Riemannian structure. This assumption reflects the practical situation that learning models are described by finitely many parameters. We denote the tangent space at a point h ∈ H by Th H and the Riemannian metric on Th H by ⟨·, ·⟩h . If H′ is another real differentiable Riemannian manifold and f : H → H′ is a differentiable function, we write Df (x)[ζ] for the derivative of f at the point x in direction ζ, that is, Df (x) : Tx H → Tf (x) H′ is the differential of f at x. From now on we assume that all objects (measures, manifolds, functions, loss etc.) are sufficiently smooth.

2.1 Natural Riemannian gradient In principle, the considerations in this subsection apply to an arbitrary smooth map R : H → R. Let Rh : Th H → H be a retraction for the manifold H at a fixed element h (see [9] for the definition). We first consider the linearization of the map R ◦ Rh : Th H → R. Let ζ ∈ Th H, then R(Rh (ζ)) = R(Rh (0)) + DR(Rh (0))[DRh (0)[ζ]] + O(∥ζ∥2 ) = R(h) + DR(Rh (0))[ζ] + O(∥ζ∥2 ) = R(h) + DR(h)[ζ] + O(∥ζ∥2 ) . Similarly to the vector space case, we can obtain a (Riemannian) gradient grad R(h) as the Riesz-representative of DR(h) with respect to the Riemannian metric ⟨·, ·⟩h at h, that is DR(h)[ζ] = ⟨grad R(h), ζ⟩h .

Specifically, for the risk R in (1.1), the Riemannian gradient at h is given by grad R(h) = grad(E(x,y)∼µ [ℓx,y ])(h) = E(x,y)∼µ [grad ℓx,y (h)] where ℓx,y : H → R denotes the map h 7→ ℓ(h, x, y). Here we deliberately interchanged the gradient with the expectation based on our assumption that all objects are sufficiently smooth. Note that different Riemannian metrics can result in different Riemannian gradients. The negative Riemannian gradient −∇R(h) at a point h provides the direction of steepest descent for the function R with respect to the geometry of the manifold H. A common difficulty in practice is that the elements of H can only be accessed through a parametrization h = F (Θ), where F : M → H and M is another Riemannian manifold (possibly a linear space) and Θ ∈ M are the parameters of the model. Hence, for practical implementation of optimization algorithms on H, we need to express all objects in terms of Θ and F . In the following, we always assume that the parametrization F is surjective (allowing for overparametrization) and a submersion, that is, DF (Θ) is surjective for all Θ. Let h = F (Θ). All tangent vectors in Th H are expressed via DF (Θ)[ζ] for some ζ ∈ TΘ M. Therefore, if we wish to express the Riemannian gradient ∇R(h) in the parameter space, we need to solve the following problem: ζ ∈ TΘ M

Find

s.t.

DF (Θ)[ζ] = grad R(F (Θ)) .

(2.1)

This problem is well-defined since F is a submersion, although ζ need not be unique. On the other hand, we usually do not have explicit access to grad R(h) = grad R(F (Θ)) but can only compute grad(R ◦ F )(Θ). Interestingly, the solution of the above problem can be expressed in terms of grad(R ◦ F )(Θ) only. Lemma 2.1. Assume DF (Θ) : TΘ M → TF (Θ) H is surjective. A vector ζ ∗ ∈ TΘ M solves (2.1) if and only if it satisfies DF (Θ)∗ DF (Θ)[ζ ∗ ] = grad(R ◦ F )(Θ)

(2.2)

where DF (Θ)∗ denotes the adjoint of DF (Θ) : TΘ M → TF (Θ) H with respect to the corresponding metrics. A particular solution is ζ ∗ = (DF (Θ)∗ DF (Θ))+ [grad(R ◦ F )(Θ)] ,

(2.3)

where (DF (Θ)∗ DF (Θ))+ denotes the Moore-Penrose inverse. Proof. The chain rule gives grad(R ◦ F )(Θ) = DF (Θ)∗ [grad R(F (Θ))] . Therefore, by applying DF (Θ)∗ to both sides of (2.1) we obtain equation (2.2). On the other hand, when applying (DF (Θ)∗ )+ (the Moore-Penrose inverse of DF (Θ)∗ ) to both sides of (2.2), we obtain (2.1) because (DF (Θ)∗ )+ DF (Θ)∗ = DF (Θ)DF (Θ)+ is the identity on TF (Θ) H = Im DF (Θ). Using this, we can also see that ζ ∗ in (2.3) satisfies DF (Θ)[ζ ∗ ] = DF (Θ)(DF (Θ)∗ DF (Θ))+ DF (Θ)∗ [grad R(F (Θ))] = DF (Θ)DF (Θ)+ [grad R(F (Θ))] = grad R(F (Θ)) , and hence solves (2.1).

Definition 2.2. A vector ζ ∗ ∈ TΘ M is called natural Riemannian gradient for R : H → R w.r.t. the parametrization F : M → H at Θ ∈ M if it satisfies (2.2). Any such vector is denoted by ngrad(R ◦ F ) := ζ ∗ . This notation might be considered slightly abusive, since by this definition, the natural Riemannian gradient is in general not unique. However, whenever we use it, we will silently assume that a particular vector has been picked. The reason why we did not use, e.g., (2.3) for obtaining a unique definition is that in practice we solve the linear system (2.1) so we do not know in advance which solution will be found. A sufficient condition for uniqueness is that F is a local diffeomorphism. However, this is not the case in the settings considered in this work, where the manifolds are parameterized by a multilinear map with inherent (scaling) indeterminacy. In some cases, if the parametrization F has a special structure, the natural gradient is identical to the standard Riemannian gradient. Corollary 2.3. If DF (Θ) : TΘ M → TF (Θ) H is isometric, that is, ⟨DF (Θ)[ζ], DF (Θ)[ξ]⟩F (Θ) = ⟨ζ, ξ⟩Θ , then the natural Riemannian gradient is unique and ngrad(R ◦ F )(Θ) = grad(R ◦ F )(Θ) . In Section 3.4 we briefly show how our definition of the natural Riemannian gradient coincides with that by Amari [4] for a statistical setting. Let us also note that in case F is not surjective, a natural gradient can still be defined as the solution to the least-squares problem min ∥DF (Θ)[ζ] − grad R(F (Θ))∥2F (Θ) ,

ζ∈TΘ

which is the approach taken in [12]. This will lead to the same formula as in Equation (2.3). Since we always assume surjectivity, we use directly the formula (2.2). In case M is a manifold embedded in a finite-dimensional Euclidean ambient space V (in particular M is then equipped with the Riemannian metric inherited from V), we can compute a natural Riemannian gradient by ngrad(R ◦ F )(Θ) = (DF (Θ)∗ DF (Θ))+ grad(R ◦ F )(Θ) = (DF (Θ)∗ DF (Θ))+ PTΘ M (∇(R ◦ F )(Θ)) , where PTΘ M : V → TΘ M is the orthogonal projection from the ambient space to the tangent space TΘ M and ∇ denotes the Euclidean gradient in V. Remark 2.4 (Pullback metric). Following the result from Definition 2.3, the natural gradient can also be interpreted as a standard gradient with respect to a special Riemannian metric, also called the pullback metric [26]. Consider the following symmetric positive semidefinite bilinear form on TΘ M: ⟨ζ, ξ⟩Θ := ⟨DF (Θ)[ζ], DF (Θ)[ξ]⟩F (Θ) . (2.4) If DF (Θ) is injective, ⟨·, ·⟩Θ is an inner product (Riemannian metric resp.) on TΘ M. By definition, DF (Θ) is an isometry with respect to the new (semi-)inner product. Since in this case we have ngrad(R ◦ F )(Θ) = grad(R ◦ F )(Θ), the natural gradient can be interpreted as gradient on M with respect to the Riemannian metric (2.4). Remark 2.5 (Reparametrization). In the case that M is a manifold and itself accessed through a parametrization F̃ : M̃ → M from another Riemannian manifold M̃ (or linear space), the natural gradient changes: One then has to solve the system (DF̃ (Θ̃)∗ DF (F̃ (Θ̃))∗ DF (F̃ (Θ̃))DF̃ (Θ̃))[χ] = grad(R ◦ F ◦ F̃ )(Θ̃) , where Θ̃ ∈ M̃ are the new parameters. However, by construction, the effective descent direction in H remains unchanged.

2.2 An idealized algorithm Recall that Riemannian gradient descent for minimizing the cost R on H constructs from a given point h ∈ H an new point hγ according to hγ = Rh (−γ grad R(h)) ,

(2.5)

where γ > 0 is an appropriate step size and Rh denotes the retraction for H at h. We can interpret this as selecting the new iterate as a point on the curve γ 7→ hγ on H, which by our assumptions is well defined and smooth for small γ > 0. It means that the function R is decreased according to R(hγ ) = R(h) − γ⟨grad R(h), grad R(h)⟩h + O(γ 2 ) . Natural Riemannian descents mimics this update rule, but instead operates on the parameters Θ and uses the natural gradient. If h = F (Θ), this results in the update Θγ = RΘ (−γ ngrad(R ◦ F )(Θ)) .

(2.6)

Note that here RΘ is a retraction on the manifold M. By construction, the curves γ 7→ F (Θγ ) and γ 7→ hγ are then equal in first order, as they they both start at h and their derivatives are equal at γ = 0 due to (2.1): d d F (Θγ ) γ=0 = −DF (Θ)[ngrad(R ◦ F )(Θ)] = − grad R(F (Θ)) = − grad R(h) = hγ γ=0 . dγ dγ This also implies R(F (Θγ )) = R(hγ ) + O(γ 2 ) = R(h) − γ⟨grad R(h), grad R(h)⟩h + O(γ 2 ). Thus, when starting at h = F (Θ), then for a single step both update rules (2.5) and (2.6) are equivalent in first order with respect to the step length. Note that they are not equal, since the higher-order terms differ. Therefore, even if initialized with matching points, both methods will deviate from each other after several iterations. (The corresponding continuous gradient flows are identical though.) On a similar note, while the natural Riemannian gradient is in theory invariant to the parametrization in every step, updates corresponding to different parametrizations will only be equal in first order, so result in different optimization dynamics in practice, as also pointed out in [12]. More importantly, different parametrization usually even lead to different empirical versions of the linear systems which are discussed next. This means that in practice, search directions are generally not equal in first order. The pseudocode of the idealized natural Riemannian gradient descent algorithm is given in Algorithm 1.

2.3 Empirical version In practice the expectation in the definition of the (true) risk R in (1.1) can be approximated with an empirical expected value of samples (xi , y i )k=1,...,m from the distribution µ. This so-called empirical risk is given as m 1 X b Rm (F (Θ)) = ℓ(F (Θ), xi , y i ) m i=1

Algorithm 1: Natural Riemannian Gradient Descent Input: Risk R : H → R, Parametrization F : M → H, Initial point Θ0 ∈ M t←1 while not converged do Compute grad(R ◦ F )(Θ(t) ) Solve DF (Θ(t) )∗ DF (Θ(t) ) [ζ (t) ] = grad(R ◦ F )(Θ(t) ) Choose step-size γ (t) > 0 Θ(t+1) ← RΘ(t) −γ (t) · ζ (t) t←t+1 end return Θ(t)

and thus its Riemannian gradient is b m ◦ F )(Θ) = 1 grad(R m

m X

grad ℓ(F (·), xi , y i )(Θ) .

(2.7)

i=1

The natural Riemannian gradient of the empirical risk is obtained by replacing the right hand side of (2.2) by (2.7). Its practical computation, however, still requires the operator DF (Θ)∗ DF (Θ) : TΘ M → TΘ M, or at least its action on tangent vectors. This can pose a challenge, since there is oftentimes no closed-form expression available for the adjoint of DF (Θ) with respect to the given Riemannian metric. In many relevant applications (such as the ones considered in this work), the hypotheses are functions h : X → Rn0 in L2 (X ,RRn0 ; µX ) where µX denotes the marginal probability measure on X w.r.t. µ, that is, µX (X) = X ×Y 1X (x)µ(dx, dy) where 1X is the indicator function of X ⊆ X . In this setting, computing the operator DF (Θ)∗ DF (Θ) can be achieved through empirical approximation based on the following familiar fact. n0 Proposition 2.6. Assume H = F (M) ⊆ L2 (X R , R ; µX ) is an embedded submanifold that is equipped with the Riemannian metric ⟨η, χ⟩ = X ⟨η(x), χ(x)⟩µX (dx). For Θ ∈ M and x ∈ X , let DFx (Θ) : TΘ M → Rn0 be defined through DFx (Θ)[ξ] = DF (Θ)[ξ](x) for all ξ ∈ TΘ M. Then Z DF (Θ)∗ DF (Θ) = DFx (Θ)∗ DFx (Θ)µX (dx) = Ex∼µX [DFx (Θ)∗ DFx (Θ)] . X

Proof. For all ζ, ξ ∈ TΘ M we have Z ⟨DF (Θ)[ζ], DF (Θ)[ξ]⟩F (Θ) =

⟨DF (Θ)[ζ](x), DF (Θ)[ξ](x)⟩µX (dx) ZX ⟨DFx (Θ)[ζ], DFx (Θ)[ξ]⟩µX (dx)

= ZX =

⟨DFx (Θ)∗ DFx (Θ)[ζ], ξ⟩Θ µX (dx)

ZX = ⟨ DFx (Θ)∗ DFx (Θ)[ζ]µX (dx), ξ⟩Θ XZ ∗ =⟨ DFx (Θ) DFx (Θ)µX (dx) [ζ], ξ⟩Θ X

where both the last and the second to last equality follow from the linearity of the integral.

The empirical approximation of DF (Θ)∗ DF (Θ) suggested by Definition 2.6 hence reads m

DF (Θ)∗ DF (Θ) ≈

1 X DFxi (Θ)∗ DFxi (Θ) , m

(2.8)

i=1

where x1 , . . . , xm are sampled from the distribution µX . Based on (2.7) and (2.8), the empirical natural Riemannian gradient at Θ is defined as the solution ζ to the linear equation m

i=1

1 X 1 X DFxi (Θ)∗ DFxi (Θ)[ζ] = grad ℓ(F (·), xi , y i )(Θ) . m m

(2.9)

b m ◦ F )(Θ). Note that the right hand side Let us abbreviate the equation with Zbm (Θ)[ζ] = grad(R bm (Θ) on TΘ M. Any solution is in the image of the symmetric positive semi-definite operator Z + bm (Θ) grad(R b m ◦ F )(Θ) + η, where η is in the null space of Zbm (Θ) and is of the form ζ = Z b hence orthogonal to grad(Rm ◦ F )(Θ). Therefore, the curve Θγ = RΘ (−γζ) satisfies b (Θγ )) = R(F b (Θ)) − γ⟨grad(R b m ◦ F )(Θ), ζ⟩ + O(γ 2 ) R(F b (Θ)) − γ⟨grad(R b m ◦ F )(Θ), Z bm (Θ)+ grad(R b m ◦ F )(Θ)⟩Θ + O(γ 2 ) = R(F b (Θ)) − γ ≤ R(F

1 b m ◦ F )(Θ), grad(R b m ◦ F )(Θ)⟩Θ + O(γ 2 ) , ⟨grad(R b ∥Zm (Θ)∥Θ

bm (Θ)∥Θ is the spectral norm on TΘ M. This shows, that −ζ is a gradient related where ∥Z descent direction for the empirical risk. Furthermore, if the number m of samples goes to infinity bm (Θ) converges to DF (Θ)∗ DF (Θ) and grad(R b m ◦ F )(Θ) converges to grad(Rm ◦ F )(Θ). then Z The limiting equation hence matches (2.2). Therefore, for fixed Θ, any accumulation point of a corresponding sequence of solutions (ζm ) to (2.9) will be a natural Riemannian gradient for R. The resulting empirical (but still somewhat idealized) version of Algorithm 1 will not be noted separately, as it essentially looks the same except for obtaining ζ (t+1) from solving (2.9) at Θ(t) .

3 Natural Gradient for functional tensor network manifolds We now specialize the above theory to functional low-rank tensor models and develop a practical version of the natural Riemannian gradient algorithm for the learning problems under consideration.

3.1 Functional tensor model Let V1 , . . . , Vd be finite-dimensional vector spaces of real-valued and continuous univariate functions with Vν = span{φν1 , . . . , φνnν } , where φνjν : Ων → R for all jν = 1, . . . , nν and ν = 1, . . . , d and Ων ⊆ R. We assume that φν1 , . . . , φνnν is a basis of Vν so that dim Vν = nν . Further, let V = V1 ⊗ · · · ⊗ Vd be the tensor product space and let n0 ∈ N.

In the functional tensor model we consider vector-functions f ∈ V n0 , that is, functions   f1 (x)   f : Ω1 × · · · × Ωd → Rn0 , f (x) =  ...  , fn0 (x) where every component belongs to V and hence can be written as fk (x) =

n1X ,...,nd

Ak,j1 ,...,jd φ1j1 (x1 )φ2j2 (x2 ) · · · φdjd (xd ) ,

k = 1, . . . , n0 .

(3.1)

j1 ,...,jd

Here A ∈ Rn0 ×n1 ×...×nd is a real-valued n0 × n1 × · · · × nd tensor containing the basis coefficients for all fk w.r.t. the tensor product basis φ1j1 ⊗ · · · ⊗ φdjd of V. Once these basis functions are fixed, the tensor A remains the sole parameter for representing functions in V n0 . The basis representation (3.1) can be written in a more compact form. Let  ν φ1  ..  Φν :=  .  φνnν

for all ν = 1, . . . , d and Φ : Ω1 × · · · × Ωd → Rn1 ×···×nd be given by Φ(x) := Φ1 (x1 ) ⊗ Φ2 (x2 ) ⊗ . . . ⊗ Φd (xd ) ,

(3.2)

which is a rank-1 tensor of functions and can be interpreted as a tensor-valued feature map. Every f ∈ V n0 may hence be equivalently written as f := ⟨A, Φ(·)⟩ := ⟨A, Φ(·)⟩1,...,d =

n1X ,...,nd

Ak,j1 ,...,jν φ1j1 (·1 )φ2j2 (·2 ) · · · φdjd (·d ) ,

j1 ,...,jd

where ⟨·, ·⟩1,...,d denotes the contraction along modes 1, . . . , d. We give a brief example for V n0 . Example 3.1 (Polynomial Basis). Let n ∈ N and V1 = V2 = . . . = Vd = R[x]n be the vector spaces of polynomials up to degree n. For Vν we can, for example, choose the monomial basis functions φνj (ξ) = ξ j−1 for j = 1, . . . , n + 1. For d = 2 and n = 2 this results in the feature map 

   1 1 vec(Φ(x)) = vec(x1  ⊗ x2 ) = (1, x1 , x2 , x21 , x22 , x1 x2 , x21 x2 , x22 x1 , x21 x22 )T . x21 x22 Note that Φ is similar, but not identical to the feature map associated to a polynomial kernel. For d = 2 and n = 2, the latter is given by Φpoly (x) = (1, x1 , x2 , x21 , x22 , x1 x2 )T see e.g. [35]. It generates all multinomial combinations of x1 and x2 of total degree ≤ 2. In contrast, the feature map Φ also generates higher-order multinomials of maximum degree ≤ 2.

3.2 Functionals tree tensor networks Since A ∈ Rn0 ×n1 ×···×nd can easily be too large to be stored, we can usually not take the full linear space V n0 as a (linear) learning model. Instead, further dimensionality reduction is required. Functional low-rank tensor models are based on restricting V n0 to a low-dimensional set H by restricting A to be an element of some tractable set T ⊂ Rn0 ×n1 ×...nd of low-rank tensors. In this way, the class of hypotheses becomes non-linear. In this work, for conceptual reasons we consider only the case that T is an embedded manifold of low-rank tensors. A general class of such manifolds are fixed-rank tree tensor networks (TTNs), which includes fixed-rank versions of the the well-known Tucker format, the tensor train (TT) format [31], or the Hierarchical Tucker (HT) format [15] as special cases. The TT format will be briefly explained further below. Manifold properties for the fixed-rank versions of these examples have been worked out in [17, 39] and are by now well understood. The limitation to fixed rank can be addressed by mixing rank-adaptive strategies with fixed-rank schemes. In the following, for any such manifold T , let G : T → V n0 ,

G(A) := ⟨A, Φ(·)⟩ ,

(3.3)

and define H := Im G. Note that H is a manifold since G is a linear isomorphism. Natural Riemannian gradient descent w.r.t. the parametrization G could in theory directly be applied on the embedded manifolds T ⊆ Rn0 ×n1 ×···×nd using the concepts of Riemannian optimization. Specifically, it requires solving Equation (2.2) on tangent spaces. This is certainly feasible for tree tensor networks. For example, for fixed-rank TT manifolds the required machinery has been applied in several works, e.g., in [24, 30, 33, 36, 40]. Notably, in [12], natural gradient descent has been applied on a TT manifold. In this work, we will follow a more direct approach which accounts for the fact that low-rank tensors are usually more conveniently accessed via another parametrization map τ :M→T ,

Θ 7→ τ (Θ) = τ (Θ1 , . . . , Θd′ ) = A ,

(3.4)

where M = U1 × U2 × · · · × Ud′ for some embedded manifolds Uk ⊆ Rmk , k = 1, . . . , d′ . The map τ is usually multilinear in the sense that there is a multilinear map τ̃ : Rm1 × · · · × Rmd′ → cl(T ) such that τ = τ̃ |T , where cl(T ) denotes set closure. In particular, the differential Dτ (Θ) is given through the Leibniz product rule: Dτ (Θ)[δΘ1 , . . . , δΘd′ ] = τ (δΘ1 , Θ2 , . . . , Θd′ ) + · · · + τ (Θ1 , . . . , Θd′ −1 , δΘd′ ). This is itself a sum of low-rank tensors and can be efficiently evaluated in practical implementations in the relevant examples. Overall, the parametrization we then use is F : M → H,

F (Θ) = (G ◦ τ )(Θ) = ⟨τ (Θ), Φ(·)⟩ .

(3.5)

For concreteness, we briefly discuss the manifold M and the map τ for the fixed-rank TT manifold and for balanced binary TTNs. Example 3.2 (TT format). In the functional TT format, we consider functions ⟨A, Φ(·)⟩ ∈ V n0 where the coefficient tensor A ∈ Rn0 ×···×nd admits an entry-wise decomposition Aj0 ,j1 ,...,jd =

r1X ,...,rd

A0 (j0 , k1 )A1 (k1 , j1 , k2 )A2 (k2 , j2 , k3 ) · · · Ad−1 (kd−1 , jd−1 , kd )Ad (kd , jd ) (3.6)

k1 ,...,kd

n0 A1,...,d

···

Φ1 (x1 )

Φ2 (x2 )

Φd (xd )

A1,2,3,4 r1,2

r3,4

A1,2 n1

Ad−1,d

A3,4 n2

Φ1 (x1 ) Φ2 (x2 ) Φ3 (x3 ) Φ4 (x4 )

···

nd−1

Φd−1 (xd−1 ) Φd (xd )

Figure 1: Functional tensor train (left) and functional (binary) tree tensor network (right).

for some Ak ∈ Rrk ×nk ×rk+1 , k = 0, . . . , d, where r0 = rd+1 = 1.1 The vector r = (r1 , . . . , rd ) is called the TT-rank of A, assuming all values rk are as small as possible for such a representation to exist. We can then conversely consider the manifold of all tensors with a fixed TT-rank r. In this case we choose r ×n ×r M = Rr∗0 ×n0 ×r1 × · · · × R∗d d d+1 , (3.7) r ×n ×r

where by R∗k k k+1 we denote the open set of third-order tensors whose (rk × nk rk+1 ) and (rk nk × rk+1 ) unfolding matrices have full row resp. column rank (equal to rk resp. rk+1 ). The map τ : M → T , (A0 , A1 , . . . , Ad ) 7→ τ (A0 , A1 , . . . , Ad ) = A is then implicitly given by Equation (3.6). We refer to, e.g., [40, Section 9.3] for further details. Due to an inherent nonuniqueness in the representation (3.6) it is in theory possible to further restrict all but one set r ×n ×r R∗k k k+1 to certain Stiefel manifolds. However, we limit the discussion of non-uniqueness to the next example of balanced binary TTNs since we do not use TT in our numerical experiments. We also remark that the core with the output mode of size n0 could in theory be put in any position in the TT. We move it to the first position purely because of notational simplicity, not because it is in any way canonical. Figure 1 (left) depicts a functional tensor train in tensor diagram notation. The cores in the top row are the parameters A0 , A1 , . . . , Ad . The tensors in the bottom row are the (evaluated) features vector Φν (xν ). Example 3.3 (Balanced binary TTN). Functional tree tensor networks (TTNs) are a generalization of functional TTs. We provide an example of a binary balanced functional TTN in Figure 1 (right). The generalization to non-balanced binary trees is straightforward. For a formal definition of the depicted TTN, let d = 2k . We label all the nodes of the balanced tree with t ∈ 2{1,...,d} , t = {(i − 1)2j + 1, . . . , i2j }, i = 1, . . . , 2k−j , j = 1, . . . , k. These index sets t form nested partitions of {1, . . . , d}. Specifically, tL and tR are the left and right children of t iff t = tL ∪ tR and tL < tR elementwise. To each node t that is not a leaf and has left and right children tL and tR , respectively, we then attach third-order tensors At ∈ RrtL ×rtR ×rt , called the cores of the TTN. This requires to fix integers rt for every node t that determine the sizes of cores. For leaves t = {ν} we enforce rt = nν , whereas for the root t = {1, . . . , d} we should take rt = n0 . The TTN then implicitly associates another set of matrices Bt to all nodes t that are not leaves according to the following recursive construction: ( Āt , if |t| = 2 (i.e. t = {ν, ν + 1}), Bt = (3.8) (BtL ⊗ BtR )Āt , if |t| ≥ 4 (i.e. t = tL ∪ tR ) . 1

Note that in this setting the mode sizes of TT cores would commonly be enumerated as r−1 , . . . , rd . In order to avoid negative indices we adopt a shifted enumeration starting at 0.

Here Āt is a reshape of the core At into an (rtL rtR ) × rt matrix. As a result, each matrix Bt is of size nt × rt , where nt is the product of dimensions nν over all ν ∈ t, e.g. n{1,2} = n1 n2 , n{1,2,3,4} = n1 n2 n3 n4 etc. In particular, the matrix B{1,...,d} ∈ R(n1 ···nd )×n0 at the root encodes a tensor A ∈ Rn0 ×n1 ×···×nd . The recursive relation (3.8) induces a multilinear map A = τ ((At )) . acting on the set ×|t|≥2 R(rtL rtR )×rt of tuples (At ) of all cores of specified size. The image of τ is the set T of all tensors A representable in the described recursive way for the fixed choice of bond dimensions (rt ). The corresponding functional TTN model H is obtained via contractions ⟨A, Φ(x)⟩ with rank-one tensors Φ(x). The key point is that performing these contractions is practically feasible using recursion, if the inner sizes rt , also called bond dimensions of the TTN, are moderate, since all computations are performed using only the cores At . The matrices Bt are never explicitly formed. Note that the described format slightly varies from the original HT format from [15] in that no separate cores (bases) were attached to the leaves. However, it corresponds to the balanced HT format for (the slices) of a n0 × (n1 n2 ) × · · · × (nd−1 nd ) reshape of A. Note that in our example the output mode of dimension n0 is attached to the root of the tree. It could, in theory also be attached to any other core, however, the root seems to be the most canonical choice. Now, in order to obtain a smooth manifold, as formally required for the purpose of this work, we would need to further restrict the cores tot satisfy rank(Āt ) = rt for all nodes except for the root. This is possible if and only if rt ≤ rtL rtR for all t with 2 ≤ |t| < d and implies that all matrices Bt have full column rank rt . Without going into detail we note that then all reshaped cores Āt except at the root can be further restricted to Stiefel manifolds without changing the image of τ . We refer to [39] or [7, Section 3.8] for details. In summary, for the reshaped cores Āt the parameter space M could be taken as a Cartesian product of Stiefel manifolds except for the root: ! M = R(r{1,...,d/2} r{d/2+1,...,d} )×n0 ×

× St(r r , r ) tL tR

(3.9)

2≤|t|<d

While this restriction to Stiefel manifolds still does not eliminate all non-uniqueness of the parametrization, it already improves the numerical stability of the TTN representation. For addressing the remaining non-uniqueness a quotient formalism can be employed, as discussed in the following remark. Remark 3.4 (Quotient formalism). The explicit multilinear parametrizations τ : M → T of low-rank tensor manifolds as in (3.4) are convenient and surjective but typically not injective. Correspondingly, the parametrization F = G ◦ τ of the functional low-rank model H will not be injective, where G is the map (3.3). For example, the representation of balanced binary TTNs vie Stiefel manifolds as in Equation (3.9) still includes orthogonal invariances between the inner nodes of the tree. As discussed in Section 2.1 the natural Riemannian gradient will then generally not be unique. This motivates employing a quotient manifold formalism on M for uniquely representing tangent vectors of T . For a detailed description of the quotient formalism for functional TTNs, we refer to [11] and [41]. Here we only state what is relevant to our context. ≡ M, and the orthogonal Concretely, we employ a Cartesian horizontal space denoted by HΘ ≡ ≡ projector onto the Cartesian horizontal space, PΘ : TΘ M → HΘ M. The restriction of Dτ to ≡ M then is bijective. This means that the restriction of F to H ≡ M is a local isomorphism HΘ Θ ≡ M. When restricting to H ≡ M, the and thus there is a unique solution to Equation (2.2) in HΘ Θ system for the natural Riemannian gradient can also be written as PΘ≡ DF (Θ)∗ DF (Θ)PΘ≡ [ζ] = PΘ≡ grad(R ◦ F )(Θ)

(3.10)

≡ M. By applying DF (Θ), it can then be easily verified that and needs to be solved for ζ ∈ HΘ the solution to this system is a natural Riemannian gradient:

DF (Θ)[ζ] = DF (Θ)[PΘ≡ ζ] = DF (Θ)[PΘ≡ ((DF (Θ)PΘ≡ )∗ (DF (Θ)PΘ≡ ))+ (DF (Θ)PΘ≡ )∗ [grad R(F (Θ))]] = (DF (Θ)PΘ≡ )(DF (Θ)PΘ≡ )+ [grad R(F (Θ))] = grad R(F (Θ)). Notably, formula (3.10) coincides with [18, Equation 3.2]. Since their framework also extends to optimization on quotient manifolds, many of their information-geometric considerations and findings apply. We conclude this subsection with a discussion on the choice of bases for the spaces Vν . The functional tensor model assumed a fixed vector of basis functions Φν = (φν1 , . . . , φνnν )T for each Vν , which enter through the feature map Φ in the map G in (3.3) and hence in the overall parametrization F = G ◦ τ . Let us indicate this by writing FΦ instead of F . While the space V n0 containing all functions of the form ⟨A, Φ(·)⟩ obviously does not depend on the particular choice of basis, one may ask how it influences the restriction to a low-rank TTN manifold. Assume the bases are changed according to Ψν = Mν−1 Φν for ν = 1, . . . , d where Mν ∈ Rnν ×nν are invertible. Let Ψ(·) = Ψ1 (·1 ) ⊗ · · · ⊗ Ψd (·d ) be the corresponding feature map. Then for any x ∈ Ω1 × · · · × Ωd we have Φ(x) = Φ1 (x1 ) ⊗ · · · ⊗ Φd (xd ) = (M1 Ψ1 (x1 ) ⊗ · · · ⊗ (Md Ψd (xd )) = (M1 ⊗ · · · ⊗ Md )Ψ(x) where M1 ⊗· · · ⊗ Md denotes the tensor product of linear maps. Correspondingly, for a functional TTN FΦ (Θ) we have FΦ (Θ) = ⟨τ (Θ), Φ(·)⟩ = ⟨(In0 ⊗ M1T ⊗ · · · ⊗ MdT )τ (Θ), Ψ(·)⟩. The considered TTN manifolds T ∈ Rn0 ×n1 ×...···×nd are invariant under tensor product of invertible linear maps, so τ (Θ) ∈ T implies (In0 ⊗ M1T ⊗ · · · ⊗ MdT )τ (Θ) ∈ T . Therefore (In0 ⊗ M1T ⊗ · · · ⊗ MdT )τ (Θ) = τ (Θ̃) for some Θ̃ ∈ M. Finding Θ̃ only requires applying the change of basis to the cores connected to the leaves of the TTN. For example, for the TT format from Definition 3.2 one easily verifies (In0 ⊗ M1T ⊗ · · · ⊗ MdT )τ (A0 , A1 , . . . , Ad ) = τ (Ã0 , Ã1 , . . . , Ãd ) with Ã0 = A0 and Ãν (kν , ·, kν+1 ) = Mν Aν (kν , ·, kν+1 ) for ν = 1, . . . , d. As a result we obtain FΦ (Θ) = FΨ (Θ̃). In conclusion, the change of the tensor product basis can be simply interpreted as a (linear!) reparametrization of the same manifold similar as discussed in Definition 2.5. By construction, the corresponding natural Riemannian gradient will automatically adjust and revert the effect, leading to the same descent direction on H (at corresponding parameters). For practical purposes, however, the choice of basis may still be relevant as can be observed in experiments; see Section 4.1. It was already mentioned that reparametrizations lead to equivalent natural Riemannian gradients only in first order, so methods will deviate over many iterations. Another aspect is that the conditioning of DF ∗ DF and its empirical counterpart can be affected by the choice of basis. This is also indicated by the discussion in Remark 3.6 on orthonormal bases.

3.3 Least-squares regression with functional tree tensor networks We now discuss how to employ the functional low-rank model for least-squares regression, which was also considered in [12]. We set Y = Rn0 and consider hypotheses f ∈ L2 (X , Rn0 ; µX ) together with the least-squares loss ℓ(f, x, y) := ∥f (x) − y∥2Y , where ∥·∥Y is any norm induced by an inner product on Y. As discussed in the introduction, we assume that for each x ∈ X there is a unique y = y(x) such that E(x,y)∼µ [ℓ(f, x, y)] = Ex∼µX [ℓ(f, x, y(x))]. The risk then becomes Z R(f ) = ∥f (x) − y(x)∥2Y µX (dx) = ∥f − y∥2L2 (X ,Y;µX ) . (3.11) X

As a learning model we choose a functional low-rank manifold H ⊆ L2 (X , Rn0 ; µX ). Concretely, we choose H = Im G as in the previous section and consider the parametrization via F : M → H, F (Θ) = (G ◦ τ )(Θ) = ⟨τ (Θ), Φ(·)⟩ as in (3.5). Note that assuming H ⊆ L2 (X , Rn0 ; µX ) in this context may require some additional conditions, but does not seem critical. For example if the marginal measure µX has an integrable density on X , assuming all basis functions φνi to be continuous and square integrable is sufficient. In particular, X = Ω1 × · · · × Ωd ⊆ Rd . We will equip H with the “trivial” Riemannian metric Z ⟨ζ, ξ⟩f := ⟨ζ(x), ξ(x)⟩Y µX (dx) , (3.12) X

where ζ, ξ ∈ Tf H. In particular, the metric is the same at any point. Of course, there are other choices for the Riemannian metric on H, but we do not consider this case for the leastsquares setting. Note that for the specific choice (3.12), the natural Riemannian gradient descent method matches the Gauss-Newton method for minimizing the risk (3.11) in parameterized form f = F (Θ). Let us now consider how to compute DF (Θ)∗ DF (Θ) by first investigating DG(A)∗ DG(A) for a tensor A. Thanks to our choice of the metric (3.12), Definition 2.6 is applicable. Proposition 3.5. It holds that DG(A)∗ DG(A) = In0 ⊗

Φ1 (x1 )Φ1 (x1 )T ⊗ . . . ⊗ Φd (xd )Φd (xd )T µX (dx) .

where ⊗ denotes the tensor product of linear operators. Proof. By Definition 2.6, we have DG(A)∗ DG(A) =

DGx (A)∗ DGx (A)µX (dx) .

(3.13)

Since Gx (A) = ⟨A, Φ(x)⟩ is linear in A, we have DGx (A) = ⟨·, Φ(x)⟩. For rank-one tensors B = b0 ⊗ b1 ⊗ · · · ⊗ bd and C = c0 ⊗ c1 ⊗ · · · ⊗ cd one directly verifies ⟨C, DGx (A)∗ DGx (A)(B)⟩0,1,...,d = ⟨DGx (A)(C), DGx (A)(B)⟩0 ⊤ T ⊤ = ⟨b0 · [bT 1 Φ1 (x1 ) · · · bd Φd (xd )], c0 · [c1 Φ1 (x1 ) · · · cd Φd (xd )]⟩0 T T T T = bT 0 c0 · b1 Φ1 (x1 )Φ1 (x1 ) c1 · · · bd Φd (xd )Φd (xd ) cd

= ⟨B, [In0 ⊗ Φ1 (x1 )Φ1 (x1 )T ⊗ . . . ⊗ Φd (xd )Φd (xd )T ]C⟩0,1,...,d . The formula then extends to all tensors B, C ∈ Rn0 ×n1 ×···×nd and the result follows by (3.13).

From the lemma we obtain Z T T ∗ ∗ Φ1 (x1 )Φ1 (x1 ) ⊗ . . . ⊗ Φd (xd )Φd (xd ) µX (dx) Dτ (Θ) DF (Θ) DF (Θ) = Dτ (Θ) In0 ⊗ X Z Dτ (Θ)∗ In0 ⊗ Φ1 (x1 )Φ1 (x1 )T ⊗ . . . ⊗ Φd (xd )Φd (xd )T Dτ (Θ)µX (dx) . = X

The integral cannot be computed exactly and is approximated empirically according to Equation (2.8), which results in m

1 X Dτ (Θ)∗ In0 ⊗ Φ1 (xi1 )Φ1 (xi1 )T ⊗ · · · ⊗ Φd (xid )Φd (xid )T Dτ (Θ) . m i=1 (3.14) The empirical linear system for the natural Riemannian gradient given by (2.9) then becomes DF (Θ)∗ DF (Θ) ≈

1 X Dτ (Θ)∗ In0 ⊗ Φ1 (xi1 )Φ1 (xi1 )T ⊗ · · · ⊗ Φd (xid )Φd (xid )T Dτ (Θ)[ζ] m i=1

1 X = grad ℓ(F (·), xi , y i )(Θ) . m i=1

Remark 3.6 (µX with product structure). Computing the Riemannian gradient becomes considerably easier if the probability measure µX is a product measure, that is, if µX factorizes into µX (X) = Q1 (X1 )Q2 (X2 ) · · · Qd (Xd ) for X = X1 × X2 × · · · × Xd . This can be the case in applications where it is possible to choose the type of sampling, such as problems involving PDEs. In this case it is possible to choose a tensor product basis {ψi11 ⊗ · · · ⊗ ψiνd } of V that is orthonormal in L2 (X ; µX ) by choosing each basis {ψ1ν ⊗ · · · ⊗ ψnν ν } orthonormal in L2 (Ων ; Qν ). Define Ψ : Rd → Rn1 ×···×nd similar to Equation (3.2) and set G̃(Θ) = ⟨Θ, Ψ(·)⟩ . It is easy to verify that Im G̃ = Im G, as already discussed in Section 3.2. By construction, G̃ is an isometry. Hence ngrad(R ◦ G̃) = grad(R ◦ G̃) and for F̃ := G̃ ◦ τ we have DF̃ (Θ)∗ DF̃ (Θ) = Dτ (Θ)∗ DG̃(τ (Θ))∗ DG̃(τ (Θ)) Dτ (Θ) = Dτ (Θ)∗ Dτ (Θ) . | {z } =I

The operator Dτ (Θ)∗ Dτ (Θ) can even be evaluated exactly and does not have to be approximated empirically. This is exactly what Da Silva and Herrmann [11] employ for their approximate Gauss-Newton method. Using Dτ (Θ)∗ Dτ (Θ) as a substitute for DF (Θ)∗ DF (Θ) when arbitrary measures and bases are involved of course completely discards functional considerations and can then really be seen as undoing the effect of the reparametrization τ . Indeed, applying Dτ (Θ)(Dτ (Θ)∗ Dτ (Θ))+ to grad(R ◦ F ) recovers grad(R ◦ G), independently of the basis used.

3.4 Multinomial logistic regression with functional tree tensor networks For classification tasks, we consider multinomial logistic regression, also called P softmax regression. Here, the task is to classify inputs x ∈ X into n0 classes. Let Sn0 := {p ∈ Rn0 | j pj = 1, pj > 0}

denote the open (probability) simplex. Recall that in this setting, for hypotheses f : X → Sn0 , the risk is generally given by Z X n0 R(f ) = − y(x)j log(f (x)j )µX (dx) , (3.15) X j=1

where y(x) ∈ {0, 1}n0 are the one-hot encoded, noise-free labels with y(x)j = 1 iff x belongs to class j. In the multinomial logistic regression setting, the hypotheses are functions f : X → Sn0 . They can be interpreted as conditional densities of a categorical distribution in the sense that the probability of a point x ∈ X to belong to the class j is given by p(y = j | x) = f (x)j . The set Sn0 can be turned into a Riemannian manifold of discrete probability mass functions. The tangent spaces given by X Tp Sn0 = {q ∈ Rn0 | qj = 0} j

will be equipped with the so called Fisher-Rao metric, given by ⟨ζ, ξ⟩FR p :=

n0 X ζj ξj j=1

1 1 = ζ T diag( , . . . , )ξ . p1 pn0 {z } | =:F (p)

Here, F(p) is the well-known Fisher Information Matrix [3, 32] of a categorical distribution with parameters p which is often defined as " # ∂ log p(y) ∂ log p(y) T F(p) = Ey∼p . ∂p ∂p Of course, there are a priori many possible choices for the metric. The motivation for the FisherRao metric originates in information geometry and statistics. This metric is (up to rescaling) the unique metric invariant under sufficient statistics. This fact is also known as Chentsov’s theorem, see [5, Theorem 2.6]. The softmax function σ : Rn0 → Sn0 is given component-wise by exp(xj ) σ(x)j = Pn0 . k=1 exp(xk ) We choose our learning model as S := σ(H) = {σ ◦ f | f ∈ H} , where H is a manifold of low-rank functional TTNs mapping to Rn0 as described in Section 3.2, together with the parametrization F : M → H, F = G ◦ τ . In other words, the learning model consists of functional TTNs composed with softmax. The manifold S is obviously parameterized by F̄ : M → S with F̄ (Θ) = (σ ◦ G ◦ τ )(Θ) = σ(⟨τ (Θ1 , . . . , Θd ), Φ(·)⟩) . We hence need to compute the natural Riemannian gradient w.r.t. F̄ . Choosing the Fisher-Rao metric on Sn0 results in the following metric for the manifold S Z Z FR ⟨ζ, ξ⟩s = ⟨ζ(x), ξ(x)⟩s(x) µX (dx) = ζ(x)T F(s(x))ξ(x)µX (dx) , (3.16) X

where s ∈ Ts S. Note that s(x) = s(· | x) is a probability mass function in Sn0 . Equation (3.16) is the Fisher-Rao metric on the space of probability densities on X × Y restricted to the densities p whose marginal distributions w.r.t. X are pX = µX . Note that similar to the previous section, we require some regularity on the functions s ∈ S. In particular, for the existence of the integral in Equation (3.15) we require that each component of s : X → Sn0 is in L1 (X , µX ) . This assumption again does not seem critical. For example, assuming all basis functions φνi to be continuous and integrable is sufficient. In particular, X = Ω1 × · · · × Ωd ⊆ Rd . Summarizing the above, we compute the natural Riemannian gradient w.r.t. F̄ and the metric in Equation (3.16). The differential of σ is easily computed as Dσ(x) = (Dσ(x))i,j ) = (σ(x)i (δi,j − σ(x)j )) . Based on this, we provide a formula for the system matrix of in Equation (2.2). Proposition 3.7. For every Θ ∈ T it holds that Z ∗ DF̄ (Θ) DF̄ (Θ) = Dτ (Θ)∗ C(fx ) ⊗ Φ1 (x1 )Φ1 (x1 )T ⊗ . . . ⊗ Φd (xd )Φd (xd )T Dτ (Θ) µX (dx) , X

where fx = G(τ (Θ))(x) = ⟨τ (Θ), Φ(x)⟩ and C(z) = Dσ(z)∗ F(σ(z))Dσ(z) ∈ Rn0 ×n0 . The proof is again an immediate consequence of Definition 2.6. The empirical equivalent is given by m 1 X DF̄ (Θ)∗ DF̄ (Θ) ≈ Dτ (Θ)∗ C(fxi ) ⊗ Φ1 (xi1 )Φ1 (xi1 )T ⊗ . . . ⊗ Φd (xid )Φd (xid )T Dτ (Θ) , m i=1 (3.17) i where fxi = ⟨τ (Θ), Φ(x )⟩. Summarizing the model and parametrization for multinomial logistic regression, we have a chain of functions between manifolds τ

M −→ T −→ H −→ S −→ R . We could in principle choose at which manifold the parametrization “ends”, or equivalently, w.r.t. which (semi-)metric we want to compute the natural Riemannian gradient. By the above reasoning, in this application, the manifold S with the Fisher-Rao metric is the canonical choice. With this choice, the natural Riemannian gradient for the multinomial logistic regression setting considered here in principle coincides with the natural gradient of Amari [4], with the difference that we consider a Riemannian gradient. Remark 3.8. In the manner of interpreting the natural Riemannian gradient as the Riemannian gradient w.r.t. a certain metric, the Fisher-Rao metric can also be interpreted as the metric induced by a log-parametrization p 7→ log ◦ p R to the space of “log-densities”. Choosing the metric ⟨ζ, ξ⟩log p = X ζ(x)ξ(x)p(dx) on Tlog p log(Sn0 ), due to Definition 2.6 leads to Z ∗ D log(p) D log(p) = D logy (p)∗ D logy (p)p(dy) Y

n0 X j=1

(0, . . . , 0,

1 1 , 0, . . . , 0)T (0, . . . , 0, , 0, . . . , 0)pj = F(p) , pj pj

where logy (p) = log p(y). This means that the natural gradient w.r.t. the Fisher-Rao metric on Sn0 points in the direction of steepest descent in the space of log-densities w.r.t. a L2 -type metric.

3.5 Further approximation strategies for practical computation A central challenge for applying natural Riemannian gradient descent to tensor networks is that the operator DF (Θ)∗ DF (Θ) : TΘ M → TΘ M in the main linear system (2.2) is in general extremely large and can neither be computed nor stored efficiently. Consider a balanced binary TTN (Definition 3.3) with d′ cores and bond dimensions r. Let r := max(r) be the maximum rank. Then DF (Θ)∗ DF (Θ) is represented by a matrix of with O(r6 (d′ )2 ) entries, since each core has size O(r3 ) and there are d′ cores. Even for rather small networks, this results in linear systems which cannot be solved efficiently (for r = 8, d′ = 15, the system matrix has ≈ 2.9 × 107 entries). In general DF (Θ)∗ DF (Θ) is not sparse and does not have low-rank structure. We propose three strategies for approximating the operator, focusing on multinomial regression, although most ideas apply more generally. For convenience, these will be presented for the analytical operator, although in practice we apply the strategies to the empirical approximations (2.8) using specific modifications for functional TTNs as discussed in the previous subsections. Block-diagonal approximation of DF (Θ)∗ DF (Θ) Instead of computing the full operator DF (Θ)∗ DF (Θ), we approximate it with a block diagonal approach. This is similar to the K-FAC as proposed in [28], where a block-diagonal approximation to the Fisher information matrix is computed in the setting of neural networks. The block diagonal approximation results in d′ operators acting on the core tensors individually, that is, we obtain d′ systems of size O(r6 ) (instead of O(r6 (d′ )2 )) which makes natural gradient descent tractable. We call the resulting algorithm BD-ngrad, its pseudocode is provided in Algorithm 2. ′ Recall that the parameter manifold M is a Cartesian product of manifolds, that is, M = ×dk=1 Uk for some embedded manifolds Uk (see Equations (3.7) and (3.9)). In the pseudocode, Dτk (Θ) : TΘk Uk → Tτ (Θ) T ,

ζk 7→ τ (Θ1 , . . . , Θk−1 , ζk , Θk+1 , . . . , Θd′ )

denotes the restriction of Dτ (Θ) to the manifold Uk , that is, the restriction to k-th (block-)column of Dτ (Θ). Clearly the block-diagonal of DF (Θ)∗ DF (Θ) is positive semidefinite. Hence, the (negative) solution of the block-diagonal system is still a descent direction for R ◦ F on M (given that the current iterate Θ is not a critical point). Note that the function f ∈ H and the system matrix Wk in Algorithm 2 are in practice never explicitly computed, which is signified by “:=” instead of “←” in the algorithm. Instead, we solve the linear system with a conjugate gradient method and compute evaluations of Wk as needed. Since for functional TTNs, M is embedded in an Euclidean ambient space, the Riemannian gradient gk can be easily computed by first computing the Euclidean gradient ∇(R ◦ F̄ )(Θ), either explicitly or with automated differentiation and subsequently projecting onto the respective tangent space TΘ M. Rank-one approximation of C(z) for multinomial logistic regression In multinomial logistic regression, we need to compute the matrix C(z) = Dσ(z)∗ F(σ(z))Dσ(z) for each sample x, where z = ⟨τ (Θ), Φ(x)⟩. Instead of computing the full matrix, it is possible

Algorithm 2: BD-ngrad for multinomial logistic regression Input: Risk R : S → R given in (3.15), Parametrization F̄ = σ ◦ G ◦ τ : M → S, Initial point Θ0 ∈ M t←1 while not converged do Compute grad(R ◦ F̄ )(Θ(t) ) for k = 1, . . . , d′ do gk ← (grad(R ◦ F̄ )(Θ(t) ))k f := G(τ (Θ(t) )) 1 Pm (t) ∗ i i i T i i T Dτ (Θ(t) ) Wk := m k i=1 Dτk (Θ ) C(f (x )) ⊗ Φ1 (x1 )Φ1 (x1 ) ⊗ . . . ⊗ Φd (xd )Φd (xd ) (t)

(t)

ζk ← solve Wk ζk = gk end (t)

(t)

ζ (t) ← (ζ1 , . . . , ζd′ )T Choose step-size γ (t) > 0 Θ(t+1) ← RΘ(t) −γ (t) · ζ (t) t←t+1 end return Θ(t)

to compute a rank-one approximation instead. This approach was also used in [29]. Note that C(z) = Dσ(z)∗ F(σ(z))Dσ(z) =

n0 X

(Dσ(z)):,j

j=1

1 (Dσ(z))T :,j . σ(z)j

We suggest to approximate C(z) with C(z) ≈

(Dσ(z)):,k (Dσ(z))T :,k σ(z)k

=: C̃(z) ,

where j ∈ {1, . . . , n0 } is sampled from the probability distribution σ(z), that is, P (j = k) = σ(z)k . This reduces the computational effort further by a factor of n0 . We call this blockdiagonal approach with one-shot sampling BDO-ngrad. The resulting algorithm is identical to Algorithm 2 with the sole difference that we swap C for C̃. Stochastic natural Riemannian gradient descent with fully-diagonal approximation For larger datasets it is not possible to compute the gradient with respect to all samples at once. Instead, we apply mini-batch stochastic gradient descent. Since gradients obtained from mini-batches are in general less exact estimates of the true gradient, we use momentum to stabilize the descent. Since we optimize on a manifold, old momenta have to be transported to the new tangent space before adding them to current gradient iterates in an exponential decay averaging fashion according to a decay parameter β1 . This can be done through a transport map TΘ2 ←Θ1 : TΘ1 M → TΘ2 M which projects the previous gradient onto the tangent space of the new point Θ2 ; for more details, see e.g. [9]. Similarly, estimating the operator DF (Θ)∗ DF (Θ) from only a few samples is not expected to produce a good approximation. We suggest to use momentum here, too, again using exponential averaging with decay parameter β2 . Like in Riemannian BFGS this would entail transporting

Algorithm 3: D-ngrad Input: Risk R : S → R given in (3.15), Parametrization F̄ = σ ◦ G ◦ τ : M → S, Initial point Θ0 ∈ M, β1 , β2 ∈ (0, 1] (0) λk ← 1 for k = 1, . . . , d′ (0)

gk ← 0 for k = 1, . . . , d′ t←1 while not converged do for batch B in batches do Compute grad(R ◦ F̄ )(Θ(t) ) for k = 1, . . . , d′ do gk ← (grad(R ◦ F̄ )(Θ(t) ))k f := F̄ (Θ(t) ) 1 P (t) )∗ C̃(f (xi )) ⊗ Φ (xi )Φ (xi )T ⊗ . . . ⊗ Φ (xi )Φ (xi )T Dτ (Θ(t) ) Wk := |B| Dτ (Θ 1 1 k d d k 1 1 i∈B d d gkT Wk gk gkT gk (t) (t−1) λk ← β2 λk + (1 − β2 )λk (t) (t−1) gk ζk ← β1 ζk + (1 − β1 ) (t+1) λk

λk ←

end (t)

(t)

ζ (t) ← (ζ1 , . . . , ζd′ )T Choose step-size γ (t) > 0 Θ(t+1) ← RΘ(t) −γ (t) · ζ (t) ζ (t) ← TΘ(t+1) ←Θ(t) (ζ (t) ) t←t+1 end end return Θ(t)

matrices to the new tangent space, too. While possible in theory, this can be difficult and inefficient in practice. In order to avoid the transport altogether, we propose a fully-diagonal approximation of the operator by approximating the (k, k)-block of DF (Θ)∗ DF (Θ) in the following way: (DF (Θ)∗ DF (Θ))k,k ≈ λmax k I , where λmax is the maximal eigenvalue of (DF (Θ)∗ DF (Θ))k,k . Note that again the (negative) k solution of the fully-diagonal system is a descent direction since the fully-diagonal system is positive definite. The obvious benefit of this fully-diagonal approximation is that we do not have to solve the natural gradient system Equation (2.2) with conjugate gradients, instead, the solution is given max itself can be estimated by simply dividing the right hand side by λmax k . The eigenvalue λk through power iteration. Numerical experiments (see Section 4) show that even just a single iteration of the power method, starting from the current gradient as the initial point, is enough to improve learning. This overall leads to a significant speed-up as well. In practice, we combine this fully-diagonal approximation with the rank-one of C(z) approximation from above. We call the combined algorithm D-ngrad, pseudocode is depicted in Algorithm 3. While this approximation of DF (Θ)∗ DF (Θ) may seem extremely rough and to show little

resemblance with the original derivation of the natural Riemannian gradient, we argue it is still derived from it in a systematic way. Note that the algorithm is similar, but not identical to the ADAM [21] and Fisher ADAM [19] algorithms. These other approaches employ the so-called empirical Fisher information matrix, which importantly is no empirical approximation in the sense of Equation (2.9), although somewhat related to it (see e.g. [27, Section 11]).

4 Numerical experiments In this section, we discuss implementation details and numerical experiments. Implementation details For our experiments we use the balanced binary TTN format as described in Definition 3.3, working directly with the TTN parameter space M, as opposed to the embedded manifold of low-rank tensors T . The evaluations of model responses and Riemannian gradients on M, which are required both for baseline algorithms as well as natural gradients, is done through forward/backpropagation on the tree tensor network, which are worked out in detail in [41, Sec. 6.2]. Concretely, forward propagation calculates the empirical risk b ◦ F )(Θ) = 1 (R m

m X

In0 ⊗ Φ1 (xi1 )T ⊗ · · · ⊗ Φd (xid )T τ (Θ) .

i=1

Naively, this would require the evaluation of m high-dimensional tensor-vector products, but by vectorizing and leveraging the binary tree structure, it can be achieved through a recursion of MTTKRP (matricized-tensor-times-Khatri-Rao-product, coined by Bader and Kolda [8]) operations acting on the individual cores of the network. Similarly, through a series of MTTKRP, backpropagation evaluates b ◦ F )(Θ) = 1 grad(R m

m X

Dτ (Θ)∗ v i ⊗ Φ1 (xi1 ) ⊗ · · · ⊗ Φd (xid ) ,

(4.1)

i=1

where v i = ∇Lyi (F (Θ)(xi )) for Ly : Rn0 → R with Ly (h(x)) = ℓ(h, x, y). Those workloads, as well as other computationally intensive subtasks are handled in parallel through a dedicated C++ library This library interfaces with a Python library which acts on higher abstraction levels, implementing the TTN network architecture and the presented descent algorithms themselves. A comment about the computation the empirical system matrices in (3.14) and (3.17) is in order. They both involve diagonal matrices, In0 and C(z i ) respectively. We denote these matrices by ∆i ∈ Rn0 ×n0 and their diagonal entries ∆ij , j = 1, . . . , n0 . Writing i h q ξji = Dτ (Θ)∗ ej ∆ij ⊗ Φ1 (xi1 ) ⊗ . . . ⊗ Φd (xid ) with ej the j-th Cartesian unit vector, the empirical version of the system matrix reads DF (Θ)∗ DF (Θ) ≈

n0 m X 1 X j=1

By setting v i = ej

ξji (ξji )T .

i=1

q ∆ij the vectors ξji can be efficiently evaluated with an additional backprop-

agation step (4.1). All of our algorithms employ this approach of representing the full system

matrix as the above sum of rank-one terms. Of course, the full matrix is never computed explicitly. For solving the the linear system we instead use the matrix-free CG solver. Note that for the BD-ngrad approximation, only the block-diagonal parts of the outer products have to be considered. For the BDO-ngrad approximation the sum over j additionally reduces to a single term. In D-ngrad the eigenvalues of the blocks are estimated by directly applying the power method in the rank-one decomposition format. In our implementation both standard and natural Riemannian gradients are further subjugated to the quotient manifold formalism described in Definition 3.4. In particular, we employ the Cartesian horizontal space H ≡ M. This means solving an empirical version of the system in Equation (3.10), that is ! ! m m X X 1 1 PΘ≡ DFxi (Θ)∗ DFxi (Θ) PΘ≡ [ζ] = PΘ≡ grad ℓ(F (·), xi , y i )(Θ) . m m i=1

i=1

The operator on the left hand side need not always have full rank. In our experiments, we almost never observed the operator on the left hand side to be of full rank. This is because the number of samples employed were too small to achieve full rank in the above rank-one formulation: m < dim(M ) = dim(range(DF (Θ)∗ )). Therefore we regularize the system according to ! m X 1 PΘ≡ DF (Θ)∗ DF (Θ)PΘ≡ ≈ PΘ≡ DFxi (Θ)∗ DFxi (Θ) + λI PΘ≡ . m i=1

In all experiments we choose the regularization parameter as λ = 5 × 10−3 and solve the system ≡ M using conjugate gradients, since the regularized operator is symmetric positive definite on HΘ on the horizontal space. ≡ M, we can use the computationally Since the computed natural Riemannian gradients are in HΘ inexpensive QR-based retraction described in [41, Sec. 4.4]. A suitable vector transport map for M, which is also compatible with the quotient structure is given by ≡ ≡ M, TΘ2 ←Θ1 : HΘ M → HΘ 1 2

TΘ2 ←Θ1 = PΘ≡2 ,

that is, the projection onto to the Cartesian horizontal space at Θ2 (see e.g. [9, Ex. 10.67]). We use this construction to transport momentum vectors to new iterates for the stochastic experiments. All experiments were conducted on an Intel Ultra 7 155H CPU complemented with 32GB of DDR5 RAM.

4.1 Least-squares recovery problem As a basic proof-of-concept, we consider a simple TTN recovery problem with least-squares loss. As training objective, we generate a randomly initialized balanced binary TTN from Definition 3.3 with 3 cores as parameter Θ ∈ M, TTN-ranks r = (5, 5) and output dimension n0 = 3. The resulting TTN therefore describes a function h∗ = F (Θ) : R4 → R3 . For the bases, we choose monomial bases of order 2, which means that h∗j ∈ V = R[x]2 ⊗ R[x]2 ⊗ R[x]2 ⊗ R[x]2 for j = 1, 2, 3, 4. The training target is to recover the function h∗ from noisy training data. The inputs of the training (m = 256) are generated by sampling uniformly from [−1, 1]4 , Q4datauniform that is, we have µX (X) = i=1 µ[−1,1] (Xi ) for X = X1 × X2 × X3 × X4 . The targets are chosen as y i = h∗ (xi ) + εi , where the components of εi ∈ R3 are sampled from a centered Gaussian distribution with variance σ 2 = 2.5 × 10−3 . A perfect recovery of h∗ would therefore result in an expected training loss of about 3σ 2 = 7.5 × 10−3 .

Figure 2: Comparison of grad and ngrad methods for a recovery problem under change of basis. The natural Riemannian gradient methods achieve the expected minimum loss (black line).

As the starting iterate for the experiments, we take a second randomly generated TTN h0 : R4 → R3 of the similar type with parameter Θ0 ∈ M. For each individual experiment we choose identical basis vectors Φν for all ν = 1, . . . , 4, but vary the concrete basis of R[x]2 . In particular, we conduct experiments using a monomial basis, a normalized Legendre basis (which is an ONB of R[x]2 ⊗ R[x]2 ⊗ R[x]2 ⊗ R[x]2 in L2 (R4 , R; µX )) and a Hermite basis. For the different bases the initial point Θ0 is rescaled such that it always represents the same h0 . In the experiments we compare standard Riemannian gradient descent (grad) and natural Riemannian gradient descent (ngrad) for the different bases. For choosing step sizes, we employ a two-way backtracking line search according to the Armijo-Goldstein criterion, see, e.g. [9, Section 4.5]. Results for one instance are displayed in Figure 2. It can be observed that the natural gradient methods react less sensitive to a change of basis than standard gradient methods, always outperforming their respective counterpart. While in theory the natural Riemannian gradient is independent of the choice of basis, this cannot exactly be observed in practice. Possible reasons for this are discussed at the end of Section 3.2. Out of the standard Riemannian gradient methods, the Legendre basis descents fastest. This is not surprising, since the orthonormal Legendre polynomials are optimally conditioned w.r.t. µX , which implies that at least the function G in the parametrization F = G ◦ τ is an isometry and should also result in a well-conditioned DF (Θ), cf. Remark 3.6.

4.2 Deterministic multinomial logistic regression We evaluate the deterministic algorithms in the multinomial classification setting on the digits dataset [2], which consists of m = 1726 grayscale images of hand-written digits that are to be classified into n0 = 10 classes. Each image consists of 8 × 8 pixels, so we pick d = 64. We employ a 80/20 train-test split. The target labels of the training data are encoded as one-hot vectors (1,x)T in R10 . We choose the basis vectors Φν (x) = ∥(1,x)∥ for all ν = 1, . . . , d, which we adapted from [37], and can be interpreted as a affine linear basis with normalization. We observed that this basis works well in practice, compared to unnormalized bases. A suitable starting iterate and TTN-ranks r were found by using the unsupervised coarse-graining method proposed by Stoudenmire [37], with maximum rank 8. In the experiments we compare standard Riemannian gradient descent, natural Riemannian gradient descent, BD-ngrad and BDO-ngrad. For completeness, we also compare a non-stochastic variant of D-ngrad, where we only have a single batch of size m and only use momentum for the

Figure 3: Comparison of standard Riemannian gradient descent (grad), natural Riemannian gradient descent (ngrad), BD-ngrad, BDO-ngrad and D-ngrad for the digits dataset.

grad 98.71%

ngrad 98.71 %

BD-ngrad 98.71 %

BDO-ngrad 98.20 %

D-ngrad 97.17 %

Table 1: Final test accuracies on digits after 500 iterations

eigenvalues but not for the gradient (which also means there is no transport of gradients). Decay parameters were chosen β1 = 0 and β2 = 0.9. In all algorithms, we use a two-way backtracking line search according to the Armijo-Goldstein criterion to choose step sizes. Results of our experiments are displayed in Figure 3. The top row shows training loss plotted against the number of iterations (left) and time (right). As expected, the proposed hierarchy of approximations subsequently reduces computational effort at the cost of deteriorating convergence. It can be observed that gradient descent overtakes the natural gradient methods at some point (top-right plot). Note however that this happens only when the natural gradient methods have already converged in terms of test accuracy (as seen in the bottom-right figure) and that this makes little to no difference in the final test accuracies after 500 iterations, which are reported in Table 1.

4.3 Stochastic multinomial logistic regression In order to investigate the stochastic setting, we conduct experiments for the larger MNIST dataset [25]. This dataset consists of pictures of handwritten digits that are again to be classified into on of n0 = 10 classes. The test setup is mostly identical to that in Section 4.3; we only highlight differences here. MNIST pictures have a resolution of 28 × 28 pixels, which would require an unbalanced binary TTN. Although unbalanced trees are both theoretically and practically viable, we scale down the samples to 16 × 16, allowing the use of a balanced tree with d = 256. We

Figure 4: Comparison of stochastic Riemannian gradient descent (grad) and D-ngrad for the MNIST dataset.

employ the same feature map as in Section 4.3 and again use the unsupervised coarse-graining method [37] for initialization, this time with maximum tree tensor ranks of 16. Again, we employ a random 80/20 train-test-split. In our experiments, we compare stochastic Riemannian gradient descent with momentum (grad) and D-ngrad. For both algorithms we use batch sizes of 128 and fixed step sizes γ = 16 for grad and γ = 4 for D-ngrad, which were found using a grid search. The momentum decay parameters were chosen as β1 = β2 = 0.9. The plots in Figure 4 compare the training loss of grad and D-ngrad for this setup. It can be observed that the natural gradient method outperforms the classical approach, both in terms of number of iterations and runtime. Furthermore, D-ngrad also achieves a better qualitative result: After 1000 iterations, the final test accuracies are 87.13% for grad and 96.64% for D-ngrad.

5 Outlook In this work we applied the concept of a natural gradient to machine learning tasks with functional TTNs as the learning model. We derived formulas for computing the natural Riemannian gradient both for least-squares regression and multinomial logistic regression and proposed several approximations to the natural gradient that lead to efficient optimization algorithms. The convergence of these methods, depending on the level of approximation, is still an open question. Since our algorithms all work with fixed manifolds M, T and H, choosing and fixing the bond dimensions of the TTN is required a priori, i.e., when designing the model and in particular before optimization. However, it is not clear how to best choose the bond dimensions of the network to achieve a given loss or accuracy. Ideally, bond dimensions would be chosen automatically during optimization, which is, however, not directly possible with the algorithms suggested in this work. The design of a rank-adaptive algorithm based on natural Riemannian gradient descent is left for future research. Functional tensor networks can in theory also be composed as layers to form larger and potentially more expressive models. For functional tensor trains, such a compositional model was considered in [12, 34]. The ideas for approximating the natural Riemannian gradient presented in this work could also be useful in a compositional functional (tree) tensor framework and lead to more efficient optimization algorithms. However, this is left for future research.

Acknowledgments The authors would like to thank Timo Felser and Tensor AI Solutions for providing the code framework that allowed the numerical evaluation of our findings. The work of A.U. was supported by the Deutsche Forschungsgemeinschaft (DFG, German Research Foundation) – Projektnummer 506561557.

References [1] P.-A. Absil, R. Mahony, and R. Sepulchre. Optimization algorithms on matrix manifolds. Princeton University Press, Princeton, NJ, 2008. [2] E. Alpaydin and C. Kaynak. Optical Recognition of Handwritten Digits. UCI Machine Learning Repository, 1998. [3] S.-I. Amari. Information geometry. In Geometry and nature (Madeira, 1995), volume 203 of Contemp. Math., pages 81–95. Amer. Math. Soc., Providence, RI, 1997. [4] Shun-ichi Amari. Natural gradient works efficiently in learning. Neural Comput., 10(2): 251–276, 1998. [5] Shun-ichi Amari and Hiroshi Nagaoka. Methods of information geometry, volume 191 of Translations of Mathematical Monographs. American Mathematical Society, Providence, RI; Oxford University Press, Oxford, 2000. [6] Markus Bachmayr. Low-rank tensor methods for partial differential equations. Acta Numer., 32:1–121, 2023. [7] Markus Bachmayr, Reinhold Schneider, and André Uschmajew. Tensor networks and hierarchical tensors for the solution of high-dimensional partial differential equations. Found. Comput. Math., 16(6):1423–1472, 2016. [8] Brett W. Bader and Tamara G. Kolda. Efficient MATLAB computations with sparse and factored tensors. SIAM J. Sci. Comput., 30(1):205–231, 2007/08. [9] Nicolas Boumal. An introduction to optimization on smooth manifolds. Cambridge University Press, Cambridge, 2023. [10] Z. Chen, K. Batselier, J. A. K. Suykens, and N. Wong. Parallelized tensor train learning of polynomial classifiers. IEEE Trans. Neural Netw. Learn. Syst., 29(10):4621–4632, 2018. [11] Curt Da Silva and Felix J. Herrmann. Optimization on the hierarchical Tucker manifold— applications to tensor completion. Linear Algebra Appl., 481:131–173, 2015. [12] Martin Eigel, Charles Miranda, Anthony Nouy, and David Sommer. Approximation and learning with compositional tensor trains. arXiv:2512.18059, 2025. [13] Alex A. Gorodetsky and John D. Jakeman. Gradient-based optimization for regression in the functional tensor-train format. J. Comput. Phys., 374:1219–1238, 2018. [14] Lars Grasedyck, Daniel Kressner, and Christine Tobler. A literature survey of low-rank tensor approximation techniques. GAMM-Mitt., 36(1):53–78, 2013.

[15] W. Hackbusch and S. Kühn. A new scheme for the tensor representation. J. Fourier Anal. Appl., 15(5):706–722, 2009. [16] Wolfgang Hackbusch. Tensor spaces and numerical tensor calculus. Springer, Cham, second edition, 2019. [17] Sebastian Holtz, Thorsten Rohwedder, and Reinhold Schneider. On manifolds of tensors of fixed TT-rank. Numer. Math., 120(4):701–731, 2012. [18] Jiang Hu, Ruicheng Ao, Anthony Man-Cho So, Minghan Yang, and Zaiwen Wen. Riemannian natural gradient methods. SIAM J. Sci. Comput., 46(1):A204–A231, 2024. [19] Dongseong Hwang. FAdam: Adam is a natural gradient optimizer using diagonal empirical Fisher information. arXiv:2405.12807, 2024. [20] Boris N. Khoromskij. Tensor numerical methods in scientific computing. De Gruyter, Berlin, 2018. [21] Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In ICLR 2015, January 2017. [22] Stefan Klus and Patrick Gelß. Tensor-based algorithms for image classification. Algorithms, 12(11):240, 2019. [23] Daniel Kressner, Michael Steinlechner, and Bart Vandereycken. Low-rank tensor completion by Riemannian optimization. BIT, 54(2):447–468, 2014. [24] Daniel Kressner, Michael Steinlechner, and Bart Vandereycken. Preconditioned low-rank Riemannian optimization for linear systems with tensor product structure. SIAM J. Sci. Comput., 38(4):A2018–A2044, 2016. [25] Yann LeCun, Corinna Cortes, and Christopher J. C. Burges. The MNIST database of handwritten digits, 1998. [26] John M. Lee. Introduction to Riemannian manifolds. Springer, Cham, second edition, 2018. [27] James Martens. New insights and perspectives on the natural gradient method. Journal of Machine Learning Research, 21(146):1–76, 2020. [28] James Martens and Roger Grosse. Optimizing neural networks with kronecker-factored approximate curvature. In Francis Bach and David Blei, editors, Proceedings of the 32nd International Conference on Machine Learning, volume 37 of Proceedings of Machine Learning Research, pages 2408–2417. PMLR, 2015. [29] James Martens and Roger Grosse. Optimizing neural networks with kronecker-factored approximate curvature. arXiv:1503.05671, 2020. [30] A. Novikov, M. Trofimov, and I. Oseledets. Exponential machines. Bull. Pol. Acad. Sci. Tech. Sci., 66(6):789–797, 2018. [31] I. V. Oseledets. Tensor-train decomposition. SIAM J. Sci. Comput., 33(5):2295–2317, 2011. [32] C. Radhakrishna Rao. Information and the accuracy attainable in the estimation of statistical parameters. Bull. Calcutta Math. Soc., 37:81–91, 1945.

[33] Maxim Rakhuba, Alexander Novikov, and Ivan Oseledets. Low-rank Riemannian eigensolver for high-dimensional Hamiltonians. J. Comput. Phys., 396:718–737, 2019. [34] R. Schneider and M. Oster. Some thoughts on compositional tensor networks. In Multiscale, nonlinear and adaptive approximation II, pages 419–447. Springer, Cham. [35] Shai Shalev-Shwartz and Shai Ben-David. Understanding Machine Learning. Cambridge University Press, 2014. [36] Michael Steinlechner. Riemannian optimization for high-dimensional tensor completion. SIAM J. Sci. Comput., 38(5):S461–S484, 2016. [37] E. Miles Stoudenmire. Learning relevant features of data with multi-scale tensor networks. Quantum Sci. Technol., 3(3):034003, 2018. [38] Edwin Stoudenmire and David J Schwab. Supervised learning with tensor networks. In D. Lee, M. Sugiyama, U. Luxburg, I. Guyon, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 29, pages 4799–4807. Curran Associates, Inc., 2016. [39] André Uschmajew and Bart Vandereycken. The geometry of algorithms using hierarchical tensors. Linear Algebra Appl., 439(1):133–166, 2013. [40] André Uschmajew and Bart Vandereycken. Geometric methods on low-rank matrix and tensor manifolds. In Handbook of variational methods for nonlinear geometric data, pages 261–313. Springer, Cham, 2020. [41] Marius Willner, Marco Trenti, and Dirk Lebiedz. Riemannian optimization on tree tensor networks with application in machine learning. arXiv:2507.21726, 2025. [42] Naoya Yamauchi, Hidekata Hontani, and Tatsuya Yokota. Expectation-maximization alternating least squares for tensor network logistic egression. Frontiers Appl. Math. Stat., 11, 2025.

Record · ID 5980 · SHA-256 132195f77643fe80
Conceptio Open Knowledge Archive — every document is proof-bundled with source, license, and retrieval metadata.