Learning Discretized Neural Networks under Ricci Flow
In this paper, we study Discretized Neural Networks (DNNs) composed of
low-precision weights and activations, which suffer from either infinite or
zero gradients due to the non-differentiable discrete function during training.
Most training-based DNNs in such scenarios employ the standard Straight-Through
Estimator (STE) to approximate the gradient w.r.t. discrete values. However,
the use of STE introduces the problem of gradient mismatch, arising from
perturbations in the approximated gradient. To address this problem, this paper
reveals that this mismatch can be interpreted as a metric perturbation in a
Riemannian manifold, viewed through the lens of duality theory. Building on
information geometry, we construct the Linearly Nearly Euclidean (LNE) manifold
for DNNs, providing a background for addressing perturbations. By introducing a
partial differential equation on metrics, i.e., the Ricci flow, we establish
the dynamical stability and convergence of the LNE metric with the $L^2$-norm
perturbation. In contrast to previous perturbation theories with convergence
rates in fractional powers, the metric perturbation under the Ricci flow
exhibits exponential decay in the LNE manifold. Experimental results across
various datasets demonstrate that our method achieves superior and more stable
performance for DNNs compared to other representative training-based methods.
Jun Chen, Hanwen Chen, Mengmeng Wang, Guang Dai, Ivor W. Tsang, Yong Liu
# Learning Discretized Neural Networks under Ricci Flow

###### Abstract

In this paper, we consider Discretized Neural Networks (DNNs) consisting of low-precision weights and activations, which suffer from either infinite or zero gradients due to the non-differentiable discrete function in the training process. In this case, most training-based DNNs employ the standard Straight-Through Estimator (STE) to approximate the gradient w.r.t. discrete values. However, the STE gives rise to the problem of gradient mismatch, due to the perturbations of the approximated gradient. To address this problem, this paper reveals that this mismatch can be viewed as a metric perturbation in a Riemannian manifold through the lens of duality theory. Further, on the basis of the information geometry, we construct the Linearly Nearly Euclidean (LNE) manifold for DNNs as a background to deal with perturbations. By introducing a partial differential equation on metrics, i.e., the Ricci flow, we prove the dynamical stability and convergence of the LNE metric with the \(L^{2}\)-norm perturbation. Unlike the previous perturbation theory whose convergence rate is the fractional powers, the metric perturbation under the Ricci flow can be exponentially decayed in the LNE manifold. The experimental results on various datasets demonstrate that our method achieves better and more stable performance for DNNs than other representative training-based methods.

D 1
Footnote 1: 1. In this paper, the continuous weight is relative to the neural network (its data type is full-precision). And the discretized weight is relative to the discretized neural network (its data type is low-precision).

D 1
Footnote 1: 1. In this paper, the continuous weight is relative to the neural network (its data type is full-precision). And the discretized weight is relative to the discretized neural network (its data type is low-precision).

## 1 Introduction

Discretized neural networks (DNNs) (Courbariaux et al., 2016; Li et al., 2016; Zhu et al., 2016) have been proven to be efficient in computing, which can significantly reduce computational complexity, storage space, power consumption, resources, etc (Chen et al., 2020). Considering a Discretized neural network (DNN) which is able to be well-trained, based on the standard chain rule, the gradient w.r.t. the continuous weight1\(\mathbf{w}\) propagating through a discrete function \(Q(\cdot)\), i.e., \(\frac{\partial L}{\partial\mathbf{w}}=\frac{\partial L}{\partial Q(\mathbf{w})}\frac{ \partial Q(\mathbf{w})}{\partial\mathbf{w}}\), suffers from either infinite or zero deriva
tives. Its root is that the derivative \(\partial Q(\mathbf{w})/\partial\mathbf{w}\) can not be calculated. In the backward, one can obtain the gradient \(\partial L/\partial Q(\mathbf{w})\), but need to update the continuous weight \(\mathbf{w}\) via the gradient \(\partial L/\partial\mathbf{w}\). Since the gradient \(\partial L/\partial\mathbf{w}\) can not be obtained explicitly, one needs the derivative \(\partial Q(\mathbf{w})/\partial\mathbf{w}\) as a bridge to calculate \(\partial L/\partial\mathbf{w}\) via the chain rule.

In order to address the problem of either infinite or zero gradients caused by the non-differentiable discrete function, Hinton (2012) first proposed the Straight-Through Estimator (STE) that yields an immediate connection between \(\partial L/\partial\mathbf{w}\) and \(\partial L/\partial Q(\mathbf{w})\) in backpropagation such that bypassing the derivative \(\partial Q(\mathbf{w})/\partial\mathbf{w}\). The definition of STE was then given by Bengio et al. (2013), which can be summarized as: the gradient w.r.t. the discretized weight can be approximated by the gradient w.r.t. the continuous weight with clipping, as shown in Figure 1(a). Subsequently, Courbariaux et al. (2016) applied STE to binarized neural networks and provided an approximated gradient as follows:

