PaperlayerView exact revision on arXiv

Source-linked AI summary

Stochastic Optimization of Tree Tensor Networks

Marius Willner, Maximilian Scharf, André Uschmajew, Timo Felser, Marco Trenti

arXiv:2609.00870v1math.OCcs.CVphysics.comp-ph

TL;DR

The paper addresses how to train TTNs stochastically while respecting their parameter and quotient-manifold geometry and preserving properties needed for downstream tensor-network use. It derives adaptive and learning-rate-free Riemannian optimizers, then evaluates them in a hybrid CNN-TTN architecture. Across image-classification tasks, the optimizers achieve performance comparable to unconstrained Adam while producing more numerically stable models for compression.

  • Problem

    Deterministic tensor-network optimizers such as DMRG require exact objectives and environments, whereas large empirical machine-learning losses motivate stochastic minibatch training.

  • Method

    The paper derives explicit stochastic Riemannian optimizers for orthogonal TTNs on both parameter and quotient manifolds, including adaptive and distance-over-gradient schemes.

  • Results

    Riemannian optimizers perform comparably to Adam in convergence, generalization, and test accuracy while learning numerically more stable TTNs for downstream compression.

  • Takeaways & Limitations

    Riemannian optimization preserves TTN structure needed by downstream compression and explainability, despite computational overhead from projections and retractions.

Abstract

from arXiv · show

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.

1 Introduction

Stochastic first-order optimization is central to machine learning, but tensor-network training inherits a mismatch between deterministic physics optimizers and large empirical objectives. The paper derives geometry-respecting stochastic methods for TTNs and evaluates them in a hybrid CNN-TTN framework.

  • Motivation: Minibatch SGD replaces expensive exact gradients with unbiased estimates, trading gradient accuracy for lower per-step cost.The gradient of an empirical average over m examples costs O(m) per step.
  • Motivation: Gradient noise can help escape saddle points and sharp minima and is associated with improved generalization in overparameterized models.
  • Motivation: DMRG locally optimizes tensors using exact environments, making its fixed-objective assumption poorly matched to shuffled, large-scale machine-learning datasets.
  • Prior approaches: Automatic differentiation enables convenient end-to-end tensor-network training but treats the network as a generic model rather than exploiting its geometry.
  • Contribution: The paper derives minibatch-computable stochastic Riemannian updates on TTN parameter and quotient manifolds and evaluates them with a hybrid CNN-TTN architecture.The method uses the TTN’s canonical gauge structure and product geometry to construct explicit updates.
  • Contribution: The hybrid CNN-TTN framework is tested on large-scale image datasets, with competitive optimization and improved conformity to TTN hierarchical structure.

2 Preliminaries

The paper models TTNs as structured collections of core tensors arranged on a rooted binary tree, then restricts them to orthogonal manifolds and quotients out residual gauge freedom. These constructions provide the geometric spaces and operations needed for Riemannian optimization.

  • Tree tensor networks: A TTN assembles order-three core tensors on a rooted binary tree, joining physical or virtual child indices and producing an outgoing output index at the root.Leaf cores join consecutive physical indices; internal cores join virtual indices.
  • Tree tensor networks: The contracted TTN represents a single tensor, while fixed bond dimensions determine model expressivity and computational complexity.Bond dimensions satisfy χ_t ≤ χ_tLχ_tR and are otherwise chosen as fixed model parameters.
  • Manifold structure: Orthogonal TTNs constrain every non-root core to be an isometry on a Stiefel manifold, forming a numerically more stable embedded submanifold.Isometrization preserves the represented full tensor, so restricting optimization to orthogonal TTNs is not a restriction of the represented model.
  • Manifold structure: The orthogonal TTN parameter space is a Cartesian product of core-tensor manifolds, with tangent spaces formed from the corresponding Stiefel tangent spaces.
  • Quotient structure: Orthogonal transformations inserted at virtual links create gauge-equivalent core parameters that leave the contracted tensor unchanged.This freedom changes neighboring core tensors without breaking their isometry constraints.
  • Quotient structure: Quotienting by the gauge group makes equivalent parameters a single abstract representative, while full multilinear-rank constraints are required for a formal quotient manifold.The quotient is theoretically useful even though its elements cannot be directly represented numerically.
  • Riemannian operations: Horizontal spaces retain directions that affect the represented tensor, and projections, retractions, and vector transports support optimization on the parameter manifold.The quotient metric is induced through horizontal lifts.

