Riemannian Optimization in Modular Systems

TL;DR

This paper introduces a Riemannian geometric framework for understanding and optimizing modular systems, including neural networks, providing new theoretical insights, efficient metrics, and stability guarantees.

Contribution

It combines Riemannian geometry with backpropagation, introduces a recursive layerwise metric, and develops a framework for analyzing modular system optimization.

Findings

Layerwise Riemannian metrics can be computed efficiently.

The framework offers stability guarantees for modular system optimization.

Applicable to neural networks and biological systems.

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(\kappa^2 L/(\xi \mu \sqrt{n}))$ where $\kappa$ and $L$ are Lipschitz constants, $\mu$ is the mass matrix scale, and $\xi$ 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.