\[\frac{\partial L}{\partial\mathbf{w}}=\frac{\partial L}{\partial\operatorname{ sign}(\mathbf{w})}\cdot\mathbb{I},\ \ \ \text{where}\ \mathbb{I}:=\left\{\begin{array}{ll}1&\text{if}\quad|\mathbf{w}|\leq 1 \\ 0&\text{otherwise}\end{array}\right., \tag{1}\]

where \(\mathbb{I}\) is the indicator function. Note that \(Q(\cdot)\) will degenerate to \(\operatorname{sign}(\cdot)\) which is equal to \(+1\) for \(\mathbf{w}\geq 0\) and \(-1\) otherwise in binarized neural networks. STE had been successfully

Figure 1: Comparison of STE and our method. We denote the arrows and points as gradients and weights, respectively. In particular, when a point falls on the grid point, it means that the weight is discretized at this time. In the forward of DNNs, the continuous weight \(\mathbf{w}\) is mapped to a discrete weight \(Q(\mathbf{w})\) via a discrete function. In the backward, the gradient is propagated from \(\partial L/\partial Q(\mathbf{w})\) to \(\partial L/\partial\mathbf{w}\). (a) The STE simply copies the gradient, i.e., \(\partial L/\partial\mathbf{w}=\partial L/\partial Q(\mathbf{w})\). (b) Our method, on the other hand, matches the gradient by introducing the proper metric \(g_{\mathbf{w}}\), i.e., \(\partial L/\partial\mathbf{w}=g_{\mathbf{w}}^{-1}\partial L/\partial Q(\mathbf{w})\), while taking into account the gradient mismatch caused by STE in a Riemannian manifold.
implemented in the training of binarized neural networks, and it was further extended to ternary neural networks (Li et al., 2016) and arbitrary bit-width discretized neural networks (Zhou et al., 2016).

On the other hand, Non-STE methods consist in all techniques that do not rely on STE, e.g., (Hou et al., 2016; Bai et al., 2018; Leng et al., 2018). However, the learning process of Non-STE methods depends heavily on hyper-parameters (Chen et al., 2019), such as weight partition portion in each iteration (Zhou et al., 2017) and penalty setting in tuning (Leng et al., 2018). Hence, STE methods are widely used in DNNs rather than Non-STE methods owing to their simplicity and versatility.

However, STE introduced into DNNs, inevitably gives rise to the problem of _gradient mismatch_: the gradient w.r.t. the continuous weight is not strictly equal to the gradient w.r.t. the discretized weight when \(|\mathbf{w}|\leq 1\)(Chen et al., 2019), which compromises the training stability of DNNs (Cai et al., 2017; Liu et al., 2018; Qin et al., 2020). Furthermore, the formula of STE tells us that this problem is able to be alleviated by modifying the gradient \(\partial L/\partial\mathbf{w}\).

Zhou et al. (2016) firstly proposed to transform the weight \(\mathbf{w}\) into the new one \(\tilde{\mathbf{w}}\) via

\[\tilde{\mathbf{w}}=\frac{\tanh(\mathbf{w})}{\max(|\tanh(\mathbf{w})|)}.\]

By discretizing the new weight \(\tilde{\mathbf{w}}\), the STE then acts on \(\tilde{\mathbf{w}}\). During back-propagation, the gradient can be further computed as follows

\[\frac{\partial L}{\partial\mathbf{w}}=\frac{\partial L}{\partial Q(\tilde{\mathbf{w} })}\frac{1-\tanh^{2}(\mathbf{w})}{\max(|\tanh(\mathbf{w})|)}.\]

The purpose of the authors is to manually redefine the indicator function \(\mathbb{I}\) as \(\frac{1-\tanh^{2}(\mathbf{w})}{\max(|\tanh(\mathbf{w})|)}\) such that the function \(\frac{1-\tanh^{2}(\mathbf{w})}{\max(|\tanh(\mathbf{w})|)}\) provides a smooth transition to avoid abrupt clipping of the indicator function near \(\pm 1\). It is remarkable that Chen et al. (2019) proposed to learn \(\partial L/\partial\mathbf{w}\) by a neural network, e.g., fully-connected layers or LSTM (Sak et al., 2014). Their specific approach is to use neural networks as a shared meta quantizer \(M_{\psi}\) parameterized by \(\psi\) across layers to replace the gradient via:

\[\frac{\partial L}{\partial\mathbf{w}}=M_{\psi}\left(\frac{\partial L}{\partial Q( \mathbf{w})},\overline{\mathbf{w}}\right)\frac{\partial\overline{\mathbf{w}}}{\partial \mathbf{w}},\]

where \(\overline{\mathbf{w}}\) is the weight from the meta quantizer. With the input of the gradient \(\partial L/\partial Q(\mathbf{w})\), the meta quantizer will output a new gradient to match \(\partial L/\partial\mathbf{w}\) by updating the weight \(\overline{\mathbf{w}}\) in the training process.

Recently, Ajanthan et al. (2021) formulated the binarization of neural networks as a constrained optimization problem by introducing a mirror descent framework (Beck and Teboulle, 2003) to perform gradient descent in the dual space (unconstrained space) with gradients computed in the primal space (discrete space). In particular, by projecting the primal space \(\mathbf{w}\) into the dual space \(\tilde{\mathbf{w}}=\tanh(\beta_{k}\mathbf{w})\), the gradient can be yielded as

\[\frac{\partial L}{\partial\mathbf{w}}=\frac{\partial L}{\partial\tilde{\mathbf{w}}} \left(1-\tanh^{2}(\beta_{k}\mathbf{w})\right).\]
As the hyper-parameter \(\beta_{k}\to\infty\), \(\tilde{\mathbf{w}}\) gradually approaches \(\text{sign}(\mathbf{w})\) until the corresponding neural network is fully binarized with an adaptive mirror map.

However, Zhou et al. (2016) obtained the new weight by manually setting the function \(\tanh\), which can only scale the gradient as a whole and do not fundamentally alleviate the gradient mismatch. On the other hand, although Chen et al. (2019) automatically matched the gradient by learning a new neural network (a meta quantizer), the additional errors are also engendered in the gradient propagation due to the meta quantizer, leading to further intensifying the problem of gradient mismatch. Subsequently, Ajanthan et al. (2021) bypassed the problem of gradient mismatch because the derivative \(\partial\tilde{\mathbf{w}}/\partial\mathbf{w}=\left(1-\tanh^{2}(\beta_{k}\mathbf{w})\right)\) can be calculated directly, which implies that this method does not maintain discrete weights during the training. Therefore, the problem of gradient mismatch still remains to be solved.

### Contributions

In this work, we regard the gradient mismatch between \(\partial L/\partial\mathbf{w}\) and \(\partial L/\partial Q(\mathbf{w})\) as a perturbation phenomenon between these two gradients. By introducing the framework of Riemannian geometry in Figure 1(b), the gradient mismatch is further viewed as a metric perturbation in a Riemannian manifold (Section 2.2) through the lens of duality theory (Amari and Nagaoka, 2000; Amari, 2016). As a partial differential equation on metrics, the Ricci flow (Sheridan and Rubinstein, 2006), is introduced, the metric perturbation can be exponentially decayed in theory such that the problem of gradient mismatch is theoretically solved. The main contributions of this paper are summarized in the following four aspects:

* We propose the LNE manifold endowed with the LNE metric, which is a special form of Ricci-flat metrics in essence. According to the information geometry (Amari, 2016), we construct LNE manifolds for neural networks as a background to deal with perturbations.
* We reveal the stability of LNE manifolds under the Ricci-DeTurck flow with the \(L^{2}\)-norm perturbation on the basis of the relationship between the Ricci-DeTurck flow and the Ricci flow. In this way, any Ricci flow starting close to the LNE metric exists for all time and converges to the LNE metric. Furthermore, unlike the previous perturbation theory whose convergence rate is the fractional powers (\(t^{-3/2}\)), the metric perturbation under the Ricci flow can give rise to exponentially decaying in the LNE manifold (\(e^{-t}\)), further bringing about theoretical assurance for effectively solving the problem of gradient mismatch.
* Based on the appealing characteristics of LNE manifolds under Ricci Flow, a novel DNNs with the acceptable complexity, i.e., Ricci Flow Discretized Neural Network (RF-DNN) is developed by calculating the Ricci curvature in such a way that the selection of coordinate systems is related to the input transformations of neural networks. In essence, the discrete Ricci flow is employed to overcome the problem of gradient mismatch in traditional DNNs.
* The experiments are implemented on several classification benchmark datasets and network structures. Experimental results demonstrate the effectiveness of our geometric method RF-DNN compared with other representative training-based methods.
### Overall Organization

This paper is organized as follows. In Section 2, we introduce the motivation and Ricci flow. According to the geometric structure measured by the LNE divergence, we deduce the corresponding LNE manifold for neural networks in Section 3. The stability of LNE manifolds under the Ricci-DeTurck is proved in Section 4. In Section 5, we calculate the approximated gradient in the LNE manifold to avoid solving the inverse of the LNE metric. In Section 6, we present how to introduce discrete Ricci flow into DNNs and yield the corresponding algorithm. The experimental results and ablation studies for RF-DNNs are presented in Section 7. Section 8 concludes the entire paper. Proofs are provided in the Appendices.

The Ricci flow on Ricci-flat metrics is known in the literature to be stable for \(C^{0}\) perturbations in the \(L^{\infty}\)-norm (Section 2.4). Based on a Bregman divergence (Bregman, 1967), the LNE metric, a special form of Ricci-flat metrics, is introduced in neural networks via the LNE divergence (Theorem 10). Stability of LNE manifolds under the Ricci-DeTurck flow is then proved (Corollary 15 and Theorem 16). A discretization of the Ricci flow is therefore proposed, that leads to a practical algorithm (RF-DNNs, Algorithm 2).

## 2 Motivation and Formulation

### Background

We start with the basic background for feed-forward DNNs that will be used throughout the paper. This background is on the basis of the work (Martens and Grosse, 2015). And we list the important notations in Appendix B.

A neural network can be regarded as a function to transform the input \(\mathbf{a}_{0}\) into the output \(\mathbf{a}_{l}\) through a series of \(l\) layers. For the \(i\)-th layer (\(i\in\{1,2,\ldots,l\}\)), we denote \(\mathbf{W}_{i}\) as the weight matrix, \(\mathbf{s}_{i}\) as the vector of these weighted sum and \(\mathbf{a}_{i}\) as the vector of output (also known as the activation). Each layer receives vectors of a weighted sum of the input from the previous layer and calculates their output via a nonlinear function.

For a DNN, we need to add a discrete function \(Q(\cdot)\) to discretize the weight matrix \(\mathbf{W}_{i}\) and the activation vector \(\mathbf{a}_{i}\). Furthermore, we mark the discretized weight matrix as \(\mathbf{W}_{i}=Q(\mathbf{W}_{i})\) and the discretized activation vector as \(\mathbf{\hat{a}}_{i}=Q(\mathbf{a}_{i})\). Then the feed-forward of DNNs at each layer is given as follows:

\[\mathbf{s}_{i} =\mathbf{\hat{W}}_{i}\mathbf{\hat{a}}_{i-1} \tag{2}\] \[\mathbf{a}_{i} =f\odot\mathbf{s}_{i}\] \[\mathbf{\hat{a}}_{i} =Q(\mathbf{a}_{i})\]

where \(f\) is a nonlinear (activation) function, \(\odot\) represents the element-wise multiplication. We vectorize \(\mathbf{\hat{W}}_{i}\) as \(\text{vec}(\mathbf{\hat{W}}_{i})\) and stack their columns together, where \(\text{vec}(*)\) is the operator that vectorizes a matrix as a column vector. Note that the vectorized weights in each layer before and after discretization are expressed as \(\mathbf{w}\) and \(\mathbf{\hat{w}}\), respectively. We define the new discretized parameter vector as \(\mathbf{\hat{\xi}}=\left[\text{vec}\left(\mathbf{\hat{W}}_{1}\right)^{\top},\text{ vec}\left(\mathbf{\hat{W}}_{2}\right)^{\top},\ldots,\text{vec}\left(\mathbf{\hat{W}}_{l} \right)^{\top}\right]^{\top}\), where we ignore the bias vector for brevity. Similarly, we can denote the parameter vector as \(\mathbf{\xi}=\left[\text{vec}\left(\mathbf{W}_{1}\right)^{\top},\text{vec}\left(\mathbf{ W}_{2}\right)^{\top},\ldots,\text{vec}\left(\mathbf{W}_{l}\right)^{\top}\right]^{\top}\) where \(\mathbf{\hat{\xi}}=Q(\mathbf{\xi})\).
In frequentist statistics, we represent the loss function as the negative log likelihood w.r.t. discretized parameter \(\hat{\mathbf{\xi}}\), where the input \(\mathbf{a}_{0}\) can be observed. The way minimizing the loss function is to maximize the likelihood:

\[L(\mathbf{y},\mathbf{z})=-\log p\left(\mathbf{y}|\mathbf{z},\hat{\mathbf{\xi}}\right) \tag{3}\]

where \(\mathbf{y}\) is the target and \(\mathbf{z}\) is the prediction calculated by DNNs. Note that \(p\left(\mathbf{y}|\mathbf{z},\hat{\mathbf{\xi}}\right)\) is the probability density function defined by the conditional probability distribution \(P_{\mathbf{y}|\mathbf{z}}\left(\hat{\mathbf{\xi}}\right)\) of DNNs. As for the back-propagation of DNNs, the details are given in later sections.

### Motivation

We consider that the root of the gradient mismatch is that there is a perturbation phenomenon between \(\partial L/\partial\mathbf{w}\) and \(\partial L/\partial Q(\mathbf{w})\), i.e.,

\[\frac{\partial L}{\partial\mathbf{w}}=\mathcal{P}\left(\frac{\partial L}{ \partial Q(\mathbf{w})}\right)\bigg{|}_{O\left(\frac{\partial L}{\partial Q(\bm {w})}\right)} \tag{4}\]

where the input of the perturbation function \(\mathcal{P}\) is the gradient \(\partial L/\partial Q(\mathbf{w})\), and the corresponding perturbation range is \(O\left(\partial L/\partial Q(\mathbf{w})\right)\). If the perturbation range \(O\left(\partial L/\partial Q(\mathbf{w})\right)\) can be significantly eliminated or decayed, then there is an elegant solution to the problem of gradient mismatch. In the framework of perturbation theory of the linear space (Kato, 2013), the rate of convergence for perturbations is the fractional powers, e.g., \(t^{-3/2}\).

Figure 2: The overview of the theoretical ideas in this paper.
Inspired by the mirror descent 2 framework (Beck and Teboulle, 2003), one can map the parameter from the primal space to the dual space, and then calculate the gradient in the dual space. When the Riemannian metric structures are introduced by means of information geometry (Amari, 2016), naturally 3, the gradient mismatch is conclusively viewed as a metric perturbation in a Riemannian manifold. Specifically, with the help of Definition 1, it is obvious that the gradient \(\partial L/\partial\mathbf{w}\) will be rewritten as \(\tilde{\partial}L/\tilde{\partial}\mathbf{w}\) in a Riemannian manifold. Therefore, based on Equation (4), the problem of gradient mismatch can be defined as

Footnote 2: Mirror descent induces non-Euclidean structure by solving iterative optimization problems using different proximity functions (Bubeck et al., 2015).

Footnote 3: Natural gradient descent selects the steepest descent along a Riemannian manifold by multiplying the standard gradient by the inverse of the metric tensor (Amari, 1998).It is worth mentioning that mirror descent and natural gradient descent are proven to be equivalent to each other (Raskutti and Mukherjee, 2015), which implies that mirror descent is the steepest descent direction along the Riemannian manifold corresponding to the choice of Bregman divergence.

\[\frac{\tilde{\partial}L}{\tilde{\partial}\mathbf{w}}=g_{\mathbf{w}}^{-1}\frac{ \partial L}{\partial Q(\mathbf{w})} \tag{5}\]

where the perturbation item is implied in the metric \(g_{\mathbf{w}}\) and the term \(\frac{\partial L}{\partial Q(\mathbf{w})}\) can be considered as the conventional gradient \(\partial L(\mathbf{w})\) in Definition 1. Then the metric perturbation phenomenon emerges, and the perturbation at this time is referred to the deviation from the original metric. In this way, we present the generalization of STE in a Riemannian manifold which will degenerate into the standard STE when the Riemannian metric \(g\) returns to the Euclidean metric \(\delta\).

**Definition 1**: _(Amari, 1998) The steepest descent direction of \(L(\mathbf{w})\) in a Riemannnian manifold, i.e., the **natural gradient descent,** is given by_

\[\tilde{\partial}L(\mathbf{w})=g_{\mathbf{w}}^{-1}\partial L(\mathbf{w})\]

_where \(g^{-1}=(g^{ij})\) is the inverse of the metric \(g=(g_{ij})\) and \(\partial L(\mathbf{w})\) is the conventional gradient,_

\[\partial L(\mathbf{w})=\left[\frac{\partial L(\mathbf{w})}{\partial w_{1}},\cdots, \frac{\partial L(\mathbf{w})}{\partial w_{n}}\right]^{\top}.\]

Subsequently, a key question is what kind of manifolds we need to construct to deal with metric perturbations naturally and easily. Or what manifold is a "good" manifold in the presence of perturbations. In practice, it seems that general relativity gives an excellent example in nature to deal with small gravitational perturbations under the framework of a Riemannian manifold (Wald, 2010). To treat the approximation in which gravity is "weak", the spacetime metric is nearly flat in the context of general relativity, which is enough for most cases, except for phenomena dealing with gravitational collapse and the large scale structure of the universe. By assuming that the deviation \(\gamma_{ij}\) of the actual spacetime metric \(g_{ij}=\eta_{ij}+\gamma_{ij}\) from a flat metric \(\eta_{ij}\) is "small", the linearized gravity is presented to approximately calculate the gravity in general relativity. An adequate definition of "smallness" in this context is that the components of \(\gamma_{ij}\) are much smaller than \(1\) in the global inertial coordinate system of \(\eta_{ij}\). Such a linearly nearly flat metric greatly simplifies
the calculation of "weak" gravity, and the manifolds constructed with such metrics are considered to be sufficient for approximating the manifold with perturbations.

Similarly, in this paper, by regarding the Euclidean metric \(\delta_{ij}\) as the flat metric \(\eta_{ij}\), we can naturally define the Linearly Nearly Euclidean (LNE) metric. If we can construct such a metric for DNNs, then a metric perturbation can be handled in the background of LNE manifolds. Motivated by the natural gradient descent (Amari, 2016) that connects neural networks with the Riemannian metric, LNE metrics can be mathematically constructed in neural networks. Specifically, our aim is to first introduce a convex function to derive the LNE divergence with the help of Bregman divergence (Bregman, 1967), and then construct the LNE metric by introducing the LNE divergence into a neural network 4. Consequently, the LNE metric appears in the gradient of the LNE manifold, similar to Definition 1. Finally, based on the constructed manifold for DNNs, the only remaining problem is how to rapidly decay the metric perturbation by employing a geometric tool, i.e., Ricci flow.

Footnote 4: The part from a convex function to the LNE divergence is in the mirror descent framework. Then, the part from the LNE divergence to the LNE metric introduces the idea of information geometry.

In addition, a series of proofs about the stability illustrates that the Ricci flow can decay the metric perturbation in the cases of Ricci-flat metrics. Therefore, as long as we can prove that a small perturbation of the LNE metric under the Ricci flow is decayed, the metric perturbation can be alleviated such that a theoretical solution for the problem of gradient mismatch in the training of DNNs is further given. Unlike the previous perturbation theory whose the convergence rate is the fractional powers, the metric perturbation under the Ricci flow can be exponentially decayed in the LNE manifold. In Figure 2, we give the overview of the theoretical ideas to facilitate sorting out the solution steps.

### Ricci Flow

**Definition 2**: _(Sheridan and Rubinstein, 2006) A **Riemannian metric** on a smooth manifold \(\mathcal{M}\) is a smoothly-varying inner product on the tangent space \(T_{p}\mathcal{M}\) at each point \(p\in\mathcal{M}\), i.e., a (0,2)-tensor which is symmetric and positive-definite at each point of \(\mathcal{M}\). One will usually write \(g\) for a Riemannian metric, and \(g_{ij}\) for it coordinate representation. A manifold together with a Riemannian metric, \((\mathcal{M},g)\), is called a **Riemannian manifold**._

The concept of Ricci flow was first proposed by Hamilton (Hamilton et al., 1982) on the Riemannian manifold \(\mathcal{M}\), based on Definition 2, of a time-dependent metric \(g(t)\). Given the initial metric \(g_{0}\), the Ricci flow is a partial differential equation that evolves the metric tensor:

\[\frac{\partial}{\partial t}g(t) =-2\operatorname{Ric}(g(t)) \tag{6}\] \[g(0) =g_{0}\]

where \(\operatorname{Ric}\) denotes the Ricci curvature tensor and the detailed definition can be found in Appendix A. The purpose of Ricci flow is to prove Thurston's Geometrization Conjecture and Poincare Conjecture, which is to evolve the metric towards certain geometric structures and topological properties (Sheridan and Rubinstein, 2006).
**Corollary 3**: _(Sheridan and Rubinstein, 2006) The Ricci flow is **strongly parabolic** if there exists \(\delta>0\) such that for all covectors \(\varphi\neq 0\) and all (symmetric 5) \(h_{ij}=\frac{\partial}{\partial t}g_{ij}(t)\neq 0\), the principal symbol of the differential operator \(-2\operatorname{Ric}\) satisfies_

Footnote 5: The Riemannian metric \(g_{ij}\) is always symmetric based on Definition 2. Hence, \(h_{ij}=\frac{\partial}{\partial t}g_{ij}(t)\) is required to be symmetric.

\[[-2\operatorname{Ric}](\varphi)(h)_{ij}h^{ij}=g^{pq}\left(\varphi_{p}\varphi_{ q}h_{ij}+\varphi_{i}\varphi_{j}h_{pq}-\varphi_{q}\varphi_{i}h_{jp}-\varphi_{q} \varphi_{j}h_{ip}\right)h^{ij}>\delta\varphi_{k}\varphi^{k}h_{rs}h^{rs}\]

_where \(h^{ij}\) is the inverse of \(h_{ij}\)._

**Theorem 4**: _(Ladyzhenskaia et al., 1988) Suppose that \(u(t):\mathcal{M}\times[0,T)\to\mathcal{E}\) is a time-dependent section of the vector bundle \(\mathcal{E}\) where \(\mathcal{M}\) is a Riemannian manifold. If the system of the Ricci flow is strongly parabolic at \(u_{0}\) where \(u_{0}=u(0):\mathcal{M}\to\mathcal{E}\), then there exists a unique solution on the time interval \([0,T)\)._

Combined with Corollary 3 and Theorem 4, one knows whether there exists a unique solution of Ricci flow for a short time by checking whether it is strongly parabolic. However, if we choose \(h_{ij}=\varphi_{i}\varphi_{j}\), it is clear that the left hand side of the inequality in Corollary 3 is 0, thus the inequality can not hold. Since the inequality can not always be satisfied, Ricci flow is not always strongly parabolic, which makes us unable to assure the existence of the solution according to Theorem 4. In what follows, the non-poarabolic root is analyzed and the solution can be found on the basis of the relationship between the Ricc flow and the Ricci-DeTurck flow. Now, it is possible to understand which parts have an impact on its non-parabolic nature by the linearization of the Ricci curvature tensor. We define the linearization of the Ricci curvature as \(\mathcal{D}[\operatorname{Ric}]\) so that

\[\mathcal{D}[\operatorname{Ric}]\left(\frac{\partial}{\partial t}g_{ij}(t) \right)=\frac{\partial}{\partial t}\operatorname{Ric}(g_{ij}(t)).\]

**Lemma 5**: _The linearization of \(-2\operatorname{Ric}\) can be rewritten as 6_

Footnote 6: In this paper, we use the Einstein summation convention (for example, \((AB)_{i}^{j}=A_{i}^{k}B_{k}^{j}\)). When the same index appears twice in one term, once as an upper index and the other time as a lower index, summation is automatically taken over this index even without the summation symbol.

\[\begin{split}&\mathcal{D}[-2\operatorname{Ric}](h)_{ij}=g^{pq} \nabla_{p}\nabla_{q}h_{ij}+\nabla_{i}V_{j}+\nabla_{j}V_{i}+O(h_{ij})\\ &\text{where }\quad V_{i}=g^{pq}\left(\frac{1}{2}\nabla_{i}h_{pq}- \nabla_{q}h_{pi}\right)\text{and }\ h_{ij}=\frac{\partial}{\partial t}g_{ij}(t).\end{split} \tag{7}\]

**Proof** The proofs can be found in Appendix C.1. \(\blacksquare\)

By carefully observing Lemma 5, the impact on the non-parabolic of Ricci flow comes from the terms \(V_{i}\) and \(V_{j}\)(Sheridan and Rubinstein, 2006), instead of the term \(g^{pq}\nabla_{p}\nabla_{q}h_{ij}\). On the other hand, the term \(O(h_{ij})\) will have no contributions to the principal symbol of \(-2\operatorname{Ric}\), so we can ignore it in this problem. Next, we try to eliminate the impact of the non-parabolic nature on the Ricci flow.
Using a time-dependent diffeomorphism \(\varphi(t):\mathcal{M}\to\mathcal{M}\) (with \(\varphi(0)=\mathrm{id}\)), the pullback metrics \(g(t)\) can be expressed as

\[g(t)=\varphi^{*}(t)\bar{g}(t), \tag{8}\]

which satisfies the Ricci flow equation, where \(\varphi^{*}(t)\) is the pullback through \(\varphi(t)\). The above formula yields the new metric \(\bar{g}(t)\) via the pullback, and the terms \(V_{i}\) and \(V_{j}\) can be reparameterized by choosing \(\varphi(t)\) to form the Ricci-DeTurck flow (w.r.t. \(\bar{g}(t)\)) that is strongly parabolic. Furthermore, the solution is followed by the DeTurck Trick (DeTurck, 1983) that has a time-dependent reparameterization of the manifold:

\[\begin{split}\frac{\partial}{\partial t}\bar{g}(t)& =-2\operatorname{Ric}(\bar{g}(t))-\mathcal{L}_{\frac{\partial \varphi(t)}{\partial t}\bar{g}(t)}\\ \bar{g}(0)&=\bar{g}_{0},\end{split} \tag{9}\]

See Appendix C.2 for details. Thus, the Ricci-DeTurck flow has a unique solution for a short time. For the long time behavior, please refer to Appendix C.3.

### Literature

For the Riemannian \(n\)-dimensional manifold \((\mathcal{M}^{n},g)\) that is isometric to the Euclidean \(n\)-dimensional space \((\mathbb{R}^{n},\delta)\), Schnurer et al. (Schnurer et al., 2007) have showed the stability of Euclidean space under the Ricci flow for a small \(C^{0}\) perturbation. Koch et al. (Koch and Lamm, 2012) have given the stability of the Euclidean space along with the Ricci flow in the \(L^{\infty}\)-norm. Moreover, for the decay of the \(L^{\infty}\)-norm on Euclidean space, Appleton (Appleton, 2018) has provided a proof using a different method.

On the other hand, for a Ricci-flat metric with small perturbations, Guenther et al. (Guenther et al., 2002) proved that such metrics would converge under Ricci flow. Considering the stability of integrable and closed Ricci-flat metrics, Sesum (Sesum, 2006) has proved that the convergence rate is exponential because the spectrum of the Lichnerowicz operator is discrete. Furthermore, Deruelle et al. (Deruelle and Kroncke, 2021) have proved that an asymptotically locally Euclidean Ricci-flat metric is dynamically stable under the Ricci flow, with the \(L^{2}\cap L^{\infty}\) perturbation on non-flat and non-compact Ricci-flat manifolds. Our work is also discussed on Ricci-flat manifolds.

## 3 Neural Networks in LNE Manifolds

The aim of this section is to build a LNE manifold via the information geometry (Amari, 2016) such that paving the way for the introduction of Ricci flow. Specifically, we first introduce a convex function (Equation (14)) to derive the LNE divergence (Theorem 10) with the help of Bregman divergence (Definition 8), and then construct the LNE metric (Equation (16)) by introducing the LNE divergence into a neural networks. Consequently, the LNE metric then appears in the steepest descent gradient (Lemma 11) of the LNE manifold. Of course, the mirror descent algorithm (Beck and Teboulle, 2003) can also equivalently construct the relationship between divergences and gradients, but it lacks the geometric meaning (manifold and metric) that is exactly what we need.
### Neural Network Manifold

A neural network is composed of a large number of neurons connected with each other. The set of all such networks forms a manifold, where the weights represented by the neuron connections can be regarded as the coordinate system.

For the \(n\times n\) matrix, all such matrices form an \(n^{2}\)-dimensional manifold. Specifically, the matrix forms the \(n(n+1)/2\)-dimensional manifold, i.e., a submanifold embedded in the manifold of all the matrices when this matrix is symmetric and positive-definite (Sheridan and Rubinstein, 2006). Note that such a matrix is described as the metric in the geometry theory. Since symmetric and positive-definite matrices have many advantages, which lead to wide implementations in real-world applications, our method will also construct such symmetric and positive-definite matrices (metrics) with some appealing geometric characteristics in neural networks.

**Remark 6**: _Comparing straight lines in Euclidean space, geodesics are the straightest possible lines that we can draw in a Riemannian manifold. Given a geodesic, there exists a unique non-Euclidean coordinate system. Once the curved coordinate system is selected in a Riemannian manifold, the symmetric and positive-definite metric is also given based on Definition 2. That geometry-based metric can describe the properties of manifolds, such as curvature (Helgason, 2001)._

### Euclidean Space and Divergence

From the viewpoint of the information geometry, the metric can be deduced by the divergence satisfying the certain criteria (Basseville, 2013), which is summarized as Definition 7.

**Definition 7**: _(Amari, 2016) \(D[P:Q]\) is called a **divergence** when it satisfies the following criteria:_

_(1) \(D[P:Q]\geq 0\),_

_(2) \(D[P:Q]=0\) when and only when \(P=Q\),_

_(3) When \(P\) and \(Q\) are sufficiently close to each other, and their coordinates are denoted by \(\boldsymbol{\xi}_{P}\) and \(\boldsymbol{\xi}_{Q}=\boldsymbol{\xi}_{P}+d\boldsymbol{\xi}\) respectively, the Taylor expansion of the divergence can be written as_

\[D[\boldsymbol{\xi}_{P}:\boldsymbol{\xi}_{Q}]=\frac{1}{2}\sum_{i,j}g_{ij}( \boldsymbol{\xi}_{P})d\xi_{i}d\xi_{j}+O(|d\boldsymbol{\xi}|^{3}), \tag{10}\]

_and the Riemannian metric \(g_{ij}\) is positive-definite, depending on \(\boldsymbol{\xi}_{P}\)._

When \(P\) and \(Q\) are sufficiently close where these two points are expressed in coordinates as \(\boldsymbol{\xi}_{P}\) and \(\boldsymbol{\xi}_{Q}\) (column vectors), the square of an infinitesimal distance \(ds^{2}\) between them by using Definition 7 can be defined as:

\[ds^{2}=2D[\boldsymbol{\xi}_{P}:\boldsymbol{\xi}_{Q}]=\sum_{i,j}g_{ij}( \boldsymbol{\xi}_{P})d\xi_{i}d\xi_{j} \tag{11}\]

where \(d\boldsymbol{\xi}\) denotes a sufficiently small coordinate variation between the coordinates \(\boldsymbol{\xi}_{P}\) and \(\boldsymbol{\xi}_{Q}\). Here we can ignore the higher order term \(O(|d\boldsymbol{\xi}|^{3})\) for the sake of simplicity and convenience. A manifold \(\mathcal{M}\) is said to be Riemannian when a postive-definite metric \(g_{ij}\) is
defined on \(\mathcal{M}\) and the square of the local distance between two nearby points \(\mathbf{\xi}_{P}\) and \(\mathbf{\xi}_{Q}\) is given by Equation (11). A divergence \(D\) provides \(\mathcal{M}\) with a Riemannian structure.

Using an orthonormal Cartesian coordinate system in the Euclidean space, the Euclidean divergence is naturally defined as a half of the square of the Euclidean distance between two nearby points \(\mathbf{\xi}\) and \(\mathbf{\xi}^{\prime}\)

\[D_{E}[\mathbf{\xi}:\mathbf{\xi}^{\prime}]=\frac{1}{2}\sum_{i}(\xi_{i}-\xi_{i}^{\prime} )^{2}. \tag{12}\]

The Riemannian metric \(g_{ij}\) degenerates into the Euclidean metric \(\delta_{ij}\) in this case, so that

\[ds^{2}=2D_{E}[\mathbf{\xi}:\mathbf{\xi}+d\mathbf{\xi}]=\sum_{i}(d\xi)^{2}=\sum_{i,j}\delta _{ij}d\xi_{i}d\xi_{j}. \tag{13}\]

Obviously, the Euclidean divergence satisfies the criteria of divergence. Note that the Euclidean metric \(\delta_{ij}\) is the identity matrix \(\mathbf{I}\) where we still retain the notation of metrics for the sake of the representation habit in the geometry theory.

### LNE Manifold and Divergence

Recall that 7, in general relativity (Wald, 2010), the complete Riemannian manifold \((\mathcal{M},g)\) that is endowed with a linearly nearly flat spacetime metric, is considered to solve the Newtonian limit by the linearized gravity. The form of this metric is \(g_{ij}=\eta_{ij}+\gamma_{ij}\) where \(\eta_{ij}\) represents a flat Minkowski metric whose background is special relativity and \(\gamma_{ij}\) is smallness. In practice, this is an excellent theory of small gravitational perturbations when gravity is "weak".

Footnote 7: The link between the LNE manifold and general relativity can be found in Section 2.2.

Similarly, we give a metric \(g_{ij}=\delta_{ij}+\gamma_{ij}\) in Riemannian \(n\)-manifold where \(\delta_{ij}\) represents a flat Euclidean metric. An adequate definition of "smallness" in this context is that the components of \(\gamma_{ij}\) are much smaller than 1 in the global inertial coordinate system of \(\delta_{ij}\). Therefore, we can systematically develop the LNE metric to cope with small perturbations.

#### 3.3.1 Convex Function and Bregman Divergence

In order to construct the LNE manifold endowed with the LNE metric in the neural network, according to Definition 7, we can introduce the divergence to obtain the expression of the LNE metric, similar to the relationship between Euclidean metric and divergence.

The premise of constructing a divergence is to find a suitable convex function. There are many applications for different convex functions in physics and optimization (Bubeck et al., 2015). Thus, we can introduce a nonlinear function \(\Phi(\mathbf{\xi})\) of coordinates \(\mathbf{\xi}\) as the convex function with a certain geometric structure to satisfy the needs of constructing the LNE divergence. When considering a differentiable function, it is said to be convex if and only if its Hessian is positive-definite

\[H(\mathbf{\xi})=\left(\frac{\partial^{2}}{\partial\xi_{i}\partial\xi_{j}}\Phi(\bm {\xi})\right).\]
**Definition 8**: _(Bregman, 1967) The **Bregman divergence**\(D_{B}[\boldsymbol{\xi}:\boldsymbol{\xi}^{\prime}]\) is defined as the difference between a convex function \(\Phi(\boldsymbol{\xi})\) and its tangent hyperplane \(z=\Phi(\boldsymbol{\xi}^{\prime})+(\boldsymbol{\xi}-\boldsymbol{\xi}^{\prime}) \nabla\Phi(\boldsymbol{\xi}^{\prime})\), depending on the Taylor expansion at the point \(\boldsymbol{\xi}^{\prime}\):_

\[D_{B}[\boldsymbol{\xi}:\boldsymbol{\xi}^{\prime}]=\Phi(\boldsymbol{\xi})-\Phi (\boldsymbol{\xi}^{\prime})-(\boldsymbol{\xi}-\boldsymbol{\xi}^{\prime}) \nabla\Phi(\boldsymbol{\xi}^{\prime}).\]

By drawing a tangent hyperplane touching it at the point \(\boldsymbol{\xi}^{\prime}\)

\[z=\Phi(\boldsymbol{\xi}^{\prime})+(\boldsymbol{\xi}-\boldsymbol{\xi}^{\prime })\nabla\Phi(\boldsymbol{\xi}^{\prime}),\]

we can write the distance between the convex function \(\Phi(\boldsymbol{\xi})\) and the tangent hyperplane \(z\) as the Bregman divergence. Considering it as a function of two points \(\boldsymbol{\xi}\) and \(\boldsymbol{\xi}^{\prime}\), we can easily know that it satisfies the criteria of divergence. Since \(\Phi(\boldsymbol{\xi})\) is convex, the graph of \(\Phi(\boldsymbol{\xi})\) is always above the tangent hyperplane, touching it at \(\boldsymbol{\xi}^{\prime}\). The graph describing their relationship is shown in Figure 3, where \(z\) is the vertical axis of the graph.

**Remark 9**: _In fact, a divergence is derived from a convex function in the form of the Bregman divergence. By choosing different convex functions, for example, we can easily obtain the Euclidean divergence from the Bregman divergence, by defining the convex function: \(\Phi(\boldsymbol{\xi})=1/2\sum_{i}\xi_{i}^{2}\) in a Euclidean space. Besides, given the convex function: \(\Phi(\boldsymbol{\xi})=\sum_{i}p(\boldsymbol{x},\xi_{i})\log p(\boldsymbol{x}, \xi_{i})\) where \(\sum_{i}p(\boldsymbol{x},\xi_{i})=\sum_{i}p(\boldsymbol{x},\xi_{i}^{\prime})= \int p(\boldsymbol{x},\boldsymbol{\xi})d\boldsymbol{\xi}=\int p(\boldsymbol{x },\boldsymbol{\xi}^{\prime})d\boldsymbol{\xi}^{\prime}=1\) is satisfied, the Bregman divergence is the same as the KL divergence._

#### 3.3.2 LNE Divergence and Gradient

Similar to the Bregman divergence with a convex function, we try to construct a new convex function to obtain the LNE divergence on which the LNE metric emerges naturally on the basis of Definition 7. Inspired by the previous work (Ajanthan et al., 2021), we propose a

Figure 3: The divergence \(D[\boldsymbol{\xi}:\boldsymbol{\xi}^{\prime}]\) is viewed as the distance between the convex function \(\Phi(\boldsymbol{\xi})\) and the tangent hyperplane \(z\). And \(z\) depends on the point \(\boldsymbol{\xi}^{\prime}\) at which the supporting hyperplane with normal vector \(\boldsymbol{n}=\nabla\Phi(\boldsymbol{\xi}^{\prime})\) is defined.
novel convex function, which satisfies symmetry and is able to geometrically construct an easy-to-compute metric with the linearly nearly Euclidean nature:

\[\Phi(\mathbf{\xi})=\sum_{i}\frac{1}{\tau^{2}}\log\left(\cosh(\tau\xi_{i})\right) \tag{14}\]

where \(\tau\) is a constant parameter that is used to control how linearly close to Euclidean structure.

**Theorem 10**: _By introducing a convex function \(\Phi\) defined by Equation (14) into Definition 8, the **LNE divergence** between two points \(\mathbf{\xi}\) and \(\mathbf{\xi}^{\prime}\) yields_

\[\begin{split} D_{LNE}[\mathbf{\xi}^{\prime}:\mathbf{\xi}]& =\sum_{i}\left[\frac{1}{\tau^{2}}\log\frac{\cosh(\tau\xi_{i}^{ \prime})}{\cosh(\tau\xi_{i})}-\frac{1}{\tau}(\xi_{i}^{\prime}-\xi_{i})\tanh( \tau\xi_{i})\right]\\ &\approx\frac{1}{2}\sum_{i,j}\left\{\delta_{ij}-\left[\tanh(\tau \mathbf{\xi})\tanh(\tau\mathbf{\xi})^{\top}\right]_{ij}d\xi_{i}d\xi_{j}\right\}.\end{split} \tag{15}\]

The detailed proofs can be found in Appendix E.1.

Comparing with Definition 7, we know that the **LNE metric** corresponding to the LNE divergence is

\[\begin{split}& g(\mathbf{\xi})=\delta_{ij}-\left[\tanh(\tau\mathbf{ \xi})\tanh(\tau\mathbf{\xi})^{\top}\right]_{ij}\\ &=\begin{bmatrix}1-\tanh(\tau\xi_{1})\tanh(\tau\xi_{1})&\cdots&-\tanh(\tau\xi_{1}) \tanh(\tau\xi_{n})\\ \vdots&\ddots&\vdots\\ -\tanh(\tau\xi_{n})\tanh(\tau\xi_{1})&\cdots&1-\tanh(\tau\xi_{n})\tanh(\tau\xi_{n}) \end{bmatrix}.\end{split} \tag{16}\]

Based on Section 3.1, we are able to employ the parameters of a neural network to construct the LNE metric (for a neural network, the coordinate \(\mathbf{\xi}\) is the parameter vector). As a result, we can describe the neural network in the LNE manifold, as measured by the LNE divergence based on Theorem 10. Naturally, the steepest descent gradient in the LNE manifold is given by Lemma 11, which is similar to the natural gradient in Definition 1.

**Lemma 11**: _The steepest descent gradient measured by the LNE divergence is defined as_

\[\tilde{\partial}_{\mathbf{\xi}}=g(\mathbf{\xi})^{-1}\partial_{\mathbf{\xi}}=\left[\delta -\tanh(\tau\mathbf{\xi})\tanh(\tau\mathbf{\xi})^{\top}\right]^{-1}\partial_{\mathbf{\xi}}. \tag{17}\]

The proofs can be found in Appendix E.2.

On the constructed LNE manifold, we can introduce the Ricci flow to decay the metric perturbation w.r.t. the LNE metric.
### Convergence Analysis

We can now describe the mirror descent via the Bregman divergence associated to \(\Phi\). Let \(\mathcal{C}\subset\mathbb{R}^{n}\) be a convex open set such that \(\mathcal{X}\subset\bar{\mathcal{C}}\) (\(\bar{\mathcal{C}}\) is the closure of \(\mathcal{C}\)), and \(\mathcal{X}\cap\mathcal{C}\neq\emptyset\). When \(\mathbf{x}^{0}\in\arg\min_{\mathbf{x}\in\mathcal{X}\cap\mathcal{C}}\Phi(\mathbf{x})\) is the initial point, for iteration \(k\geq 0\) and step size \(\eta\), the update of the mirror descent can be written as:

\[\begin{split}\nabla\Phi(\mathbf{y}^{k+1})&=\nabla\Phi( \mathbf{x}^{k})-\eta\partial^{k}\\ \mathbf{x}^{k+1}&=\arg\min_{\mathbf{x}\in\mathcal{X}\cap \mathcal{C}}D_{B}[\mathbf{x},\mathbf{y}^{k+1}]\end{split} \tag{18}\]

where \(\partial^{k}\in\partial F(\mathbf{x}^{k})\) and \(\mathbf{y}^{k+1}\in\mathcal{C}\).

**Theorem 12**: _(Bubeck et al., 2015) Let \(\Phi\) be a mirror map \(\rho\)-strongly convex on \(\mathcal{X}\cap\mathcal{C}\) w.r.t. \(\|\cdot\|\). Let \(R^{2}=\sup_{\mathbf{x}\in\mathcal{X}\cap\mathcal{C}}\Phi(\mathbf{x})-\Phi(\mathbf{x}^{0})\), and \(F\) be convex and \(L\)-Lipschitz w.r.t. \(\|\cdot\|\). Then mirror descent with \(\eta=\frac{R}{L}\sqrt{\frac{2\rho}{t}}\) satisfies_

\[F\left(\frac{1}{t}\sum_{k=0}^{t-1}\mathbf{x}^{k}\right)-F(\mathbf{x}^{*})\leq RL\sqrt {\frac{2}{\rho t}}.\]

In this paper, the mirror descent is obtained by taking Equation (14) on \(\mathcal{C}=\mathbb{R}^{n}\). The function \(\Phi(\mathbf{\xi})\) is a mirror map \(1\)-strongly convex w.r.t. \(\|\cdot\|_{2}\), and furthermore the associated Bregman divergence is given by Theorem 10. In that case, the mirror descent is equivalent to projected subgradient descent, and the rate of convergence is \(O(1/\sqrt{t})\), which is consistent with the standard mirror descent (Theorem 12). Of course, this gives the same convergence as (Ajanthan et al., 2021).

It is easy to verify that the projection w.r.t. this Bregman divergence on the simplex \(\Delta_{n}=\{\xi\in\mathbb{R}^{n}:\sum_{i=1}^{n}\xi_{i}=1\}\) amounts to a simple renormalization (Bubeck et al., 2015). For \(\mathcal{X}=\Delta_{n}\), one has \(\mathbf{x}^{0}=(1/n,\cdots,1/n)\) and \(R^{2}=\log(\cosh n)\). In that case, the mirror descent with the Equation (14) achieves a rate of convergence of order \(\sqrt{\frac{\log(\cosh n)}{t}}\).

## 4 Evolution of LNE Manifolds under Ricci Flow

In this section, we will focus on LNE metrics under Ricci flow, and prove that the evolution of LNE manifolds can achieve good performances in terms of stability (all time), i.e., the Ricci flow can exponentially decay the \(L^{2}\)-norm perturbation to the LNE metric.

### Relationship between LNE Metrics and Ricci Flow

To facilitate the handling of the metric perturbation, we have given the LNE metric in Equation (16), a special form of Ricci-flat metrics (Guenther et al., 2002; Deruelle and Kroncke, 2021), which means that the Ricci-flat metrics on noncompact manifolds are LNE. In particular, the constructed metric \(g(\mathbf{\xi})\) in Equation (16) is LNE, which is equivalent to either \(g_{0}\) under the Ricci flow or \(\bar{g}_{0}\) under the Ricci-DeTurck flow since \(g_{0}\) and \(\bar{g}_{0}\) are diffeomorphic 8 to each other on the basis of Equation (8).
The full statement of the LNE metric is the linearly nearly Euclidean Ricci-flat metric in Definition 13 given on the basis of the previous work (Deruelle and Kroncke, 2021).

A complete Riemannian \(n\)-manifold \((\mathcal{M}^{n},\bar{g}_{0})\) is said to be LNE with one end of order \(\iota>0\) if there exists a compact set \(K\subset\mathcal{M}\), a radius \(r\), a point \(x\) in \(\mathcal{M}\) and a diffeomorphism satisfying \(\phi:\mathcal{M}\backslash K\rightarrow(\mathbb{R}^{n}\backslash B(x,r))/SO(n)\). Note that \(B(x,r)\) is the ball and \(SO(n)\) is a finite group acting freely on \(\mathbb{R}^{n}\backslash\{0\}\). Then

\[\left|\partial^{k}(\phi_{*}\gamma)\right|_{\delta}=O(r^{-\iota-k})\ \ \forall k\geq 0 \tag{19}\]

holds on \((\mathbb{R}^{n}\backslash B(x,r))/SO(n)\). The LNE metric \(\bar{g}_{0}\) can be linearly decomposed into a form containing the Euclidean metric \(\delta\) and the deviation \(\gamma\):

\[\bar{g}_{0}=\delta+\gamma. \tag{20}\]

### Finite Time Existence

In this subsection, we consider the linear stability and integrability of the LNE manifold \((\mathcal{M}^{n},g_{0})\). Fortunately, similar to the previous proofs (Koiso, 1983; Besse, 2007), we can know that \((\mathcal{M}^{n},g_{0})\) is integral and linearly stable based on Definition 29 and Definition 30.

Due to the diffeomorphic relationship between the Ricci flow and the Ricci-DeTurck flow, we introduce a metric perturbation for the Ricci-DeTurck flow, and Equation (9) can be further reformulated as follows:

\[\begin{split}\frac{\partial}{\partial t}\bar{g}(t)& =-2\operatorname{Ric}(\bar{g}(t))-\mathcal{L}_{\frac{\partial \epsilon(t)}{\partial t}}\bar{g}(t)\\ \bar{g}(0)&=\bar{g}_{0}+d\end{split} \tag{21}\]

where \(d=\bar{g}(0)-\bar{g}_{0}\) is a metric perturbation deviated from the LNE metric \(\bar{g}_{0}\). Corollary 15 guarantees the finite time existence of the Ricci-DeTurck flow w.r.t. \(L^{2}\)-norm perturbations, and then provides the necessary premise for proving its all time convergence.

(Bamler, 2010, 2011) Let \((\mathcal{M}^{n},\bar{g}_{0})\) be a complete Ricci-flat \(n\)-manifold. If \(\bar{g}(0)\) is a metric satisfying \(\|\bar{g}(0)-\bar{g}_{0}\|_{L^{\infty}}<\epsilon\) where \(\epsilon>0\), then there exists a constant \(C<\infty\) and a unique Ricci-DeTurck flow \(\bar{g}(t)\) that satisfies

\[\|\bar{g}(t)-\bar{g}_{0}\|_{L^{\infty}}<C\|\bar{g}(0)-\bar{g}_{0}\|_{L^{ \infty}}<C\cdot\epsilon. \tag{22}\]

Let \((\mathcal{M}^{n},\bar{g}_{0})\) be the LNE \(n\)-manifold. For a Ricci-DeTurck flow \(\bar{g}(t)\) on a maximal time interval \(t\in[0,T)\) and \(k\in\mathbb{N}\), there exists constants \(C_{k}=C_{k}(\bar{g}_{0},T)\) such that

\[\|\nabla^{k}d(t)\|_{L^{2}}\leq C_{k}\cdot t^{-k/2} \tag{23}\]

where \(d(t)=\bar{g}(t)-\bar{g}_{0}\) is the time-evolving perturbation.

ProofWhen Lemma 14 satisfies in a finite time, according to (Deruelle and Kroncke, 2021), the Ricci-DeTurck flow with the LNE metric w.r.t. the \(L^{2}\)-norm perturbation exists.
### All Time Convergence for \(L^{2}\)-norm Perturbations

In order to prove the all time stability of LNE metrics under Ricci-DeTurck flow, we need to construct \(\bar{g}_{0}(t)\) that is a family of Ricci-flat reference metrics with \(\frac{\partial}{\partial t}\bar{g}_{0}(t)=O((\bar{g}(t)-\bar{g}_{0}(t))^{2})\). Let

\[\mathcal{F}=\left\{\bar{g}(t)\in\mathcal{M}^{n}\;\middle|\;2\operatorname{Ric} (\bar{g}(t))+\mathcal{L}_{\frac{\partial\varphi(t)}{\partial t}}\bar{g}(t)=0\right\}\]

be the set of stationary points under the Ricci-DeTurck flow. Then, we are able to establish a manifold via an \(L^{2}\)-neighbourhood \(\mathcal{U}\) of integral \(\bar{g}_{0}\) in the space of metrics

\[\tilde{\mathcal{F}}=\mathcal{F}\cap\mathcal{U}. \tag{24}\]

For all \(\bar{g}\in\tilde{\mathcal{F}}\), the terms \(\operatorname{Ric}(\bar{g}(t))=0\) and \(\mathcal{L}_{\frac{\partial\varphi(t)}{\partial t}}\bar{g}(t)=0\) hold individually, based on the previous work (Deruelle and Kroncke, 2021). In this way, \(d(t)-d_{0}(t)=\bar{g}(t)-\bar{g}_{0}(t)\) holds because we write \(d_{0}(t)=\bar{g}_{0}(t)-\bar{g}_{0}\).

According to Theorem 3.1, it is obvious that the \(L^{2}\)-norm metric perturbation w.r.t. the LNE metric can be dynamically decayed by the Ricci-DeTurck flow. See the details in Appendix D.

**Theorem 3.1**: _Let \((\mathcal{M}^{n},\bar{g}_{0})\) be the LNE \(n\)-manifold which is linearly stable and integrable. For any metric \(\bar{g}(t)\in\mathcal{B}_{L^{2}}(\bar{g}_{0},\epsilon_{2})\) where a constant \(\epsilon_{2}>0\), there is a complete Ricci-DeTurck flow \((\mathcal{M}^{n},\bar{g}(t))\) starting from \(\bar{g}(t)\) converging to the LNE metric \(\bar{g}(\infty)\in\mathcal{B}_{L^{2}}(\bar{g}_{0},\epsilon_{1})\) where \(\epsilon_{1}\) is a small enough constant._

**Proof** By Lemma 3.1, we have a constant \(\epsilon_{2}>0\) such that \(d(t)\in\mathcal{B}_{L^{2}}(0,\epsilon_{2})\) holds. By Lemma 3.1 (in the second step) and Corollary 3.1 (in the third step), we can obtain

\[\left\|d_{0}(T)\right\|_{L^{2}}\leq C\int_{1}^{T}\left\|\frac{ \partial}{\partial t}d_{0}(t)\right\|_{L^{2}}\mathrm{d}t\] \[\quad\leq C\int_{1}^{T}\left\|\nabla^{\bar{g}_{0}}\left(d(t)-d_{ 0}(t)\right)\right\|_{L^{2}}^{2}\mathrm{d}t\] \[\quad\leq C\left\|d(1)-d_{0}(1)\right\|_{L^{2}}^{2}\leq C\|d(1) \|_{L^{2}}^{2}\leq C\cdot\left(\epsilon_{2}\right)^{2}.\]

Furthermore, we can obtain from the above formulas

\[\left\|d(T)-d_{0}(T)\right\|_{L^{2}}\leq\left\|d(1)-d_{0}(1)\right\|_{L^{2}} \leq C\cdot\epsilon_{2}.\]

By the triangle inequality, we get

\[\left\|d(T)\right\|_{L^{2}}\leq C\cdot\left(\epsilon_{2}\right)^{2}+C\cdot \epsilon_{2}.\]

Followed by Corollary 3.1 and Lemma 3.1, \(T\) should be pushed further outward, i.e.,

\[\lim_{t\to+\infty}\sup\left\|\frac{\partial}{\partial t}d_{0}(t)\right\|_{L^ {2}}\leq\lim_{t\to+\infty}\sup\left\|\nabla^{\bar{g}_{0}}\left(d(t)-d_{0}(t) \right)\right\|_{L^{2}}^{2}=0.\]

Thus, \(\bar{g}(t)\) will converge to \(\bar{g}(\infty)=\bar{g}_{0}+d_{0}(\infty)\) as \(t\) approaches \(+\infty\) based on the elliptic regularity. In other words, \(d(t)-d_{0}(t)\) will converge to \(0\) as \(t\) goes to \(+\infty\) w.r.t. all Sobolev norms (Minerbe, 2009),

\[\lim_{t\to+\infty}\left\|d(t)-d_{0}(t)\right\|_{L^{2}}\leq\lim_{t\to+\infty}C \left\|\nabla^{\bar{g}_{0}}\left(d(t)-d_{0}(t)\right)\right\|_{L^{2}}=0.\]
Any Ricci-DeTurck flow that starting close to the LNE metric exists for all time, and the Ricci-DeTurck flow will converge to the LNE metric following with (Deruelle and Kroncke, 2021).

### Perturbation Analysis

By proving the finite time existence of the Ricci-DeTurck flow with \(L^{2}\)-norm perturbations (Corollary 15), we then prove that \(L^{2}\)-norm perturbations w.r.t. the LNE metric converges for all time under the Ricci-DeTurck flow (Theorem 16). Based on (Sesum, 2006), we further yield \(|\bar{g}(t)-\bar{g}_{0}(\infty)|<Ce^{-\epsilon_{2}t}\) such that the metric perturbation exponentially converges.

Recall that we can reparameterize \(\bar{g}(t)\) to \(g(t)=\varphi^{*}(t)\bar{g}(t)\) via the pullback. Since the perturbation entirely comes from \(\bar{g}(t)\) and has nothing to do with the time-dependent diffeomorphism \(\varphi^{*}(t)\) in essence, it also exponentially converges for \(g(t)\) under the Ricci flow when the existence of the solution of Ricci flow satisfies. And in Section 3, the metric \(g(\mathbf{\xi})=\delta_{ij}-\left[\tanh(\tau\mathbf{\xi})\tanh(\tau\mathbf{\xi})^{\top} \right]_{ij}\) we construct for the neural network is a kind of LNE metrics (based on Definition 13), so the perturbation for this metric satisfies the exponential decay under the Ricci flow. In what follows, the gradient mismatch in DNNs can be effectively overcome by employing the LNE manifolds under Ricci flow.

## 5 Discretized Neural Networks in LNE Manifolds

Up to now, we have theoretically solved the problem of gradient mismatch. Specifically, we have completed the construction of LNE manifolds for neural networks in Section 3 and exponentially decayed the metric perturbation in LNE manifolds in Section 4.

However, based on Lemma 11, there is the inverse of the LNE metric for the steepest descent gradient in the LNE manifold, which makes it difficult to be calculated in practice. Therefore, in this section, our aim is to approximate the inverse of the LNE metric, and further obtain the approximated gradient in the LNE manifold, that leads to the practical algorithm for training DNNs in the LNE manifold.

### Gradient Computation in Discretized Neural Networks

Recall that Courbariaux et al. (2016) applied STE to binarized neural networks, whose form is given in Equation (1). Then, Zhou et al. (2016) applied STE to arbitrary bit-width discretized neural networks. Similarly, the generalization of STE in discretized neural networks yields

\[\frac{\partial L}{\partial\mathbf{w}}=\frac{\partial L}{\partial Q(\mathbf{w})}\cdot \mathbb{I}. \tag{25}\]

There is still a contradiction to be solved before the LNE manifold is introduced into the DNN. Based on Lemma 11, the LNE manifold is defined on the parameter \(\mathbf{\xi}\) of all layers in a neural network. However, the gradient in back-propagation is defined layer-by-layer, i.e., on the weight \(\mathbf{w}\) of each layer, which causes that the gradient update can not be associated with the LNE manifold. Fortunately, the LNE manifold can be layer-by-layer redefined by replacing \(\mathbf{\xi}\) with \(\mathbf{w}\), which is equivalent to defining the LNE manifold for each layer. And
on the basis of Lemma 11, the steepest descent gradient is rewritten as

\[\tilde{\partial}_{\mathbf{w}}=g^{-1}(\mathbf{w})\partial_{\mathbf{w}}=\left[\delta-\tanh( \tau\mathbf{w})\tanh(\tau\mathbf{w})^{\top}\right]^{-1}\partial_{\mathbf{w}}, \tag{26}\]

which can be used for the gradient computation in DNNs, i.e.,

\[\frac{\tilde{\partial}L}{\tilde{\partial}\mathbf{w}}=\left[\delta-\tanh(\tau\mathbf{w} )\tanh(\tau\mathbf{w})^{\top}\right]^{-1}\frac{\partial L}{\partial Q(\mathbf{w})}. \tag{27}\]

Further, the proposed gradient in above means that we are in the framework of solving the problem of gradient mismatch, similar to Equation (5). Note that the metric at this time is layer-by-layer LNE. On the other hand, gradient computation involves the inverse of the LNE metric, which greatly consumes computing resources. Hence, we propose two methods for approximating the gradient of DNNs in LNE manifolds: weak approximation and strong approximation, respectively. The approximated gradient is defined as the direction in parameter space that gives the largest variation in the objective per unit of variation along the layer-by-layer LNE manifold.

### Strong Approximation

For the \(n\times n\) symmetric metric \(g(\mathbf{w})\), it can be decomposed in terms of the combination of entries \(P\) and \(A\), where \(P\) is the entries made up of the elements of the lower triangular matrix that contains \(n(n-1)/2\) real parameters and \(A\) is the entries made up of the elements of the diagonal matrix that contains \(n\) real parameters.

Our aim is to approximate the inverse of the LNE metric, and further approximate the gradient in Equation (27). Based on the universal approximation theorem (Cybenko, 1989;

Figure 4: Flow chart of strong approximation of \(g^{-1}(\mathbf{w})\). The new entries \(\tilde{P}\) and \(\tilde{A}\) produced by neural network form a matrix \(\mathbf{G}\), which will multiply by the metric \(g(\mathbf{w})\). As the loss function defined by Equation (28) decreases, the matrix \(\mathbf{G}\) can be used to approximate the inverse of the metric \(g(\mathbf{w})\).
Hornik, 1991), a continuous function on compact subsets can be approximated by a neural network with a single hidden layer and a finite number of neurons (Jejjala et al., 2020). Inspired by this work, we also introduce a multi-layer perceptron (MLP) neural network, as shown in Figure 4, to minimize the loss function

\[\tilde{L}=\|\mathbf{I}-g(\mathbf{w})\mathbf{G}\|^{2}. \tag{28}\]

Hence, the matrix \(\mathbf{G}\) can be used to strongly approximate the inverse of the metric \(g(\mathbf{w})\).

### Weak Approximation

In this subsection, we can also present a weak approximation for the inverse of the LNE metric with the efficient calculation.

**Definition 17**: _(Bhatia, 2013) For \(\mathbf{A}\in\mathcal{R}^{n\times n}\), \(\mathbf{A}\) is called **diagonally dominant** when it satisfies_

\[\big{|}a_{ii}\big{|}>\sum_{j=1,j\neq i}^{n}\big{|}a_{ij}\big{|},\ \ \ i=1,2, \ldots,n.\]

**Definition 18**: _(Bhatia, 2013) If \(\mathbf{A}\in\mathcal{R}^{n\times n}\) is a diagonally dominant matrix, then \(\mathbf{A}\) is a nonsingular matrix together, i.e., \(\mathbf{A}^{-1}\) exists._

Considering the properties of the LNE metric, we can easily make \(g^{-1}(\mathbf{w})\) satisfy Definition 17 by adjusting the parameter \(\tau\). Furthermore, the existence of the inverse of the LNE metric can be guaranteed on the basis of Definition 18. By placing a higher requirement, the LNE metric is diagonally dominant. According to Corollary 19, the weak approximation of the gradient in the LNE manifold can be calculated, giving rise to a nice feature to facilitate the fast computation of the inverse.

**Corollary 19**: _Based on Definition 17 and Definition 18, the weak approximation of the gradient in the LNE manifold is defined as_

\[\tilde{\partial}_{\mathbf{w}}=\left[\delta-\tanh(\tau\mathbf{w})\tanh(\tau\mathbf{w})^{ \top}\right]^{-1}\partial_{\mathbf{w}}\approx\left[\delta+\tanh(\tau\mathbf{w})\tanh (\tau\mathbf{w})^{\top}\right]\partial_{\mathbf{w}} \tag{29}\]

_if the LNE metric is diagonally dominant._

**Proof** Considering the inverse of the LNE metric, due to the diagonally dominant property in Definition 17 and Definition 18, we can approximate \(\left[\delta-\tanh(\tau\mathbf{w})\tanh(\tau\mathbf{w})^{\top}\right]^{-1}\) by ignoring the fourth-order small quantity \(\sum O(\rho_{a}\rho_{b}\rho_{c}\rho_{d})\), i.e.,

\[\left[\delta-\tanh(\tau\mathbf{w})\tanh(\tau\mathbf{w})^{\top}\right] \left[\delta+\tanh(\tau\mathbf{w})\tanh(\tau\mathbf{w})^{\top}\right]\] \[=\begin{bmatrix}1-\rho_{1}\rho_{1}&-\rho_{1}\rho_{2}&\cdots\\ -\rho_{2}\rho_{1}&1-\rho_{2}\rho_{2}&\cdots\\ \vdots&\vdots&\ddots\end{bmatrix}\begin{bmatrix}1+\rho_{1}\rho_{1}&\rho_{1} \rho_{2}&\cdots\\ \rho_{2}\rho_{1}&1+\rho_{2}\rho_{2}&\cdots\\ \vdots&\vdots&\ddots\end{bmatrix}\] \[=\begin{bmatrix}1-\sum O(\rho_{a}\rho_{b}\rho_{c}\rho_{d})&\rho_ {1}\rho_{2}-\rho_{1}\rho_{2}-\sum O(\rho_{a}\rho_{b}\rho_{c}\rho_{d})&\cdots \\ -\rho_{2}\rho_{1}+\rho_{2}\rho_{1}-\sum O(\rho_{a}\rho_{b}\rho_{c}\rho_{d})&1- \sum O(\rho_{a}\rho_{b}\rho_{c}\rho_{d})&\cdots\\ \vdots&\vdots&\ddots\end{bmatrix}\approx\mathbf{I}.\]
### Training

```
0: A minibatch of inputs and targets \((\mathbf{x}=\mathbf{a_{0}},\mathbf{y})\), \(\mathbf{\xi}\) mapped to \((\mathbf{W}_{1},\mathbf{W}_{2},\ldots,\mathbf{W}_{l})\), \(\mathbf{\hat{\xi}}\) mapped to \(\left(\mathbf{\hat{W}}_{1},\mathbf{\hat{W}}_{2},\ldots,\mathbf{\hat{W}}_{l}\right)\), a nonlinear function \(f\), a constant factor \(\tau\) and a learning rate \(\eta\).
0: The updated discretized parameters \(\mathbf{\hat{\xi}}\).
1: {Forward propagation}
2:for\(i=1;i\leq l;i++\)do
3: Discretize \(\mathbf{\hat{W}}_{i}=Q(\mathbf{W}_{i})\);
4: Compute \(\mathbf{s}_{i}=\mathbf{\hat{W}}_{i}\mathbf{\hat{a}}_{i-1}\);
5: Discretize \(\mathbf{\hat{a}}_{i}=Q\left(f\odot\mathbf{s}_{i}\right)\);
6:endfor
7: {Loss derivative}
8: Compute \(L=L(\mathbf{y},\mathbf{z})\);
9: Compute \(\partial_{\mathbf{a}_{l}}L=\frac{\partial L(\mathbf{y},\mathbf{z})}{\partial\mathbf{z}}|_{\bm {z}=\hat{\mathbf{a}}_{l}}\);
10: {Backward propagation}
11:for\(i=l;i\geq 1;i--\)do
12: Compute \(\partial_{\mathbf{s}_{i}}L=\partial_{\mathbf{a}_{i}}L\odot f^{\prime}(\mathbf{s}_{i})\);
13: Compute \(\partial_{\hat{\mathbf{W}}_{i}}L=(\nabla_{\mathbf{s}_{i}}L)\,\mathbf{\hat{a}}_{i-1}^{\top}\);
14: Compute \(\tilde{\partial}_{\mathbf{W}_{i}}L=g^{-1}(\mathbf{W}_{i})\partial_{\hat{\mathbf{W}}_{i}}L\) based on Equation (27);
15: Compute \(\partial_{\hat{\mathbf{a}}_{i-1}}L=\hat{\mathbf{W}}_{i}^{\top}\left(\partial_{\mathbf{s} _{i}}L\right)\);
16:endfor
17: {The parameters update}
18:for\(i=l;i\geq 1;i--\)do
19: Update \(\mathbf{W}_{i}\leftarrow\mathbf{W}_{i}-\eta\cdot\tilde{\partial}_{\mathbf{W}_{i}}L\);
20:endfor
21: Update \(\mathbf{\hat{\xi}}=\left[\operatorname{vec}\left(\mathbf{\hat{W}}_{1}\right)^{\top}, \operatorname{vec}\left(\mathbf{\hat{W}}_{2}\right)^{\top},\ldots,\operatorname{ vec}\left(\mathbf{\hat{W}}_{l}\right)^{\top}\right]^{\top}\);
```

**Algorithm 1** An algorithm for the training of DNNs in the LNE manifold. We represent the gradient in the LNE manifold as \(\tilde{\partial}\). For brevity, we ignore the normalization operation (Ioffe and Szegedy, 2015; Ba et al., 2016).

Consequently, based on the previous work (Courbariaux et al., 2016), we give the practical algorithm of DNNs in the LNE manifold. As shown in Algorithm 1, this algorithm is similar to the general algorithm of DNNs, except for Line 14. Recall that, in Figure 1, the general DNNs use STE to simply copy the gradient, i.e., \(\tilde{\partial}_{\mathbf{W}_{i}}L=\partial_{\hat{\mathbf{W}}_{i}}L\), but we match the gradient by introducing the LNE metric. Furthermore, we can practically calculate this gradient in Line 14 via the strong approximation or the weak approximation in above.
## 6 Ricci Flow Discretized Neural Networks

In this section, we propose Ricci flow discretized neural networks (RF-DNNs). When the Ricci flow is introduced, it implies that the background of DNNs discussed here is the LNE manifold. Our aim is to provide a practical solution for the metric perturbation and further for the problem of gradient mismatch. Therefore, we will focus on how to calculate the discrete Ricci flow in practice rather than just staying on theoretical analysis.

To relate the Ricci flow to neural networks, we need to discretize the Ricci flow and choose a suitable coordinate system. For the left hand side of the Ricci flow, we have established the connection between the LNE metric and neural networks in Section 3. Note that we use the form of the LNE metric in the calculation, and such metrics at this time are with perturbations. For the right hand side of the Ricci flow, we need to calculate the Ricci curvature tensor with the coordinate system. And this coordinate system is the key to linking the Ricci curvature and neural networks. Specifically, we define a method for calculating the Ricci curvature in such a way that the selection of coordinate systems is related to the input transformations, which implies that the Ricci curvature in neural networks reflects the effect of different input transformations on the parameter.

### Ricci Curvature in Neural Networks

Now we consider the Ricci curvature tensor on the Riemannian metric \(g\). According to Appendix A, its coordinate form can be expressed as

\[-2\operatorname{Ric}(g)=-2R^{i}_{ikj}=2R^{i}_{kij} \tag{30}\] \[=g^{ip}\left(\partial_{i}\partial_{k}g_{pj}-\partial_{i} \partial_{j}g_{pk}+\partial_{p}\partial_{j}g_{ik}+\partial_{p}\partial_{k}g_{ ij}\right).\]

In order to relate the Ricci curvature to neural networks, we define a method for calculating the Ricci curvature in such a way that the selection of coordinate systems is related to the input transformations. When the Ricci curvature is equal to zero, it means that different input transformations will not cause variations of the parameter.

Inspired by the previous work (Kaul and Lall, 2019), we can treat the terms \(\partial_{i}\) and \(\partial_{p}\) as variations for the translation and rotation of each input, respectively. In general, since the data augmentation does not involve the rotation in real-world applications such as image classification tasks (He et al., 2016; Shorten and Khoshgoftaar, 2019), we consider the translation instead of the rotation by discarding the index \(p\), i.e., \(\partial_{p}(\partial_{j}g_{ik}+\partial_{k}g_{ij})=0\) for the fairness of ablation studies. When considering one of the translation and rotation, \(g^{ip}\) will degenerate into \(\delta^{ip}\) (identity matrix). Subsequently, \(\partial_{i}\partial_{k}g\) and \(\partial_{i}\partial_{j}g\) can be treated as variations for the row and column transformation of the input data w.r.t. the metric \(g\), respectively. Further the Ricci curvature can be rewritten as

\[-2\operatorname{Ric}(g)=\partial_{i}\partial_{k}g_{pj}-\partial_{i}\partial_{ j}g_{pk}. \tag{31}\]

**Remark 20**: _The choice of \(i\) and \(p\) is arbitrary (including \(k\) and \(j\)), and can even be other coordinate representations. Here, we just give it a specific geometric meaning by considering the characteristic of the image classification task._
As shown in Figure 5, with the help of Equation (31), we yield the Ricci curvature with coordinate systems in terms of the difference equation:

\[-2\operatorname{Ric}(g)=\frac{g|_{k_{1}}-g|_{k_{2}}}{k_{1}-k_{2}}-\frac{g|_{j_{1} }-g|_{j_{2}}}{j_{1}-j_{2}} \tag{32}\]

where we approximate partial derivatives with difference equations (Kaul and Lall, 2019), i.e., \(\partial_{i}\partial_{k}g=(g|_{k_{1}}-g|_{k_{2}})/(k_{1}-k_{2})\) and \(\partial_{i}\partial_{j}g=(g|_{j_{1}}-g|_{j_{2}})/(j_{1}-j_{2})\) corresponding to the input translation dimensions \(k\) and \(j\), respectively. Note that \(g|_{k_{1}}\), \(g|_{k_{2}}\), \(g|_{j_{1}}\), and \(g|_{j_{2}}\) are four metric structures under different small translation transformations \(k_{1}\), \(k_{2}\), \(j_{1}\), and \(j_{2}\) respectively. In general, \((k_{1}-k_{2})\) and \((j_{1}-j_{2})\) are translations less than 4 pixels, which is consistent with data augmentation (He et al., 2016).

### Existence of Discrete Ricci Flow in Neural Networks

Recall that we previously considered the Ricci-DeTurck flow instead of the Ricci flow as the solution of the Ricci flow does not always exist based on Section 2.3. If we can guarantee that the solution of Ricci flow exists in neural networks, then we can use the Ricci flow to exponentially decay the metric perturbation based on Section 4.4.

We have considered the right hand side of the Ricci flow in Equation (6), i.e., the Ricci curvature tensor. Now we define the equivalent form of the left hand side of the Ricci flow

Figure 5: By feeding the original image into the neural network and performing a forward and backward on the linear layer to update the weights \(\mathbf{w}\), we can construct the metric structure \(g(\mathbf{w})\) based on Section 5.1. In addition, we first process 4 different small translation transformations (\(k_{1}\), \(k_{2}\), \(j_{1}\), and \(j_{2}\)) on the original image, and then input them into the neural network. By sequentially performing a forward and backward, we can obtain four metric structures (\(g|_{k_{1}}\), \(g|_{k_{2}}\), \(g|_{j_{1}}\), and \(g|_{j_{2}}\)) under the corresponding translations. Combined with these metric structures, we can present the Ricci curvature structure \(\operatorname{Ric}(g)\).
in terms of the difference equation:

\[\frac{\partial}{\partial t}g(t):=g(t+1)-g(t). \tag{33}\]

where \(t\in\{0,1,\cdots,T-1\}\) is a uniform partition of the interval \([0,T]\). In the training of neural networks, \(T\) is the total number of iterations. As the limit \(T\rightarrow\infty\), the above formula holds on.

Combining Equation (32) and Equation (33) in neural networks, we present the expression of the discrete Ricci flow in terms of the difference equation:

\[\begin{split} g(t+1)|_{k_{1}}-g(t)|_{k_{1}}&=\frac {g(t)|_{k_{1}}-g(t)|_{k_{2}}}{k_{1}-k_{2}}-\frac{g(t)|_{j_{1}}-g(t)|_{j_{2}}}{ j_{1}-j_{2}}\\ g(0)|_{k_{1}}&=\delta-\tanh(\tau\mathbf{w})\tanh(\tau \mathbf{w})^{\top}\end{split} \tag{34}\]

To ensure the existence of the solution of the discrete Ricci flow, we are able to achieve this goal by adding a regularization into the loss function to constrain the discrete Ricci flow in DNNs. Following Equation (34), we further present the regularization:

\[N=\left\|g(t+1)|_{k_{1}}-g(t)|_{k_{1}}-\frac{g(t)|_{k_{1}}-g(t)|_{k_{2}}}{k_ {1}-k_{2}}+\frac{g(t)|_{j_{1}}-g(t)|_{j_{2}}}{j_{1}-j_{2}}\right\|_{L^{2}}^{ 2}, \tag{35}\]

where \(g(t)\) is \(\epsilon\)-close to the LNE metric \(g_{0}\) based on Definition 21. In other words, \(g(t)\) is LNE metric with perturbations.

**Definition 21**: _(Sheridan and Rubinstein, 2006) Let \(g(t)\) be the metrics on the LNE manifold. For \(\epsilon>0\), \(\mathcal{B}_{L^{2}}(g_{0},\epsilon)\) is the \(\epsilon\)-ball with respect to the \(L^{2}\)-norm induced by \(g_{0}\) and centred at \(g_{0}\), where any metric \(g(t)\in\mathcal{B}_{L^{2}}(g_{0},\epsilon)\) is \(\epsilon\)-close to \(g_{0}\) if_

\[(1+\epsilon)^{-1}g_{0}\leq g(t)\leq(1+\epsilon)g_{0}\]

_in the sense of matrices._

By constraining the regularization \(N\) in DNNs, the solution of the discrete Ricci flow exists when \(N\to 0\). Simultaneously, the metric perturbation will exponentially converges (\(g(t)\to g_{0}\)) as the evolution of the discrete Ricci flow.

### Algorithm Design

By applying the constraints of the discrete Ricci Flow in layer-by-layer LNE manifold, we can finally alleviate the problem of gradient mismatch. Since the background is the LNE manifold, we can construct the satisfied gradient on the basis of Equation (27). Note that, at this time, the metric is time-dependent under the Ricci flow, i.e., \(g_{\mathbf{w}}(t)\).

With the indicator function \(\mathbb{I}\) to represent the constraints of the discrete Ricci flow, we yield the gradient under the discrete Ricci flow as follows:

\[\tilde{\partial}_{\mathbf{w}}L=g_{\mathbf{w}}^{-1}(t)\partial_{Q(\mathbf{w})}\cdot \mathbb{I},\ \ \text{where}\ \ \mathbb{I}:=\left\{\begin{array}{ll}1&\text{ if }\left|N\right|\leq \varepsilon\\ 0&\text{ otherwise}\end{array}\right. \tag{36}\]
where we define a small constant \(\varepsilon\). Furthermore, due to the error in estimating the the regularization with finite training time, we allow an \(\varepsilon\)-bounded error in the above formulation, which also implies that the cases are less than or equal to \(\varepsilon\) (rather than strictly equal to \(0\)), all satisfy the discrete Ricci flow.

The overall process is shown in Algorithm 2. Compared with Algorithm 1, we add Line 7 and Line 15. In Line 7, we need to calculate the regularization to ensure that the solution of discrete Ricci flow exists. On the other hand, in Line 15, we calculate the gradient in the LNE manifold under the discrete Ricci flow, unlike Algorithm 1 which just calculates the gradient in the LNE manifold with perturbations. When we apply the Ricci flow, it means that the LNE manifold at this time is dynamic and anti-perturbation.
**Remark 22**: _In addition to using the discretized weight and activation, DNNs need to save the non-discretized weight and activation for gradient update. Note that the gradient of a DNN is non-discretized._

### Complexity Analysis

Based on Algorithm 2, one knows that the forward time complexity is about \(\mathcal{O}(n^{2})\) where the time complexity of Line 4 and Line 5 are about \(\mathcal{O}(n^{2})\) and \(\mathcal{O}(n)\), respectively. In the backward pass, the time complexity of Line 13 is about \(\mathcal{O}(n)\). Then the time complexity of the gradients w.r.t. the weight (Line 15) and activation (Line 17) both are about \(\mathcal{O}(n^{2})\). Therefore, the backward time complexity is about \(\mathcal{O}(2n^{2})\). For a training process of a neural network, its total complexity is \(\mathcal{O}(n^{2})\).

Since the computation of the Ricci curvature involves four different translations of input data w.r.t. the metric, its time complexity is about \(\mathcal{O}(n^{2})\). In this way, the updated weights are only used to calculate the constraints of discrete Ricci flow, and the final weights can be obtained by a backward pass again. The time complexity of Line 16 is \(\mathcal{O}(n^{2})\) when we use the weak approximation to calculate the gradient. Thus, the total complexity of our RF-DNN is still \(\mathcal{O}(n^{2})\). Hence, the complexity of the RF-DNN keep consistency with the general neural network.

## 7 Experiments

In this section, we design ablation studies to compare our RF-DNN 9 trained from scratch with other STE methods. Furthermore, given a pre-trained model, we evaluate the performance of the RF-DNN in comparison with several representative training-based methods on classification benchmark datasets. All the experiments implemented in Python are conducted with PyTorch (Paszke et al., 2019). The hardware environment is conducted on a Workstation with an Intel(R) Xeon(R) Silver 4214 CPU(2.20 GHz), GeForce GTX 2080Ti GPU and 128GB RAM.

Footnote 9: In order to calculate the gradient conveniently, we use the weak approximation of the inverse of the LNE metric in all experiments.

### Experimental Settings

The two datasets used in our experiments are introduced as follows.

**CIFAR datasets:** There are two CIFAR benchmarks (Krizhevsky et al., 2009) consisting of natural color images with 32 \(\times\) 32 pixels, respectively, 50k training and 10k test images, and we pick out 5k training images as a validation set from the training set. CIFAR-10 consists of images organized into 10 classes and CIFAR-100 into 100 classes. We adopt a standard data augmentation scheme (random corner cropping and random flipping) that is widely used for these two datasets. We normalize the images with the means of the channel and standard deviations in preprocessing.

**ImageNet dataset:** The ImageNet benchmark (Russakovsky et al., 2015) consists of 1.2 million high-resolution natural images, where the validation set contains 50k images. These images are organized into 1000 categories of the objects for training, which are resized
to 224 \(\times\) 224 pixels before fed into the network. In the next experiments, we report our single-crop evaluation results using top-1 and top-5 accuracies.

We specify the discrete function, the composition of which has a significant influence on the performance and computation of DNNs. Specifically, the discrete function is able to simplify the calculations which are also varied depending on the different discrete values, e.g., fixed-point multiplication, SHIFT operation (Elhoushi et al., 2019) and XNOR operation (Rastegari et al., 2016) etc.

We mark \(Q^{1}\) as the 1-bit discrete function:

\[Q^{1}(\cdot)=\operatorname{sign}(\cdot)=\{-1,+1\}. \tag{37}\]

The \(k\)-bit, over 1-bit, discrete function can be marked as \(Q^{k}\),

\[Q^{k>1}(\cdot)=\frac{2}{2^{k}-1}\operatorname{round}\left[(2^{k}-1)\left( \frac{\cdot}{2\max\lvert\cdot\rvert}+\frac{1}{2}\right)\right]-1 \tag{38}\]

where \(\operatorname{round}[\cdot]\) is the rounding function and \(\max\lvert\cdot\rvert\) means to calculate the absolute value of the input first, and then find its maximum value. In this way, a DNN using discrete function \(Q^{1}(\cdot)\) can be calculated with XNOR operation and using discrete function \(Q^{k>1}(\cdot)\) can be calculated with fixed-point multiplication.

### Ablation Studies with STE Methods

In order to illustrate the superiority of RF-DNN against the problem of gradient mismatch, we compare our RF-DNN with three other methods by training from scratch. In Table 1, Table 2 and Table 3, we mark \(\{-1,+1\}\) in '**Forward**', which indicates that the weights are binarized using Equation (37), i.e., \(-1\) or \(+1\), in the forward of DNNs. In the backward, the methods (Dorefa (Zhou et al., 2016), MultiFCG (Chen et al., 2019) and FCGrad (Chen et al., 2019)) use the different approximated gradients to update the weights. Here, we apply different ResNet models (He et al., 2016) to compare ablation studies.

The batch normalization with a batch size of 128 is used in the learning strategy, and Nesterov momentum of 0.9 (Dozat, 2016) is used in SGD optimization. For CIFAR, we set total training epochs as 200 and set a weight decay of 0.0005 where the learning rate is lowered by 10 times at epoch 80, 150, and 190, with the initial value of 0.1. For ImageNet, we set total training epochs as 100 and set the learning rate of each parameter group using a cosine annealing schedule with a weight decay of 0.0001. All experiments are conducted for 5 times, and the statistics of the last 10/5 epochs' test accuracies are reported for a fair comparison. Hence, we evaluate the accuracy performance in terms of (mean \(\pm\) std). Note that we perform standard data augmentation and pre-processing on CIFAR and ImageNet datasets.

In Table 1, Table 2 and Table 3, we use the same the discrete function \(Q^{1}(\cdot)\), parameter setting and optimizer for fairness in the forward. The only difference is the gradient in the backward propagation. The performance in various models and datasets shows that our RF-DNN has significant improvement over other STE methods. The average results of multiple experiments are better than that of other methods, which may benefit from the alleviation of the gradient mismatch to make the loss function of DNNs more fully descended. In addition, the minor variances show that our training method is relatively stable, which also confirms our point of view.
### Convergence and Stability Analysis

Since the standard deviations can reflect the convergence and stability of training to a certain extent, we visualize the Table 1 as shown in Figure 6(a). Intuitively, compared to Dorefa and MultiFCG, our proposed RF-DNN can better alleviate the perturbations caused by the gradient mismatch to achieve more stable performance. Furthermore, we also present the accuracy performance of RF-DNN with different bit width weight representations in Figure 6(b). And we see fairly consistent stability across different bit width and backbone models.

As shown in Figure 7, the RF-DNN can achieve higher accuracies than Dorefa on CIFAR100 dataset, i.e., 1.25% higher on the training dataset with ResNet56, 1.85% higher on the test dataset with ResNet56, 1.97% higher on the training dataset with ResNet110

\begin{table}
\begin{tabular}{c c c c c} \hline \hline
**Network** & **Forward** & **Backward** & **Test Acc (\%)** & **FP Acc (\%)** \\ \hline \multirow{3}{*}{ResNet20} & & Dorefa & 88.28\(\pm\)0.81 & \multirow{3}{*}{91.50} \\  & & MultiFCG & 88.94\(\pm\)0.46 & & \\  & & RF-DNN & **89.83\(\pm\)0.23** & \\ \hline \multirow{3}{*}{ResNet32} & & Dorefa & 90.23\(\pm\)0.63 & \multirow{3}{*}{92.13} \\  & & MultiFCG & 89.63\(\pm\)0.38 & & \\  & & RF-DNN & **90.75\(\pm\)0.19** & \\ \hline \multirow{3}{*}{ResNet44} & & Dorefa & 90.71\(\pm\)0.58 & \multirow{3}{*}{93.56} \\  & & MultiFCG & 90.54\(\pm\)0.21 & & \\ \cline{1-1}  & & RF-DNN & **91.63\(\pm\)0.11** & \\ \hline \hline \end{tabular}
\end{table}
Table 1: The experimental results on CIFAR10 with ResNet20/32/44. The accuracy of full-precision (FP) baseline is reported by (Chen et al., 2019).

\begin{table}
\begin{tabular}{c c c c c} \hline \hline
**Network** & **Forward** & **Backward** & **Test Acc (\%)** & **FP Acc (\%)** \\ \hline \multirow{3}{*}{ResNet56} & \multirow{3}{*}{\{\(-\)1,+1\}} & Dorefa & 66.71\(\pm\)2.32 & \multirow{3}{*}{71.22} \\  & & MultiFCG & 66.58\(\pm\)0.37 & \\  & & FCGrad & 66.56\(\pm\)0.35 & \\  & & RF-DNN & **68.56\(\pm\)0.32** & \\ \hline \multirow{3}{*}{ResNet110} & \multirow{3}{*}{\{\(-\)1,+1\}} & Dorefa & 68.15\(\pm\)0.50 & \multirow{3}{*}{72.54} \\  & & MultiFCG & 68.27\(\pm\)0.14 & \\ \cline{1-1}  & & FCGrad & 68.74\(\pm\)0.36 & \\ \cline{1-1}  & & RF-DNN & **69.20\(\pm\)0.28** & \\ \hline \hline \end{tabular}
\end{table}
Table 2: The experimental results on CIFAR100 with ResNet56/110. The accuracy of full-precision (FP) baseline is reported by (Chen et al., 2019).
Figure 6: Accuracy performance (mean \(\pm\) std) for ResNet20/32/44 on CIFAR10. The line and bar represent the mean and standard deviation of the different random seed results, respectively. (a) We compare our RF-DNN with Dorefa and MultiFCG via the 1bit weight representation, which is also the visualization of Table 1. (b) We present our RF-DNN via different bit width weight representations. Note that higher mean and lower deviation always implies better convergence and stability.

Figure 7: Training and test curves of ResNet56/110 on CIFAR100 compared between Dorefa and RF-DNN. Intuitively, RF-DNN has a more stable training performance than Dorefa.
and 1.05% higher on the test dataset with ResNet110. In addition, it can be seen from the fluctuation of the test curves in Figure 7 that the RF-DNN has a tremendous improvement compared with the Dorefa on the stability of training. From the training curve, our method significantly outperforms Dorefa, of course, this needs to be combined with the test curve to explain the superiority of RF-DNN. From the test curve, the accuracy of our method is always higher than that of Dorefa, and the corresponding volatility has improved. The experimental results verify that our theoretical framework is an effective solution against the gradient mismatch, further improving the training performance of DNNs.

### Comparisons with Training-based Methods

In this experiment, based on a full-precision pre-trained model, we compare the RF-DNN with several state-of-the-art DNNs, e.g., DeepShift (Elhoushi et al., 2019), QN (Yang et al., 2019), ADMM (Leng et al., 2018), MetaQuant (Chen et al., 2019), INT8 (Zhu et al., 2020), SR+DR (Gysel et al., 2018), ELQ (Zhou et al., 2018), MD (Ajanthan et al., 2021) and RQ (Louizos et al., 2019), under the same bit width using Equation (37) or Equation (38). Note that **W** and **A** represent the bit width of weights and activations respectively in Table 4. Consequently, the experimental results show that our RF-DNN is able to achieve better performance than other recent state-of-the-art training-based methods, which seems to benefit from our effective solution of the gradient mismatch.

## 8 Conclusion and Future Work

Traditional discretized neural networks (DNNs) tell us that the weights can only take low-precision discrete values, as well as the activations. This makes DNNs' memory footprint very light compared to full-precision floating point networks. However, it makes their training difficult: to maintain discrete weights, in general, the gradient w.r.t. discrete weights is approximated with the Straight-Through Estimator (STE), which causes the update weight to differ from the gradient w.r.t. continuous weights (_gradient mismatch_).

This paper proposes a novel analysis of the gradient mismatch phenomenon through the lens of duality theory. This mismatch is then viewed as a metric perturbation in a Riemannian manifold. In theory, on the basis of the information geometry, we construct the LNE manifold for neural networks, such that forming the background to effectively

\begin{table}
\begin{tabular}{c c c c c} \hline \hline
**Network** & **Forward** & **Backward** & **Test Top1/Top5 (\%)** & **FP Top1/Top5 (\%)** \\ \hline \multirow{4}{*}{ResNet18} & & Dorefa & 58.34\(\pm\)2.07/81.47\(\pm\)1.56 & \multirow{4}{*}{69.76/89.08} \\  & & MultiFCG & 59.47\(\pm\)0.02/82.41\(\pm\)0.01 & & \\ \cline{1-1}  & & FCGrad & 59.83\(\pm\)0.36/82.67\(\pm\)0.23 & & \\ \cline{1-1}  & & RF-DNN & **60.83\(\pm\)0.41/83.54\(\pm\)0.18** & & \\ \hline \hline \end{tabular}
\end{table}
Table 3: The experimental results on ImageNet with ResNet18. The accuracy of full-precision (FP) baseline is reported by (Chen et al., 2019).
deal with a metric perturbation. By revealing the stability of LNE metrics with the \(L^{2}\)-norm perturbation under the Ricci-DeTurck flow, we subsequently introduce the Ricci flow Discretized Neural Network (RF-DNN) in practice by using the constraints of the discrete Ricci flow in the LNE manifold to alleviate the metric perturbation with the exponential convergence rate, giving rise to an appealing solution for STE in DNNs. Experimentally, our RF-DNN achieves improvements in both the stability and performance of DNNs.

In this paper, information geometry is a very important part of combining the geometric tool (Ricci flow) and neural networks. Our future research will continue to explore the

\begin{table}
\begin{tabular}{l c c c c c c} \hline \hline \multirow{2}{*}{**Method**} & \multirow{2}{*}{**W**} & \multirow{2}{*}{**A**} & \multicolumn{2}{c}{**Top-1**} & \multicolumn{2}{c}{**Top-5**} \\ \cline{5-7}  & & & Accuracy & Gap & Accuracy & Gap \\ \hline
**AlexNet** (Original) & 32 & 32 & 56.52\% & - & 79.07\% & - \\ \hline RF-DNN (ours) & 6 & 32 & **56.39\%** & \(\mathbf{-0.13\%}\) & **78.78\%** & \(\mathbf{-0.29\%}\) \\ DeepShift (Elhoushi et al., 2019) & 6 & 32 & 54.97\% & \(\mathbf{-1.55\%}\) & 78.26\% & \(\mathbf{-0.81\%}\) \\ \hline
**ResNet18** (Original) & 32 & 32 & 69.76\% & - & 89.08\% & - \\ \hline RF-DNN (ours) & 1 & 32 & **67.05\%** & \(\mathbf{-2.71\%}\) & **88.09\%** & \(\mathbf{-0.99\%}\) \\ MD (Ajanthan et al., 2021) & 1 & 32 & 66.78\% & \(\mathbf{-2.98\%}\) & 87.01\% & \(\mathbf{-2.07\%}\) \\ ELQ (Zhou et al., 2018) & 1 & 32 & 66.21\% & \(\mathbf{-3.55\%}\) & 86.43\% & \(\mathbf{-2.65\%}\) \\ ADMM (Leng et al., 2018) & 1 & 32 & 64.80\% & \(\mathbf{-4.96\%}\) & 86.20\% & \(\mathbf{-2.88\%}\) \\ QN (Yang et al., 2019) & 1 & 32 & 66.50\% & \(\mathbf{-3.26\%}\) & 87.30\% & \(\mathbf{-1.78\%}\) \\ MetaQuant (Chen et al., 2019) & 1 & 32 & 63.44\% & \(\mathbf{-6.32\%}\) & 84.77\% & \(\mathbf{-4.31\%}\) \\ RF-DNN (ours) & 4 & 4 & **66.75\%** & \(\mathbf{-3.01\%}\) & **87.02\%** & \(\mathbf{-2.06\%}\) \\ RQ ST (Louizos et al., 2019) & 4 & 4 & 62.46\% & \(\mathbf{-7.30\%}\) & 84.78\% & \(\mathbf{-4.30\%}\) \\ \hline
**ResNet50** (Original) & 32 & 32 & 76.13\% & - & 92.86\% & - \\ \hline RF-DNN (ours) & 8 & 8 & **76.07\%** & \(\mathbf{-0.06\%}\) & **92.87\%** & \(\mathbf{+0.01\%}\) \\ INT8 (Zhu et al., 2020) & 8 & 8 & 75.87\% & \(\mathbf{-0.26\%}\) & - & - \\ \hline
**MobileNet** (Original) & 32 & 32 & 70.61\% & - & 89.47\% & - \\ \hline RF-DNN (ours) & 5 & 5 & **61.32\%** & \(\mathbf{-9.29\%}\) & **84.08\%** & \(\mathbf{-5.39\%}\) \\ SR+DR (Gysel et al., 2018) & 5 & 5 & 59.39\% & \(\mathbf{-11.22\%}\) & 82.35\% & \(\mathbf{-7.12\%}\) \\ RQ ST (Louizos et al., 2019) & 5 & 5 & 56.85\% & \(\mathbf{-13.76\%}\) & 80.35\% & \(\mathbf{-9.12\%}\) \\ RF-DNN (ours) & 8 & 8 & **70.76\%** & \(\mathbf{+0.15\%}\) & **89.54\%** & \(\mathbf{+0.07\%}\) \\ RQ (Louizos et al., 2019) & 8 & 8 & 70.43\% & \(\mathbf{-0.18\%}\) & 89.42\% & \(\mathbf{-0.05\%}\) \\ \hline \hline \end{tabular}
\end{table}
Table 4: The classification accuracy results on ImageNet and comparison with other training-based methods, with AlexNet (Krizhevsky et al., 2012), ResNet18, ResNet50 and MobileNet (Howard et al., 2017). Note that the accuracy of full-precision baseline is reported by (Elhoushi et al., 2019).
connection between neural networks and manifolds and aim to introduce geometric ideas to solve practical problems in machine learning.

## Acknowledgments

We thank all reviewers and the editor for excellent contributions.
## Appendix A Differential Geometry

1. Riemann curvature tensor (Rm) is a (1,3)-tensor defined for a 1-form \(\omega\):

\[R^{l}_{ijk}\omega_{l}=\nabla_{i}\nabla_{j}\omega_{k}-\nabla_{j}\nabla_{i}\omega_ {k}\]

where the covariant derivative of \(F\) satisfies

\[\nabla_{p}F^{j_{1}\dots j_{l}}_{i_{1}\dots i_{k}}=\partial_{p}F^{j_{1}\dots j_ {l}}_{i_{1}\dots i_{k}}+\sum_{s=1}^{l}F^{j_{1}\dots q\dots j_{l}}_{i_{1}\dots i _{k}}\Gamma^{j_{s}}_{pq}-\sum_{s=1}^{k}F^{j_{1}\dots j_{l}}_{i_{1}\dots q\dots i _{k}}\Gamma^{q}_{pi}.\]

In particular, coordinate form of the Riemann curvature tensor is:

\[R^{l}_{ijk}=\partial_{i}\Gamma^{l}_{jk}-\partial_{j}\Gamma^{l}_{ik}+\Gamma^{p }_{jk}\Gamma^{l}_{ip}-\Gamma^{p}_{ik}\Gamma^{l}_{jp}.\]

2. Christoffel symbol in terms of an ordinary derivative operator is:

\[\Gamma^{k}_{ij}=\frac{1}{2}g^{kl}(\partial_{i}g_{jl}+\partial_{j}g_{il}- \partial_{l}g_{ij}).\]

3. Ricci curvature tensor (Ric) is a (0,2)-tensor:

\[R_{ij}=R^{p}_{pij}.\]

4. Scalar curvature is the trace of the Ricci curvature tensor:

\[R=g^{ij}R_{ij}.\]

5. Lie derivative of \(F\) in the direction \(\frac{d\varphi(t)}{dt}\):

\[\mathcal{L}_{\frac{d\varphi(t)}{dt}}F=\left(\frac{d}{dt}\varphi^{*}(t)F\right) _{t=0}\]

where \(\varphi(t):\mathcal{M}\rightarrow\mathcal{M}\) for \(t\in(-\epsilon,\epsilon)\) is a time-dependent diffeomorphism of \(\mathcal{M}\) to \(\mathcal{M}\).

## Appendix B Notation

For clarity of definitions in this paper, we list the important notations as shown in Table 5.

## Appendix C Proof for the Ricci Flow

### Proof for Lemma 5

**Lemma 23**: _The linearization of the Ricci curvature tensor is given by_

\[\mathcal{D}[\mathrm{Ric}](h)_{ij}=-\frac{1}{2}g^{pq}(\nabla_{p}\nabla_{q}h_{ ij}+\nabla_{i}\nabla_{j}h_{pq}-\nabla_{q}\nabla_{i}h_{jp}-\nabla_{q}\nabla_{j}h_{ip}).\]

Based on Appendix A, we have

\[\nabla_{q}\nabla_{i}h_{jp}=\nabla_{i}\nabla_{q}h_{jp}-R^{r}_{qij}h_{rp}-R^{r} _{qip}h_{jm}.\]
Combining with Lemma 23, we can obtain the deformation equation because of \(\nabla g=0\),

\[\mathcal{D}[-2\text{Ric}](h)_{ij}= g^{pq}\nabla_{p}\nabla_{q}h_{ij}+\nabla_{i}\left(\frac{1}{2} \nabla_{j}h_{pq}-\nabla_{q}h_{jp}\right)+\nabla_{j}\left(\frac{1}{2}\nabla_{i}h _{pq}-\nabla_{q}h_{ip}\right)+O(h_{ij})\] \[= g^{pq}\nabla_{p}\nabla_{q}h_{ij}+\nabla_{i}V_{j}+\nabla_{j}V_{i} +O(h_{ij}).\]

\begin{table}
\begin{tabular}{|p{113.8pt} p{113.8pt} p{113.8pt}|p{113.8pt}|} \hline \(\boldsymbol{W}_{i}\): & weight matrix for the \(i\)-th layer & \(\boldsymbol{\hat{W}}_{i}\): & discretized weight matrix for the \(i\)-th layer \\ \hline \(\boldsymbol{w}\): & vectorized weights in each layer & \(\boldsymbol{\hat{w}}\): & discretized vectorized weights in each layer \\ \hline \(\boldsymbol{a}_{i}\): & activation vector for the \(i\)-th layer & \(\boldsymbol{\hat{a}}_{i}\): & discretized activation vector for the \(i\)-th layer \\ \hline \(\boldsymbol{\xi}\): & parameter vector & \(\boldsymbol{\bar{\xi}}\): & discretized parameter vector \\ \hline \(Q^{1}\): & 1-bit discrete function & \(Q^{k>1}\): & k-bit discrete function (over 1-bit) \\ \hline \(\delta\): & Euclidean metric (identity matrix) & \(\Phi\): & convex function \\ \hline \(g_{0}\): & LNE metric under Ricci flow & \(\bar{g}_{0}\): & LNE metric under Ricci-DeTurck flow \\ \hline \(g\) or \(g(t)\): & the metrics under Ricci flow & \(\bar{g}\) or \(\bar{g}(t)\): & the metrics under Ricci-DeTurck flow \\ \hline \(g(0)\): & initial metric under Ricci flow & \(\bar{g}(0)\): & initial metric under Ricci-DeTurck flow \\ \hline \(d(0)\): & the initial perturbation & \(d(t)\): & the time-evolving perturbation \\ \hline \(D\): & divergence & \(L\) or \(\bar{L}\): & loss function \\ \hline \(L_{g_{0}}\): & Lichnerowicz operator & \(L^{2}\) or \(L^{\infty}\): & norm \\ \hline \(\partial\): & partial derivative & \(\nabla\): & covariant derivative \\ \hline \(\mathcal{L}\): & Lie derivative & \(\Delta_{g_{0}}\): & the Laplacian \\ \hline Rm: & Riemann curvature tensor & \(f\): & nonlinear function \\ \hline Ric: & Ricci curvature tensor & \(\mathcal{D}[\text{Ric}]\): & the linearization of the Ricci curvature tensor \\ \hline \(\varphi^{*}\): & pullback & \(\phi_{*}\): & pushforward \\ \hline \(B(x,r)\): & the ball with a radius \(r\) and a point \(x\in\mathcal{M}\) & \(\mathcal{B}_{L^{2}}(\bar{g}_{0},\epsilon)\): & the \(\epsilon\)-ball with respect to the \(L^{2}\)-norm induced by \(\bar{g}_{0}\) and centred at \(\bar{g}_{0}\) \\ \hline \end{tabular}
\end{table}
Table 5: Definitions of notations
### Description of the DeTurck Trick

Based on the chain rule for the Lie derivative in Appendix A, we can calculate

\[\frac{\partial}{\partial t}g(t) =\frac{\partial\left(\varphi^{*}(t)\bar{g}(t)\right)}{\partial t}\] \[=\left(\frac{\partial\left(\varphi^{*}(t+\tau)\bar{g}(t+\tau) \right)}{\partial\tau}\right)_{\tau=0}\] \[=\left(\varphi^{*}(t)\frac{\partial\bar{g}(t+\tau)}{\partial \tau}\right)_{\tau=0}+\left(\frac{\partial\left(\varphi^{*}(t+\tau)\bar{g}(t) \right)}{\partial\tau}\right)_{\tau=0}\] \[=\varphi^{*}(t)\frac{\partial}{\partial t}\bar{g}(t)+\varphi^{*} (t)\mathcal{L}_{\frac{\partial\varphi(t)}{\partial t}}\bar{g}(t)\]

where \(\frac{\partial\varphi(t)}{\partial t}\) is equal to \(V(t)\)(Sheridan and Rubinstein, 2006). With the help of Equation (6), we have the following expression for the pullback metric \(g(t)\)

\[\frac{\partial}{\partial t}g(t)=\varphi^{*}(t)\frac{\partial}{\partial t} \bar{g}(t)+\varphi^{*}(t)\mathcal{L}_{\frac{\partial\varphi(t)}{\partial t}} \bar{g}(t)=-2\operatorname{Ric}(\varphi^{*}(t)\bar{g}(t))=-2\varphi^{*}(t) \operatorname{Ric}(\bar{g}(t)). \tag{39}\]

The diffeomorphism invariance of the Ricci curvature tensor is used in the last step. The above equation is equivalent to

\[\frac{\partial}{\partial t}\bar{g}(t)=-2\operatorname{Ric}(\bar{g}(t))- \mathcal{L}_{\frac{\partial\varphi(t)}{\partial t}}\bar{g}(t).\]

Based on Definition 24, we further yield

\[\frac{\partial}{\partial t}\bar{g}(t)=-2\operatorname{Ric}(\bar{g}(t))-\nabla _{i}V_{j}-\nabla_{j}V_{i}.\]

**Definition 24**: _(Sheridan and Rubinstein, 2006) On a Riemannian manifold \((\mathcal{M},g)\), we have_

\[(\mathcal{L}_{X}g)_{ij}=\nabla_{i}X_{j}+\nabla_{j}X_{i},\]

_where \(\nabla\) denotes the Levi-Civita connection of the metric \(g\), for any vector field \(X\)._

### Curvature Explosion at Singularity

In general, we present the behavior of Ricci flow in finite time and show that the evolution of the curvature is close to divergence. The core demonstration is followed with Theorem 28.

**Theorem 25**: _(Sheridan and Rubinstein, 2006) Given a smooth Riemannian metric \(g_{0}\) on a closed manifold \(\mathcal{M}\), there exists a maximal time interval \([0,T)\) such that a solution \(g(t)\) of the Ricci flow, with \(g(0)=g_{0}\), exists and is smooth on \([0,T)\), and this solution is unique._

**Theorem 26**: _Let \(\mathcal{M}\) be a closed manifold and \(g(t)\) a smooth time-dependent metric on \(\mathcal{M}\), defined for \(t\in[0,T)\). If there exists a constant \(C<\infty\) for all \(x\in\mathcal{M}\) such that_

\[\int_{0}^{T}\left|\frac{\partial}{\partial t}g_{x}(t)\right|_{g(t)}dt\leq C, \tag{40}\]
_then the metrics \(g(t)\) converge uniformly as \(t\) approaches \(T\) to a continuous metric \(g(T)\) that is uniformly equivalent to \(g(0)\) and satisfies_

\[e^{-C}g_{x}(0)\leq g_{x}(T)\leq e^{C}g_{x}(0). \tag{41}\]

**Proof** Considering any \(x\in\mathcal{M}\), \(t_{0}\in[0,T)\), \(V\in T_{x}\mathcal{M}\), we have

\[\left|\log\left(\frac{g_{x}(t_{0})(V,V)}{g_{x}(0)(V,V)}\right)\right| =\left|\int_{0}^{t_{0}}\frac{\partial}{\partial t}\left[\log g_{x} (t)(V,V)\right]dt\right|\] \[=\left|\int_{0}^{t_{0}}\frac{\frac{\partial}{\partial t}g_{x}(t)( V,V)}{g_{x}(t)(V,V)}dt\right|\] \[\leq\int_{0}^{t_{0}}\left|\frac{\partial}{\partial t}g_{x}(t) \left(\frac{V}{|V|_{g(t)}},\frac{V}{|V|_{g(t)}}\right)\right|dt\] \[\leq\int_{0}^{t_{0}}\left|\frac{\partial}{\partial t}g_{x}(t) \right|_{g(t)}dt\] \[\leq C.\]

By exponentiating both sides of the above inequality, we have

\[e^{-C}g_{x}(0)(V,V)\leq g_{x}(t_{0})(V,V)\leq e^{C}g_{x}(0)(V,V).\]

This inequality can be rewritten as

\[e^{-C}g_{x}(0)\leq g_{x}(t_{0})(V,V)\leq e^{C}g_{x}(0)(V,V)\]

because it holds for any \(V\). Thus, the metrics \(g(t)\) are uniformly equivalent to \(g(0)\).

Consequently, we have the well-defined integral:

\[g_{x}(T)-g_{x}(0)=\int_{0}^{T}\frac{\partial}{\partial t}g_{x}(t)dt.\]

We can show that this integral is well-defined from two perspectives. Firstly, as long as the metrics are smooth, the integral exists. Secondly, the integral is absolutely integrable. Based on the norm inequality induced by \(g(0)\), we can obtain

\[\left|g_{x}(T)-g_{x}(t)\right|_{g(0)}\leq\int_{t}^{T}\left|\frac{\partial}{ \partial t}g_{x}(t)\right|_{g(0)}dt.\]

For each \(x\in\mathcal{M}\), the above integral will approach zero as \(t\) approaches \(T\). Since \(\mathcal{M}\) is compact, the metrics \(g(t)\) converge uniformly to a continuous metric \(g(T)\) which is uniformly equivalent to \(g(0)\) on \(\mathcal{M}\). Moreover, we can show that

\[e^{-C}g_{x}(0)\leq g_{x}(T)\leq e^{C}g_{x}(0).\]
**Corollary 27**: _Let \((\mathcal{M},g(t))\) be a solution of the Ricci flow on a closed manifold. If \(|\operatorname{Rm}|_{g(t)}\) is bounded on a finite time \([0,T)\), then \(g(t)\) converges uniformly as \(t\) approaches \(T\) to a continuous metric \(g(T)\) which is uniformly equivalent to \(g(0)\)._

**Proof** The bound on \(|\operatorname{Rm}|_{g(t)}\) implies one on \(|\operatorname{Ric}|_{g(t)}\). Based on Equation (6), we can extend the bound on \(|\frac{\partial}{\partial t}g(t)|_{g(t)}\). Therefore, we obtain an integral of a bounded quantity over a finite interval is also bounded, by Theorem 26.

**Theorem 28**: _If \(g_{0}\) is a smooth metric on a compact manifold \(\mathcal{M}\), the Ricci flow with \(g(0)=g_{0}\) has a unique solution \(g(t)\) on a maximal time interval \(t\in[0,T)\). If \(T<\infty\), then_

\[\lim_{t\to T}\left(\sup_{x\in\mathcal{M}}|\operatorname{Rm}_{x}(t)| \right)=\infty. \tag{42}\]

**Proof** For a contradiction, we assume that \(|\operatorname{Rm}_{x}(t)|\) is bounded by a constant. It follows from Corollary 27 that the metrics \(g(t)\) converges smoothly to a smooth metric \(g(T)\). Based on Theorem 25, it is possible to find a solution to the Ricci flow on \(t\in[0,T)\), as the smooth metric \(g(T)\) is uniformly equivalent to the initial metric \(g(0)\).

Hence, we can extend the solution of the Ricci flow after the time point \(t=T\), which contradicts the choice of \(T\) as the maximal time for the existence of the Ricci flow on \([0,T)\). In other words, \(|\operatorname{Rm}_{x}(t)|\) is unbounded.

As approaching the singular time \(T\), the Riemann curvature \(|\operatorname{Rm}|_{g(t)}\) becomes no longer convergent and tends to explode.

## Appendix D Proof for All Time Convergence in LNE Manifolds

**Definition 29**: _(Deruelle and Kroncke, 2021) A complete LNE \(n\)-manifold \((\mathcal{M}^{n},g_{0})\) is said to be linearly stable if the \(L^{2}\) spectrum of the Lichnerowicz operator \(L_{g_{0}}:=\Delta_{g_{0}}+2\operatorname{Rm}(g_{0})*\) is in \((-\infty,0]\) where \(\Delta_{g_{0}}\) is the Laplacian, when \(L_{g_{0}}\) acting on \(d_{ij}\) satisfies_

\[L_{g_{0}}(d) =\Delta_{g_{0}}d+2\operatorname{Rm}(g_{0})*d \tag{43}\] \[=\Delta_{g_{0}}d+2\operatorname{Rm}(g_{0})_{iklj}d_{mn}g_{0}^{km }g_{0}^{ln}.\]

**Definition 30**: _(Deruelle and Kroncke, 2021) A \(n\)-manifold \((\mathcal{M}^{n},g_{0})\) is said to be integrable if a neighbourhood of \(g_{0}\) has a smooth structure._

We rewrite the Ricci-DeTurck flow (21) as an evolution of the difference \(d(t):=\bar{g}(t)-\bar{g}_{0}\), such that

\[\frac{\partial}{\partial t}d(t) =\frac{\partial}{\partial t}\bar{g}(t)=-2\operatorname{Ric}(\bar{ g}(t))+2\operatorname{Ric}(\bar{g}_{0})+\mathcal{L}_{\frac{\partial\varphi^{\prime}(t) }{\partial t}}\bar{g}_{0}-\mathcal{L}_{\frac{\partial\varphi(t)}{\partial t}} \bar{g}(t) \tag{44}\] \[=\Delta d(t)+\operatorname{Rm}*d(t)+F_{\bar{g}^{-1}}*\nabla^{\bar {g}_{0}}d(t)*\nabla^{\bar{g}_{0}}d(t)+\nabla^{\bar{g}_{0}}\left(G_{\Gamma(\bar {g}_{0})}*d(t)*\nabla^{\bar{g}_{0}}d(t)\right),\]
where the tensors \(F\) and \(G\) depend on \(\bar{g}^{-1}\) and \(\Gamma(\bar{g}_{0})\). Note that \(\bar{g}_{0}\) is the LNE metric which satisfies the above formula.

In the follwing, we denote \(\|\cdot\|_{L^{2}}\) or \(\|\cdot\|_{L^{\infty}}\) as the \(L^{2}\)-norm or \(L^{\infty}\)-norm w.r.t. the LNE metric \(\bar{g}_{0}\), and mark generic constants as \(C\) or \(C_{1}\).

**Lemma 31**: _Let \(\bar{g}(t)\) be a Ricci-DeTurck flow on a maximal time interval \(t\in(0,T)\) in an \(L^{2}\)-neighbourhood of \(\bar{g}_{0}\). We have the following estimate:_

\[\left\|\frac{\partial}{\partial t}d_{0}(t)\right\|_{L^{2}}\leq C\left\|\nabla^ {\bar{g}_{0}(t)}\left(d(t)-d_{0}(t)\right)\right\|_{L^{2}}^{2}. \tag{45}\]

**Proof** According to the Hardy inequality (Minerbe, 2009), we have the same proofs by referring the details (Deruelle and Kroncke, 2021). \(\blacksquare\)

**Theorem 32**: _Let \((\mathcal{M}^{n},\bar{g}_{0})\) be the LNE \(n\)-manifold which is linearly stable and integrable. Then, there exists a constant \(\alpha_{\bar{g}_{0}}\) satisfying_

\[\left(\Delta d(t)+\mathrm{Rm}(\bar{g}_{0})\ast d(t),d(t)\right)_{L^{2}}\leq- \alpha_{\bar{g}_{0}}\left\|\nabla^{\bar{g}_{0}}d(t)\right\|_{L^{2}}^{2} \tag{46}\]

_for all \(\bar{g}(t)\in\tilde{\mathcal{F}}\) whose definition is given in Equation (24)._

**Proof** The similar proofs can be found in (Devyyer, 2014) with some minor modifications. Due to the linear stability requirement of LNE manifolds in Definition 29 and Definition 30, \(-L_{\bar{g}_{0}}\) is non-negative. Then there exists a positive constant \(\alpha_{\bar{g}_{0}}\) satisfying

\[\alpha_{\bar{g}_{0}}\left(-\Delta d(t),d(t)\right)_{L^{2}}\leq\left(-\Delta d (t)-\mathrm{Rm}(\bar{g}_{0})\ast d(t),d(t)\right)_{L^{2}}.\]

By Taylor expansion, we repeatedly use elliptic regularity and Sobolev embedding (Pacini, 2010) to obtain the estimate. \(\blacksquare\)

**Corollary 33**: _Let \((\mathcal{M}^{n},\bar{g}_{0})\) be the LNE \(n\)-manifold which is integrable. For a Ricci-DeTurck flow \(\bar{g}(t)\) on a maximal time interval \(t\in[0,T]\), if it satisfies \(\|\bar{g}(t)-\bar{g}_{0}\|_{L^{\infty}}<\epsilon\) where \(\epsilon>0\), then there exists a constant \(C<\infty\) for \(t\in[0,T]\) such that the evolution inequality satisfies_

\[\|d(t)-d_{0}(t)\|_{L^{2}}^{2}\geq C\int_{0}^{T}\left\|\nabla^{\bar{g}_{0}(t)} \left(d(t)-d_{0}(t)\right)\right\|_{L^{2}}^{2}\mathrm{d}t. \tag{47}\]

**Proof** Based on Equation (44), we know

\[\frac{\partial}{\partial t}(d(t)-d_{0})= \Delta(d(t)-d_{0})+\mathrm{Rm}\ast(d(t)-d_{0})\] \[+F_{\bar{g}^{-1}}\ast\nabla^{\bar{g}_{0}}(d(t)-d_{0})\ast\nabla^{ \bar{g}_{0}}(d(t)-d_{0})\] \[+\nabla^{\bar{g}_{0}}\left(G_{\Gamma(\bar{g}_{0})}\ast(d(t)-d_{0} )\ast\nabla^{\bar{g}_{0}}(d(t)-d_{0})\right).\]
Followed by Lemma 31 and Theorem 32, we further obtain

\[\frac{\partial}{\partial t}\|d(t)-d_{0}\|_{L^{2}}^{2}= 2\left(\Delta(d(t)-d_{0})+\operatorname{Rm}\ast(d(t)-d_{0}),d(t)-d _{0}\right)_{L^{2}}\] \[+\left(F_{\bar{g}^{-1}}\ast\nabla^{\bar{g}_{0}}(d(t)-d_{0})\ast \nabla^{\bar{g}_{0}}(d(t)-d_{0}),d(t)-d_{0}\right)_{L^{2}}\] \[+\left(\nabla^{\bar{g}_{0}}\left(G_{\Gamma(\bar{g}_{0})}\ast(d(t) -d_{0})\ast\nabla^{\bar{g}_{0}}(d(t)-d_{0})\right),d(t)-d_{0}\right)_{L^{2}}\] \[+\left(d(t)-d_{0},\frac{\partial}{\partial t}d_{0}(t)\right)_{L^{ 2}}+\int_{\mathcal{M}}\left(d(t)-d_{0}\right)\ast(d(t)-d_{0})\ast\frac{ \partial}{\partial t}d_{0}(t)\mathrm{d}\mu\] \[\leq -2\alpha_{\bar{g}_{0}}\left\|\nabla^{\bar{g}_{0}}\left(d(t)-d_{0 }\right)\right\|_{L^{2}}^{2}\] \[+C\left\|(d(t)-d_{0})\right\|_{L^{\infty}}\left\|\nabla^{\bar{g} _{0}}\left(d(t)-d_{0}\right)\right\|_{L^{2}}^{2}\] \[+\left\|\frac{\partial}{\partial t}d_{0}(t)\right\|_{L^{2}}\left\| d(t)-d_{0}\right\|_{L^{2}}\] \[\leq \left(-2\alpha_{\bar{g}_{0}}+C\cdot\epsilon\right)\left\|\nabla^{ \bar{g}_{0}}\left(d(t)-d_{0}\right)\right\|_{L^{2}}^{2}.\]

Let \(\epsilon\) be a small enough constant that \(-2\alpha_{\bar{g}_{0}}+C\cdot\epsilon<0\) holds, we can find

\[\frac{\partial}{\partial t}\|d(t)-d_{0}\|_{L^{2}}^{2}\leq-C\left\|\nabla^{ \bar{g}_{0}}\left(d(t)-d_{0}\right)\right\|_{L^{2}}^{2}\]

holds.

## Appendix E Proof for the Information Geometry

### Proof for Theorem 10

The LNE divergence can be defined between two nearby points \(\boldsymbol{\xi}\) and \(\boldsymbol{\xi}^{\prime}\), where the first derivative of the LNE divergence w.r.t. \(\boldsymbol{\xi}^{\prime}\) is:

\[\partial_{\boldsymbol{\xi}^{\prime}}D_{LNE}[\boldsymbol{\xi}^{ \prime}:\boldsymbol{\xi}]\] \[=\sum_{i}\left[\partial_{\boldsymbol{\xi}^{\prime}}\frac{1}{\tau^ {2}}\log\cosh(\tau\xi_{i}^{\prime})-\partial_{\boldsymbol{\xi}^{\prime}}\frac {1}{\tau^{2}}\log\cosh(\tau\xi_{i})-\frac{1}{\tau}\partial_{\boldsymbol{\xi}^ {\prime}}(\xi_{i}^{\prime}-\xi_{i})\tanh(\tau\xi_{i})\right]\] \[=\sum_{i}\partial_{\boldsymbol{\xi}^{\prime}}\frac{1}{\tau^{2}} \log\cosh(\tau\xi_{i}^{\prime})-\frac{1}{\tau}\tanh(\tau\boldsymbol{\xi}).\]

The second derivative of the LNE divergence w.r.t. \(\boldsymbol{\xi}^{\prime}\) is:

\[\partial_{\boldsymbol{\xi}^{\prime}}^{2}D_{LNE}[\boldsymbol{\xi}^{\prime}: \boldsymbol{\xi}]=\sum_{i}\partial_{\boldsymbol{\xi}^{\prime}}^{2}\frac{1}{ \tau^{2}}\log\cosh(\tau\xi_{i}^{\prime}).\]
We deduce the Taylor expansion of the LNE divergence at \(\mathbf{\xi}^{\prime}=\mathbf{\xi}\):

\[D_{LNE}[\mathbf{\xi}^{\prime}:\mathbf{\xi}] \approx D_{LNE}[\mathbf{\xi}:\mathbf{\xi}]+\left(\sum_{i}\partial_{\mathbf{\xi} ^{\prime}}\frac{1}{\tau^{2}}\log\cosh(\tau\xi_{i}^{\prime})-\frac{1}{\tau} \tanh(\tau\mathbf{\xi})\right)^{\top}\bigg{|}_{\mathbf{\xi}^{\prime}=\mathbf{\xi}}d\mathbf{\xi}\] \[+\frac{1}{2}d\mathbf{\xi}^{\top}\left(\sum_{i}\partial_{\mathbf{\xi}^{ \prime}}^{2}\frac{1}{\tau^{2}}\log\cosh(\tau\xi_{i}^{\prime})\right)\bigg{|}_{ \mathbf{\xi}^{\prime}=\mathbf{\xi}}d\mathbf{\xi}\] \[=0+0+\frac{1}{2\tau^{2}}d\mathbf{\xi}^{\top}\partial\left[\frac{ \partial\cosh(\tau\mathbf{\xi})}{\cosh(\tau\mathbf{\xi})}\right]d\mathbf{\xi}\] \[=\frac{1}{2\tau^{2}}d\mathbf{\xi}^{\top}\frac{\partial^{2}\cosh(\tau \mathbf{\xi})\cosh(\tau\mathbf{\xi})-\partial\cosh(\tau\mathbf{\xi})\partial\cosh(\tau\mathbf{ \xi})^{\top}}{\cosh^{2}(\tau\mathbf{\xi})}d\mathbf{\xi}\] \[=\frac{1}{2\tau^{2}}d\mathbf{\xi}^{\top}\left(\frac{\partial^{2}\cosh (\tau\mathbf{\xi})}{\cosh(\tau\mathbf{\xi})}-\tau^{2}\left[\frac{\sinh(\tau\mathbf{\xi})}{ \cosh(\tau\mathbf{\xi})}\right]\left[\frac{\sinh(\tau\mathbf{\xi})}{\cosh(\tau\mathbf{\xi})} \right]^{\top}\right)d\mathbf{\xi}\] \[=\frac{1}{2}\sum_{i,j}\left\{\delta_{ij}-\left[\tanh(\tau\mathbf{\xi} )\tanh(\tau\mathbf{\xi})^{\top}\right]_{ij}d\xi_{i}d\xi_{j}\right\}.\]

### Proof for Lemma 11

We would like to know in which direction minimizes the loss function with the constraints of the LNE divergence, so that we do the minimization:

\[d\mathbf{\xi}^{*}=\operatorname*{arg\,min}_{d\mathbf{\xi}\text{ s.t. }D_{LNE}[\mathbf{\xi}:\mathbf{\xi} +d\mathbf{\xi}]=c}L(\mathbf{\xi}+d\mathbf{\xi})\]

where \(c\) is the constant. The loss function descends along the manifold with constant speed, regardless the curvature.

Furthermore, we can write the minimization in Lagrangian form. Combined with Theorem 10, the LNE divergence can be approximated by its second order Taylor expansion. Approximating \(L(\mathbf{\xi}+d\mathbf{\xi})\) with it first order Taylor expansion, we get:

\[d\mathbf{\xi}^{*} =\operatorname*{arg\,min}_{d\mathbf{\xi}}\,L(\mathbf{\xi}+d\mathbf{\xi})+ \lambda\left(D_{LNE}[\mathbf{\xi}:\mathbf{\xi}+d\mathbf{\xi}]-c\right)\] \[\approx\operatorname*{arg\,min}_{d\mathbf{\xi}}\,L(\mathbf{\xi})+ \partial_{\mathbf{\xi}}L(\mathbf{\xi})^{\top}d\mathbf{\xi}+\frac{\lambda}{2}d\mathbf{\xi}^{\top }g(\mathbf{\xi})d\mathbf{\xi}-c\lambda.\]

To solve this minimization, we set its derivative w.r.t. \(d\mathbf{\xi}\) to zero:

\[0 =\frac{\partial}{\partial d\mathbf{\xi}}L(\mathbf{\xi})+\partial_{\mathbf{\xi} }L(\mathbf{\xi})^{\top}d\mathbf{\xi}+\frac{\lambda}{2}d\mathbf{\xi}^{\top}\left[\delta- \tanh(\tau\mathbf{\xi})\tanh(\tau\mathbf{\xi})^{\top}\right]d\mathbf{\xi}-c\lambda\] \[=\partial_{\mathbf{\xi}}L(\mathbf{\xi})+\lambda\left[\delta-\tanh(\tau\bm {\xi})\tanh(\tau\mathbf{\xi})^{\top}\right]d\mathbf{\xi}\] \[d\mathbf{\xi} =-\frac{1}{\lambda}\left[\delta-\tanh(\tau\mathbf{\xi})\tanh(\tau\mathbf{ \xi})^{\top}\right]^{-1}\partial_{\mathbf{\xi}}L(\mathbf{\xi})\]

where a constant factor \(1/\lambda\) can be absorbed into learning rate. Therefore, we get the optimal descent direction, i.e., the opposite direction of gradient, which takes into account the local curvature defined by \(\left[\delta-\tanh(\tau\mathbf{\xi})\tanh(\tau\mathbf{\xi})^{\top}\right]^{-1}\).