2.2 Adaptive Riemannian optimization

Adaptive Riemannian optimization extends coordinate-wise gradient rescaling and momentum to product manifolds by treating manifold factors as block coordinates. The framework includes ADAM-like methods and connects to matrix-wise MUON updates.

  • Adaptive schemes: Adaptive Euclidean methods rescale current gradients coordinate-wise using information accumulated from past iterates.
  • Stochastic optimization: Stochastic adaptive methods are especially useful when gradient estimates are sparse and potentially unstable.Momentum and adaptive rescaling can yield substantial improvements over standard gradient descent in this setting.
  • Adaptive Riemannian schemes: On a product Riemannian manifold, points and tangent vectors separate into factor components that serve as block coordinates for adaptive updates.The update direction is followed through a manifold retraction.
  • Adaptive Riemannian schemes: Momentum-based updates can be incorporated through vector transport, allowing past tangent information to move between manifold points.
  • MUON connection: MUON treats each neural-network weight matrix and its gradient as a separate matrix-space block and orthogonalizes each update using polar decomposition.Because these matrix spaces are Euclidean, the basic MUON update does not require a retraction.

2.3 Riemannian distance over gradients

Distance-over-gradient methods adapt stochastic step sizes without manual learning-rate tuning by combining gradient information with distances from the initial iterate. Their manifold extensions require curvature-aware distances and practical distance estimates.

  • Distance over gradients: Distance over gradients provides a tuning-free dynamic step-size formula for stochastic gradient descent and has been extended to manifold optimization.The motivation is to reduce the overhead of selecting α for a machine-learning task.
  • Distance over gradients: The Riemannian step-size schedule uses the geodesic distance between the initial point and the current iterate together with the current gradient.
  • Curvature: A curvature function ζ_κ adjusts the schedule for manifolds with negative sectional curvature.
  • Practical use: The formula applies to stochastic Riemannian gradient descent for geodesically convex functions, with ε preventing zero-magnitude step sizes.Experiments reported in the cited work find DOG-like algorithms relatively insensitive to ε.
  • Distance computation: Grassmann manifolds quotient Stiefel manifolds by orthogonal transformations, and their distances can be expressed through singular values of Stiefel representatives.Grassmann sectional curvature is non-negative.

3 Tools

The TTN parameter space supports block-separable adaptive optimization, while quotient optimization can use lifted representatives despite lacking a product structure. DOG-like methods use geodesic-distance proxies, with quotient distances requiring a representation caveat.

  • TTN geometry: The TTN parameter manifold factorizes as a Cartesian product, so core tensors serve as block coordinates for adaptive Riemannian schemes.The metric acts separably on block-coordinate variations.
  • Quotient geometry: The quotient manifold couples all network nodes, but quotient tangents and metrics lift uniquely to decomposable total-space representatives.This makes adaptive schemes practical through horizontal lifts and total-space computations.
  • Distance over gradients: Geodesic distance is the main obstacle to DOG-like algorithms on both the parameter and quotient TTN manifolds.The distance computation is required for learning-rate-free step-size schemes.
  • Distance over gradients: On the parameter manifold, node-wise distances combine into a proxy that underestimates the true distance and therefore produces smaller, more conservative step-sizes.The proxy uses Stiefel distances for non-root tensors and Euclidean distance at the root.
  • Distance over gradients: On the quotient, the distance expression is derived from a connecting geodesic and horizontal lift, using the quotient’s shared geodesic spray with Grassmann products.The practical formula depends on compatible representatives of quotient points.
  • Distance over gradients: Equation (17) is exact only for a specific endpoint representation; arbitrary retractions make it an approximation, while exponential-map updates preserve compatible representations.DOG requires distances between iterates, so first-order geodesic approximations are treated as practically acceptable.

