Covariant Gradient Descent (2504.05279v2)

Abstract: We present a manifestly covariant formulation of the gradient descent method, ensuring consistency across arbitrary coordinate systems and general curved trainable spaces. The optimization dynamics is defined using a covariant force vector and a covariant metric tensor, both computed from the first and second statistical moments of the gradients. These moments are estimated through time-averaging with an exponential weight function, which preserves linear computational complexity. We show that commonly used optimization methods such as RMSProp, Adam and AdaBelief correspond to special limits of the covariant gradient descent (CGD) and demonstrate how these methods can be further generalized and improved.

• The paper introduces Covariant Gradient Descent as a unifying framework that leverages gradient statistics to define a dynamic metric tensor.
• It formulates optimization dynamics using exponential moving averages of gradient moments to capture full gradient covariances and generalize standard methods.
• Experiments on the Rosenbrock function and neural networks show that CGD can outperform traditional optimizers in convergence speed and loss reduction.

The paper "Covariant Gradient Descent" (2504.05279) introduces a unified geometric framework for gradient-based optimization methods called Covariant Gradient Descent (CGD). The core idea is to formulate the optimization dynamics in a way that is consistent across different coordinate systems, drawing concepts from differential geometry.

The standard gradient descent update step for parameters $q^\mu$ is typically written as $$\dot{q}^\mu = -\gamma \frac{\partial H}{\partial q^\mu}$$, where $H$ is the loss function and $\gamma$ is the learning rate. The authors note that in a general coordinate system and a potentially curved parameter space, a covariant formulation is $$\dot{q}^\mu = -\gamma g^{\mu\nu}(q) \frac{\partial H}{\partial q^\nu}$$, where $g^{\mu\nu}(q)$ is the inverse metric tensor of the parameter space.

CGD generalizes this further by proposing the dynamics:

$\dot{q}^\mu(t) = -\gamma g^{\mu\nu}(t) F_\nu(t)$

Here, the metric tensor $g_{\mu\nu}(t)$ and the covariant force vector $F_\nu(t)$ are not fixed tensors of the parameter space geometry, but rather dynamic quantities constructed from the statistical moments of the gradients of the loss function, $\frac{\partial H}{\partial q^\mu}$. This accounts for the emergent geometry influenced by the fluctuations in gradients during the learning process.

The statistical moments are estimated using temporal averaging, specifically exponential moving averages. The first moment is $$M^{(1)}_\mu(t) = \left \langle \frac{\partial H}{\partial q^\mu} \right \rangle(t)$$ and the second moment is $$M^{(2)}_{\mu\nu}(t) = \left \langle \frac{\partial H}{\partial q^\mu} \frac{\partial H}{\partial q^\nu} \right \rangle(t)$$. These averages are computed iteratively, for example:

$$\langle f\rangle(t) = \frac{1}{1+\tau} f(t) +\frac{\tau}{1+\tau} \langle{f}\rangle(t-1)$$

where $\tau$ is the averaging timescale. The parameters of this discrete-time averaging correspond to the $\beta$ parameters in optimizers like Adam.

The paper shows that common optimization methods can be seen as special cases of CGD by defining the force and metric based on these moments:

• Gradient Descent (GD): $F_\mu$ is the instantaneous gradient, $g_{\mu\nu}$ is the identity matrix (Euclidean metric). This corresponds to $F_\mu = M^{(1)}_\mu$ with $\tau_1=0$, and $g_{\mu\nu} = \delta_{\mu\nu}$.
• Stochastic Gradient Descent (SGD): Similar to GD but with gradient noise, often incorporating momentum which can be captured by using $M^{(1)}_\mu$ with $\tau_1 > 0$.
• RMSProp and Adam: These methods use the element-wise square root of the gradient's squared magnitude (variance) to scale the updates. In the CGD framework, this corresponds to setting the metric tensor as diagonal, $g_{\mu\nu} = G(\diag(M^{(2)}))_{\mu\nu}$, where $G(x) = 1/\sqrt{\epsilon + x}$. RMSProp uses $\tau_1=0, \tau_2 > 0$, while Adam uses $\tau_1 > 0, \tau_2 > 0$.

A key contribution is highlighting that RMSProp and Adam only utilize the diagonal elements of the second moment matrix $M^{(2)}$, which represent the variances of individual gradient components. The off-diagonal elements, representing covariances between gradient components, are ignored. The CGD framework naturally allows for using the full covariance matrix $M^{(2)}$ to define the metric tensor $g_{\mu\nu} = G(M^{(2)})_{\mu\nu}$, potentially capturing richer geometric information about the loss landscape, such as correlations and anisotropy. The paper explores a general form $G(M^{(2)}) = (\epsilon I + M^{(2)})^a$.

Numerical experiments on the 2D Rosenbrock function and a simple neural network trained for multiplication demonstrate that CGD variants, particularly those using the full covariance matrix ("full CGD"), can outperform standard optimizers like SGD, RMSProp, and Adam in terms of convergence speed and final loss. The analysis of the covariance matrix eigenvalues during training shows that they decrease significantly, suggesting that the CGD optimizer learns to navigate the parameter space by focusing updates in an effectively lower-dimensional subspace, adapting to the emerging anisotropy of the loss landscape.

The paper concludes by discussing practical challenges, primarily the computational cost of computing and inverting the full covariance matrix for high-dimensional models. It suggests future work on efficient approximations or low-rank representations of the metric tensor. Theoretically, the emergence of a curved metric in the parameter space is linked to concepts of emergent geometry in physical systems, suggesting a deeper connection between machine learning and physics.

In summary, CGD provides a unifying, geometrically-principled view of gradient-based optimization, emphasizing the role of gradient statistics in defining an emergent metric tensor that guides the learning dynamics. By explicitly incorporating gradient covariances (via the full $M^{(2)}$ matrix), CGD offers a promising direction for developing more effective adaptive optimization methods, despite current computational limitations for very large models.