Natural Riemannian gradient for learning functional tensor networks
Authors: Nikolas Klug, Michael Ulbrich, André Uschmajew, Marius Willner
Organizations: Institute of Mathematics, University of Augsburg, 86159 Augsburg, Germany · Department of Mathematics, Technical University of Munich, 85748 Garching b. München, Germany · Institute of Mathematics & Centre for Advanced Analytics and Predictive Sciences, University of Augsburg, 86159 Augsburg, Germany
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.
Figures & tables
Figure 1: Functional tensor train (left) and functional (binary) tree tensor network (right).
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).
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
ngrad
BD-ngrad
BDO-ngrad
D-ngrad
98.71%
98.71 %
98.71 %
98.20 %
97.17 %
Table 1: Final test accuracies on digits after 500 iterations
Figure 4: Comparison of stochastic Riemannian gradient descent (grad) and D-ngrad for the MNIST dataset.
Tensor networks represent high-dimensional data and quantum states as interconnected low-rank tensors to enable efficient computation, learning, and simulation.
Tree tensor networks (TTNs) are widely used in low-rank approximation and quantum many-body simulation. In this work, we present a formal analysis of the quotient geometry underlying the TTN parameter space. Our framework allows for arbitrary horizontal distributions, and we develop efficient first- and second-order optimization algorithms that exploit this geometry. Additionally, we devise a backpropagation algorithm for training TTNs in a kernel learning setting. We validate our methods through numerical experiments on a representative digit classification task and reveal an important tradeoff between two different horizontal distributions that are available for TTNs: while one offers cleaner geometric statements, the other ultimately leads to more efficient algorithms.
Marius Willner, Marco Trenti, Dirk Lebiedz
Chair of Mathematical Data Science, University of Augsburg, Augsburg, Germany · Tensor AI Solutions GmbH, Pfaffenhofen a.d. Roth, Germany · Institute for Numerical Mathematics, University of Ulm, Ulm, Germany
Tensor networks, originally developed for quantum many-body physics, are promising models for machine learning. We derive stochastic Riemannian optimizers for tree tensor networks (TTNs) on both their parameter and quotient manifolds, including adaptive and learning-rate-free schemes suitable for minibatch training. Using a hybrid CNN-TTN architecture, we evaluate the methods on Fashion-MNIST, CIFAR10, and Imagenette. The proposed optimizers achieve predictive performance comparable to unconstrained optimization while enabling numerically stable downstream compression.
Marius Willner, Maximilian Scharf, André Uschmajew +2
Institute of Mathematics, University of Augsburg, 86159 Augsburg, Germany · Centre for Advanced Analytics and Predictive Sciences, University of Augsburg, 86159 Augsburg, Germany · Tensor AI Solutions GmbH, 89284 Pfaffenhofen an der Roth, Germany +1
We explore Riemannian optimization techniques for rank-factored matrix parameters, targeting contemporary deep learning applications. We examine ten points in the algorithm design space: two geometries for rank-r matrices, three geometries for rank-r partial isometries, and block-matrix variants of these five, where factors are shared across block-rows and block-columns. We apply our methods to the multihead attention parameters in small language models. After tuning learning rates, our methods do not conclusively outperform an AdamW baseline. Our implementations are available online.