Riemannian Optimization in Modular Systems

2026-03-03

AI Summary

This work addresses the lack of theoretical understanding regarding the geometric and dynamical nature of backpropagation in the joint optimization of modular systems such as neural networks. By modeling backpropagation as a constrained optimization problem, the paper introduces a composable “Riemannian module” framework that integrates Riemannian geometry, optimal control, and theoretical physics. The approach recursively defines inter-layer Riemannian metrics and leverages the Woodbury identity for computational efficiency. Coupled with nonlinear contraction theory, it provides rigorous guarantees on convergence and stability. The proposed method achieves stable optimization at a computational complexity lower than that of classical natural gradient methods, yielding an algorithmic stability bound of $O(\kappa^2 L/(\xi \mu \sqrt{n}))$, and is applicable not only to neural networks but also to broader classes of sequential modular systems.

Technology Category

• Machine Learning: Optimization
• Search and Optimization: Non-convex Optimization
• Natural Language Processing: Learning & Optimization for NLP

Application Category

• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• Responsible Web: Human-perceived consequences of algorithmic deployment on the web
• Economics, Online Markets and Human Computation: Incentives in network design for Web infrastructures and ecosystems

Abstract

Understanding how systems built out of modular components can be jointly optimized is an important problem in biology, engineering, and machine learning. The backpropagation algorithm is one such solution and has been instrumental in the success of neural networks. Despite its empirical success, a strong theoretical understanding of it is lacking. Here, we combine tools from Riemannian geometry, optimal control theory, and theoretical physics to advance this understanding. We make three key contributions: First, we revisit the derivation of backpropagation as a constrained optimization problem and combine it with the insight that Riemannian gradient descent trajectories can be understood as the minimum of an action. Second, we introduce a recursively defined layerwise Riemannian metric that exploits the modular structure of neural networks and can be efficiently computed using the Woodbury matrix identity, avoiding the $O(n^3)$ cost of full metric inversion. Third, we develop a framework of composable ``Riemannian modules'' whose convergence properties can be quantified using nonlinear contraction theory, providing algorithmic stability guarantees of order $O(κ^2 L/(ξμ\sqrt{n}))$ where $κ$ and $L$ are Lipschitz constants, $μ$ is the mass matrix scale, and $ξ$ bounds the condition number. Our layerwise metric approach provides a practical alternative to natural gradient descent. While we focus here on studying neural networks, our approach more generally applies to the study of systems made of modules that are optimized over time, as it occurs in biology during both evolution and development.

Problem

• modular systems
• joint optimization
• backpropagation
• Riemannian optimization
• theoretical understanding

Innovation

• Riemannian optimization
• modular systems
• layerwise metric
• natural gradient descent
• nonlinear contraction theory