4 Algorithms

The proposed TTN optimizers instantiate adaptive, distance-over-gradients, and MUON-style updates on either the total or quotient space. RADAM and RDOG use node-wise block operations, while RMUON orthogonalizes matrix-shaped gradient components.

  • RADAM: RADAM applies Riemannian ADAM to TTNs on both the total space and quotient space.Quotient optimization uses representatives, horizontal update directions, and the projector-vector-transport pair specified for the quotient.
  • RADAM: Total-space iterations ignore the quotient structure, whereas quotient iterations operate implicitly through horizontal lifts of quotient gradients.The projector choice determines which geometry the iteration follows.
  • RDOG: RDOG maintains a separate distance tracker for each TTN core tensor rather than one distance for the entire network.It uses node-wise distances derived from the total and quotient geodesic-distance proxies.
  • RDOG: In experiments, RDOG uses Frobenius distance at the root and Grassmann- or Stiefel-based node distances elsewhere, with projectors selected by optimization space.The same total-space versus quotient-space distinction applies to RDOG’s projector.
  • RMUON: Two Riemannian MUON variants are introduced: RMUON combines MUON updates with momentum, while another feeds MUON updates into a DOG-like scheme.Both are presented as TTN optimizers.
  • RMUON: MUON-style directions are formed by reshaping TTN tensors into matrices and orthogonalizing node-wise gradients through polar decomposition.These orthogonalized directions are compatible with the quotient geometry.

5 Numerical Experiments

The experiments compare stochastic Riemannian optimizers with ADAM on Fashion-MNIST, CIFAR10, and Imagenette, including their behavior under compression. Riemannian methods achieve comparable predictive accuracy while preserving numerical stability for downstream orthogonalization-dependent tasks.

  • Experimental setup: The numerical evaluation implements the proposed TTN optimizers with PyTorch automatic differentiation on random minibatches and a fixed batch size of 256.Initial parameters use an unsupervised construction algorithm, while learning rates and step-size estimates are selected by grid search.
  • Optimizer implementation: POGO is used componentwise as a cheaper, GPU-friendly approximation to polar retraction that stays close to the orthogonal manifold.The experiments omit a detailed treatment of Riemannian algorithms that ignore the TTN quotient structure.
  • Fashion-MNIST: After 30 epochs on Fashion-MNIST, all optimizers reach approximately 89% test accuracy, while RMUONDOG achieves the lowest training loss.RADAM, ADAM, and RMUON follow closely in training loss; RDOG converges most slowly but retains comparable final accuracy.
  • Fashion-MNIST: Riemannian optimizers support heavy Fashion-MNIST TTN compression without significant accuracy loss, unlike ADAM-trained models.ADAM causes the TTN norm to grow and moves the model away from its orthogonal initialization as training proceeds.
  • CIFAR10 and Imagenette: The CIFAR10 and Imagenette experiments use a hybrid CNN-TTN architecture in which a shared CNN processes image patches into 64-dimensional embeddings for the TTN classifier.The CNN feature extractor is optimized jointly with the TTN, with image augmentation applied during training.
  • CIFAR10 and Imagenette: On CIFAR10 and Imagenette, Riemannian optimizers match ADAM in test accuracy, reaching just above 80% and just below 80%, respectively, while ADAM runs faster.ADAM’s downstream compression damage is smaller on CIFAR10 but renders the deeper Imagenette model incompressible.

6 Conclusion

The paper develops stochastic optimization methods for orthogonal TTNs and evaluates them in a hybrid CNN-TTN framework. Riemannian optimizers achieve convergence and generalization comparable to ADAM while producing more numerically stable models for downstream tasks, though efficiency and the causes of unconstrained-optimization instabilities remain open issues.

  • Contributions: The work develops a theoretical framework for stochastic optimization on the manifold of orthogonal TTNs for machine-learning applications.
  • Contributions: The hybrid CNN-TTN framework preserves tensor-network explainability and compressibility while enabling efficient processing of large-scale images.
  • Findings: Across three image-classification tasks, stochastic Riemannian optimizers are comparable to unconstrained ADAM in convergence and generalization while learning more numerically stable downstream models.
  • Limitations and outlook: Riemannian optimization adds computational overhead, although approximate GPU-accelerated POGO retraction reduces it; further efficiency improvements remain necessary.
  • Limitations and outlook: Stochastic optimization with adaptive TTN bond dimensions remains an open direction because changing bond dimensions alters the manifold structure.

A Geodesics of the TTN quotient

The TTN quotient’s geodesics can be represented through geodesics on a Cartesian product of Grassmann manifolds, with shared horizontal representations in the TTN parameter space.

  • Geometric construction: The correspondence follows because horizontal distributions can coincide even though the quotient and Grassmann product have different vertical spaces.The shared horizontal structure lets both quotient constructions use the same horizontal parametrization in the TTN space.
  • Geodesic correspondence: The TTN quotient and the Grassmannian product share geodesics of the same length through a common horizontal curve in the TTN parameter space.This correspondence maps quotient geodesics to Grassmann-product geodesics while preserving their representation and length.
  • Geometric limitation: The equivalence is one-way because not every Grassmann-product geodesic lifts to a TTN geodesic satisfying the full multilinear-rank constraints.Thus, Grassmannian geodesics do not all have corresponding curves in the TTN quotient.
  • Exponential map: The quotient exponential map is expressed through Grassmannian exponential maps applied to the non-root tensor components.A horizontal lift at a representative θ supplies the componentwise tangent directions used by the exponential map.

B Convergence of RADAGRAD

RADAGRAD convergence on the TTN quotient is established under bounded-gradient, retraction, and geodesic-convexity assumptions, with objective values converging in expectation as iterations grow.

  • Algorithm and setting: RADAGRAD operates on minibatch loss oracles over the TTN quotient and uses horizontal quotient gradients in its updates.Algorithm 5 uses the horizontal projector ProjH and a retraction that descends to the quotient exponential map.
  • Scope of analysis: The convergence proof is presented for RADAGRAD, while extensions to more advanced adaptive algorithms are indicated as possible.The paper focuses on RADAGRAD to emphasize quotient-specific features more clearly.
  • Strong convexity: For twice differentiable strongly convex losses, a quotient point mapping to a critical point has positive-definite Riemannian Hessian.The result follows by restricting the loss to the fixed-hierarchical-rank tensor manifold and comparing its Hessian on the tangent space.
  • Convergence assumptions: Under assumptions (A1–A3), the regret of RADAGRAD on the TTN quotient is bounded.The assumptions require bounded gradients, a suitable exponential-map retraction, and a compact geodesically convex region containing the iterates and oracle minimizer.
  • Convergence guarantee: For K → ∞, the regret result implies convergence of objective values in expectation.This is the stated asymptotic consequence of the quotient regret bound.

C Additional numerical results

Numerical comparisons on CIFAR10 show that quotient-respecting optimizers are slightly more accurate but can have different per-step costs, while Riemannian optimization stabilizes tensor norms.

  • Quotient optimization: Optimizers operating directly on T perform slightly worse in final loss and test accuracy than optimizers respecting the quotient T/G.The comparison is illustrated on CIFAR10.
  • Runtime: RADAM and RMUON are slightly faster on T/G because ProjH is cheaper than Proj.The speed difference is attributed to the lower cost of horizontal projection.
  • Runtime: DOG-like algorithms are faster on T because evaluating q_k is cheaper than evaluating d_G.Their relative speed advantage comes from cheaper distance-related evaluation on the parameter manifold.
  • Numerical stability: Riemannian optimizers such as RADAM keep the tensor norm stable, whereas exploding norms indicate numerical problems for optimizers acting on the tensor norm.Norm stability matters for downstream operations such as compression.
Loading 2609.00870v1…